1"""Walk real ephys along the s_* axis by changing the quantization step.
3The filtered ephys variant that benchcompress benchmarks fixes the step at
4v = 0.25 noise units, which puts the noise at 4 quantization steps — the
5high-resolution corner where integer-residual coding is already near optimal.
6Coarser steps move the same recording down the s_* axis into the regime where
7the fractional-phase loss is supposed to bite.
9For each step size: requantize, fit the model, synthesize a surrogate, and
10measure the prediction-based methods on both.
12Usage: python sweep.py [--n N] [--order O] [--zstd]
13"""
14import argparse
15import glob
16import os
17import sys
19import numpy as np
20from scipy.signal import butter, lfilter
22sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
23import codec_suite as cc # noqa: E402
24from fitmodel import Fit # noqa: E402
26CACHE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cache")
27RATE = 30000.0
28STEPS = [0.125, 0.25, 0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0]
31def bandpass(x, lowcut=300.0, highcut=6000.0, rate=RATE):
32 nyq = 0.5 * rate
33 b, a = butter(5, [lowcut / nyq, highcut / nyq], btype="band")
34 return lfilter(b, a, x)
37def highpass(x, lowcut=300.0, rate=RATE):
38 nyq = 0.5 * rate
39 b, a = butter(5, lowcut / nyq, btype="high")
40 return lfilter(b, a, x)
43def noise_units(raw, rate=RATE, do_bandpass=True):
44 """Trace scaled so that the MAD noise level is 1.0, optionally bandpassed
45 first. With do_bandpass=False nothing shapes the spectrum but the
46 acquisition hardware — the test of whether the model needs an explicit
47 filtering step or only a spectrum."""
48 x = np.asarray(raw, dtype=np.float64) - np.median(raw)
49 xf = bandpass(x, rate=rate) if do_bandpass else x
50 ref = xf if do_bandpass else highpass(x, rate=rate)
51 nl = float(np.median(np.abs(ref - np.median(ref))) / 0.6745)
52 return xf / nl
55def main():
56 p = argparse.ArgumentParser()
57 p.add_argument("--n", type=int, default=200_000)
58 p.add_argument("--order", type=int, default=32)
59 p.add_argument("--zstd", action="store_true",
60 help="also measure zstd+delta as the practical baseline")
61 p.add_argument("--nofilter", action="store_true",
62 help="skip the bandpass: sweep the unfiltered trace")
63 args = p.parse_args()
65 for path in sorted(glob.glob(os.path.join(CACHE, "*.raw.npy"))):
66 name = os.path.basename(path)[: -len(".raw.npy")]
67 y = noise_units(np.load(path), do_bandpass=not args.nofilter)[: args.n]
68 shaping = ("unfiltered (acquisition spectrum only)" if args.nofilter
69 else "bandpass 300-6000 Hz")
70 print(f"\n{'=' * 96}\n{name} n = {y.size} "
71 f"({shaping}, noise std = 1.0 before quantization)\n{'=' * 96}")
72 print(f" {'step':>6}{'noise/step':>11}{'s_*':>8}{'G(s_*)':>9}"
73 f"{'| LPC+ANS':>12}{'condG':>8}{'gap':>7}"
74 f"{'| LPC+ANS':>12}{'condG':>8}{'gap':>7}{'| model-real':>13}")
75 print(f" {' ' * 34}{'real':>27}{'model surrogate':>27}")
76 print(f" {'-' * 94}")
77 for v in STEPS:
78 z = np.round(y / v).astype(np.int16)
79 if np.unique(z).size < 3:
80 continue
81 fit = Fit(z, nfft=4096)
82 zs = fit.synthesize(seed=0)
84 def block(zz):
85 nb, h0 = cc.lpc_ans_bytes(zz, args.order)
86 return (8.0 * nb / zz.size,
87 cc.conditional_gaussian_rate(zz, args.order))
89 ra, rc = block(z)
90 ma, mc = block(zs)
91 extra = ""
92 if args.zstd:
93 nz = cc.GENERIC["zstd"](cc.delta(z).tobytes())
94 ns = cc.GENERIC["zstd"](cc.delta(zs).tobytes())
95 extra = (f" zstd+delta real {8.0 * nz / z.size:6.3f} "
96 f"model {8.0 * ns / zs.size:6.3f}")
97 print(f" {v:6.3f}{1.0 / v:11.2f}{fit.s_star:8.3f}"
98 f"{fit.predicted_rate:9.4f}"
99 f"{ra:12.4f}{rc:8.4f}{ra - rc:7.4f}"
100 f"{ma:12.4f}{mc:8.4f}{ma - mc:7.4f}"
101 f"{ma - ra:+13.4f}{extra}")
104if __name__ == "__main__":
105 main()