fix alg closure bugs
2 changed files+42−34
python/ephys_compression_tests/algorithms/ans/__init__.pymodified+35−31View file
@@ -273,9 +273,9 @@ for a in algorithm_dicts_base:
273273 return reconstructed
274274 algorithm_dicts.append({
275275 "name": a["name"] + f"-ar{order}",
276- "version": a["version"],
277- "encode": encode0_ar_lossy,
278- "decode": decode0_ar_lossy,
276+ "version": "2",
277+ "encode": encode0_ar,
278+ "decode": decode0_ar,
279279 "description": a["description"] + f" with auto-regressive prediction encoding of order {order}",
280280 "tags": a["tags"] + [f"ar{order}"],
281281 "source_file": a["source_file"],
@@ -285,36 +285,40 @@ for a in algorithm_dicts_base:
285285 # Add lossy ar
286286 for ar_order in [2, 8]:
287287 for tolerance in [1, 2, 3, 4, 5]:
288- def encode0_ar_lossy(x: np.ndarray, tolerance=tolerance, order=ar_order) -> bytes:
289- assert x.ndim == 2, f"Input array must be 2D (timepoints x channels)"
290- coeffs, residuals, initial_values = encode_ar_lossy(x, order=order, step=tolerance * 2 + 1)
291- # coeffs: (n_channels x order), residuals: (n_timepoints-order x n_channels), initial_values: (order x n_channels)
292- encoded_residuals = ans_encode_0(residuals)
293- coeffs_bytes = coeffs.astype(np.float32).tobytes()
294- initial_values_bytes = initial_values.astype(np.int16).tobytes()
295- return coeffs_bytes + initial_values_bytes + encoded_residuals
296- def decode0_ar_lossy(x: bytes, dtype: str, shape: tuple, order=ar_order) -> np.ndarray:
297- assert len(shape) == 2, f"Shape must be 2D (timepoints x channels)"
298- dtype_np = np.dtype(dtype)
299- n_channels = shape[1]
300- # coeffs is (n_channels x order)
301- num_bytes_coeffs = n_channels * order * np.dtype(np.float32).itemsize
302- coeffs_bytes = x[:num_bytes_coeffs]
303- coeffs = np.frombuffer(coeffs_bytes, dtype=np.float32).reshape((n_channels, order))
304- # initial_values is (order x n_channels)
305- num_bytes_initial_values = order * n_channels * dtype_np.itemsize
306- initial_values_bytes = x[num_bytes_coeffs : num_bytes_coeffs + num_bytes_initial_values]
307- initial_values = np.frombuffer(initial_values_bytes, dtype=dtype_np).reshape((order, n_channels))
308- encoded_residuals = x[num_bytes_coeffs + num_bytes_initial_values :]
309- # residuals is ((shape[0]-order) x n_channels)
310- residuals = ans_decode_0(encoded_residuals, dtype, (shape[0]-order, n_channels))
311- reconstructed = decode_ar(coeffs, residuals, initial_values)
312- return reconstructed
288+ def make_encode_ar_lossy(tolerance=tolerance, order=ar_order):
289+ def encode0_ar_lossy(x: np.ndarray) -> bytes:
290+ assert x.ndim == 2, f"Input array must be 2D (timepoints x channels)"
291+ coeffs, residuals, initial_values = encode_ar_lossy(x, order=order, step=tolerance * 2 + 1)
292+ # coeffs: (n_channels x order), residuals: (n_timepoints-order x n_channels), initial_values: (order x n_channels)
293+ encoded_residuals = ans_encode_0(residuals)
294+ coeffs_bytes = coeffs.astype(np.float32).tobytes()
295+ initial_values_bytes = initial_values.astype(np.int16).tobytes()
296+ return coeffs_bytes + initial_values_bytes + encoded_residuals
297+ return encode0_ar_lossy
298+ def make_decode_ar_lossy(order=ar_order):
299+ def decode0_ar_lossy(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
300+ assert len(shape) == 2, f"Shape must be 2D (timepoints x channels)"
301+ dtype_np = np.dtype(dtype)
302+ n_channels = shape[1]
303+ # coeffs is (n_channels x order)
304+ num_bytes_coeffs = n_channels * order * np.dtype(np.float32).itemsize
305+ coeffs_bytes = x[:num_bytes_coeffs]
306+ coeffs = np.frombuffer(coeffs_bytes, dtype=np.float32).reshape((n_channels, order))
307+ # initial_values is (order x n_channels)
308+ num_bytes_initial_values = order * n_channels * dtype_np.itemsize
309+ initial_values_bytes = x[num_bytes_coeffs : num_bytes_coeffs + num_bytes_initial_values]
310+ initial_values = np.frombuffer(initial_values_bytes, dtype=dtype_np).reshape((order, n_channels))
311+ encoded_residuals = x[num_bytes_coeffs + num_bytes_initial_values :]
312+ # residuals is ((shape[0]-order) x n_channels)
313+ residuals = ans_decode_0(encoded_residuals, dtype, (shape[0]-order, n_channels))
314+ reconstructed = decode_ar(coeffs, residuals, initial_values)
315+ return reconstructed
316+ return decode0_ar_lossy
313317 algorithm_dicts.append({
314318 "name": f"ans-ar{ar_order}-lossy-tol{tolerance}",
315- "version": "1",
316- "encode": encode0_ar_lossy,
317- "decode": decode0_ar_lossy,
319+ "version": "2",
320+ "encode": make_encode_ar_lossy(),
321+ "decode": make_decode_ar_lossy(),
318322 "description": f"ANS with lossy auto-regressive prediction encoding of order {ar_order} and tolerance {tolerance}",
319323 "tags": ["ans", "lossy", f"ar{ar_order}"],
320324 "source_file": SOURCE_FILE,
python/ephys_compression_tests/algorithms/wavpack/__init__.pymodified+7−3View file
@@ -83,11 +83,15 @@ for a in algorithm_dicts_base:
8383
8484 # Add lossy versions
8585 for bps in [3, 4, 5, 6]:
86+ def encode_lossy(x: np.ndarray, bps=bps) -> bytes:
87+ return wavpack_encode(x, bps=bps)
88+ def decode_lossy(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
89+ return wavpack_decode(x, dtype, shape)
8690 algorithm_dicts.append({
8791 "name": f"wavpack-lossy-{bps}",
88- "version": "1",
89- "encode": lambda x: wavpack_encode(x, bps=bps),
90- "decode": lambda x, dtype, shape: wavpack_decode(x, dtype, shape),
92+ "version": "2",
93+ "encode": encode_lossy,
94+ "decode": decode_lossy,
9195 "description": f"WavPack lossy with {bps} bits per sample",
9296 "tags": ["wavpack", "lossy"],
9397 "source_file": SOURCE_FILE,