Jungfraujoch/image_analysis/geom_refinement/XtalOptimizer.cpp

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "../../common/JFJochMath.h"
#include <algorithm>
#include <Eigen/Dense>

#include "XtalOptimizer.h"
#include "XtalResidual.h"
#include "ceres/ceres.h"
#include "ceres/rotation.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<typename T>
    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<typename T>
    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<T> rot(rot_aa);
        rot.Rotate(astar_unrot, astar_rot);
        rot.Rotate(bstar_unrot, bstar_rot);
        rot.Rotate(cstar_unrot, cstar_rot);

        const Eigen::Matrix<T, 3, 1> 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<typename T>
    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
// those are legitimately weaker, so a raw intensity weight would suppress exactly the wrong ones; the
// shell normalisation makes the weight resolution-neutral by construction.
//
// The weight enters as w^2 on the squared residual, w^2 = r/(1+r): the shell median contributes half,
// a 4x-median spot 0.8, a quarter-median spot 0.2. Weak spots still pull, they just do not drive. Unlike
// a robust loss this is a PRIOR - it never looks at the current residual, so it cannot mistake a genuine
// spot for an outlier when the starting geometry is far off and leave the fit unable to move.
static std::vector<double> SpotConfidenceWeights(const std::vector<SpotToSave> &spots) {
    constexpr size_t spots_per_shell = 32;

    // Resolution order. Sorting a packed (resolution, index) array rather than an index vector with a
    // projection into the spots keeps the comparisons off the 80-byte records - the same keys in the
    // same order, so introsort makes the same comparisons and the same swaps, and the order it leaves
    // is the same.
    struct SpotByRes {
        float d_A;
        uint32_t index;
    };
    std::vector<SpotByRes> by_res(spots.size());
    for (size_t i = 0; i < spots.size(); i++)
        by_res[i] = {spots[i].d_A, static_cast<uint32_t>(i)};
    std::ranges::sort(by_res, {}, &SpotByRes::d_A);

    const size_t nshells = std::max<size_t>(1, spots.size() / spots_per_shell);
    std::vector<double> weight(spots.size());
    std::vector<float> shell_intensity;

    for (size_t s = 0; s < nshells; s++) {
        const size_t begin = s * spots.size() / nshells;
        const size_t end = (s + 1) * spots.size() / nshells;

        shell_intensity.clear();
        for (size_t i = begin; i < end; i++)
            shell_intensity.push_back(spots[by_res[i].index].intensity);
        std::ranges::nth_element(shell_intensity, shell_intensity.begin() + shell_intensity.size() / 2);
        const double median = std::max(1e-3f, shell_intensity[shell_intensity.size() / 2]);

        for (size_t i = begin; i < end; i++) {
            const double r = std::max(0.0f, spots[by_res[i].index].intensity) / median;
            weight[by_res[i].index] = std::sqrt(r / (1.0 + r));
        }
    }
    return weight;
}

// The oscillation width at which the acceptance gate starts profiling out the rotation coordinate.
// The BAND is principled: the dead zone below matters once the exposure's own rms rotation ambiguity,
// wedge/sqrt(12) = 0.29*wedge, is comparable to the crystal's intrinsic along-u rocking spread, which
// measures ~0.26 deg, and that puts the boundary somewhere between 0.25 and 1.0 deg. The POINT is
// empirical and is taken at the conservative end of that band, because fine slicing is the core case
// and coarse slicing is compatibility: below this the gate is left exactly as it was.
constexpr float COARSE_SLICING_WEDGE_DEG = 0.5f;

// The dead zone's half-width as a fraction of the exposure: the rms of a rotation coordinate uniform
// over the frame, which is the width a least-squares is calibrated on. Half the exposure - the worst
// case a spot could sit at - and forgiving the direction outright were both measured worse.
const double DEAD_ZONE_K = 1.0 / std::sqrt(12.0);

