diff --git a/RELEASE.md b/RELEASE.md index f1dd8a78..a5d3284b 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -50,6 +50,21 @@ erf from the Gaussian's exponential, a rounding-level change. One iteration costs one model evaluation of the ``max_calls`` budget, so pixels that used to exhaust the default budget of 100 may now converge. +- Added ``Minimizer.VarPro``, a built-in variable projection + solver. Every bundled model now declares the parameters it is linear in + (``linear_par`` and ``basis_and_grad`` in C++, checked against the model + derivatives by the tests). Each trial point solves those exactly and the + Levenberg-Marquardt iteration runs over the nonlinear parameters only, so + the fit needs fewer evaluations and does not depend on start values of the + linear parameters. Limits on linear parameters, pixels that do not converge + and models without the separable structure fall back to + ``LevenbergMarquardt``. Parameter errors are the same Gauss-Newton + estimates. The Gaussian, plateau, charge-sharing and S-curve models also + provide their basis functions for all scan points at once + (``basis_columns``), a loop that compiles to vector instructions with the + new ``model::fast_exp``; ``eval`` and ``eval_and_grad``, and with them the + other minimizers, are unchanged. ``Minimizer.Migrad`` remains the + default. - ``NDArray``/``NDView`` expressions now support scalar operands (``2 * a + b / 4``) and can be assigned to an existing ``NDArray``, reusing its buffer. diff --git a/benchmarks/fit_benchmark.cpp b/benchmarks/fit_benchmark.cpp index 02bef6ec..8e522f80 100644 --- a/benchmarks/fit_benchmark.cpp +++ b/benchmarks/fit_benchmark.cpp @@ -138,6 +138,15 @@ static void BM_FitGausLMErrors(benchmark::State &state) { run_gaussian_fit(state, aare::Minimizer::LevenbergMarquardt, true); } +// Built-in variable projection: A solved exactly, iteration over mu and sigma +static void BM_FitGausVarPro(benchmark::State &state) { + run_gaussian_fit(state, aare::Minimizer::VarPro, false); +} + +static void BM_FitGausVarProErrors(benchmark::State &state) { + run_gaussian_fit(state, aare::Minimizer::VarPro, true); +} + // Gaussian with sigma fixed: exercises the fixed-parameter path of the // minimizers on the Moderate_noise case. static void run_gaussian_fixed_sigma_fit(benchmark::State &state, @@ -167,6 +176,9 @@ static void BM_FitGausFixedSigmaFumili(benchmark::State &state) { static void BM_FitGausFixedSigmaLM(benchmark::State &state) { run_gaussian_fixed_sigma_fit(state, aare::Minimizer::LevenbergMarquardt); } +static void BM_FitGausFixedSigmaVarPro(benchmark::State &state) { + run_gaussian_fixed_sigma_fit(state, aare::Minimizer::VarPro); +} // ---------------------------------------------------------------- // Rising S-curve (6 parameters), typical for threshold scans @@ -221,6 +233,10 @@ static void BM_FitScurveLM(benchmark::State &state) { run_scurve_fit(state, aare::Minimizer::LevenbergMarquardt); } +static void BM_FitScurveVarPro(benchmark::State &state) { + run_scurve_fit(state, aare::Minimizer::VarPro); +} + // Weighted fits take the per-point uncertainties into the residuals. static void BM_FitScurveWeightedMigrad(benchmark::State &state) { run_scurve_fit(state, aare::Minimizer::Migrad, true); @@ -231,6 +247,9 @@ static void BM_FitScurveWeightedFumili(benchmark::State &state) { static void BM_FitScurveWeightedLM(benchmark::State &state) { run_scurve_fit(state, aare::Minimizer::LevenbergMarquardt, true); } +static void BM_FitScurveWeightedVarPro(benchmark::State &state) { + run_scurve_fit(state, aare::Minimizer::VarPro, true); +} // ---------------------------------------------------------------- // Pol2 (3 parameters, cheap model): exposes the solver overhead per @@ -271,6 +290,9 @@ static void BM_FitPol2Fumili(benchmark::State &state) { static void BM_FitPol2LM(benchmark::State &state) { run_pol2_fit(state, aare::Minimizer::LevenbergMarquardt); } +static void BM_FitPol2VarPro(benchmark::State &state) { + run_pol2_fit(state, aare::Minimizer::VarPro); +} // ---------------------------------------------------------------- // GaussianChargeSharingKb (8 parameters, two exponentials and two erfs per @@ -318,6 +340,9 @@ static void BM_FitChargeSharingKbFumili(benchmark::State &state) { static void BM_FitChargeSharingKbLM(benchmark::State &state) { run_charge_sharing_kb_fit(state, aare::Minimizer::LevenbergMarquardt); } +static void BM_FitChargeSharingKbVarPro(benchmark::State &state) { + run_charge_sharing_kb_fit(state, aare::Minimizer::VarPro); +} // ---------------------------------------------------------------- // Data cube through fit_3d on one thread: exposes the per-pixel overhead @@ -393,10 +418,17 @@ static void BM_FitCubeFumiliErrors(benchmark::State &state) { static void BM_FitCubeLMErrors(benchmark::State &state) { run_cube_fit(state, aare::Minimizer::LevenbergMarquardt, true); } +static void BM_FitCubeVarPro(benchmark::State &state) { + run_cube_fit(state, aare::Minimizer::VarPro, false); +} +static void BM_FitCubeVarProErrors(benchmark::State &state) { + run_cube_fit(state, aare::Minimizer::VarPro, true); +} BENCHMARK(BM_FitGausMigrad)->DenseRange(0, 5)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitGausFumili)->DenseRange(0, 5)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitGausLM)->DenseRange(0, 5)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitGausVarPro)->DenseRange(0, 5)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitGausMigradHesse) ->DenseRange(0, 5) @@ -405,26 +437,34 @@ BENCHMARK(BM_FitGausFumiliErrors) ->DenseRange(0, 5) ->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitGausLMErrors)->DenseRange(0, 5)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitGausVarProErrors) + ->DenseRange(0, 5) + ->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitGausFixedSigmaMigrad)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitGausFixedSigmaFumili)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitGausFixedSigmaLM)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitGausFixedSigmaVarPro)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitScurveMigrad)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitScurveFumili)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitScurveLM)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitScurveVarPro)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitScurveWeightedMigrad)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitScurveWeightedFumili)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitScurveWeightedLM)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitScurveWeightedVarPro)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitPol2Migrad)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitPol2Fumili)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitPol2LM)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitPol2VarPro)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitChargeSharingKbMigrad)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitChargeSharingKbFumili)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitChargeSharingKbLM)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitChargeSharingKbVarPro)->Unit(benchmark::kMicrosecond); // fit_3d runs in worker threads, so measure wall time. BENCHMARK(BM_FitCubeMigrad)->Unit(benchmark::kMillisecond)->UseRealTime(); @@ -433,5 +473,7 @@ BENCHMARK(BM_FitCubeLM)->Unit(benchmark::kMillisecond)->UseRealTime(); BENCHMARK(BM_FitCubeMigradErrors)->Unit(benchmark::kMillisecond)->UseRealTime(); BENCHMARK(BM_FitCubeFumiliErrors)->Unit(benchmark::kMillisecond)->UseRealTime(); BENCHMARK(BM_FitCubeLMErrors)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK(BM_FitCubeVarPro)->Unit(benchmark::kMillisecond)->UseRealTime(); +BENCHMARK(BM_FitCubeVarProErrors)->Unit(benchmark::kMillisecond)->UseRealTime(); BENCHMARK_MAIN(); diff --git a/docs/src/pyFit.rst b/docs/src/pyFit.rst index 4bcafedb..2dbb5569 100644 --- a/docs/src/pyFit.rst +++ b/docs/src/pyFit.rst @@ -41,7 +41,7 @@ The result dictionary contains ``par`` and ``chi2``. It also contains Choosing the minimizer ~~~~~~~~~~~~~~~~~~~~~~ -Fits use Minuit2's Migrad by default. Two alternatives share the same +Fits use Minuit2's Migrad by default. Three alternatives share the same interface and are selected with the ``minimizer`` argument or property:: model = Gaussian(minimizer=Minimizer.LevenbergMarquardt) @@ -60,8 +60,18 @@ interface and are selected with the ``minimizer`` argument or property:: the analytic Jacobian of the model. It reflects steps at parameter limits and reuses its buffers between pixels, so data cubes are fitted without per-pixel allocations. Parameter errors are ``sqrt(diag((J^T J)^-1))``. +- ``Minimizer.VarPro``: a built-in variable projection solver. + Every bundled model is linear in some of its parameters (the amplitude of a + Gaussian, the baseline and amplitude of an S-curve, every parameter of a + polynomial). Each trial point solves those exactly and the + Levenberg-Marquardt iteration runs over the remaining nonlinear parameters + only, so it needs fewer evaluations than LevenbergMarquardt and cannot be + misled by poor start values of the linear parameters, which it ignores. + A pixel whose linear solution violates a limit on a linear parameter, or + that does not converge, is refitted with LevenbergMarquardt. Parameter + errors are the same Gauss-Newton estimates. -All three converge to the same minimum on well-behaved data; the ``[fit]`` +All four converge to the same minimum on well-behaved data; the ``[fit]`` test suite checks that. The bundled ``fit_benchmark`` compares their speed on Gaussian peaks, S-curves and a data cube. For very noisy data the Gauss-Newton steps of Fumili and LevenbergMarquardt converge more slowly, so compare the @@ -78,15 +88,16 @@ free start values are clamped into their limits. The constructor arguments control the minimizer: - ``max_calls`` caps the work per pixel: Minuit2 function calls for Migrad and - Fumili, model evaluations for LevenbergMarquardt, which spends one - evaluation per iteration (two when a step crosses a limit). ``0`` selects - Minuit's default budget. A pixel that does not converge within the budget - returns zeros. -- ``tolerance`` is the EDM tolerance in Minuit's convention. Migrad and - LevenbergMarquardt stop when the estimated distance to the minimum is below - ``0.002 * tolerance``, Fumili below ``1e-4 * tolerance``. + Fumili, model evaluations for LevenbergMarquardt and VarPro, + which spend one evaluation per iteration (two when a step crosses a limit). + ``0`` selects Minuit's default budget. A pixel that does not converge within + the budget returns zeros. +- ``tolerance`` is the EDM tolerance in Minuit's convention. Migrad, + LevenbergMarquardt and VarPro stop when the estimated distance + to the minimum is below ``0.002 * tolerance``, Fumili below + ``1e-4 * tolerance``. - ``strategy`` is the Minuit2 strategy (0 is fast, 1 is Minuit2's default). - LevenbergMarquardt ignores it. + LevenbergMarquardt and VarPro ignore it. - ``compute_errors`` adds ``par_err`` to the result. Fixed parameters and parameters that end on a limit report an error of 0 with every minimizer. @@ -96,4 +107,4 @@ and ``SetParLimits`` rejects a lower limit at or above the upper one. .. autofunction:: fit .. autoclass:: Minimizer - :members: Migrad, Fumili, LevenbergMarquardt + :members: Migrad, Fumili, LevenbergMarquardt, VarPro diff --git a/include/aare/FitModel.hpp b/include/aare/FitModel.hpp index b5acbc4d..6b2e0ceb 100644 --- a/include/aare/FitModel.hpp +++ b/include/aare/FitModel.hpp @@ -26,17 +26,30 @@ namespace aare { * steps at parameter limits and reuses its buffers between pixels, so a * data cube is fitted without per-pixel allocations. Parameter errors are * sqrt(diag((J^T J)^-1)) over the free parameters. + * - VarPro: the built-in variable projection solver + * (src/VariableProjection.hpp) for models that declare their linear + * parameters (model::is_separable). Every trial point solves the linear + * parameters exactly and the Levenberg-Marquardt iteration runs over the + * nonlinear ones only, which needs fewer evaluations and avoids poor + * start values of the linear parameters. Start values of free linear + * parameters are ignored. A limit on a linear parameter cannot be + * enforced by the linear solve: a pixel whose solution violates one, a + * pixel that does not converge, and every pixel of a model without the + * separable structure are fitted with LevenbergMarquardt instead. + * Parameter errors are the same Gauss-Newton estimates as with + * LevenbergMarquardt. * * With every minimizer, fixed parameters and parameters that end on a limit * report an error of 0 and a failed fit returns zeros. `tolerance` is the - * EDM tolerance in Minuit's convention (Migrad and LevenbergMarquardt stop - * below 0.002 * tolerance, Fumili below 1e-4 * tolerance). `max_calls` - * bounds the work per pixel in the minimizer's own units (Minuit2 function - * calls, or model evaluations for LevenbergMarquardt, which spends one - * evaluation per iteration and two when a step crosses a limit). `strategy` - * only affects the Minuit2 minimizers. + * EDM tolerance in Minuit's convention (Migrad, LevenbergMarquardt and + * VarPro stop below 0.002 * tolerance, Fumili below + * 1e-4 * tolerance). `max_calls` bounds the work per pixel in the + * minimizer's own units (Minuit2 function calls, or model evaluations for + * LevenbergMarquardt and VarPro, which spend one evaluation per + * iteration and two when a step crosses a limit). `strategy` only affects + * the Minuit2 minimizers. */ -enum class Minimizer { Migrad, Fumili, LevenbergMarquardt }; +enum class Minimizer { Migrad, Fumili, LevenbergMarquardt, VarPro }; /** * @brief Fit configuration for one model: parameter limits, fixed diff --git a/include/aare/Models.hpp b/include/aare/Models.hpp index 092937e8..0a744d28 100644 --- a/include/aare/Models.hpp +++ b/include/aare/Models.hpp @@ -5,7 +5,11 @@ #include #include #include +#include +#include +#include #include +#include #include namespace aare::model { @@ -41,6 +45,50 @@ inline double fast_erf(double x) { return fast_erf_from_exp(x, std::exp(-x * x)); } +/** + * @brief exp(x) without a call into libm, for loops that should vectorise. + * + * Range reduction x = k ln2 + r with the integer k taken from the mantissa + * of x * log2(e) + 1.5 * 2^52, a degree-13 Taylor polynomial for exp(r) on + * |r| <= ln2 / 2, and 2^k assembled in the exponent bits, so that a loop + * calling it compiles to vector instructions on SSE2, AVX and NEON. The + * relative error is below 4e-16 for finite x; |x| is clamped to 700. The + * bulk basis functions of the models (basis_columns) use it; eval and + * eval_and_grad keep std::exp. + */ +inline double fast_exp(double x) { + constexpr double log2e = 1.4426950408889634074; + constexpr double magic = 6755399441055744.0; // 1.5 * 2^52 + constexpr double ln2_hi = 0.693145751953125; + constexpr double ln2_lo = 1.42860682030941723212e-6; + x = std::clamp(x, -700.0, 700.0); + const double t = x * log2e + magic; + const double k = t - magic; + const double r = (x - k * ln2_hi) - k * ln2_lo; + double p = 1.0 / 6227020800.0; + p = p * r + 1.0 / 479001600.0; + p = p * r + 1.0 / 39916800.0; + p = p * r + 1.0 / 3628800.0; + p = p * r + 1.0 / 362880.0; + p = p * r + 1.0 / 40320.0; + p = p * r + 1.0 / 5040.0; + p = p * r + 1.0 / 720.0; + p = p * r + 1.0 / 120.0; + p = p * r + 1.0 / 24.0; + p = p * r + 1.0 / 6.0; + p = p * r + 0.5; + p = p * r + 1.0; + p = p * r + 1.0; + std::int64_t ti = 0; + std::int64_t mi = 0; + std::memcpy(&ti, &t, sizeof(ti)); + std::memcpy(&mi, &magic, sizeof(mi)); + const std::int64_t bits = ((ti - mi) + 1023) << 52; + double scale = 0.0; + std::memcpy(&scale, &bits, sizeof(scale)); + return p * scale; +} + /** * @brief Per-parameter metadata: name and optional default bounds. * @@ -55,6 +103,86 @@ struct ParamInfo { inline constexpr double no_bound = std::numeric_limits::infinity(); +/** + * @brief Detects a separable model, one that declares its linear parameters + * for Minimizer::VarPro. + * + * A model whose function is linear in some of its parameters lists them in + * ascending order in `linear_par` and provides + * + * static void basis_and_grad(double x, const std::vector &par, + * std::array &phi, + * std::array, nnl> &dphi) + * + * with nlin = linear_par.size() and nnl = npar - nlin, such that + * + * f(x; par) = sum_j par[linear_par[j]] * phi[j] + * + * and dphi[k][j] = d phi[j] / d par[nonlinear[k]], where nonlinear lists the + * remaining parameter indices in ascending order (see separable_traits). + * basis_and_grad ignores the linear entries of par. The [fit] tests check + * phi and dphi against eval and eval_and_grad for every model. + * + * A model may additionally provide the same quantities for all scan points + * at once, see has_basis_columns: + * + * static void basis_columns(const double *x, ssize_t n, + * const std::vector &par, + * double *phi, double *dphi) + * + * writes phi[j * n + i] = phi_j(x_i) and dphi[(k * nlin + j) * n + i] = + * d phi_j / d par[nonlinear[k]] (x_i), one contiguous column per function. + * Written as a loop free of reductions with the parameters hoisted and + * fast_exp instead of std::exp, it compiles to vector instructions, which + * is where VarPro spends most of its time. The tests check it + * against basis_and_grad. + */ +template +struct is_separable : std::false_type {}; + +template +struct is_separable> + : std::true_type {}; + +/** @brief Detects a separable model that provides basis_columns. */ +template +struct has_basis_columns : std::false_type {}; + +template +struct has_basis_columns> + : std::true_type {}; + +/** @brief The parameter indices not listed in linear, in ascending order. */ +template +constexpr std::array +nonlinear_indices(const std::array &linear) { + std::array out{}; + std::size_t m = 0; + for (std::size_t k = 0; k < Npar; ++k) { + bool is_linear = false; + for (std::size_t j = 0; j < Nlin; ++j) + is_linear = is_linear || linear[j] == k; + if (!is_linear) + out[m++] = k; + } + return out; +} + +/** @brief Index bookkeeping of a separable model, see is_separable. */ +template struct separable_traits { + static constexpr std::size_t npar = Model::npar; + static constexpr std::size_t nlin = Model::linear_par.size(); + static constexpr std::size_t nnl = npar - nlin; + static constexpr std::array nonlinear_par = + nonlinear_indices(Model::linear_par); + static constexpr std::size_t linear_index(int j) { + return Model::linear_par[static_cast(j)]; + } + using Basis = std::array; + using BasisGrad = std::array, nnl>; +}; + // _____________________________________________________________________ // // Pol1 @@ -93,6 +221,17 @@ struct Pol1 { g[1] = x; // df/dp1 } + /** @brief Both parameters are linear: f = p0 * 1 + p1 * x. */ + static constexpr std::array linear_par = {0, 1}; + + static void basis_and_grad( + double x, [[maybe_unused]] const std::vector &par, + std::array &phi, + [[maybe_unused]] std::array, 0> &dphi) { + phi[0] = 1.0; + phi[1] = x; + } + static bool is_valid([[maybe_unused]] const std::vector &par) { return true; // always valid } @@ -152,6 +291,18 @@ struct Pol2 { g[2] = x * x; // df/dp2 } + /** @brief All parameters are linear: f = p0 * 1 + p1 * x + p2 * x^2. */ + static constexpr std::array linear_par = {0, 1, 2}; + + static void basis_and_grad( + double x, [[maybe_unused]] const std::vector &par, + std::array &phi, + [[maybe_unused]] std::array, 0> &dphi) { + phi[0] = 1.0; + phi[1] = x; + phi[2] = x * x; + } + static bool is_valid([[maybe_unused]] const std::vector &par) { return true; // always valid } @@ -259,6 +410,44 @@ struct Gaussian { sig; // df/dsigma = A*e*(x-mu)^2/sig^3 } + /** @brief A is linear: f = A * exp(-(x - mu)^2 / (2 sigma^2)). */ + static constexpr std::array linear_par = {0}; + + static void basis_and_grad(double x, const std::vector &par, + std::array &phi, + std::array, 2> &dphi) { + const double mu = par[1]; + const double sig = par[2]; + + const double dx = x - mu; + const double inv_sig = 1.0 / sig; + const double inv_sig2 = inv_sig * inv_sig; + const double e = std::exp(-0.5 * dx * dx * inv_sig2); + + phi[0] = e; + dphi[0][0] = e * dx * inv_sig2; // d/dmu + dphi[1][0] = e * dx * dx * inv_sig2 * inv_sig; // d/dsigma + } + + /** @brief basis_and_grad for all points, one column per function. */ + static void basis_columns(const double *x, ssize_t n, + const std::vector &par, double *phi, + double *dphi) { + const double mu = par[1]; + const double inv_sig = 1.0 / par[2]; + const double inv_sig2 = inv_sig * inv_sig; + double *e_col = phi; + double *dmu = dphi; + double *dsig = dphi + n; + for (ssize_t i = 0; i < n; ++i) { + const double dx = x[i] - mu; + const double e = fast_exp(-0.5 * dx * dx * inv_sig2); + e_col[i] = e; + dmu[i] = e * dx * inv_sig2; + dsig[i] = e * dx * dx * inv_sig2 * inv_sig; + } + } + /** @brief Reject degenerate sigma (zero width). */ static bool is_valid(const std::vector &par) { return par[2] != 0.0; @@ -385,6 +574,63 @@ struct GaussianErfcPlateau { S * inv_sqrt_2pi * e * dx * inv_sig2; // df/dsigma } + /** @brief A and S are linear: f = A * e + S * step. */ + static constexpr std::array linear_par = {0, 1}; + + static void basis_and_grad(double x, const std::vector &par, + std::array &phi, + std::array, 2> &dphi) { + const double mu = par[2]; + const double sig = par[3]; + + const double dx = x - mu; + const double inv_sig = 1.0 / sig; + const double inv_sig2 = inv_sig * inv_sig; + const double inv_sig3 = inv_sig2 * inv_sig; + const double z = dx * inv_sqrt2 * inv_sig; + + const double e = std::exp(-z * z); + const double step = 0.5 * (1.0 - fast_erf_from_exp(z, e)); + + phi[0] = e; + phi[1] = step; + dphi[0][0] = e * dx * inv_sig2; // de/dmu + dphi[0][1] = inv_sqrt_2pi * e * inv_sig; // dstep/dmu + dphi[1][0] = e * dx * dx * inv_sig3; // de/dsigma + dphi[1][1] = inv_sqrt_2pi * e * dx * inv_sig2; // dstep/dsigma + } + + /** @brief basis_and_grad for all points, one column per function. */ + static void basis_columns(const double *x, ssize_t n, + const std::vector &par, double *phi, + double *dphi) { + const double mu = par[2]; + const double inv_sig = 1.0 / par[3]; + const double inv_sig2 = inv_sig * inv_sig; + const double inv_sig3 = inv_sig2 * inv_sig; + const double cz = inv_sqrt2 * inv_sig; + const double ce = inv_sqrt_2pi * inv_sig; + const double ce2 = inv_sqrt_2pi * inv_sig2; + double *e_col = phi; + double *step_col = phi + n; + double *de_dmu = dphi; + double *dstep_dmu = dphi + n; + double *de_dsig = dphi + 2 * n; + double *dstep_dsig = dphi + 3 * n; + for (ssize_t i = 0; i < n; ++i) { + const double dx = x[i] - mu; + const double z = dx * cz; + const double e = fast_exp(-z * z); + const double step = 0.5 * (1.0 - fast_erf_from_exp(z, e)); + e_col[i] = e; + step_col[i] = step; + de_dmu[i] = e * dx * inv_sig2; + dstep_dmu[i] = ce * e; + de_dsig[i] = e * dx * dx * inv_sig3; + dstep_dsig[i] = ce2 * e * dx; + } + } + /** @brief Reject degenerate or negative width. */ static bool is_valid(const std::vector &par) { return par[3] > 0.0; @@ -545,6 +791,84 @@ struct GaussianChargeSharing { g[5] = N * H; } + /** + * @brief p0, p1 and N are linear: f = p0 * 1 + p1 * (-x) + N * (G + C H). + * C multiplies N, so it stays a nonlinear parameter. + */ + static constexpr std::array linear_par = {0, 1, 4}; + + static void basis_and_grad(double x, const std::vector &par, + std::array &phi, + std::array, 3> &dphi) { + const double mu = par[2]; + const double sig = par[3]; + const double C = par[5]; + + const double inv_sig = 1.0 / sig; + const double inv_sig2 = inv_sig * inv_sig; + const double inv_sig3 = inv_sig2 * inv_sig; + + const double dx = x - mu; + const double u = dx * inv_sig; + + const double G = std::exp(-0.5 * u * u); + const double H = 0.5 * (1.0 - fast_erf_from_exp(u * inv_sqrt2, G)); + + const double dG_dmu = G * dx * inv_sig2; + const double dG_dsig = G * dx * dx * inv_sig3; + const double dH_dmu = inv_sqrt_2pi * G * inv_sig; + const double dH_dsig = inv_sqrt_2pi * G * dx * inv_sig2; + + phi[0] = 1.0; + phi[1] = -x; + phi[2] = G + C * H; + for (int k = 0; k < 3; ++k) { + dphi[k][0] = 0.0; + dphi[k][1] = 0.0; + } + dphi[0][2] = dG_dmu + C * dH_dmu; // d/dmu + dphi[1][2] = dG_dsig + C * dH_dsig; // d/dsigma + dphi[2][2] = H; // d/dC + } + + /** @brief basis_and_grad for all points, one column per function. */ + static void basis_columns(const double *x, ssize_t n, + const std::vector &par, double *phi, + double *dphi) { + const double mu = par[2]; + const double inv_sig = 1.0 / par[3]; + const double C = par[5]; + const double inv_sig2 = inv_sig * inv_sig; + const double inv_sig3 = inv_sig2 * inv_sig; + double *one = phi; + double *minus_x = phi + n; + double *shape = phi + 2 * n; + // d phi_0 and d phi_1 vanish for every nonlinear parameter. + std::fill(dphi, dphi + 2 * n, 0.0); + std::fill(dphi + 3 * n, dphi + 5 * n, 0.0); + std::fill(dphi + 6 * n, dphi + 8 * n, 0.0); + double *d_mu = dphi + 2 * n; + double *d_sig = dphi + 5 * n; + double *d_C = dphi + 8 * n; + for (ssize_t i = 0; i < n; ++i) { + const double xi = x[i]; + const double dx = xi - mu; + const double u = dx * inv_sig; + const double G = fast_exp(-0.5 * u * u); + const double H = 0.5 * (1.0 - fast_erf_from_exp(u * inv_sqrt2, G)); + const double dG_dmu = G * dx * inv_sig2; + const double dG_dsig = G * dx * dx * inv_sig3; + const double dH_dmu = inv_sqrt_2pi * G * inv_sig; + const double dH_dsig = inv_sqrt_2pi * G * dx * inv_sig2; + one[i] = 1.0; + minus_x[i] = -xi; + shape[i] = G + C * H; + d_mu[i] = dG_dmu + C * dH_dmu; + d_sig[i] = dG_dsig + C * dH_dsig; + d_C[i] = H; + } + } + static bool is_valid(const std::vector &par) { return par[3] > 0.0; } @@ -744,6 +1068,133 @@ struct GaussianChargeSharingKb { g[7] = N * kb; } + /** + * @brief p0, p1 and N are linear: f = p0 * 1 + p1 * (-x) + N * (ka + q kb). + * C, r and q multiply N or enter the peak shapes, so they stay nonlinear. + */ + static constexpr std::array linear_par = {0, 1, 4}; + + static void basis_and_grad(double x, const std::vector &par, + std::array &phi, + std::array, 5> &dphi) { + const double mu = par[2]; + const double sig = par[3]; + const double C = par[5]; + const double r = par[6]; + const double q = par[7]; + + const double inv_sig = 1.0 / sig; + const double inv_sig2 = inv_sig * inv_sig; + const double inv_sig3 = inv_sig2 * inv_sig; + + const double mu_b = r * mu; + + const double dx = x - mu; + const double dxb = x - mu_b; + + const double u = dx * inv_sig; + const double ub = dxb * inv_sig; + + const double G = std::exp(-0.5 * u * u); + const double Gb = std::exp(-0.5 * ub * ub); + + const double H = 0.5 * (1.0 - fast_erf_from_exp(u * inv_sqrt2, G)); + const double Hb = 0.5 * (1.0 - fast_erf_from_exp(ub * inv_sqrt2, Gb)); + + const double ka = G + C * H; + const double kb = Gb + C * Hb; + + const double dG_dmu = G * dx * inv_sig2; + const double dG_dsig = G * dx * dx * inv_sig3; + const double dH_dmu = inv_sqrt_2pi * G * inv_sig; + const double dH_dsig = inv_sqrt_2pi * G * dx * inv_sig2; + + const double dGb_dmu_b = Gb * dxb * inv_sig2; + const double dGb_dsig = Gb * dxb * dxb * inv_sig3; + const double dHb_dmu_b = inv_sqrt_2pi * Gb * inv_sig; + const double dHb_dsig = inv_sqrt_2pi * Gb * dxb * inv_sig2; + + const double dka_dmu = dG_dmu + C * dH_dmu; + const double dka_dsig = dG_dsig + C * dH_dsig; + + const double dkb_dmu_b = dGb_dmu_b + C * dHb_dmu_b; + const double dkb_dsig = dGb_dsig + C * dHb_dsig; + + phi[0] = 1.0; + phi[1] = -x; + phi[2] = ka + q * kb; + for (int k = 0; k < 5; ++k) { + dphi[k][0] = 0.0; + dphi[k][1] = 0.0; + } + dphi[0][2] = dka_dmu + q * r * dkb_dmu_b; // d/dmu, mu_b = r * mu + dphi[1][2] = dka_dsig + q * dkb_dsig; // d/dsigma + dphi[2][2] = H + q * Hb; // d/dC + dphi[3][2] = q * mu * dkb_dmu_b; // d/dr + dphi[4][2] = kb; // d/dq + } + + /** @brief basis_and_grad for all points, one column per function. */ + static void basis_columns(const double *x, ssize_t n, + const std::vector &par, double *phi, + double *dphi) { + const double mu = par[2]; + const double inv_sig = 1.0 / par[3]; + const double C = par[5]; + const double r = par[6]; + const double q = par[7]; + const double inv_sig2 = inv_sig * inv_sig; + const double inv_sig3 = inv_sig2 * inv_sig; + const double mu_b = r * mu; + double *one = phi; + double *minus_x = phi + n; + double *shape = phi + 2 * n; + // d phi_0 and d phi_1 vanish for every nonlinear parameter; the + // derivative of phi_2 with respect to nonlinear parameter k is + // column k * 3 + 2. + for (int k = 0; k < 5; ++k) + std::fill(dphi + (k * 3) * n, dphi + (k * 3 + 2) * n, 0.0); + double *d_mu = dphi + 2 * n; + double *d_sig = dphi + 5 * n; + double *d_C = dphi + 8 * n; + double *d_r = dphi + 11 * n; + double *d_q = dphi + 14 * n; + for (ssize_t i = 0; i < n; ++i) { + const double xi = x[i]; + const double dx = xi - mu; + const double dxb = xi - mu_b; + const double u = dx * inv_sig; + const double ub = dxb * inv_sig; + const double G = fast_exp(-0.5 * u * u); + const double Gb = fast_exp(-0.5 * ub * ub); + const double H = 0.5 * (1.0 - fast_erf_from_exp(u * inv_sqrt2, G)); + const double Hb = + 0.5 * (1.0 - fast_erf_from_exp(ub * inv_sqrt2, Gb)); + const double ka = G + C * H; + const double kb = Gb + C * Hb; + const double dG_dmu = G * dx * inv_sig2; + const double dG_dsig = G * dx * dx * inv_sig3; + const double dH_dmu = inv_sqrt_2pi * G * inv_sig; + const double dH_dsig = inv_sqrt_2pi * G * dx * inv_sig2; + const double dGb_dmu_b = Gb * dxb * inv_sig2; + const double dGb_dsig = Gb * dxb * dxb * inv_sig3; + const double dHb_dmu_b = inv_sqrt_2pi * Gb * inv_sig; + const double dHb_dsig = inv_sqrt_2pi * Gb * dxb * inv_sig2; + const double dka_dmu = dG_dmu + C * dH_dmu; + const double dka_dsig = dG_dsig + C * dH_dsig; + const double dkb_dmu_b = dGb_dmu_b + C * dHb_dmu_b; + const double dkb_dsig = dGb_dsig + C * dHb_dsig; + one[i] = 1.0; + minus_x[i] = -xi; + shape[i] = ka + q * kb; + d_mu[i] = dka_dmu + q * r * dkb_dmu_b; + d_sig[i] = dka_dsig + q * dkb_dsig; + d_C[i] = H + q * Hb; + d_r[i] = q * mu * dkb_dmu_b; + d_q[i] = kb; + } + } + static bool is_valid(const std::vector &par) { return par[3] > 0.0; } @@ -848,6 +1299,79 @@ struct RisingScurve { g[5] = step * dx; // df/dp5 } + /** + * @brief p0, p1, A and C are linear: + * f = p0 * 1 + p1 * x + A * S + C * S * (x - mu). + */ + static constexpr std::array linear_par = {0, 1, 4, 5}; + + static void basis_and_grad(double x, const std::vector &par, + std::array &phi, + std::array, 2> &dphi) { + const double p2 = par[2]; + const double p3 = par[3]; + + const double inv_p3 = 1.0 / p3; + const double dx = x - p2; + const double z = dx * inv_sqrt2 * inv_p3; + const double e = std::exp(-z * z); + const double step = 0.5 * (1.0 + fast_erf_from_exp(z, e)); + + const double dSdp2 = -inv_sqrt_2pi * e * inv_p3; + const double dSdp3 = dSdp2 * dx * inv_p3; + + phi[0] = 1.0; + phi[1] = x; + phi[2] = step; + phi[3] = step * dx; + dphi[0][0] = 0.0; + dphi[0][1] = 0.0; + dphi[0][2] = dSdp2; // dS/dmu + dphi[0][3] = dSdp2 * dx - step; // d(S dx)/dmu + dphi[1][0] = 0.0; + dphi[1][1] = 0.0; + dphi[1][2] = dSdp3; // dS/dsigma + dphi[1][3] = dSdp3 * dx; // d(S dx)/dsigma + } + + /** @brief basis_and_grad for all points, one column per function. */ + static void basis_columns(const double *x, ssize_t n, + const std::vector &par, double *phi, + double *dphi) { + const double p2 = par[2]; + const double inv_p3 = 1.0 / par[3]; + const double cz = inv_sqrt2 * inv_p3; + const double ce = -inv_sqrt_2pi * inv_p3; + double *one = phi; + double *x_col = phi + n; + double *step_col = phi + 2 * n; + double *step_dx = phi + 3 * n; + // d phi_0 and d phi_1 vanish for mu and sigma. + std::fill(dphi, dphi + 2 * n, 0.0); + std::fill(dphi + 4 * n, dphi + 6 * n, 0.0); + double *dmu_step = dphi + 2 * n; + double *dmu_step_dx = dphi + 3 * n; + double *dsig_step = dphi + 6 * n; + double *dsig_step_dx = dphi + 7 * n; + for (ssize_t i = 0; i < n; ++i) { + const double xi = x[i]; + const double dx = xi - p2; + const double z = dx * cz; + const double e = fast_exp(-z * z); + const double step = 0.5 * (1.0 + fast_erf_from_exp(z, e)); + const double dSdp2 = ce * e; + const double dSdp3 = dSdp2 * dx * inv_p3; + one[i] = 1.0; + x_col[i] = xi; + step_col[i] = step; + step_dx[i] = step * dx; + dmu_step[i] = dSdp2; + dmu_step_dx[i] = dSdp2 * dx - step; + dsig_step[i] = dSdp3; + dsig_step_dx[i] = dSdp3 * dx; + } + } + /** @brief Reject degenerate width (zero transition width). */ static bool is_valid(const std::vector &par) { return par[3] != 0.0; @@ -989,6 +1513,80 @@ struct FallingScurve { g[5] = step * dx; } + /** + * @brief p0, p1, A and C are linear: + * f = p0 * 1 + p1 * x + A * S + C * S * (x - mu). + */ + static constexpr std::array linear_par = {0, 1, 4, 5}; + + static void basis_and_grad(double x, const std::vector &par, + std::array &phi, + std::array, 2> &dphi) { + const double p2 = par[2]; + const double p3 = par[3]; + + const double inv_p3 = 1.0 / p3; + const double dx = x - p2; + const double z = dx * inv_sqrt2 * inv_p3; + const double e = std::exp(-z * z); + const double step = 0.5 * (1.0 - fast_erf_from_exp(z, e)); + + const double dSdp2 = + inv_sqrt_2pi * e * inv_p3; // sign flipped vs rising + const double dSdp3 = dSdp2 * dx * inv_p3; + + phi[0] = 1.0; + phi[1] = x; + phi[2] = step; + phi[3] = step * dx; + dphi[0][0] = 0.0; + dphi[0][1] = 0.0; + dphi[0][2] = dSdp2; // dS/dmu + dphi[0][3] = dSdp2 * dx - step; // d(S dx)/dmu + dphi[1][0] = 0.0; + dphi[1][1] = 0.0; + dphi[1][2] = dSdp3; // dS/dsigma + dphi[1][3] = dSdp3 * dx; // d(S dx)/dsigma + } + + /** @brief basis_and_grad for all points, one column per function. */ + static void basis_columns(const double *x, ssize_t n, + const std::vector &par, double *phi, + double *dphi) { + const double p2 = par[2]; + const double inv_p3 = 1.0 / par[3]; + const double cz = inv_sqrt2 * inv_p3; + const double ce = inv_sqrt_2pi * inv_p3; // sign flipped vs rising + double *one = phi; + double *x_col = phi + n; + double *step_col = phi + 2 * n; + double *step_dx = phi + 3 * n; + // d phi_0 and d phi_1 vanish for mu and sigma. + std::fill(dphi, dphi + 2 * n, 0.0); + std::fill(dphi + 4 * n, dphi + 6 * n, 0.0); + double *dmu_step = dphi + 2 * n; + double *dmu_step_dx = dphi + 3 * n; + double *dsig_step = dphi + 6 * n; + double *dsig_step_dx = dphi + 7 * n; + for (ssize_t i = 0; i < n; ++i) { + const double xi = x[i]; + const double dx = xi - p2; + const double z = dx * cz; + const double e = fast_exp(-z * z); + const double step = 0.5 * (1.0 - fast_erf_from_exp(z, e)); + const double dSdp2 = ce * e; + const double dSdp3 = dSdp2 * dx * inv_p3; + one[i] = 1.0; + x_col[i] = xi; + step_col[i] = step; + step_dx[i] = step * dx; + dmu_step[i] = dSdp2; + dmu_step_dx[i] = dSdp2 * dx - step; + dsig_step[i] = dSdp3; + dsig_step_dx[i] = dSdp3 * dx; + } + } + /** @brief Reject degenerate width (zero transition width). */ static bool is_valid(const std::vector &par) { return par[3] != 0.0; diff --git a/python/src/fit.hpp b/python/src/fit.hpp index 9fb9ac79..f26cc20c 100644 --- a/python/src/fit.hpp +++ b/python/src/fit.hpp @@ -34,23 +34,27 @@ template void bind_fit_model(py::module &m, const char *name) { ---------- strategy : int Minuit2 strategy (0 = fast, 1 = Minuit2's default). Ignored - by ``Minimizer.LevenbergMarquardt``. + by ``Minimizer.LevenbergMarquardt`` and + ``Minimizer.VarPro``. max_calls : int Work budget per pixel: Minuit2 function calls, or model - evaluations for ``Minimizer.LevenbergMarquardt`` (one per - iteration, two when a step crosses a limit). 0 selects - Minuit's default budget. + evaluations for ``Minimizer.LevenbergMarquardt`` and + ``Minimizer.VarPro`` (one per iteration, two + when a step crosses a limit). 0 selects Minuit's default + budget. tolerance : float - EDM tolerance in Minuit's convention. Migrad and - LevenbergMarquardt stop below ``0.002 * tolerance``, Fumili - below ``1e-4 * tolerance``. + EDM tolerance in Minuit's convention. Migrad, + LevenbergMarquardt and VarPro stop below + ``0.002 * tolerance``, Fumili below ``1e-4 * tolerance``. compute_errors : bool Also return parameter errors: from Hesse with Migrad, from - the linearised covariance with Fumili and LevenbergMarquardt. - Fixed parameters and parameters ending on a limit report 0. + the linearised covariance with Fumili, LevenbergMarquardt + and VarPro. Fixed parameters and parameters + ending on a limit report 0. minimizer : Minimizer - ``Minimizer.Migrad`` (default), ``Minimizer.Fumili`` or - ``Minimizer.LevenbergMarquardt``. + ``Minimizer.Migrad`` (default), ``Minimizer.Fumili``, + ``Minimizer.LevenbergMarquardt`` or + ``Minimizer.VarPro``. )doc") .def("SetParLimits", py::overload_cast(&FM::SetParLimits), @@ -118,8 +122,9 @@ template void bind_fit_model(py::module &m, const char *name) { Fit this model to 1D or 3D data using Minuit2. The minimizer is selected by the ``minimizer`` property - (``Minimizer.Migrad`` by default, ``Minimizer.Fumili`` or - ``Minimizer.LevenbergMarquardt``). + (``Minimizer.Migrad`` by default, ``Minimizer.Fumili``, + ``Minimizer.LevenbergMarquardt`` or + ``Minimizer.VarPro``). Parameters ---------- @@ -319,12 +324,26 @@ void define_fit_bindings(py::module &m) { ``sqrt(diag((J^T J)^-1))``; ``max_calls`` counts model evaluations and ``strategy`` is ignored. + ``VarPro`` is the built-in variable projection solver. + Every bundled model is linear in some of its parameters (the + amplitude of a Gaussian, the baseline and amplitude of an S-curve, + every parameter of a polynomial); each trial point solves those + exactly and the Levenberg-Marquardt iteration runs over the + remaining nonlinear parameters only. It needs fewer evaluations than + ``LevenbergMarquardt`` and cannot be misled by poor start values of + the linear parameters, which it ignores. A pixel whose linear + solution violates a limit on a linear parameter, or that does not + converge, is refitted with ``LevenbergMarquardt``. Its parameter + errors are the same Gauss-Newton estimates; ``max_calls`` counts + evaluations and ``strategy`` is ignored. + With every minimizer, fixed parameters and parameters that end on a limit report an error of 0 and a failed fit returns zeros. )doc") .value("Migrad", aare::Minimizer::Migrad) .value("Fumili", aare::Minimizer::Fumili) - .value("LevenbergMarquardt", aare::Minimizer::LevenbergMarquardt); + .value("LevenbergMarquardt", aare::Minimizer::LevenbergMarquardt) + .value("VarPro", aare::Minimizer::VarPro); // ── Bind model classes ────────────────────────────────────────── bind_fit_model(m, "Gaussian"); diff --git a/python/tests/test_Fit.py b/python/tests/test_Fit.py index 77df6f81..a5436b60 100644 --- a/python/tests/test_Fit.py +++ b/python/tests/test_Fit.py @@ -8,6 +8,7 @@ MINIMIZERS = [ aare.Minimizer.Migrad, aare.Minimizer.Fumili, aare.Minimizer.LevenbergMarquardt, + aare.Minimizer.VarPro, ] @@ -16,7 +17,10 @@ def gaussian(x, A, mu, sigma): def uses_minuit2(minimizer): - return minimizer != aare.Minimizer.LevenbergMarquardt + return minimizer not in ( + aare.Minimizer.LevenbergMarquardt, + aare.Minimizer.VarPro, + ) @pytest.fixture(params=MINIMIZERS, ids=lambda m: m.name) @@ -45,7 +49,12 @@ def test_gaussian_model_evaluates_and_fits_data(minimizer): def test_minimizer_can_be_selected(): assert aare.Gaussian().minimizer == aare.Minimizer.Migrad - assert set(aare.Minimizer.__members__) == {"Migrad", "Fumili", "LevenbergMarquardt"} + assert set(aare.Minimizer.__members__) == { + "Migrad", + "Fumili", + "LevenbergMarquardt", + "VarPro", + } for minimizer in MINIMIZERS: assert aare.Gaussian(minimizer=minimizer).minimizer == minimizer diff --git a/src/DampedGaussNewton.hpp b/src/DampedGaussNewton.hpp new file mode 100644 index 00000000..b7bdd698 --- /dev/null +++ b/src/DampedGaussNewton.hpp @@ -0,0 +1,520 @@ +// SPDX-License-Identifier: MPL-2.0 +#pragma once + +#include +#include +#include +#include +#include + +namespace aare::detail { + +/** @brief Outcome of one of the built-in minimizers. */ +struct MinimizerResult { + bool valid = false; + int calls = 0; + int iterations = 0; + double chi2 = 0.0; +}; + +/** + * @brief Cholesky factorisation of the leading m x m block of the symmetric + * positive-definite matrix A. Writes the lower triangle of L with + * A = L L^T. Returns false when the block is not positive definite. + */ +template +bool cholesky_factor(int m, const double A[N][N], double L[N][N]) { + for (int i = 0; i < m; ++i) { + for (int j = 0; j <= i; ++j) { + double sum = A[i][j]; + for (int k = 0; k < j; ++k) + sum -= L[i][k] * L[j][k]; + if (i == j) { + if (!(sum > 0.0)) + return false; + L[i][i] = std::sqrt(sum); + } else { + L[i][j] = sum / L[j][j]; + } + } + } + return true; +} + +/** @brief Solves L L^T z = b for the leading m x m block of the factor L. */ +template +void cholesky_substitute(int m, const double L[N][N], const double *b, + double *z) { + double t[N]; + for (int i = 0; i < m; ++i) { + double sum = b[i]; + for (int k = 0; k < i; ++k) + sum -= L[i][k] * t[k]; + t[i] = sum / L[i][i]; + } + for (int i = m - 1; i >= 0; --i) { + double sum = t[i]; + for (int k = i + 1; k < m; ++k) + sum -= L[k][i] * z[k]; + z[i] = sum / L[i][i]; + } +} + +/** @brief Solves A z = b for the leading m x m block of the SPD matrix A. */ +template +bool cholesky_solve(int m, const double A[N][N], const double *b, double *z) { + double L[N][N]; + if (!cholesky_factor(m, A, L)) + return false; + cholesky_substitute(m, L, b, z); + return true; +} + +/** @brief Inverse of the leading m x m block of the SPD matrix A. */ +template +bool invert_spd(int m, const double A[N][N], double inv[N][N]) { + double e[N]; + double col[N]; + for (int j = 0; j < m; ++j) { + for (int i = 0; i < m; ++i) + e[i] = (i == j) ? 1.0 : 0.0; + if (!cholesky_solve(m, A, e, col)) + return false; + for (int i = 0; i < m; ++i) + inv[i][j] = col[i]; + } + return true; +} + +/** + * @brief Gauss-Newton parameter errors sqrt(diag(H^-1)) over the parameters + * that are not skipped, where H is the upper triangle of the Gauss-Newton + * Hessian of chi2 / 2 over all N parameters. Skipped parameters (fixed, or + * frozen on a limit) report 0. Returns false, leaving all errors 0, when the + * block cannot be inverted. + */ +template +bool gauss_newton_errors(const double H[N][N], const bool *skip, double *err) { + std::fill(err, err + N, 0.0); + double A[N][N]; + double inv[N][N]; + int m = 0; + for (int a = 0; a < N; ++a) { + if (skip[a]) + continue; + int mb = 0; + for (int c = 0; c < N; ++c) { + if (skip[c]) + continue; + A[m][mb++] = a <= c ? H[a][c] : H[c][a]; + } + ++m; + } + if (m == 0 || !invert_spd(m, A, inv)) + return false; + int i = 0; + for (int a = 0; a < N; ++a) { + if (skip[a]) + continue; + err[a] = std::sqrt(std::max(inv[i][i], 0.0)); + ++i; + } + return true; +} + +/** + * @brief Damped Gauss-Newton iteration shared by the built-in minimizers. + * + * The problem is described by an evaluator called as + * `eval(const double *p, Normal &out)`: at the point p (all NP parameters, + * fixed ones at their values) it fills out.F = chi2 / 2, the gradient + * out.g = dF/dp and the upper triangle of the Gauss-Newton Hessian out.H, + * both over all NP parameters, and returns false when the point is not + * admissible. LevenbergMarquardt evaluates residuals and the Jacobian of + * the full model, VariableProjection the projected problem of the nonlinear + * parameters. The driver extracts the free parameters and implements + * + * - Marquardt damping u * diag(H), starting at u = damping0, with the + * update rule of Nielsen (Madsen, Nielsen and Tingleff, Methods for + * Non-Linear Least Squares Problems, 2004). The damping also bounds the + * accuracy of the last accepted step, so a solver that converges in few + * iterations should start with a small damping. + * - Limits: a trial step that leaves the allowed box is reflected at the + * limit. The step truncated at the limit is evaluated as well and the + * better of the two is kept. A parameter sitting on a limit whose gradient + * points outward is frozen for that iteration. + * - Convergence in Minuit's convention: the iteration stops when the + * estimated distance to the minimum of chi2 (EDM, from the Gauss-Newton + * Hessian) drops below 0.002 * tolerance, or when the step becomes + * negligible. + * - max_calls bounds the number of evaluations: the initial point and every + * trial point count, so an iteration costs one evaluation, or two when a + * limit is crossed. + * - With reject_degenerate, a trial point at which the model has lost its + * sensitivity to a free parameter (its Hessian diagonal collapsed by more + * than eight orders of magnitude) is rejected like an inadmissible point, + * so the damping grows and a shorter step is tried instead. + * - With polish, a converged fit takes one more undamped Gauss-Newton step + * and keeps it when it lowers chi2, see polish_step(). + * + * After fit() the final point, its Normal and the frozen parameters remain + * available for the error computation. + */ +template class DampedGaussNewton { + public: + static_assert(NP >= 1 && NP <= 16, + "DampedGaussNewton keeps NP x NP matrices on the stack"); + + struct Normal { + double F = 0.0; // chi2 / 2 + double g[NP] = {}; // dF/dp + double H[NP][NP] = {}; // Gauss-Newton Hessian, upper triangle + std::array aux{}; // evaluator payload, kept with the point + }; + + struct Options { + double tolerance = 0.5; // EDM tolerance, Minuit's convention + int max_calls = 100; // budget of evaluations + double damping0 = 0.1; // initial Marquardt damping + bool reject_degenerate = false; + bool polish = false; + }; + + template + MinimizerResult fit(Eval &eval, const double *start, const double *lower, + const double *upper, const bool *fixed, + const Options &opt) { + nfree_ = 0; + for (int k = 0; k < NP; ++k) { + lower_[k] = lower[k]; + upper_[k] = upper[k]; + p_[k] = start[k]; + fixed_[k] = fixed[k]; + free_pos_[k] = -1; + if (!fixed[k]) { + free_pos_[k] = nfree_; + free_[nfree_++] = k; + } + } + for (int a = 0; a < nfree_; ++a) { + q_[a] = p_[free_[a]]; + active_[a] = false; + } + + MinimizerResult res; + if (!evaluate(eval, q_, cur_)) + return res; + res.calls = 1; + const double edm_target = 0.002 * opt.tolerance; + const int max_calls = opt.max_calls; + const bool reject_degenerate = opt.reject_degenerate; + double u = opt.damping0; // damping relative to diag(H) + double nu = 2.0; + + for (;; ++res.iterations) { + extract_free(); + mark_active(); + if (gmax_ < 1e-12) { + res.valid = true; + break; + } + // Solve for the damped step first. The solve yields + // g^T (H + u D)^-1 g, a lower bound on the EDM, so the undamped + // factorisation behind edm() is only needed when that bound + // leaves room for convergence. + const bool have_step = solve_step(u); + if (!have_step || bound_ < edm_margin * edm_target) { + const double e = edm(); + if (e < edm_target) { + // A point already deep inside the tolerance needs no + // polishing. + if (opt.polish && e >= 0.01 * edm_target && + res.calls < max_calls) + polish_step(eval, res, reject_degenerate); + res.valid = true; + break; + } + } + if (res.calls >= max_calls) + break; + if (!have_step) { + u *= nu; + nu *= 2.0; + if (u > 1e32) + break; + continue; + } + double step2 = 0.0; + double q2 = 0.0; + bool crossed = false; + for (int a = 0; a < nfree_; ++a) { + const int k = free_[a]; + const double v = q_[a] + delta_[a]; + const double truncated = std::clamp(v, lower_[k], upper_[k]); + double reflected = v; + if (truncated != v) { + crossed = true; + if (v < lower_[k]) + reflected = 2.0 * lower_[k] - v; + if (v > upper_[k]) + reflected = 2.0 * upper_[k] - v; + reflected = std::clamp(reflected, lower_[k], upper_[k]); + } + q_new_[a] = reflected; + q_alt_[a] = truncated; + const double d = q_new_[a] - q_[a]; + step2 += d * d; + q2 += q_[a] * q_[a]; + } + if (std::sqrt(step2) <= 1e-8 * (std::sqrt(q2) + 1e-8)) { + res.valid = true; + break; + } + ++res.calls; + bool ok = evaluate(eval, q_new_, trial_) && + !(reject_degenerate && degenerate(trial_)); + if (crossed) { + // Also try the step truncated at the limit, keep the better. + ++res.calls; + const bool ok_alt = evaluate(eval, q_alt_, alt_) && + !(reject_degenerate && degenerate(alt_)); + if (ok_alt && (!ok || alt_.F < trial_.F)) { + q_new_ = q_alt_; + std::swap(trial_, alt_); + ok = true; + } + } + const double rho = ok ? (cur_.F - trial_.F) / predicted_ : -1.0; + if (rho > 0.0) { + // The trial becomes the current point. + q_ = q_new_; + std::swap(cur_, trial_); + for (int a = 0; a < nfree_; ++a) + p_[free_[a]] = q_[a]; + const double t = 2.0 * rho - 1.0; + u *= std::max(1.0 / 3.0, 1.0 - t * t * t); + nu = 2.0; + } else { + u *= nu; + nu *= 2.0; + if (u > 1e32) + break; + } + } + if (!res.valid) + return res; + + // p_ tracks the last evaluated point, which may be a rejected trial. + for (int a = 0; a < nfree_; ++a) + p_[free_[a]] = q_[a]; + res.chi2 = 2.0 * cur_.F; + return res; + } + + /** @brief The final point, all NP parameters. */ + const double *point() const { return p_; } + + /** @brief F, gradient and Hessian at the final point. */ + const Normal &normal() const { return cur_; } + + /** @brief True for a fixed parameter or one frozen on a limit. */ + bool skipped(int k) const { + return fixed_[k] || (free_pos_[k] >= 0 && active_[free_pos_[k]]); + } + + /** @brief Gauss-Newton errors at the final point, see + * gauss_newton_errors(). */ + void errors(double *err_out) const { + bool skip[NP]; + for (int k = 0; k < NP; ++k) + skip[k] = skipped(k); + gauss_newton_errors(cur_.H, skip, err_out); + } + + private: + template + bool evaluate(Eval &eval, const std::array &q, Normal &out) { + for (int a = 0; a < nfree_; ++a) + p_[free_[a]] = q[a]; + return eval(static_cast(p_), out); + } + + // Gradient and Hessian of the free parameters at the current point. + void extract_free() { + gmax_ = 0.0; + for (int a = 0; a < nfree_; ++a) { + const int ka = free_[a]; + g_[a] = cur_.g[ka]; + for (int b = 0; b < nfree_; ++b) { + const int kb = free_[b]; + jtj_[a][b] = ka <= kb ? cur_.H[ka][kb] : cur_.H[kb][ka]; + } + gmax_ = std::max(gmax_, std::abs(g_[a])); + } + } + + // A parameter sitting on a limit with the descent direction pointing + // outward is frozen for this iteration. The test deliberately looks at + // the last evaluated point p_, which after a rejected step is the trial + // rather than the current point q_: a rejected trial that came off the + // limit unfreezes the parameter, and the reflected step can then leave + // the limit. Testing q_ instead was measured to leave more fits stuck on + // a limit with a higher chi2 (FallingScurve and GaussianChargeSharing + // cubes), although it helped Pol2. + void mark_active() { + for (int a = 0; a < nfree_; ++a) { + const int k = free_[a]; + active_[a] = (p_[k] <= lower_[k] && g_[a] > 0.0) || + (p_[k] >= upper_[k] && g_[a] < 0.0); + } + } + + // A trial point where the Hessian diagonal of a free parameter collapsed + // relative to the current point: the model no longer responds to that + // parameter there, so the iteration could stop at a spurious stationary + // point (a step function with the width driven to zero, for example). + bool degenerate(const Normal &trial) const { + for (int a = 0; a < nfree_; ++a) { + if (active_[a]) + continue; + const int k = free_[a]; + if (!(trial.H[k][k] >= 1e-8 * cur_.H[k][k])) + return true; + } + return false; + } + + // Minuit's estimated distance to the minimum of chi2, using the + // Gauss-Newton Hessian: g^T H^-1 g with g the gradient of chi2/2. + // Also keeps the undamped Gauss-Newton step for polish_step(). + double edm() { + double A[NP][NP]; + double b[NP]; + double z[NP]; + const int m = pack(A, b, 0.0); + for (int a = 0; a < nfree_; ++a) + gn_step_[a] = 0.0; + if (m == 0) + return 0.0; + if (!cholesky_solve(m, A, b, z)) + return std::numeric_limits::infinity(); + double e = 0.0; + int i = 0; + for (int a = 0; a < nfree_; ++a) { + if (active_[a]) + continue; + gn_step_[a] = -z[i]; + e += b[i] * z[i]; + ++i; + } + return e; + } + + // One undamped Gauss-Newton step from a point that passed the EDM test, + // kept when it lowers F. The test stops the iteration within the + // tolerance, but a solver that converges in few iterations stops while + // the damping is still noticeable, so its last accepted step fell short + // of the minimum by about the damping factor. The extra evaluation lands + // at the minimum to the accuracy of a full Gauss-Newton step. + template + void polish_step(Eval &eval, MinimizerResult &res, bool reject_degenerate) { + for (int a = 0; a < nfree_; ++a) { + const int k = free_[a]; + q_new_[a] = std::clamp(q_[a] + gn_step_[a], lower_[k], upper_[k]); + } + ++res.calls; + const bool ok = evaluate(eval, q_new_, trial_) && + !(reject_degenerate && degenerate(trial_)) && + trial_.F < cur_.F; + if (ok) { + q_ = q_new_; + std::swap(cur_, trial_); + extract_free(); + } + for (int a = 0; a < nfree_; ++a) + p_[free_[a]] = q_[a]; + if (ok) + mark_active(); + } + + // Copy the non-frozen block of H (+ u * diag) and g into A and b. + int pack(double A[NP][NP], double b[NP], double u) const { + double mean_diag = 0.0; + for (int a = 0; a < nfree_; ++a) + mean_diag += jtj_[a][a]; + mean_diag = std::max(mean_diag / std::max(nfree_, 1), 1e-300); + int m = 0; + for (int a = 0; a < nfree_; ++a) { + if (active_[a]) + continue; + int mb = 0; + for (int c = 0; c < nfree_; ++c) { + if (active_[c]) + continue; + A[m][mb++] = jtj_[a][c]; + } + A[m][m] += u * std::max(jtj_[a][a], 1e-10 * mean_diag); + b[m] = g_[a]; + ++m; + } + return m; + } + + // Damped Gauss-Newton step (H + u D) delta = -g for the non-frozen + // parameters, the predicted decrease of F for that step, and + // bound_ = g^T (H + u D)^-1 g, which is at most the EDM because u D is + // positive semi-definite. + bool solve_step(double u) { + double A[NP][NP]; + double b[NP]; + double z[NP]; + const int m = pack(A, b, u); + for (int a = 0; a < nfree_; ++a) + delta_[a] = 0.0; + predicted_ = 0.0; + bound_ = 0.0; + if (m == 0) + return true; + if (!cholesky_solve(m, A, b, z)) + return false; + int i = 0; + for (int a = 0; a < nfree_; ++a) { + if (active_[a]) + continue; + delta_[a] = -z[i]; + bound_ += b[i] * z[i]; + predicted_ += + 0.5 * delta_[a] * + (u * std::max(jtj_[a][a], 1e-300) * delta_[a] - g_[a]); + ++i; + } + return predicted_ > 0.0; + } + + int nfree_ = 0; + int free_[NP] = {}; + int free_pos_[NP] = {}; + bool fixed_[NP] = {}; + bool active_[NP] = {}; + double lower_[NP] = {}; + double upper_[NP] = {}; + double p_[NP] = {}; + std::array q_{}; + std::array q_new_{}; + std::array q_alt_{}; + double g_[NP] = {}; + double delta_[NP] = {}; + double gn_step_[NP] = {}; + double jtj_[NP][NP] = {}; + // The two solves that bound the EDM differ by rounding only when the + // damping is negligible; the margin keeps the stop decision exact. + static constexpr double edm_margin = 1.0 + 1e-6; + double gmax_ = 0.0; + double predicted_ = 0.0; + double bound_ = 0.0; + Normal cur_; + Normal trial_; + Normal alt_; +}; + +} // namespace aare::detail diff --git a/src/Fit.cpp b/src/Fit.cpp index 855f3486..37187483 100644 --- a/src/Fit.cpp +++ b/src/Fit.cpp @@ -3,6 +3,7 @@ #include "FitHelpers.hpp" #include "FitMinuit2.hpp" #include "LevenbergMarquardt.hpp" +#include "VariableProjection.hpp" #include "aare/Models.hpp" #include "aare/utils/par.hpp" #include "aare/utils/task.hpp" @@ -10,6 +11,7 @@ #include #include #include +#include #include namespace aare { @@ -117,11 +119,13 @@ std::vector FitModel::GetParNames() const { namespace detail { +struct NotSeparable {}; + /** * @brief Fits pixels with the minimizer selected by the model. * - * Keep one instance per thread: the Levenberg-Marquardt buffers are reused - * between pixels so that a data cube is fitted without per-pixel + * Keep one instance per thread: the buffers of the built-in minimizers are + * reused between pixels so that a data cube is fitted without per-pixel * allocations. The Minuit2 minimizers create their state per pixel. */ template class PixelFitter { @@ -139,18 +143,40 @@ template class PixelFitter { double *err_out, double &chi2) { const auto start = start_values(model, x, y); - if (model.minimizer() == Minimizer::LevenbergMarquardt) { - const auto res = - lm_.fit(model, x, y, y_err, start, par_out, err_out); - chi2 = res.valid ? res.chi2 : 0.0; - return res.valid; + if (model.minimizer() == Minimizer::VarPro) { + // A model without the separable structure, a pixel that does + // not converge and a pixel whose linear solution violates a + // limit fall back to the full solver. + if constexpr (model::is_separable::value) { + const auto res = + vp_.fit(model, x, y, y_err, start, par_out, err_out); + if (res.valid) { + chi2 = res.chi2; + return true; + } + } + return fit_lm(model, x, y, y_err, start, par_out, err_out, chi2); } + if (model.minimizer() == Minimizer::LevenbergMarquardt) + return fit_lm(model, x, y, y_err, start, par_out, err_out, chi2); return fit_pixel_minuit2(model, x, y, y_err, start, par_out, err_out, chi2); } private: + bool fit_lm(const FitModel &model, NDView x, + NDView y, NDView y_err, + const std::array &start, double *par_out, + double *err_out, double &chi2) { + const auto res = lm_.fit(model, x, y, y_err, start, par_out, err_out); + chi2 = res.valid ? res.chi2 : 0.0; + return res.valid; + } + LevenbergMarquardt lm_; + std::conditional_t::value, + VariableProjection, NotSeparable> + vp_; }; } // namespace detail diff --git a/src/Fit.test.cpp b/src/Fit.test.cpp index fc6b7226..75b9cf8d 100644 --- a/src/Fit.test.cpp +++ b/src/Fit.test.cpp @@ -76,11 +76,15 @@ const char *name(Minimizer m) { return "Fumili"; case Minimizer::LevenbergMarquardt: return "LevenbergMarquardt"; + case Minimizer::VarPro: + return "VarPro"; } return "unknown"; } -bool uses_minuit2(Minimizer m) { return m != Minimizer::LevenbergMarquardt; } +bool uses_minuit2(Minimizer m) { + return m != Minimizer::LevenbergMarquardt && m != Minimizer::VarPro; +} template aare::FitModel make_model(Minimizer minimizer, bool errors = false, @@ -137,7 +141,8 @@ void check_minimizers_agree(const std::string &label, Data d, bool weighted, const auto ref = run(Minimizer::Migrad); for (const auto candidate : - {Minimizer::Fumili, Minimizer::LevenbergMarquardt}) { + {Minimizer::Fumili, Minimizer::LevenbergMarquardt, + Minimizer::VarPro}) { INFO(label << ": " << name(candidate) << " vs Migrad"); const auto res = run(candidate); REQUIRE(res.size() == ref.size()); @@ -185,8 +190,9 @@ void check_minimizers_agree(const std::string &label, Data d, bool weighted, } // namespace TEST_CASE("Fit unweighted and weighted noise-free Gaussian data", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); auto d = gaussian_data(); @@ -217,8 +223,9 @@ TEST_CASE("Fit unweighted and weighted noise-free Gaussian data", "[fit]") { } TEST_CASE("Fit weighted Pol1 errors match the analytic formula", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); constexpr ssize_t n = 20; NDArray x({n}); @@ -248,8 +255,9 @@ TEST_CASE("Fit weighted Pol1 errors match the analytic formula", "[fit]") { } TEST_CASE("Fit with a fixed parameter", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); auto d = gaussian_data(); @@ -289,8 +297,9 @@ TEST_CASE("Fit with a fixed parameter", "[fit]") { } TEST_CASE("Fit with parameter limits", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); auto d = gaussian_data(); const double margin = limit_margin(minimizer); @@ -349,8 +358,9 @@ TEST_CASE("Fit with parameter limits", "[fit]") { } TEST_CASE("Fit honours user start values", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); auto d = gaussian_data(); @@ -381,8 +391,9 @@ TEST_CASE("Fit honours user start values", "[fit]") { } TEST_CASE("Fit reports failure as zeros when max_calls is exhausted", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); auto d = gaussian_data(); @@ -440,8 +451,9 @@ TEST_CASE("Copying a FitModel keeps the minimizer", "[fit]") { } TEST_CASE("Fit rising and falling S-curves", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); const std::vector truth{10.0, 0.1, 50.0, 3.0, 1000.0, 2.0}; @@ -498,8 +510,9 @@ TEST_CASE("Fit rising and falling S-curves", "[fit]") { } TEST_CASE("Fit a Gaussian data cube in parallel", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); NDArray x({n_points}); NDArray values({n_points}); @@ -536,8 +549,9 @@ TEST_CASE("Fit a Gaussian data cube in parallel", "[fit]") { } TEST_CASE("Fit a data cube with errors matches per-pixel fits", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); constexpr ssize_t rows = 4; constexpr ssize_t cols = 3; @@ -594,8 +608,9 @@ TEST_CASE("Fit a data cube with errors matches per-pixel fits", "[fit]") { TEST_CASE("Fitting with every parameter fixed returns the fixed values", "[fit]") { - const auto minimizer = GENERATE(Minimizer::Migrad, Minimizer::Fumili, - Minimizer::LevenbergMarquardt); + const auto minimizer = + GENERATE(Minimizer::Migrad, Minimizer::Fumili, + Minimizer::LevenbergMarquardt, Minimizer::VarPro); INFO("minimizer " << name(minimizer)); auto d = gaussian_data(); @@ -725,7 +740,8 @@ TEST_CASE("All minimizers agree on a noisy data cube", "[fit]") { const auto ref = fit_cube(Minimizer::Migrad); for (const auto candidate : - {Minimizer::Fumili, Minimizer::LevenbergMarquardt}) { + {Minimizer::Fumili, Minimizer::LevenbergMarquardt, + Minimizer::VarPro}) { INFO(name(candidate) << " vs Migrad"); const auto res = fit_cube(candidate); for (ssize_t row = 0; row < rows; ++row) { @@ -742,3 +758,118 @@ TEST_CASE("All minimizers agree on a noisy data cube", "[fit]") { } } } + +namespace { + +// The separable structure of a model must reproduce eval and eval_and_grad: +// f = sum_j par[linear_par[j]] * phi[j], df/dpar[linear_par[j]] = phi[j] and +// df/dpar[nonlinear[k]] = sum_j par[linear_par[j]] * dphi[k][j]. +template +void check_separable(const std::vector &par, + const std::vector &xs) { + using Traits = aare::model::separable_traits; + static_assert(aare::model::is_separable::value); + typename Traits::Basis phi{}; + typename Traits::BasisGrad dphi{}; + std::array g{}; + double f = 0.0; + for (const double x : xs) { + INFO("x = " << x); + Model::basis_and_grad(x, par, phi, dphi); + Model::eval_and_grad(x, par, f, g); + double sum = 0.0; + for (std::size_t j = 0; j < Traits::nlin; ++j) { + const auto idx = Model::linear_par[j]; + sum += par[idx] * phi[j]; + CHECK(g[idx] == Approx(phi[j]).epsilon(1e-12).margin(1e-14)); + } + CHECK(Model::eval(x, par) == Approx(sum).epsilon(1e-12).margin(1e-14)); + CHECK(f == Approx(sum).epsilon(1e-12).margin(1e-14)); + for (std::size_t k = 0; k < Traits::nnl; ++k) { + double d = 0.0; + for (std::size_t j = 0; j < Traits::nlin; ++j) + d += par[Model::linear_par[j]] * dphi[k][j]; + CHECK(g[Traits::nonlinear_par[k]] == + Approx(d).epsilon(1e-12).margin(1e-14)); + } + } + // The bulk basis pass must agree with the per-point one. + if constexpr (aare::model::has_basis_columns::value) { + const auto n = static_cast(xs.size()); + std::vector cols(Traits::nlin * xs.size()); + std::vector dcols(Traits::nnl * Traits::nlin * xs.size()); + Model::basis_columns(xs.data(), n, par, cols.data(), dcols.data()); + for (ssize_t i = 0; i < n; ++i) { + INFO("x = " << xs[static_cast(i)]); + Model::basis_and_grad(xs[static_cast(i)], par, phi, + dphi); + for (std::size_t j = 0; j < Traits::nlin; ++j) + CHECK(cols[j * xs.size() + static_cast(i)] == + Approx(phi[j]).epsilon(1e-12).margin(1e-14)); + for (std::size_t k = 0; k < Traits::nnl; ++k) + for (std::size_t j = 0; j < Traits::nlin; ++j) + CHECK(dcols[(k * Traits::nlin + j) * xs.size() + + static_cast(i)] == + Approx(dphi[k][j]).epsilon(1e-12).margin(1e-14)); + } + } +} + +} // namespace + +TEST_CASE("Separable models are consistent with eval and eval_and_grad", + "[fit]") { + using namespace aare::model; + const std::vector xs = {-4.0, -1.5, 0.3, 2.0, + 7.5, 40.0, 55.0, 120.0}; + check_separable({2.0, -0.5}, xs); + check_separable({2.0, -0.5, 0.05}, xs); + check_separable({120.0, 0.8, 1.3}, xs); + check_separable({80.0, 15.0, 2.0, 1.5}, xs); + check_separable({20.0, 0.05, 50.0, 4.0, 1500.0, 0.3}, + xs); + check_separable( + {20.0, 0.05, 50.0, 4.0, 1500.0, 0.3, 1.1, 0.1}, xs); + check_separable({10.0, 0.1, 50.0, 3.0, 1000.0, 2.0}, xs); + check_separable({10.0, 0.1, 50.0, 3.0, 1000.0, 2.0}, xs); +} + +TEST_CASE("VarPro falls back to LevenbergMarquardt for a limit " + "on a linear parameter", + "[fit]") { + using aare::model::FallingScurve; + const std::vector truth{10.0, 0.1, 50.0, 3.0, 1000.0, 2.0}; + auto d = scurve_data(truth); + + auto vp = make_model(Minimizer::VarPro, true, 500); + auto lm = + make_model(Minimizer::LevenbergMarquardt, true, 500); + // The unconstrained amplitude is 1000, so the linear solve cannot + // honour this limit and the pixel is refitted with the full solver. + vp.SetParLimits("A", 0.0, 500.0); + lm.SetParLimits("A", 0.0, 500.0); + + const auto res_vp = + aare::fit_pixel(vp, d.x.view(), d.y.view(), d.y_err.view()); + const auto res_lm = + aare::fit_pixel(lm, d.x.view(), d.y.view(), d.y_err.view()); + REQUIRE(res_vp.size() == res_lm.size()); + for (ssize_t k = 0; k < res_vp.size(); ++k) + CHECK(res_vp(k) == res_lm(k)); + CHECK(res_vp(4) == Approx(500.0).margin(1e-9)); + CHECK(res_vp(12) > 0.0); +} + +TEST_CASE("fast_exp matches std::exp", "[fit]") { + double max_rel = 0.0; + for (int i = 0; i <= 14000; ++i) { + const double x = -700.0 + 1400.0 * static_cast(i) / 14000.0; + const double reference = std::exp(x); + max_rel = + std::max(max_rel, std::abs(aare::model::fast_exp(x) - reference) / + reference); + } + CHECK(max_rel < 1e-15); + CHECK(aare::model::fast_exp(0.0) == 1.0); + CHECK(aare::model::fast_exp(-1000.0) == aare::model::fast_exp(-700.0)); +} diff --git a/src/LevenbergMarquardt.hpp b/src/LevenbergMarquardt.hpp index c40cf9e6..c63d04ac 100644 --- a/src/LevenbergMarquardt.hpp +++ b/src/LevenbergMarquardt.hpp @@ -1,6 +1,7 @@ // SPDX-License-Identifier: MPL-2.0 #pragma once +#include "DampedGaussNewton.hpp" #include "aare/FitModel.hpp" #include "aare/NDView.hpp" @@ -8,7 +9,6 @@ #include #include #include -#include #include namespace aare::detail { @@ -21,29 +21,17 @@ namespace aare::detail { * are ignored). Keep one instance per thread: the buffers are reused between * pixels so that after the first fit no memory is allocated. * - * - Fixed parameters are left out of the linear system. - * - Every evaluation computes the residuals and the Jacobian, so an accepted - * trial point is ready for the next iteration without a second pass over - * the data. Rejected trials are rare with the damping rule below, and a - * rejected one only wastes the derivative part of a single pass. - * - Limits: a trial step that leaves the allowed box is reflected at the - * limit. The step truncated at the limit is evaluated as well and the - * better of the two is kept. A parameter sitting on a limit whose gradient - * points outward is frozen for that iteration. - * - Damping: Marquardt scaling u * diag(J^T J), starting at u = 0.1, with the - * update rule of Nielsen (Madsen, Nielsen and Tingleff, Methods for - * Non-Linear Least Squares Problems, 2004). - * - Convergence follows Minuit's convention: the fit stops when the - * estimated distance to the minimum of chi2 (EDM, computed from the - * Gauss-Newton Hessian) drops below 0.002 * tolerance, or when the step - * becomes negligible. - * - max_calls limits the number of model evaluations: the initial point and - * every trial point count, so an iteration costs one evaluation, or two - * when a limit is crossed. As in Minuit2, 0 selects the default budget - * 200 + 100 * npar + 5 * npar^2. - * - Errors are Gauss-Newton estimates, sqrt(diag((J^T J)^-1)), over the free - * parameters that do not sit on a limit. Fixed and limit-bound parameters - * report 0. + * Every evaluation computes the residuals r = (y - f) / s and the Jacobian + * J = dr/dp in one pass over the data and reduces them to the normal + * equations g = J^T r and J^T J, which DampedGaussNewton iterates on: it + * handles fixed parameters, limits, damping and the convergence test, see + * there. Rejected trials are rare with its damping rule, and a rejected one + * only wastes a single pass. As in Minuit2, max_calls = 0 selects the + * default budget 200 + 100 * npar + 5 * npar^2. + * + * Errors are Gauss-Newton estimates, sqrt(diag((J^T J)^-1)), over the free + * parameters that do not sit on a limit. Fixed and limit-bound parameters + * report 0. */ template class LevenbergMarquardt { public: @@ -51,12 +39,9 @@ template class LevenbergMarquardt { static_assert(Model::npar >= 1 && Model::npar <= 16, "LevenbergMarquardt keeps npar x npar matrices on the stack"); - struct Result { - bool valid = false; - int calls = 0; - int iterations = 0; - double chi2 = 0.0; - }; + using Result = MinimizerResult; + using Driver = DampedGaussNewton; + using Normal = typename Driver::Normal; /** * @brief Fit one pixel. @@ -78,23 +63,17 @@ template class LevenbergMarquardt { s_ = s; weighted_ = s.size() > 0; const auto n = static_cast(x.size()); - nfree_ = 0; + std::array lower{}; + std::array upper{}; + std::array fixed{}; for (int k = 0; k < npar; ++k) { const auto idx = static_cast(k); - lower_[k] = model.lower_limit(idx); - upper_[k] = model.upper_limit(idx); - p_[k] = start[static_cast(k)]; - if (!model.is_user_fixed(idx)) - free_[nfree_++] = k; + lower[k] = model.lower_limit(idx); + upper[k] = model.upper_limit(idx); + fixed[k] = model.is_user_fixed(idx); } - for (int a = 0; a < nfree_; ++a) - q_[a] = p_[free_[a]]; r_.resize(n); - r_new_.resize(n); - r_alt_.resize(n); J_.resize(n * static_cast(npar)); - J_new_.resize(n * static_cast(npar)); - J_alt_.resize(n * static_cast(npar)); if (weighted_) { // 1 / s_i once per pixel rather than once per evaluation. w_.resize(n); @@ -103,139 +82,30 @@ template class LevenbergMarquardt { ? 1.0 / s_[static_cast(i)] : 0.0; } + const int max_calls = model.max_calls() > 0 + ? static_cast(model.max_calls()) + : 200 + 100 * npar + 5 * npar * npar; - Result res; - auto fail = [&]() { + auto eval = [this](const double *p, Normal &out) { + return evaluate(p, out); + }; + typename Driver::Options opt; + opt.tolerance = model.tolerance(); + opt.max_calls = max_calls; + // Start cautiously: the estimated start values of the linear + // parameters can be far off, and the first steps are then wild. + opt.damping0 = 0.1; + const Result res = driver_.fit(eval, start.data(), lower.data(), + upper.data(), fixed.data(), opt); + if (!res.valid) { std::fill(par_out, par_out + npar, 0.0); if (err_out) std::fill(err_out, err_out + npar, 0.0); return res; - }; - - if (!evaluate(q_, r_, J_, F_)) - return fail(); - res.calls = 1; - const double edm_target = 0.002 * model.tolerance(); - const int max_calls = model.max_calls() > 0 - ? static_cast(model.max_calls()) - : 200 + 100 * npar + 5 * npar * npar; - double u = 0.1; // damping relative to diag(J^T J) - double nu = 2.0; - - for (;; ++res.iterations) { - build_normal_equations(); - mark_active(); - if (gmax_ < 1e-12) { - res.valid = true; - break; - } - // Solve for the damped step first. The solve yields - // g^T (J^T J + u D)^-1 g, a lower bound on the EDM, so the - // undamped factorisation behind edm() is only needed when that - // bound leaves room for convergence. - const bool have_step = solve_step(u); - if ((!have_step || bound_ < edm_margin * edm_target) && - edm() < edm_target) { - res.valid = true; - break; - } - if (res.calls >= max_calls) - break; - if (!have_step) { - u *= nu; - nu *= 2.0; - if (u > 1e32) - break; - continue; - } - double step2 = 0.0; - double q2 = 0.0; - bool crossed = false; - for (int a = 0; a < nfree_; ++a) { - const int k = free_[a]; - const double v = q_[a] + delta_[a]; - const double truncated = std::clamp(v, lower_[k], upper_[k]); - double reflected = v; - if (truncated != v) { - crossed = true; - if (v < lower_[k]) - reflected = 2.0 * lower_[k] - v; - if (v > upper_[k]) - reflected = 2.0 * upper_[k] - v; - reflected = std::clamp(reflected, lower_[k], upper_[k]); - } - q_new_[a] = reflected; - q_alt_[a] = truncated; - const double d = q_new_[a] - q_[a]; - step2 += d * d; - q2 += q_[a] * q_[a]; - } - if (std::sqrt(step2) <= 1e-8 * (std::sqrt(q2) + 1e-8)) { - res.valid = true; - break; - } - double F_new = 0.0; - ++res.calls; - bool ok = evaluate(q_new_, r_new_, J_new_, F_new); - if (crossed) { - // Also try the step truncated at the limit, keep the better. - double F_alt = 0.0; - ++res.calls; - if (evaluate(q_alt_, r_alt_, J_alt_, F_alt) && - (!ok || F_alt < F_new)) { - q_new_ = q_alt_; - F_new = F_alt; - std::swap(r_new_, r_alt_); - std::swap(J_new_, J_alt_); - ok = true; - } - } - const double rho = ok ? (F_ - F_new) / predicted_ : -1.0; - if (rho > 0.0) { - // The trial becomes the current point together with its - // residuals and Jacobian. - q_ = q_new_; - F_ = F_new; - std::swap(r_, r_new_); - std::swap(J_, J_new_); - for (int a = 0; a < nfree_; ++a) - p_[free_[a]] = q_[a]; - const double t = 2.0 * rho - 1.0; - u *= std::max(1.0 / 3.0, 1.0 - t * t * t); - nu = 2.0; - } else { - u *= nu; - nu *= 2.0; - if (u > 1e32) - break; - } - } - if (!res.valid) - return fail(); - - // p_ tracks the last evaluated point, which may be a rejected trial. - for (int a = 0; a < nfree_; ++a) - p_[free_[a]] = q_[a]; - for (int k = 0; k < npar; ++k) - par_out[k] = p_[k]; - res.chi2 = 2.0 * F_; - - if (err_out) { - std::fill(err_out, err_out + npar, 0.0); - double A[npar][npar]; - double b[npar]; - double inv[npar][npar]; - const int nact = pack(A, b, 0.0); - if (nact > 0 && invert_spd(nact, A, inv)) { - int i = 0; - for (int a = 0; a < nfree_; ++a) { - if (active_[a]) - continue; - err_out[free_[a]] = std::sqrt(std::max(inv[i][i], 0.0)); - ++i; - } - } } + std::copy(driver_.point(), driver_.point() + npar, par_out); + if (err_out) + driver_.errors(err_out); return res; } @@ -244,38 +114,28 @@ template class LevenbergMarquardt { return weighted_ ? w_[static_cast(i)] : 1.0; } - // Residuals r = (y - f) / s and the Jacobian J = dr/dp for the free - // parameters at the point q. F = chi2 / 2. Returns false when the model - // rejects the parameters. - bool evaluate(const std::array &q, - std::vector &r, std::vector &J, double &F) { - for (int a = 0; a < nfree_; ++a) - p_[free_[a]] = q[a]; - pvec_.assign(p_.begin(), p_.end()); + // Residuals r = (y - f) / s and the Jacobian J = dr/dp at the point p, + // reduced to F = chi2 / 2, g = J^T r and the upper triangle of J^T J. + // The sums run over all parameters with fixed-size loops, so the + // compiler can unroll them and keep the accumulators in registers. + // Returns false when the model rejects the parameters. + bool evaluate(const double *p, Normal &out) { + pvec_.assign(p, p + npar); if (!Model::is_valid(pvec_)) return false; const ssize_t n = x_.size(); - F = 0.0; + double F = 0.0; std::array g{}; double f = 0.0; for (ssize_t i = 0; i < n; ++i) { const double w = weight(i); Model::eval_and_grad(x_[i], pvec_, f, g); - r[i] = (y_[i] - f) * w; - F += 0.5 * r[i] * r[i]; - double *Ji = &J[static_cast(i) * npar]; + r_[i] = (y_[i] - f) * w; + F += 0.5 * r_[i] * r_[i]; + double *Ji = &J_[static_cast(i) * npar]; for (int k = 0; k < npar; ++k) Ji[k] = -g[k] * w; } - return true; - } - - // Gradient g = J^T r and Gauss-Newton Hessian J^T J of chi2 / 2 for the - // free parameters. The sums run over all parameters with fixed-size - // loops, so the compiler can unroll them and keep the accumulators in - // registers; the entries of fixed parameters are simply not copied out. - void build_normal_equations() { - const ssize_t n = x_.size(); double G[npar] = {}; double H[npar][npar] = {}; for (ssize_t i = 0; i < n; ++i) { @@ -287,182 +147,20 @@ template class LevenbergMarquardt { H[a][b] += Ji[a] * Ji[b]; } } - gmax_ = 0.0; - for (int a = 0; a < nfree_; ++a) { - const int ka = free_[a]; - g_[a] = G[ka]; - for (int b = 0; b < nfree_; ++b) { - const int kb = free_[b]; - jtj_[a][b] = ka <= kb ? H[ka][kb] : H[kb][ka]; - } - gmax_ = std::max(gmax_, std::abs(g_[a])); - } - } - - // A parameter sitting on a limit with the descent direction pointing - // outward is frozen for this iteration. The test deliberately looks at - // the last evaluated point p_, which after a rejected step is the trial - // rather than the current point q_: a rejected trial that came off the - // limit unfreezes the parameter, and the reflected step can then leave - // the limit. Testing q_ instead was measured to leave more fits stuck on - // a limit with a higher chi2 (FallingScurve and GaussianChargeSharing - // cubes), although it helped Pol2. - void mark_active() { - for (int a = 0; a < nfree_; ++a) { - const int k = free_[a]; - active_[a] = (p_[k] <= lower_[k] && g_[a] > 0.0) || - (p_[k] >= upper_[k] && g_[a] < 0.0); - } - } - - // Minuit's estimated distance to the minimum of chi2, using the - // Gauss-Newton Hessian: g^T (J^T J)^-1 g with g the gradient of chi2/2. - double edm() { - double A[npar][npar]; - double b[npar]; - double z[npar]; - const int m = pack(A, b, 0.0); - if (m == 0) - return 0.0; - if (!cholesky_solve(m, A, b, z)) - return std::numeric_limits::infinity(); - double e = 0.0; - for (int i = 0; i < m; ++i) - e += b[i] * z[i]; - return e; - } - - // Copy the non-frozen block of J^T J (+ u * diag) and g into A and b. - int pack(double A[npar][npar], double b[npar], double u) const { - double mean_diag = 0.0; - for (int a = 0; a < nfree_; ++a) - mean_diag += jtj_[a][a]; - mean_diag = std::max(mean_diag / std::max(nfree_, 1), 1e-300); - int m = 0; - for (int a = 0; a < nfree_; ++a) { - if (active_[a]) - continue; - int mb = 0; - for (int c = 0; c < nfree_; ++c) { - if (active_[c]) - continue; - A[m][mb++] = jtj_[a][c]; - } - A[m][m] += u * std::max(jtj_[a][a], 1e-10 * mean_diag); - b[m] = g_[a]; - ++m; - } - return m; - } - - // Damped Gauss-Newton step (J^T J + u D) delta = -g for the non-frozen - // parameters, the predicted decrease of F for that step, and - // bound_ = g^T (J^T J + u D)^-1 g, which is at most the EDM because - // u D is positive semi-definite. - bool solve_step(double u) { - double A[npar][npar]; - double b[npar]; - double z[npar]; - const int m = pack(A, b, u); - for (int a = 0; a < nfree_; ++a) - delta_[a] = 0.0; - predicted_ = 0.0; - bound_ = 0.0; - if (m == 0) - return true; - if (!cholesky_solve(m, A, b, z)) - return false; - int i = 0; - for (int a = 0; a < nfree_; ++a) { - if (active_[a]) - continue; - delta_[a] = -z[i]; - bound_ += b[i] * z[i]; - predicted_ += - 0.5 * delta_[a] * - (u * std::max(jtj_[a][a], 1e-300) * delta_[a] - g_[a]); - ++i; - } - return predicted_ > 0.0; - } - - static bool cholesky_solve(int m, const double A[npar][npar], - const double *b, double *z) { - double L[npar][npar]; - for (int i = 0; i < m; ++i) { - for (int j = 0; j <= i; ++j) { - double sum = A[i][j]; - for (int k = 0; k < j; ++k) - sum -= L[i][k] * L[j][k]; - if (i == j) { - if (!(sum > 0.0)) - return false; - L[i][i] = std::sqrt(sum); - } else { - L[i][j] = sum / L[j][j]; - } - } - } - double t[npar]; - for (int i = 0; i < m; ++i) { - double sum = b[i]; - for (int k = 0; k < i; ++k) - sum -= L[i][k] * t[k]; - t[i] = sum / L[i][i]; - } - for (int i = m - 1; i >= 0; --i) { - double sum = t[i]; - for (int k = i + 1; k < m; ++k) - sum -= L[k][i] * z[k]; - z[i] = sum / L[i][i]; - } - return true; - } - - static bool invert_spd(int m, const double A[npar][npar], - double inv[npar][npar]) { - double e[npar]; - double col[npar]; - for (int j = 0; j < m; ++j) { - for (int i = 0; i < m; ++i) - e[i] = (i == j) ? 1.0 : 0.0; - if (!cholesky_solve(m, A, e, col)) - return false; - for (int i = 0; i < m; ++i) - inv[i][j] = col[i]; + out.F = F; + for (int a = 0; a < npar; ++a) { + out.g[a] = G[a]; + for (int b = a; b < npar; ++b) + out.H[a][b] = H[a][b]; } return true; } NDView x_, y_, s_; bool weighted_ = false; - int nfree_ = 0; - int free_[npar] = {}; - bool active_[npar] = {}; - double lower_[npar] = {}; - double upper_[npar] = {}; - std::array p_{}; - std::array q_{}; - std::array q_new_{}; - std::array q_alt_{}; - double g_[npar] = {}; - double delta_[npar] = {}; - double jtj_[npar][npar] = {}; - // The two solves that bound the EDM differ by rounding only when the - // damping is negligible; the margin keeps the stop decision exact. - static constexpr double edm_margin = 1.0 + 1e-6; - double F_ = 0.0; - double gmax_ = 0.0; - double predicted_ = 0.0; - double bound_ = 0.0; - // Residuals and Jacobian of the current point, of the trial point and of - // the candidate truncated at a limit; swapped, never copied. + Driver driver_; std::vector r_; - std::vector r_new_; - std::vector r_alt_; std::vector J_; - std::vector J_new_; - std::vector J_alt_; std::vector w_; // 1 / s_i for weighted fits std::vector pvec_; }; diff --git a/src/VariableProjection.hpp b/src/VariableProjection.hpp new file mode 100644 index 00000000..742a4a88 --- /dev/null +++ b/src/VariableProjection.hpp @@ -0,0 +1,425 @@ +// SPDX-License-Identifier: MPL-2.0 +#pragma once + +#include "DampedGaussNewton.hpp" +#include "aare/FitModel.hpp" +#include "aare/Models.hpp" +#include "aare/NDView.hpp" + +#include +#include +#include +#include +#include + +namespace aare::detail { + +/** + * @brief Variable projection least-squares minimiser for a separable + * FitModel (Golub and Pereyra, 1973). + * + * The model declares the parameters it is linear in, see + * model::is_separable. For a given point of the nonlinear parameters the + * linear ones are solved exactly from their normal equations, and the + * damped Gauss-Newton iteration of DampedGaussNewton runs over the + * nonlinear parameters only, on the projected residual + * r(alpha) = y - Phi(alpha) beta(alpha). Its Jacobian uses Kaufman's + * approximation, the derivatives of the model at fixed beta projected out + * of the span of the basis functions, and the reduced Gauss-Newton Hessian + * is the Schur complement of the full one. Compared with + * LevenbergMarquardt the iteration needs fewer evaluations, cannot be misled + * by poor start values of the linear parameters, and each evaluation + * accumulates far fewer sums. The driver rejects trial points at which the + * model loses its sensitivity to a nonlinear parameter and polishes the + * converged point with one undamped step. + * + * One evaluation is a pass over the data that writes the basis functions + * and their derivatives as columns (the model's vectorisable basis_columns + * when it has one, otherwise basis_and_grad per point), a reduction to the + * normal equations of the linear parameters, the small solve, and a second + * reduction over the columns for the projected residual. + * + * Keep one instance per thread: the buffers are reused between pixels so + * that after the first fit no memory is allocated. + * + * - Fixed nonlinear parameters are left out of the iteration. A fixed + * linear parameter keeps its value and moves to the data side of the + * linear solve. + * - Limits on nonlinear parameters are handled by the driver. A limit on a + * linear parameter cannot be enforced by the linear solve: when the + * solution violates one, the fit reports failure so that the caller can + * fall back to LevenbergMarquardt. + * - Start values of free linear parameters are ignored. + * - max_calls counts evaluations of the projected problem, as in + * LevenbergMarquardt; 0 selects 200 + 100 * npar + 5 * npar^2. + * - Errors are the Gauss-Newton estimates of the full model at the + * solution, sqrt(diag((J^T J)^-1)) over the free parameters that do not + * sit on a limit, the same as LevenbergMarquardt reports. + */ +template class VariableProjection { + public: + static_assert(model::is_separable::value, + "VariableProjection needs a model with linear_par and " + "basis_and_grad"); + using Traits = model::separable_traits; + static constexpr int npar = static_cast(Model::npar); + static constexpr int nlin = static_cast(Traits::nlin); + static constexpr int nnl = static_cast(Traits::nnl); + static_assert(Model::npar >= 1 && Model::npar <= 16, + "VariableProjection keeps npar x npar matrices on the stack"); + static_assert(nlin >= 1, "a separable model has at least one linear " + "parameter"); + + using Result = MinimizerResult; + // The driver iterates over the nonlinear parameters. A fully linear + // model has none and is solved without iterating. + static constexpr int ndrv = nnl > 0 ? nnl : 1; + static constexpr int ncol = nnl * nlin > 0 ? nnl *nlin : 1; + using Driver = DampedGaussNewton; + using Normal = typename Driver::Normal; + + /** + * @brief Fit one pixel, see LevenbergMarquardt::fit for the arguments. + */ + Result fit(const FitModel &model, NDView x, + NDView y, NDView s, + const std::array &start, double *par_out, + double *err_out) { + x_ = x; + y_ = y; + const auto n = static_cast(x.size()); + phi_.resize(n * nlin); + dphi_.resize(n * nlin * nnl); + weighted_ = s.size() > 0; + if (weighted_) { + // 1 / s_i^2 once per pixel; points with s_i == 0 are ignored. + w2_.resize(n); + for (std::size_t i = 0; i < n; ++i) { + const double si = s[static_cast(i)]; + w2_[i] = si != 0.0 ? 1.0 / (si * si) : 0.0; + } + } + pvec_.assign(start.begin(), start.end()); + have_last_ = false; + nfree_lin_ = 0; + for (int j = 0; j < nlin; ++j) { + const auto k = static_cast(Traits::linear_index(j)); + lin_fixed_[j] = model.is_user_fixed(k); + if (!lin_fixed_[j]) + free_lin_[nfree_lin_++] = j; + } + std::array alpha{}; + std::array lower{}; + std::array upper{}; + std::array fixed{}; + for (int k = 0; k < nnl; ++k) { + const auto idx = + static_cast(Traits::nonlinear_par[k]); + alpha[k] = start[idx]; + lower[k] = model.lower_limit(idx); + upper[k] = model.upper_limit(idx); + fixed[k] = model.is_user_fixed(idx); + } + const int max_calls = model.max_calls() > 0 + ? static_cast(model.max_calls()) + : 200 + 100 * npar + 5 * npar * npar; + + Result res; + auto fail = [&]() { + res.valid = false; + res.chi2 = 0.0; + std::fill(par_out, par_out + npar, 0.0); + if (err_out) + std::fill(err_out, err_out + npar, 0.0); + return res; + }; + + if constexpr (nnl == 0) { + Normal nrm; + if (!evaluate(alpha.data(), nrm)) + return fail(); + res.calls = 1; + res.valid = true; + res.chi2 = 2.0 * nrm.F; + for (int j = 0; j < nlin; ++j) + pvec_[Traits::linear_index(j)] = + nrm.aux[static_cast(j)]; + } else { + auto eval = [this](const double *a, Normal &out) { + return evaluate(a, out); + }; + typename Driver::Options opt; + opt.tolerance = model.tolerance(); + opt.max_calls = max_calls; + // The same cautious start as LevenbergMarquardt: a large first + // step can still drive a width to zero. Rejecting degenerate + // trial points guards against that, and the polish step makes + // up for the few iterations, which stop while the damping still + // shortens the steps. + opt.damping0 = 0.1; + opt.reject_degenerate = true; + opt.polish = true; + res = driver_.fit(eval, alpha.data(), lower.data(), upper.data(), + fixed.data(), opt); + if (!res.valid) + return fail(); + for (int k = 0; k < nnl; ++k) + pvec_[Traits::nonlinear_par[k]] = driver_.point()[k]; + for (int j = 0; j < nlin; ++j) + pvec_[Traits::linear_index(j)] = + driver_.normal().aux[static_cast(j)]; + } + + // The linear solve knows nothing about limits. + for (int j = 0; j < nlin; ++j) { + if (lin_fixed_[j]) + continue; + const auto k = static_cast(Traits::linear_index(j)); + if (pvec_[k] < model.lower_limit(k) || + pvec_[k] > model.upper_limit(k)) + return fail(); + } + std::copy(pvec_.begin(), pvec_.end(), par_out); + if (err_out) + errors(model, err_out); + return res; + } + + private: + // F = chi2 / 2, gradient and Kaufman Hessian of the projected residual + // at the nonlinear point alpha, and the linear solution in out.aux. + // Returns false when the model rejects the point or the basis functions + // are linearly dependent there. + bool evaluate(const double *alpha, Normal &out) { + for (int k = 0; k < nnl; ++k) + pvec_[Traits::nonlinear_par[k]] = alpha[k]; + if (!Model::is_valid(pvec_)) + return false; + const ssize_t n = x_.size(); + + // Basis functions and their derivatives, one column per function: + // phi_[j * n + i] and dphi_[(k * nlin + j) * n + i]. + if constexpr (model::has_basis_columns::value) { + Model::basis_columns(x_.data(), n, pvec_, phi_.data(), + dphi_.data()); + } else { + typename Traits::Basis phi{}; + typename Traits::BasisGrad dphi{}; + double *pc = phi_.data(); + double *dc = dphi_.data(); + const auto un = static_cast(n); + for (ssize_t i = 0; i < n; ++i) { + Model::basis_and_grad(x_[i], pvec_, phi, dphi); + const auto ui = static_cast(i); + for (int j = 0; j < nlin; ++j) + pc[static_cast(j) * un + ui] = + phi[static_cast(j)]; + for (int k = 0; k < nnl; ++k) + for (int j = 0; j < nlin; ++j) + dc[static_cast(k * nlin + j) * un + ui] = + dphi[static_cast(k)] + [static_cast(j)]; + } + } + return weighted_ ? reduce(alpha, out) : reduce(alpha, out); + } + + // The reductions over the columns: normal equations and solve of the + // linear parameters, then the projected residual. All loops over the + // linear and nonlinear parameters have compile-time bounds; the free + // linear block is packed out of the full sums afterwards. + template bool reduce(const double *alpha, Normal &out) { + const ssize_t n = x_.size(); + const auto un = static_cast(n); + const double *pc[nlin]; + for (int j = 0; j < nlin; ++j) + pc[j] = phi_.data() + static_cast(j) * un; + const double *dc[ncol]; + for (int c = 0; c < nnl * nlin; ++c) + dc[c] = dphi_.data() + static_cast(c) * un; + const double *yd = y_.data(); + const double *w2 = Weighted ? w2_.data() : nullptr; + + // Pass 1: normal equations M = Phi^T W Phi, t = Phi^T W y of all + // linear parameters. + double M[nlin][nlin] = {}; + double t[nlin] = {}; + for (ssize_t i = 0; i < n; ++i) { + const double wi = Weighted ? w2[i] : 1.0; + const double yi = yd[i]; + double f[nlin]; + for (int j = 0; j < nlin; ++j) + f[j] = pc[j][i]; + for (int a = 0; a < nlin; ++a) { + const double fa = f[a] * wi; + t[a] += fa * yi; + for (int b = a; b < nlin; ++b) + M[a][b] += fa * f[b]; + } + } + + // Free linear parameters: their block of M, with the fixed ones + // moved to the right-hand side, v_a = t_a - sum_fixed c_j M_aj. + double A[nlin][nlin]; + double v[nlin]; + for (int a = 0; a < nfree_lin_; ++a) { + const int ja = free_lin_[a]; + double va = t[ja]; + for (int j = 0; j < nlin; ++j) + if (lin_fixed_[j]) + va -= pvec_[Traits::linear_index(j)] * + (ja <= j ? M[ja][j] : M[j][ja]); + v[a] = va; + for (int b = 0; b < nfree_lin_; ++b) { + const int jb = free_lin_[b]; + A[a][b] = ja <= jb ? M[ja][jb] : M[jb][ja]; + } + } + double L[nlin][nlin]; + if (!cholesky_factor(nfree_lin_, A, L)) + return false; + double beta_free[nlin]; + cholesky_substitute(nfree_lin_, L, v, beta_free); + for (int a = 0; a < nfree_lin_; ++a) + pvec_[Traits::linear_index(free_lin_[a])] = beta_free[a]; + double beta[nlin]; + for (int j = 0; j < nlin; ++j) { + beta[j] = pvec_[Traits::linear_index(j)]; + out.aux[static_cast(j)] = beta[j]; + } + + // Pass 2: residual r = y - Phi beta and u_k = df/dalpha_k at fixed + // beta; F, b = U^T W r, N = U^T W U and Q = Phi^T W U. + double F = 0.0; + double b[ndrv] = {}; + double N[ndrv][ndrv] = {}; + double Q[ndrv][nlin] = {}; + for (ssize_t i = 0; i < n; ++i) { + const double wi = Weighted ? w2[i] : 1.0; + double f[nlin]; + for (int j = 0; j < nlin; ++j) + f[j] = pc[j][i]; + double r = yd[i]; + for (int j = 0; j < nlin; ++j) + r -= beta[j] * f[j]; + double u[ndrv] = {}; + for (int k = 0; k < nnl; ++k) + for (int j = 0; j < nlin; ++j) + u[k] += beta[j] * dc[k * nlin + j][i]; + const double rw = r * wi; + F += 0.5 * r * rw; + for (int k = 0; k < nnl; ++k) { + const double uw = u[k] * wi; + b[k] += uw * r; + for (int l = k; l < nnl; ++l) + N[k][l] += uw * u[l]; + for (int j = 0; j < nlin; ++j) + Q[k][j] += uw * f[j]; + } + } + + // The full Gauss-Newton Hessian of the model at this point is + // [[M, Q^T], [Q, N]]; kept for the errors of the final point. + for (int k = 0; k < nnl; ++k) + last_alpha_[k] = alpha[k]; + for (int a = 0; a < nlin; ++a) + for (int c = a; c < nlin; ++c) + last_M_[a][c] = M[a][c]; + for (int k = 0; k < nnl; ++k) { + for (int j = 0; j < nlin; ++j) + last_Q_[k][j] = Q[k][j]; + for (int l = k; l < nnl; ++l) + last_N_[k][l] = N[k][l]; + } + have_last_ = true; + + // Kaufman's Jacobian: U projected out of the span of the free basis, + // H = N - Q_free^T A^-1 Q_free. + double qf[ndrv][nlin]; + double z[ndrv][nlin]; + for (int k = 0; k < nnl; ++k) { + for (int a = 0; a < nfree_lin_; ++a) + qf[k][a] = Q[k][free_lin_[a]]; + cholesky_substitute(nfree_lin_, L, qf[k], z[k]); + } + out.F = F; + for (int k = 0; k < nnl; ++k) { + out.g[k] = -b[k]; + for (int l = k; l < nnl; ++l) { + double h = N[k][l]; + for (int a = 0; a < nfree_lin_; ++a) + h -= qf[k][a] * z[l][a]; + out.H[k][l] = h; + } + } + return true; + } + + // Gauss-Newton errors from the Jacobian of the full model at the + // solution, the estimate LevenbergMarquardt reports as well. The + // Jacobian's normal matrix is assembled from the sums of the last + // evaluation when that was the final point, which it is unless the last + // trial was rejected; otherwise one pass over the data computes it. + void errors(const FitModel &model, double *err_out) { + double H[npar][npar] = {}; + bool at_solution = have_last_; + for (int k = 0; k < nnl; ++k) + at_solution = at_solution && + pvec_[Traits::nonlinear_par[k]] == last_alpha_[k]; + if (at_solution) { + for (int a = 0; a < nlin; ++a) + for (int c = a; c < nlin; ++c) + H[Traits::linear_index(a)][Traits::linear_index(c)] = + last_M_[a][c]; + for (int k = 0; k < nnl; ++k) { + const auto m = static_cast(Traits::nonlinear_par[k]); + for (int j = 0; j < nlin; ++j) { + const auto i = static_cast(Traits::linear_index(j)); + H[std::min(i, m)][std::max(i, m)] = last_Q_[k][j]; + } + for (int l = k; l < nnl; ++l) + H[m][Traits::nonlinear_par[l]] = last_N_[k][l]; + } + } else { + const ssize_t n = x_.size(); + std::array g{}; + double f = 0.0; + for (ssize_t i = 0; i < n; ++i) { + Model::eval_and_grad(x_[i], pvec_, f, g); + const double wi = + weighted_ ? w2_[static_cast(i)] : 1.0; + for (int a = 0; a < npar; ++a) + for (int c = a; c < npar; ++c) + H[a][c] += g[a] * g[c] * wi; + } + } + bool skip[npar]; + for (int k = 0; k < npar; ++k) + skip[k] = model.is_user_fixed(static_cast(k)); + if constexpr (nnl > 0) { + for (int k = 0; k < nnl; ++k) + skip[Traits::nonlinear_par[k]] = + skip[Traits::nonlinear_par[k]] || driver_.skipped(k); + } + gauss_newton_errors(H, skip, err_out); + } + + NDView x_, y_; + Driver driver_; + bool weighted_ = false; + // Sums of the last evaluation, see errors(). + bool have_last_ = false; + double last_alpha_[ndrv] = {}; + double last_M_[nlin][nlin] = {}; + double last_Q_[ndrv][nlin] = {}; + double last_N_[ndrv][ndrv] = {}; + int nfree_lin_ = 0; + int free_lin_[nlin] = {}; + bool lin_fixed_[nlin] = {}; + std::vector phi_; // nlin columns of n basis values + std::vector dphi_; // nnl * nlin columns of their derivatives + std::vector w2_; // 1 / s_i^2 for weighted fits + std::vector pvec_; +}; + +} // namespace aare::detail