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]:
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
61 Returns:
62 Tuple containing:
63 - result: Dictionary with benchmark metrics
64 - encoded: Compressed data bytes
65 """
66 original_size = len(data.tobytes())
67 dtype = str(data.dtype)
69 if verbose:
70 print(" Encoding...")
71 encode_time, encode_mb_per_sec, encoded = run_timed_trials(data, encode_fn, data)
72 compressed_size = len(encoded)
73 compression_ratio = original_size / compressed_size
75 if verbose:
76 print(" Compression complete:")
77 print(f" Compressed size: {compressed_size:,} bytes")
78 print(f" Compression ratio: {compression_ratio:.2f}x")
79 print(f" Encode time: {encode_time*1000:.2f}ms")
80 print(f" Encode throughput: {encode_mb_per_sec:.2f} MB/s")
81 print(" Decoding...")
83 decode_time, decode_mb_per_sec, decoded = run_timed_trials(
84 data, decode_fn, encoded, dtype, data.shape
85 )
87 if verbose:
88 print(f" Decode time: {decode_time*1000:.2f}ms")
89 print(f" Decode throughput: {decode_mb_per_sec:.2f} MB/s")
91 # Verify correctness
92 if len(data) != len(decoded):
93 raise ValueError(
94 f"Decompression failed: decoded length {len(decoded)} != original length {len(data)}"
95 )
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}")
114 if verbose:
115 print(" Verification successful!")
117 result = {
118 "compression_ratio": compression_ratio,
119 "encode_time": encode_time,
120 "decode_time": decode_time,
121 "encode_mb_per_sec": encode_mb_per_sec,
122 "decode_mb_per_sec": decode_mb_per_sec,
123 "original_size": original_size,
124 "compressed_size": compressed_size,
125 "array_shape": data.shape,
126 "array_dtype": dtype,
127 "timestamp": time.time(),
128 "cache_status": "new",
129 "rmse": rmse,
130 "max_error": max_error,
131 }
133 return result, encoded