aboutsummaryrefslogtreecommitdiff
path: root/lib/tests/test_chunking.py
blob: f00c04e63d2d3b804f34f7c35706cc0c62119633 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
import numpy as np
import pytest

from producer.engines.base import blend
from producer.engines.chunking import apply_chunked, plan_chunks, stitch


def test_plan_chunks_basic():
    assert plan_chunks(0, 5, 2) == []
    assert plan_chunks(4, 10, 2) == [(0, 4)]
    assert plan_chunks(10, 5, 0) == [(0, 5), (5, 10)]
    assert plan_chunks(10, 5, 2) == [(0, 5), (3, 8), (6, 10)]
    assert plan_chunks(7, 5, 2) == [(0, 5), (3, 7)]


def test_plan_chunks_last_span_covers_tail():
    spans = plan_chunks(11, 5, 4)
    assert spans[-1][1] == 11
    assert spans[-1][1] - spans[-1][0] >= 5


def test_plan_chunks_rejects_overlap_ge_chunk():
    with pytest.raises(ValueError):
        plan_chunks(100, 5, 5)
    with pytest.raises(ValueError):
        plan_chunks(100, 5, 6)


def test_stitch_identity_pieces_reconstruct_signal():
    rng = np.random.default_rng(0)
    x = rng.standard_normal(10_000).astype(np.float32)
    spans = plan_chunks(x.size, 2500, 300)
    y = stitch(spans, x.size, (x[a:b] for a, b in spans))
    assert y.dtype == np.float32
    assert y.shape == x.shape
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_stitch_single_span_passthrough():
    x = np.arange(100, dtype=np.float32)
    y = stitch([(0, 100)], 100, [x])
    np.testing.assert_array_equal(y, x)


def test_stitch_crossfade_of_complementary_pieces():
    # piece A is silence, piece B is a full-scale ramp; the crossfade zone
    # must be a smooth blend, not a jump.
    spans = [(0, 10), (5, 15)]
    pieces = [np.zeros(10, dtype=np.float32), np.ones(10, dtype=np.float32)]
    y = stitch(spans, 15, pieces)
    mid = y[7]  # 50% through the overlap
    assert 0.4 < mid < 0.6
    assert y[0] == 0.0
    assert y[14] == 1.0


def test_stitch_pads_short_pieces():
    spans = [(0, 10), (5, 15)]
    pieces = [np.arange(10, dtype=np.float32), np.arange(8, dtype=np.float32)]
    y = stitch(spans, 15, pieces)
    assert y.shape == (15,)
    assert np.all(np.isfinite(y))


def test_apply_chunked_identity_matches_whole_file():
    rng = np.random.default_rng(1)
    x = rng.standard_normal(50_000).astype(np.float32)
    sr = 8000
    calls = []

    def fn(chunk):
        calls.append(chunk.size)
        return chunk

    y = apply_chunked(x, sr, 7.0, 0.05, fn)
    assert calls == [50_000]
    np.testing.assert_array_equal(y, x)

    calls.clear()
    y = apply_chunked(x, sr, 0.5, 0.05, fn)
    assert len(calls) >= 10
    assert all(c <= 4000 for c in calls)
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_apply_chunked_zero_disables_chunking():
    x = np.zeros(10, dtype=np.float32)
    calls = []

    def fn(chunk):
        calls.append(chunk.size)
        return chunk

    apply_chunked(x, 8000, 0.0, 0.5, fn)
    assert calls == [10]


def test_apply_chunked_context_reconstructs_identity():
    rng = np.random.default_rng(3)
    x = rng.standard_normal(30_000).astype(np.float32)
    y = apply_chunked(x, 8000, 1.5, 0.25, lambda c: c, context_s=0.5)
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_apply_chunked_context_feeds_padded_chunks():
    x = np.ones(5_000, dtype=np.float32)
    sizes = []

    def fn(chunk):
        sizes.append(chunk.size)
        return chunk

    apply_chunked(x, 1000, 2.0, 0.0, fn, context_s=0.5)
    # plan_chunks(5000, 2000, 0) -> (0,2000), (2000,4000), (4000,5000)
    assert sizes == [2500, 3000, 1500]


def test_apply_chunked_context_trims_cold_start_artifacts():
    x = np.ones(6_000, dtype=np.float32)

    def fn(chunk):
        out = chunk.copy()
        out[0] = 0.0  # cold-start artifact at the start of every model call
        return out

    broken = apply_chunked(x, 1000, 2.0, 0.0, fn)
    assert int(np.sum(broken == 0.0)) > 1

    y = apply_chunked(x, 1000, 2.0, 0.0, fn, context_s=0.5)
    assert y[0] == 0.0  # only the true file start stays degraded
    assert np.all(y[1:] == 1.0)


