Jungfraujoch/image_analysis/geom_refinement/XtalResidual.h

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

#pragma once

#include <cmath>

#include <Eigen/Dense>

#include "ceres/ceres.h"
#include "ceres/rotation.h"
#include "gemmi/symmetry.hpp"

#include "../../common/JFJochException.h"

// 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),
// p0 = lattice orientation angle-axis(3), p1 = unit-cell lengths(3), p2 = unit-cell angles(3, rad).
// The residual is the difference between the observed reciprocal vector (from the detector position via
// the current beam/distance/tilt) and the predicted one (h*a* + k*b* + l*c* for the current cell +
// orientation), in Angstrom^-1. rot3 and the goniometer angle_rad are baked in as constants.
struct XtalResidual {
    XtalResidual(double x, double y,
                 double lambda,
                 double pixel_size,
                 double rot3,
                 double angle_rad,
                 double exp_h, double exp_k,
                 double exp_l,
                 gemmi::CrystalSystem symmetry)
        : obs_x(x), obs_y(y),
          inv_lambda(1.0/lambda),
          pixel_size(pixel_size),
          rot3(rot3),
          exp_h(exp_h),
          exp_k(exp_k),
          exp_l(exp_l),
          angle_rad(angle_rad),
          symmetry(symmetry) {
        if (std::fabs(lambda) < 1e-6)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                "Lambda cannot be close to zero");
    }

    template<typename T>
    bool operator()(const T *const beam,
                    const T *const distance_mm,
                    const T *const detector_rot,
                    const T *const rotation_axis,
                    const T *const p0,
                    const T *const p1,
                    const T *const p2,
                    T *residual) const {
        // PyFAI convention: poni_rot = Rz(-rot3) * Rx(-rot2) * Ry(+rot1).
        // 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.
        const T rot1 = detector_rot[0];
        const T rot2 = detector_rot[1];

        // Ry(+rot1): rotation around Y-axis
        const T c1 = ceres::cos(rot1);
        const T s1 = ceres::sin(rot1);

        // Rx(-rot2): rotation around X-axis with inverted sign (PyFAI left-handed)
        const T c2 = ceres::cos(rot2);
        const T s2 = ceres::sin(rot2);

        // Rz(-rot3): rotation around Z (beam); constant, identity when rot3 == 0
        const T c3 = T(cos(rot3));
        const T s3 = T(sin(rot3));

        // Detector coordinates in mm
        const T det_x = (T(obs_x) - beam[0]) * T(pixel_size);
        const T det_y = (T(obs_y) - beam[1]) * T(pixel_size);
        const T det_z = T(distance_mm[0]);

        // Apply Ry(rot1) first: rotate around Y
        const T t1_x = c1 * det_x + s1 * det_z;
        const T t1_y = det_y;
        const T t1_z = -s1 * det_x + c1 * det_z;

        // Then apply Rx(-rot2): rotate around X
        const T t2_x = t1_x;
        const T t2_y = c2 * t1_y + s2 * t1_z;
        const T t2_z = -s2 * t1_y + c2 * t1_z;

        // Then apply Rz(-rot3): rotate around Z (beam)
        const T x = c3 * t2_x + s3 * t2_y;
        const T y = -s3 * t2_x + c3 * t2_y;
        const T z = t2_z;

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

        T recip_raw[3];
        recip_raw[0] = x * inv_norm * T(inv_lambda);
        recip_raw[1] = y * inv_norm * T(inv_lambda);
        recip_raw[2] = (z * inv_norm - T(1.0)) * T(inv_lambda);

        // Apply goniometer "back-to-start" rotation:
        // brings observed reciprocal from image orientation into reference crystal frame
        const T aa_back[3] = {
            T(angle_rad) * rotation_axis[0],
            T(angle_rad) * rotation_axis[1],
            T(angle_rad) * rotation_axis[2]
        };

        T recip_obs[3];
        ceres::AngleAxisRotatePoint(aa_back, recip_raw, recip_obs);

        const Eigen::Matrix<T, 3, 1> e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]);

        // Build unit cell lengths and B (convention: columns are a, b, c prior to global rotation)
        Eigen::Matrix<T, 3, 1> e_uc_len = Eigen::Matrix<T, 3, 1>::Zero();
        Eigen::Matrix<T, 3, 3> B = Eigen::Matrix<T, 3, 3>::Identity();

        if (symmetry == gemmi::CrystalSystem::Hexagonal) {
            e_uc_len << p1[0], p1[0], p1[2];
            B(0, 1) = T(-0.5); // cos(120)
            B(1, 1) = T(sqrt(3.0) / 2.0); // sin(120)
        } else if (symmetry == gemmi::CrystalSystem::Orthorhombic) {
            e_uc_len << p1[0], p1[1], p1[2];
        } else if (symmetry == gemmi::CrystalSystem::Tetragonal) {
            e_uc_len << p1[0], p1[0], p1[2];
        } else if (symmetry == gemmi::CrystalSystem::Cubic) {
            e_uc_len << p1[0], p1[0], p1[0];
        } 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]);
        } else {
            // Triclinic: p1 = (a,b,c), p2 = (alpha, beta, gamma) in radians
            const T ca = ceres::cos(p2[0]);
            const T cb = ceres::cos(p2[1]);
            const T cg = ceres::cos(p2[2]);
            const T sg = ceres::sin(p2[2]);

            e_uc_len << p1[0], p1[1], p1[2];

            B(0, 0) = T(1);
            B(1, 0) = T(0);
            B(2, 0) = T(0);
            B(0, 1) = cg;
            B(1, 1) = sg;
            B(2, 1) = T(0);

            // c vector components:
            const T cx = cb;
            const T cy = (ca - cb * cg) / sg;
            const T v = T(1) - cx * cx - cy * cy;
            const T cz = (v >= T(0)) ? ceres::sqrt(v) : T(0);

            B(0, 2) = cx;
            B(1, 2) = cy;
            B(2, 2) = cz;
        }

        // Build unrotated direct lattice columns: (B * D), then rotate them by p0.
        // This avoids AngleAxisToRotationMatrix + matrix multiplications.
        const T L0 = e_uc_len[0];
        const T L1 = e_uc_len[1];
        const T L2 = e_uc_len[2];

        T col0_unrot[3] = {B(0, 0) * L0, B(1, 0) * L0, B(2, 0) * L0};
        T col1_unrot[3] = {B(0, 1) * L1, B(1, 1) * L1, B(2, 1) * L1};
        T col2_unrot[3] = {B(0, 2) * L2, B(1, 2) * L2, B(2, 2) * L2};

        T col0_rot[3], col1_rot[3], col2_rot[3];
        ceres::AngleAxisRotatePoint(p0, col0_unrot, col0_rot);
        ceres::AngleAxisRotatePoint(p0, col1_unrot, col1_rot);
        ceres::AngleAxisRotatePoint(p0, col2_unrot, col2_rot);

        const Eigen::Matrix<T, 3, 1> A(col0_rot[0], col0_rot[1], col0_rot[2]);
        const Eigen::Matrix<T, 3, 1> Bv(col1_rot[0], col1_rot[1], col1_rot[2]);
        const Eigen::Matrix<T, 3, 1> C(col2_rot[0], col2_rot[1], col2_rot[2]);

        const Eigen::Matrix<T, 3, 1> BxC = Bv.cross(C);
        const Eigen::Matrix<T, 3, 1> CxA = C.cross(A);
        const Eigen::Matrix<T, 3, 1> AxB = A.cross(Bv);

        const T V = A.dot(BxC);
        const T invV = T(1) / V;

        const Eigen::Matrix<T, 3, 1> Astar = BxC * invV;
        const Eigen::Matrix<T, 3, 1> Bstar = CxA * invV;
        const Eigen::Matrix<T, 3, 1> Cstar = AxB * invV;

        const T h = T(exp_h);
        const T k = T(exp_k);
        const T l = T(exp_l);

        const Eigen::Matrix<T, 3, 1> e_pred_recip = Astar * h + Bstar * k + Cstar * l;

        residual[0] = e_obs_recip[0] - e_pred_recip[0];
        residual[1] = e_obs_recip[1] - e_pred_recip[1];
        residual[2] = e_obs_recip[2] - e_pred_recip[2];

        return true;
    }

    const double obs_x, obs_y;
    const double inv_lambda;
    const double pixel_size;
    const double rot3;
    const double exp_h;
    const double exp_k;
    const double exp_l;
    const double angle_rad;
    gemmi::CrystalSystem symmetry;
};

