zia_benchmark
20 changed files+830−10
README.mdmodified+1−1View file
@@ -1,3 +1,3 @@
11 # zia
22
3-Compression of integer arrays
3+Benchmarking compression of integer arrays
scripts/run_benchmarks.pyadded+74−0View file
@@ -0,0 +1,74 @@
1+#!/usr/bin/env python3
2+
3+import json
4+from pathlib import Path
5+import zia_benchmark
6+
7+def format_size(size_bytes: float) -> str:
8+ """Format size in bytes to human readable string"""
9+ for unit in ['B', 'KB', 'MB', 'GB']:
10+ if size_bytes < 1024:
11+ return f"{size_bytes:.1f} {unit}"
12+ size_bytes /= 1024
13+ return f"{size_bytes:.1f} TB"
14+
15+def format_time(seconds: float) -> str:
16+ """Format time in seconds to human readable string"""
17+ if seconds < 0.001:
18+ return f"{seconds * 1_000_000:.1f} µs"
19+ if seconds < 1:
20+ return f"{seconds * 1_000:.1f} ms"
21+ return f"{seconds:.2f} s"
22+
23+def main():
24+ # Run benchmarks
25+ print("Running benchmarks...")
26+ results = zia_benchmark.run_benchmarks()
27+
28+ # Group results by dataset
29+ datasets = {}
30+ for result in results['results']:
31+ dataset_name = result['dataset']
32+ if dataset_name not in datasets:
33+ datasets[dataset_name] = []
34+ datasets[dataset_name].append(result)
35+
36+ # Print results
37+ print("\nBenchmark Results:")
38+ print("=================")
39+
40+ for dataset_name, dataset_results in sorted(datasets.items()):
41+ print(f"\nDataset: {dataset_name}")
42+ print("-" * (len(dataset_name) + 9))
43+ print(f"Original size: {format_size(dataset_results[0]['original_size'])}")
44+
45+ # Sort algorithms by compression ratio
46+ dataset_results.sort(key=lambda x: x['compression_ratio'], reverse=True)
47+
48+ # Print table header
49+ print("\n{:<15} {:>12} {:>12} {:>12}".format(
50+ "Algorithm", "Ratio", "Encode", "Decode"
51+ ))
52+ print("-" * 53)
53+
54+ # Print results for each algorithm
55+ for result in dataset_results:
56+ print("{:<15} {:>11.2f}x {:>12} {:>12}".format(
57+ result['algorithm'],
58+ result['compression_ratio'],
59+ format_time(result['encode_time']),
60+ format_time(result['decode_time'])
61+ ))
62+
63+ # Save detailed results to JSON
64+ output_dir = Path("benchmark_results")
65+ output_dir.mkdir(exist_ok=True)
66+ output_file = output_dir / "results.json"
67+
68+ with open(output_file, "w") as f:
69+ json.dump(results, f, indent=2)
70+
71+ print(f"\nDetailed results saved to {output_file}")
72+
73+if __name__ == "__main__":
74+ main()
test1.pyadded+176−0View file
@@ -0,0 +1,176 @@
1+# %%
2+import numpy as np
3+from zia_benchmark._filters import bandpass_filter, highpass_filter
4+from zia_benchmark._data_loaders import load_real_000876, load_real_000409, load_real_001290
5+from zia_benchmark._compress_ints_lossless import compress_ints_lossless
6+from zia_benchmark._analysis import linear_fit, compute_entropy_per_sample, estimate_noise_level
7+import matplotlib.pyplot as plt
8+
9+# %%
10+N = 500_000
11+
12+channel_number = 101
13+X = load_real_000409(num_samples=N, num_channels=1, start_channel=channel_number).flatten()
14+
15+# X = load_real_001290(num_samples=N, num_channels=1, start_channel=0).flatten()
16+# X = load_real_000876(num_samples=N, num_channels=1, start_channel=45).flatten()
17+
18+X = X.astype(np.int16)
19+
20+# %%
21+plt.figure(figsize=(12, 4))
22+plt.plot(X[:2400])
23+# %%
24+def print_ideal_compression_ratio(X):
25+ ee = compute_entropy_per_sample(X)
26+ print(f'Ideal compression ratio: {X.itemsize * 8 / ee:.2f} ({ee:.2f} bits per sample)')
27+
28+def print_actual_compression_ratios(X):
29+ buf_zstd = compress_ints_lossless(X, method='zstd')
30+ buf_zlib = compress_ints_lossless(X, method='zlib')
31+ buf_lzma = compress_ints_lossless(X, method='lzma')
32+ buf_ans = compress_ints_lossless(X, method='simple_ans')
33+ print(f'Zstd compression ratio: {len(X) * X.itemsize / len(buf_zstd):.2f}')
34+ print(f'Zlib compression ratio: {len(X) * X.itemsize / len(buf_zlib):.2f}')
35+ print(f'Lzma compression ratio: {len(X) * X.itemsize / len(buf_lzma):.2f}')
36+ print(f'simple_ans compression ratio: {len(X) * X.itemsize / len(buf_ans):.2f}')
37+
38+def get_marcovian_prediction_residual(X, M):
39+ sequences = np.array([X[i:i+M] for i in range(len(X) - 2 * M + 1)])
40+ predictors = sequences[:, :M - 1]
41+ target = sequences[:, M - 1]
42+
43+ coeffs, predict = linear_fit(predictors, target)
44+ predictions = predict(predictors)
45+ predictions = np.round(predictions)
46+ residuals = target - predictions
47+ residuals = residuals.astype(np.int16)
48+ return residuals
49+
50+# %%
51+print('RAW')
52+print_ideal_compression_ratio(X)
53+
54+# %%
55+print('RAW DELTA ENCODING')
56+print_ideal_compression_ratio(np.diff(X))
57+
58+# %%
59+print('RAW DELTA ENCODING - actual compression ratios')
60+print_actual_compression_ratios(np.diff(X))
61+print_ideal_compression_ratio(np.diff(X))
62+
63+# %%
64+X_mr = get_marcovian_prediction_residual(X, 20)
65+print('RAW MARCOVIAN')
66+print_ideal_compression_ratio(X_mr)
67+
68+# %%
69+v = 0.25 # step size for quantization
70+lowcut = 300
71+highcut = 6000
72+X_filt = bandpass_filter(X - np.median(X), sampling_frequency=30000, lowcut=lowcut, highcut=highcut)
73+noise_level = estimate_noise_level(X_filt, sampling_frequency=30000)
74+X_filt_normalized = X_filt / noise_level
75+X2b = X_filt_normalized / v
76+X2 = np.round(X2b).astype(np.int16)
77+
78+# %%
79+plt.figure(figsize=(12, 4))
80+plt.plot(X[:2400])
81+
82+# %%
83+print('FILTERED (and quantized)')
84+print_ideal_compression_ratio(X2)
85+
86+# %%
87+print('FILTERED DELTA ENCODING')
88+print_ideal_compression_ratio(np.diff(X2))
89+
90+# %%
91+residuals2 = get_marcovian_prediction_residual(X2, 20)
92+print('FILTERED MARCOVIAN')
93+print_ideal_compression_ratio(residuals2)
94+
95+# %%
96+import matplotlib.pyplot as plt
97+plt.figure(figsize=(12, 4))
98+plt.plot(X[:2400])
99+plt.title('RAW')
100+
101+plt.figure(figsize=(12, 4))
102+plt.plot(X2[:2400])
103+plt.title('FILTERED')
104+
105+plt.figure(figsize=(12, 4))
106+plt.plot(residuals2[:2400])
107+plt.title('FILTERED MARCOVIAN')
108+
109+# %%
110+def sliding_max(x, delta):
111+ y = np.zeros_like(x)
112+ for i in range(len(x)):
113+ y[i] = np.max(x[max(0, i - delta):min(len(x), i + delta + 1)])
114+ return y
115+
116+def smoothed(x, delta):
117+ y = np.zeros_like(x)
118+ for i in range(len(x)):
119+ y[i] = np.mean(x[max(0, i - delta):min(len(x), i + delta + 1)])
120+ return y
121+
122+cc = [3, 6]
123+Y = sliding_max(np.abs(X_filt_normalized), 50)
124+Y = smoothed(Y, 20)
125+Y = np.minimum(1, np.maximum(0, (Y - cc[0]) / (cc[1] - cc[0])))
126+# Y = highpass_filter(Y, sampling_frequency=30000, lowcut=3)
127+Y_scaled = X_filt_normalized * Y
128+X3b = Y_scaled / v
129+X3 = np.round(X3b).astype(np.int16)
130+
131+plt.figure(figsize=(12, 4))
132+plt.plot(X2[:2400], color='lightgray')
133+# plt.plot(Y[4000:5000])
134+plt.plot(X3b[:2400])
135+# %%
136+print('FILTERED (and quantized) with suppression')
137+print_ideal_compression_ratio(X3)
138+print('')
139+print('FILTERED DELTA ENCODING with suppression')
140+print_ideal_compression_ratio(np.diff(X3))
141+print('')
142+print('FILTERED MARCOVIAN with suppression')
143+residuals3 = get_marcovian_prediction_residual(X3, 20)
144+print_ideal_compression_ratio(residuals3)
145+print('ACTUAL FILTERED MARCOVIAN with suppression')
146+print_actual_compression_ratios(residuals3)
147+# %%
148+def get_run_lengths(x):
149+ runs = []
150+ i = 0
151+ current_nonzero_run_length = 0
152+ while i < len(x):
153+ if np.all(x[i:i+10] == 0):
154+ runs.append(current_nonzero_run_length)
155+ current_nonzero_run_length = 0
156+ j = i
157+ while j < len(x) and x[j] == 0:
158+ j += 1
159+ runs.append(j - i)
160+ i = j
161+ else:
162+ current_nonzero_run_length += 1
163+ i += 1
164+ if np.max(runs) < 256:
165+ return np.array(runs, dtype=np.uint8)
166+ if np.max(runs) < 2 ** 16:
167+ return np.array(runs, dtype=np.uint16)
168+ return np.array(runs, dtype=np.uint32)
169+
170+AA = residuals3[residuals3 != 0]
171+run_lengths = get_run_lengths(residuals3)
172+print(run_lengths, run_lengths.itemsize)
173+ee = compute_entropy_per_sample(AA)
174+theoretical_size = (len(AA) * ee / 8) + run_lengths.nbytes
175+theoretical_compression_ratio = len(residuals3) * X.itemsize / theoretical_size
176+print(f'Theoretical compression ratio: {theoretical_compression_ratio:.2f}')
zia_benchmark/README.mdadded+1−0View file
@@ -0,0 +1 @@
1+# zia_benchmark
zia_benchmark/setup.pymodified+6−9View file
@@ -1,24 +1,21 @@
11 from setuptools import setup, find_packages
22
33 setup(
4- name="zia",
4+ name="zia_benchmark",
55 version="0.1.0",
66 packages=find_packages(where="src"),
77 package_dir={"": "src"},
88 install_requires=[
9- "numpy"
9+ "numpy",
10+ "zstandard",
11+ "simple_ans"
1012 ],
11- extras_require={
12- "dev": [
13- "pytest",
14- ],
15- },
1613 python_requires=">=3.8",
1714 author="Jeremy Magland",
18- description="Compression of integer arrays",
15+ description="Benchmarking compression methods for integer arrays",
1916 long_description=open("README.md").read(),
2017 long_description_content_type="text/markdown",
21- keywords="neuroscience, compression, signal processing",
18+ keywords="compression, signal processing",
2219 classifiers=[
2320 "Development Status :: 3 - Alpha",
2421 "Intended Audience :: Science/Research",
zia_benchmark/src/zia/__init__.pydeleted+0−0View file
No changes to the file's content.
zia_benchmark/src/zia_benchmark/__init__.pyadded+5−0View file
@@ -0,0 +1,5 @@
1+from .algorithms import algorithms
2+from .datasets import datasets
3+from .run_benchmarks import run_benchmarks
4+
5+__all__ = ['algorithms', 'datasets', 'run_benchmarks']
zia_benchmark/src/zia/_analysis.py →zia_benchmark/src/zia_benchmark/_analysis.pyrenamed+0−0View file
No changes to the file's content.
zia_benchmark/src/zia/_compress_ints_lossless.py →zia_benchmark/src/zia_benchmark/_compress_ints_lossless.pyrenamed+0−0View file
No changes to the file's content.
zia_benchmark/src/zia/_data_loaders.py →zia_benchmark/src/zia_benchmark/_data_loaders.pyrenamed+0−0View file
No changes to the file's content.
zia_benchmark/src/zia/_filters.py →zia_benchmark/src/zia_benchmark/_filters.pyrenamed+0−0View file
No changes to the file's content.
zia_benchmark/src/zia_benchmark/_memobin.pyadded+119−0View file
@@ -0,0 +1,119 @@
1+import json
2+import requests
3+from typing import Optional
4+
5+def create_signed_upload_url(url: str, size: int, user_id: str, memobin_api_key: str) -> str:
6+ """Create a signed upload URL for memobin.
7+
8+ Args:
9+ url: The target URL for the file
10+ size: Size of the file in bytes
11+ user_id: User ID for memobin
12+ memobin_api_key: API key for memobin authentication
13+
14+ Returns:
15+ The signed upload URL
16+
17+ Raises:
18+ ValueError: If the URL prefix is invalid
19+ requests.RequestException: If the API request fails
20+ """
21+ prefix = "https://tempory.net/f/memobin/"
22+ if not url.startswith(prefix):
23+ raise ValueError("Invalid url. Does not have proper prefix")
24+
25+ file_path = url[len(prefix):]
26+ tempory_api_url = "https://hub.tempory.net/api/uploadFile"
27+
28+ response = requests.post(
29+ tempory_api_url,
30+ headers={
31+ "Content-Type": "application/json",
32+ "Authorization": f"Bearer {memobin_api_key}"
33+ },
34+ json={
35+ "appName": "memobin",
36+ "filePath": file_path,
37+ "size": size,
38+ "userId": user_id
39+ }
40+ )
41+
42+ if not response.ok:
43+ raise requests.RequestException("Failed to get signed url")
44+
45+ result = response.json()
46+ upload_url = result["uploadUrl"]
47+ download_url = result["downloadUrl"]
48+
49+ if download_url != url:
50+ raise ValueError("Mismatch between download url and url")
51+
52+ return upload_url
53+
54+def construct_memobin_url(alg_name: str, dataset_name: str, alg_version: str,
55+ dataset_version: str, system_version: str) -> str:
56+ """Construct the memobin URL for a specific benchmark result.
57+
58+ Args:
59+ alg_name: Name of the algorithm
60+ dataset_name: Name of the dataset
61+ alg_version: Version of the algorithm
62+ dataset_version: Version of the dataset
63+ system_version: Version of the system
64+
65+ Returns:
66+ The constructed memobin URL
67+ """
68+ path = f"{alg_name}/{dataset_name}/{alg_version}/{dataset_version}/{system_version}/metadata.json"
69+ return f"https://tempory.net/f/memobin/{path}"
70+
71+def upload_to_memobin(metadata: dict, url: str, user_id: str, memobin_api_key: str) -> None:
72+ """Upload metadata to memobin.
73+
74+ Args:
75+ metadata: The metadata to upload
76+ url: The target URL for the file
77+ user_id: User ID for memobin
78+ memobin_api_key: API key for memobin authentication
79+
80+ Raises:
81+ requests.RequestException: If the upload fails
82+ """
83+ metadata_bytes = json.dumps(metadata).encode('utf-8')
84+ size = len(metadata_bytes)
85+
86+ upload_url = create_signed_upload_url(url, size, user_id, memobin_api_key)
87+
88+ response = requests.put(
89+ upload_url,
90+ data=metadata_bytes,
91+ headers={"Content-Type": "application/json"}
92+ )
93+
94+ if not response.ok:
95+ raise requests.RequestException("Failed to upload metadata to memobin")
96+
97+def download_from_memobin(url: str) -> Optional[dict]:
98+ """Download metadata from memobin.
99+
100+ Args:
101+ url: The URL to download from
102+
103+ Returns:
104+ The downloaded metadata as a dictionary, or None if not found
105+
106+ Raises:
107+ requests.RequestException: If the download fails for a reason other than 404
108+ """
109+ response = None
110+ try:
111+ response = requests.get(url)
112+ if response.status_code == 404:
113+ return None
114+ response.raise_for_status()
115+ return response.json()
116+ except requests.RequestException as e:
117+ if response and response.status_code == 404:
118+ return None
119+ raise e
zia_benchmark/src/zia_benchmark/algorithms/__init__.pyadded+5−0View file
@@ -0,0 +1,5 @@
1+from .zlib import algorithms as zlib_algorithms
2+from .zstd import algorithms as zstd_algorithms
3+from .simple_ans import algorithms as simple_ans_algorithms
4+
5+algorithms = zlib_algorithms + zstd_algorithms + simple_ans_algorithms
zia_benchmark/src/zia_benchmark/algorithms/simple_ans/__init__.pyadded+70−0View file
@@ -0,0 +1,70 @@
1+import numpy as np
2+
3+
4+def simple_ans_encode(x: np.ndarray) -> bytes:
5+ from simple_ans import ans_encode
6+ assert x.ndim == 1
7+ encoded = ans_encode(x)
8+ if x.dtype == np.uint8:
9+ dtype_code = 0
10+ elif x.dtype == np.uint16:
11+ dtype_code = 1
12+ elif x.dtype == np.uint32:
13+ dtype_code = 2
14+ elif x.dtype == np.int16:
15+ dtype_code = 3
16+ elif x.dtype == np.int32:
17+ dtype_code = 4
18+ else:
19+ raise ValueError(f"Unsupported dtype: {x.dtype}")
20+ header = [
21+ dtype_code,
22+ encoded.num_bits,
23+ encoded.signal_length,
24+ encoded.state,
25+ len(encoded.symbol_counts)
26+ ] + [c for c in encoded.symbol_counts] + [v for v in encoded.symbol_values]
27+ header_bytes = np.array(header, dtype=np.int64).tobytes()
28+ header_size = np.uint32(len(header_bytes))
29+ return header_size.tobytes() + header_bytes + encoded.bitstream
30+
31+
32+def simple_ans_decode(x: bytes, dtype: str) -> np.ndarray:
33+ from simple_ans import ans_decode, EncodedSignal
34+ header_size = np.frombuffer(x[:4], dtype=np.uint32)[0]
35+ header = np.frombuffer(x[4:4 + header_size], dtype=np.int64)
36+ dtype_code, num_bits, signal_length, state, num_symbols = header[:5]
37+ symbol_counts = header[5:5 + num_symbols]
38+ symbol_values = header[5 + num_symbols:]
39+ bitstream = x[4 + header_size:]
40+ if dtype_code == 0:
41+ assert dtype == 'uint8'
42+ elif dtype_code == 1:
43+ assert dtype == 'uint16'
44+ elif dtype_code == 2:
45+ assert dtype == 'uint32'
46+ elif dtype_code == 3:
47+ assert dtype == 'int16'
48+ elif dtype_code == 4:
49+ assert dtype == 'int32'
50+ else:
51+ raise ValueError(f"Unsupported dtype code: {dtype_code}")
52+
53+ encoded = EncodedSignal(
54+ num_bits=int(num_bits),
55+ signal_length=int(signal_length),
56+ state=int(state),
57+ symbol_counts=symbol_counts.astype(np.uint32),
58+ symbol_values=symbol_values.astype(dtype),
59+ bitstream=bitstream
60+ )
61+ return ans_decode(encoded)
62+
63+algorithms = [
64+ {
65+ 'name': 'simple-ans',
66+ 'version': '1',
67+ 'encode': lambda x: simple_ans_encode(x),
68+ 'decode': lambda x, dtype: simple_ans_decode(x, dtype)
69+ }
70+]
zia_benchmark/src/zia_benchmark/algorithms/zlib/__init__.pyadded+48−0View file
@@ -0,0 +1,48 @@
1+import numpy as np
2+
3+
4+def zlib_encode(x: np.ndarray, level: int) -> bytes:
5+ import zlib
6+ assert x.ndim == 1
7+ buf = x.tobytes()
8+ compressed = zlib.compress(buf, level=level)
9+ return compressed
10+
11+def zlib_decode(x: bytes, dtype: str) -> np.ndarray:
12+ import zlib
13+ buf = zlib.decompress(x)
14+ y = np.frombuffer(buf, dtype=dtype)
15+ return y
16+
17+algorithms = [
18+ {
19+ 'name': 'zlib-1',
20+ 'version': '1',
21+ 'encode': lambda x: zlib_encode(x, level=1),
22+ 'decode': lambda x, dtype: zlib_decode(x, dtype)
23+ },
24+ {
25+ 'name': 'zlib-3',
26+ 'version': '1',
27+ 'encode': lambda x: zlib_encode(x, level=3),
28+ 'decode': lambda x, dtype: zlib_decode(x, dtype)
29+ },
30+ {
31+ 'name': 'zlib-5',
32+ 'version': '1',
33+ 'encode': lambda x: zlib_encode(x, level=5),
34+ 'decode': lambda x, dtype: zlib_decode(x, dtype)
35+ },
36+ {
37+ 'name': 'zlib-7',
38+ 'version': '1',
39+ 'encode': lambda x: zlib_encode(x, level=7),
40+ 'decode': lambda x, dtype: zlib_decode(x, dtype)
41+ },
42+ {
43+ 'name': 'zlib-9',
44+ 'version': '1',
45+ 'encode': lambda x: zlib_encode(x, level=9),
46+ 'decode': lambda x, dtype: zlib_decode(x, dtype)
47+ }
48+]
zia_benchmark/src/zia_benchmark/algorithms/zstd/__init__.pyadded+62−0View file
@@ -0,0 +1,62 @@
1+import numpy as np
2+
3+
4+def zstd_encode(x: np.ndarray, level: int) -> bytes:
5+ import zstandard as zstd
6+ assert x.ndim == 1
7+ buf = x.tobytes()
8+ compressor = zstd.ZstdCompressor(level=level)
9+ compressed = compressor.compress(buf)
10+ return compressed
11+
12+def zstd_decode(x: bytes, dtype: str) -> np.ndarray:
13+ import zstandard as zstd
14+ decompressor = zstd.ZstdDecompressor()
15+ buf = decompressor.decompress(x)
16+ y = np.frombuffer(buf, dtype=dtype)
17+ return y
18+
19+algorithms = [
20+ {
21+ 'name': 'zstd-4',
22+ 'version': '1',
23+ 'encode': lambda x: zstd_encode(x, level=4),
24+ 'decode': lambda x, dtype: zstd_decode(x, dtype)
25+ },
26+ {
27+ 'name': 'zstd-7',
28+ 'version': '1',
29+ 'encode': lambda x: zstd_encode(x, level=7),
30+ 'decode': lambda x, dtype: zstd_decode(x, dtype)
31+ },
32+ {
33+ 'name': 'zstd-10',
34+ 'version': '1',
35+ 'encode': lambda x: zstd_encode(x, level=10),
36+ 'decode': lambda x, dtype: zstd_decode(x, dtype)
37+ },
38+ {
39+ 'name': 'zstd-13',
40+ 'version': '1',
41+ 'encode': lambda x: zstd_encode(x, level=13),
42+ 'decode': lambda x, dtype: zstd_decode(x, dtype)
43+ },
44+ {
45+ 'name': 'zstd-16',
46+ 'version': '1',
47+ 'encode': lambda x: zstd_encode(x, level=16),
48+ 'decode': lambda x, dtype: zstd_decode(x, dtype)
49+ },
50+ {
51+ 'name': 'zstd-19',
52+ 'version': '1',
53+ 'encode': lambda x: zstd_encode(x, level=19),
54+ 'decode': lambda x, dtype: zstd_decode(x, dtype)
55+ },
56+ {
57+ 'name': 'zstd-22',
58+ 'version': '1',
59+ 'encode': lambda x: zstd_encode(x, level=22),
60+ 'decode': lambda x, dtype: zstd_decode(x, dtype)
61+ }
62+]
zia_benchmark/src/zia_benchmark/datasets/__init__.pyadded+4−0View file
@@ -0,0 +1,4 @@
1+from .bernoulli import datasets as bernoulli_datasets
2+from .gaussian import datasets as gaussian_datasets
3+
4+datasets = bernoulli_datasets + gaussian_datasets
zia_benchmark/src/zia_benchmark/datasets/bernoulli/__init__.pyadded+35−0View file
@@ -0,0 +1,35 @@
1+import numpy as np
2+
3+
4+def create_bernoulli(*, n_samples: int, p: float, seed: int) -> np.ndarray:
5+ rng = np.random.default_rng(seed)
6+ x = rng.binomial(1, p, n_samples).astype(np.uint8)
7+ return x
8+
9+datasets = [
10+ {
11+ 'name': 'bernoulli-0.1',
12+ 'version': '1',
13+ 'create': lambda: create_bernoulli(n_samples=1_000_000, p=0.1, seed=0)
14+ },
15+ {
16+ 'name': 'bernoulli-0.2',
17+ 'version': '1',
18+ 'create': lambda: create_bernoulli(n_samples=1_000_000, p=0.2, seed=0)
19+ },
20+ {
21+ 'name': 'bernoulli-0.3',
22+ 'version': '1',
23+ 'create': lambda: create_bernoulli(n_samples=1_000_000, p=0.3, seed=0)
24+ },
25+ {
26+ 'name': 'bernoulli-0.4',
27+ 'version': '1',
28+ 'create': lambda: create_bernoulli(n_samples=1_000_000, p=0.4, seed=0)
29+ },
30+ {
31+ 'name': 'bernoulli-0.5',
32+ 'version': '1',
33+ 'create': lambda: create_bernoulli(n_samples=1_000_000, p=0.5, seed=0)
34+ }
35+]
zia_benchmark/src/zia_benchmark/datasets/gaussian/__init__.pyadded+35−0View file
@@ -0,0 +1,35 @@
1+import numpy as np
2+
3+
4+def create_gaussian(*, n_samples: int, stddev: float, seed: int) -> np.ndarray:
5+ rng = np.random.default_rng(seed)
6+ x = np.round(rng.normal(0, stddev, n_samples)).astype(np.int16)
7+ return x
8+
9+datasets = [
10+ {
11+ 'name': 'gaussian-1',
12+ 'version': '1',
13+ 'create': lambda: create_gaussian(n_samples=1_000_000, stddev=1, seed=0)
14+ },
15+ {
16+ 'name': 'gaussian-2',
17+ 'version': '1',
18+ 'create': lambda: create_gaussian(n_samples=1_000_000, stddev=2, seed=0)
19+ },
20+ {
21+ 'name': 'gaussian-3',
22+ 'version': '1',
23+ 'create': lambda: create_gaussian(n_samples=1_000_000, stddev=3, seed=0)
24+ },
25+ {
26+ 'name': 'gaussian-5',
27+ 'version': '1',
28+ 'create': lambda: create_gaussian(n_samples=1_000_000, stddev=5, seed=0)
29+ },
30+ {
31+ 'name': 'gaussian-8',
32+ 'version': '1',
33+ 'create': lambda: create_gaussian(n_samples=1_000_000, stddev=8, seed=0)
34+ }
35+]
zia_benchmark/src/zia_benchmark/run_benchmarks.pyadded+189−0View file
@@ -0,0 +1,189 @@
1+import time
2+import json
3+import os
4+from typing import Dict, Any, Tuple
5+import numpy as np
6+from statistics import median
7+from .algorithms import algorithms
8+from .datasets import datasets
9+from ._memobin import construct_memobin_url, upload_to_memobin, download_from_memobin
10+
11+
12+system_version = 'v3'
13+
14+def run_benchmarks(cache_dir: str = '.benchmark_cache', verbose: bool = True) -> Dict[str, Any]:
15+ """Run all benchmarks, with caching based on algorithm and dataset versions.
16+
17+ Results are stored in separate directories for each dataset/algorithm combination:
18+ cache_dir/
19+ dataset_name/
20+ algorithm_name/
21+ metadata.json # Contains algorithm version, dataset version, and results
22+ compressed.dat # The actual compressed data
23+
24+ Args:
25+ cache_dir: Directory to store cached results
26+
27+ Returns:
28+ Dictionary containing benchmark results and metadata
29+ """
30+ print("\n=== Starting Benchmark Run ===")
31+ print(f"Cache directory: {cache_dir}")
32+
33+ os.makedirs(cache_dir, exist_ok=True)
34+
35+ results = []
36+ print("\nRunning benchmarks for all dataset-algorithm combinations...")
37+
38+ # Run benchmarks for each dataset and algorithm combination
39+ for dataset in datasets:
40+ print(f"\n--- Dataset: {dataset['name']} ---")
41+ # Create dataset once for all algorithms
42+ data = dataset['create']()
43+ dtype = str(data.dtype)
44+ original_size = len(data.tobytes())
45+ print(f"Created dataset: shape={data.shape}, dtype={dtype}")
46+ print(f"Original size: {original_size:,} bytes")
47+
48+ for algorithm in algorithms:
49+ alg_name = algorithm['name']
50+ print(f"\nTesting algorithm: {alg_name}")
51+
52+ # Check if we can use cached result
53+ test_dir = os.path.join(cache_dir, dataset['name'], alg_name)
54+ metadata_file = os.path.join(test_dir, 'metadata.json')
55+ compressed_file = os.path.join(test_dir, 'compressed.dat')
56+
57+ # First try local cache
58+ cached_data = None
59+ if os.path.exists(metadata_file) and os.path.exists(compressed_file):
60+ with open(metadata_file, 'r') as f:
61+ cached_data = json.load(f)
62+
63+ # If not in local cache, try memobin
64+ if cached_data is None:
65+ memobin_url = construct_memobin_url(
66+ alg_name, dataset['name'],
67+ algorithm['version'], dataset['version'],
68+ system_version
69+ )
70+ if verbose:
71+ print(" Looking for cached result in memobin...")
72+ cached_data = download_from_memobin(memobin_url)
73+ if cached_data is not None:
74+ if verbose:
75+ print(" Found result in memobin, saving locally...")
76+ # Save to local cache
77+ os.makedirs(test_dir, exist_ok=True)
78+ with open(metadata_file, 'w') as f:
79+ json.dump(cached_data, f, indent=2)
80+
81+ if cached_data is not None and (
82+ cached_data['result']['algorithm_version'] == algorithm['version'] and
83+ cached_data['result']['dataset_version'] == dataset['version'] and
84+ cached_data['result'].get('system_version', '') == system_version
85+ ):
86+ print(" Using cached result:")
87+ results.append(cached_data['result'])
88+ continue
89+
90+ print(" Running new benchmark...")
91+
92+ def run_timed_trials(operation, *args) -> Tuple[float, float]:
93+ """Run multiple trials of an operation until total time exceeds 1 second.
94+ Returns (median_time, mb_per_sec)"""
95+ times = []
96+ total_time = 0
97+ array_size_mb = data.nbytes / (1024 * 1024) # Convert to MB
98+
99+ while total_time < 1.0:
100+ start_time = time.perf_counter()
101+ _ = operation(*args) # Execute operation but discard result
102+ trial_time = time.perf_counter() - start_time
103+ times.append(trial_time)
104+ total_time += trial_time
105+
106+ median_time = median(times)
107+ mb_per_sec = array_size_mb / median_time
108+ return median_time, mb_per_sec
109+
110+ # Measure encoding with multiple trials
111+ encode_time, encode_mb_per_sec = run_timed_trials(algorithm['encode'], data)
112+ encoded = algorithm['encode'](data) # One final encode to get the result
113+ compressed_size = len(encoded)
114+ compression_ratio = original_size / compressed_size
115+ print(" Compression complete:")
116+ print(f" Compressed size: {compressed_size:,} bytes")
117+ print(f" Compression ratio: {compression_ratio:.2f}x")
118+ print(f" Encode time: {encode_time*1000:.2f}ms")
119+ print(f" Encode throughput: {encode_mb_per_sec:.2f} MB/s")
120+
121+ print(" Verifying decompression...")
122+ # Measure decoding with multiple trials
123+ decode_time, decode_mb_per_sec = run_timed_trials(algorithm['decode'], encoded, dtype)
124+ decoded = algorithm['decode'](encoded, dtype) # One final decode to verify
125+ print(f" Decode time: {decode_time*1000:.2f}ms")
126+ print(f" Decode throughput: {decode_mb_per_sec:.2f} MB/s")
127+
128+ # Verify correctness
129+ if not np.array_equal(data, decoded):
130+ raise ValueError(
131+ f"Decompression verification failed for {alg_name} on {dataset['name']}"
132+ )
133+ print(" Verification successful!")
134+
135+ # Store result
136+ result = {
137+ 'dataset': dataset['name'],
138+ 'algorithm': alg_name,
139+ 'algorithm_version': algorithm['version'],
140+ 'dataset_version': dataset['version'],
141+ 'system_version': system_version,
142+ 'compression_ratio': compression_ratio,
143+ 'encode_time': encode_time,
144+ 'decode_time': decode_time,
145+ 'encode_mb_per_sec': encode_mb_per_sec,
146+ 'decode_mb_per_sec': decode_mb_per_sec,
147+ 'original_size': original_size,
148+ 'compressed_size': compressed_size,
149+ 'array_shape': data.shape,
150+ 'array_dtype': dtype,
151+ 'timestamp': time.time()
152+ }
153+ results.append(result)
154+
155+ # Save result and compressed data
156+ os.makedirs(test_dir, exist_ok=True)
157+ cache_data = {
158+ 'result': result
159+ }
160+ with open(metadata_file, 'w') as f:
161+ json.dump(cache_data, f, indent=2)
162+ with open(compressed_file, 'wb') as f:
163+ f.write(encoded)
164+ print(f" Results saved to: {test_dir}")
165+
166+ # Upload to memobin if API key is set
167+ memobin_api_key = os.environ.get('MEMOBIN_API_KEY')
168+ if memobin_api_key:
169+ if verbose:
170+ print(" Uploading results to memobin...")
171+ try:
172+ memobin_url = construct_memobin_url(
173+ alg_name, dataset['name'],
174+ algorithm['version'], dataset['version'],
175+ system_version
176+ )
177+ upload_to_memobin(
178+ cache_data,
179+ memobin_url,
180+ os.environ.get('MEMOBIN_USER_ID', 'default'),
181+ memobin_api_key
182+ )
183+ if verbose:
184+ print(" Successfully uploaded to memobin")
185+ except Exception as e:
186+ print(f" Warning: Failed to upload to memobin: {str(e)}")
187+
188+ print("\n=== Benchmark Run Complete ===\n")
189+ return {'results': results}