/ concept-collection / ephys_compression_tests
Sign in
concept-collection / ephys_compression_tests
support multi-channel
Jeremy Magland <jmagland@flatironinstitute.org> committed commit 68647938bceb parent 721e8d8 Browse files
11 changed files+407−130
prepare_datasets/.gitignoreadded+2−0View file
@@ -0,0 +1,2 @@
1+*.si
2+*.npy
\ No newline at end of file
prepare_datasets/aind/__pycache__/s3_utils.cpython-312.pycadded+0−0View file
Binary file not shown.
prepare_datasets/aind/prepare_aind_compression_np2_probeB.pyadded+34−0View file
@@ -0,0 +1,34 @@
1+import os
2+import spikeinterface as si
3+import numpy as np
4+from s3_utils import download_s3_folder
5+
6+s3_folder_name = "s3://aind-benchmark-data/ephys-compression/aind-np2/612962_2022-04-13_19-18-04_ProbeB"
7+local_folder_name = "612962_2022-04-13_19-18-04_ProbeB.si"
8+
9+if not os.path.exists(local_folder_name):
10+ download_s3_folder(s3_folder_name, local_folder_name)
11+
12+recording = si.load(
13+ local_folder_name
14+)
15+
16+channel_ids = [
17+ 'CH101',
18+ 'CH102',
19+ 'CH103',
20+ 'CH104',
21+ 'CH105',
22+ 'CH106',
23+ 'CH107',
24+ 'CH108',
25+ 'CH109',
26+ 'CH110'
27+]
28+
29+fname = f'aind_compression_np2_probeB_ch101-110.raw.npy'
30+if not os.path.exists(fname):
31+ print(f'Writing {fname}...')
32+ X = recording.get_traces(channel_ids=channel_ids, start_frame=30000, end_frame=30000 + 30000 * 10)
33+ print(f'X.shape = {X.shape}')
34+ np.save(fname, X)
prepare_datasets/aind/s3_utils.pyadded+130−0View file
@@ -0,0 +1,130 @@
1+"""
2+Utility functions for downloading from S3 public buckets
3+"""
4+import boto3
5+from botocore import UNSIGNED
6+from botocore.config import Config
7+import os
8+from pathlib import Path
9+import sys
10+
11+
12+class ProgressCallback:
13+ """Callback to show download progress"""
14+ def __init__(self, filename, filesize):
15+ self._filename = filename
16+ self._size = filesize
17+ self._seen_so_far = 0
18+
19+ def __call__(self, bytes_amount):
20+ self._seen_so_far += bytes_amount
21+ percentage = (self._seen_so_far / self._size) * 100 if self._size > 0 else 0
22+ sys.stdout.write(
23+ f"\r Progress: {self._seen_so_far:,} / {self._size:,} bytes ({percentage:.1f}%)"
24+ )
25+ sys.stdout.flush()
26+
27+
28+def download_s3_folder(s3_url: str, local_dir: str, skip_confirmation: bool = False):
29+ """
30+ Download entire folder from S3 public bucket
31+
32+ Args:
33+ s3_url: S3 URL in format s3://bucket-name/path/to/folder/
34+ local_dir: Local directory path to download files to
35+ skip_confirmation: If True, skip user confirmation prompt
36+ """
37+ # Parse S3 URL
38+ if not s3_url.startswith('s3://'):
39+ raise ValueError(f"S3 URL must start with 's3://': {s3_url}")
40+
41+ s3_url = s3_url.rstrip('/')
42+ s3_parts = s3_url[5:].split('/', 1)
43+ bucket_name = s3_parts[0]
44+ prefix = s3_parts[1] + '/' if len(s3_parts) > 1 else ''
45+
46+ # S3 configuration for public bucket (no credentials needed)
47+ s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED))
48+
49+ # Create local directory
50+ Path(local_dir).mkdir(parents=True, exist_ok=True)
51+
52+ print(f"Scanning s3://{bucket_name}/{prefix}")
53+ print(f"Will download to local directory: {local_dir}/")
54+ print()
55+
56+ # List all objects in the folder to calculate total size
57+ paginator = s3.get_paginator('list_objects_v2')
58+ pages = paginator.paginate(Bucket=bucket_name, Prefix=prefix)
59+
60+ files_to_download = []
61+ total_size = 0
62+
63+ for page in pages:
64+ if 'Contents' not in page:
65+ print("No files found in the specified path")
66+ return
67+
68+ for obj in page['Contents']:
69+ s3_key = obj['Key']
70+ file_size = obj['Size']
71+
72+ # Skip if it's just the directory itself
73+ if s3_key == prefix or s3_key == prefix.rstrip('/'):
74+ continue
75+
76+ # Get the relative path (remove the prefix)
77+ relative_path = s3_key[len(prefix):]
78+ if not relative_path:
79+ continue
80+
81+ files_to_download.append({
82+ 's3_key': s3_key,
83+ 'relative_path': relative_path,
84+ 'size': file_size
85+ })
86+ total_size += file_size
87+
88+ print(f"Found {len(files_to_download)} files")
89+ print(f"Total size: {total_size:,} bytes ({total_size / (1024**2):.2f} MB, {total_size / (1024**3):.2f} GB)")
90+ print()
91+
92+ if not skip_confirmation:
93+ response = input("Continue with download? (y/n): ")
94+ if response.lower() != 'y':
95+ print("Download cancelled")
96+ return
97+
98+ print()
99+ print("Starting download...")
100+ print()
101+
102+ # Download all files
103+ downloaded_bytes = 0
104+ for idx, file_info in enumerate(files_to_download, 1):
105+ s3_key = file_info['s3_key']
106+ relative_path = file_info['relative_path']
107+ file_size = file_info['size']
108+
109+ # Local file path
110+ local_file = os.path.join(local_dir, relative_path)
111+
112+ # Create subdirectories if needed
113+ local_file_dir = os.path.dirname(local_file)
114+ if local_file_dir:
115+ Path(local_file_dir).mkdir(parents=True, exist_ok=True)
116+
117+ # Download the file with progress callback
118+ print(f"[{idx}/{len(files_to_download)}] {relative_path} ({file_size:,} bytes)")
119+ progress = ProgressCallback(relative_path, file_size)
120+ s3.download_file(bucket_name, s3_key, local_file, Callback=progress)
121+ print() # New line after progress
122+
123+ downloaded_bytes += file_size
124+ overall_progress = (downloaded_bytes / total_size) * 100
125+ print(f" Overall progress: {downloaded_bytes:,} / {total_size:,} bytes ({overall_progress:.1f}%)")
126+ print()
127+
128+ print(f"Download complete!")
129+ print(f"Total files: {len(files_to_download)}")
130+ print(f"Total size: {total_size:,} bytes ({total_size / (1024**3):.2f} GB)")
python/ephys_compression_tests/algorithms/ans/__init__.pymodified+82−49View file
@@ -21,16 +21,24 @@ def create_ans_header(
2121 signal_length: int,
2222 state: np.uint64,
2323 symbol_counts: np.ndarray,
24- symbol_values: np.ndarray
24+ symbol_values: np.ndarray,
25+ shape: tuple
2526 ) -> bytes:
27+ ndim = len(shape)
28+ section0 = np.array([ndim] + list(shape), dtype=np.uint32)
2629 section1 = np.array([dtype_code, num_words, signal_length, len(symbol_counts)], dtype=np.uint32)
2730 section2 = np.array([state], dtype=np.uint64)
2831 symbol_counts_bytes = symbol_counts.astype(np.uint32).tobytes()
2932 symbol_values_bytes = symbol_values.tobytes()
3033
31- return section1.tobytes() + section2.tobytes() + symbol_counts_bytes + symbol_values_bytes
34+ return section0.tobytes() + section1.tobytes() + section2.tobytes() + symbol_counts_bytes + symbol_values_bytes
3235
3336 def unpack_ans_header(header_bytes: bytes) -> dict:
37+ # read section 0
38+ ndim = np.frombuffer(header_bytes[:4], dtype=np.uint32)[0]
39+ shape = tuple(np.frombuffer(header_bytes[4 : 4 + ndim * 4], dtype=np.uint32))
40+ offset = 4 + ndim * 4
41+ header_bytes = header_bytes[offset:]
3442 # read section 1
3543 section1_size = 4 * 4 # 4 uint32
3644 section1 = np.frombuffer(header_bytes[:section1_size], dtype=np.uint32)
@@ -62,12 +70,18 @@ def unpack_ans_header(header_bytes: bytes) -> dict:
6270 "state": state,
6371 "symbol_counts": symbol_counts,
6472 "symbol_values": symbol_values,
73+ "shape": shape
6574 }
6675
6776
6877 def ans_encode_0(x: np.ndarray) -> bytes:
6978 from simple_ans import ans_encode
7079
80+ shape0 = x.shape
81+ if x.ndim == 2:
82+ # flatten
83+ x = x.reshape(-1)
84+
7185 encoded = ans_encode(x)
7286 if x.dtype == np.uint8:
7387 dtype_code = 0
@@ -90,6 +104,7 @@ def ans_encode_0(x: np.ndarray) -> bytes:
90104 state=encoded.state,
91105 symbol_counts=encoded.symbol_counts,
92106 symbol_values=encoded.symbol_values,
107+ shape=shape0
93108 )
94109
95110 header_size = np.array([len(header_bytes)], dtype="uint32")
@@ -112,6 +127,9 @@ def ans_decode_0(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
112127 state = header_dict["state"]
113128 symbol_counts = header_dict["symbol_counts"]
114129 symbol_values = header_dict["symbol_values"]
130+ shape_from_header = header_dict["shape"]
131+ if shape != shape_from_header:
132+ raise ValueError("Shape mismatch between provided shape and shape in header")
115133
116134 words_bytes = x[4 + header_size :]
117135
@@ -156,29 +174,30 @@ for a in algorithm_dicts_base:
156174
157175 # add delta encoding
158176 for a in algorithm_dicts_base:
159- def encode0_ar_lossy(x: np.ndarray, a=a) -> bytes:
160- x_diff = np.diff(x)
161- x0 = x[0:1]
177+ def encode0(x: np.ndarray, a=a) -> bytes:
178+ assert x.ndim == 2 and x.shape[0] > 1, "Input array must be 2D with more than one timepoint"
179+ x_diff = np.diff(x, axis=0)
180+ first_timepoint = x[0:1, :].flatten()
162181 encoded_diff = a["encode"](x_diff)
163182 # Store the first value at the start
164- first_value_bytes = x0.tobytes()
165- return first_value_bytes + encoded_diff
166- def decode0_ar_lossy(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
183+ first_timepoint_bytes = first_timepoint.tobytes()
184+ return first_timepoint_bytes + encoded_diff
185+ def decode0(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
167186 dtype_np = np.dtype(dtype)
168- num_bytes_first_value = dtype_np.itemsize
169- first_value_bytes = x[:num_bytes_first_value]
170- x0 = np.frombuffer(first_value_bytes, dtype=dtype_np)
171- encoded_diff = x[num_bytes_first_value:]
172- x_diff = a["decode"](encoded_diff, dtype, (shape[0]-1,))
187+ num_bytes_first_timepoint = dtype_np.itemsize * shape[1]
188+ first_timepoint_bytes = x[:num_bytes_first_timepoint]
189+ first_timepoint = np.frombuffer(first_timepoint_bytes, dtype=dtype_np)
190+ encoded_diff = x[num_bytes_first_timepoint:]
191+ x_diff = a["decode"](encoded_diff, dtype, (shape[0]-1, shape[1]))
173192 x_reconstructed = np.empty(shape, dtype=dtype_np)
174- x_reconstructed[0] = x0
175- x_reconstructed[1:] = x0 + np.cumsum(x_diff)
193+ x_reconstructed[0] = first_timepoint
194+ x_reconstructed[1:] = first_timepoint + np.cumsum(x_diff, axis=0)
176195 return x_reconstructed
177196 algorithm_dicts.append({
178197 "name": a["name"] + "-delta",
179198 "version": a["version"],
180- "encode": encode0_ar_lossy,
181- "decode": decode0_ar_lossy,
199+ "encode": encode0,
200+ "decode": decode0,
182201 "description": a["description"] + " with delta encoding",
183202 "tags": a["tags"] + ["delta"],
184203 "source_file": a["source_file"],
@@ -188,28 +207,30 @@ for a in algorithm_dicts_base:
188207 # add delta2 encoding
189208 for a in algorithm_dicts_base:
190209 def encode0_ar_lossy(x: np.ndarray, a=a) -> bytes:
191- x_diff = np.diff(np.diff(x))
192- x0 = x[0:1]
210+ assert x.ndim == 2 and x.shape[0] > 2, "Input array must be 2D with more than two timepoints"
211+ x_diff = np.diff(np.diff(x, axis=0), axis=0)
212+ first_timepoint = x[0:1, :].flatten()
213+ second_timepoint = x[1:2, :].flatten()
193214 encoded_diff = a["encode"](x_diff)
194215 # Store the first value at the start
195- first_value_bytes = x0.tobytes()
196- second_value_bytes = x[1:2].tobytes()
197- return first_value_bytes + second_value_bytes + encoded_diff
216+ first_timepoint_bytes = first_timepoint.tobytes()
217+ second_timepoint_bytes = second_timepoint.tobytes()
218+ return first_timepoint_bytes + second_timepoint_bytes + encoded_diff
198219 def decode0_ar_lossy(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
199220 dtype_np = np.dtype(dtype)
200- num_bytes_first_value = dtype_np.itemsize
201- first_value_bytes = x[:num_bytes_first_value]
202- second_value_bytes = x[num_bytes_first_value:2*num_bytes_first_value]
203- x0 = np.frombuffer(first_value_bytes, dtype=dtype_np)
204- x1 = np.frombuffer(second_value_bytes, dtype=dtype_np)
205- encoded_diff2 = x[2*num_bytes_first_value:]
206- x_diff2 = a["decode"](encoded_diff2, dtype, (shape[0]-2,))
207- x_recon1 = np.empty((shape[0]-1,), dtype=dtype_np)
221+ num_bytes_first_timepoint = dtype_np.itemsize * shape[1]
222+ first_timepoint_bytes = x[:num_bytes_first_timepoint]
223+ second_timepoint_bytes = x[num_bytes_first_timepoint:2*num_bytes_first_timepoint]
224+ x0 = np.frombuffer(first_timepoint_bytes, dtype=dtype_np)
225+ x1 = np.frombuffer(second_timepoint_bytes, dtype=dtype_np)
226+ encoded_diff2 = x[2*num_bytes_first_timepoint:]
227+ x_diff2 = a["decode"](encoded_diff2, dtype, (shape[0]-2, shape[1]))
228+ x_recon1 = np.empty((shape[0]-1,shape[1]), dtype=dtype_np)
208229 x_recon1[0] = x1 - x0
209- x_recon1[1:] = x_recon1[0] + np.cumsum(x_diff2)
230+ x_recon1[1:] = x_recon1[0] + np.cumsum(x_diff2, axis=0)
210231 x_reconstructed = np.empty(shape, dtype=dtype_np)
211232 x_reconstructed[0] = x0
212- x_reconstructed[1:] = x0 + np.cumsum(x_recon1)
233+ x_reconstructed[1:] = x0 + np.cumsum(x_recon1, axis=0)
213234 return x_reconstructed
214235 algorithm_dicts.append({
215236 "name": a["name"] + "-delta2",
@@ -225,25 +246,31 @@ for a in algorithm_dicts_base:
225246 # Add auto-regressive prediction encoding
226247 for a in algorithm_dicts_base:
227248 for order in [2, 8]:
228- def encode0_ar_lossy(x: np.ndarray, a=a, order=order) -> bytes:
249+ def encode0_ar(x: np.ndarray, a=a, order=order) -> bytes:
250+ assert x.ndim == 2 and x.shape[0] > order, f"Input array must be 2D (timepoints x channels) with more than {order} timepoints"
229251 coeffs, residuals, initial_values = encode_ar(x, order=order)
252+ # coeffs: (n_channels x order), residuals: (n_timepoints-order x n_channels), initial_values: (order x n_channels)
230253 encoded_residuals = a["encode"](residuals)
231254 coeffs_bytes = coeffs.astype(np.float32).tobytes()
232255 initial_values_bytes = initial_values.astype(np.int16).tobytes()
233256 return coeffs_bytes + initial_values_bytes + encoded_residuals
234- def decode0_ar_lossy(x: bytes, dtype: str, shape: tuple, a=a, order=order) -> np.ndarray:
257+ def decode0_ar(x: bytes, dtype: str, shape: tuple, a=a, order=order) -> np.ndarray:
258+ assert len(shape) == 2, f"Shape must be 2D (timepoints x channels)"
235259 dtype_np = np.dtype(dtype)
236- num_bytes_coeffs = order * np.dtype(np.float32).itemsize
260+ n_channels = shape[1]
261+ # coeffs is (n_channels x order)
262+ num_bytes_coeffs = n_channels * order * np.dtype(np.float32).itemsize
237263 coeffs_bytes = x[:num_bytes_coeffs]
238- coeffs = np.frombuffer(coeffs_bytes, dtype=np.float32)
239- num_initial_values = len(coeffs)
240- num_bytes_initial_values = num_initial_values * dtype_np.itemsize
264+ coeffs = np.frombuffer(coeffs_bytes, dtype=np.float32).reshape((n_channels, order))
265+ # initial_values is (order x n_channels)
266+ num_bytes_initial_values = order * n_channels * dtype_np.itemsize
241267 initial_values_bytes = x[num_bytes_coeffs : num_bytes_coeffs + num_bytes_initial_values]
242- initial_values = np.frombuffer(initial_values_bytes, dtype=dtype_np)
268+ initial_values = np.frombuffer(initial_values_bytes, dtype=dtype_np).reshape((order, n_channels))
243269 encoded_residuals = x[num_bytes_coeffs + num_bytes_initial_values :]
244- residuals = a["decode"](encoded_residuals, dtype, (shape[0]-num_initial_values,))
270+ # residuals is ((shape[0]-order) x n_channels)
271+ residuals = a["decode"](encoded_residuals, dtype, (shape[0]-order, n_channels))
245272 reconstructed = decode_ar(coeffs, residuals, initial_values)
246- return reconstructed.reshape(shape)
273+ return reconstructed
247274 algorithm_dicts.append({
248275 "name": a["name"] + f"-ar{order}",
249276 "version": a["version"],
@@ -255,28 +282,34 @@ for a in algorithm_dicts_base:
255282 "long_description": a["long_description"]
256283 })
257284
258-# Add lossy ar2
285+# Add lossy ar
259286 for ar_order in [2, 8]:
260287 for tolerance in [1, 2, 3, 4, 5]:
261288 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)"
262290 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)
263292 encoded_residuals = ans_encode_0(residuals)
264293 coeffs_bytes = coeffs.astype(np.float32).tobytes()
265294 initial_values_bytes = initial_values.astype(np.int16).tobytes()
266295 return coeffs_bytes + initial_values_bytes + encoded_residuals
267296 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)"
268298 dtype_np = np.dtype(dtype)
269- num_bytes_coeffs = order * np.dtype(np.float32).itemsize
299+ n_channels = shape[1]
300+ # coeffs is (n_channels x order)
301+ num_bytes_coeffs = n_channels * order * np.dtype(np.float32).itemsize
270302 coeffs_bytes = x[:num_bytes_coeffs]
271- coeffs = np.frombuffer(coeffs_bytes, dtype=np.float32)
272- num_initial_values = len(coeffs)
273- num_bytes_initial_values = num_initial_values * dtype_np.itemsize
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
274306 initial_values_bytes = x[num_bytes_coeffs : num_bytes_coeffs + num_bytes_initial_values]
275- initial_values = np.frombuffer(initial_values_bytes, dtype=dtype_np)
307+ initial_values = np.frombuffer(initial_values_bytes, dtype=dtype_np).reshape((order, n_channels))
276308 encoded_residuals = x[num_bytes_coeffs + num_bytes_initial_values :]
277- residuals = ans_decode_0(encoded_residuals, dtype, (shape[0]-num_initial_values,))
309+ # residuals is ((shape[0]-order) x n_channels)
310+ residuals = ans_decode_0(encoded_residuals, dtype, (shape[0]-order, n_channels))
278311 reconstructed = decode_ar(coeffs, residuals, initial_values)
279- return reconstructed.reshape(shape)
312+ return reconstructed
280313 algorithm_dicts.append({
281314 "name": f"ans-ar{ar_order}-lossy-tol{tolerance}",
282315 "version": "1",
python/ephys_compression_tests/algorithms/ans/ar.pymodified+101−43View file
@@ -8,23 +8,23 @@ from numba import njit
88 def _warmup_numba_functions():
99 """Warmup numba JIT compilation with small test data."""
1010 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)
11+ # Create small test data for 2D arrays (timepoints x channels)
12+ test_data = np.array([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7], [7, 8], [8, 9]], dtype=np.int16)
13+ test_coeffs = np.array([[0.5, 0.3], [0.4, 0.2]], dtype=np.float32) # channels x order
14+ test_residuals = np.array([[1, 2], [2, 3], [3, 4], [4, 5]], dtype=np.int16)
15+ test_initial = np.array([[1, 2], [2, 3]], dtype=np.int16)
1616 test_step = 2
1717
1818 # 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)
19+ _create_design_matrix_channel(test_data[:, 0], 2)
20+ _apply_ar_residuals_kernel_channel(test_data[:, 0], test_coeffs[0])
21+ _apply_ar_residuals_lossy_kernel_channel(test_data[:, 0], test_coeffs[0], test_step)
22+ _decode_ar_kernel_channel(test_coeffs[0], test_residuals[:, 0], test_initial[:, 0])
2323
2424
2525 @njit
26-def _create_design_matrix(data: np.ndarray, order: int) -> Tuple[np.ndarray, np.ndarray]:
27- """Numba-optimized design matrix creation for AR model."""
26+def _create_design_matrix_channel(data: np.ndarray, order: int) -> Tuple[np.ndarray, np.ndarray]:
27+ """Numba-optimized design matrix creation for AR model (single channel)."""
2828 n = len(data)
2929 X_design = np.zeros((n - order, order))
3030 y_target = data[order:]
@@ -41,31 +41,35 @@ def fit_ar_model(data: np.ndarray, order: int) -> np.ndarray:
4141 Fit an autoregressive model of given order using least squares.
4242
4343 Args:
44- data: Input data array
44+ data: Input data array (timepoints x channels)
4545 order: AR model order
4646
4747 Returns:
48- AR coefficients as numpy array
48+ AR coefficients as numpy array (channels x order)
4949 """
50- n = len(data)
51- if order >= n:
52- raise ValueError(f"AR order {order} must be less than data length {n}")
50+ n_timepoints, n_channels = data.shape
51+ if order >= n_timepoints:
52+ raise ValueError(f"AR order {order} must be less than data length {n_timepoints}")
5353
54- # Create design matrix using numba-optimized function
55- X_design, y_target = _create_design_matrix(data, order)
54+ # Fit AR model for each channel separately
55+ coeffs = np.zeros((n_channels, order), dtype=np.float32)
5656
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)
57+ for ch in range(n_channels):
58+ # Create design matrix using numba-optimized function
59+ X_design, y_target = _create_design_matrix_channel(data[:, ch], order)
60+
61+ # Use faster solve via normal equations: (X^T X) coeffs = X^T y
62+ # This is faster than lstsq for overdetermined systems
63+ XtX = X_design.T @ X_design
64+ Xty = X_design.T @ y_target
65+ coeffs[ch] = np.linalg.solve(XtX, Xty)
6266
6367 return coeffs
6468
6569
6670 @njit
67-def _apply_ar_residuals_kernel(data: np.ndarray, coeffs: np.ndarray) -> np.ndarray:
68- """Numba-optimized kernel for computing AR residuals."""
71+def _apply_ar_residuals_kernel_channel(data: np.ndarray, coeffs: np.ndarray) -> np.ndarray:
72+ """Numba-optimized kernel for computing AR residuals (single channel)."""
6973 order = len(coeffs)
7074 n = len(data)
7175 residuals = np.empty(n - order, dtype=data.dtype)
@@ -84,9 +88,10 @@ def _apply_ar_residuals_kernel(data: np.ndarray, coeffs: np.ndarray) -> np.ndarr
8488
8589 return residuals
8690
91+
8792 @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."""
93+def _apply_ar_residuals_lossy_kernel_channel(data: np.ndarray, coeffs: np.ndarray, step: int) -> np.ndarray:
94+ """Numba-optimized kernel for computing AR residuals with lossy quantization (single channel)."""
9095 order = len(coeffs)
9196 n = len(data)
9297 residuals = np.empty(n - order, dtype=data.dtype)
@@ -126,21 +131,47 @@ def apply_ar_residuals(data: np.ndarray, coeffs: np.ndarray) -> np.ndarray:
126131 Apply AR model with given coefficients and return residuals.
127132
128133 Args:
129- data: Input data array
130- coeffs: AR coefficients
134+ data: Input data array (timepoints x channels)
135+ coeffs: AR coefficients (channels x order)
131136
132137 Returns:
133- Residuals array
138+ Residuals array (timepoints x channels)
134139 """
135140 # Ensure coeffs is float32
136141 coeffs = np.array(coeffs, dtype=np.float32)
137142
138- return _apply_ar_residuals_kernel(data, coeffs)
143+ n_timepoints, n_channels = data.shape
144+ order = coeffs.shape[1]
145+ residuals = np.zeros((n_timepoints - order, n_channels), dtype=data.dtype)
146+
147+ for ch in range(n_channels):
148+ residuals[:, ch] = _apply_ar_residuals_kernel_channel(data[:, ch], coeffs[ch])
149+
150+ return residuals
139151
140152
141153 def apply_ar_residuals_lossy(data: np.ndarray, coeffs: np.ndarray, step: int) -> np.ndarray:
154+ """
155+ Apply AR model with given coefficients and return lossy residuals.
156+
157+ Args:
158+ data: Input data array (timepoints x channels)
159+ coeffs: AR coefficients (channels x order)
160+ step: Quantization step size
161+
162+ Returns:
163+ Residuals array (timepoints x channels)
164+ """
142165 coeffs = np.array(coeffs, dtype=np.float32)
143- return _apply_ar_residuals_lossy_kernel(data, coeffs, step)
166+
167+ n_timepoints, n_channels = data.shape
168+ order = coeffs.shape[1]
169+ residuals = np.zeros((n_timepoints - order, n_channels), dtype=data.dtype)
170+
171+ for ch in range(n_channels):
172+ residuals[:, ch] = _apply_ar_residuals_lossy_kernel_channel(data[:, ch], coeffs[ch], step)
173+
174+ return residuals
144175
145176
146177 def encode_ar(data: np.ndarray, order: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -148,11 +179,14 @@ def encode_ar(data: np.ndarray, order: int) -> Tuple[np.ndarray, np.ndarray, np.
148179 Encode data using AR model - returns coefficients, residuals, and initial values.
149180
150181 Args:
151- data: Input data array (int16)
182+ data: Input data array (timepoints x channels, int16)
152183 order: AR model order
153184
154185 Returns:
155186 Tuple of (coefficients, residuals, initial_values)
187+ - coefficients: channels x order (float32)
188+ - residuals: (timepoints - order) x channels (int16)
189+ - initial_values: order x channels (int16)
156190 """
157191 # Fit AR model
158192 coeffs = fit_ar_model(data, order)
@@ -164,11 +198,26 @@ def encode_ar(data: np.ndarray, order: int) -> Tuple[np.ndarray, np.ndarray, np.
164198 residuals = apply_ar_residuals(data, coeffs)
165199
166200 # Store initial values
167- initial_values = data[:order]
201+ initial_values = data[:order, :]
168202
169203 return coeffs, residuals, initial_values
170204
205+
171206 def encode_ar_lossy(data: np.ndarray, order: int, step: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
207+ """
208+ Encode data using AR model with lossy quantization.
209+
210+ Args:
211+ data: Input data array (timepoints x channels, int16)
212+ order: AR model order
213+ step: Quantization step size
214+
215+ Returns:
216+ Tuple of (coefficients, residuals, initial_values)
217+ - coefficients: channels x order (float32)
218+ - residuals: (timepoints - order) x channels (int16)
219+ - initial_values: order x channels (int16)
220+ """
172221 # Fit AR model
173222 coeffs = fit_ar_model(data, order)
174223
@@ -179,14 +228,14 @@ def encode_ar_lossy(data: np.ndarray, order: int, step: int) -> Tuple[np.ndarray
179228 residuals = apply_ar_residuals_lossy(data, coeffs, step=step)
180229
181230 # Store initial values
182- initial_values = data[:order]
231+ initial_values = data[:order, :]
183232
184233 return coeffs, residuals, initial_values
185234
186235
187236 @njit
188-def _decode_ar_kernel(coeffs: np.ndarray, residuals: np.ndarray, initial_values: np.ndarray) -> np.ndarray:
189- """Numba-optimized kernel for AR decoding."""
237+def _decode_ar_kernel_channel(coeffs: np.ndarray, residuals: np.ndarray, initial_values: np.ndarray) -> np.ndarray:
238+ """Numba-optimized kernel for AR decoding (single channel)."""
190239 order = len(coeffs)
191240 n = len(residuals) + order
192241 reconstructed = np.empty(n, dtype=np.int16)
@@ -213,19 +262,28 @@ def decode_ar(coeffs: np.ndarray, residuals: np.ndarray, initial_values: np.ndar
213262 Decode AR encoded data.
214263
215264 Args:
216- coeffs: AR coefficients (float32)
217- residuals: Residuals array
218- initial_values: Initial values (first 'order' samples)
265+ coeffs: AR coefficients (channels x order, float32)
266+ residuals: Residuals array ((timepoints - order) x channels)
267+ initial_values: Initial values (order x channels)
219268
220269 Returns:
221- Reconstructed data array
270+ Reconstructed data array (timepoints x channels)
222271 """
223272 # Ensure coeffs is float32
224273 coeffs = np.array(coeffs, dtype=np.float32)
225274
226- return _decode_ar_kernel(coeffs, residuals, initial_values)
275+ n_channels = coeffs.shape[0]
276+ order = coeffs.shape[1]
277+ n_residuals = residuals.shape[0]
278+ n_timepoints = n_residuals + order
279+
280+ reconstructed = np.zeros((n_timepoints, n_channels), dtype=np.int16)
281+
282+ for ch in range(n_channels):
283+ reconstructed[:, ch] = _decode_ar_kernel_channel(coeffs[ch], residuals[:, ch], initial_values[:, ch])
284+
285+ return reconstructed
227286
228287
229288 # Warmup numba functions on module import
230289 _warmup_numba_functions()
231-
python/ephys_compression_tests/algorithms/lzma/__init__.pymodified+13−12View file
@@ -92,25 +92,26 @@ algorithm_dicts = []
9292 for a in algorithm_dicts_base:
9393 algorithm_dicts.append(a)
9494
95-# Add delta encoding
95+# add delta encoding
9696 for a in algorithm_dicts_base:
9797 def encode0(x: np.ndarray, a=a) -> bytes:
98- x_diff = np.diff(x)
99- x0 = x[0:1]
98+ assert x.ndim == 2 and x.shape[0] > 1, "Input array must be 2D with more than one timepoint"
99+ x_diff = np.diff(x, axis=0)
100+ first_timepoint = x[0:1, :].flatten()
100101 encoded_diff = a["encode"](x_diff)
101102 # Store the first value at the start
102- first_value_bytes = x0.tobytes()
103- return first_value_bytes + encoded_diff
103+ first_timepoint_bytes = first_timepoint.tobytes()
104+ return first_timepoint_bytes + encoded_diff
104105 def decode0(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
105106 dtype_np = np.dtype(dtype)
106- num_bytes_first_value = dtype_np.itemsize
107- first_value_bytes = x[:num_bytes_first_value]
108- x0 = np.frombuffer(first_value_bytes, dtype=dtype_np)
109- encoded_diff = x[num_bytes_first_value:]
110- x_diff = a["decode"](encoded_diff, dtype, (shape[0]-1,))
107+ num_bytes_first_timepoint = dtype_np.itemsize * shape[1]
108+ first_timepoint_bytes = x[:num_bytes_first_timepoint]
109+ first_timepoint = np.frombuffer(first_timepoint_bytes, dtype=dtype_np)
110+ encoded_diff = x[num_bytes_first_timepoint:]
111+ x_diff = a["decode"](encoded_diff, dtype, (shape[0]-1, shape[1]))
111112 x_reconstructed = np.empty(shape, dtype=dtype_np)
112- x_reconstructed[0] = x0
113- x_reconstructed[1:] = x0 + np.cumsum(x_diff)
113+ x_reconstructed[0] = first_timepoint
114+ x_reconstructed[1:] = first_timepoint + np.cumsum(x_diff, axis=0)
114115 return x_reconstructed
115116 algorithm_dicts.append({
116117 "name": a["name"] + "-delta",
python/ephys_compression_tests/algorithms/wavpack/__init__.pymodified+12−11View file
@@ -52,22 +52,23 @@ for a in algorithm_dicts_base:
5252 # add delta encoding
5353 for a in algorithm_dicts_base:
5454 def encode0(x: np.ndarray, a=a) -> bytes:
55- x_diff = np.diff(x)
56- x0 = x[0:1]
55+ assert x.ndim == 2 and x.shape[0] > 1, "Input array must be 2D with more than one timepoint"
56+ x_diff = np.diff(x, axis=0)
57+ first_timepoint = x[0:1, :].flatten()
5758 encoded_diff = a["encode"](x_diff)
5859 # Store the first value at the start
59- first_value_bytes = x0.tobytes()
60- return first_value_bytes + encoded_diff
60+ first_timepoint_bytes = first_timepoint.tobytes()
61+ return first_timepoint_bytes + encoded_diff
6162 def decode0(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
6263 dtype_np = np.dtype(dtype)
63- num_bytes_first_value = dtype_np.itemsize
64- first_value_bytes = x[:num_bytes_first_value]
65- x0 = np.frombuffer(first_value_bytes, dtype=dtype_np)
66- encoded_diff = x[num_bytes_first_value:]
67- x_diff = a["decode"](encoded_diff, dtype, (shape[0]-1,))
64+ num_bytes_first_timepoint = dtype_np.itemsize * shape[1]
65+ first_timepoint_bytes = x[:num_bytes_first_timepoint]
66+ first_timepoint = np.frombuffer(first_timepoint_bytes, dtype=dtype_np)
67+ encoded_diff = x[num_bytes_first_timepoint:]
68+ x_diff = a["decode"](encoded_diff, dtype, (shape[0]-1, shape[1]))
6869 x_reconstructed = np.empty(shape, dtype=dtype_np)
69- x_reconstructed[0] = x0
70- x_reconstructed[1:] = x0 + np.cumsum(x_diff)
70+ x_reconstructed[0] = first_timepoint
71+ x_reconstructed[1:] = first_timepoint + np.cumsum(x_diff, axis=0)
7172 return x_reconstructed
7273 algorithm_dicts.append({
7374 "name": a["name"] + "-delta",
python/ephys_compression_tests/algorithms/zlib/__init__.pymodified+13−12View file
@@ -92,25 +92,26 @@ algorithm_dicts = []
9292 for a in algorithm_dicts_base:
9393 algorithm_dicts.append(a)
9494
95-# Add delta encoding
95+# add delta encoding
9696 for a in algorithm_dicts_base:
9797 def encode0(x: np.ndarray, a=a) -> bytes:
98- x_diff = np.diff(x)
99- x0 = x[0:1]
98+ assert x.ndim == 2 and x.shape[0] > 1, "Input array must be 2D with more than one timepoint"
99+ x_diff = np.diff(x, axis=0)
100+ first_timepoint = x[0:1, :].flatten()
100101 encoded_diff = a["encode"](x_diff)
101102 # Store the first value at the start
102- first_value_bytes = x0.tobytes()
103- return first_value_bytes + encoded_diff
103+ first_timepoint_bytes = first_timepoint.tobytes()
104+ return first_timepoint_bytes + encoded_diff
104105 def decode0(x: bytes, dtype: str, shape: tuple, a=a) -> np.ndarray:
105106 dtype_np = np.dtype(dtype)
106- num_bytes_first_value = dtype_np.itemsize
107- first_value_bytes = x[:num_bytes_first_value]
108- x0 = np.frombuffer(first_value_bytes, dtype=dtype_np)
109- encoded_diff = x[num_bytes_first_value:]
110- x_diff = a["decode"](encoded_diff, dtype, (shape[0]-1,))
107+ num_bytes_first_timepoint = dtype_np.itemsize * shape[1]
108+ first_timepoint_bytes = x[:num_bytes_first_timepoint]
109+ first_timepoint = np.frombuffer(first_timepoint_bytes, dtype=dtype_np)
110+ encoded_diff = x[num_bytes_first_timepoint:]
111+ x_diff = a["decode"](encoded_diff, dtype, (shape[0]-1, shape[1]))
111112 x_reconstructed = np.empty(shape, dtype=dtype_np)
112- x_reconstructed[0] = x0
113- x_reconstructed[1:] = x0 + np.cumsum(x_diff)
113+ x_reconstructed[0] = first_timepoint
114+ x_reconstructed[1:] = first_timepoint + np.cumsum(x_diff, axis=0)
114115 return x_reconstructed
115116 algorithm_dicts.append({
116117 "name": a["name"] + "-delta",
python/ephys_compression_tests/datasets/aind_compression/__init__.pymodified+18−3View file
@@ -19,7 +19,7 @@ def _load_long_description():
1919
2020 LONG_DESCRIPTION = _load_long_description()
2121
22-tags = ["real", "ecephys", "timeseries", "1d", "integer", "correlated"]
22+tags = ["real", "ecephys", "timeseries", "integer", "correlated"]
2323
2424
2525 def load_aind_ch101() -> np.ndarray:
@@ -34,8 +34,14 @@ def load_aind_ch101() -> np.ndarray:
3434 response.raise_for_status()
3535 data = np.load(io.BytesIO(response.content)).flatten()
3636 return data
37-
3837
38+def load_aind_ch101_110() -> np.ndarray:
39+ url = "https://tempory.net/ephys-compression-tests/aind/aind_compression_np2_probeB_ch101-110.raw.npy"
40+ print(f'Loading AIND dataset from {url}...')
41+ response = requests.get(url)
42+ response.raise_for_status()
43+ data = np.load(io.BytesIO(response.content))
44+ return data
3945
4046 dataset_dicts_base = [
4147 {
@@ -43,7 +49,16 @@ dataset_dicts_base = [
4349 "version": "1",
4450 "description": "AIND CH101 dataset",
4551 "create": load_aind_ch101,
46- "tags": tags,
52+ "tags": tags + ["single-channel"],
53+ "source_file": SOURCE_FILE,
54+ "long_description": LONG_DESCRIPTION,
55+ },
56+ {
57+ "name": "aind-compression-np2-ProbeB-ch101-110",
58+ "version": "1",
59+ "description": "AIND CH101-110 dataset",
60+ "create": load_aind_ch101_110,
61+ "tags": tags + ["multi-channel"],
4762 "source_file": SOURCE_FILE,
4863 "long_description": LONG_DESCRIPTION,
4964 }
python/ephys_compression_tests/run_benchmarks/benchmark_timing.pymodified+2−0View file
@@ -63,6 +63,8 @@ def run_compression_benchmark(
6363 - result: Dictionary with benchmark metrics
6464 - encoded: Compressed data bytes
6565 """
66+ if data.ndim == 1:
67+ data = data[:, np.newaxis]
6668 original_size = len(data.tobytes())
6769 dtype = str(data.dtype)
6870