quantization correction
1 changed file+39−5
python/ephys_compression_tests/datasets/aind_compression/__init__.pymodified+39−5View file
@@ -22,6 +22,36 @@ LONG_DESCRIPTION = _load_long_description()
2222 tags = ["real", "ecephys", "timeseries", "integer", "correlated"]
2323
2424
25+# It's important to correct the quantization levels before using these datasets
26+# Wavpack in particular will do a lot worse if the data is not properly quantized
27+def correct_quantization_for_channel(data: np.ndarray) -> np.ndarray:
28+ closed_to_zero_val = np.argmin(np.abs(data))
29+ print(f'Value closest to zero: {data[closed_to_zero_val]} at index {closed_to_zero_val}')
30+ data = data - data[closed_to_zero_val]
31+ unique_vals = np.unique(data)
32+ # differences between unique values
33+ diffs = np.diff(unique_vals)
34+ # minimum diff is the quantization step size
35+ diff0 = np.min(diffs[diffs > 0])
36+ print(f'Identified quantization step size: {diff0}')
37+ # divide by quantization step size
38+ data = data / diff0
39+ data = data.astype(np.int16)
40+ return data
41+
42+def correct_quantization(data: np.ndarray) -> np.ndarray:
43+ if data.ndim == 1:
44+ return correct_quantization_for_channel(data)
45+ elif data.ndim == 2:
46+ corrected_channels = []
47+ for ch in range(data.shape[1]):
48+ print(f'Correcting quantization for channel {ch}...')
49+ corrected_ch = correct_quantization_for_channel(data[:, ch])
50+ corrected_channels.append(corrected_ch)
51+ return np.stack(corrected_channels, axis=1)
52+ else:
53+ raise ValueError(f'Unsupported data ndim: {data.ndim}')
54+
2555 def load_aind_np2_probeB_ch101() -> np.ndarray:
2656 """Load AIND CH101 dataset from external URL.
2757
@@ -33,6 +63,7 @@ def load_aind_np2_probeB_ch101() -> np.ndarray:
3363 response = requests.get(url)
3464 response.raise_for_status()
3565 data = np.load(io.BytesIO(response.content)).flatten()
66+ data = correct_quantization(data)
3667 return data
3768
3869 def load_aind_np2_probeB_ch101_110() -> np.ndarray:
@@ -41,6 +72,7 @@ def load_aind_np2_probeB_ch101_110() -> np.ndarray:
4172 response = requests.get(url)
4273 response.raise_for_status()
4374 data = np.load(io.BytesIO(response.content))
75+ data = correct_quantization(data)
4476 return data
4577
4678 def load_aind_np1_probeA_101_110() -> np.ndarray:
@@ -49,6 +81,7 @@ def load_aind_np1_probeA_101_110() -> np.ndarray:
4981 response = requests.get(url)
5082 response.raise_for_status()
5183 data = np.load(io.BytesIO(response.content))
84+ data = correct_quantization(data)
5285 return data
5386
5487 # ibl-np1-probe00
@@ -58,12 +91,13 @@ def load_ibl_np1_probe00_101_110() -> np.ndarray:
5891 response = requests.get(url)
5992 response.raise_for_status()
6093 data = np.load(io.BytesIO(response.content))
94+ data = correct_quantization(data)
6195 return data
6296
6397 dataset_dicts_base = [
6498 {
6599 "name": "aind-compression-np2-ProbeB-ch101",
66- "version": "1",
100+ "version": "2",
67101 "description": "AIND CH101 dataset",
68102 "create": load_aind_np2_probeB_ch101,
69103 "tags": tags + ["single-channel"],
@@ -72,7 +106,7 @@ dataset_dicts_base = [
72106 },
73107 {
74108 "name": "aind-compression-np2-ProbeB-ch101-110",
75- "version": "1",
109+ "version": "2",
76110 "description": "AIND CH101-110 dataset",
77111 "create": load_aind_np2_probeB_ch101_110,
78112 "tags": tags + ["multi-channel"],
@@ -81,7 +115,7 @@ dataset_dicts_base = [
81115 },
82116 {
83117 "name": "aind-compression-np1-ProbeA-ch101-110",
84- "version": "1",
118+ "version": "2",
85119 "description": "AIND NP1 ProbeA CH101-110 dataset",
86120 "create": load_aind_np1_probeA_101_110,
87121 "tags": tags + ["multi-channel"],
@@ -90,7 +124,7 @@ dataset_dicts_base = [
90124 },
91125 {
92126 "name": "ibl-compression-np1-Probe00-ch101-110",
93- "version": "1",
127+ "version": "2",
94128 "description": "IBL NP1 Probe00 CH101-110 dataset",
95129 "create": load_ibl_np1_probe00_101_110,
96130 "tags": tags + ["multi-channel"],
@@ -114,7 +148,7 @@ for d in dataset_dicts_base:
114148 dataset_dicts.append(
115149 {
116150 "name": f'{d["name"]}-filtered',
117- "version": "1",
151+ "version": "2",
118152 "description": f'{d["description"]} (bandpass filtered 300-4000 Hz)',
119153 "create": create0,
120154 "tags": d["tags"] + ["filtered", "bandpass"],