fix markov dtypes
4 changed files+118−94
benchcompress/src/benchcompress/algorithms/ans/__init__.pymodified+6−7View file
@@ -217,9 +217,9 @@ def ans_markov_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
217217 pos += num_symbols
218218 symbol_values = header[pos : pos + num_symbols]
219219 pos += num_symbols
220- coeffs = header[pos : pos + num_coeffs]
220+ coeffs = header[pos : pos + num_coeffs].astype(np.float32)
221221 pos += num_coeffs
222- initial = header[pos : pos + num_initial]
222+ initial = header[pos : pos + num_initial].astype(dtype)
223223 pos += num_initial
224224 bitstream = x[4 + header_size :]
225225 if dtype_code == 0:
@@ -240,10 +240,9 @@ def ans_markov_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
240240 signal_length=int(signal_length),
241241 state=int(state),
242242 symbol_counts=symbol_counts.astype(np.uint32),
243- symbol_values=symbol_values.astype(np.int16), # resid is always int16
243+ symbol_values=symbol_values.astype(dtype),
244244 bitstream=bitstream,
245245 )
246- import time
247246
248247 resid = ans_decode(encoded)
249248 output = markov_reconstruct_cpp(coeffs, initial, resid)
@@ -361,9 +360,9 @@ def ans_markov_sparse_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
361360 pos += num_symbols
362361 symbol_values = header[pos : pos + num_symbols]
363362 pos += num_symbols
364- coeffs = header[pos : pos + num_coeffs]
363+ coeffs = header[pos : pos + num_coeffs].astype(np.float32)
365364 pos += num_coeffs
366- initial = header[pos : pos + num_initial]
365+ initial = header[pos : pos + num_initial].astype(dtype)
367366 pos += num_initial
368367
369368 bitstream_end = 4 + header_size + bitstream_length
@@ -402,7 +401,7 @@ def ans_markov_sparse_decode(x: bytes, dtype: str, shape: tuple) -> np.ndarray:
402401 signal_length=int(signal_length),
403402 state=int(state),
404403 symbol_counts=symbol_counts.astype(np.uint32),
405- symbol_values=symbol_values.astype(np.int16), # resid is always int16
404+ symbol_values=symbol_values.astype(dtype),
406405 bitstream=bitstream,
407406 )
408407
benchcompress/src/benchcompress/algorithms/ans/get_run_lengths.cppmodified+12−83View file
@@ -1,92 +1,21 @@
1-#include <cstdint>
2-#include <pybind11/numpy.h>
3-#include <pybind11/pybind11.h>
4-#include <vector>
1+#include "get_run_lengths.hpp"
52
63 namespace py = pybind11;
74
8-py::array get_run_lengths_cpp(py::array_t<int16_t> x) {
9- auto x_buf = x.request();
10- int16_t *x_ptr = static_cast<int16_t *>(x_buf.ptr);
11- size_t N = x_buf.shape[0];
12-
13- std::vector<uint32_t> runs;
14- size_t i = 0;
15- uint32_t current_nonzero_run_length = 0;
16-
17- while (i < N) {
18- // Check for a sequence of at least 10 zeros
19- bool has_zeros = true;
20- for (size_t j = 0; j < 10 && i + j < N; j++) {
21- if (x_ptr[i + j] != 0) {
22- has_zeros = false;
23- break;
24- }
25- }
26-
27- if (has_zeros) {
28- // Add current non-zero run if any
29- runs.push_back(current_nonzero_run_length);
30- current_nonzero_run_length = 0;
31-
32- // Count consecutive zeros
33- size_t j = i;
34- while (j < N && x_ptr[j] == 0) {
35- j++;
36- }
37- runs.push_back(j - i);
38- i = j;
39- } else {
40- current_nonzero_run_length++;
41- i++;
42- }
43- }
44-
45- // Add final non-zero run if any
46- if (current_nonzero_run_length > 0) {
47- runs.push_back(current_nonzero_run_length);
48- }
49-
50- // Determine appropriate dtype based on max run length
51- uint32_t max_run = 0;
52- for (const auto &run : runs) {
53- if (run > max_run) {
54- max_run = run;
55- }
56- }
57-
58- // Create numpy array with appropriate dtype
59- std::vector<ssize_t> shape = {static_cast<ssize_t>(runs.size())};
5+// Explicit instantiation for int16_t
6+py::array get_run_lengths_int16(py::array_t<int16_t> x) {
7+ return get_run_lengths_impl<int16_t>(x);
8+}
609
61- if (max_run < 256) {
62- py::array_t<uint8_t> result(shape);
63- auto result_buf = result.request();
64- uint8_t *result_ptr = static_cast<uint8_t *>(result_buf.ptr);
65- for (size_t i = 0; i < runs.size(); i++) {
66- result_ptr[i] = static_cast<uint8_t>(runs[i]);
67- }
68- return result;
69- } else if (max_run < 65536) {
70- py::array_t<uint16_t> result(shape);
71- auto result_buf = result.request();
72- uint16_t *result_ptr = static_cast<uint16_t *>(result_buf.ptr);
73- for (size_t i = 0; i < runs.size(); i++) {
74- result_ptr[i] = static_cast<uint16_t>(runs[i]);
75- }
76- return result;
77- } else {
78- py::array_t<uint32_t> result(shape);
79- auto result_buf = result.request();
80- uint32_t *result_ptr = static_cast<uint32_t *>(result_buf.ptr);
81- for (size_t i = 0; i < runs.size(); i++) {
82- result_ptr[i] = runs[i];
83- }
84- return result;
85- }
10+// Explicit instantiation for int32_t
11+py::array get_run_lengths_int32(py::array_t<int32_t> x) {
12+ return get_run_lengths_impl<int32_t>(x);
8613 }
8714
8815 PYBIND11_MODULE(get_run_lengths_cpp_ext, m) {
8916 m.doc() = "C++ implementation of get_run_lengths using pybind11";
90- m.def("get_run_lengths_cpp", &get_run_lengths_cpp,
91- "Calculate run lengths of zeros and non-zeros in a signal");
17+ m.def("get_run_lengths_int16", &get_run_lengths_int16,
18+ "Calculate run lengths of zeros and non-zeros in a signal (int16)");
19+ m.def("get_run_lengths_int32", &get_run_lengths_int32,
20+ "Calculate run lengths of zeros and non-zeros in a signal (int32)");
9221 }
benchcompress/src/benchcompress/algorithms/ans/get_run_lengths.hppadded+88−0View file
@@ -0,0 +1,88 @@
1+#pragma once
2+
3+#include <cstdint>
4+#include <pybind11/numpy.h>
5+#include <pybind11/pybind11.h>
6+#include <vector>
7+
8+namespace py = pybind11;
9+
10+template <typename T> py::array get_run_lengths_impl(py::array_t<T> x) {
11+ auto x_buf = x.request();
12+ T *x_ptr = static_cast<T *>(x_buf.ptr);
13+ size_t N = x_buf.shape[0];
14+
15+ std::vector<uint32_t> runs;
16+ size_t i = 0;
17+ uint32_t current_nonzero_run_length = 0;
18+
19+ while (i < N) {
20+ // Check for a sequence of at least 10 zeros
21+ bool has_zeros = true;
22+ for (size_t j = 0; j < 10 && i + j < N; j++) {
23+ if (x_ptr[i + j] != 0) {
24+ has_zeros = false;
25+ break;
26+ }
27+ }
28+
29+ if (has_zeros) {
30+ // Add current non-zero run if any
31+ runs.push_back(current_nonzero_run_length);
32+ current_nonzero_run_length = 0;
33+
34+ // Count consecutive zeros
35+ size_t j = i;
36+ while (j < N && x_ptr[j] == 0) {
37+ j++;
38+ }
39+ runs.push_back(j - i);
40+ i = j;
41+ } else {
42+ current_nonzero_run_length++;
43+ i++;
44+ }
45+ }
46+
47+ // Add final non-zero run if any
48+ if (current_nonzero_run_length > 0) {
49+ runs.push_back(current_nonzero_run_length);
50+ }
51+
52+ // Determine appropriate dtype based on max run length
53+ uint32_t max_run = 0;
54+ for (const auto &run : runs) {
55+ if (run > max_run) {
56+ max_run = run;
57+ }
58+ }
59+
60+ // Create numpy array with appropriate dtype
61+ std::vector<ssize_t> shape = {static_cast<ssize_t>(runs.size())};
62+
63+ if (max_run < 256) {
64+ py::array_t<uint8_t> result(shape);
65+ auto result_buf = result.request();
66+ uint8_t *result_ptr = static_cast<uint8_t *>(result_buf.ptr);
67+ for (size_t i = 0; i < runs.size(); i++) {
68+ result_ptr[i] = static_cast<uint8_t>(runs[i]);
69+ }
70+ return result;
71+ } else if (max_run < 65536) {
72+ py::array_t<uint16_t> result(shape);
73+ auto result_buf = result.request();
74+ uint16_t *result_ptr = static_cast<uint16_t *>(result_buf.ptr);
75+ for (size_t i = 0; i < runs.size(); i++) {
76+ result_ptr[i] = static_cast<uint16_t>(runs[i]);
77+ }
78+ return result;
79+ } else {
80+ py::array_t<uint32_t> result(shape);
81+ auto result_buf = result.request();
82+ uint32_t *result_ptr = static_cast<uint32_t *>(result_buf.ptr);
83+ for (size_t i = 0; i < runs.size(); i++) {
84+ result_ptr[i] = runs[i];
85+ }
86+ return result;
87+ }
88+}
benchcompress/src/benchcompress/algorithms/ans/get_run_lengths.pymodified+12−4View file
@@ -1,16 +1,24 @@
11 import numpy as np
2-from .get_run_lengths_cpp_ext import get_run_lengths_cpp
2+from .get_run_lengths_cpp_ext import get_run_lengths_int16, get_run_lengths_int32
33
44
55 def get_run_lengths(x: np.ndarray) -> np.ndarray:
66 """Calculate run lengths of zeros and non-zeros in a signal using C++ implementation.
77
88 Args:
9- x: Input signal (will be converted to int16)
9+ x: Input signal (must be int16 or int32)
1010
1111 Returns:
1212 np.ndarray: Array of run lengths alternating between non-zero and zero runs.
1313 The dtype will be uint8, uint16, or uint32 depending on the maximum run length.
14+
15+ Raises:
16+ ValueError: If input array is not int16 or int32
1417 """
15- # Call C++ implementation with proper type conversion
16- return get_run_lengths_cpp(x.astype(np.int16))
18+ # Check input dtype and call appropriate implementation
19+ if x.dtype == np.int16:
20+ return get_run_lengths_int16(x)
21+ elif x.dtype == np.int32:
22+ return get_run_lengths_int32(x)
23+ else:
24+ raise ValueError(f"Input array must be int16 or int32, got {x.dtype}")