Files
Jungfraujoch/image_analysis/geom_refinement/LMSolver.cpp
T
leonarski_fandClaude Opus 5.5 a4bb6f74f1 XtalOptimizer: own Levenberg-Marquardt solver in place of Ceres
The crystal refinement (XtalOptimizer, both the seven-block and the reduced
beam+orientation form, and XtalOptimizerRotationOnly) no longer builds a
ceres::Problem. XtalRefine holds the problem as data and solves it with
LMSolver, which follows Ceres' trust-region LM step for step - Jacobi scaling,
damping and radius updates, stopping rules, box projection, the projected
Armijo line search with cubic interpolation on bounded problems, the
SphereManifold for the spindle - but takes J^T J and J^T r directly instead of
a Jacobian. The residual is the same XtalResidual code, now Ceres-free and
evaluated on a forward-mode Dual (Dual.h); everything that depends on
parameters alone (detector-angle trig, per-frame back-rotation, reciprocal
basis, orientation rotation) is worked out once per evaluation, and the
observed and predicted halves carry 6 and 9 derivative lanes rather than 16.
The sums are cut into blocks that depend on the residual count alone, so the
answer does not depend on the thread count. Because the line-search trial point
is the candidate point, a bounded iteration costs one evaluation instead of
Ceres' three.

Validation (rc174 + this, -march=x86-64-v3):
- p.mtz md5 identical to the Ceres build on myob/cytc/thau x10sa, GPU and CPU
  builds, and on the lyso8 stills reference.
- Solve corpus (every 16-parameter solve and every 10th per-image solve of the
  three sets, 8.3k problems, inputs and Ceres results dumped from a run that
  reproduced the md5s): usable/failed agree on all, iteration counts identical
  on all, parameters agree to <2e-11 (in px / rad / 0.01 A units), costs to
  1e-13.
- Same process, same threads: 7-9x faster per solve than Ceres.
- In-run (GPU, loaded box): xtal 16-parameter solves myob 22.2 -> 6.4 core-s,
  cytc 88 -> 26 core-s; per-image solves 5.4 -> 1.2 core-s (myob); cytc first
  pass indexing windows 1.9 -> 0.85 s, myob 1.1 -> 0.45 s; solver share of the
  whole cytc run 17% -> 4% of CPU samples.

New tests compare the solver with Ceres on synthetic rotation problems
(full/weighted/reduced) and check thread-count independence.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01K5K8jvPPbmCrbqnWkddTuB
2026-10-03 09:46:26 +02:00

