1from typing import Any, Tuple, Callable, Dict
2from statistics import median
3import time
4import numpy as np
7def run_timed_trials(
8 data: np.ndarray, operation: Callable, *args
9) -> Tuple[float, float, Any]:
10 """Run multiple trials of an operation until total time exceeds 1 second.
12 Args:
13 data: Input numpy array for calculating throughput
14 operation: Function to benchmark
15 *args: Arguments to pass to the operation
17 Returns:
18 Tuple containing:
19 - median_time: Median execution time across trials
20 - mb_per_sec: Throughput in MB/s
21 - result: Result from the last trial execution
22 """
23 times = []
24 total_time = 0
25 array_size_mb = data.nbytes / (1024 * 1024) # Convert to MB
27 operation(
28 *args
29 ) # execute once prior to timing in case there's any initial overhead
31 ret = None
32 while total_time < 1.0:
33 start_time = time.perf_counter()
34 ret = operation(*args) # Execute operation
35 trial_time = time.perf_counter() - start_time
36 times.append(trial_time)
37 total_time += trial_time
39 median_time = median(times)
40 mb_per_sec = array_size_mb / median_time
41 return median_time, mb_per_sec, ret
44def run_compression_benchmark(
45 data: np.ndarray,
46 algorithm_name: str,
47 encode_fn: Callable,
48 decode_fn: Callable,
49 verbose: bool = True,
50 lossy: bool = False,
51) -> Tuple[Dict[str, Any], bytes, np.ndarray]:
52 """Run compression and decompression benchmarks for an algorithm.
54 Args:
55 data: Input numpy array to compress
56 algorithm_name: Name of the algorithm being benchmarked
57 encode_fn: Compression function
58 decode_fn: Decompression function
59 verbose: Whether to print progress messages
60 lossy: Whether the algorithm is lossy
62 Returns:
63 Tuple containing:
64 - result: Dictionary with benchmark metrics
65 - encoded: Compressed data bytes
66 - decoded: Decompressed data array
67 """
68 if data.ndim == 1:
69 data = data[:, np.newaxis]
70 original_size = len(data.tobytes())
71 dtype = str(data.dtype)
73 if verbose:
74 print(" Encoding...")
75 encode_time, encode_mb_per_sec, encoded = run_timed_trials(data, encode_fn, data)
76 compressed_size = len(encoded)
77 compression_ratio = original_size / compressed_size
79 if verbose:
80 print(" Compression complete:")
81 print(f" Compressed size: {compressed_size:,} bytes")
82 print(f" Compression ratio: {compression_ratio:.2f}x")
83 print(f" Encode time: {encode_time*1000:.2f}ms")
84 print(f" Encode throughput: {encode_mb_per_sec:.2f} MB/s")
85 print(" Decoding...")
87 decode_time, decode_mb_per_sec, decoded = run_timed_trials(
88 data, decode_fn, encoded, dtype, data.shape
89 )
91 if verbose:
92 print(f" Decode time: {decode_time*1000:.2f}ms")
93 print(f" Decode throughput: {decode_mb_per_sec:.2f} MB/s")
95 # Verify correctness
96 if len(data) != len(decoded):
97 raise ValueError(
98 f"Decompression failed: decoded length {len(decoded)} != original length {len(data)}"
99 )
101 if not lossy:
102 if not np.array_equal(data, decoded):
103 print(data[:100])
104 print(decoded[:100])
105 for j in range(len(data)):
106 if data[j] != decoded[j]:
107 print(f"Error at index {j}: {data[j]} != {decoded[j]}")
108 break
109 raise ValueError(f"Decompression verification failed for {algorithm_name}")
110 rmse = 0.0
111 max_error = 0.0
112 else:
113 # compute RMSE and max error
114 rmse = float(np.sqrt(np.mean((data - decoded) ** 2)))
115 max_error = float(np.max(np.abs(data - decoded)))
116 print(f" RMSE: {rmse:.4f}, Max error: {max_error:.4f}")
118 if verbose:
119 print(" Verification successful!")
121 result = {
122 "compression_ratio": compression_ratio,
123 "encode_time": encode_time,
124 "decode_time": decode_time,
125 "encode_mb_per_sec": encode_mb_per_sec,
126 "decode_mb_per_sec": decode_mb_per_sec,
127 "original_size": original_size,
128 "compressed_size": compressed_size,
129 "array_shape": data.shape,
130 "array_dtype": dtype,
131 "timestamp": time.time(),
132 "cache_status": "new",
133 "rmse": rmse,
134 "max_error": max_error,
135 }
137 return result, encoded, decoded