From 4c7841e630aa3ea3eeb0b8a99823cf216461f64b Mon Sep 17 00:00:00 2001 From: Alice Date: Sat, 3 Oct 2026 17:27:55 +0200 Subject: [PATCH] pass Fitrange argument to fit function --- include/aare/Fit.hpp | 8 ++++++- python/src/fit.hpp | 50 +++++++++++++++++++++++++++++++------------- src/fit/Fit.cpp | 30 +++++++++++++++++--------- 3 files changed, 63 insertions(+), 25 deletions(-) diff --git a/include/aare/Fit.hpp b/include/aare/Fit.hpp index cc7edce1..cd9c44e6 100644 --- a/include/aare/Fit.hpp +++ b/include/aare/Fit.hpp @@ -4,6 +4,8 @@ #include "aare/FitModel.hpp" #include "aare/NDArray.hpp" +#include + namespace aare { // --------------------------------------------------------------------------- @@ -62,11 +64,15 @@ NDArray fit_pixel(const FitModel &model, NDView x, * @param chi2_out Output chi-squared / objective values, shape `(rows, * cols)`. * @param n_threads Number of threads used to split rows. + * @param fit_range Optional view of shape `(rows, cols, 2)` with the start + * and end indices of the scan points to fit for each pixel. If not provided, + * the full range is used. */ template void fit_3d(const FitModel &model, NDView x, NDView y, NDView y_err, NDView par_out, NDView err_out, - NDView chi2_out, int n_threads); + NDView chi2_out, int n_threads, + std::optional> fit_range = std::nullopt); } // namespace aare diff --git a/python/src/fit.hpp b/python/src/fit.hpp index efe3024d..c72e875b 100644 --- a/python/src/fit.hpp +++ b/python/src/fit.hpp @@ -17,7 +17,8 @@ py::object fit_dispatch(const aare::FitModel &model, py::array_t x, py::array_t y, - py::object y_err_obj, int n_threads); + py::object y_err_obj, int n_threads, + std::optional> fit_range = std::nullopt); template void bind_fit_model(py::module &m, const char *name) { using FM = aare::FitModel; @@ -117,8 +118,11 @@ template void bind_fit_model(py::module &m, const char *name) { [](const FM &self, py::array_t x, py::array_t y, - py::object y_err_obj, int n_threads) -> py::object { - return fit_dispatch(self, x, y, y_err_obj, n_threads); + py::object y_err_obj, int n_threads, + std::optional> fit_range = + std::nullopt) -> py::object { + return fit_dispatch(self, x, y, y_err_obj, n_threads, + fit_range); }, R"doc( Fit this model to 1D or 3D data using Minuit2. @@ -138,9 +142,11 @@ template void bind_fit_model(py::module &m, const char *name) { Per-point uncertainties. None for unweighted fit. n_threads : int Number of threads for 3D parallel loop. + fit_range : array_like, shape (rows, cols, 2), optional + Start and end indices of the scan points to fit for each pixel. If not provided, the full range is used. )doc", py::arg("x"), py::arg("y"), py::arg("y_err") = py::none(), - py::arg("n_threads") = 4); + py::arg("n_threads") = 4, py::arg("fit_range") = py::none()); if constexpr (std::is_same_v or std::is_same_v) { @@ -226,7 +232,8 @@ py::object fit_dispatch(const aare::FitModel &model, py::array_t x, py::array_t y, - py::object y_err_obj, int n_threads) { + py::object y_err_obj, int n_threads, + std::optional> fit_range) { constexpr std::size_t npar = Model::npar; if (y.ndim() == 3) { @@ -237,6 +244,11 @@ fit_dispatch(const aare::FitModel &model, auto x_view = make_view_1d(x); auto y_view = make_view_3d(y); + const auto fit_range_view = + fit_range.has_value() + ? std::optional(make_view_3d(fit_range.value())) + : std::nullopt; + if (!y_err_obj.is_none()) { auto y_err = py::cast< py::array_t>( @@ -250,7 +262,8 @@ fit_dispatch(const aare::FitModel &model, py::gil_scoped_release release; // release GIL for parallel loop aare::fit_3d(model, x_view, y_view, y_view_err, par_out->view(), err_out->view(), - chi2_out->view(), n_threads); + chi2_out->view(), n_threads, + fit_range_view); } if (model.compute_errors()) { @@ -273,7 +286,8 @@ fit_dispatch(const aare::FitModel &model, aare::fit_3d(model, x_view, y_view, dummy_err, par_out->view(), err_out->view(), - chi2_out->view(), n_threads); + chi2_out->view(), n_threads, + fit_range_view); return py::dict("par"_a = return_image_data(par_out), "par_err"_a = return_image_data(err_out), @@ -284,9 +298,9 @@ fit_dispatch(const aare::FitModel &model, { py::gil_scoped_release release; // release GIL for parallel loop - aare::fit_3d(model, x_view, y_view, dummy_err, - par_out->view(), dummy_err_out, - chi2_out->view(), n_threads); + aare::fit_3d( + model, x_view, y_view, dummy_err, par_out->view(), + dummy_err_out, chi2_out->view(), n_threads, fit_range_view); } return py::dict("par"_a = return_image_data(par_out), @@ -298,6 +312,12 @@ fit_dispatch(const aare::FitModel &model, auto x_view = make_view_1d(x); auto y_view = make_view_1d(y); + if (fit_range.has_value()) { + throw std::runtime_error( + "fit_range is only supported for 3D data. Directly pass the " + "clipped data to fit() for 1D data."); + } + if (!y_err_obj.is_none()) { auto y_err = py::cast< py::array_t>( @@ -328,13 +348,15 @@ py::object dispatch_any_model( py::object model_obj, py::array_t x, py::array_t y, - py::object y_err_obj, int n_threads) { + py::object y_err_obj, int n_threads, + std::optional> fit_range = std::nullopt) { using namespace aare::model; #define AARE_DISPATCH_MODEL(Model) \ if (py::isinstance>(model_obj)) { \ const auto &mdl = model_obj.cast &>(); \ - return fit_dispatch(mdl, x, y, y_err_obj, n_threads); \ + return fit_dispatch(mdl, x, y, y_err_obj, n_threads, \ + fit_range); \ } AARE_DISPATCH_MODEL(Pol1) @@ -423,7 +445,7 @@ void define_fit_bindings(py::module &m) { R"( Fit a model to 1D or 3D data with the minimizer selected by the model's ``minimizer`` property. - + Parameters ---------- model : object @@ -439,7 +461,7 @@ void define_fit_bindings(py::module &m) { Per-point uncertainties. Same shape as y. None → unweighted fit. n_threads : int Number of threads for the 3D parallel loop. - + Returns ------- dict diff --git a/src/fit/Fit.cpp b/src/fit/Fit.cpp index 8b34c177..f8d0b2dc 100644 --- a/src/fit/Fit.cpp +++ b/src/fit/Fit.cpp @@ -267,7 +267,8 @@ template void fit_3d(const FitModel &model, NDView x, NDView y, NDView y_err, NDView par_out, NDView err_out, - NDView chi2_out, int n_threads) { + NDView chi2_out, int n_threads, + std::optional> fit_range) { constexpr std::size_t npar = Model::npar; detail::check_fit_3d_shapes(x, y, y_err, par_out, err_out, chi2_out); @@ -280,14 +281,23 @@ void fit_3d(const FitModel &model, NDView x, std::array err{}; for (ssize_t row = first_row; row < last_row; row++) { for (ssize_t col = 0; col < y.shape(1); col++) { - NDView values(&y(row, col, 0), {y.shape(2)}); + const size_t start_range = + fit_range ? (*fit_range)(row, col, 0) : 0; + const ssize_t num_y_samples = + fit_range ? (*fit_range)(row, col, 1) - start_range + : y.shape(2); + NDView values(&y(row, col, start_range), + {num_y_samples}); NDView errors = - has_errors ? NDView(&y_err(row, col, 0), - {y_err.shape(2)}) - : NDView{}; + has_errors + ? NDView(&y_err(row, col, start_range), + {num_y_samples}) + : NDView{}; + + NDView x_clipped(&x(start_range), {num_y_samples}); double chi2 = 0.0; - fitter.fit(model, x, values, errors, par.data(), + fitter.fit(model, x_clipped, values, errors, par.data(), want_par_errors ? err.data() : nullptr, chi2); for (std::size_t k = 0; k < npar; ++k) @@ -317,10 +327,10 @@ void fit_3d(const FitModel &model, NDView x, NDView); \ template NDArray fit_pixel( \ const FitModel &, NDView, NDView); \ - template void fit_3d(const FitModel &, NDView, \ - NDView, NDView, \ - NDView, NDView, \ - NDView, int); + template void fit_3d( \ + const FitModel &, NDView, NDView, \ + NDView, NDView, NDView, \ + NDView, int, std::optional>); AARE_INSTANTIATE_FIT(model::Gaussian) AARE_INSTANTIATE_FIT(model::GaussianErfcPlateau)