fix lpc_numba
2 changed files+37−73
python/ephys_compression_tests/algorithms/ans/__init__.pymodified+4−4View file
@@ -7,7 +7,7 @@ from ...types import Algorithm
77 # Adapter functions
88 def encode_lpc(data: np.ndarray, order: int):
99 """Encode using LPC model - adapter for lpc_numba."""
10- coeffs, initial_points = lpc_numba.fit_lpc_model(data, k=order, subsample_factor=100, min_samples=2000)
10+ coeffs, initial_points = lpc_numba.fit_lpc_model(data, k=order, subsample_factor=100, min_samples=1000)
1111 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:, :]
@@ -19,7 +19,7 @@ def encode_lpc(data: np.ndarray, order: int):
1919 def encode_lpc_lossy(data: np.ndarray, order: int, step: int):
2020 """Encode using LPC model with lossy quantization - adapter for lpc_numba."""
2121 # Fit the LPC model
22- coeffs, initial_points = lpc_numba.fit_lpc_model(data, k=order, subsample_factor=100, min_samples=2000)
22+ coeffs, initial_points = lpc_numba.fit_lpc_model(data, k=order, subsample_factor=100, min_samples=1000)
2323
2424 # Compute residuals with quantization
2525 residuals_full = lpc_numba.compute_residuals_lossy(data, coeffs, initial_points, step=step)
@@ -318,7 +318,7 @@ for a in algorithm_dicts_base:
318318 return reconstructed
319319 algorithm_dicts.append({
320320 "name": a["name"] + f"-lpc{order}",
321- "version": a["version"] + f".4",
321+ "version": a["version"] + f".5",
322322 "encode": encode0_lpc,
323323 "decode": decode0_lpc,
324324 "description": a["description"] + f" with auto-regressive prediction encoding of order {order}",
@@ -361,7 +361,7 @@ for lpc_order in [2, 8]:
361361 return decode0_lpc_lossy
362362 algorithm_dicts.append({
363363 "name": f"ans-lpc{lpc_order}-lossy-tol{tolerance}",
364- "version": "13",
364+ "version": "14",
365365 "encode": make_encode_lpc_lossy(),
366366 "decode": make_decode_lpc_lossy(),
367367 "description": f"ANS with lossy linear predictive coding of order {lpc_order} and tolerance {tolerance}",
python/ephys_compression_tests/algorithms/ans/lpc_numba.pymodified+33−69View file
@@ -4,75 +4,25 @@ All operations work with int16 data.
44 """
55
66 import numpy as np
7-from numba import jit, prange
7+from numba import jit, prange, njit
88
99
10-@jit(nopython=True, parallel=False, fastmath=True)
11-def _fit_lpc_model_channel(channel_data: np.ndarray, k: int, subsample_factor: int,
12- min_samples: int) -> np.ndarray:
13- """
14- Fit LPC model for a single channel using least squares with subsampling.
15-
16- Args:
17- channel_data: 1D array for a single channel (int16)
18- k: LPC model order
19- subsample_factor: Use every Nth sample for fitting
20- min_samples: Minimum number of samples to use
21-
22- Returns:
23- coefficients: 1D array of k coefficients (float32)
24- """
25- n = len(channel_data)
26-
27- # Determine subsampling stride
28- max_samples = n - k
29- # Use subsample_factor, but ensure we don't skip so much that we get fewer than min_samples
30- if subsample_factor * min_samples > max_samples:
31- # If subsample_factor would give us too few samples, reduce stride
32- stride = max(1, max_samples // min_samples)
33- else:
34- stride = max(1, subsample_factor)
35-
36- # Number of samples we'll actually use
37- n_samples = (max_samples + stride - 1) // stride
38-
39- # Build design matrix X and target vector y using subsampled data
40- # X has shape (n_samples, k), y has shape (n_samples,)
41- X = np.zeros((n_samples, k), dtype=np.float32)
42- y = np.zeros(n_samples, dtype=np.float32)
10+@njit
11+def _create_design_matrix_channel(data: np.ndarray, order: int, subsample_factor: int = 1) -> tuple[np.ndarray, np.ndarray]:
12+ """Numba-optimized design matrix creation for LPC model (single channel) with optional subsampling."""
13+ n = len(data)
14+ n_samples = (n - order + subsample_factor - 1) // subsample_factor
15+ X_design = np.zeros((n_samples, order))
16+ y_target = np.zeros(n_samples)
4317
4418 sample_idx = 0
45- for t in range(k, n, stride):
46- if sample_idx >= n_samples:
47- break
48- for i in range(k):
49- X[sample_idx, i] = np.float32(channel_data[t - 1 - i])
50- y[sample_idx] = np.float32(channel_data[t])
19+ for i in range(0, n - order, subsample_factor):
20+ for j in range(order):
21+ X_design[sample_idx, j] = data[i + order - j - 1]
22+ y_target[sample_idx] = data[i + order]
5123 sample_idx += 1
5224
53- # Truncate if needed
54- actual_samples = sample_idx
55- if actual_samples < n_samples:
56- X = X[:actual_samples, :]
57- y = y[:actual_samples]
58-
59- # Solve normal equations: X.T @ X @ coef = X.T @ y
60- XtX = np.zeros((k, k), dtype=np.float32)
61- Xty = np.zeros(k, dtype=np.float32)
62-
63- for i in range(k):
64- for j in range(k):
65- for t in range(actual_samples):
66- XtX[i, j] += X[t, i] * X[t, j]
67-
68- for i in range(k):
69- for t in range(actual_samples):
70- Xty[i] += X[t, i] * y[t]
71-
72- # Solve the system using numpy's solver
73- coefficients = np.linalg.solve(XtX, Xty)
74-
75- return coefficients
25+ return X_design[:sample_idx], y_target[:sample_idx]
7626
7727
7828 def fit_lpc_model(data: np.ndarray, k: int, subsample_factor: int = 1,
@@ -91,18 +41,32 @@ def fit_lpc_model(data: np.ndarray, k: int, subsample_factor: int = 1,
9141 initial_points: Array of shape (channels, k) with dtype int16 (first k samples per channel)
9242 """
9343 n_timepoints, n_channels = data.shape
94-
95- if n_timepoints <= k:
96- raise ValueError(f"Need at least {k+1} timepoints for LPC({k}) model")
44+ if k >= n_timepoints:
45+ raise ValueError(f"LPC order {k} must be less than data length {n_timepoints}")
9746
9847 # Store initial k points for each channel
99- initial_points = data[:k, :].T.copy() # (channels, k)
48+ initial_points = data[:k, :].T.copy().astype(np.int16) # Shape: (n_channels, k)
10049
101- # Fit coefficients for each channel
50+ # Adjust subsample_factor if needed to ensure we have at least min_samples
51+ effective_subsample_factor = subsample_factor
52+ n_subsampled = (n_timepoints - k) // effective_subsample_factor
53+ if n_subsampled < min_samples:
54+ # Adjust subsample_factor to meet min_samples requirement
55+ effective_subsample_factor = max(1, (n_timepoints - k) // min_samples)
56+
57+ # Fit LPC model for each channel separately
10258 coefficients = np.zeros((n_channels, k), dtype=np.float32)
10359
10460 for ch in range(n_channels):
105- coefficients[ch, :] = _fit_lpc_model_channel(data[:, ch], k, subsample_factor, min_samples)
61+ # Create design matrix using numba-optimized function
62+ # This subsamples the target points but uses full history for predictors
63+ X_design, y_target = _create_design_matrix_channel(data[:, ch], k, effective_subsample_factor)
64+
65+ # Use faster solve via normal equations: (X^T X) coeffs = X^T y
66+ # This is faster than lstsq for overdetermined systems
67+ XtX = X_design.T @ X_design
68+ Xty = X_design.T @ y_target
69+ coefficients[ch] = np.linalg.solve(XtX, Xty).astype(np.float32)
10670
10771 return coefficients, initial_points
10872