use numba for markov decode
6 changed files+21−9
devel/markov_bench.pymodified+3−1View file
@@ -243,7 +243,9 @@ def benchmark_markov_model(N=1_000_000, M=5, seed=0):
243243 print("\n--- Prediction (Naive) ---")
244244 print(f" Predict: {dt_pred_naive:.3f}s, ~{throughput_pred_naive:.2f} MB/s")
245245
246- # 2) Numba-Accelerated Prediction
246+ # Warm up Numba
247+ _ = predict_numba(c_vec, seed_vals, len(seed_vals) + 1)
248+
247249 t0 = time.perf_counter()
248250 pred_numba = predict_numba(c_vec, seed_vals, N)
249251 dt_pred_numba = time.perf_counter() - t0
zia_benchmark/setup.pymodified+2−1View file
@@ -13,7 +13,8 @@ setup(
1313 "requests",
1414 "lindi",
1515 "brotli",
16- "click"
16+ "click",
17+ "numba"
1718 ],
1819 python_requires=">=3.8",
1920 entry_points={
zia_benchmark/src/zia_benchmark/algorithms/simple_ans/__init__.pymodified+1−1View file
@@ -257,7 +257,7 @@ algorithms = [
257257 },
258258 {
259259 "name": "simple-ans-markov",
260- "version": "3",
260+ "version": "4",
261261 "encode": lambda x: simple_ans_markov_encode(x),
262262 "decode": lambda x, dtype: simple_ans_markov_decode(x, dtype),
263263 "description": "ANS compression with Markov prediction for exploiting temporal correlations in the data.",
zia_benchmark/src/zia_benchmark/algorithms/simple_ans/markov.pymodified+2−0View file
@@ -1,4 +1,5 @@
11 import numpy as np
2+import numba
23 from zia_benchmark._analysis import linear_fit
34
45
@@ -34,6 +35,7 @@ def markov_predict(x: np.ndarray, M: int) -> tuple:
3435 return coeffs, initial, residuals
3536
3637
38+@numba.jit(nopython=True)
3739 def markov_reconstruct(
3840 coeffs: np.ndarray, initial: np.ndarray, resid: np.ndarray
3941 ) -> np.ndarray:
zia_benchmark/src/zia_benchmark/cli.pymodified+3−1View file
@@ -99,7 +99,8 @@ def list():
9999 type=click.Path(),
100100 )
101101 @click.option("--quiet", "-q", is_flag=True, help="Reduce output verbosity")
102-def run(algorithm, dataset, cache_dir, quiet):
102+@click.option("--force", "-f", is_flag=True, help="Force re-run without using cache")
103+def run(algorithm, dataset, cache_dir, quiet, force):
103104 """Run benchmarks with specified options"""
104105 # Filter algorithms and datasets
105106 filtered_algorithms = filter_algorithms(algorithm)
@@ -120,6 +121,7 @@ def run(algorithm, dataset, cache_dir, quiet):
120121 verbose=not quiet,
121122 selected_algorithms=filtered_algorithms,
122123 selected_datasets=filtered_datasets,
124+ force=force,
123125 )
124126
125127 # Print summary
zia_benchmark/src/zia_benchmark/run_benchmarks.pymodified+10−5View file
@@ -13,6 +13,10 @@ from ._memobin import (
1313 download_from_memobin,
1414 exists_in_memobin,
1515 )
16+from .algorithms.simple_ans.markov import markov_reconstruct
17+
18+# warm up the JIT
19+markov_reconstruct(np.array([1, 2, 3]), np.array([1, 2]), np.array([1]))
1620
1721
1822 system_version = "v5"
@@ -45,6 +49,7 @@ def run_benchmarks(
4549 verbose: bool = True,
4650 selected_algorithms: Optional[List[dict]] = None,
4751 selected_datasets: Optional[List[dict]] = None,
52+ force: bool = False,
4853 ) -> Dict[str, Any]:
4954 """Run all benchmarks, with caching based on algorithm and dataset versions.
5055
@@ -99,14 +104,14 @@ def run_benchmarks(
99104
100105 print(f"\nTesting algorithm: {alg_name} (tags: {alg_tags})")
101106
102- # Check if we can use cached result
107+ # Check if we can use cached result (unless force flag is set)
103108 test_dir = os.path.join(cache_dir, dataset["name"], alg_name)
104109 metadata_file = os.path.join(test_dir, "metadata.json")
105110 compressed_file = os.path.join(test_dir, "compressed.dat")
106111
107- # First try local cache
112+ # First try local cache (unless force flag is set)
108113 cached_data = None
109- if os.path.exists(metadata_file):
114+ if not force and os.path.exists(metadata_file):
110115 with open(metadata_file, "r") as f:
111116 cached_data = json.load(f)
112117 # if versions do not match, then set to None
@@ -119,8 +124,8 @@ def run_benchmarks(
119124 ):
120125 cached_data = None
121126
122- # If not in local cache, try memobin
123- if cached_data is None:
127+ # If not in local cache, try memobin (unless force flag is set)
128+ if cached_data is None and not force:
124129 memobin_url = construct_memobin_url(
125130 alg_name,
126131 dataset["name"],