bool XtalOptimizerInternal(XtalOptimizerData &data,
                           std::span<const std::vector<SpotToSave>> spots,
                           const std::vector<std::vector<double>> &weights,
                           const float tolerance,
                           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.
        if (data.latt.VolumeFraction() < MIN_BASIS_VOLUME_FRACTION)
            return false;

        Coord vec0 = data.latt.Vec0();
        Coord vec1 = data.latt.Vec1();
        Coord vec2 = data.latt.Vec2();
        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();

        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<XtalFrameConstants> 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};

        switch (data.crystal_system) {
            case gemmi::CrystalSystem::Orthorhombic:
                LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
                break;
            case gemmi::CrystalSystem::Tetragonal:
                LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
                latt_vec1[0] = (latt_vec1[0] + latt_vec1[1]) / 2.0;
                break;
            case gemmi::CrystalSystem::Cubic:
                LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
                latt_vec1[0] = (latt_vec1[0] + latt_vec1[1] + latt_vec1[2]) / 3.0;
                break;
            case gemmi::CrystalSystem::Hexagonal:
                LatticeToRodriguesAndLengths_Hex(data.latt, latt_vec0, latt_vec1);
                break;
            case gemmi::CrystalSystem::Monoclinic:
                LatticeToRodriguesLengthsBeta_Mono(data.latt, latt_vec0, latt_vec1, beta);
                latt_vec2[0] = beta;
                latt_vec2[1] = 0.0;
                latt_vec2[2] = 0.0;
                break;
            default:
                // Triclinic: initialize a,b,c and α,β,γ from current unit cell
                LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
                auto uc = data.latt.GetUnitCell();
                latt_vec2[0] = uc.alpha * PI / 180.0;
                latt_vec2[1] = uc.beta * PI / 180.0;
                latt_vec2[2] = uc.gamma * PI / 180.0;
                break;
        }

        // The spindle. `rocking_spindle` is the fallback for a caller that holds one frame and so
        // passes no `axis` to back-rotate by: the back-rotation is the identity there either way
        // (angle_rad is zero and an AngleAxisRotator of a zero angle-axis ignores the vector, so the
        // block is also held constant), but leaving the {1,0,0} initialiser standing would hand any
        // later reader of this vector the LAB X AXIS in place of the spindle.
        if (const auto spindle = data.axis ? std::optional(data.axis->GetAxis()) : data.rocking_spindle) {
            rot_vec[0] = spindle->x;
            rot_vec[1] = spindle->y;
            rot_vec[2] = spindle->z;
        }

        // The exposure this refinement's spots are spread over, and the spindle they are spread
        // along. Taken from the explicit rocking fields where the caller set them - the per-frame
        // refinement, which does not back-rotate but whose spots still span an exposure - and
        // otherwise from the axis this call does back-rotate by.
        const float rocking_wedge_deg = data.rocking_wedge_deg > 0.0f
                                            ? data.rocking_wedge_deg
                                            : ((data.axis && data.axis->IsScanning())
                                                   ? data.axis->GetWedge_deg() : 0.0f);
        const Coord rocking_spindle = data.rocking_spindle.value_or(
            data.axis ? data.axis->GetAxis() : Coord());
        // Zero everywhere below the trigger, which switches the dead zone off and leaves the gate
        // computing the plain fractional-index miss.
        const double dead_zone_rad = rocking_wedge_deg >= COARSE_SLICING_WEDGE_DEG
                                         ? rocking_wedge_deg * PI / 180.0 * DEAD_ZONE_K
                                         : 0.0;

        const float tolerance_sq = tolerance * tolerance;

        // The same for every spot of every frame, so taken once here rather than per residual.
        const double cos_rot3 = std::cos(data.geom.GetPoniRot3_rad());
        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.
        const bool beam_and_orientation_only = data.refine_beam_center
                                               && !data.refine_detector_angles
                                               && !data.refine_rotation_axis
                                               && !data.refine_unit_cell;

        // 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
        // when the spots are unweighted.
        double effective_spots = 0.0;

        for (int i = 0; i < spots.size(); i++) {
            if (spots[i].empty())
                continue;

            const std::vector<double> &weight = weights[i]; // empty = unweighted

            double angle_rad = 0.0;
            std::optional<RotMatrix> rot_matr;

            if (data.axis) {
                const float angle_deg = data.axis->GetAngle_deg(i) + data.axis->GetWedge_deg() / 2.0;
                angle_rad = angle_deg * PI / 180.0;
                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);

            // Add residuals for each point
            for (size_t j = 0; j < spots[i].size(); j++) {
                const auto &pt = spots[i][j];
                if (!data.index_ice_rings && pt.ice_ring)
                    continue;

                Coord recip = pt.ReciprocalCoord(data.geom);

                if (rot_matr)
                    recip = rot_matr.value() * recip;

                double h_fp = recip * vec0;
                double k_fp = recip * vec1;
                double l_fp = recip * vec2;

                double h = std::round(h_fp);
                double k = std::round(k_fp);
                double l = std::round(l_fp);

                double norm_sq = (h - h_fp) * (h - h_fp) + (k - k_fp) * (k - k_fp) + (l - l_fp) * (l - l_fp);

                // At coarse slicing the spot diffracted somewhere inside the exposure, not at its
                // midpoint, and that unknown angle is a real part of the miss. Charge only the part
                // of it the exposure cannot supply: a rotation delta about the spindle moves the
                // fractional index along u = m x q, so the component of the miss along u is free up
                // to the exposure's rms half-width and only the excess counts. Every other direction
                // is untouched - |q| among them, so every d-spacing is unaffected. Without this the
                // gate is a resolution cut that tightens with the frame width, since the miss grows
                // as a/d.
                if (dead_zone_rad > 0.0) {
                    const Coord u = rocking_spindle % recip;
                    const double u0 = u * vec0, u1 = u * vec1, u2 = u * vec2;
                    const double u_sq = u0 * u0 + u1 * u1 + u2 * u2;
                    if (u_sq > 1e-24) {
                        const double inv_u = 1.0 / std::sqrt(u_sq);
                        const double d_par = ((h - h_fp) * u0 + (k - k_fp) * u1 + (l - l_fp) * u2) * inv_u;
                        const double dead = dead_zone_rad * std::sqrt(u_sq);
                        const double excess = std::max(0.0, std::fabs(d_par) - dead);
                        norm_sq = std::max(0.0, norm_sq - d_par * d_par) + excess * excess;
                    }
                }

                if (norm_sq > tolerance_sq)
                    continue;

                const double weight_sq = weight.empty() ? 1.0 : weight[j] * weight[j];
                effective_spots += weight_sq;

                const XtalResidual residual(pt.x, pt.y,
                                            data.geom.GetWavelength_A(),
                                            data.geom.GetPixelSize_mm(),
                                            cos_rot3, sin_rot3,
                                            angle_rad,
                                            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<XtalResidualBeamOrientation, 3, 2, 3>(
                            new XtalResidualBeamOrientation(residual, distance_mm, frame_const.back())),
                        loss,
                        beam,
                        latt_vec0
                    );
                else
                    problem.AddResidualBlock(
                        new ceres::AutoDiffCostFunction<XtalResidualFixedDistance, 3, 2, 2, 3, 3, 3, 3>(
                            new XtalResidualFixedDistance(residual, distance_mm)),
                        loss,
                        beam,
                        detector_rot,
                        rot_vec,
                        latt_vec0,
                        latt_vec1,
                        latt_vec2
                    );
            }
        }

        if (problem.NumResidualBlocks() < data.min_spots)
            return false;

        // The gauge direction of a single-axis rotation experiment - parallel to the spindle - written
        // once, for both of the parameter pairs it applies to. The two need it in DIFFERENT frames and
        // that is the whole difficulty:
        //
        //   beam[0]/beam[1] are PIXEL columns and rows. The pixel axes reach the laboratory through
        //   det_matrix = PoniRotMatrix * DetectorOrientation::Matrix(), so on a quarter turn of 1 or 3
        //   the pixel X axis IS the laboratory Y axis. Comparing the goniometer vector's laboratory
        //   components against a beam index is therefore only right when that orientation is the
        //   identity; elsewhere it pins the determined component and frees the gauge one. Project the
        //   spindle onto the pixel axes' own laboratory images instead - exact for any orientation,
        //   any tilt and a spindle at any angle, and equal to picking the dominant component when the
        //   orientation is the identity and the spindle lies along a detector axis.
        //
        //   detector_rot[0]/[1] are rotations about the LABORATORY y and x axes (see PoniRotMatrix),
        //   applied outside that orientation matrix, and they move the direct beam along laboratory x
        //   and y respectively by D/pixel per radian. So the tilt's gauge combination is the spindle's
        //   own laboratory x and y components, with no orientation in it.
        //
        // Same spindle, same physical direction, each in the frame its parameters live in.
        double gauge_beam_x = 0.0, gauge_beam_y = 0.0;
        double gauge_rot_x = 0.0, gauge_rot_y = 0.0;
        if (data.axis) {
            const Coord spindle = data.axis->GetAxis().Normalize();
            const Coord fast = data.geom.GetFastAxis();
            const Coord slow = data.geom.GetSlowAxis();
            const double bx = spindle * fast, by = spindle * slow;
            const double bn = std::hypot(bx, by);
            if (bn > 0.0) {
                gauge_beam_x = bx / bn;
                gauge_beam_y = by / bn;
            }
            const double rn = std::hypot(spindle.x, spindle.y);
            if (rn > 0.0) {
                gauge_rot_x = spindle.x / rn;
                gauge_rot_y = spindle.y / rn;
            }
        }
        // Weight so a gauge prior is a sigma_px-pixel restraint that competes with the positional
        // residuals. k = d|recip|/d(beam_px) ~ pixel/(distance*lambda) [A^-1/px]; scaling by
        // sqrt(#residuals) makes the prior's curvature ~ (1/9) of the well-constrained-data curvature
        // at sigma_px=3, i.e. data wins the perpendicular direction, the prior wins the gauge one.
        // Note what that scaling means: the prior's curvature grows with the number of spots exactly
        // as the data's does, so the split it picks between two aliased parameters is the same however
        // much data the stage has. More frames, a longer sweep or a later stage cannot break it.
        constexpr double sigma_px = 3.0;
        // The tilt's budgets, in those same direct-beam pixels: one for the spindle-parallel
        // combination and one for the perpendicular one. Zero means no restraint at all, so which
        // component is held and which is refined is these two numbers and nothing else.
        //
        // The parallel one is TIGHTER than the beam's on purpose: the data determine the SUM of the
        // two, so with equal budgets the shift splits evenly and half of a beam-centre error still
        // arrives as an angle (measured: the coupling to the starting beam centre falls only from
        // 79% to 41% of one-for-one at equal budgets, and to 8% at this one). The detector tilt is a
        // property of the mounting, re-measured when the detector is calibrated; the beam centre
        // drifts between runs. When both ends of an alias have to be restrained, the tighter
        // restraint belongs on the one that moves less.
        //
        // The perpendicular one is free. That is the arrangement the data support today: it is the
        // component whose conditioning tracks the 2theta the fit reaches, i.e. the one the data speak
        // about, while the parallel one's does not move with 2theta at all.
        constexpr double SIGMA_TILT_PARALLEL_PX = 1.0;
        constexpr double SIGMA_TILT_PERPENDICULAR_PX = 0.0;
        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) {
            // 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
            // beam that the co-refined orientation keeps position-consistent. Rather than freeze it (the beam
            // 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<GaugeDirectionPrior, 1, 2>(
                    new GaugeDirectionPrior(gauge_beam_x, gauge_beam_y,
                                            gauge_beam_x * beam[0] + gauge_beam_y * beam[1], gauge_w)),
                nullptr, beam);
        }

        // 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 {
                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);
                }
                // 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
                // its gauge combination free lets a beam-centre error the prior refuses to absorb
                // reappear as an angle - measured at 0.072 deg per pixel of the STARTING beam centre,
                // against a geometric one-for-one of 0.080, while the refined beam never leaves its
                // anchor by more than a quarter of a pixel.
                //
                // Restraining it does not make the tilt a measurement, and nothing here should be read
                // that way. In THIS fit the restrained component carries no information of its own:
                // the crystal orientation is refined alongside it and absorbs the difference, so it
                // ends up as accurate as the file's beam centre and no more. The free component does
                // carry information, and is separately known to sit ~0.06 deg from a powder
                // calibration on one measured detector, which is many times its formal error - so a
                // single crystal's tilt is not a number to feed back into a file. What this buys is
                // that a beam-centre error is no longer laundered into a reported angle.
                //
                // "In this fit" is the load-bearing part: a later stage that FREEZES the orientation
                // has no such compensator, and whether the parallel component is measurable there is a
                // different question with a different answer. This restraint is local to the fit that
                // co-refines the orientation and does not speak for any other.
                if (data.axis) {
                    const double lever = distance_mm / data.geom.GetPixelSize_mm();
                    // Parallel first, then the perpendicular direction (-gy, gx). Both go through the
                    // same restraint, so swapping which one is held is a change to the two budgets.
                    const double dirs[2][2] = {{gauge_rot_x, gauge_rot_y}, {-gauge_rot_y, gauge_rot_x}};
                    const double budget[2] = {SIGMA_TILT_PARALLEL_PX, SIGMA_TILT_PERPENDICULAR_PX};
                    for (int i = 0; i < 2; ++i) {
                        if (budget[i] <= 0.0)
                            continue;
                        problem.AddResidualBlock(
                            new ceres::AutoDiffCostFunction<GaugeDirectionPrior, 1, 2>(
                                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);
                    }
                }
            }

            if (!data.refine_rotation_axis) {
                problem.SetParameterBlockConstant(rot_vec);
            } else {
                // 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
                // RotationIndexer::accumulate() under a rotation model the fit did not use. Measured
                // over the corpus, that length reached 1.2 % and the fit/score disagreement a whole
                // degree of goniometer angle. It is not a usable measurement either: on synthetic
                // 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>);
            }

            if (!data.refine_unit_cell) {
                problem.SetParameterBlockConstant(latt_vec1);
                problem.SetParameterBlockConstant(latt_vec2);
            } else {
                // 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);
                }

                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);
                } 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);
                    }
                } else {
                    // Orthorhombic / Tetragonal / Cubic / Hexagonal:
                    // latt_vec2 has no meaning for these systems — always freeze it.
                    problem.SetParameterBlockConstant(latt_vec2);
                }
            }
        }

        // 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;
        if (data.max_iterations > 0)
            options.max_num_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);

        // 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
        // relies on when it sets max_solver_time. Checked before anything is written back, so a
        // failed refinement leaves data untouched rather than committing half a fit.
        if (!summary.IsSolutionUsable())
            return false;

        if (data.refine_beam_center) {
            data.beam_corr_x = data.geom.GetBeamX_pxl() - beam[0];
            data.beam_corr_y = data.geom.GetBeamY_pxl() - beam[1];
            data.geom.BeamX_pxl(beam[0]).BeamY_pxl(beam[1]);
        }


        if (data.refine_detector_angles)
            data.geom.PoniRot1_rad(detector_rot[0]).PoniRot2_rad(detector_rot[1]);

        if (data.axis && data.refine_rotation_axis)
            data.axis.value().Axis(Coord(rot_vec[0], rot_vec[1], rot_vec[2]));

        if (data.crystal_system == gemmi::CrystalSystem::Orthorhombic)
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, PI / 2.0, PI / 2.0);
        else if (data.crystal_system == gemmi::CrystalSystem::Tetragonal) {
            latt_vec1[1] = latt_vec1[0];
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, PI / 2.0, PI / 2.0);
        } else if (data.crystal_system == gemmi::CrystalSystem::Cubic) {
            latt_vec1[1] = latt_vec1[0];
            latt_vec1[2] = latt_vec1[0];
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, PI / 2.0, PI / 2.0);
        } else if (data.crystal_system == gemmi::CrystalSystem::Hexagonal) {
            latt_vec1[1] = latt_vec1[0];
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1,PI / 2.0, PI / 2.0, 2.0 * PI / 3.0);
        } else if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) {
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, PI / 2.0, latt_vec2[0], PI / 2.0);
        } else {
            // Triclinic via the same generic builder
            data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, latt_vec2[0], latt_vec2[1], latt_vec2[2]);
        }
        return true;
    } catch (...) {
        // Convergence problems, likely not updated
        return false;
    }
}

