implement lossy algs
6 changed files+155−23
python/ephys_compression_tests/algorithms/ans/__init__.pymodified+44−13View file
@@ -1,6 +1,6 @@
11 import numpy as np
22 import os
3-from .ar import encode_ar, decode_ar
3+from .ar import encode_ar, encode_ar_lossy, decode_ar
44 from ...types import Algorithm
55
66 SOURCE_FILE = "ans/__init__.py"
@@ -156,14 +156,14 @@ for a in algorithm_dicts_base:
156156
157157 # add delta encoding
158158 for a in algorithm_dicts_base:
159- def encode0(x: np.ndarray, a=a) -> bytes:
159+ def encode0_ar2_lossy(x: np.ndarray, a=a) -> bytes:
160160 x_diff = np.diff(x)
161161 x0 = x[0:1]
162162 encoded_diff = a["encode"](x_diff)
163163 # Store the first value at the start
164164 first_value_bytes = x0.tobytes()
165165 return first_value_bytes + encoded_diff
166- def decode0(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
166+ def decode0_ar2_lossy(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
167167 dtype_np = np.dtype(dtype)
168168 num_bytes_first_value = dtype_np.itemsize
169169 first_value_bytes = x[:num_bytes_first_value]
@@ -177,8 +177,8 @@ for a in algorithm_dicts_base:
177177 algorithm_dicts.append({
178178 "name": a["name"] + "-delta",
179179 "version": a["version"],
180- "encode": encode0,
181- "decode": decode0,
180+ "encode": encode0_ar2_lossy,
181+ "decode": decode0_ar2_lossy,
182182 "description": a["description"] + " with delta encoding",
183183 "tags": a["tags"] + ["delta"],
184184 "source_file": a["source_file"],
@@ -187,7 +187,7 @@ for a in algorithm_dicts_base:
187187
188188 # add delta2 encoding
189189 for a in algorithm_dicts_base:
190- def encode0(x: np.ndarray, a=a) -> bytes:
190+ def encode0_ar2_lossy(x: np.ndarray, a=a) -> bytes:
191191 x_diff = np.diff(np.diff(x))
192192 x0 = x[0:1]
193193 encoded_diff = a["encode"](x_diff)
@@ -195,7 +195,7 @@ for a in algorithm_dicts_base:
195195 first_value_bytes = x0.tobytes()
196196 second_value_bytes = x[1:2].tobytes()
197197 return first_value_bytes + second_value_bytes + encoded_diff
198- def decode0(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
198+ def decode0_ar2_lossy(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
199199 dtype_np = np.dtype(dtype)
200200 num_bytes_first_value = dtype_np.itemsize
201201 first_value_bytes = x[:num_bytes_first_value]
@@ -214,8 +214,8 @@ for a in algorithm_dicts_base:
214214 algorithm_dicts.append({
215215 "name": a["name"] + "-delta2",
216216 "version": a["version"],
217- "encode": encode0,
218- "decode": decode0,
217+ "encode": encode0_ar2_lossy,
218+ "decode": decode0_ar2_lossy,
219219 "description": a["description"] + " with delta2 encoding",
220220 "tags": a["tags"] + ["delta2"],
221221 "source_file": a["source_file"],
@@ -225,13 +225,13 @@ for a in algorithm_dicts_base:
225225 # Add auto-regressive prediction encoding
226226 for a in algorithm_dicts_base:
227227 for order in [2, 8]:
228- def encode0(x: np.ndarray, a=a, order=order) -> bytes:
228+ def encode0_ar2_lossy(x: np.ndarray, a=a, order=order) -> bytes:
229229 coeffs, residuals, initial_values = encode_ar(x, order=order)
230230 encoded_residuals = a["encode"](residuals)
231231 coeffs_bytes = coeffs.astype(np.float32).tobytes()
232232 initial_values_bytes = initial_values.astype(np.int16).tobytes()
233233 return coeffs_bytes + initial_values_bytes + encoded_residuals
234- def decode0(x: bytes, dtype: str, shape: tuple, a=a, order=order) -> np.ndarray:
234+ def decode0_ar2_lossy(x: bytes, dtype: str, shape: tuple, a=a, order=order) -> np.ndarray:
235235 dtype_np = np.dtype(dtype)
236236 num_bytes_coeffs = order * np.dtype(np.float32).itemsize
237237 coeffs_bytes = x[:num_bytes_coeffs]
@@ -247,14 +247,45 @@ for a in algorithm_dicts_base:
247247 algorithm_dicts.append({
248248 "name": a["name"] + f"-ar{order}",
249249 "version": a["version"],
250- "encode": encode0,
251- "decode": decode0,
250+ "encode": encode0_ar2_lossy,
251+ "decode": decode0_ar2_lossy,
252252 "description": a["description"] + f" with auto-regressive prediction encoding of order {order}",
253253 "tags": a["tags"] + [f"ar{order}"],
254254 "source_file": a["source_file"],
255255 "long_description": a["long_description"]
256256 })
257257
258+# Add lossy ar2
259+def encode0_ar2_lossy(x: np.ndarray) -> bytes:
260+ coeffs, residuals, initial_values = encode_ar_lossy(x, order=2, step=2 * 2 + 1)
261+ encoded_residuals = ans_encode_0(residuals)
262+ coeffs_bytes = coeffs.astype(np.float32).tobytes()
263+ initial_values_bytes = initial_values.astype(np.int16).tobytes()
264+ return coeffs_bytes + initial_values_bytes + encoded_residuals
265+def decode0_ar2_lossy(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
266+ dtype_np = np.dtype(dtype)
267+ num_bytes_coeffs = 2 * np.dtype(np.float32).itemsize
268+ coeffs_bytes = x[:num_bytes_coeffs]
269+ coeffs = np.frombuffer(coeffs_bytes, dtype=np.float32)
270+ num_initial_values = len(coeffs)
271+ num_bytes_initial_values = num_initial_values * dtype_np.itemsize
272+ initial_values_bytes = x[num_bytes_coeffs : num_bytes_coeffs + num_bytes_initial_values]
273+ initial_values = np.frombuffer(initial_values_bytes, dtype=dtype_np)
274+ encoded_residuals = x[num_bytes_coeffs + num_bytes_initial_values :]
275+ residuals = ans_decode_0(encoded_residuals, dtype, (shape[0]-num_initial_values,))
276+ reconstructed = decode_ar(coeffs, residuals, initial_values)
277+ return reconstructed.reshape(shape)
278+algorithm_dicts.append({
279+ "name": "ans-ar2-lossy-2",
280+ "version": "1",
281+ "encode": encode0_ar2_lossy,
282+ "decode": decode0_ar2_lossy,
283+ "description": "ANS with lossy auto-regressive prediction encoding of order 2",
284+ "tags": ["ans", "lossy", "ar2"],
285+ "source_file": SOURCE_FILE,
286+ "long_description": LONG_DESCRIPTION
287+})
288+
258289 algorithms = [
259290 Algorithm(**a)
260291 for a in algorithm_dicts
python/ephys_compression_tests/algorithms/ans/ar.pymodified+58−0View file
@@ -13,10 +13,12 @@ def _warmup_numba_functions():
1313 test_coeffs = np.array([0.5, 0.3], dtype=np.float32)
1414 test_residuals = np.array([1, 2, 3, 4], dtype=np.int16)
1515 test_initial = np.array([1, 2], dtype=np.int16)
16+ test_step = 2
1617
1718 # Warmup each numba function
1819 _create_design_matrix(test_data, 2)
1920 _apply_ar_residuals_kernel(test_data, test_coeffs)
21+ _apply_ar_residuals_lossy_kernel(test_data, test_coeffs, test_step)
2022 _decode_ar_kernel(test_coeffs, test_residuals, test_initial)
2123
2224
@@ -82,6 +84,42 @@ def _apply_ar_residuals_kernel(data: np.ndarray, coeffs: np.ndarray) -> np.ndarr
8284
8385 return residuals
8486
87+@njit
88+def _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)
93+
94+ # Pre-allocate reconstructed array for efficiency
95+ reconstructed = np.empty(n, dtype=np.int16)
96+ reconstructed[:order] = data[:order]
97+
98+ # Convert step to float32 for consistent float arithmetic
99+ step_f32 = np.float32(step)
100+
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])
107+
108+ # Round to nearest integer
109+ prediction_int = np.int16(np.round(prediction))
110+
111+ # Compute residual from original data
112+ residual = data[i] - prediction_int
113+
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
117+
118+ # Reconstruct sample using quantized residual for future predictions
119+ reconstructed[i] = prediction_int + quantized_residual
120+
121+ return residuals
122+
85123
86124 def apply_ar_residuals(data: np.ndarray, coeffs: np.ndarray) -> np.ndarray:
87125 """
@@ -100,6 +138,11 @@ def apply_ar_residuals(data: np.ndarray, coeffs: np.ndarray) -> np.ndarray:
100138 return _apply_ar_residuals_kernel(data, coeffs)
101139
102140
141+def 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)
144+
145+
103146 def encode_ar(data: np.ndarray, order: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
104147 """
105148 Encode data using AR model - returns coefficients, residuals, and initial values.
@@ -125,6 +168,21 @@ def encode_ar(data: np.ndarray, order: int) -> Tuple[np.ndarray, np.ndarray, np.
125168
126169 return coeffs, residuals, initial_values
127170
171+def 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)
174+
175+ # Convert coefficients to float32 to match what will be deserialized
176+ coeffs = coeffs.astype(np.float32)
177+
178+ # Compute residuals using float32 coefficients
179+ residuals = apply_ar_residuals_lossy(data, coeffs, step=step)
180+
181+ # Store initial values
182+ initial_values = data[:order]
183+
184+ return coeffs, residuals, initial_values
185+
128186
129187 @njit
130188 def _decode_ar_kernel(coeffs: np.ndarray, residuals: np.ndarray, initial_values: np.ndarray) -> np.ndarray:
python/ephys_compression_tests/algorithms/wavpack/__init__.pymodified+27−2View file
@@ -15,9 +15,12 @@ def _load_long_description():
1515 LONG_DESCRIPTION = _load_long_description()
1616
1717
18-def wavpack_encode(x: np.ndarray) -> bytes:
18+def wavpack_encode(x: np.ndarray, bps: float=None) -> bytes:
1919 from wavpack_numcodecs import WavPack
20- codec = WavPack()
20+ if bps is not None:
21+ codec = WavPack(bps=bps)
22+ else:
23+ codec = WavPack()
2124 encoded = codec.encode(x)
2225 assert isinstance(encoded, bytes)
2326 return encoded
@@ -77,6 +80,28 @@ for a in algorithm_dicts_base:
7780 "long_description": a["long_description"]
7881 })
7982
83+# Add lossy version
84+algorithm_dicts.append({
85+ "name": "wavpack-lossy-3",
86+ "version": "1",
87+ "encode": lambda x: wavpack_encode(x, bps=3),
88+ "decode": lambda x, dtype, shape: wavpack_decode(x, dtype, shape),
89+ "description": "WavPack lossy with 3 bits per sample",
90+ "tags": ["wavpack", "lossy"],
91+ "source_file": SOURCE_FILE,
92+ "long_description": LONG_DESCRIPTION,
93+})
94+algorithm_dicts.append({
95+ "name": "wavpack-lossy-4",
96+ "version": "1",
97+ "encode": lambda x: wavpack_encode(x, bps=4),
98+ "decode": lambda x, dtype, shape: wavpack_decode(x, dtype, shape),
99+ "description": "WavPack lossy with 4 bits per sample",
100+ "tags": ["wavpack", "lossy"],
101+ "source_file": SOURCE_FILE,
102+ "long_description": LONG_DESCRIPTION,
103+})
104+
80105 algorithms = [
81106 Algorithm(**a)
82107 for a in algorithm_dicts
python/ephys_compression_tests/cli.pymodified+5−0View file
@@ -133,6 +133,11 @@ def run(algorithm, dataset, cache_dir, quiet, force):
133133 f"\n Encode speed: {result['encode_mb_per_sec']:.2f} MB/s"
134134 f"\n Decode speed: {result['decode_mb_per_sec']:.2f} MB/s"
135135 )
136+ if result["rmse"] != 0.0 or result["max_error"] != 0.0:
137+ click.echo(
138+ f" RMSE: {result['rmse']:.4f}"
139+ f"\n Max error: {result['max_error']:.4f}"
140+ )
136141
137142
138143 def main():
python/ephys_compression_tests/run_benchmarks/benchmark_timing.pymodified+19−8View file
@@ -47,6 +47,7 @@ def run_compression_benchmark(
4747 encode_fn: Callable,
4848 decode_fn: Callable,
4949 verbose: bool = True,
50+ lossy: bool = False,
5051 ) -> Tuple[Dict[str, Any], bytes]:
5152 """Run compression and decompression benchmarks for an algorithm.
5253
@@ -93,14 +94,22 @@ def run_compression_benchmark(
9394 f"Decompression failed: decoded length {len(decoded)} != original length {len(data)}"
9495 )
9596
96- if not np.array_equal(data, decoded):
97- print(data[:100])
98- print(decoded[:100])
99- for j in range(len(data)):
100- if data[j] != decoded[j]:
101- print(f"Error at index {j}: {data[j]} != {decoded[j]}")
102- break
103- raise ValueError(f"Decompression verification failed for {algorithm_name}")
97+ if not lossy:
98+ if not np.array_equal(data, decoded):
99+ print(data[:100])
100+ print(decoded[:100])
101+ for j in range(len(data)):
102+ if data[j] != decoded[j]:
103+ print(f"Error at index {j}: {data[j]} != {decoded[j]}")
104+ break
105+ raise ValueError(f"Decompression verification failed for {algorithm_name}")
106+ rmse = 0.0
107+ max_error = 0.0
108+ else:
109+ # compute RMSE and max error
110+ rmse = float(np.sqrt(np.mean((data - decoded) ** 2)))
111+ max_error = float(np.max(np.abs(data - decoded)))
112+ print(f" RMSE: {rmse:.4f}, Max error: {max_error:.4f}")
104113
105114 if verbose:
106115 print(" Verification successful!")
@@ -117,6 +126,8 @@ def run_compression_benchmark(
117126 "array_dtype": dtype,
118127 "timestamp": time.time(),
119128 "cache_status": "new",
129+ "rmse": rmse,
130+ "max_error": max_error,
120131 }
121132
122133 return result, encoded
python/ephys_compression_tests/run_benchmarks/run_benchmarks.pymodified+2−0View file
@@ -151,12 +151,14 @@ def run_benchmarks(
151151 print(f" Warning: Failed to upload dataset to memobin: {str(e)}")
152152
153153 # Run the benchmark
154+ lossy = "lossy" in alg_tags
154155 result, encoded = run_compression_benchmark(
155156 data,
156157 alg_name,
157158 algorithm.encode,
158159 algorithm.decode,
159160 verbose,
161+ lossy=lossy
160162 )
161163
162164 # Add metadata to result