diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 8dcb538ac..3122fff96 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -50,7 +50,7 @@ either way. Eigen is header-only: only its headers reach the binaries, and no Ei ## Vendored directly in the repository -These live in the source tree (see the path) rather than being fetched; traccc is the exception - code adapted into first-party files rather than a vendored directory, see the note at the end of this file. +These live in the source tree (see the path) rather than being fetched; traccc and the Ceres-derived minimiser are the exceptions - code adapted into first-party files rather than a vendored directory, see the notes at the end of this file. | Component | Path | Copyright | License (SPDX) | License text | |---|---|---|---|---| @@ -69,6 +69,7 @@ These live in the source tree (see the path) rather than being fetched; traccc i | [pocketfft](https://github.com/mreineck/pocketfft) | `gemmi_gph/gemmi/third_party/pocketfft_hdronly.h` | Max-Planck-Society; Peter Bell; MIT (FFTW-derived parts) | BSD-3-Clause | [pocketfft.txt](licenses/pocketfft.txt) | | [tinydir](https://github.com/cxong/tinydir) | `gemmi_gph/gemmi/third_party/tinydir.h` | Cong Xu, Lautis Sun, Baudouin Feildel, Andargor | BSD-2-Clause | [tinydir.txt](licenses/tinydir.txt) | | [traccc (ACTS)](https://github.com/acts-project/traccc) | `image_analysis/spot_finding/StrongPixelSet.cpp`, `SpotExtractorGPU.cu` | CERN, for the benefit of the ACTS project | MPL-2.0 | [traccc.txt](licenses/traccc.txt) | +| [Ceres Solver](https://github.com/ceres-solver/ceres-solver) (adapted) | `image_analysis/geom_refinement/LMSolver.h`, `LMSolver.cpp` | Google Inc. | BSD-3-Clause | [ceres-solver.txt](licenses/ceres-solver.txt) | | [xbflash.qspi](https://github.com/Xilinx/XRT) | `tools/xbflash.qspi/` | Xilinx / AMD | Apache-2.0 | [xbflash-qspi.txt](licenses/xbflash-qspi.txt) | | [wingetopt](https://github.com/alex85k/wingetopt) | `tools/wingetopt/` | Todd C. Miller; The NetBSD Foundation | ISC AND BSD-2-Clause | [wingetopt.txt](licenses/wingetopt.txt) | @@ -107,6 +108,10 @@ served frontend, so the shipped web UI carries its own attribution. and `SpotExtractorGPU.cu` follows the design of its GPU counterpart. MPL-2.0 is file-level, so both files name the origin at the top and are covered by `licenses/traccc.txt`. See [ACKNOWLEDGEMENT.md](docs/ACKNOWLEDGEMENT.md) for the citation. +* **Ceres Solver** is also fetched and linked (table above); separately, `LMSolver.h`/`.cpp` re-implement + its trust-region Levenberg-Marquardt minimiser, projected line search, polynomial step choice and + sphere manifold for the crystal refinement, following its source. Both files name the origin at the + top and are covered by `licenses/ceres-solver.txt`. * **FFTW** is GPL-2.0-or-later — compatible with, and absorbed by, this project's GPL-3.0 license. * **Apache-2.0** components: where upstream ships a `NOTICE` file, it is reproduced in the corresponding `licenses/` text. diff --git a/image_analysis/geom_refinement/CMakeLists.txt b/image_analysis/geom_refinement/CMakeLists.txt index b004e8d4e..2330dbb93 100644 --- a/image_analysis/geom_refinement/CMakeLists.txt +++ b/image_analysis/geom_refinement/CMakeLists.txt @@ -24,6 +24,11 @@ ADD_LIBRARY(JFJochGeomRefinement STATIC XtalOptimizer.cpp XtalOptimizer.h XtalResidual.h + XtalRefine.cpp + XtalRefine.h + LMSolver.cpp + LMSolver.h + Dual.h PostRefine.cpp PostRefine.h GeometryRefiner.cpp diff --git a/image_analysis/geom_refinement/Dual.h b/image_analysis/geom_refinement/Dual.h new file mode 100644 index 000000000..932828058 --- /dev/null +++ b/image_analysis/geom_refinement/Dual.h @@ -0,0 +1,142 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// A forward-mode dual number with N derivative lanes: a value and its gradient with respect to N +// parameters. The residuals of the crystal refinement are written as templates over their scalar type, +// so the same code runs on a plain double and on this. The value part of every operation is the plain +// double arithmetic of the same expression - written the way ceres::Jet writes it, division through the +// reciprocal - so a residual evaluated on a Dual has the same value as on a Jet. + +#include +#include + +#include + +template +struct Dual { + double a = 0.0; + double v[N] = {}; + + Dual() = default; + Dual(double value) : a(value) {} // NOLINT: implicit, a constant is a dual with zero derivatives + + static Dual Variable(double value, int lane) { + Dual d(value); + d.v[lane] = 1.0; + return d; + } + + Dual &operator+=(const Dual &o) { a += o.a; for (int i = 0; i < N; i++) v[i] += o.v[i]; return *this; } + Dual &operator-=(const Dual &o) { a -= o.a; for (int i = 0; i < N; i++) v[i] -= o.v[i]; return *this; } + Dual &operator*=(const Dual &o) { *this = *this * o; return *this; } + Dual &operator/=(const Dual &o) { *this = *this / o; return *this; } + + friend Dual operator+(const Dual &x) { return x; } + friend Dual operator-(const Dual &x) { + Dual r(-x.a); + for (int i = 0; i < N; i++) r.v[i] = -x.v[i]; + return r; + } + + friend Dual operator+(const Dual &x, const Dual &y) { + Dual r(x.a + y.a); + for (int i = 0; i < N; i++) r.v[i] = x.v[i] + y.v[i]; + return r; + } + friend Dual operator+(const Dual &x, double s) { Dual r = x; r.a += s; return r; } + friend Dual operator+(double s, const Dual &x) { Dual r = x; r.a += s; return r; } + + friend Dual operator-(const Dual &x, const Dual &y) { + Dual r(x.a - y.a); + for (int i = 0; i < N; i++) r.v[i] = x.v[i] - y.v[i]; + return r; + } + friend Dual operator-(const Dual &x, double s) { Dual r = x; r.a -= s; return r; } + friend Dual operator-(double s, const Dual &x) { + Dual r(s - x.a); + for (int i = 0; i < N; i++) r.v[i] = -x.v[i]; + return r; + } + + friend Dual operator*(const Dual &x, const Dual &y) { + Dual r(x.a * y.a); + for (int i = 0; i < N; i++) r.v[i] = x.a * y.v[i] + x.v[i] * y.a; + return r; + } + friend Dual operator*(const Dual &x, double s) { + Dual r(x.a * s); + for (int i = 0; i < N; i++) r.v[i] = x.v[i] * s; + return r; + } + friend Dual operator*(double s, const Dual &x) { return x * s; } + + friend Dual operator/(const Dual &x, const Dual &y) { + const double y_inv = 1.0 / y.a; + const double q = x.a * y_inv; + Dual r(q); + for (int i = 0; i < N; i++) r.v[i] = (x.v[i] - q * y.v[i]) * y_inv; + return r; + } + friend Dual operator/(const Dual &x, double s) { + const double s_inv = 1.0 / s; + return x * s_inv; + } + friend Dual operator/(double s, const Dual &y) { + const double y_inv = 1.0 / y.a; + const double d = -s * y_inv * y_inv; + Dual r(s * y_inv); + for (int i = 0; i < N; i++) r.v[i] = d * y.v[i]; + return r; + } + + friend bool operator<(const Dual &x, const Dual &y) { return x.a < y.a; } + friend bool operator>(const Dual &x, const Dual &y) { return x.a > y.a; } + friend bool operator<=(const Dual &x, const Dual &y) { return x.a <= y.a; } + friend bool operator>=(const Dual &x, const Dual &y) { return x.a >= y.a; } + friend bool operator==(const Dual &x, const Dual &y) { return x.a == y.a; } + friend bool operator!=(const Dual &x, const Dual &y) { return x.a != y.a; } + + // The chain rule for a function of one argument: value f, derivative df. + Dual Chain(double f, double df) const { + Dual r(f); + for (int i = 0; i < N; i++) r.v[i] = df * v[i]; + return r; + } + + friend Dual sqrt(const Dual &x) { + const double s = std::sqrt(x.a); + return x.Chain(s, 0.5 / s); + } + friend Dual cos(const Dual &x) { return x.Chain(std::cos(x.a), -std::sin(x.a)); } + friend Dual sin(const Dual &x) { return x.Chain(std::sin(x.a), std::cos(x.a)); } + friend Dual hypot(const Dual &x, const Dual &y, const Dual &z) { + // As ceres::hypot(Jet, Jet, Jet): the value is std::hypot, the derivative x/h dx + y/h dy + z/h dz. + const double h = std::hypot(x.a, y.a, z.a); + Dual r(h); + for (int i = 0; i < N; i++) r.v[i] = x.a / h * x.v[i] + y.a / h * y.v[i] + z.a / h * z.v[i]; + return r; + } + friend int fpclassify(const Dual &x) { return std::fpclassify(x.a); } +}; + +// What Eigen needs to hold a Dual in a fixed-size matrix (the reciprocal basis is built in one). +namespace Eigen { + template + struct NumTraits> : GenericNumTraits { + typedef Dual Real; + typedef Dual NonInteger; + typedef Dual Nested; + typedef Dual Literal; + enum { + IsComplex = 0, IsInteger = 0, IsSigned = 1, RequireInitialization = 1, + ReadCost = 1, AddCost = 1, MulCost = 1 + }; + static inline Real epsilon() { return Real(std::numeric_limits::epsilon()); } + static inline Real dummy_precision() { return Real(1e-12); } + static inline Real highest() { return Real(std::numeric_limits::max()); } + static inline Real lowest() { return Real(-std::numeric_limits::max()); } + static inline int digits10() { return NumTraits::digits10(); } + }; +} diff --git a/image_analysis/geom_refinement/LMSolver.cpp b/image_analysis/geom_refinement/LMSolver.cpp new file mode 100644 index 000000000..967ae965e --- /dev/null +++ b/image_analysis/geom_refinement/LMSolver.cpp @@ -0,0 +1,241 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// 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 + +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::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(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(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 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(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 &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 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 ¤t, 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 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; +} diff --git a/image_analysis/geom_refinement/LMSolver.h b/image_analysis/geom_refinement/LMSolver.h new file mode 100644 index 000000000..ee3b9657d --- /dev/null +++ b/image_analysis/geom_refinement/LMSolver.h @@ -0,0 +1,365 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +// The minimiser below follows Ceres Solver's trust-region Levenberg-Marquardt step for step - its +// options, its Jacobi scaling, its damping and radius updates, its stopping rules, its box projection, +// its projected Armijo line search on bounded problems and its SphereManifold - so that a problem moved +// off Ceres takes the same path to the same answer. Adapted from +// https://github.com/ceres-solver/ceres-solver (internal/ceres/trust_region_minimizer.cc, +// levenberg_marquardt_strategy.cc, trust_region_step_evaluator.cc, line_search.cc, 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 +// +// What it does NOT take from Ceres is the Jacobian: the caller hands over J^T J and J^T r directly, +// accumulated however it likes, and the minimiser never sees a row of J. Everything Ceres computes from +// the Jacobian - the column norms, the normal equations, the model cost change - is a function of those +// two alone. Only the options the crystal refinements use are reproduced: monotonic steps, no inner +// iterations, a dense Cholesky of the normal equations. + +#pragma once + +#include +#include +#include +#include + +#include + +struct LMBlock { + int offset = 0; // into the ambient parameter vector + int size = 0; // ambient size, at most 3 + bool constant = false; + bool sphere = false; // Ceres' SphereManifold: the norm is kept, the tangent has size - 1 coordinates + double lower[3] = {-std::numeric_limits::max(), -std::numeric_limits::max(), + -std::numeric_limits::max()}; + double upper[3] = {std::numeric_limits::max(), std::numeric_limits::max(), + std::numeric_limits::max()}; + + int TangentSize() const { return constant ? 0 : (sphere ? size - 1 : size); } +}; + +struct LMOptions { + int max_iterations = 50; + double max_time_s = 1e9; +}; + +enum class LMTermination { Convergence, NoConvergence, Failure }; + +struct LMSummary { + LMTermination termination = LMTermination::Failure; + // Iterations as Ceres counts them in Summary::iterations: iteration 0 included. + int iterations = 0; + int evaluations = 0; + int line_search_steps = 0; + double initial_cost = 0.0; + double final_cost = 0.0; + + bool IsSolutionUsable() const { return termination != LMTermination::Failure; } +}; + +// Ceres' SphereManifold<3> for a three-vector: the Householder reflection that takes x to the pole, the +// Plus that walks a tangent step along the sphere, and the 3x2 Jacobian of that Plus at zero step. +void SpherePlus3(const double x[3], const double delta[2], double out[3]); +void SpherePlusJacobian3(const double x[3], double jacobian[3][2]); + +// The polynomial step-size choice of Ceres' Armijo line search, cubic interpolation. +struct LMLineSample { + double x = 0.0; + double value = 0.0; + double gradient = 0.0; + bool value_is_valid = false; + bool gradient_is_valid = false; +}; +double LMInterpolatedStepSize(const LMLineSample &lowerbound, const LMLineSample &previous, + const LMLineSample ¤t, double min_step, double max_step); + +// Evaluate is called as eval(x, cost, g, H): x the ambient parameters, cost 1/2 sum of squared +// residuals, and - where g and H are not null - the gradient J^T r and J^T J in the TANGENT coordinates +// of the non-constant blocks, in block order. It returns false where anything came out non-finite. +// x is updated in place on success; it is left untouched on failure. +template +LMSummary SolveLM(std::vector &x_io, const std::vector &blocks, const LMOptions &options, + Evaluate &&eval) { + using Vec = Eigen::VectorXd; + using Mat = Eigen::MatrixXd; + constexpr double kMax = std::numeric_limits::max(); + + // Ceres' defaults, which is what the callers always ran with. + constexpr double initial_radius = 1e4; + constexpr double max_radius = 1e16; + constexpr double min_radius = 1e-32; + constexpr double min_relative_decrease = 1e-3; + constexpr double min_lm_diagonal = 1e-6; + constexpr double max_lm_diagonal = 1e32; + constexpr int max_consecutive_invalid_steps = 5; + constexpr double function_tolerance = 1e-6; + constexpr double gradient_tolerance = 1e-10; + constexpr double parameter_tolerance = 1e-8; + constexpr double sufficient_decrease = 1e-4; + constexpr double max_step_contraction = 1e-3; + constexpr double min_step_contraction = 0.6; + constexpr double min_line_search_step = 1e-9; + constexpr int max_line_search_iterations = 20; + + const auto start = std::chrono::steady_clock::now(); + LMSummary summary; + + int n = 0; + bool constrained = false; + for (const auto &b: blocks) { + for (int j = 0; j < b.size; j++) + if (!std::isfinite(x_io[b.offset + j])) + return summary; + n += b.TangentSize(); + for (int j = 0; j < b.size; j++) { + if (b.constant) { + if (x_io[b.offset + j] < b.lower[j] || x_io[b.offset + j] > b.upper[j]) + return summary; + } else { + if (b.lower[j] >= b.upper[j]) + return summary; + if (b.lower[j] > -kMax || b.upper[j] < kMax) + constrained = true; + } + } + } + + // x (+) delta, block by block, projected onto the bounds - Ceres' ParameterBlock::Plus. + const auto plus = [&](const Vec &x, const Vec &delta, Vec &out) { + out = x; + int t = 0; + for (const auto &b: blocks) { + if (b.constant) + continue; + if (b.sphere) { + SpherePlus3(x.data() + b.offset, delta.data() + t, out.data() + b.offset); + } else { + for (int j = 0; j < b.size; j++) + out[b.offset + j] = x[b.offset + j] + delta[t + j]; + } + for (int j = 0; j < b.size; j++) { + out[b.offset + j] = std::max(out[b.offset + j], b.lower[j]); + out[b.offset + j] = std::min(out[b.offset + j], b.upper[j]); + } + t += b.TangentSize(); + } + }; + // Norms over the parameters Ceres keeps in its state: the non-constant blocks only. + const auto free_norm = [&](const Vec &v) { + double s = 0.0; + for (const auto &b: blocks) + if (!b.constant) + for (int j = 0; j < b.size; j++) + s += v[b.offset + j] * v[b.offset + j]; + return std::sqrt(s); + }; + const auto free_max_norm = [&](const Vec &v) { + double m = 0.0; + for (const auto &b: blocks) + if (!b.constant) + for (int j = 0; j < b.size; j++) + m = std::max(m, std::fabs(v[b.offset + j])); + return m; + }; + + Vec x = Eigen::Map(x_io.data(), static_cast(x_io.size())); + if (constrained) { + Vec projected; + plus(x, Vec::Zero(n), projected); + x = projected; + } + + // One evaluation point with everything Ceres computes there. + struct Point { + Vec x; + double cost = kMax; + Vec g; + Mat H; + bool valid = false; + }; + const auto evaluate = [&](const Vec &at, Point &p) { + p.x = at; + p.g.setZero(n); + p.H.setZero(n, n); + summary.evaluations++; + p.valid = eval(at.data(), p.cost, &p.g, &p.H) && std::isfinite(p.cost); + if (!p.valid) + p.cost = kMax; + }; + + Point cur; + evaluate(x, cur); + if (!cur.valid) + return summary; + summary.initial_cost = cur.cost; + + // Jacobi scaling, fixed from the Jacobian at the starting point. + Vec scale(n); + for (int i = 0; i < n; i++) + scale[i] = 1.0 / (1.0 + std::sqrt(cur.H(i, i))); + + Vec gs, neg_g, projected; + Mat Hs; + double gradient_max_norm = 0.0; + const auto take_point = [&]() { + gs = scale.cwiseProduct(cur.g); + Hs = scale.asDiagonal() * cur.H * scale.asDiagonal(); + neg_g = -cur.g; + plus(cur.x, neg_g, projected); + gradient_max_norm = free_max_norm(cur.x - projected); + }; + take_point(); + + double radius = initial_radius; + double decrease_factor = 2.0; + bool reuse_diagonal = false; + Vec diagonal(n); + int consecutive_invalid = 0; + bool any_successful_step = false; + bool step_successful = true; // iteration 0 + int iteration = 0; + + Point trial; // the last point the line search evaluated, reused as the candidate when it is one + + const auto step_rejected = [&]() { + radius = radius / decrease_factor; + decrease_factor *= 2.0; + reuse_diagonal = true; + }; + + const auto finish = [&](LMTermination t) { + summary.termination = t; + summary.final_cost = cur.cost; + if (t != LMTermination::Failure) + for (int i = 0; i < x.size(); i++) + x_io[i] = cur.x[i]; + return summary; + }; + + for (;;) { + // FinalizeIterationAndCheckIfMinimizerCanContinue + summary.iterations++; + if (std::chrono::duration(std::chrono::steady_clock::now() - start).count() + >= options.max_time_s) + return finish(LMTermination::NoConvergence); + if (iteration >= options.max_iterations) + return finish(LMTermination::NoConvergence); + if (step_successful && gradient_max_norm <= gradient_tolerance) + return finish(LMTermination::Convergence); + if (radius <= min_radius) + return finish(LMTermination::Convergence); + + iteration++; + step_successful = false; + + // ComputeTrustRegionStep: the damped normal equations of the scaled Jacobian. + if (!reuse_diagonal) + for (int i = 0; i < n; i++) + diagonal[i] = std::min(std::max(Hs(i, i), min_lm_diagonal), max_lm_diagonal); + Mat lhs = Hs; + for (int i = 0; i < n; i++) { + const double d = std::sqrt(diagonal[i] / radius); + lhs(i, i) += d * d; + } + reuse_diagonal = true; + Eigen::LLT llt(lhs); + bool step_valid = false; + Vec step; + double model_cost_change = 0.0; + if (llt.info() == Eigen::Success) { + step = -llt.solve(gs); + if (step.allFinite()) { + model_cost_change = -(step.dot(gs) + 0.5 * step.dot(Hs * step)); + step_valid = model_cost_change > 0.0; + } + } + if (!step_valid) { + if (++consecutive_invalid >= max_consecutive_invalid_steps) + return finish(LMTermination::Failure); + step_rejected(); + continue; + } + consecutive_invalid = 0; + Vec delta = step.cwiseProduct(scale); + + bool have_trial = false; + if (constrained) { + // Projected Armijo line search along delta, cubic interpolation. + const double initial_gradient = cur.g.dot(delta); + const double direction_max_norm = delta.lpNorm(); + LMLineSample initial{0.0, cur.cost, initial_gradient, true, true}; + LMLineSample previous, current; + const auto line_eval = [&](double alpha, LMLineSample &s) { + s = LMLineSample{}; + s.x = alpha; + Vec moved; + plus(cur.x, Vec(alpha * delta), moved); + evaluate(moved, trial); + have_trial = true; + if (!trial.valid) + return; + s.value = trial.cost; + s.value_is_valid = true; + s.gradient = delta.dot(trial.g); + s.gradient_is_valid = std::isfinite(s.gradient); + }; + line_eval(1.0, current); + bool success = true; + int ls_iterations = 0; + while (!current.value_is_valid + || current.value > initial.value + sufficient_decrease * initial_gradient * current.x) { + ++ls_iterations; + if (ls_iterations >= max_line_search_iterations) { + success = false; + break; + } + const double alpha = LMInterpolatedStepSize(initial, previous, current, + max_step_contraction * current.x, + min_step_contraction * current.x); + if (alpha * direction_max_norm < min_line_search_step) { + success = false; + break; + } + previous = current; + line_eval(alpha, current); + } + summary.line_search_steps += ls_iterations; + if (success) + delta *= current.x; + } + + // ComputeCandidatePointAndEvaluateCost + Vec candidate_x; + plus(cur.x, delta, candidate_x); + Point cand; + if (have_trial && trial.x == candidate_x) + cand = std::move(trial); + else + evaluate(candidate_x, cand); + + if (any_successful_step) { + const double step_norm = free_norm(cur.x - cand.x); + if (step_norm <= parameter_tolerance * (free_norm(cur.x) + parameter_tolerance)) + return finish(LMTermination::Convergence); + } + if (std::fabs(cur.cost - cand.cost) <= function_tolerance * cur.cost) + return finish(LMTermination::Convergence); + + const double relative_decrease = (cand.cost >= kMax) + ? std::numeric_limits::lowest() + : (cur.cost - cand.cost) / model_cost_change; + if (relative_decrease > min_relative_decrease) { + any_successful_step = true; + step_successful = true; + cur = std::move(cand); + take_point(); + radius = radius / std::max(1.0 / 3.0, 1.0 - std::pow(2.0 * relative_decrease - 1.0, 3)); + radius = std::min(max_radius, radius); + decrease_factor = 2.0; + reuse_diagonal = false; + } else { + step_rejected(); + } + } +} diff --git a/image_analysis/geom_refinement/XtalOptimizer.cpp b/image_analysis/geom_refinement/XtalOptimizer.cpp index 58b049e10..23e671a2f 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.cpp +++ b/image_analysis/geom_refinement/XtalOptimizer.cpp @@ -7,84 +7,11 @@ #include "XtalOptimizer.h" #include "XtalResidual.h" -#include "ceres/ceres.h" +#include "XtalRefine.h" #include "ceres/rotation.h" +#include "Dual.h" #include "LatticeReduction.h" -// Soft header prior on ONE beam-centre component (the spindle-parallel, gauge-weak one). Residual = w*(b - b0); -// the caller sets w so the prior behaves like a sigma-pixel restraint that competes with the (unit-weight) -// positional residuals - strong enough to pin the gauge direction, negligible in the well-constrained one. -// Soft restraint on one direction of a two-component block: g.(p - p0), weighted. Used for the beam -// centre and for the detector tilt, which are the same gauge seen twice (see the gauge block below), -// so they take the same direction g and cannot disagree about it. -struct GaugeDirectionPrior { - GaugeDirectionPrior(double gx, double gy, double p0, double weight) - : gx(gx), gy(gy), p0(p0), weight(weight) {} - template - bool operator()(const T *const p, T *residual) const { - residual[0] = T(weight) * (T(gx) * p[0] + T(gy) * p[1] - T(p0)); - return true; - } - double gx, gy, p0, weight; -}; - -struct XtalResidualRotationOnlyPrecomp { - XtalResidualRotationOnlyPrecomp(const Coord &recip_obs, - const CrystalLattice &latt, - double h, double k, double l) - : s_obs(recip_obs), - astar(latt.Astar()), bstar(latt.Bstar()), cstar(latt.Cstar()), - h(h), k(k), l(l) { - } - - template - bool operator()(const T *const rot_aa, T *residual) const { - const T astar_unrot[3] = {T(astar.x), T(astar.y), T(astar.z)}; - const T bstar_unrot[3] = {T(bstar.x), T(bstar.y), T(bstar.z)}; - const T cstar_unrot[3] = {T(cstar.x), T(cstar.y), T(cstar.z)}; - - T astar_rot[3], bstar_rot[3], cstar_rot[3]; - - const AngleAxisRotator rot(rot_aa); - rot.Rotate(astar_unrot, astar_rot); - rot.Rotate(bstar_unrot, bstar_rot); - rot.Rotate(cstar_unrot, cstar_rot); - - const Eigen::Matrix s_pred(T(h) * astar_rot[0] + T(k) * bstar_rot[0] + T(l) * cstar_rot[0], - T(h) * astar_rot[1] + T(k) * bstar_rot[1] + T(l) * cstar_rot[1], - T(h) * astar_rot[2] + T(k) * bstar_rot[2] + T(l) * cstar_rot[2] - ); - - // Residual in reciprocal space - residual[0] = T(s_obs.x) - s_pred[0]; - residual[1] = T(s_obs.y) - s_pred[1]; - residual[2] = T(s_obs.z) - s_pred[2]; - return true; - } - - const Coord s_obs; - const Coord astar, bstar, cstar; - const double h, k, l; -}; - -// Regularizer: penalises ||rot_aa|| to prefer the smallest rotation that -// explains the data. Weight should be chosen in the same units as the -// reciprocal-space residuals (Å⁻¹ per radian). A value of ~0.01–0.1 is -// typically enough to break degeneracy without biasing the solution. -struct RotationNormRegularizer { - explicit RotationNormRegularizer(double weight) : weight(weight) {} - - template - bool operator()(const T *const rot_aa, T *residual) const { - residual[0] = T(weight) * rot_aa[0]; - residual[1] = T(weight) * rot_aa[1]; - residual[2] = T(weight) * rot_aa[2]; - return true; - } - - const double weight; -}; - // Prior confidence weight per spot: how strong the spot is FOR ITS RESOLUTION. The frame's spots are // ordered by resolution and cut into equal-count shells, and each intensity is divided by its shell // median. Refinement needs the high-resolution spots (they carry the cell and distance information) and @@ -153,12 +80,11 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, const int num_threads) { try { // A coplanar basis has no reciprocal cell: 1/V is infinite, every predicted reciprocal vector - // comes out NaN, and Ceres fails on the very first evaluation - after dumping the offending - // block to stderr. There is nothing for the refinement to recover here, so refuse the lattice - // before the problem is built rather than let the solver discover it. The check has to be on - // the vectors: this close to flat, float cell angles no longer carry even the SIGN of the - // metric determinant, and the triclinic branch of XtalResidual then clamps c into the a-b - // plane and divides by the zero volume that makes. + // comes out NaN, and the solver fails on the very first evaluation. There is nothing for the + // refinement to recover here, so refuse the lattice before the problem is built rather than let + // the solver discover it. The check has to be on the vectors: this close to flat, float cell + // angles no longer carry even the SIGN of the metric determinant, and the triclinic branch of + // XtalResidual then clamps c into the a-b plane and divides by the zero volume that makes. if (data.latt.VolumeFraction() < MIN_BASIS_VOLUME_FRACTION) return false; @@ -168,24 +94,21 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, double beta = data.latt.GetUnitCell().beta; // Initial guess for the parameters - double beam[2] = {data.geom.GetBeamX_pxl(), data.geom.GetBeamY_pxl()}; - double distance_mm = data.geom.GetDetectorDistance_mm(); + const double distance_mm = data.geom.GetDetectorDistance_mm(); - double detector_rot[2] = {data.geom.GetPoniRot1_rad(), data.geom.GetPoniRot2_rad()}; - - // The per-frame constants of the reduced residual (see XtalFrameConstants), one entry per frame - // that contributes. Reserved up front and never grown past that, so the residual blocks' pointers - // into it stay valid, and declared before the problem so that it outlives it. - std::vector frame_const; - frame_const.reserve(spots.size()); - - ceres::Problem problem; - - double latt_vec0[3] = {0.0, 0.0, 0.0}; - double latt_vec1[3] = {0.0, 0.0, 0.0}; - double latt_vec2[3] = {0.0, 0.0, 0.0}; - - double rot_vec[3] = {1, 0, 0}; + XtalRefineProblem problem; + problem.crystal_system = data.crystal_system; + problem.distance_mm = distance_mm; + double *beam = problem.beam; + beam[0] = data.geom.GetBeamX_pxl(); + beam[1] = data.geom.GetBeamY_pxl(); + double *detector_rot = problem.detector_rot; + detector_rot[0] = data.geom.GetPoniRot1_rad(); + detector_rot[1] = data.geom.GetPoniRot2_rad(); + double *latt_vec0 = problem.latt_vec0; + double *latt_vec1 = problem.latt_vec1; + double *latt_vec2 = problem.latt_vec2; + double *rot_vec = problem.rot_vec; switch (data.crystal_system) { case gemmi::CrystalSystem::Orthorhombic: @@ -252,14 +175,15 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, const double sin_rot3 = std::sin(data.geom.GetPoniRot3_rad()); // Per-image rotation refinement frees only the beam and the orientation and holds the other five - // blocks constant, so the seven-block residual makes Ceres differentiate 17 parameters to use 5. - // Where that is the configuration, use the reduced residual instead - identical fit, Jet<5> - // autodiff. Any other combination (stills also free the cell, the offline refiner frees distance - // and detector angles) keeps the general form below. + // blocks constant. Where that is the configuration, the solver uses the reduced residual - the + // identical fit, with the crystal half worked out once (see XtalResidualBeamOrientation). Any + // other combination (stills also free the cell, the rotation indexer frees detector angles and + // spindle) keeps the general form. const bool beam_and_orientation_only = data.refine_beam_center && !data.refine_detector_angles && !data.refine_rotation_axis && !data.refine_unit_cell; + problem.beam_and_orientation_only = beam_and_orientation_only; // Sum of w^2 over the spots that entered - the beam prior below is scaled by it so that its // strength relative to the data is the same weighted or not. Equals the residual block count @@ -281,9 +205,8 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, rot_matr = data.axis->GetTransformationAngle(angle_deg); } - if (beam_and_orientation_only) - frame_const.emplace_back(detector_rot, rot_vec, angle_rad, latt_vec1, latt_vec2, - data.crystal_system); + const int frame_index = static_cast(problem.frame_angle_rad.size()); + problem.frame_angle_rad.push_back(angle_rad); // Add residuals for each point for (size_t j = 0; j < spots[i].size(); j++) { @@ -333,7 +256,7 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, const double weight_sq = weight.empty() ? 1.0 : weight[j] * weight[j]; effective_spots += weight_sq; - const XtalResidual residual(pt.x, pt.y, + problem.residuals.emplace_back(pt.x, pt.y, data.geom.GetWavelength_A(), data.geom.GetPixelSize_mm(), cos_rot3, sin_rot3, @@ -341,38 +264,14 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, h, k, l, data.crystal_system, data.geom.GetOrientation()); - - // Ceres has no per-residual weight; ScaledLoss(nullptr, a) multiplies the squared - // residual by the constant a, i.e. it applies a weight of sqrt(a) to the residual. - ceres::LossFunction *loss = weight.empty() - ? nullptr - : new ceres::ScaledLoss(nullptr, weight_sq, - ceres::TAKE_OWNERSHIP); - - if (beam_and_orientation_only) - problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new XtalResidualBeamOrientation(residual, distance_mm, frame_const.back())), - loss, - beam, - latt_vec0 - ); - else - problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new XtalResidualFixedDistance(residual, distance_mm)), - loss, - beam, - detector_rot, - rot_vec, - latt_vec0, - latt_vec1, - latt_vec2 - ); + problem.frame.push_back(frame_index); + // A per-residual weight w enters the squared residual as w^2. + if (!weight.empty()) + problem.weight_sq.push_back(weight_sq); } } - if (problem.NumResidualBlocks() < data.min_spots) + if (static_cast(problem.residuals.size()) < data.min_spots) return false; // The gauge direction of a single-axis rotation experiment - parallel to the spindle - written @@ -440,9 +339,8 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, const double gauge_w = data.geom.GetPixelSize_mm() / (distance_mm * data.geom.GetWavelength_A()) * std::sqrt(effective_spots) / sigma_px; - if (!data.refine_beam_center) - problem.SetParameterBlockConstant(beam); - else if (data.axis) { + problem.beam_constant = !data.refine_beam_center; + if (data.refine_beam_center && data.axis) { // Gauge handling (single-axis rotation): rotating the whole experiment about the spindle leaves every // spot position unchanged, so the beam-centre component PARALLEL to the spindle is a null/gauge-weak // direction. Refining it freely lets it wander (~+3 px) and absorb centroid systematics into a wrong @@ -450,23 +348,19 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, // does drift - it is only LaB6-monitored to ~a few px), RESTRAIN it toward the header with a soft // prior: the gauge direction has ~zero data sensitivity so the prior pins it near the header, while a // real, well-supported drift can still overcome it. - problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new GaugeDirectionPrior(gauge_beam_x, gauge_beam_y, - gauge_beam_x * beam[0] + gauge_beam_y * beam[1], gauge_w)), - nullptr, beam); + problem.priors.push_back({XtalRefinePrior::Block::Beam, gauge_beam_x, gauge_beam_y, + gauge_beam_x * beam[0] + gauge_beam_y * beam[1], gauge_w}); } // Distance, detector angles, rotation axis and cell are parameter blocks only in the general // seven-block residual; the reduced one bakes them in, so there is nothing left to configure. if (!beam_and_orientation_only) { - if (!data.refine_detector_angles) { - problem.SetParameterBlockConstant(detector_rot); - } else { + problem.detector_rot_constant = !data.refine_detector_angles; + if (data.refine_detector_angles) { const double rot_range = 3.0 / 180.0 * PI; for (int i = 0; i < 2; ++i) { - problem.SetParameterLowerBound(detector_rot, i, detector_rot[i] - rot_range); - problem.SetParameterUpperBound(detector_rot, i, detector_rot[i] + rot_range); + problem.detector_rot_lower[i] = detector_rot[i] - rot_range; + problem.detector_rot_upper[i] = detector_rot[i] + rot_range; } // The same gauge as the beam prior above, described a second time: the tilt moves the // direct beam exactly as the beam centre does, at D/pixel px per radian, so leaving @@ -497,20 +391,15 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, for (int i = 0; i < 2; ++i) { if (budget[i] <= 0.0) continue; - problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new GaugeDirectionPrior(dirs[i][0], dirs[i][1], - dirs[i][0] * detector_rot[0] - + dirs[i][1] * detector_rot[1], - gauge_w * (sigma_px / budget[i]) * lever)), - nullptr, detector_rot); + problem.priors.push_back({XtalRefinePrior::Block::DetectorRot, dirs[i][0], dirs[i][1], + dirs[i][0] * detector_rot[0] + dirs[i][1] * detector_rot[1], + gauge_w * (sigma_px / budget[i]) * lever}); } } } - if (!data.refine_rotation_axis) { - problem.SetParameterBlockConstant(rot_vec); - } else { + problem.rot_vec_constant = !data.refine_rotation_axis; + if (data.refine_rotation_axis) { // Only the DIRECTION of the goniometer axis is a parameter. The residual applies // angle_rad * |rot_vec|, so a free three-vector also fits a rotation SCALE - which // GoniometerAxis::Axis() then normalises away, leaving the candidate scored by @@ -520,63 +409,46 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, // data it recovers 54 % of a known scale error, repeated first passes on one dataset // disagree with each other in SIGN, and on the one dataset with a real 1.3 % stage // fault it comes out negative. The rotation scale is measured properly, once, with - // four gates and a jackknife, in PostRefine. - problem.SetManifold(rot_vec, new ceres::SphereManifold<3>); + // four gates and a jackknife, in PostRefine. Refined on the sphere (see SolveXtalRefine). } - if (!data.refine_unit_cell) { - problem.SetParameterBlockConstant(latt_vec1); - problem.SetParameterBlockConstant(latt_vec2); - } else { + problem.latt_vec1_constant = !data.refine_unit_cell; + problem.latt_vec2_constant = !data.refine_unit_cell; + if (data.refine_unit_cell) { // Parameter bounds // Lengths for (int i = 0; i < 3; ++i) { - problem.SetParameterLowerBound(latt_vec1, i, data.min_length_A); - problem.SetParameterUpperBound(latt_vec1, i, data.max_length_A); + problem.latt_vec1_lower[i] = data.min_length_A; + problem.latt_vec1_upper[i] = data.max_length_A; } if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) { - const double beta_lo = std::max(1e-6, PI * (data.min_angle_deg / 180.0)); - const double beta_hi = std::min(PI - 1e-6, PI * (data.max_angle_deg / 180.0)); - problem.SetParameterLowerBound(latt_vec2, 0, beta_lo); - problem.SetParameterUpperBound(latt_vec2, 0, beta_hi); + problem.latt_vec2_constant = false; + problem.latt_vec2_lower[0] = std::max(1e-6, PI * (data.min_angle_deg / 180.0)); + problem.latt_vec2_upper[0] = std::min(PI - 1e-6, PI * (data.max_angle_deg / 180.0)); } else if (data.crystal_system == gemmi::CrystalSystem::Triclinic) { // α, β, γ bounds (radians) const double alo = PI * (data.min_angle_deg / 180.0); const double ahi = PI * (data.max_angle_deg / 180.0); for (int i = 0; i < 3; ++i) { - problem.SetParameterLowerBound(latt_vec2, i, alo); - problem.SetParameterUpperBound(latt_vec2, i, ahi); + problem.latt_vec2_lower[i] = alo; + problem.latt_vec2_upper[i] = ahi; } } else { // Orthorhombic / Tetragonal / Cubic / Hexagonal: // latt_vec2 has no meaning for these systems — always freeze it. - problem.SetParameterBlockConstant(latt_vec2); + problem.latt_vec2_constant = true; } } } - // Configure solver - ceres::Solver::Options options; - // Normal equations, not QR. The problem is very tall and thin - thousands of spots against at - // most 17 parameters - and that is the shape DENSE_QR handles worst: it copies the Jacobian out - // of Ceres' row-major storage into a column-major buffer on every solve, and Eigen's blocked - // Householder then degenerates to the unblocked path because its block size is min(48, columns). - // Accumulating J^T J reads the Jacobian once instead. Both solve the same damped system, so the - // step is the same to round-off; the column scaling Ceres applies by default and the LM diagonal - // keep the squared condition number in hand. - options.linear_solver_type = ceres::DENSE_NORMAL_CHOLESKY; - options.minimizer_progress_to_stdout = false; + // Stopping rule: a bound on iterations is reproducible, a bound on wall-clock time is not (see + // XtalOptimizerData::max_iterations). if (data.max_iterations > 0) - options.max_num_iterations = data.max_iterations; + problem.options.max_iterations = data.max_iterations; else - options.max_solver_time_in_seconds = data.max_time; - options.logging_type = ceres::LoggingType::SILENT; - options.num_threads = num_threads; // usually 1 (called from many threads); caller may raise it - ceres::Solver::Summary summary; - - // Run optimization - ceres::Solve(options, &problem, &summary); + problem.options.max_time_s = data.max_time; + const LMSummary summary = SolveXtalRefine(problem, num_threads); // Only a genuine numerical failure is rejected here: a solve that ran out of iterations or // out of time but still descended counts as usable, which is what the real-time caller @@ -653,7 +525,7 @@ bool XtalOptimizerRotationOnly(XtalOptimizerData &data, return false; // Parameter: angle-axis for the extra rotation. Identity == {0,0,0}. - double rot_aa[3] = {0.0, 0.0, 0.0}; + std::vector rot_aa = {0.0, 0.0, 0.0}; // Spot selection by current indexing (same approach as XtalOptimizerInternal) const Coord a0 = data.latt.Vec0(); @@ -662,7 +534,12 @@ bool XtalOptimizerRotationOnly(XtalOptimizerData &data, const float tol_sq = tolerance * tolerance; - ceres::Problem problem; + // Each selected spot: its observed reciprocal vector and the indices it is fitted to. + struct Observation { + Coord s_obs; + double h, k, l; + }; + std::vector observations; for (const auto &pt : spots) { if (!data.index_ice_rings && pt.ice_ring) @@ -697,42 +574,68 @@ bool XtalOptimizerRotationOnly(XtalOptimizerData &data, if (data.axis.has_value()) s_obs = data.axis->GetTransformationAngle(pt.phi) * s_obs; - auto *cost = - new ceres::AutoDiffCostFunction( - new XtalResidualRotationOnlyPrecomp(s_obs, data.latt, h, k, l) - ); - - problem.AddResidualBlock(cost, nullptr, rot_aa); + observations.push_back({s_obs, h, k, l}); } - if (problem.NumResidualBlocks() < data.min_spots) + if (static_cast(observations.size()) < data.min_spots) return false; - // Regularization: prefer the smallest rotation correction that fits the - // data. This is essential when spots are nearly coplanar in reciprocal - // space (e.g. still images), where the rotation component perpendicular - // to the scattering plane is otherwise underdetermined. - // The weight is in Å⁻¹ rad⁻¹; tune relative to your typical residual. - { - const double reg_weight = 0.05; // e.g. 0.05 - problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new RotationNormRegularizer(reg_weight)), - nullptr, rot_aa); - } + // Residual: s_obs - R(rot_aa) (h a* + k b* + l c*), the reciprocal basis rotated once per + // evaluation and shared by every spot. + // + // Regularization: prefer the smallest rotation correction that fits the data, w * rot_aa. This is + // essential when spots are nearly coplanar in reciprocal space (e.g. still images), where the + // rotation component perpendicular to the scattering plane is otherwise underdetermined. The + // weight is in A^-1 rad^-1, relative to the typical residual. + const double reg_weight = 0.05; + const Coord astar = data.latt.Astar(), bstar = data.latt.Bstar(), cstar = data.latt.Cstar(); + const auto evaluate = [&](const double *aa, double &cost, Eigen::VectorXd *g, Eigen::MatrixXd *H) { + using D = Dual<3>; + const D aa_d[3] = {D::Variable(aa[0], 0), D::Variable(aa[1], 1), D::Variable(aa[2], 2)}; + const AngleAxisRotator rot(aa_d); + const double astar_unrot[3] = {astar.x, astar.y, astar.z}; + const double bstar_unrot[3] = {bstar.x, bstar.y, bstar.z}; + const double cstar_unrot[3] = {cstar.x, cstar.y, cstar.z}; + D astar_rot[3], bstar_rot[3], cstar_rot[3]; + rot.Rotate(astar_unrot, astar_rot); + rot.Rotate(bstar_unrot, bstar_rot); + rot.Rotate(cstar_unrot, cstar_rot); - ceres::Solver::Options options; - options.linear_solver_type = ceres::DENSE_NORMAL_CHOLESKY; // tall and thin, as above - options.minimizer_progress_to_stdout = false; + cost = 0.0; + const auto add = [&](double r, const double *J) { + cost += 0.5 * r * r; + if (!g) + return; + for (int i = 0; i < 3; i++) { + (*g)[i] += J[i] * r; + for (int j = 0; j < 3; j++) + (*H)(i, j) += J[i] * J[j]; + } + }; + for (const auto &o: observations) { + const double s_obs[3] = {o.s_obs.x, o.s_obs.y, o.s_obs.z}; + for (int c = 0; c < 3; c++) { + const D pred = o.h * astar_rot[c] + o.k * bstar_rot[c] + o.l * cstar_rot[c]; + const double J[3] = {-pred.v[0], -pred.v[1], -pred.v[2]}; + add(s_obs[c] - pred.a, J); + } + } + for (int c = 0; c < 3; c++) { + double J[3] = {0.0, 0.0, 0.0}; + J[c] = reg_weight; + add(reg_weight * aa[c], J); + } + return std::isfinite(cost) && (!g || (g->allFinite() && H->allFinite())); + }; + + std::vector blocks(1); + blocks[0].size = 3; + LMOptions options; if (data.max_iterations > 0) - options.max_num_iterations = data.max_iterations; + options.max_iterations = data.max_iterations; else - options.max_solver_time_in_seconds = data.max_time; - options.logging_type = ceres::LoggingType::SILENT; - options.num_threads = 1; - - ceres::Solver::Summary summary; - ceres::Solve(options, &problem, &summary); + options.max_time_s = data.max_time; + const LMSummary summary = SolveLM(rot_aa, blocks, options, evaluate); if (!summary.IsSolutionUsable()) return false; @@ -747,7 +650,7 @@ bool XtalOptimizerRotationOnly(XtalOptimizerData &data, // rotating the reciprocal vectors (a*, b*, c*) by the same R. No // transpose or inversion of R is needed here. double R_raw[9]; - ceres::AngleAxisToRotationMatrix(rot_aa, R_raw); // row-major 3x3 + ceres::AngleAxisToRotationMatrix(rot_aa.data(), R_raw); // row-major 3x3 Eigen::Matrix3d R; R << R_raw[0], R_raw[3], R_raw[6], diff --git a/image_analysis/geom_refinement/XtalOptimizer.h b/image_analysis/geom_refinement/XtalOptimizer.h index 1bf9084de..0b64937bb 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.h +++ b/image_analysis/geom_refinement/XtalOptimizer.h @@ -62,7 +62,7 @@ struct XtalOptimizerData { std::optional angle_axis; }; -// num_threads sets the Ceres solver thread count for the internal least-squares refine. It defaults +// num_threads sets the thread count of the internal least-squares refine (the answer does not depend on it). It defaults // to 1 because XtalOptimizer is usually called from many threads at once; raise it only when a caller // runs a small number of refinements concurrently and wants each to use several cores. bool XtalOptimizer(XtalOptimizerData &data, std::span> spots, diff --git a/image_analysis/geom_refinement/XtalRefine.cpp b/image_analysis/geom_refinement/XtalRefine.cpp new file mode 100644 index 000000000..3b466b97d --- /dev/null +++ b/image_analysis/geom_refinement/XtalRefine.cpp @@ -0,0 +1,334 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "XtalRefine.h" + +#include + +#include "Dual.h" +#include "../../common/ParallelFor.h" + +namespace { + // Ambient layout of the parameter vector and the order of the blocks in it. + constexpr int OFF_BEAM = 0, OFF_ROT = 2, OFF_AXIS = 4, OFF_P0 = 7, OFF_LEN = 10, OFF_ANG = 13, N_AMBIENT = 16; + + // Derivative lanes. The observed half of a residual depends on beam, detector angles and spindle + // (OBS lanes), the predicted half on orientation and cell (PRED lanes); each half is carried on a + // dual number of its own width and the two are put side by side only in the sums. + constexpr int OBS = 6, PRED = 9, LANES = OBS + PRED; + using DO = Dual; + using DP = Dual; + + // The reduced problem: beam (2) and orientation (3) only. + constexpr int R_OBS = 2, R_PRED = 3, R_LANES = R_OBS + R_PRED; + + // Sums over one block of residuals, in lane coordinates. + template + struct Sums { + double cost = 0.0; + double g[L] = {}; + double H[L][L] = {}; // upper triangle + + void Add(double r, const double *J, double w2) { + cost += 0.5 * w2 * r * r; + for (int i = 0; i < L; i++) { + const double wj = w2 * J[i]; + g[i] += wj * r; + for (int j = i; j < L; j++) + H[i][j] += wj * J[j]; + } + } + + void Add(const Sums &o) { + cost += o.cost; + for (int i = 0; i < L; i++) { + g[i] += o.g[i]; + for (int j = i; j < L; j++) + H[i][j] += o.H[i][j]; + } + } + }; + + std::vector MakeBlocks(const XtalRefineProblem &p) { + std::vector blocks(6); + blocks[0].offset = OFF_BEAM; + blocks[0].size = 2; + blocks[0].constant = p.beam_constant; + + blocks[1].offset = OFF_ROT; + blocks[1].size = 2; + blocks[1].constant = p.beam_and_orientation_only || p.detector_rot_constant; + for (int j = 0; j < 2; j++) { + blocks[1].lower[j] = p.detector_rot_lower[j]; + blocks[1].upper[j] = p.detector_rot_upper[j]; + } + + blocks[2].offset = OFF_AXIS; + blocks[2].size = 3; + blocks[2].constant = p.beam_and_orientation_only || p.rot_vec_constant; + blocks[2].sphere = true; + + blocks[3].offset = OFF_P0; + blocks[3].size = 3; + + blocks[4].offset = OFF_LEN; + blocks[4].size = 3; + blocks[4].constant = p.beam_and_orientation_only || p.latt_vec1_constant; + + blocks[5].offset = OFF_ANG; + blocks[5].size = 3; + blocks[5].constant = p.beam_and_orientation_only || p.latt_vec2_constant; + + for (int j = 0; j < 3; j++) { + blocks[4].lower[j] = p.latt_vec1_lower[j]; + blocks[4].upper[j] = p.latt_vec1_upper[j]; + blocks[5].lower[j] = p.latt_vec2_lower[j]; + blocks[5].upper[j] = p.latt_vec2_upper[j]; + } + return blocks; + } + + // Where each lane lands among the tangent coordinates of the free blocks; -1 for a held one. + template + std::array LaneToTangent(const std::vector &blocks, const std::array &lane_block, + const std::array &lane_index) { + std::array map{}; + std::array tangent_offset{}; + int t = 0; + for (size_t b = 0; b < blocks.size(); b++) { + tangent_offset[b] = t; + t += blocks[b].TangentSize(); + } + for (size_t l = 0; l < L; l++) + map[l] = blocks[lane_block[l]].constant ? -1 : tangent_offset[lane_block[l]] + lane_index[l]; + return map; + } + + template + void ToTangent(const Sums &s, const std::array &map, Eigen::VectorXd &g, Eigen::MatrixXd &H) { + for (int i = 0; i < L; i++) { + if (map[i] < 0) + continue; + g[map[i]] += s.g[i]; + for (int j = i; j < L; j++) { + if (map[j] < 0) + continue; + H(map[i], map[j]) += s.H[i][j]; + if (map[i] != map[j]) + H(map[j], map[i]) += s.H[i][j]; + } + } + } + + void AddPriors(const XtalRefineProblem &p, const double *x, double &cost, Eigen::VectorXd *g, + Eigen::MatrixXd *H, int beam_tangent, int rot_tangent) { + for (const auto &prior: p.priors) { + const bool on_beam = prior.block == XtalRefinePrior::Block::Beam; + const double *v = x + (on_beam ? OFF_BEAM : OFF_ROT); + const double r = prior.weight * (prior.gx * v[0] + prior.gy * v[1] - prior.p0); + cost += 0.5 * r * r; + const int t = on_beam ? beam_tangent : rot_tangent; + if (!g || t < 0) + continue; + const double J[2] = {prior.weight * prior.gx, prior.weight * prior.gy}; + for (int i = 0; i < 2; i++) { + (*g)[t + i] += J[i] * r; + for (int j = 0; j < 2; j++) + (*H)(t + i, t + j) += J[i] * J[j]; + } + } + } + + template + D Seed(double value, int lane, bool free) { + return free ? D::Variable(value, lane) : D(value); + } + + // Residual blocks of at least this many residuals; the cut depends on the count alone. + constexpr int MIN_RESIDUALS_PER_BLOCK = 256; + + template + Sums SumResiduals(const XtalRefineProblem &p, int num_threads, Fn &&residual) { + const int n = static_cast(p.residuals.size()); + std::vector> partial(ReductionBlocks(n, MIN_RESIDUALS_PER_BLOCK)); + ParallelBlocks(n, std::max(1, num_threads), [&](int b, int lo, int hi) { + for (int i = lo; i < hi; i++) + residual(i, partial[b]); + }, MIN_RESIDUALS_PER_BLOCK); + Sums total; + for (const auto &s: partial) + total.Add(s); + return total; + } + + double WeightSq(const XtalRefineProblem &p, int i) { + return p.weight_sq.empty() ? 1.0 : p.weight_sq[i]; + } + + bool AllFinite(double cost, const Eigen::VectorXd *g, const Eigen::MatrixXd *H) { + return std::isfinite(cost) && (!g || (g->allFinite() && H->allFinite())); + } + + // Seven-block residual (XtalResidualFixedDistance), with every block that depends on parameters + // alone - detector-angle sines and cosines, the spindle back-rotation of each frame, the reciprocal + // basis of the cell, the orientation's rotation - worked out once per evaluation. + bool EvaluateGeneral(const XtalRefineProblem &p, const std::array &map, int num_threads, + const double *x, double &cost, Eigen::VectorXd *g, Eigen::MatrixXd *H) { + const bool beam_free = map[0] >= 0, rot_free = map[2] >= 0, axis_free = map[4] >= 0; + const bool p0_free = map[OBS] >= 0, len_free = map[OBS + 3] >= 0, ang_free = map[OBS + 6] >= 0; + + const DO beam[2] = {Seed(x[OFF_BEAM], 0, beam_free), + Seed(x[OFF_BEAM + 1], 1, beam_free)}; + const DO rot1 = Seed(x[OFF_ROT], 2, rot_free); + const DO rot2 = Seed(x[OFF_ROT + 1], 3, rot_free); + const DO c1 = cos(rot1), s1 = sin(rot1), c2 = cos(rot2), s2 = sin(rot2); + + DO axis[3] = {x[OFF_AXIS], x[OFF_AXIS + 1], x[OFF_AXIS + 2]}; + if (axis_free) { + double J[3][2]; + SpherePlusJacobian3(x + OFF_AXIS, J); + for (int k = 0; k < 3; k++) { + axis[k].v[4] = J[k][0]; + axis[k].v[5] = J[k][1]; + } + } + std::vector> rot_back; + rot_back.reserve(p.frame_angle_rad.size()); + for (const double angle: p.frame_angle_rad) { + const DO aa_back[3] = {angle * axis[0], angle * axis[1], angle * axis[2]}; + rot_back.emplace_back(aa_back); + } + + DP p0[3], len[3], ang[3]; + for (int k = 0; k < 3; k++) { + p0[k] = Seed(x[OFF_P0 + k], k, p0_free); + len[k] = Seed(x[OFF_LEN + k], 3 + k, len_free); + ang[k] = Seed(x[OFF_ANG + k], 6 + k, ang_free); + } + Eigen::Matrix bxc, cxa, axb; + DP invV; + XtalResidual::ReciprocalBasis(len, ang, p.crystal_system, bxc, cxa, axb, invV); + const AngleAxisRotator rot_p0(p0); + + const Sums s = SumResiduals(p, num_threads, [&](int i, Sums &acc) { + const XtalResidual &res = p.residuals[i]; + DO obs[3]; + res.ObservedRecipCore(beam, p.distance_mm, c1, s1, c2, s2, rot_back[p.frame[i]], obs); + DP unrot[3], pred[3]; + res.CombineRecipUnrot(bxc, cxa, axb, invV, unrot); + rot_p0.Rotate(unrot, pred); + const double w2 = WeightSq(p, i); + for (int k = 0; k < 3; k++) { + double J[LANES]; + for (int l = 0; l < OBS; l++) + J[l] = obs[k].v[l]; + for (int l = 0; l < PRED; l++) + J[OBS + l] = -pred[k].v[l]; + acc.Add(obs[k].a - pred[k].a, J, w2); + } + }); + + cost = s.cost; + if (g) + ToTangent(s, map, *g, *H); + AddPriors(p, x, cost, g, H, map[0], map[2]); + return AllFinite(cost, g, H); + } + + // The reduced residual (XtalResidualBeamOrientation): detector, spindle and cell held, so the + // back-rotation of each frame and the unrotated prediction of each residual are constants. + bool EvaluateBeamOrientation(const XtalRefineProblem &p, const std::array &map, + const std::vector> &rot_back, + const std::vector> &unrot, double c1, double s1, + double c2, double s2, int num_threads, + const double *x, double &cost, Eigen::VectorXd *g, Eigen::MatrixXd *H) { + using DB = Dual; + using DR = Dual; + const bool beam_free = map[0] >= 0; + const DB beam[2] = {beam_free ? DB::Variable(x[OFF_BEAM], 0) : DB(x[OFF_BEAM]), + beam_free ? DB::Variable(x[OFF_BEAM + 1], 1) : DB(x[OFF_BEAM + 1])}; + const DR p0[3] = {DR::Variable(x[OFF_P0], 0), DR::Variable(x[OFF_P0 + 1], 1), + DR::Variable(x[OFF_P0 + 2], 2)}; + const AngleAxisRotator rot_p0(p0); + + const Sums s = SumResiduals(p, num_threads, [&](int i, Sums &acc) { + const XtalResidual &res = p.residuals[i]; + DB obs[3]; + res.ObservedRecipCore(beam, p.distance_mm, c1, s1, c2, s2, rot_back[p.frame[i]], obs); + DR pred[3]; + rot_p0.Rotate(unrot[i].data(), pred); + const double w2 = WeightSq(p, i); + for (int k = 0; k < 3; k++) { + double J[R_LANES]; + for (int l = 0; l < R_OBS; l++) + J[l] = obs[k].v[l]; + for (int l = 0; l < R_PRED; l++) + J[R_OBS + l] = -pred[k].v[l]; + acc.Add(obs[k].a - pred[k].a, J, w2); + } + }); + + cost = s.cost; + if (g) + ToTangent(s, map, *g, *H); + AddPriors(p, x, cost, g, H, map[0], -1); + return AllFinite(cost, g, H); + } +} + +LMSummary SolveXtalRefine(XtalRefineProblem &p, int num_threads) { + const std::vector blocks = MakeBlocks(p); + + std::vector x(N_AMBIENT); + const auto put = [&](int off, const double *v, int n) { for (int i = 0; i < n; i++) x[off + i] = v[i]; }; + put(OFF_BEAM, p.beam, 2); + put(OFF_ROT, p.detector_rot, 2); + put(OFF_AXIS, p.rot_vec, 3); + put(OFF_P0, p.latt_vec0, 3); + put(OFF_LEN, p.latt_vec1, 3); + put(OFF_ANG, p.latt_vec2, 3); + + LMSummary summary; + if (p.beam_and_orientation_only) { + const std::array lane_block = {0, 0, 3, 3, 3}; + const std::array lane_index = {0, 1, 0, 1, 2}; + const auto map = LaneToTangent(blocks, lane_block, lane_index); + + std::vector> rot_back; + rot_back.reserve(p.frame_angle_rad.size()); + for (const double angle: p.frame_angle_rad) + rot_back.push_back(XtalFrameConstants::BackRotator(angle, p.rot_vec)); + Eigen::Matrix bxc, cxa, axb; + double invV; + XtalResidual::ReciprocalBasis(p.latt_vec1, p.latt_vec2, p.crystal_system, bxc, cxa, axb, invV); + std::vector> unrot(p.residuals.size()); + for (size_t i = 0; i < p.residuals.size(); i++) + p.residuals[i].CombineRecipUnrot(bxc, cxa, axb, invV, unrot[i].data()); + const double c1 = std::cos(p.detector_rot[0]), s1 = std::sin(p.detector_rot[0]); + const double c2 = std::cos(p.detector_rot[1]), s2 = std::sin(p.detector_rot[1]); + + summary = SolveLM(x, blocks, p.options, [&](const double *at, double &cost, Eigen::VectorXd *g, + Eigen::MatrixXd *H) { + return EvaluateBeamOrientation(p, map, rot_back, unrot, c1, s1, c2, s2, num_threads, at, cost, g, H); + }); + } else { + const std::array lane_block = {0, 0, 1, 1, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5}; + const std::array lane_index = {0, 1, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2, 0, 1, 2}; + const auto map = LaneToTangent(blocks, lane_block, lane_index); + summary = SolveLM(x, blocks, p.options, [&](const double *at, double &cost, Eigen::VectorXd *g, + Eigen::MatrixXd *H) { + return EvaluateGeneral(p, map, num_threads, at, cost, g, H); + }); + } + + if (summary.IsSolutionUsable()) { + const auto get = [&](int off, double *v, int n) { for (int i = 0; i < n; i++) v[i] = x[off + i]; }; + get(OFF_BEAM, p.beam, 2); + get(OFF_ROT, p.detector_rot, 2); + get(OFF_AXIS, p.rot_vec, 3); + get(OFF_P0, p.latt_vec0, 3); + get(OFF_LEN, p.latt_vec1, 3); + get(OFF_ANG, p.latt_vec2, 3); + } + return summary; +} diff --git a/image_analysis/geom_refinement/XtalRefine.h b/image_analysis/geom_refinement/XtalRefine.h new file mode 100644 index 000000000..4ee88870e --- /dev/null +++ b/image_analysis/geom_refinement/XtalRefine.h @@ -0,0 +1,63 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include + +#include "XtalResidual.h" +#include "LMSolver.h" + +// A soft restraint w * (g . p - p0) on one direction g of a two-component block - the beam centre or +// the detector tilt, which are the same gauge seen twice and so take the same direction (see +// XtalOptimizer). +struct XtalRefinePrior { + enum class Block { Beam, DetectorRot } block = Block::Beam; + double gx = 0.0, gy = 0.0, p0 = 0.0, weight = 0.0; +}; + +// The least-squares problem XtalOptimizer solves, as data: the residuals with the frame each belongs +// to, the parameter blocks with what is held and what is bounded, and the priors. The parameter arrays +// are in/out. Parameter blocks: beam(2), detector_rot(2), rot_vec(3, the spindle, refined on the sphere), +// latt_vec0(3, orientation angle-axis), latt_vec1(3, cell lengths), latt_vec2(3, cell angles) - see +// XtalResidual. The distance is a constant of the problem. +struct XtalRefineProblem { + static constexpr double kNoBound = std::numeric_limits::max(); + + gemmi::CrystalSystem crystal_system = gemmi::CrystalSystem::Triclinic; + // Only beam and orientation free, everything else held: the reduced residual (see + // XtalResidualBeamOrientation), whose crystal half is a constant of the problem. + bool beam_and_orientation_only = false; + double distance_mm = 0.0; + + std::vector residuals; + std::vector frame; // per residual, an index into frame_angle_rad + std::vector frame_angle_rad; + std::vector weight_sq; // per residual; empty = unweighted + + double beam[2] = {0, 0}; + double detector_rot[2] = {0, 0}; + double rot_vec[3] = {1, 0, 0}; + double latt_vec0[3] = {0, 0, 0}; + double latt_vec1[3] = {0, 0, 0}; + double latt_vec2[3] = {0, 0, 0}; + + bool beam_constant = false; + bool detector_rot_constant = true; + bool rot_vec_constant = true; + bool latt_vec1_constant = true; + bool latt_vec2_constant = true; + + double detector_rot_lower[2] = {-kNoBound, -kNoBound}, detector_rot_upper[2] = {kNoBound, kNoBound}; + double latt_vec1_lower[3] = {-kNoBound, -kNoBound, -kNoBound}, latt_vec1_upper[3] = {kNoBound, kNoBound, kNoBound}; + double latt_vec2_lower[3] = {-kNoBound, -kNoBound, -kNoBound}, latt_vec2_upper[3] = {kNoBound, kNoBound, kNoBound}; + + std::vector priors; + + LMOptions options; +}; + +// Solves the problem in place. The residual sums are cut into blocks that depend on the problem +// alone, so the answer is the same at any thread count. +LMSummary SolveXtalRefine(XtalRefineProblem &problem, int num_threads); diff --git a/image_analysis/geom_refinement/XtalResidual.h b/image_analysis/geom_refinement/XtalResidual.h index b67e055c8..26c05a26a 100644 --- a/image_analysis/geom_refinement/XtalResidual.h +++ b/image_analysis/geom_refinement/XtalResidual.h @@ -7,8 +7,6 @@ #include -#include "ceres/ceres.h" -#include "ceres/rotation.h" #include "gemmi/symmetry.hpp" #include "../../common/JFJochException.h" @@ -110,6 +108,9 @@ inline void EffectiveCellFromParams(gemmi::CrystalSystem symmetry, const double } } +// The scalar type is a template parameter throughout: a double, a ceres::Jet or a Dual (Dual.h). The +// mathematical functions are called unqualified, so each type's own overload is found by lookup. +// // Detector -> reciprocal geometry residual, shared by the per-image XtalOptimizer (one lattice, one // frame) and the offline GeometryRefiner (shared beam/distance/cell blocks, one orientation block per // frame). Parameter blocks: beam(2), distance_mm(1), detector_rot(2 = rot1,rot2), rotation_axis(3), @@ -164,16 +165,18 @@ struct XtalResidual { // detector_rot[0] = rot1, detector_rot[1] = rot2 are refined; rot3 is fixed // (e.g. from a PONI import) and baked in here as a constant so that a non-zero // rot3 is not silently dropped during refinement. + using std::cos; + using std::sin; const C rot1 = detector_rot[0]; const C rot2 = detector_rot[1]; // Ry(+rot1): rotation around Y-axis - const C c1 = ceres::cos(rot1); - const C s1 = ceres::sin(rot1); + const C c1 = cos(rot1); + const C s1 = sin(rot1); // Rx(-rot2): rotation around X-axis with inverted sign (PyFAI left-handed) - const C c2 = ceres::cos(rot2); - const C s2 = ceres::sin(rot2); + const C c2 = cos(rot2); + const C s2 = sin(rot2); // Apply the goniometer "back-to-start" rotation of this frame's angle. const C aa_back[3] = { @@ -227,7 +230,8 @@ struct XtalResidual { const T z = t2_z; // convert to recip space - const T lab_norm = ceres::sqrt(x * x + y * y + z * z); + using std::sqrt; + const T lab_norm = sqrt(x * x + y * y + z * z); const T inv_norm = T(1) / lab_norm; T recip_raw[3]; @@ -258,6 +262,9 @@ struct XtalResidual { static void ReciprocalBasis(const C *const p1, const C *const p2, gemmi::CrystalSystem symmetry, Eigen::Matrix &bxc, Eigen::Matrix &cxa, Eigen::Matrix &axb, C &invV) { + using std::cos; + using std::sin; + using std::sqrt; // Build unit cell lengths and B (convention: columns are a, b, c prior to global rotation) Eigen::Matrix e_uc_len = Eigen::Matrix::Zero(); Eigen::Matrix B = Eigen::Matrix::Identity(); @@ -275,14 +282,14 @@ struct XtalResidual { } else if (symmetry == gemmi::CrystalSystem::Monoclinic) { // Unique axis b: alpha = gamma = 90°, beta free (angle between a and c) e_uc_len << p1[0], p1[1], p1[2]; - B(0, 2) = ceres::cos(p2[0]); - B(2, 2) = ceres::sin(p2[0]); + B(0, 2) = cos(p2[0]); + B(2, 2) = sin(p2[0]); } else { // Triclinic: p1 = (a,b,c), p2 = (alpha, beta, gamma) in radians - const C ca = ceres::cos(p2[0]); - const C cb = ceres::cos(p2[1]); - const C cg = ceres::cos(p2[2]); - const C sg = ceres::sin(p2[2]); + const C ca = cos(p2[0]); + const C cb = cos(p2[1]); + const C cg = cos(p2[2]); + const C sg = sin(p2[2]); e_uc_len << p1[0], p1[1], p1[2]; @@ -297,7 +304,7 @@ struct XtalResidual { const C cx = cb; const C cy = (ca - cb * cg) / sg; const C v = C(1) - cx * cx - cy * cy; - const C cz = (v >= C(0)) ? ceres::sqrt(v) : C(0); + const C cz = (v >= C(0)) ? sqrt(v) : C(0); B(0, 2) = cx; B(1, 2) = cy; @@ -399,12 +406,12 @@ struct XtalResidual { // reciprocal basis of the fixed cell. They are constants of the whole problem there - one frame per // image, one cell - and deriving them inside each residual costs six trigonometric calls, a hypot and a // division per evaluation for numbers that never change. Built by the caller, which has to keep it alive -// as long as the ceres::Problem that points at it. +// as long as the residuals that point at it. struct XtalFrameConstants { XtalFrameConstants(const double *detector_rot, const double *rotation_axis, double angle_rad, const double *uc_len, const double *uc_angle, gemmi::CrystalSystem symmetry) - : c1(ceres::cos(detector_rot[0])), s1(ceres::sin(detector_rot[0])), - c2(ceres::cos(detector_rot[1])), s2(ceres::sin(detector_rot[1])), + : c1(cos(detector_rot[0])), s1(sin(detector_rot[0])), + c2(cos(detector_rot[1])), s2(sin(detector_rot[1])), rot_back(BackRotator(angle_rad, rotation_axis)) { XtalResidual::ReciprocalBasis(uc_len, uc_angle, symmetry, bxc, cxa, axb, invV); } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index f01b89506..799c06c0d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -112,6 +112,8 @@ ADD_EXECUTABLE(jfjoch_test RingsFromProfileTest.cpp CalibrationTest.cpp XtalOptimizerTest.cpp + XtalRefineTest.cpp + XtalRefineCeres.h CrystalLatticeTest.cpp FPGAPTPTest.cpp ResolutionShellsTest.cpp diff --git a/tests/XtalRefineCeres.h b/tests/XtalRefineCeres.h new file mode 100644 index 000000000..77f34b4f9 --- /dev/null +++ b/tests/XtalRefineCeres.h @@ -0,0 +1,92 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// The reference the LM solver of XtalRefine is checked against: the same XtalRefineProblem handed to +// Ceres exactly as XtalOptimizer used to build it - the same residual functors, losses, priors, bounds, +// manifold and options. + +#include "ceres/ceres.h" +#include "../image_analysis/geom_refinement/XtalRefine.h" + +struct XtalRefineCeresPrior { + XtalRefineCeresPrior(double gx, double gy, double p0, double weight) + : gx(gx), gy(gy), p0(p0), weight(weight) {} + template + bool operator()(const T *const p, T *residual) const { + residual[0] = T(weight) * (T(gx) * p[0] + T(gy) * p[1] - T(p0)); + return true; + } + double gx, gy, p0, weight; +}; + +inline ceres::Solver::Summary SolveXtalRefineCeres(XtalRefineProblem &p, int num_threads) { + std::vector frame_const; + frame_const.reserve(p.frame_angle_rad.size()); + if (p.beam_and_orientation_only) + for (const double angle: p.frame_angle_rad) + frame_const.emplace_back(p.detector_rot, p.rot_vec, angle, p.latt_vec1, p.latt_vec2, p.crystal_system); + + ceres::Problem problem; + for (size_t i = 0; i < p.residuals.size(); i++) { + ceres::LossFunction *loss = p.weight_sq.empty() + ? nullptr + : new ceres::ScaledLoss(nullptr, p.weight_sq[i], ceres::TAKE_OWNERSHIP); + if (p.beam_and_orientation_only) + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction( + new XtalResidualBeamOrientation(p.residuals[i], p.distance_mm, frame_const[p.frame[i]])), + loss, p.beam, p.latt_vec0); + else + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction( + new XtalResidualFixedDistance(p.residuals[i], p.distance_mm)), + loss, p.beam, p.detector_rot, p.rot_vec, p.latt_vec0, p.latt_vec1, p.latt_vec2); + } + for (const auto &prior: p.priors) + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction( + new XtalRefineCeresPrior(prior.gx, prior.gy, prior.p0, prior.weight)), + nullptr, prior.block == XtalRefinePrior::Block::Beam ? p.beam : p.detector_rot); + + const auto bounds = [&](double *block, const double *lo, const double *hi, int n) { + for (int i = 0; i < n; i++) { + if (lo[i] > -XtalRefineProblem::kNoBound) + problem.SetParameterLowerBound(block, i, lo[i]); + if (hi[i] < XtalRefineProblem::kNoBound) + problem.SetParameterUpperBound(block, i, hi[i]); + } + }; + if (p.beam_constant) + problem.SetParameterBlockConstant(p.beam); + if (!p.beam_and_orientation_only) { + if (p.detector_rot_constant) + problem.SetParameterBlockConstant(p.detector_rot); + else + bounds(p.detector_rot, p.detector_rot_lower, p.detector_rot_upper, 2); + if (p.rot_vec_constant) + problem.SetParameterBlockConstant(p.rot_vec); + else + problem.SetManifold(p.rot_vec, new ceres::SphereManifold<3>); + if (p.latt_vec1_constant) + problem.SetParameterBlockConstant(p.latt_vec1); + else + bounds(p.latt_vec1, p.latt_vec1_lower, p.latt_vec1_upper, 3); + if (p.latt_vec2_constant) + problem.SetParameterBlockConstant(p.latt_vec2); + else + bounds(p.latt_vec2, p.latt_vec2_lower, p.latt_vec2_upper, 3); + } + + ceres::Solver::Options options; + options.linear_solver_type = ceres::DENSE_NORMAL_CHOLESKY; + options.minimizer_progress_to_stdout = false; + options.max_num_iterations = p.options.max_iterations; + options.max_solver_time_in_seconds = p.options.max_time_s; + options.logging_type = ceres::LoggingType::SILENT; + options.num_threads = num_threads; + ceres::Solver::Summary summary; + ceres::Solve(options, &problem, &summary); + return summary; +} diff --git a/tests/XtalRefineTest.cpp b/tests/XtalRefineTest.cpp new file mode 100644 index 000000000..a26021f91 --- /dev/null +++ b/tests/XtalRefineTest.cpp @@ -0,0 +1,127 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +// Ceres first: its logging header defines a CHECK macro of its own, which Catch's must replace here. +#include "XtalRefineCeres.h" +#undef CHECK +#include + +#include "../image_analysis/geom_refinement/LatticeReduction.h" +#include "../image_analysis/bragg_prediction/BraggPrediction.h" + +// The LM solver of XtalRefine is meant to take Ceres' path to Ceres' answer: same steps accepted, same +// stopping rule, same point. These cases hand one problem to both and compare. +namespace { + // A monoclinic crystal rotated about X over ten 3-degree frames, its predicted spots as observations; + // the refinement starts from a perturbed beam, tilt and cell. + XtalRefineProblem RotationProblem(bool reduced, bool weighted) { + DiffractionExperiment exp; + exp.IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(1000).BeamY_pxl(1000) + .PoniRot1_rad(0.01).PoniRot2_rad(0.02).DetectorDistance_mm(200); + const auto geom = exp.GetDiffractionGeometry(); + const CrystalLattice latt(40, 50, 80, 90, 95, 90); + const GoniometerAxis axis("omega", 0.0f, 3.0f, Coord(1, 0, 0), std::nullopt); + const gemmi::CrystalSystem sys = gemmi::CrystalSystem::Monoclinic; + + XtalRefineProblem p; + p.crystal_system = sys; + p.beam_and_orientation_only = reduced; + p.distance_mm = 200; + + BraggPrediction prediction; + const BraggPredictionSettings settings{.high_res_A = 1.5, .ewald_dist_cutoff = 0.002}; + for (int img = 0; img < 10; img++) { + const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f; + const auto n = prediction.Calc(exp, latt.Multiply(axis.GetTransformationAngle(angle_deg).transpose()), + settings); + p.frame_angle_rad.push_back(angle_deg * PI / 180.0); + for (int i = 0; i < n; i++) { + const auto &r = prediction.GetReflections().at(i); + p.residuals.emplace_back(r.predicted_x + 0.3 * std::sin(i), r.predicted_y + 0.3 * std::cos(i), + geom.GetWavelength_A(), geom.GetPixelSize_mm(), 1.0, 0.0, + angle_deg * PI / 180.0, r.h, r.k, r.l, sys); + p.frame.push_back(img); + if (weighted) + p.weight_sq.push_back(0.2 + 0.6 * (i % 5) / 4.0); + } + } + + p.beam[0] = 1000.0; + p.beam[1] = 997.0; + p.detector_rot[0] = 0.012; + p.detector_rot[1] = 0.018; + p.rot_vec[0] = 1.0; + p.rot_vec[1] = 0.0; + p.rot_vec[2] = 0.0; + double beta = 0; + LatticeToRodriguesLengthsBeta_Mono(CrystalLattice(39.7f, 50.6f, 79.6f, 90.0f, 94.5f, 90.0f), + p.latt_vec0, p.latt_vec1, beta); + p.latt_vec2[0] = beta; + + if (!reduced) { + p.detector_rot_constant = false; + p.rot_vec_constant = false; + p.latt_vec1_constant = false; + p.latt_vec2_constant = false; + for (int i = 0; i < 2; i++) { + p.detector_rot_lower[i] = p.detector_rot[i] - 0.05; + p.detector_rot_upper[i] = p.detector_rot[i] + 0.05; + } + for (int i = 0; i < 3; i++) { + p.latt_vec1_lower[i] = 5.0; + p.latt_vec1_upper[i] = 100.0; + } + p.latt_vec2_lower[0] = PI / 3; + p.latt_vec2_upper[0] = 2 * PI / 3; + p.priors.push_back({XtalRefinePrior::Block::Beam, 1.0, 0.0, p.beam[0], 0.5}); + p.priors.push_back({XtalRefinePrior::Block::DetectorRot, 0.0, 1.0, p.detector_rot[1], 50.0}); + } + p.options.max_iterations = 50; + return p; + } + + void CompareWithCeres(XtalRefineProblem p) { + XtalRefineProblem q = p; + const LMSummary lm = SolveXtalRefine(p, 4); + const ceres::Solver::Summary ref = SolveXtalRefineCeres(q, 4); + + REQUIRE(lm.IsSolutionUsable() == ref.IsSolutionUsable()); + CHECK(lm.iterations == static_cast(ref.iterations.size())); + CHECK(lm.final_cost == Catch::Approx(ref.final_cost).epsilon(1e-9)); + const auto same = [](const double *a, const double *b, int n, double tol) { + for (int i = 0; i < n; i++) + CHECK(a[i] == Catch::Approx(b[i]).margin(tol)); + }; + same(p.beam, q.beam, 2, 1e-7); + same(p.detector_rot, q.detector_rot, 2, 1e-10); + same(p.rot_vec, q.rot_vec, 3, 1e-10); + same(p.latt_vec0, q.latt_vec0, 3, 1e-10); + same(p.latt_vec1, q.latt_vec1, 3, 1e-8); + same(p.latt_vec2, q.latt_vec2, 3, 1e-10); + } +} + +TEST_CASE("XtalRefine_matches_Ceres_full", "[XtalOptimizer]") { + CompareWithCeres(RotationProblem(false, false)); +} + +TEST_CASE("XtalRefine_matches_Ceres_full_weighted", "[XtalOptimizer]") { + CompareWithCeres(RotationProblem(false, true)); +} + +TEST_CASE("XtalRefine_matches_Ceres_beam_orientation", "[XtalOptimizer]") { + CompareWithCeres(RotationProblem(true, false)); +} + +TEST_CASE("XtalRefine_same_answer_at_any_thread_count", "[XtalOptimizer]") { + XtalRefineProblem a = RotationProblem(false, false); + XtalRefineProblem b = a; + SolveXtalRefine(a, 1); + SolveXtalRefine(b, 7); + for (int i = 0; i < 3; i++) { + CHECK(a.latt_vec0[i] == b.latt_vec0[i]); + CHECK(a.latt_vec1[i] == b.latt_vec1[i]); + } + CHECK(a.beam[0] == b.beam[0]); + CHECK(a.beam[1] == b.beam[1]); +}