bool XtalOptimizer(XtalOptimizerData &data, std::span<const std::vector<SpotToSave>> spots,
                   int num_threads) {
    // A spot's confidence weight is set by its resolution and its intensity, neither of which the solver
    // touches, so the three passes below all get the same weights: take them once.
    std::vector<std::vector<double>> weights(spots.size());
    if (data.weight_spots_by_confidence)
        for (size_t i = 0; i < spots.size(); i++)
            if (!spots[i].empty())
                weights[i] = SpotConfidenceWeights(spots[i]);

    if (!XtalOptimizerInternal(data, spots, weights, 0.3, num_threads))
        return false;
    XtalOptimizerInternal(data, spots, weights, 0.2, num_threads);
    return XtalOptimizerInternal(data, spots, weights, 0.1, num_threads);
}

bool XtalOptimizer(XtalOptimizerData &data, const std::vector<SpotToSave> &spots, int num_threads) {
    return XtalOptimizer(data, std::span(&spots, 1), num_threads);
}

bool XtalOptimizerRotationOnly(XtalOptimizerData &data,
                               const std::vector<SpotToSave> &spots,
                               const float tolerance) {
    try {
        // Same refusal as XtalOptimizerInternal: the residual here is built from Astar/Bstar/Cstar,
        // which divide by the cell volume, so a coplanar basis makes every one of them infinite.
        if (data.latt.VolumeFraction() < MIN_BASIS_VOLUME_FRACTION)
            return false;

        // Parameter: angle-axis for the extra rotation. Identity == {0,0,0}.
        double rot_aa[3] = {0.0, 0.0, 0.0};

        // Spot selection by current indexing (same approach as XtalOptimizerInternal)
        const Coord a0 = data.latt.Vec0();
        const Coord b0 = data.latt.Vec1();
        const Coord c0 = data.latt.Vec2();

        const float tol_sq = tolerance * tolerance;

        ceres::Problem problem;

        for (const auto &pt : spots) {
            if (!data.index_ice_rings && pt.ice_ring)
                continue;

            // Compute fractional HKL using the CURRENT lattice
            Coord recip_index = pt.ReciprocalCoord(data.geom);
            if (data.axis.has_value())
                recip_index = data.axis->GetTransformationAngle(pt.phi) * recip_index;

            const double h_fp = static_cast<double>(recip_index * a0);
            const double k_fp = static_cast<double>(recip_index * b0);
            const double l_fp = static_cast<double>(recip_index * c0);

            const double h = std::round(h_fp);
            const double k = std::round(k_fp);
            const double l = std::round(l_fp);

            const double norm_sq =
                (h - h_fp) * (h - h_fp) +
                (k - k_fp) * (k - k_fp) +
                (l - l_fp) * (l - l_fp);

            if (norm_sq > static_cast<double>(tol_sq))
                continue;

            // s_obs must be in the same reference frame as the
            // predicted reciprocal vector (h·a* + k·b* + l·c*), which is the
            // phi=0 crystal frame.  Apply the same goniometer back-rotation
            // that was used above for the HKL assignment.
            Coord s_obs = data.geom.DetectorToRecip(pt.x, pt.y);
            if (data.axis.has_value())
                s_obs = data.axis->GetTransformationAngle(pt.phi) * s_obs;

            auto *cost =
                new ceres::AutoDiffCostFunction<XtalResidualRotationOnlyPrecomp, 3, 3>(
                    new XtalResidualRotationOnlyPrecomp(s_obs, data.latt, h, k, l)
                );

            problem.AddResidualBlock(cost, nullptr, rot_aa);
        }

        if (problem.NumResidualBlocks() < 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<RotationNormRegularizer, 3, 3>(
                    new RotationNormRegularizer(reg_weight)),
                nullptr, rot_aa);
        }

        ceres::Solver::Options options;
        options.linear_solver_type = ceres::DENSE_NORMAL_CHOLESKY;   // tall and thin, as above
        options.minimizer_progress_to_stdout = false;
        if (data.max_iterations > 0)
            options.max_num_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);

        if (!summary.IsSolutionUsable())
            return false;

        // Apply rotation to direct-lattice vectors.
        // ceres::AngleAxisToRotationMatrix writes a **row-major** 3×3 matrix,
        // and Eigen's << operator also fills row-by-row, so the assignment
        // below is correct without any transposing.
        //
        // Note: for a pure orthogonal rotation R, R⁻ᵀ = R, so rotating the
        // direct-lattice vectors (A, B, C) by R is exactly equivalent to
        // 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

        Eigen::Matrix3d R;
        R << R_raw[0], R_raw[3], R_raw[6],
             R_raw[1], R_raw[4], R_raw[7],
             R_raw[2], R_raw[5], R_raw[8];

        const Eigen::Vector3d A(a0.x, a0.y, a0.z);
        const Eigen::Vector3d B(b0.x, b0.y, b0.z);
        const Eigen::Vector3d C(c0.x, c0.y, c0.z);

        const Eigen::Vector3d A2 = R * A;
        const Eigen::Vector3d B2 = R * B;
        const Eigen::Vector3d C2 = R * C;

        data.latt = CrystalLattice(
            Coord(static_cast<float>(A2.x()), static_cast<float>(A2.y()), static_cast<float>(A2.z())),
            Coord(static_cast<float>(B2.x()), static_cast<float>(B2.y()), static_cast<float>(B2.z())),
            Coord(static_cast<float>(C2.x()), static_cast<float>(C2.y()), static_cast<float>(C2.z()))
        );

        double theta = std::sqrt(rot_aa[0] * rot_aa[0] + rot_aa[1] * rot_aa[1] + rot_aa[2] * rot_aa[2]);
        data.angle_corr = theta;
        if (theta > 1e-6) {
            Coord rot;
            rot.x = rot_aa[0] / theta;
            rot.y = rot_aa[1] / theta;
            rot.z = rot_aa[2] / theta;
            data.angle_axis = rot;
        } else
            data.angle_axis.reset();

        return true;
    } catch (...) {
        return false;
    }
}