1import numpy as np
2from ..ans.markov_reconstruct import markov_reconstruct as markov_reconstruct_cpp
3from ..ans.markov_predict import markov_predict as markov_predict_cpp
4from ..ans.get_run_lengths import get_run_lengths
7SOURCE_FILE = "zstd/__init__.py"
10def zstd_delta_encode(x: np.ndarray, level: int) -> bytes:
11 import zstandard as zstd
13 assert x.ndim == 1
15 y = np.diff(x)
16 y = np.insert(y, 0, x[0])
17 buf = y.tobytes()
18 compressor = zstd.ZstdCompressor(level=level)
19 compressed = compressor.compress(buf)
20 return compressed
23def zstd_delta_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
24 import zstandard as zstd
26 assert len(shape) == 1
28 decompressor = zstd.ZstdDecompressor()
29 buf = decompressor.decompress(x)
30 y = np.frombuffer(buf, dtype=dtype)
31 return np.cumsum(y)
34def zstd_encode(x: np.ndarray, level: int) -> bytes:
35 import zstandard as zstd
37 buf = x.tobytes()
38 compressor = zstd.ZstdCompressor(level=level)
39 compressed = compressor.compress(buf)
40 return compressed
43def zstd_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
44 import zstandard as zstd
46 decompressor = zstd.ZstdDecompressor()
47 buf = decompressor.decompress(x)
48 y = np.frombuffer(buf, dtype=dtype)
49 return y.reshape(shape)
52def zstd_markov_encode(x: np.ndarray, level: int) -> bytes:
53 import zstandard as zstd
54 import struct
56 assert x.ndim == 1
57 coeffs, initial, resid = markov_predict_cpp(x, M=6, num_training_samples=10000)
59 # Convert coeffs and initial to bytes
60 coeffs_bytes = coeffs.tobytes()
61 initial_bytes = initial.tobytes()
63 # Create header with lengths
64 header = struct.pack("QQ", len(coeffs_bytes), len(initial_bytes))
66 # Compress residuals
67 resid_bytes = resid.tobytes()
68 compressor = zstd.ZstdCompressor(level=level)
69 compressed_resid = compressor.compress(resid_bytes)
71 # Combine all parts
72 return header + coeffs_bytes + initial_bytes + compressed_resid
75def zstd_markov_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
76 import zstandard as zstd
77 import struct
79 assert len(shape) == 1
81 # Extract header
82 header_size = struct.calcsize("QQ")
83 coeffs_len, initial_len = struct.unpack("QQ", x[:header_size])
85 # Extract coefficients and initial values
86 pos = header_size
87 coeffs = np.frombuffer(x[pos : pos + coeffs_len], dtype=np.float32)
88 pos += coeffs_len
89 initial = np.frombuffer(x[pos : pos + initial_len], dtype=dtype)
90 pos += initial_len
92 # Decompress residuals
93 decompressor = zstd.ZstdDecompressor()
94 resid_buf = decompressor.decompress(x[pos:])
95 resid = np.frombuffer(resid_buf, dtype=dtype)
97 # Reconstruct signal
98 output = markov_reconstruct_cpp(coeffs, initial, resid)
99 return output
102def zstd_markov_zrle_encode(x: np.ndarray, level: int) -> bytes:
103 import zstandard as zstd
104 import struct
106 assert x.ndim == 1
108 # Get run lengths for zero/non-zero sequences
109 run_lengths = get_run_lengths(x)
111 # Determine run length dtype code
112 if run_lengths.dtype == np.uint8:
113 run_length_dtype_code = 0
114 elif run_lengths.dtype == np.uint16:
115 run_length_dtype_code = 1
116 elif run_lengths.dtype == np.uint32:
117 run_length_dtype_code = 2
118 else:
119 raise ValueError(f"Unsupported run length dtype: {run_lengths.dtype}")
121 # Extract non-zero data
122 non_zero_arrays = []
123 array_pos = 0
124 i = 0
125 while i < len(run_lengths):
126 non_zero_len = int(run_lengths[i])
127 if non_zero_len > 0:
128 non_zero_arrays.append(x[array_pos : array_pos + non_zero_len])
129 array_pos += non_zero_len
130 i += 1
131 if i < len(run_lengths):
132 array_pos += int(run_lengths[i]) # Skip zeros
133 i += 1
135 non_zero_data = np.concatenate(non_zero_arrays)
137 # Apply Markov prediction on non-zero data
138 coeffs, initial, resid = markov_predict_cpp(
139 non_zero_data, M=6, num_training_samples=10000
140 )
142 # Convert data to bytes
143 coeffs_bytes = coeffs.tobytes()
144 initial_bytes = initial.tobytes()
145 run_lengths_bytes = run_lengths.tobytes()
147 # Create header with lengths and dtype code
148 header = struct.pack(
149 "QQQQB",
150 len(coeffs_bytes),
151 len(initial_bytes),
152 len(run_lengths_bytes),
153 len(run_lengths),
154 run_length_dtype_code,
155 )
157 # Compress residuals
158 resid_bytes = resid.tobytes()
159 compressor = zstd.ZstdCompressor(level=level)
160 compressed_resid = compressor.compress(resid_bytes)
162 # Combine all parts
163 return header + coeffs_bytes + initial_bytes + run_lengths_bytes + compressed_resid
166def zstd_markov_zrle_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
167 import zstandard as zstd
168 import struct
170 assert len(shape) == 1
172 # Extract header
173 header_size = struct.calcsize("QQQQB")
174 coeffs_len, initial_len, run_lengths_len, num_run_lengths, run_length_dtype_code = (
175 struct.unpack("QQQQB", x[:header_size])
176 )
178 # Extract components
179 pos = header_size
180 coeffs = np.frombuffer(x[pos : pos + coeffs_len], dtype=np.float32)
181 pos += coeffs_len
182 initial = np.frombuffer(x[pos : pos + initial_len], dtype=dtype)
183 pos += initial_len
185 # Get run lengths with proper dtype
186 if run_length_dtype_code == 0:
187 run_lengths = np.frombuffer(x[pos : pos + run_lengths_len], dtype=np.uint8)
188 elif run_length_dtype_code == 1:
189 run_lengths = np.frombuffer(x[pos : pos + run_lengths_len], dtype=np.uint16)
190 elif run_length_dtype_code == 2:
191 run_lengths = np.frombuffer(x[pos : pos + run_lengths_len], dtype=np.uint32)
192 else:
193 raise ValueError(f"Unsupported run length dtype code: {run_length_dtype_code}")
194 pos += run_lengths_len
196 assert len(run_lengths) == num_run_lengths
198 # Decompress residuals
199 decompressor = zstd.ZstdDecompressor()
200 resid_buf = decompressor.decompress(x[pos:])
201 resid = np.frombuffer(resid_buf, dtype=dtype)
203 # Reconstruct non-zero data
204 non_zero_data = markov_reconstruct_cpp(coeffs, initial, resid)
206 # Reconstruct full array using run lengths
207 non_zero_pos = 0
208 i = 0
209 segments = []
210 while i < len(run_lengths):
211 non_zero_len = int(run_lengths[i])
212 if non_zero_len > 0:
213 segment = non_zero_data[non_zero_pos : non_zero_pos + non_zero_len]
214 segments.append(segment)
215 non_zero_pos += non_zero_len
216 i += 1
217 if i < len(run_lengths):
218 segments.append(np.zeros(int(run_lengths[i]), dtype=non_zero_data.dtype))
219 i += 1
221 return np.concatenate(segments)
224algorithms = [
225 {
226 "name": "zstd-4",
227 "version": "1",
228 "encode": lambda x: zstd_encode(x, level=4),
229 "decode": lambda x, dtype, shape: zstd_decode(x, dtype, shape),
230 "description": "Zstandard compression at level 4 (fast compression).",
231 "tags": ["zstd"],
232 "source_file": SOURCE_FILE,
233 },
234 {
235 "name": "zstd-7",
236 "version": "1",
237 "encode": lambda x: zstd_encode(x, level=7),
238 "decode": lambda x, dtype, shape: zstd_decode(x, dtype, shape),
239 "description": "Zstandard compression at level 7 (balanced speed/compression).",
240 "tags": ["zstd"],
241 "source_file": SOURCE_FILE,
242 },
243 {
244 "name": "zstd-10",
245 "version": "1",
246 "encode": lambda x: zstd_encode(x, level=10),
247 "decode": lambda x, dtype, shape: zstd_decode(x, dtype, shape),
248 "description": "Zstandard compression at level 10 (better compression).",
249 "tags": ["zstd"],
250 "source_file": SOURCE_FILE,
251 },
252 {
253 "name": "zstd-13",
254 "version": "1",
255 "encode": lambda x: zstd_encode(x, level=13),
256 "decode": lambda x, dtype, shape: zstd_decode(x, dtype, shape),
257 "description": "Zstandard compression at level 13 (high compression).",
258 "tags": ["zstd"],
259 "source_file": SOURCE_FILE,
260 },
261 {
262 "name": "zstd-16",
263 "version": "1",
264 "encode": lambda x: zstd_encode(x, level=16),
265 "decode": lambda x, dtype, shape: zstd_decode(x, dtype, shape),
266 "description": "Zstandard compression at level 16 (very high compression).",
267 "tags": ["zstd"],
268 "source_file": SOURCE_FILE,
269 },
270 {
271 "name": "zstd-19",
272 "version": "1",
273 "encode": lambda x: zstd_encode(x, level=19),
274 "decode": lambda x, dtype, shape: zstd_decode(x, dtype, shape),
275 "description": "Zstandard compression at level 19 (ultra high compression).",
276 "tags": ["zstd"],
277 "source_file": SOURCE_FILE,
278 },
279 {
280 "name": "zstd-22",
281 "version": "1",
282 "encode": lambda x: zstd_encode(x, level=22),
283 "decode": lambda x, dtype, shape: zstd_decode(x, dtype, shape),
284 "description": "Zstandard compression at maximum level 22 (highest compression).",
285 "tags": ["zstd"],
286 "source_file": SOURCE_FILE,
287 },
288 {
289 "name": "zstd-22-delta",
290 "version": "1",
291 "encode": lambda x: zstd_delta_encode(x, level=22),
292 "decode": lambda x, dtype, shape: zstd_delta_decode(x, dtype, shape),
293 "description": "Zstandard compression at level 22 with delta encoding for improved compression of sequential data.",
294 "tags": ["zstd", "delta_encoding", "1d"],
295 "source_file": SOURCE_FILE,
296 },
297 {
298 "name": "zstd-22-markov",
299 "version": "1",
300 "encode": lambda x: zstd_markov_encode(x, level=22),
301 "decode": lambda x, dtype, shape: zstd_markov_decode(x, dtype, shape),
302 "description": "Zstandard compression at level 22 with Markov prediction for exploiting temporal correlations in the data.",
303 "tags": ["zstd", "markov_prediction", "1d"],
304 "source_file": SOURCE_FILE,
305 },
306 {
307 "name": "zstd-22-markov-zrle",
308 "version": "1",
309 "encode": lambda x: zstd_markov_zrle_encode(x, level=22),
310 "decode": lambda x, dtype, shape: zstd_markov_zrle_decode(x, dtype, shape),
311 "description": "Zstandard compression at level 22 with Markov prediction and zero run-length encoding for sparse data.",
312 "tags": ["zstd", "markov_prediction", "zero_rle", "1d"],
313 "source_file": SOURCE_FILE,
314 },
315]