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 <numeric>
#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.
struct BeamComponentPrior {
    BeamComponentPrior(int component, double b0, double weight)
        : component(component), b0(b0), weight(weight) {}
    template<typename T>
    bool operator()(const T *const beam, T *residual) const {
        residual[0] = T(weight) * (beam[component] - T(b0));
        return true;
    }
    int component;
    double b0, 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];

        ceres::AngleAxisRotatePoint(rot_aa, astar_unrot, astar_rot);
        ceres::AngleAxisRotatePoint(rot_aa, bstar_unrot, bstar_rot);
        ceres::AngleAxisRotatePoint(rot_aa, 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;

    std::vector<size_t> by_res(spots.size());
    std::iota(by_res.begin(), by_res.end(), 0);
    std::ranges::sort(by_res, {}, [&](size_t i) { return spots[i].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]].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]].intensity) / median;
            weight[by_res[i]] = std::sqrt(r / (1.0 + r));
        }
    }
    return weight;
}

bool XtalOptimizerInternal(XtalOptimizerData &data,
                           const std::vector<std::vector<SpotToSave>> &spots,
                           const float tolerance,
                           const int num_threads) {
    try {
        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()};

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

        if (data.axis) {
            rot_vec[0] = data.axis->GetAxis().x;
            rot_vec[1] = data.axis->GetAxis().y;
            rot_vec[2] = data.axis->GetAxis().z;
        }

        const float tolerance_sq = tolerance * tolerance;

        // 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;

            std::vector<double> weight; // empty = unweighted
            if (data.weight_spots_by_confidence)
                weight = SpotConfidenceWeights(spots[i]);

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

            // 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);

                if (norm_sq > tolerance_sq)
                    continue;

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

                problem.AddResidualBlock(
                    new ceres::AutoDiffCostFunction<XtalResidual, 3, 2, 1, 2, 3, 3, 3, 3>(
                        new XtalResidual(pt.x, pt.y,
                                         data.geom.GetWavelength_A(),
                                         data.geom.GetPixelSize_mm(),
                                         data.geom.GetPoniRot3_rad(),
                                         angle_rad,
                                         h, k, l,
                                         data.crystal_system)),
                    // 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.
                    weight.empty()
                        ? nullptr
                        : new ceres::ScaledLoss(nullptr, weight_sq, ceres::TAKE_OWNERSHIP),
                    beam,
                    &distance_mm,
                    detector_rot,
                    rot_vec,
                    latt_vec0,
                    latt_vec1,
                    latt_vec2
                );
            }
        }

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

        if (!data.refine_distance_mm)
            problem.SetParameterBlockConstant(&distance_mm);
        else {
            const double dist_range = 0.1;
            problem.SetParameterLowerBound(&distance_mm, 0, distance_mm * (1.0 - dist_range));
            problem.SetParameterUpperBound(&distance_mm, 0, distance_mm * (1.0 + dist_range));
        }

        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. The spindle is along a detector axis in standard
            // geometry, so restrain the dominant of X / Y.
            const Coord spindle = data.axis->GetAxis();
            const int parallel = (std::fabs(spindle.x) >= std::fabs(spindle.y)) ? 0 : 1;
            // Weight so the 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.
            constexpr double sigma_px = 3.0;
            const double k = data.geom.GetPixelSize_mm() / (distance_mm * data.geom.GetWavelength_A());
            const double w = k * std::sqrt(effective_spots) / sigma_px;
            problem.AddResidualBlock(
                new ceres::AutoDiffCostFunction<BeamComponentPrior, 1, 2>(
                    new BeamComponentPrior(parallel, beam[parallel], w)),
                nullptr, beam);
        }

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

        if (!data.refine_rotation_axis) {
            problem.SetParameterBlockConstant(rot_vec);
        }

        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;
        options.linear_solver_type = ceres::DENSE_QR;
        options.minimizer_progress_to_stdout = false;
        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);

        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_distance_mm)
            data.geom.DetectorDistance_mm(distance_mm);

        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, const std::vector<std::vector<SpotToSave>> &spots,
                   int num_threads) {
    if (!XtalOptimizerInternal(data, spots, 0.3, num_threads))
        return false;
    XtalOptimizerInternal(data, spots, 0.2, num_threads);
    return XtalOptimizerInternal(data, spots, 0.1, num_threads);
}

bool XtalOptimizerRotationOnly(XtalOptimizerData &data,
                               const std::vector<SpotToSave> &spots,
                               const float tolerance) {
    try {
        // 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_QR;
        options.minimizer_progress_to_stdout = false;
        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);

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