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 if data.ndim == 1:
67 data = data[:, np.newaxis]
68 original_size = len(data.tobytes())
69 dtype = str(data.dtype)
71 if verbose:
72 print(" Encoding...")
73 encode_time, encode_mb_per_sec, encoded = run_timed_trials(data, encode_fn, data)
74 compressed_size = len(encoded)
75 compression_ratio = original_size / compressed_size
77 if verbose:
78 print(" Compression complete:")
79 print(f" Compressed size: {compressed_size:,} bytes")
80 print(f" Compression ratio: {compression_ratio:.2f}x")
81 print(f" Encode time: {encode_time*1000:.2f}ms")
82 print(f" Encode throughput: {encode_mb_per_sec:.2f} MB/s")
83 print(" Decoding...")
85 decode_time, decode_mb_per_sec, decoded = run_timed_trials(
86 data, decode_fn, encoded, dtype, data.shape
87 )
89 if verbose:
90 print(f" Decode time: {decode_time*1000:.2f}ms")
91 print(f" Decode throughput: {decode_mb_per_sec:.2f} MB/s")
93 # Verify correctness
94 if len(data) != len(decoded):
95 raise ValueError(
96 f"Decompression failed: decoded length {len(decoded)} != original length {len(data)}"
97 )
99 if not lossy:
100 if not np.array_equal(data, decoded):
101 print(data[:100])
102 print(decoded[:100])
103 for j in range(len(data)):
104 if data[j] != decoded[j]:
105 print(f"Error at index {j}: {data[j]} != {decoded[j]}")
106 break
107 raise ValueError(f"Decompression verification failed for {algorithm_name}")
108 rmse = 0.0
109 max_error = 0.0
110 else:
111 # compute RMSE and max error
112 rmse = float(np.sqrt(np.mean((data - decoded) ** 2)))
113 max_error = float(np.max(np.abs(data - decoded)))
114 print(f" RMSE: {rmse:.4f}, Max error: {max_error:.4f}")
116 if verbose:
117 print(" Verification successful!")
119 result = {
120 "compression_ratio": compression_ratio,
121 "encode_time": encode_time,
122 "decode_time": decode_time,
123 "encode_mb_per_sec": encode_mb_per_sec,
124 "decode_mb_per_sec": decode_mb_per_sec,
125 "original_size": original_size,
126 "compressed_size": compressed_size,
127 "array_shape": data.shape,
128 "array_dtype": dtype,
129 "timestamp": time.time(),
130 "cache_status": "new",
131 "rmse": rmse,
132 "max_error": max_error,
133 }
135 return result, encoded