concept-collection / benchcompress
add zstd-22-markov-zrle
Jeremy Magland <jmagland@flatironinstitute.org> committed commit b0fb9dd991c1 parent c1f542a Browse files
3 changed files+113−5
benchcompress/pyproject.tomlmodified+1−1View file
@@ -26,7 +26,7 @@ dependencies = [
2626
2727 [tool.scikit-build]
2828 cmake.minimum-version = "3.15"
29-cmake.source-dir = "src/benchcompress/algorithms/simple_ans"
29+cmake.source-dir = "src/benchcompress/algorithms/ans"
3030 cmake.build-type = "Release"
3131 wheel.packages = ["src/benchcompress"]
3232
benchcompress/src/benchcompress/algorithms/zstd/__init__.pymodified+110−2View file
@@ -1,6 +1,7 @@
11 import numpy as np
2-from ..simple_ans.markov_reconstruct import markov_reconstruct as markov_reconstruct_cpp
3-from ..simple_ans.markov_predict import markov_predict as markov_predict_cpp
2+from ..ans.markov_reconstruct import markov_reconstruct as markov_reconstruct_cpp
3+from ..ans.markov_predict import markov_predict as markov_predict_cpp
4+from ..ans.get_run_lengths import get_run_lengths
45
56
67 SOURCE_FILE = "zstd/__init__.py"
@@ -94,6 +95,104 @@ def zstd_markov_decode(x: bytes, dtype: str) -> np.ndarray:
9495 return output
9596
9697
98+def zstd_markov_zrle_encode(x: np.ndarray, level: int) -> bytes:
99+ import zstandard as zstd
100+ import struct
101+
102+ assert x.ndim == 1
103+
104+ # Get run lengths for zero/non-zero sequences
105+ run_lengths = get_run_lengths(x)
106+
107+ # Extract non-zero data
108+ non_zero_arrays = []
109+ array_pos = 0
110+ i = 0
111+ while i < len(run_lengths):
112+ non_zero_len = int(run_lengths[i])
113+ if non_zero_len > 0:
114+ non_zero_arrays.append(x[array_pos : array_pos + non_zero_len])
115+ array_pos += non_zero_len
116+ i += 1
117+ if i < len(run_lengths):
118+ array_pos += int(run_lengths[i]) # Skip zeros
119+ i += 1
120+
121+ non_zero_data = np.concatenate(non_zero_arrays)
122+
123+ # Apply Markov prediction on non-zero data
124+ coeffs, initial, resid = markov_predict_cpp(
125+ non_zero_data, M=6, num_training_samples=10000
126+ )
127+
128+ # Convert data to bytes
129+ coeffs_bytes = coeffs.tobytes()
130+ initial_bytes = initial.tobytes()
131+ run_lengths_bytes = run_lengths.tobytes()
132+
133+ # Create header with lengths
134+ header = struct.pack(
135+ "QQQQ",
136+ len(coeffs_bytes),
137+ len(initial_bytes),
138+ len(run_lengths_bytes),
139+ len(run_lengths),
140+ )
141+
142+ # Compress residuals
143+ resid_bytes = resid.tobytes()
144+ compressor = zstd.ZstdCompressor(level=level)
145+ compressed_resid = compressor.compress(resid_bytes)
146+
147+ # Combine all parts
148+ return header + coeffs_bytes + initial_bytes + run_lengths_bytes + compressed_resid
149+
150+
151+def zstd_markov_zrle_decode(x: bytes, dtype: str) -> np.ndarray:
152+ import zstandard as zstd
153+ import struct
154+
155+ # Extract header
156+ header_size = struct.calcsize("QQQQ")
157+ coeffs_len, initial_len, run_lengths_len, num_run_lengths = struct.unpack(
158+ "QQQQ", x[:header_size]
159+ )
160+
161+ # Extract components
162+ pos = header_size
163+ coeffs = np.frombuffer(x[pos : pos + coeffs_len], dtype=np.float32)
164+ pos += coeffs_len
165+ initial = np.frombuffer(x[pos : pos + initial_len], dtype=dtype)
166+ pos += initial_len
167+ run_lengths = np.frombuffer(x[pos : pos + run_lengths_len], dtype=np.uint32)
168+ pos += run_lengths_len
169+
170+ # Decompress residuals
171+ decompressor = zstd.ZstdDecompressor()
172+ resid_buf = decompressor.decompress(x[pos:])
173+ resid = np.frombuffer(resid_buf, dtype=dtype)
174+
175+ # Reconstruct non-zero data
176+ non_zero_data = markov_reconstruct_cpp(coeffs, initial, resid)
177+
178+ # Reconstruct full array using run lengths
179+ non_zero_pos = 0
180+ i = 0
181+ segments = []
182+ while i < len(run_lengths):
183+ non_zero_len = int(run_lengths[i])
184+ if non_zero_len > 0:
185+ segment = non_zero_data[non_zero_pos : non_zero_pos + non_zero_len]
186+ segments.append(segment)
187+ non_zero_pos += non_zero_len
188+ i += 1
189+ if i < len(run_lengths):
190+ segments.append(np.zeros(int(run_lengths[i]), dtype=non_zero_data.dtype))
191+ i += 1
192+
193+ return np.concatenate(segments)
194+
195+
97196 algorithms = [
98197 {
99198 "name": "zstd-4",
@@ -176,4 +275,13 @@ algorithms = [
176275 "tags": ["zstd", "markov_prediction"],
177276 "source_file": SOURCE_FILE,
178277 },
278+ {
279+ "name": "zstd-22-markov-zrle",
280+ "version": "1",
281+ "encode": lambda x: zstd_markov_zrle_encode(x, level=22),
282+ "decode": lambda x, dtype: zstd_markov_zrle_decode(x, dtype),
283+ "description": "Zstandard compression at level 22 with Markov prediction and zero run-length encoding for sparse data.",
284+ "tags": ["zstd", "markov_prediction", "zero_rle"],
285+ "source_file": SOURCE_FILE,
286+ },
179287 ]
test1.pymodified+2−2View file
@@ -33,11 +33,11 @@ def print_actual_compression_ratios(X):
3333 buf_zstd = compress_ints_lossless(X, method='zstd')
3434 buf_zlib = compress_ints_lossless(X, method='zlib')
3535 buf_lzma = compress_ints_lossless(X, method='lzma')
36- buf_ans = compress_ints_lossless(X, method='simple_ans')
36+ buf_ans = compress_ints_lossless(X, method='ans')
3737 print(f'Zstd compression ratio: {len(X) * X.itemsize / len(buf_zstd):.2f}')
3838 print(f'Zlib compression ratio: {len(X) * X.itemsize / len(buf_zlib):.2f}')
3939 print(f'Lzma compression ratio: {len(X) * X.itemsize / len(buf_lzma):.2f}')
40- print(f'simple_ans compression ratio: {len(X) * X.itemsize / len(buf_ans):.2f}')
40+ print(f'ANS compression ratio: {len(X) * X.itemsize / len(buf_ans):.2f}')
4141
4242 def get_marcovian_prediction_residual(X, M):
4343 sequences = np.array([X[i:i+M] for i in range(len(X) - 2 * M + 1)])