1"""Auto-regressive model utilities for ANS compression."""
3import numpy as np
4from typing import Tuple
5from numba import njit
8def _warmup_numba_functions():
9 """Warmup numba JIT compilation with small test data."""
10 print("Warming up numba functions for AR model...")
11 # Create small test data
12 test_data = np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int16)
13 test_coeffs = np.array([0.5, 0.3], dtype=np.float32)
14 test_residuals = np.array([1, 2, 3, 4], dtype=np.int16)
15 test_initial = np.array([1, 2], dtype=np.int16)
16 test_step = 2
18 # Warmup each numba function
19 _create_design_matrix(test_data, 2)
20 _apply_ar_residuals_kernel(test_data, test_coeffs)
21 _apply_ar_residuals_lossy_kernel(test_data, test_coeffs, test_step)
22 _decode_ar_kernel(test_coeffs, test_residuals, test_initial)
26def _create_design_matrix(data: np.ndarray, order: int) -> Tuple[np.ndarray, np.ndarray]:
27 """Numba-optimized design matrix creation for AR model."""
28 n = len(data)
29 X_design = np.zeros((n - order, order))
30 y_target = data[order:]
32 for i in range(n - order):
33 for j in range(order):
34 X_design[i, j] = data[i + order - j - 1]
36 return X_design, y_target
39def fit_ar_model(data: np.ndarray, order: int) -> np.ndarray:
40 """
41 Fit an autoregressive model of given order using least squares.
43 Args:
44 data: Input data array
45 order: AR model order
47 Returns:
48 AR coefficients as numpy array
49 """
50 n = len(data)
51 if order >= n:
52 raise ValueError(f"AR order {order} must be less than data length {n}")
54 # Create design matrix using numba-optimized function
55 X_design, y_target = _create_design_matrix(data, order)
57 # Use faster solve via normal equations: (X^T X) coeffs = X^T y
58 # This is faster than lstsq for overdetermined systems
59 XtX = X_design.T @ X_design
60 Xty = X_design.T @ y_target
61 coeffs = np.linalg.solve(XtX, Xty)
63 return coeffs
67def _apply_ar_residuals_kernel(data: np.ndarray, coeffs: np.ndarray) -> np.ndarray:
68 """Numba-optimized kernel for computing AR residuals."""
69 order = len(coeffs)
70 n = len(data)
71 residuals = np.empty(n - order, dtype=data.dtype)
73 for i in range(order, n):
74 # Predict using previous 'order' samples
75 # Use float32 accumulation
76 prediction = np.float32(0.0)
77 for j in range(order):
78 prediction += coeffs[j] * np.float32(data[i - j - 1])
80 # Round to nearest integer using numpy's round (banker's rounding)
81 prediction_int = np.int16(np.round(prediction))
82 residual = data[i] - prediction_int
83 residuals[i - order] = residual
85 return residuals
88def _apply_ar_residuals_lossy_kernel(data: np.ndarray, coeffs: np.ndarray, step: int) -> np.ndarray:
89 """Numba-optimized kernel for computing AR residuals with lossy quantization."""
90 order = len(coeffs)
91 n = len(data)
92 residuals = np.empty(n - order, dtype=data.dtype)
94 # Pre-allocate reconstructed array for efficiency
95 reconstructed = np.empty(n, dtype=np.int16)
96 reconstructed[:order] = data[:order]
98 # Convert step to float32 for consistent float arithmetic
99 step_f32 = np.float32(step)
101 for i in range(order, n):
102 # Predict using previous 'order' samples from reconstructed data
103 # Use float32 accumulation
104 prediction = np.float32(0.0)
105 for j in range(order):
106 prediction += coeffs[j] * np.float32(reconstructed[i - j - 1])
108 # Round to nearest integer
109 prediction_int = np.int16(np.round(prediction))
111 # Compute residual from original data
112 residual = data[i] - prediction_int
114 # Quantize residual to nearest multiple of step
115 quantized_residual = np.int16(np.round(np.float32(residual) / step_f32) * step_f32)
116 residuals[i - order] = quantized_residual
118 # Reconstruct sample using quantized residual for future predictions
119 reconstructed[i] = prediction_int + quantized_residual
121 return residuals
124def apply_ar_residuals(data: np.ndarray, coeffs: np.ndarray) -> np.ndarray:
125 """
126 Apply AR model with given coefficients and return residuals.
128 Args:
129 data: Input data array
130 coeffs: AR coefficients
132 Returns:
133 Residuals array
134 """
135 # Ensure coeffs is float32
136 coeffs = np.array(coeffs, dtype=np.float32)
138 return _apply_ar_residuals_kernel(data, coeffs)
141def apply_ar_residuals_lossy(data: np.ndarray, coeffs: np.ndarray, step: int) -> np.ndarray:
142 coeffs = np.array(coeffs, dtype=np.float32)
143 return _apply_ar_residuals_lossy_kernel(data, coeffs, step)
146def encode_ar(data: np.ndarray, order: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
147 """
148 Encode data using AR model - returns coefficients, residuals, and initial values.
150 Args:
151 data: Input data array (int16)
152 order: AR model order
154 Returns:
155 Tuple of (coefficients, residuals, initial_values)
156 """
157 # Fit AR model
158 coeffs = fit_ar_model(data, order)
160 # Convert coefficients to float32 to match what will be deserialized
161 coeffs = coeffs.astype(np.float32)
163 # Compute residuals using float32 coefficients
164 residuals = apply_ar_residuals(data, coeffs)
166 # Store initial values
167 initial_values = data[:order]
169 return coeffs, residuals, initial_values
171def encode_ar_lossy(data: np.ndarray, order: int, step: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
172 # Fit AR model
173 coeffs = fit_ar_model(data, order)
175 # Convert coefficients to float32 to match what will be deserialized
176 coeffs = coeffs.astype(np.float32)
178 # Compute residuals using float32 coefficients
179 residuals = apply_ar_residuals_lossy(data, coeffs, step=step)
181 # Store initial values
182 initial_values = data[:order]
184 return coeffs, residuals, initial_values
188def _decode_ar_kernel(coeffs: np.ndarray, residuals: np.ndarray, initial_values: np.ndarray) -> np.ndarray:
189 """Numba-optimized kernel for AR decoding."""
190 order = len(coeffs)
191 n = len(residuals) + order
192 reconstructed = np.empty(n, dtype=np.int16)
193 reconstructed[:order] = initial_values
195 for i in range(order, n):
196 # Predict using AR model
197 # Use float32 accumulation
198 prediction = np.float32(0.0)
199 for j in range(order):
200 prediction += coeffs[j] * np.float32(reconstructed[i - j - 1])
202 # Round to nearest integer using numpy's round (banker's rounding)
203 prediction_int = np.int16(np.round(prediction))
205 # Add residual
206 reconstructed[i] = prediction_int + residuals[i - order]
208 return reconstructed
211def decode_ar(coeffs: np.ndarray, residuals: np.ndarray, initial_values: np.ndarray) -> np.ndarray:
212 """
213 Decode AR encoded data.
215 Args:
216 coeffs: AR coefficients (float32)
217 residuals: Residuals array
218 initial_values: Initial values (first 'order' samples)
220 Returns:
221 Reconstructed data array
222 """
223 # Ensure coeffs is float32
224 coeffs = np.array(coeffs, dtype=np.float32)
226 return _decode_ar_kernel(coeffs, residuals, initial_values)
229# Warmup numba functions on module import
230_warmup_numba_functions()