// Same residual with distance, detector angles, rotation axis and unit cell baked in as constants, so
// that only beam(2) and orientation(3) remain parameter blocks. Ceres sizes its autodiff dual numbers
// from the DECLARED blocks, not from which of them the caller then holds constant, so the seven-block
// form above differentiates 17 parameters even when 5 are free; this one runs on Jet<double, 5>. It
// forwards to XtalResidual with the very same values, so the residual and the beam/orientation columns
// of its Jacobian are unchanged.
struct XtalResidualBeamOrientation {
    XtalResidualBeamOrientation(const XtalResidual &residual,
                                double distance_mm,
                                const double *detector_rot,
                                const double *rotation_axis,
                                const double *uc_len,
                                const double *uc_angle)
        : residual(residual),
          distance_mm(distance_mm),
          detector_rot{detector_rot[0], detector_rot[1]},
          rotation_axis{rotation_axis[0], rotation_axis[1], rotation_axis[2]},
          uc_len{uc_len[0], uc_len[1], uc_len[2]},
          uc_angle{uc_angle[0], uc_angle[1], uc_angle[2]} {
    }

    template<typename T>
    bool operator()(const T *const beam, const T *const p0, T *residual_out) const {
        const T distance[1] = {T(distance_mm)};
        const T rot[2] = {T(detector_rot[0]), T(detector_rot[1])};
        const T axis[3] = {T(rotation_axis[0]), T(rotation_axis[1]), T(rotation_axis[2])};
        const T p1[3] = {T(uc_len[0]), T(uc_len[1]), T(uc_len[2])};
        const T p2[3] = {T(uc_angle[0]), T(uc_angle[1]), T(uc_angle[2])};
        return residual(beam, distance, rot, axis, p0, p1, p2, residual_out);
    }

    const XtalResidual residual;
    const double distance_mm;
    const double detector_rot[2];
    const double rotation_axis[3];
    const double uc_len[3];
    const double uc_angle[3];
};