mirror of
https://github.com/slsdetectorgroup/aare.git
synced 2026-10-31 02:11:28 +01:00
Erik FröjdhandClaude Fable 5.1
a211a3e8e2
Every bundled model now declares the parameters it is linear in (linear_par, basis_and_grad) and, except the polynomials, a vectorisable bulk basis pass (basis_columns) built on the new model::fast_exp. The Levenberg-Marquardt iteration is factored into a shared DampedGaussNewton driver with unchanged results; VarPro solves the linear parameters exactly at every trial point and iterates over the nonlinear ones, rejecting degenerate trial points and polishing the converged point with one undamped step. 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 [fit] tests run VarPro through every configuration and check the separable structure of each model against its derivatives. Measured on an M1 Max, single thread: FallingScurve population 5.5 us per fit against 11.7 us with LevenbergMarquardt, and 2 wrong minima in 4096 pixels against 30; Gaussian 3.3 vs 4.3 us; GaussianChargeSharingKb 33 vs 49 us; Pol2 0.34 vs 1.52 us. Migrad stays the default. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
426 lines
18 KiB
C++
426 lines
18 KiB
C++
// SPDX-License-Identifier: MPL-2.0
|
|
#pragma once
|
|
|
|
#include "DampedGaussNewton.hpp"
|
|
#include "aare/FitModel.hpp"
|
|
#include "aare/Models.hpp"
|
|
#include "aare/NDView.hpp"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cmath>
|
|
#include <cstddef>
|
|
#include <vector>
|
|
|
|
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 <typename Model> class VariableProjection {
|
|
public:
|
|
static_assert(model::is_separable<Model>::value,
|
|
"VariableProjection needs a model with linear_par and "
|
|
"basis_and_grad");
|
|
using Traits = model::separable_traits<Model>;
|
|
static constexpr int npar = static_cast<int>(Model::npar);
|
|
static constexpr int nlin = static_cast<int>(Traits::nlin);
|
|
static constexpr int nnl = static_cast<int>(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<ndrv>;
|
|
using Normal = typename Driver::Normal;
|
|
|
|
/**
|
|
* @brief Fit one pixel, see LevenbergMarquardt::fit for the arguments.
|
|
*/
|
|
Result fit(const FitModel<Model> &model, NDView<double, 1> x,
|
|
NDView<double, 1> y, NDView<double, 1> s,
|
|
const std::array<double, Model::npar> &start, double *par_out,
|
|
double *err_out) {
|
|
x_ = x;
|
|
y_ = y;
|
|
const auto n = static_cast<std::size_t>(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<ssize_t>(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<unsigned int>(Traits::linear_index(j));
|
|
lin_fixed_[j] = model.is_user_fixed(k);
|
|
if (!lin_fixed_[j])
|
|
free_lin_[nfree_lin_++] = j;
|
|
}
|
|
std::array<double, ndrv> alpha{};
|
|
std::array<double, ndrv> lower{};
|
|
std::array<double, ndrv> upper{};
|
|
std::array<bool, ndrv> fixed{};
|
|
for (int k = 0; k < nnl; ++k) {
|
|
const auto idx =
|
|
static_cast<unsigned int>(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<int>(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<std::size_t>(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<std::size_t>(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<unsigned int>(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<Model>::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<std::size_t>(n);
|
|
for (ssize_t i = 0; i < n; ++i) {
|
|
Model::basis_and_grad(x_[i], pvec_, phi, dphi);
|
|
const auto ui = static_cast<std::size_t>(i);
|
|
for (int j = 0; j < nlin; ++j)
|
|
pc[static_cast<std::size_t>(j) * un + ui] =
|
|
phi[static_cast<std::size_t>(j)];
|
|
for (int k = 0; k < nnl; ++k)
|
|
for (int j = 0; j < nlin; ++j)
|
|
dc[static_cast<std::size_t>(k * nlin + j) * un + ui] =
|
|
dphi[static_cast<std::size_t>(k)]
|
|
[static_cast<std::size_t>(j)];
|
|
}
|
|
}
|
|
return weighted_ ? reduce<true>(alpha, out) : reduce<false>(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 Weighted> bool reduce(const double *alpha, Normal &out) {
|
|
const ssize_t n = x_.size();
|
|
const auto un = static_cast<std::size_t>(n);
|
|
const double *pc[nlin];
|
|
for (int j = 0; j < nlin; ++j)
|
|
pc[j] = phi_.data() + static_cast<std::size_t>(j) * un;
|
|
const double *dc[ncol];
|
|
for (int c = 0; c < nnl * nlin; ++c)
|
|
dc[c] = dphi_.data() + static_cast<std::size_t>(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<nlin>(nfree_lin_, A, L))
|
|
return false;
|
|
double beta_free[nlin];
|
|
cholesky_substitute<nlin>(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<std::size_t>(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<nlin>(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> &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<int>(Traits::nonlinear_par[k]);
|
|
for (int j = 0; j < nlin; ++j) {
|
|
const auto i = static_cast<int>(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<double, Model::npar> 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<std::size_t>(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<unsigned int>(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<npar>(H, skip, err_out);
|
|
}
|
|
|
|
NDView<double, 1> 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<double> phi_; // nlin columns of n basis values
|
|
std::vector<double> dphi_; // nnl * nlin columns of their derivatives
|
|
std::vector<double> w2_; // 1 / s_i^2 for weighted fits
|
|
std::vector<double> pvec_;
|
|
};
|
|
|
|
} // namespace aare::detail
|