/ concept-collection / benchcompress
Sign in
concept-collection / benchcompress
handle int32 for markov
Jeremy Magland <jmagland@flatironinstitute.org> committed commit 6cc72656444a parent dc74cd0 Browse files
6 changed files+228−154
benchcompress/src/benchcompress/algorithms/ans/markov_predict.cppmodified+17−83View file
@@ -1,95 +1,29 @@
1-#include <Eigen/Dense>
2-#include <cmath>
3-#include <iostream>
4-#include <pybind11/eigen.h>
5-#include <pybind11/numpy.h>
6-#include <pybind11/pybind11.h>
1+#include "markov_predict.hpp"
72
83 namespace py = pybind11;
94
5+// Explicit instantiation for int16_t
106 std::tuple<py::array_t<float>, py::array_t<int16_t>, py::array_t<int16_t>>
11-markov_predict_cpp(py::array_t<int16_t> x, size_t M,
12- size_t num_training_samples) {
13- // Get array buffer
14- auto x_buf = x.request();
15- int16_t *x_ptr = static_cast<int16_t *>(x_buf.ptr);
16- size_t N = x_buf.shape[0];
17-
18- // Keep initial values for reconstruction
19- size_t initial_size = M - 1;
20- std::vector<ssize_t> initial_shape = {static_cast<ssize_t>(initial_size)};
21- py::array_t<int16_t> initial(initial_shape);
22- py::buffer_info initial_buf = initial.request(true);
23- int16_t *initial_ptr = static_cast<int16_t *>(initial_buf.ptr);
24-
25- // Copy initial values
26- for (size_t i = 0; i < initial_size; i++) {
27- initial_ptr[i] = x_ptr[i];
28- }
29-
30- // Create sequences matrix for linear regression
31- size_t resid_size = N - M + 1;
32- // Use only num_training_samples sequences for model fitting
33- size_t num_samples_for_fit = std::min(resid_size, num_training_samples);
34- // Take evenly spaced samples for training
35- size_t stride =
36- resid_size > num_training_samples ? resid_size / num_training_samples : 1;
37-
38- Eigen::MatrixXf predictors(num_samples_for_fit, M - 1);
39- Eigen::VectorXf target(num_samples_for_fit);
40-
41- // Fill predictors matrix and target vector with strided training samples
42- for (size_t i = 0; i < num_samples_for_fit; i++) {
43- size_t idx = i * stride;
44- for (size_t j = 0; j < M - 1; j++) {
45- predictors(i, j) = static_cast<float>(x_ptr[idx + j]);
46- }
47- target(i) = static_cast<float>(x_ptr[idx + M - 1]);
48- }
49-
50- // Add constant term column (ones) to predictors
51- Eigen::MatrixXf X(predictors.rows(), predictors.cols() + 1);
52- X << predictors, Eigen::VectorXf::Ones(predictors.rows());
53-
54- // Solve least squares problem: X * coeffs = target
55- Eigen::VectorXf coeffs = X.colPivHouseholderQr().solve(target);
56-
57- // Create coefficients array
58- std::vector<ssize_t> coeffs_shape = {static_cast<ssize_t>(M)};
59- py::array_t<float> coeffs_array(coeffs_shape);
60- py::buffer_info coeffs_buf = coeffs_array.request(true);
61- float *coeffs_ptr = static_cast<float *>(coeffs_buf.ptr);
62-
63- // Copy coefficients
64- for (size_t i = 0; i < M - 1; i++) {
65- coeffs_ptr[i] = coeffs(i);
66- }
67- coeffs_ptr[M - 1] = coeffs(M - 1); // bias term
68-
69- // Calculate residuals
70- std::vector<ssize_t> resid_shape = {static_cast<ssize_t>(resid_size)};
71- py::array_t<int16_t> residuals(resid_shape);
72- py::buffer_info resid_buf = residuals.request(true);
73- int16_t *resid_ptr = static_cast<int16_t *>(resid_buf.ptr);
74-
75- for (size_t i = 0; i < resid_size; i++) {
76- float prediction = 0.0f;
77- for (size_t j = 0; j < M - 1; j++) {
78- float term = coeffs_ptr[j] * static_cast<float>(x_ptr[i + j]);
79- prediction += term;
80- }
81- prediction += coeffs_ptr[M - 1]; // bias term
82- float rounded_prediction = std::round(prediction);
83- resid_ptr[i] = x_ptr[i + M - 1] - static_cast<int16_t>(rounded_prediction);
84- }
7+markov_predict_int16(py::array_t<int16_t> x, size_t M,
8+ size_t num_training_samples) {
9+ return markov_predict_impl<int16_t>(x, M, num_training_samples);
10+}
8511
86- return std::make_tuple(coeffs_array, initial, residuals);
12+// Explicit instantiation for int32_t
13+std::tuple<py::array_t<float>, py::array_t<int32_t>, py::array_t<int32_t>>
14+markov_predict_int32(py::array_t<int32_t> x, size_t M,
15+ size_t num_training_samples) {
16+ return markov_predict_impl<int32_t>(x, M, num_training_samples);
8717 }
8818
8919 PYBIND11_MODULE(markov_predict_cpp_ext, m) {
9020 m.doc() = "C++ implementation of markov_predict using pybind11";
91- m.def("markov_predict_cpp", &markov_predict_cpp,
21+ m.def("markov_predict_int16", &markov_predict_int16,
22+ "Predict signal using Markov model and return coefficients, initial "
23+ "values and residuals (int16)",
24+ py::arg("x"), py::arg("M"), py::arg("num_training_samples") = 10000);
25+ m.def("markov_predict_int32", &markov_predict_int32,
9226 "Predict signal using Markov model and return coefficients, initial "
93- "values and residuals",
27+ "values and residuals (int32)",
9428 py::arg("x"), py::arg("M"), py::arg("num_training_samples") = 10000);
9529 }
benchcompress/src/benchcompress/algorithms/ans/markov_predict.hppadded+89−0View file
@@ -0,0 +1,89 @@
1+#pragma once
2+
3+#include <Eigen/Dense>
4+#include <cmath>
5+#include <iostream>
6+#include <pybind11/eigen.h>
7+#include <pybind11/numpy.h>
8+#include <pybind11/pybind11.h>
9+
10+namespace py = pybind11;
11+
12+template <typename T>
13+std::tuple<py::array_t<float>, py::array_t<T>, py::array_t<T>>
14+markov_predict_impl(py::array_t<T> x, size_t M, size_t num_training_samples) {
15+ // Get array buffer
16+ auto x_buf = x.request();
17+ T *x_ptr = static_cast<T *>(x_buf.ptr);
18+ size_t N = x_buf.shape[0];
19+
20+ // Keep initial values for reconstruction
21+ size_t initial_size = M - 1;
22+ std::vector<ssize_t> initial_shape = {static_cast<ssize_t>(initial_size)};
23+ py::array_t<T> initial(initial_shape);
24+ py::buffer_info initial_buf = initial.request(true);
25+ T *initial_ptr = static_cast<T *>(initial_buf.ptr);
26+
27+ // Copy initial values
28+ for (size_t i = 0; i < initial_size; i++) {
29+ initial_ptr[i] = x_ptr[i];
30+ }
31+
32+ // Create sequences matrix for linear regression
33+ size_t resid_size = N - M + 1;
34+ // Use only num_training_samples sequences for model fitting
35+ size_t num_samples_for_fit = std::min(resid_size, num_training_samples);
36+ // Take evenly spaced samples for training
37+ size_t stride =
38+ resid_size > num_training_samples ? resid_size / num_training_samples : 1;
39+
40+ Eigen::MatrixXf predictors(num_samples_for_fit, M - 1);
41+ Eigen::VectorXf target(num_samples_for_fit);
42+
43+ // Fill predictors matrix and target vector with strided training samples
44+ for (size_t i = 0; i < num_samples_for_fit; i++) {
45+ size_t idx = i * stride;
46+ for (size_t j = 0; j < M - 1; j++) {
47+ predictors(i, j) = static_cast<float>(x_ptr[idx + j]);
48+ }
49+ target(i) = static_cast<float>(x_ptr[idx + M - 1]);
50+ }
51+
52+ // Add constant term column (ones) to predictors
53+ Eigen::MatrixXf X(predictors.rows(), predictors.cols() + 1);
54+ X << predictors, Eigen::VectorXf::Ones(predictors.rows());
55+
56+ // Solve least squares problem: X * coeffs = target
57+ Eigen::VectorXf coeffs = X.colPivHouseholderQr().solve(target);
58+
59+ // Create coefficients array
60+ std::vector<ssize_t> coeffs_shape = {static_cast<ssize_t>(M)};
61+ py::array_t<float> coeffs_array(coeffs_shape);
62+ py::buffer_info coeffs_buf = coeffs_array.request(true);
63+ float *coeffs_ptr = static_cast<float *>(coeffs_buf.ptr);
64+
65+ // Copy coefficients
66+ for (size_t i = 0; i < M - 1; i++) {
67+ coeffs_ptr[i] = coeffs(i);
68+ }
69+ coeffs_ptr[M - 1] = coeffs(M - 1); // bias term
70+
71+ // Calculate residuals
72+ std::vector<ssize_t> resid_shape = {static_cast<ssize_t>(resid_size)};
73+ py::array_t<T> residuals(resid_shape);
74+ py::buffer_info resid_buf = residuals.request(true);
75+ T *resid_ptr = static_cast<T *>(resid_buf.ptr);
76+
77+ for (size_t i = 0; i < resid_size; i++) {
78+ float prediction = 0.0f;
79+ for (size_t j = 0; j < M - 1; j++) {
80+ float term = coeffs_ptr[j] * static_cast<float>(x_ptr[i + j]);
81+ prediction += term;
82+ }
83+ prediction += coeffs_ptr[M - 1]; // bias term
84+ float rounded_prediction = std::round(prediction);
85+ resid_ptr[i] = x_ptr[i + M - 1] - static_cast<T>(rounded_prediction);
86+ }
87+
88+ return std::make_tuple(coeffs_array, initial, residuals);
89+}
benchcompress/src/benchcompress/algorithms/ans/markov_predict.pymodified+13−5View file
@@ -1,18 +1,26 @@
11 import numpy as np
2-from .markov_predict_cpp_ext import markov_predict_cpp
2+from .markov_predict_cpp_ext import markov_predict_int16, markov_predict_int32
33
44
55 def markov_predict(x: np.ndarray, M: int, num_training_samples: int = 10000) -> tuple:
66 """Predict signal using Markov model and return coefficients, initial values and residuals 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 M: Number of previous samples to use for prediction
1111 num_training_samples: Maximum number of samples to use for fitting the model coefficients (default: 10000).
1212 Using fewer samples speeds up model fitting on large inputs while maintaining accuracy.
1313
1414 Returns:
15- tuple: (coefficients (float32), initial_values (int16), residuals (int16))
15+ tuple: (coefficients (float32), initial_values (same dtype as input), residuals (same dtype as input))
16+
17+ Raises:
18+ ValueError: If input array is not int16 or int32
1619 """
17- # Call C++ implementation with proper type conversion
18- return markov_predict_cpp(x.astype(np.int16), M, num_training_samples)
20+ # Check input dtype and call appropriate implementation
21+ if x.dtype == np.int16:
22+ return markov_predict_int16(x, M, num_training_samples)
23+ elif x.dtype == np.int32:
24+ return markov_predict_int32(x, M, num_training_samples)
25+ else:
26+ raise ValueError(f"Input array must be int16 or int32, got {x.dtype}")
benchcompress/src/benchcompress/algorithms/ans/markov_reconstruct.cppmodified+17−58View file
@@ -1,68 +1,27 @@
1-#include <cmath>
2-#include <iostream>
3-#include <pybind11/numpy.h>
4-#include <pybind11/pybind11.h>
1+#include "markov_reconstruct.hpp"
52
63 namespace py = pybind11;
74
8-py::array_t<int16_t> markov_reconstruct_cpp(py::array_t<float> coeffs,
9- py::array_t<int16_t> initial,
10- py::array_t<int16_t> resid) {
11- // Get array buffers
12- auto coeffs_buf = coeffs.request();
13- auto initial_buf = initial.request();
14- auto resid_buf = resid.request();
15-
16- // Get raw pointers to data
17- float *coeffs_ptr = static_cast<float *>(coeffs_buf.ptr);
18- int16_t *initial_ptr = static_cast<int16_t *>(initial_buf.ptr);
19- int16_t *resid_ptr = static_cast<int16_t *>(resid_buf.ptr);
20-
21- // Calculate dimensions
22- size_t M = initial_buf.shape[0] + 1; // Number of samples used in prediction
23- size_t output_size = resid_buf.shape[0] + initial_buf.shape[0];
24-
25- // Create output array with explicit shape and memory ownership
26- std::vector<ssize_t> shape = {static_cast<ssize_t>(output_size)};
27- py::array_t<int16_t> output(shape);
28- py::buffer_info output_buf = output.request(true); // Request writable buffer
29- int16_t *output_ptr = static_cast<int16_t *>(output_buf.ptr);
30-
31- // Copy initial values with bounds check
32- for (size_t i = 0; i < initial_buf.shape[0] && i < output_size; i++) {
33- output_ptr[i] = initial_ptr[i];
34- }
35-
36- size_t resid_size = resid_buf.shape[0];
37-
38- // Reconstruct signal iteratively
39- for (size_t i = 0; i < resid_size; i++) {
40- float prediction = 0.0f;
41-
42- // Calculate prediction using coefficients (excluding bias term)
43- for (size_t j = 0; j < M - 1; j++) {
44- float term = coeffs_ptr[j] * static_cast<float>(output_ptr[i + j]);
45- prediction += term;
46- }
47-
48- // Add bias term separately
49- prediction += coeffs_ptr[coeffs_buf.shape[0] - 1];
50-
51- // Round prediction to nearest integer
52- float rounded_prediction = std::round(prediction);
53-
54- // Add residual and store result
55- int16_t final_value = static_cast<int16_t>(
56- rounded_prediction + static_cast<float>(resid_ptr[i]));
57- output_ptr[i + M - 1] = final_value;
58- }
5+// Explicit instantiation for int16_t
6+py::array_t<int16_t> markov_reconstruct_int16(py::array_t<float> coeffs,
7+ py::array_t<int16_t> initial,
8+ py::array_t<int16_t> resid) {
9+ return markov_reconstruct_impl<int16_t>(coeffs, initial, resid);
10+}
5911
60- return output;
12+// Explicit instantiation for int32_t
13+py::array_t<int32_t> markov_reconstruct_int32(py::array_t<float> coeffs,
14+ py::array_t<int32_t> initial,
15+ py::array_t<int32_t> resid) {
16+ return markov_reconstruct_impl<int32_t>(coeffs, initial, resid);
6117 }
6218
6319 PYBIND11_MODULE(markov_reconstruct_cpp_ext, m) {
6420 m.doc() = "C++ implementation of markov_reconstruct using pybind11";
65- m.def("markov_reconstruct_cpp", &markov_reconstruct_cpp,
66- "Reconstruct signal from Markov model parameters and residuals",
21+ m.def("markov_reconstruct_int16", &markov_reconstruct_int16,
22+ "Reconstruct signal from Markov model parameters and residuals (int16)",
23+ py::arg("coeffs"), py::arg("initial"), py::arg("resid"));
24+ m.def("markov_reconstruct_int32", &markov_reconstruct_int32,
25+ "Reconstruct signal from Markov model parameters and residuals (int32)",
6726 py::arg("coeffs"), py::arg("initial"), py::arg("resid"));
6827 }
benchcompress/src/benchcompress/algorithms/ans/markov_reconstruct.hppadded+64−0View file
@@ -0,0 +1,64 @@
1+#pragma once
2+
3+#include <cmath>
4+#include <iostream>
5+#include <pybind11/numpy.h>
6+#include <pybind11/pybind11.h>
7+
8+namespace py = pybind11;
9+
10+template <typename T>
11+py::array_t<T> markov_reconstruct_impl(py::array_t<float> coeffs,
12+ py::array_t<T> initial,
13+ py::array_t<T> resid) {
14+ // Get array buffers
15+ auto coeffs_buf = coeffs.request();
16+ auto initial_buf = initial.request();
17+ auto resid_buf = resid.request();
18+
19+ // Get raw pointers to data
20+ float *coeffs_ptr = static_cast<float *>(coeffs_buf.ptr);
21+ T *initial_ptr = static_cast<T *>(initial_buf.ptr);
22+ T *resid_ptr = static_cast<T *>(resid_buf.ptr);
23+
24+ // Calculate dimensions
25+ size_t M = initial_buf.shape[0] + 1; // Number of samples used in prediction
26+ size_t output_size = resid_buf.shape[0] + initial_buf.shape[0];
27+
28+ // Create output array with explicit shape and memory ownership
29+ std::vector<ssize_t> shape = {static_cast<ssize_t>(output_size)};
30+ py::array_t<T> output(shape);
31+ py::buffer_info output_buf = output.request(true); // Request writable buffer
32+ T *output_ptr = static_cast<T *>(output_buf.ptr);
33+
34+ // Copy initial values with bounds check
35+ for (size_t i = 0; i < initial_buf.shape[0] && i < output_size; i++) {
36+ output_ptr[i] = initial_ptr[i];
37+ }
38+
39+ size_t resid_size = resid_buf.shape[0];
40+
41+ // Reconstruct signal iteratively
42+ for (size_t i = 0; i < resid_size; i++) {
43+ float prediction = 0.0f;
44+
45+ // Calculate prediction using coefficients (excluding bias term)
46+ for (size_t j = 0; j < M - 1; j++) {
47+ float term = coeffs_ptr[j] * static_cast<float>(output_ptr[i + j]);
48+ prediction += term;
49+ }
50+
51+ // Add bias term separately
52+ prediction += coeffs_ptr[coeffs_buf.shape[0] - 1];
53+
54+ // Round prediction to nearest integer
55+ float rounded_prediction = std::round(prediction);
56+
57+ // Add residual and store result
58+ T final_value =
59+ static_cast<T>(rounded_prediction + static_cast<float>(resid_ptr[i]));
60+ output_ptr[i + M - 1] = final_value;
61+ }
62+
63+ return output;
64+}
benchcompress/src/benchcompress/algorithms/ans/markov_reconstruct.pymodified+28−8View file
@@ -1,5 +1,8 @@
11 import numpy as np
2-from .markov_reconstruct_cpp_ext import markov_reconstruct_cpp
2+from .markov_reconstruct_cpp_ext import (
3+ markov_reconstruct_int16,
4+ markov_reconstruct_int32,
5+)
36
47
58 def markov_reconstruct(coeffs, initial, resid):
@@ -7,13 +10,30 @@ def markov_reconstruct(coeffs, initial, resid):
710
811 Args:
912 coeffs: Model coefficients from linear regression (float32)
10- initial: Initial values needed for prediction (int16)
11- resid: Prediction residuals (int16)
13+ initial: Initial values needed for prediction (int16 or int32)
14+ resid: Prediction residuals (must match initial dtype)
1215
1316 Returns:
14- np.ndarray: Reconstructed signal (int16)
17+ np.ndarray: Reconstructed signal (same dtype as input)
18+
19+ Raises:
20+ ValueError: If initial/resid arrays are not int16 or int32, or if their dtypes don't match
1521 """
16- # Call C++ implementation
17- return markov_reconstruct_cpp(
18- coeffs.astype(np.float32), initial.astype(np.int16), resid.astype(np.int16)
19- )
22+ # Convert coeffs to float32 if needed
23+ coeffs = coeffs.astype(np.float32)
24+
25+ # Check input dtypes
26+ if initial.dtype != resid.dtype:
27+ raise ValueError(
28+ f"Initial and residual arrays must have same dtype, got {initial.dtype} and {resid.dtype}"
29+ )
30+
31+ # Call appropriate implementation based on dtype
32+ if initial.dtype == np.int16:
33+ return markov_reconstruct_int16(coeffs, initial, resid)
34+ elif initial.dtype == np.int32:
35+ return markov_reconstruct_int32(coeffs, initial, resid)
36+ else:
37+ raise ValueError(
38+ f"Initial/residual arrays must be int16 or int32, got {initial.dtype}"
39+ )
moveopenescclose