/ concept-collection / ephys_compression_tests
Sign in
concept-collection / ephys_compression_tests
127 lines · 4.3 KBBlameHistoryRaw
1import os
2import json
3from typing import Optional, Dict, Any
4import numpy as np
5from ._memobin import (
6 construct_memobin_url,
7 download_from_memobin,
8)
11def check_cached_result(
12 cache_dir: str,
13 dataset_name: str,
14 algorithm_name: str,
15 algorithm_version: str,
16 dataset_version: str,
17 system_version: str,
18 force: bool = False,
19 verbose: bool = True,
20) -> Optional[Dict[str, Any]]:
21 """Check for cached benchmark results locally and in memobin.
23 Args:
24 cache_dir: Directory containing cached results
25 dataset_name: Name of the dataset
26 algorithm_name: Name of the algorithm
27 algorithm_version: Version of the algorithm
28 dataset_version: Version of the dataset
29 system_version: Version of the system
30 force: If True, ignore cached results
31 verbose: Whether to print progress messages
33 Returns:
34 Cached result dictionary if found and valid, None otherwise
35 """
36 test_dir = os.path.join(cache_dir, dataset_name, algorithm_name)
37 metadata_file = os.path.join(test_dir, "metadata.json")
39 # First try local cache (unless force flag is set)
40 cached_data = None
41 if not force and os.path.exists(metadata_file):
42 with open(metadata_file, "r") as f:
43 cached_data = json.load(f)
44 # if versions do not match, then set to None
45 if isinstance(cached_data, dict) and "result" in cached_data:
46 result = cached_data["result"]
47 if (
48 result["algorithm_version"] != algorithm_version
49 or result["dataset_version"] != dataset_version
50 or result.get("system_version", "") != system_version
51 ):
52 cached_data = None
54 # If not in local cache, try memobin (unless force flag is set)
55 if cached_data is None and not force:
56 memobin_url = construct_memobin_url(
57 algorithm_name,
58 dataset_name,
59 algorithm_version,
60 dataset_version,
61 system_version,
62 "metadata.json",
63 )
64 if verbose:
65 print(" Looking for cached result in memobin...")
66 cached_data = download_from_memobin(memobin_url)
67 if cached_data is not None:
68 if verbose:
69 print(" Found result in memobin, saving locally...")
70 # Save to local cache
71 os.makedirs(test_dir, exist_ok=True)
72 with open(metadata_file, "w") as f:
73 json.dump(cached_data, f, indent=2)
75 if (
76 cached_data is not None
77 and isinstance(cached_data, dict)
78 and "result" in cached_data
79 ):
80 result = cached_data["result"]
81 if (
82 isinstance(result, dict)
83 and result.get("algorithm_version") == algorithm_version
84 and result.get("dataset_version") == dataset_version
85 and result.get("system_version", "") == system_version
86 ):
87 result["cache_status"] = "cached"
88 return result
90 return None
93def save_result_to_cache(
94 result: Dict[str, Any],
95 encoded_data: bytes,
96 cache_dir: str,
97 dataset_name: str,
98 algorithm_name: str,
99 reconstructed_data: Optional[np.ndarray] = None,
100) -> None:
101 """Save benchmark result and compressed data to cache.
103 Args:
104 result: Benchmark result dictionary
105 encoded_data: Compressed data bytes
106 cache_dir: Directory to store cached results
107 dataset_name: Name of the dataset
108 algorithm_name: Name of the algorithm
109 reconstructed_data: Optional reconstructed array for lossy algorithms
110 """
111 test_dir = os.path.join(cache_dir, dataset_name, algorithm_name)
112 metadata_file = os.path.join(test_dir, "metadata.json")
113 compressed_file = os.path.join(test_dir, "compressed.dat")
114 reconstructed_file = os.path.join(test_dir, "reconstructed.dat")
116 os.makedirs(test_dir, exist_ok=True)
117 cache_data = {"result": result}
119 with open(metadata_file, "w") as f:
120 json.dump(cache_data, f, indent=2)
121 with open(compressed_file, "wb") as f:
122 f.write(encoded_data)
124 # Save reconstructed data for lossy algorithms
125 if reconstructed_data is not None:
126 with open(reconstructed_file, "wb") as f:
127 f.write(reconstructed_data.tobytes())
moveopenescclose