Jungfraujoch/image_analysis/geom_refinement/RingOptimizer.cpp

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

#include <algorithm>

#include "../../common/DetectorOrientation.h"
#include "../../common/JFJochMath.h"
#include "RingOptimizer.h"
#include "ceres/ceres.h"

struct RingResidual {
    RingResidual(double x, double y, double lambda,
                 double pixel_size,
                 double expected_q,
                 const DetectorOrientation &orientation)
            : obs_x(x), obs_y(y),
              lambda(lambda),
              pixel_size(pixel_size),
              expected_len_recip_sq(expected_q * expected_q / (4.0 * PI * PI)) {
        const RotMatrix delta = orientation.Matrix();
        det_m00 = delta.Column(0).x;
        det_m01 = delta.Column(1).x;
        det_m10 = delta.Column(0).y;
        det_m11 = delta.Column(1).y;
    }

    template<typename T>
    bool operator()(const T* const center_x, const T* const center_y,
                    const T* const distance, const T* const rot1,
                    const T* const rot2, T* residual) const {

        // Calculate lab coordinates from observed pixel coordinates
        T u_lab = (T(obs_x) - center_x[0]) * T(pixel_size); // convert to mm
        T v_lab = (T(obs_y) - center_y[0]) * T(pixel_size);
        // The discrete image orientation, which turns the offset from the PONI before the tilt acts.
        // Identity unless a detector says otherwise. It cannot change a ring's radius, but it does
        // change which way the tilt tips the ring, which is exactly what this fits.
        T x_lab = det_m00 * u_lab + det_m01 * v_lab;
        T y_lab = det_m10 * u_lab + det_m11 * v_lab;
        T z_lab = distance[0];

// Apply rotations around y and x axes
        T c1 = ceres::cos(rot1[0]);
        T c2 = ceres::cos(rot2[0]);
        T s1 = ceres::sin(rot1[0]);
        T s2 = ceres::sin(rot2[0]);

        T x = x_lab * c1 + z_lab * s1;
        T y = y_lab * c2 + (-x_lab * s1 + z_lab * c1) * s2;
        T z = -y_lab * s2 + (-x_lab * s1 + z_lab * c1) * c2;

        // convert to recip space
        T lab_norm = ceres::sqrt(x*x + y*y + z*z);

        T R_x = x / (lab_norm * T(lambda));
        T R_y = y / (lab_norm * T(lambda));
        T R_z = (z / lab_norm - T(1.0)) / T(lambda);

        T predicted_len_recip_sq = R_x * R_x + R_y * R_y + R_z * R_z;

        residual[0] = predicted_len_recip_sq - T(expected_len_recip_sq);
        return true;
    }

    const double obs_x, obs_y;
    const double lambda;
    const double pixel_size;
    const double expected_len_recip_sq;
    double det_m00, det_m01, det_m10, det_m11;
};

RingOptimizer::RingOptimizer(const DiffractionGeometry& geom, bool refine_tilt)
        : reference(geom), refine_tilt(refine_tilt) {}

DiffractionGeometry RingOptimizer::Run(const std::vector<RingOptimizerInput> &input) {
    // Initial guess for the parameters
    double center_x = reference.GetBeamX_pxl();
    double center_y = reference.GetBeamY_pxl();
    double distance = reference.GetDetectorDistance_mm();
    double rot1 = reference.GetPoniRot1_rad();
    double rot2 = reference.GetPoniRot2_rad();

    ceres::Problem problem;

    // Add residuals for each point
    for (const auto& pt : input) {
        problem.AddResidualBlock(
                new ceres::AutoDiffCostFunction<RingResidual, 1, 1, 1, 1, 1, 1>(
                        new RingResidual(pt.x, pt.y,
                                         reference.GetWavelength_A(),
                                         reference.GetPixelSize_mm(),
                                         pt.q_expected,
                                         reference.GetOrientation())),
                nullptr,
                &center_x,
                &center_y,
                &distance,
                &rot1,
                &rot2
        );
    }

    // A single ring cannot tell the beam centre from the detector tilt: both displace its radius as
    // cos(phi), and what separates them is only how that amplitude scales with the ring's radius, which
    // takes two rings. Hold the tilt where it was given, so the one thing a single ring does fix - where
    // its centre lies - comes out rather than being traded away against an unconstrained tilt.
    const bool single_ring = !input.empty()
            && std::all_of(input.begin(), input.end(), [&](const RingOptimizerInput &p) {
                   return p.q_expected == input.front().q_expected;
               });
    if (single_ring || !refine_tilt) {
        problem.SetParameterBlockConstant(&rot1);
        problem.SetParameterBlockConstant(&rot2);
    }

    // Configure solver
    ceres::Solver::Options options;
    options.linear_solver_type = ceres::DENSE_QR;
    options.minimizer_progress_to_stdout = false;
    options.logging_type = ceres::LoggingType::SILENT;
    options.num_threads = 1;

    ceres::Solver::Summary summary;

    // Run optimization
    ceres::Solve(options, &problem, &summary);

    // A failed fit must not move the detector geometry. Both callers assign the result straight
    // back over the geometry they passed in, so handing back the reference leaves the calibration
    // where it was instead of committing a diverged beam centre and distance.
    if (!summary.IsSolutionUsable())
        return reference;

    DiffractionGeometry refined_geom(reference);
    refined_geom.BeamX_pxl(center_x).BeamY_pxl(center_y).DetectorDistance_mm(distance)
                .PoniRot1_rad(rot1).PoniRot2_rad(rot2);

    return refined_geom;
}