def test_apply_chunked_context_survives_oom_retry():
    x = np.ones(8_000, dtype=np.float32)
    attempts = []

    def fn(chunk):
        attempts.append(chunk.size)
        if chunk.size > 3000:
            raise MemoryError("simulated oom")
        out = chunk.copy()
        out[0] = 0.0
        return out

    y = apply_chunked(x, 1000, 6.0, 0.0, fn, min_chunk_s=1.0, context_s=0.5)
    assert max(attempts) > 3000
    assert y[0] == 0.0
    assert np.all(y[1:] == 1.0)


def test_apply_chunked_retries_smaller_on_oom():
    x = np.random.default_rng(2).standard_normal(40_000).astype(np.float32)
    sr = 8000
    attempts = []

    def fn(chunk):
        attempts.append(chunk.size)
        if chunk.size > 3000:
            raise MemoryError("simulated oom")
        return chunk

    y = apply_chunked(x, sr, 5.0, 0.0, fn, min_chunk_s=0.25)
    assert attempts[0] == 40_000
    assert attempts[-1] == 2500
    assert max(a for a in attempts if a <= 3000) == 2500
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_apply_chunked_oom_gives_up_at_min_chunk():
    x = np.zeros(40_000, dtype=np.float32)

    def fn(chunk):
        raise MemoryError("always")

    with pytest.raises(MemoryError):
        apply_chunked(x, 8000, 1.0, 0.0, fn, min_chunk_s=0.5)


def test_apply_chunked_progress_callback():
    x = np.random.default_rng(3).standard_normal(20_000).astype(np.float32)
    sr = 8000
    seen: list[tuple[int, int]] = []
    y = apply_chunked(x, sr, 0.5, 0.0, lambda c: c, on_progress=lambda d, t: seen.append((d, t)))
    total = len(plan_chunks(x.size, round(0.5 * sr), 0))
    assert seen[0] == (0, total)
    assert seen[-1] == (total, total)
    assert [d for d, _t in seen[1:]] == list(range(1, total + 1))
    np.testing.assert_allclose(y, x, atol=1e-6)


def test_apply_chunked_progress_resets_after_oom_retry():
    x = np.random.default_rng(4).standard_normal(40_000).astype(np.float32)
    sr = 8000
    seen: list[tuple[int, int]] = []

    def fn(chunk):
        if chunk.size > 3000:
            raise MemoryError("simulated oom")
        return chunk

    y = apply_chunked(
        x, sr, 5.0, 0.0, fn, min_chunk_s=0.25, on_progress=lambda d, t: seen.append((d, t))
    )
    assert seen[0] == (0, 1)  # first attempt is one whole-file chunk
    done, total = seen[-1]
    assert done == total and total > 2  # retried into smaller chunks
    np.testing.assert_allclose(y, x, atol=1e-5)


def test_apply_chunked_progress_absent_when_disabled():
    x = np.zeros(10_000, dtype=np.float32)
    y = apply_chunked(x, 8000, 0.5, 0.0, lambda c: c, on_progress=None)
    assert y.shape == x.shape


def test_apply_chunked_whole_file_mode_falls_back_on_oom():
    x = np.random.default_rng(5).standard_normal(10_000).astype(np.float32)
    sr = 1000

    def fn(chunk):
        if chunk.size > 1500:
            raise MemoryError("simulated oom")
        return chunk * 2

    y = apply_chunked(x, sr, 0.0, 0.0, fn, min_chunk_s=0.5)
    np.testing.assert_allclose(y, x * 2, atol=1e-6)


def test_blend_float32_no_float64_temporaries():
    x = np.full(1000, 0.25, dtype=np.float32)
    y = np.full(1000, 0.75, dtype=np.float32)
    out = blend(x, y, 0.5)
    assert out.dtype == np.float32
    np.testing.assert_allclose(out, 0.5, atol=1e-7)
    assert blend(x, y, 1.0) is y
    np.testing.assert_allclose(blend(x, y, 0.0), x, atol=1e-7)


def test_blend_realigns_resample_drift():
    # a 48k -> 16k -> 48k round trip can come back a sample or two long
    # (resample_poly emits ceil(n * up/down) per hop); blend must cope
    x = np.ones(5, dtype=np.float32)
    long_y = np.full(7, 0.5, dtype=np.float32)
    np.testing.assert_allclose(blend(x, long_y, 0.5), np.full(5, 0.75), atol=1e-7)

    short_y = np.full(4, 0.5, dtype=np.float32)
    out = blend(x, short_y, 0.5)
    assert out.shape == x.shape
    np.testing.assert_allclose(out, [0.75, 0.75, 0.75, 0.75, 0.5], atol=1e-7)