From 5beabf18888bd91a4102bdccf5cf80698ed2a574 Mon Sep 17 00:00:00 2001 From: Erik Frojdh Date: Fri, 25 Sep 2026 18:05:43 +0200 Subject: [PATCH] optimizations --- RELEASE.md | 11 +++ benchmarks/fit_benchmark.cpp | 154 ++++++++++++++++++++++++++++++++++- docs/src/pyFit.rst | 7 +- include/aare/FitModel.hpp | 5 +- include/aare/Models.hpp | 51 ++++++------ python/src/fit.hpp | 20 ++++- python/tests/test_Fit.py | 9 +- src/LevenbergMarquardt.hpp | 143 +++++++++++++++++++++----------- 8 files changed, 320 insertions(+), 80 deletions(-) diff --git a/RELEASE.md b/RELEASE.md index eb8eb1fc..a2523ee3 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -41,6 +41,15 @@ from their linearised covariance instead of running Hesse. Pixels for which Fumili does not reach a valid minimum, which Minuit2's implementation cannot once a two-sided limit becomes active, are refitted with Migrad. +- ``Minimizer.LevenbergMarquardt`` is about twice as fast for every model: it + evaluates the Jacobian at the trial point instead of re-evaluating accepted + points, accumulates the normal equations with fixed-size loops, needs a + single factorisation per iteration, and the erf-based models compute one + exponential per point instead of two. Results are unchanged, except that + ``GaussianChargeSharing`` and ``GaussianChargeSharingKb`` now evaluate the + 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. - ``NDArray``/``NDView`` expressions now support scalar operands (``2 * a + b / 4``) and can be assigned to an existing ``NDArray``, reusing its buffer. @@ -106,6 +115,8 @@ - ``TimingMode::Auto`` changed to ``TimingMode::AUTO_TIMING``, ``TimingMode::Trigger`` changed to ``TimingMode::TRIGGER_EXPOSURE`` ### Bugfixes: +- Fitting a 3D data cube in Python without ``y_err`` now returns ``par_err`` + when ``compute_errors`` is set, as fitting a single pixel already did. - Mismatched operators inhibited vectorization in gcc of NDArray math operators - ``NDArray``/``NDView`` math operators and expressions were not vectorized for ``uint8_t`` and 64 bit integers, up to 30x slower than a plain loop. diff --git a/benchmarks/fit_benchmark.cpp b/benchmarks/fit_benchmark.cpp index df3ea412..02bef6ec 100644 --- a/benchmarks/fit_benchmark.cpp +++ b/benchmarks/fit_benchmark.cpp @@ -138,17 +138,49 @@ static void BM_FitGausLMErrors(benchmark::State &state) { run_gaussian_fit(state, aare::Minimizer::LevenbergMarquardt, 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, + aare::Minimizer minimizer) { + const auto &tc = get_test_cases()[1]; + auto data = generate_gaussian_data(tc); + auto xv = data.x.view(); + auto yv = data.y.view(); + + aare::FitModel model(0, 500, 0.5, false, minimizer); + model.FixParameter(2, tc.true_sig); + + aare::NDArray result; + for (auto _ : state) { + result = aare::fit_pixel(model, xv, yv); + benchmark::DoNotOptimize(result.data()); + } + report_accuracy(state, tc, result); +} + +static void BM_FitGausFixedSigmaMigrad(benchmark::State &state) { + run_gaussian_fixed_sigma_fit(state, aare::Minimizer::Migrad); +} +static void BM_FitGausFixedSigmaFumili(benchmark::State &state) { + run_gaussian_fixed_sigma_fit(state, aare::Minimizer::Fumili); +} +static void BM_FitGausFixedSigmaLM(benchmark::State &state) { + run_gaussian_fixed_sigma_fit(state, aare::Minimizer::LevenbergMarquardt); +} + // ---------------------------------------------------------------- // Rising S-curve (6 parameters), typical for threshold scans // ---------------------------------------------------------------- static constexpr ssize_t N_SCURVE_POINTS = 100; -static void run_scurve_fit(benchmark::State &state, aare::Minimizer minimizer) { +static void run_scurve_fit(benchmark::State &state, aare::Minimizer minimizer, + bool weighted = false) { const std::vector truth = {5.0, 0.1, 50.0, 4.0, 200.0, 0.5}; const double noise_sigma = 0.02 * truth[4]; aare::NDArray x({N_SCURVE_POINTS}); aare::NDArray y({N_SCURVE_POINTS}); + aare::NDArray y_err({N_SCURVE_POINTS}, noise_sigma); std::mt19937 rng(SEED); std::normal_distribution noise(0.0, noise_sigma); for (ssize_t i = 0; i < N_SCURVE_POINTS; ++i) { @@ -157,13 +189,17 @@ static void run_scurve_fit(benchmark::State &state, aare::Minimizer minimizer) { } auto xv = x.view(); auto yv = y.view(); + auto ev = y_err.view(); const aare::FitModel model(0, 500, 0.5, false, minimizer); aare::NDArray result; for (auto _ : state) { - result = aare::fit_pixel(model, xv, yv); + result = + weighted + ? aare::fit_pixel(model, xv, yv, ev) + : aare::fit_pixel(model, xv, yv); benchmark::DoNotOptimize(result.data()); } @@ -185,6 +221,104 @@ static void BM_FitScurveLM(benchmark::State &state) { run_scurve_fit(state, aare::Minimizer::LevenbergMarquardt); } +// 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); +} +static void BM_FitScurveWeightedFumili(benchmark::State &state) { + run_scurve_fit(state, aare::Minimizer::Fumili, true); +} +static void BM_FitScurveWeightedLM(benchmark::State &state) { + run_scurve_fit(state, aare::Minimizer::LevenbergMarquardt, true); +} + +// ---------------------------------------------------------------- +// Pol2 (3 parameters, cheap model): exposes the solver overhead per +// iteration rather than the model evaluation. +// ---------------------------------------------------------------- +static void run_pol2_fit(benchmark::State &state, aare::Minimizer minimizer) { + constexpr ssize_t n = 51; + const std::vector truth = {12.0, -0.8, 0.05}; + + aare::NDArray x({n}); + aare::NDArray y({n}); + std::mt19937 rng(SEED); + std::normal_distribution noise(0.0, 2.0); + for (ssize_t i = 0; i < n; ++i) { + x[i] = static_cast(i); + y[i] = aare::model::Pol2::eval(x[i], truth) + noise(rng); + } + auto xv = x.view(); + auto yv = y.view(); + + const aare::FitModel model(0, 500, 0.5, false, + minimizer); + + aare::NDArray result; + for (auto _ : state) { + result = aare::fit_pixel(model, xv, yv); + benchmark::DoNotOptimize(result.data()); + } + state.counters["dp2"] = result(2) - truth[2]; +} + +static void BM_FitPol2Migrad(benchmark::State &state) { + run_pol2_fit(state, aare::Minimizer::Migrad); +} +static void BM_FitPol2Fumili(benchmark::State &state) { + run_pol2_fit(state, aare::Minimizer::Fumili); +} +static void BM_FitPol2LM(benchmark::State &state) { + run_pol2_fit(state, aare::Minimizer::LevenbergMarquardt); +} + +// ---------------------------------------------------------------- +// GaussianChargeSharingKb (8 parameters, two exponentials and two erfs per +// point): the most expensive model per evaluation. +// ---------------------------------------------------------------- +static void run_charge_sharing_kb_fit(benchmark::State &state, + aare::Minimizer minimizer) { + constexpr ssize_t n = 200; + const std::vector truth = {20.0, 0.05, 110.0, 6.0, + 1500.0, 0.1, 1.1, 0.1}; + + aare::NDArray x({n}); + aare::NDArray y({n}); + std::mt19937 rng(SEED); + std::normal_distribution noise(0.0, 1.0); + for (ssize_t i = 0; i < n; ++i) { + x[i] = static_cast(i); + const double clean = + aare::model::GaussianChargeSharingKb::eval(x[i], truth); + y[i] = clean + std::sqrt(std::max(clean, 1.0)) * noise(rng); + } + auto xv = x.view(); + auto yv = y.view(); + + const aare::FitModel model( + 0, 500, 0.5, false, minimizer); + + aare::NDArray result; + for (auto _ : state) { + result = aare::fit_pixel(model, + xv, yv); + benchmark::DoNotOptimize(result.data()); + } + state.counters["dMu"] = result(2) - truth[2]; + state.counters["dN"] = result(4) - truth[4]; + state.counters["chi2"] = result(8); +} + +static void BM_FitChargeSharingKbMigrad(benchmark::State &state) { + run_charge_sharing_kb_fit(state, aare::Minimizer::Migrad); +} +static void BM_FitChargeSharingKbFumili(benchmark::State &state) { + run_charge_sharing_kb_fit(state, aare::Minimizer::Fumili); +} +static void BM_FitChargeSharingKbLM(benchmark::State &state) { + run_charge_sharing_kb_fit(state, aare::Minimizer::LevenbergMarquardt); +} + // ---------------------------------------------------------------- // Data cube through fit_3d on one thread: exposes the per-pixel overhead // of each minimizer on top of the minimisation itself. @@ -272,10 +406,26 @@ BENCHMARK(BM_FitGausFumiliErrors) ->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitGausLMErrors)->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_FitScurveMigrad)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitScurveFumili)->Unit(benchmark::kMicrosecond); BENCHMARK(BM_FitScurveLM)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitScurveWeightedMigrad)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitScurveWeightedFumili)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitScurveWeightedLM)->Unit(benchmark::kMicrosecond); + +BENCHMARK(BM_FitPol2Migrad)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitPol2Fumili)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitPol2LM)->Unit(benchmark::kMicrosecond); + +BENCHMARK(BM_FitChargeSharingKbMigrad)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitChargeSharingKbFumili)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_FitChargeSharingKbLM)->Unit(benchmark::kMicrosecond); + // fit_3d runs in worker threads, so measure wall time. BENCHMARK(BM_FitCubeMigrad)->Unit(benchmark::kMillisecond)->UseRealTime(); BENCHMARK(BM_FitCubeFumili)->Unit(benchmark::kMillisecond)->UseRealTime(); diff --git a/docs/src/pyFit.rst b/docs/src/pyFit.rst index 36c59033..92755398 100644 --- a/docs/src/pyFit.rst +++ b/docs/src/pyFit.rst @@ -68,9 +68,10 @@ 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. ``0`` selects Minuit's - default budget. A pixel that does not converge within the budget returns - zeros. + 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``. diff --git a/include/aare/FitModel.hpp b/include/aare/FitModel.hpp index ca77985e..b5acbc4d 100644 --- a/include/aare/FitModel.hpp +++ b/include/aare/FitModel.hpp @@ -32,8 +32,9 @@ namespace aare { * 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). `strategy` only - * affects the Minuit2 minimizers. + * 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. */ enum class Minimizer { Migrad, Fumili, LevenbergMarquardt }; diff --git a/include/aare/Models.hpp b/include/aare/Models.hpp index 1c673e99..092937e8 100644 --- a/include/aare/Models.hpp +++ b/include/aare/Models.hpp @@ -13,10 +13,14 @@ namespace aare::model { inline constexpr double inv_sqrt2 = 0.70710678118654752440; inline constexpr double inv_sqrt_2pi = 0.39894228040143267794; -inline double fast_erf(double x) { - // Abramowitz–Stegun Handbook of Mathematical Functions - // erf approximation with max error ~1.5e-7, faster than std::erf. - +/** + * @brief erf(z) given exp(-z^2), from the Abramowitz-Stegun approximation + * 7.1.26 (Handbook of Mathematical Functions, max error about 1.5e-7). + * + * Models whose derivatives need exp(-z^2) anyway pass it in, so that a + * single exponential serves both the value and the gradient. + */ +inline double fast_erf_from_exp(double z, double exp_minus_z2) { const double a1 = 0.254829592; const double a2 = -0.284496736; const double a3 = 1.421413741; @@ -24,16 +28,17 @@ inline double fast_erf(double x) { const double a5 = 1.061405429; const double p = 0.3275911; - const int sign = x < 0 ? -1 : 1; - x = std::abs(x); - - // 7.1.26 - const double t = 1.0 / (1.0 + p * x); + const double t = 1.0 / (1.0 + p * std::abs(z)); const double y = 1.0 - - (((((a5 * t + a4) * t + a3) * t + a2) * t + a1) * t) * std::exp(-x * x); + (((((a5 * t + a4) * t + a3) * t + a2) * t + a1) * t) * exp_minus_z2; - return sign * y; + return z < 0 ? -y : y; +} + +/** @brief erf approximation, faster than std::erf; see fast_erf_from_exp. */ +inline double fast_erf(double x) { + return fast_erf_from_exp(x, std::exp(-x * x)); } /** @@ -347,7 +352,7 @@ struct GaussianErfcPlateau { const double z = dx * inv_sqrt2 / sig; const double e = std::exp(-z * z); - const double step = 0.5 * (1.0 - fast_erf(z)); + const double step = 0.5 * (1.0 - fast_erf_from_exp(z, e)); return A * e + S * step; } @@ -363,7 +368,7 @@ struct GaussianErfcPlateau { const double z = dx * inv_sqrt2 / sig; const double e = std::exp(-z * z); - const double step = 0.5 * (1.0 - fast_erf(z)); + const double step = 0.5 * (1.0 - fast_erf_from_exp(z, e)); f = A * e + S * step; @@ -500,7 +505,7 @@ struct GaussianChargeSharing { const double u = dx / sig; const double G = std::exp(-0.5 * u * u); - const double H = 0.5 * (1.0 - fast_erf(u * inv_sqrt2)); + const double H = 0.5 * (1.0 - fast_erf_from_exp(u * inv_sqrt2, G)); return p0 - p1 * x + N * (G + C * H); } @@ -518,7 +523,7 @@ struct GaussianChargeSharing { const double u = dx / sig; const double G = std::exp(-0.5 * u * u); - const double H = 0.5 * (1.0 - fast_erf(u * inv_sqrt2)); + const double H = 0.5 * (1.0 - fast_erf_from_exp(u * inv_sqrt2, G)); f = p0 - p1 * x + N * (G + C * H); @@ -662,8 +667,8 @@ struct GaussianChargeSharingKb { 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(u * inv_sqrt2)); - const double Hb = 0.5 * (1.0 - fast_erf(ub * inv_sqrt2)); + 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; @@ -693,8 +698,8 @@ struct GaussianChargeSharingKb { 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(u * inv_sqrt2)); - const double Hb = 0.5 * (1.0 - fast_erf(ub * inv_sqrt2)); + 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; @@ -826,12 +831,12 @@ struct RisingScurve { const double dx = x - p2; const double z = dx * inv_sqrt2 / p3; - const double step = 0.5 * (1.0 + fast_erf(z)); + const double e = std::exp(-z * z); + const double step = 0.5 * (1.0 + fast_erf_from_exp(z, e)); const double amp = p4 + p5 * dx; f = (p0 + p1 * x) + step * amp; - const double e = std::exp(-z * z); const double dSdp2 = -inv_sqrt_2pi * e / p3; const double dSdp3 = -inv_sqrt_2pi * e * dx / (p3 * p3); @@ -967,12 +972,12 @@ struct FallingScurve { const double dx = x - p2; const double z = dx * inv_sqrt2 / p3; - const double step = 0.5 * (1.0 - fast_erf(z)); + const double e = std::exp(-z * z); + const double step = 0.5 * (1.0 - fast_erf_from_exp(z, e)); const double amp = p4 + p5 * dx; f = (p0 + p1 * x) + step * amp; - const double e = std::exp(-z * z); const double dSdp2 = +inv_sqrt_2pi * e / p3; // sign flipped vs rising const double dSdp3 = +inv_sqrt_2pi * e * dx / (p3 * p3); diff --git a/python/src/fit.hpp b/python/src/fit.hpp index 5c28af40..c0ab4914 100644 --- a/python/src/fit.hpp +++ b/python/src/fit.hpp @@ -37,7 +37,8 @@ template void bind_fit_model(py::module &m, const char *name) { by ``Minimizer.LevenbergMarquardt``. max_calls : int Work budget per pixel: Minuit2 function calls, or model - evaluations for ``Minimizer.LevenbergMarquardt``. 0 selects + evaluations for ``Minimizer.LevenbergMarquardt`` (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 @@ -217,8 +218,23 @@ fit_dispatch(const aare::FitModel &model, "chi2"_a = return_image_data(chi2_out)); } } else { - + // Unweighted fit; parameter errors are still available when + // requested, as in the 1D path. NDView dummy_err{}; + + if (model.compute_errors()) { + auto err_out = + new NDArray({y.shape(0), y.shape(1), npar}, 0.0); + + aare::fit_3d(model, x_view, y_view, dummy_err, + par_out->view(), err_out->view(), + chi2_out->view(), n_threads); + + return py::dict("par"_a = return_image_data(par_out), + "par_err"_a = return_image_data(err_out), + "chi2"_a = return_image_data(chi2_out)); + } + NDView dummy_err_out{}; aare::fit_3d(model, x_view, y_view, dummy_err, diff --git a/python/tests/test_Fit.py b/python/tests/test_Fit.py index 3a89c654..77df6f81 100644 --- a/python/tests/test_Fit.py +++ b/python/tests/test_Fit.py @@ -179,20 +179,25 @@ def test_scurve_fit(cls, minimizer): @pytest.mark.parametrize("with_errors", [False, True]) -def test_fit_3d_matches_fit_1d(with_errors, minimizer): +@pytest.mark.parametrize("weighted", [False, True]) +def test_fit_3d_matches_fit_1d(with_errors, weighted, minimizer): x = np.linspace(-6.0, 6.0, 61) rows, cols = 4, 3 y = np.empty((rows, cols, x.size)) for r in range(rows): for c in range(cols): y[r, c] = gaussian(x, 80.0 + 10.0 * r, -0.6 + 0.3 * c, 0.9 + 0.1 * r) - y_err = np.ones_like(y) if with_errors else None + y_err = np.ones_like(y) if weighted else None model = aare.Gaussian(compute_errors=with_errors, minimizer=minimizer) cube = model.fit(x, y, y_err, n_threads=3) + expected_keys = {"par", "par_err", "chi2"} if with_errors else {"par", "chi2"} + assert set(cube) == expected_keys assert cube["par"].shape == (rows, cols, 3) assert cube["chi2"].shape == (rows, cols) + if with_errors: + assert cube["par_err"].shape == (rows, cols, 3) for r in range(rows): for c in range(cols): single = model.fit(x, y[r, c], None if y_err is None else y_err[r, c]) diff --git a/src/LevenbergMarquardt.hpp b/src/LevenbergMarquardt.hpp index 1c88c396..c40cf9e6 100644 --- a/src/LevenbergMarquardt.hpp +++ b/src/LevenbergMarquardt.hpp @@ -22,6 +22,10 @@ namespace aare::detail { * 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 @@ -33,9 +37,10 @@ namespace aare::detail { * 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, - * every trial point and every accepted point count. As in Minuit2, 0 - * selects the default budget 200 + 100 * npar + 5 * npar^2. + * - 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. @@ -86,7 +91,18 @@ template class LevenbergMarquardt { 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); + for (std::size_t i = 0; i < n; ++i) + w_[i] = s_[static_cast(i)] != 0.0 + ? 1.0 / s_[static_cast(i)] + : 0.0; + } Result res; auto fail = [&]() { @@ -96,7 +112,7 @@ template class LevenbergMarquardt { return res; }; - if (!evaluate(q_, r_, true, F_)) + if (!evaluate(q_, r_, J_, F_)) return fail(); res.calls = 1; const double edm_target = 0.002 * model.tolerance(); @@ -109,13 +125,23 @@ template class LevenbergMarquardt { for (;; ++res.iterations) { build_normal_equations(); mark_active(); - if (edm() < edm_target || gmax_ < 1e-12) { + 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 (!solve_step(u)) { + if (!have_step) { u *= nu; nu *= 2.0; if (u > 1e32) @@ -150,24 +176,30 @@ template class LevenbergMarquardt { } double F_new = 0.0; ++res.calls; - bool ok = evaluate(q_new_, r_new_, false, F_new); + 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_new_, false, F_alt) && + 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_; - ++res.calls; - if (!evaluate(q_, r_, true, F_)) - break; // same point as the accepted trial, cannot happen + 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; @@ -209,14 +241,14 @@ template class LevenbergMarquardt { private: double weight(ssize_t i) const { - return weighted_ ? (s_[i] != 0.0 ? 1.0 / s_[i] : 0.0) : 1.0; + return weighted_ ? w_[static_cast(i)] : 1.0; } - // Residuals r = (y - f) / s and, optionally, the Jacobian dr/dp for the - // free parameters at the point q. F = chi2 / 2. Returns false when the - // model rejects the parameters. + // 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, bool with_jacobian, double &F) { + 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()); @@ -224,52 +256,57 @@ template class LevenbergMarquardt { return false; const ssize_t n = x_.size(); F = 0.0; - if (with_jacobian) { - 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]; - for (int a = 0; a < nfree_; ++a) - Ji[a] = -g[free_[a]] * w; - } - } else { - for (ssize_t i = 0; i < n; ++i) { - r[i] = (y_[i] - Model::eval(x_[i], pvec_)) * weight(i); - F += 0.5 * r[i] * r[i]; - } + 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]; + 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(); - for (int a = 0; a < nfree_; ++a) { - g_[a] = 0.0; - for (int b = 0; b < nfree_; ++b) - jtj_[a][b] = 0.0; - } + double G[npar] = {}; + double H[npar][npar] = {}; for (ssize_t i = 0; i < n; ++i) { const double *Ji = &J_[static_cast(i) * npar]; - for (int a = 0; a < nfree_; ++a) { - g_[a] += Ji[a] * r_[i]; - for (int b = a; b < nfree_; ++b) - jtj_[a][b] += Ji[a] * Ji[b]; + const double ri = r_[i]; + for (int a = 0; a < npar; ++a) { + G[a] += Ji[a] * ri; + for (int b = a; b < npar; ++b) + H[a][b] += Ji[a] * Ji[b]; } } gmax_ = 0.0; for (int a = 0; a < nfree_; ++a) { - for (int b = 0; b < a; ++b) - jtj_[a][b] = jtj_[b][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. + // 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]; @@ -319,7 +356,9 @@ template class LevenbergMarquardt { } // Damped Gauss-Newton step (J^T J + u D) delta = -g for the non-frozen - // parameters; also the predicted decrease of F for that step. + // 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]; @@ -328,6 +367,7 @@ template class LevenbergMarquardt { 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)) @@ -337,6 +377,7 @@ template class LevenbergMarquardt { 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]); @@ -407,12 +448,22 @@ template class LevenbergMarquardt { 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. 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_; };