diff options
Diffstat (limited to 'lib/tests/test_chunking.py')
| -rw-r--r-- | lib/tests/test_chunking.py | 250 |
1 files changed, 250 insertions, 0 deletions
diff --git a/lib/tests/test_chunking.py b/lib/tests/test_chunking.py new file mode 100644 index 0000000..f00c04e --- /dev/null +++ b/lib/tests/test_chunking.py @@ -0,0 +1,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) |