/ concept-collection / ephys_compression_tests
concept-collection / ephys_compression_tests
ar -> lpc
Jeremy Magland <jmagland@flatironinstitute.org> committed commit f3fef9e797c2 parent 2ee3596 Browse files
2 changed files+57−57
python/ephys_compression_tests/algorithms/ans/__init__.pymodified+42−42View file
@@ -1,14 +1,14 @@
11 import numpy as np
22 import os
3-from . import ar_numba
3+from . import lpc_numba
44 from ...types import Algorithm
55
66
7-# Adapter functions to match the old ar.py API
8-def encode_ar(data: np.ndarray, order: int):
9- """Encode using AR model - adapter for ar_numba."""
10- coeffs, initial_points = ar_numba.fit_ar_model(data, k=order)
11- residuals_full = ar_numba.compute_residuals(data, coeffs, initial_points)
7+# Adapter functions
8+def encode_lpc(data: np.ndarray, order: int):
9+ """Encode using LPC model - adapter for lpc_numba."""
10+ coeffs, initial_points = lpc_numba.fit_lpc_model(data, k=order)
11+ residuals_full = lpc_numba.compute_residuals(data, coeffs, initial_points)
1212 # Extract residuals excluding the initial points (first 'order' rows)
1313 residuals = residuals_full[order:, :]
1414 # Transpose initial_points to match old API: (order, channels)
@@ -16,13 +16,13 @@ def encode_ar(data: np.ndarray, order: int):
1616 return coeffs, residuals, initial_values
1717
1818
19-def encode_ar_lossy(data: np.ndarray, order: int, step: int):
20- """Encode using AR model with lossy quantization - adapter for ar_numba."""
21- # Fit the AR model
22- coeffs, initial_points = ar_numba.fit_ar_model(data, k=order)
23-
19+def encode_lpc_lossy(data: np.ndarray, order: int, step: int):
20+ """Encode using LPC model with lossy quantization - adapter for lpc_numba."""
21+ # Fit the LPC model
22+ coeffs, initial_points = lpc_numba.fit_lpc_model(data, k=order)
23+
2424 # Compute residuals with quantization
25- residuals_full = ar_numba.compute_residuals_lossy(data, coeffs, initial_points, step=step)
25+ residuals_full = lpc_numba.compute_residuals_lossy(data, coeffs, initial_points, step=step)
2626
2727 # Extract residuals excluding the initial points (first 'order' rows)
2828 residuals = residuals_full[order:, :]
@@ -32,8 +32,8 @@ def encode_ar_lossy(data: np.ndarray, order: int, step: int):
3232 return coeffs, residuals, initial_values
3333
3434
35-def decode_ar(coeffs: np.ndarray, residuals: np.ndarray, initial_values: np.ndarray):
36- """Decode AR encoded data - adapter for ar_numba."""
35+def decode_lpc(coeffs: np.ndarray, residuals: np.ndarray, initial_values: np.ndarray):
36+ """Decode LPC encoded data - adapter for lpc_numba."""
3737 # Transpose initial_values from (order, channels) to (channels, order)
3838 initial_points = initial_values.T
3939
@@ -46,7 +46,7 @@ def decode_ar(coeffs: np.ndarray, residuals: np.ndarray, initial_values: np.ndar
4646 residuals_full[:order, :] = initial_points.T
4747 residuals_full[order:, :] = residuals
4848
49- return ar_numba.reconstruct_from_residuals(residuals_full, coeffs, initial_points)
49+ return lpc_numba.reconstruct_from_residuals(residuals_full, coeffs, initial_points)
5050
5151 SOURCE_FILE = "ans/__init__.py"
5252
@@ -251,7 +251,7 @@ for a in algorithm_dicts_base:
251251
252252 # add delta2 encoding
253253 for a in algorithm_dicts_base:
254- def encode0_ar_lossy(x: np.ndarray, a=a) -> bytes:
254+ def encode0_lpc_lossy(x: np.ndarray, a=a) -> bytes:
255255 assert x.ndim == 2 and x.shape[0] > 2, "Input array must be 2D with more than two timepoints"
256256 x_diff = np.diff(np.diff(x, axis=0), axis=0)
257257 first_timepoint = x[0:1, :].flatten()
@@ -261,7 +261,7 @@ for a in algorithm_dicts_base:
261261 first_timepoint_bytes = first_timepoint.tobytes()
262262 second_timepoint_bytes = second_timepoint.tobytes()
263263 return first_timepoint_bytes + second_timepoint_bytes + encoded_diff
264- def decode0_ar_lossy(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
264+ def decode0_lpc_lossy(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
265265 dtype_np = np.dtype(dtype)
266266 num_bytes_first_timepoint = dtype_np.itemsize * shape[1]
267267 first_timepoint_bytes = x[:num_bytes_first_timepoint]
@@ -280,8 +280,8 @@ for a in algorithm_dicts_base:
280280 algorithm_dicts.append({
281281 "name": a["name"] + "-delta2",
282282 "version": a["version"],
283- "encode": encode0_ar_lossy,
284- "decode": decode0_ar_lossy,
283+ "encode": encode0_lpc_lossy,
284+ "decode": decode0_lpc_lossy,
285285 "description": a["description"] + " with delta2 encoding",
286286 "tags": a["tags"] + ["delta2"],
287287 "source_file": a["source_file"],
@@ -291,15 +291,15 @@ for a in algorithm_dicts_base:
291291 # Add auto-regressive prediction encoding
292292 for a in algorithm_dicts_base:
293293 for order in [2, 8]:
294- def encode0_ar(x: np.ndarray, a=a, order=order) -> bytes:
294+ def encode0_lpc(x: np.ndarray, a=a, order=order) -> bytes:
295295 assert x.ndim == 2 and x.shape[0] > order, f"Input array must be 2D (timepoints x channels) with more than {order} timepoints"
296- coeffs, residuals, initial_values = encode_ar(x, order=order)
296+ coeffs, residuals, initial_values = encode_lpc(x, order=order)
297297 # coeffs: (n_channels x order), residuals: (n_timepoints-order x n_channels), initial_values: (order x n_channels)
298298 encoded_residuals = a["encode"](residuals)
299299 coeffs_bytes = coeffs.astype(np.float32).tobytes()
300300 initial_values_bytes = initial_values.astype(np.int16).tobytes()
301301 return coeffs_bytes + initial_values_bytes + encoded_residuals
302- def decode0_ar(x: bytes, dtype: str, shape: tuple, a=a, order=order) -> np.ndarray:
302+ def decode0_lpc(x: bytes, dtype: str, shape: tuple, a=a, order=order) -> np.ndarray:
303303 assert len(shape) == 2, f"Shape must be 2D (timepoints x channels)"
304304 dtype_np = np.dtype(dtype)
305305 n_channels = shape[1]
@@ -314,34 +314,34 @@ for a in algorithm_dicts_base:
314314 encoded_residuals = x[num_bytes_coeffs + num_bytes_initial_values :]
315315 # residuals is ((shape[0]-order) x n_channels)
316316 residuals = a["decode"](encoded_residuals, dtype, (shape[0]-order, n_channels))
317- reconstructed = decode_ar(coeffs, residuals, initial_values)
317+ reconstructed = decode_lpc(coeffs, residuals, initial_values)
318318 return reconstructed
319319 algorithm_dicts.append({
320- "name": a["name"] + f"-ar{order}",
320+ "name": a["name"] + f"-lpc{order}",
321321 "version": a["version"] + f".3",
322- "encode": encode0_ar,
323- "decode": decode0_ar,
322+ "encode": encode0_lpc,
323+ "decode": decode0_lpc,
324324 "description": a["description"] + f" with auto-regressive prediction encoding of order {order}",
325- "tags": a["tags"] + [f"ar{order}"],
325+ "tags": a["tags"] + [f"lpc{order}"],
326326 "source_file": a["source_file"],
327327 "long_description": a["long_description"]
328328 })
329329
330-# Add lossy ar
331-for ar_order in [2, 8]:
330+# Add lossy lpc
331+for lpc_order in [2, 8]:
332332 for tolerance in [1, 2, 3, 4, 6, 8, 12, 16]:
333- def make_encode_ar_lossy(tolerance=tolerance, order=ar_order):
334- def encode0_ar_lossy(x: np.ndarray) -> bytes:
333+ def make_encode_lpc_lossy(tolerance=tolerance, order=lpc_order):
334+ def encode0_lpc_lossy(x: np.ndarray) -> bytes:
335335 assert x.ndim == 2, f"Input array must be 2D (timepoints x channels)"
336- coeffs, residuals, initial_values = encode_ar_lossy(x, order=order, step=tolerance * 2 + 1)
336+ coeffs, residuals, initial_values = encode_lpc_lossy(x, order=order, step=tolerance * 2 + 1)
337337 # coeffs: (n_channels x order), residuals: (n_timepoints-order x n_channels), initial_values: (order x n_channels)
338338 encoded_residuals = ans_encode_0(residuals)
339339 coeffs_bytes = coeffs.astype(np.float32).tobytes()
340340 initial_values_bytes = initial_values.astype(np.int16).tobytes()
341341 return coeffs_bytes + initial_values_bytes + encoded_residuals
342- return encode0_ar_lossy
343- def make_decode_ar_lossy(order=ar_order):
344- def decode0_ar_lossy(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
342+ return encode0_lpc_lossy
343+ def make_decode_lpc_lossy(order=lpc_order):
344+ def decode0_lpc_lossy(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
345345 assert len(shape) == 2, f"Shape must be 2D (timepoints x channels)"
346346 dtype_np = np.dtype(dtype)
347347 n_channels = shape[1]
@@ -356,16 +356,16 @@ for ar_order in [2, 8]:
356356 encoded_residuals = x[num_bytes_coeffs + num_bytes_initial_values :]
357357 # residuals is ((shape[0]-order) x n_channels)
358358 residuals = ans_decode_0(encoded_residuals, dtype, (shape[0]-order, n_channels))
359- reconstructed = decode_ar(coeffs, residuals, initial_values)
359+ reconstructed = decode_lpc(coeffs, residuals, initial_values)
360360 return reconstructed
361- return decode0_ar_lossy
361+ return decode0_lpc_lossy
362362 algorithm_dicts.append({
363- "name": f"ans-ar{ar_order}-lossy-tol{tolerance}",
363+ "name": f"ans-lpc{lpc_order}-lossy-tol{tolerance}",
364364 "version": "12",
365- "encode": make_encode_ar_lossy(),
366- "decode": make_decode_ar_lossy(),
367- "description": f"ANS with lossy auto-regressive prediction encoding of order {ar_order} and tolerance {tolerance}",
368- "tags": ["ans", "lossy", f"ar{ar_order}"],
365+ "encode": make_encode_lpc_lossy(),
366+ "decode": make_decode_lpc_lossy(),
367+ "description": f"ANS with lossy linear predictive coding of order {lpc_order} and tolerance {tolerance}",
368+ "tags": ["ans", "lossy", f"lpc{lpc_order}"],
369369 "source_file": SOURCE_FILE,
370370 "long_description": LONG_DESCRIPTION
371371 })
python/ephys_compression_tests/algorithms/ans/ar_numba.py →python/ephys_compression_tests/algorithms/ans/lpc_numba.pyrenamed+15−15View file
@@ -1,5 +1,5 @@
11 """
2-Numba-accelerated implementation of autoregressive model operations.
2+Numba-accelerated implementation of linear predictive coding (LPC) model operations.
33 All operations work with int16 data.
44 """
55
@@ -8,14 +8,14 @@ from numba import jit, prange
88
99
1010 @jit(nopython=True, parallel=False, fastmath=True)
11-def _fit_ar_model_channel(channel_data: np.ndarray, k: int, subsample_factor: int,
11+def _fit_lpc_model_channel(channel_data: np.ndarray, k: int, subsample_factor: int,
1212 min_samples: int) -> np.ndarray:
1313 """
14- Fit AR model for a single channel using least squares with subsampling.
14+ Fit LPC model for a single channel using least squares with subsampling.
1515
1616 Args:
1717 channel_data: 1D array for a single channel (int16)
18- k: AR model order
18+ k: LPC model order
1919 subsample_factor: Use every Nth sample for fitting
2020 min_samples: Minimum number of samples to use
2121
@@ -75,14 +75,14 @@ def _fit_ar_model_channel(channel_data: np.ndarray, k: int, subsample_factor: in
7575 return coefficients
7676
7777
78-def fit_ar_model(data: np.ndarray, k: int, subsample_factor: int = 1,
78+def fit_lpc_model(data: np.ndarray, k: int, subsample_factor: int = 1,
7979 min_samples: int = 1000) -> tuple[np.ndarray, np.ndarray]:
8080 """
81- Fit an autoregressive model of order k to multi-channel time series data.
82-
81+ Fit a linear predictive coding (LPC) model of order k to multi-channel time series data.
82+
8383 Args:
8484 data: 2D array of shape (timepoints, channels) with dtype int16
85- k: Order of the autoregressive model
85+ k: Order of the LPC model
8686 subsample_factor: Use every Nth sample for fitting (default: 1)
8787 min_samples: Minimum number of samples to use for fitting (default: 1000)
8888
@@ -93,7 +93,7 @@ def fit_ar_model(data: np.ndarray, k: int, subsample_factor: int = 1,
9393 n_timepoints, n_channels = data.shape
9494
9595 if n_timepoints <= k:
96- raise ValueError(f"Need at least {k+1} timepoints for AR({k}) model")
96+ raise ValueError(f"Need at least {k+1} timepoints for LPC({k}) model")
9797
9898 # Store initial k points for each channel
9999 initial_points = data[:k, :].T.copy() # (channels, k)
@@ -102,7 +102,7 @@ def fit_ar_model(data: np.ndarray, k: int, subsample_factor: int = 1,
102102 coefficients = np.zeros((n_channels, k), dtype=np.float32)
103103
104104 for ch in range(n_channels):
105- coefficients[ch, :] = _fit_ar_model_channel(data[:, ch], k, subsample_factor, min_samples)
105+ coefficients[ch, :] = _fit_lpc_model_channel(data[:, ch], k, subsample_factor, min_samples)
106106
107107 return coefficients, initial_points
108108
@@ -179,7 +179,7 @@ def _compute_residuals_lossy_jit(data: np.ndarray, coefficients: np.ndarray,
179179 def compute_residuals(data: np.ndarray, coefficients: np.ndarray,
180180 initial_points: np.ndarray) -> np.ndarray:
181181 """
182- Compute residuals given data and AR model coefficients.
182+ Compute residuals given data and LPC model coefficients.
183183
184184 Args:
185185 data: 2D array of shape (timepoints, channels) with dtype int16
@@ -242,7 +242,7 @@ def _reconstruct_from_residuals_jit(residuals: np.ndarray, coefficients: np.ndar
242242 def reconstruct_from_residuals(residuals: np.ndarray, coefficients: np.ndarray,
243243 initial_points: np.ndarray) -> np.ndarray:
244244 """
245- Reconstruct original data from residuals and AR model coefficients.
245+ Reconstruct original data from residuals and LPC model coefficients.
246246
247247 Args:
248248 residuals: 2D array of shape (timepoints, channels) with dtype int16
@@ -262,14 +262,14 @@ def warmup(n_channels: int = 10, k: int = 10):
262262
263263 Args:
264264 n_channels: Number of channels for warmup data
265- k: AR model order for warmup
265+ k: LPC model order for warmup
266266 """
267267 print("Warming up JIT...", end="", flush=True)
268268 # Create small warmup data
269269 warmup_data = np.random.randint(-1000, 1000, size=(1000, n_channels), dtype=np.int16)
270270
271- # Warm up fit_ar_model
272- coefficients, initial_points = fit_ar_model(warmup_data, k)
271+ # Warm up fit_lpc_model
272+ coefficients, initial_points = fit_lpc_model(warmup_data, k)
273273
274274 # Warm up compute_residuals
275275 residuals = compute_residuals(warmup_data, coefficients, initial_points)