242 lines
9.1 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
// Adapted from https://github.com/ceres-solver/ceres-solver (internal/ceres/polynomial.cc,
// include/ceres/internal/sphere_manifold_functions.h, householder_vector.h)
// Copyright 2023 Google Inc. All rights reserved.
// BSD-3-Clause, see licenses/ceres-solver.txt
#include "LMSolver.h"
#include <Eigen/Eigenvalues>
namespace {
void HouseholderVector3(const double x[3], double v[3], double &beta) {
const double sigma = x[0] * x[0] + x[1] * x[1];
v[0] = x[0];
v[1] = x[1];
v[2] = 1.0;
beta = 0.0;
const double x_pivot = x[2];
if (sigma <= std::numeric_limits<double>::epsilon()) {
if (x_pivot < 0.0)
beta = 2.0;
return;
}
const double mu = std::sqrt(x_pivot * x_pivot + sigma);
const double v_pivot = (x_pivot <= 0.0) ? x_pivot - mu : -sigma / (x_pivot + mu);
beta = 2.0 * v_pivot * v_pivot / (sigma + v_pivot * v_pivot);
v[0] /= v_pivot;
v[1] /= v_pivot;
}
double Norm3(const double x[3]) {
return std::sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]);
}
using Vector = Eigen::VectorXd;
using Matrix = Eigen::MatrixXd;
double EvaluatePolynomial(const Vector &polynomial, double x) {
double v = 0.0;
for (int i = 0; i < polynomial.size(); ++i)
v = v * x + polynomial(i);
return v;
}
void BalanceCompanionMatrix(Matrix &companion_matrix) {
Matrix offdiagonal = companion_matrix;
offdiagonal.diagonal().setZero();
const int degree = static_cast<int>(companion_matrix.rows());
const double gamma = 0.9;
bool scaling_has_changed;
do {
scaling_has_changed = false;
for (int i = 0; i < degree; ++i) {
const double col_norm = offdiagonal.col(i).lpNorm<1>();
if (std::fpclassify(col_norm) != FP_ZERO) {
const double row_norm = offdiagonal.row(i).lpNorm<1>();
int exponent = 0;
std::frexp(row_norm / col_norm, &exponent);
exponent /= 2;
if (exponent != 0) {
const double scaled_col_norm = std::ldexp(col_norm, exponent);
const double scaled_row_norm = std::ldexp(row_norm, -exponent);
if (scaled_col_norm + scaled_row_norm < gamma * (col_norm + row_norm)) {
scaling_has_changed = true;
offdiagonal.row(i) *= std::ldexp(1.0, -exponent);
offdiagonal.col(i) *= std::ldexp(1.0, exponent);
}
}
}
}
} while (scaling_has_changed);
offdiagonal.diagonal() = companion_matrix.diagonal();
companion_matrix = offdiagonal;
}
// Real parts of the roots, as Ceres' FindPolynomialRoots (the imaginary parts are not used here).
bool FindPolynomialRoots(const Vector &polynomial_in, Vector &real) {
if (polynomial_in.size() == 0)
return false;
int lead = 0;
while (lead < polynomial_in.size() - 1 && polynomial_in(lead) == 0.0)
++lead;
Vector polynomial = polynomial_in.tail(polynomial_in.size() - lead);
const int degree = static_cast<int>(polynomial.size()) - 1;
if (degree == 0) {
real.resize(0);
return true;
}
if (degree == 1) {
real.resize(1);
real(0) = -polynomial(1) / polynomial(0);
return true;
}
if (degree == 2) {
const double a = polynomial(0);
const double b = polynomial(1);
const double c = polynomial(2);
const double D = b * b - 4 * a * c;
const double sqrt_D = std::sqrt(std::fabs(D));
real.setZero(2);
if (D >= 0) {
if (b >= 0) {
real(0) = (-b - sqrt_D) / (2.0 * a);
real(1) = (2.0 * c) / (-b - sqrt_D);
} else {
real(0) = (2.0 * c) / (-b + sqrt_D);
real(1) = (-b + sqrt_D) / (2.0 * a);
}
} else {
real(0) = -b / (2.0 * a);
real(1) = -b / (2.0 * a);
}
return true;
}
polynomial /= polynomial(0);
Matrix companion = Matrix::Zero(degree, degree);
companion.diagonal(-1).setOnes();
companion.col(degree - 1) = -polynomial.reverse().head(degree);
BalanceCompanionMatrix(companion);
Eigen::EigenSolver<Matrix> solver(companion, false);
if (solver.info() != Eigen::Success)
return false;
real = solver.eigenvalues().real();
return true;
}
void MinimizePolynomial(const Vector &polynomial, double x_min, double x_max,
double &optimal_x, double &optimal_value) {
optimal_x = (x_min + x_max) / 2.0;
optimal_value = EvaluatePolynomial(polynomial, optimal_x);
const double x_min_value = EvaluatePolynomial(polynomial, x_min);
if (x_min_value < optimal_value) {
optimal_value = x_min_value;
optimal_x = x_min;
}
const double x_max_value = EvaluatePolynomial(polynomial, x_max);
if (x_max_value < optimal_value) {
optimal_value = x_max_value;
optimal_x = x_max;
}
if (polynomial.rows() <= 2)
return;
const int degree = static_cast<int>(polynomial.rows()) - 1;
Vector derivative(degree);
for (int i = 0; i < degree; ++i)
derivative(i) = (degree - i) * polynomial(i);
Vector roots_real;
if (!FindPolynomialRoots(derivative, roots_real))
return;
for (int i = 0; i < roots_real.rows(); ++i) {
const double root = roots_real(i);
if (root < x_min || root > x_max)
continue;
const double value = EvaluatePolynomial(polynomial, root);
if (value < optimal_value) {
optimal_value = value;
optimal_x = root;
}
}
}
Vector FindInterpolatingPolynomial(const std::vector<LMLineSample> &samples) {
int num_constraints = 0;
for (const auto &s: samples)
num_constraints += (s.value_is_valid ? 1 : 0) + (s.gradient_is_valid ? 1 : 0);
const int degree = num_constraints - 1;
Matrix lhs = Matrix::Zero(num_constraints, num_constraints);
Vector rhs = Vector::Zero(num_constraints);
int row = 0;
for (const auto &s: samples) {
if (s.value_is_valid) {
for (int j = 0; j <= degree; ++j)
lhs(row, j) = std::pow(s.x, degree - j);
rhs(row) = s.value;
++row;
}
if (s.gradient_is_valid) {
for (int j = 0; j < degree; ++j)
lhs(row, j) = (degree - j) * std::pow(s.x, degree - j - 1);
rhs(row) = s.gradient;
++row;
}
}
Eigen::FullPivLU<Matrix> lu(lhs);
return lu.setThreshold(0.0).solve(rhs);
}
}
void SpherePlus3(const double x[3], const double delta[2], double out[3]) {
const double norm_delta = std::sqrt(delta[0] * delta[0] + delta[1] * delta[1]);
if (norm_delta == 0.0) {
out[0] = x[0];
out[1] = x[1];
out[2] = x[2];
return;
}
double v[3], beta;
HouseholderVector3(x, v, beta);
const double sin_delta_by_delta = std::sin(norm_delta) / norm_delta;
const double y[3] = {sin_delta_by_delta * delta[0], sin_delta_by_delta * delta[1], std::cos(norm_delta)};
const double vy = v[0] * y[0] + v[1] * y[1] + v[2] * y[2];
const double x_norm = Norm3(x);
for (int i = 0; i < 3; i++)
out[i] = x_norm * (y[i] - v[i] * (beta * vy));
}
void SpherePlusJacobian3(const double x[3], double jacobian[3][2]) {
double v[3], beta;
HouseholderVector3(x, v, beta);
const double x_norm = Norm3(x);
for (int i = 0; i < 2; ++i)
for (int r = 0; r < 3; ++r)
jacobian[r][i] = (-beta * v[i] * v[r] + (r == i ? 1.0 : 0.0)) * x_norm;
}
double LMInterpolatedStepSize(const LMLineSample &lowerbound, const LMLineSample &previous,
const LMLineSample &current, double min_step, double max_step) {
if (!current.value_is_valid)
return std::min(std::max(current.x * 0.5, min_step), max_step);
std::vector<LMLineSample> samples{lowerbound, current};
if (previous.value_is_valid)
samples.push_back(previous);
const Vector polynomial = FindInterpolatingPolynomial(samples);
double step = 0.0, value = 0.0;
MinimizePolynomial(polynomial, min_step, max_step, step, value);
for (const auto &s: samples) {
if (s.x < min_step || s.x > max_step)
continue;
const double v = EvaluatePolynomial(polynomial, s.x);
if (v < value) {
step = s.x;
value = v;
}
}
return step;
}