Jungfraujoch/image_analysis/bragg_prediction/BraggPredictionRot.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 "BraggPredictionRot.h"
#include "../bragg_integration/SystematicAbsence.h"


int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const CrystalLattice &lattice,
    const BraggPredictionSettings &settings) {

    const auto geom = experiment.GetDiffractionGeometry();
    const auto det_width_pxl = static_cast<float>(experiment.GetXPixelsNum());
    const auto det_height_pxl = static_cast<float>(experiment.GetYPixelsNum());

    const float one_over_dmax = 1.0f / settings.high_res_A;
    const float one_over_dmax_sq = one_over_dmax * one_over_dmax;

    float one_over_wavelength = 1.0f / geom.GetWavelength_A();

    const Coord Astar = lattice.Astar();
    const Coord Bstar = lattice.Bstar();
    const Coord Cstar = lattice.Cstar();
    const Coord S0 = geom.GetScatteringVector();

    std::vector<float> rot = geom.GetPoniRotMatrix().transpose().arr();

    // Precompute detector geometry constants
    float beam_x = geom.GetBeamX_pxl();
    float beam_y = geom.GetBeamY_pxl();
    float det_distance = geom.GetDetectorDistance_mm();
    float pixel_size = geom.GetPixelSize_mm();
    float F = det_distance / pixel_size;

    const auto gon_opt = experiment.GetGoniometer();
    if (!gon_opt.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "BraggPredictionRotationCPU requires a goniometer axis");
    const GoniometerAxis& gon = *gon_opt;

    const Coord m2 = gon.GetAxis().Normalize();
    const Coord m1 = (m2 % S0).Normalize();
    const Coord m3 = (m1 % m2).Normalize();

    const float m2_S0 = m2 * S0;
    const float m3_S0 = m3 * S0;

    int i = 0;

    const float mos_angle_rad = settings.mosaicity_deg * static_cast<float>(PI) / 180.f;
    const float half_wedge_angle_rad = settings.wedge_deg * static_cast<float>(PI) / 180.f / 2.0f ;

    // Energy bandwidth widens the rocking curve. Differentiating Bragg's law at fixed d gives
    // dtheta = (dlambda/lambda) tan(theta_B), a spread in the same glancing angle the mosaic spread
    // smears, so it adds to sigma_M in quadrature. It is NOT divided by zeta here: rotating the
    // crystal by dphi changes theta by zeta*dphi, so the 1/zeta that turns an angular width into a
    // rotation width is already the one c1 (and the epsilon3 cutoff) applies to sigma_M. The fitted
    // sigma_M has this term deconvolved out (CalcMosaicityXDS), so it is not counted twice.
    // sin(theta_B) = lambda/(2d) = lambda*|p0|/2. Zero bandwidth leaves every reflection untouched.
    const float bandwidth_sigma = settings.bandwidth_sigma;
    const float half_wavelength_A = geom.GetWavelength_A() / 2.0f;

    for (int h = -settings.max_h; h <= settings.max_h; h++) {
        // Precompute A* h contribution

        for (int k = -settings.max_k; k <= settings.max_k; k++) {
            // Accumulate B* k contribution

            for (int l = -settings.max_l; l <= settings.max_l; l++) {
                if (systematic_absence(h, k, l, settings.centering))
                    continue;

                if (i >= max_reflections)
                    continue;
                Coord p0 = Astar * h + Bstar * k + Cstar * l;

                float p0_sq = p0 * p0;
                if (p0_sq <= 0.0f || p0_sq > one_over_dmax_sq)
                    continue;

                const float p0_m1 = p0 * m1;
                const float p0_m2 = p0 * m2;
                const float p0_m3 = p0 * m3;

                const float rho_sq = p0_sq - (p0_m2 * p0_m2);

                const float p_m3 = (- p0_sq / 2 - p0_m2 * m2_S0) / m3_S0;
                const float p_m2 = p0_m2;
                const float p_m1_opt[2] = {
                    std::sqrt(rho_sq - p_m3 * p_m3),
                    -std::sqrt(rho_sq - p_m3 * p_m3)
                };

                // No solution for Laue equations
                if ((rho_sq < p_m3 * p_m3) || (p0_sq > 4 * S0 * S0))
                    continue;

                // Effective rocking width for this reflection: mosaicity broadened by the bandwidth
                // term. sin(theta_B) <= 1 is guaranteed by the p0_sq test just above.
                float mos_eff_rad = mos_angle_rad;
                if (bandwidth_sigma > 0.0f) {
                    const float sin_theta = half_wavelength_A * std::sqrt(p0_sq);
                    const float dphi_bw = bandwidth_sigma * sin_theta / std::sqrt(1.0f - sin_theta * sin_theta);
                    mos_eff_rad = std::sqrt(mos_angle_rad * mos_angle_rad + dphi_bw * dphi_bw);
                }

                for (const auto& p_m1 : p_m1_opt) {
                    if (i >= max_reflections)
                        continue;

                    const float cosphi = (p_m1 * p0_m1 + p_m3 * p0_m3) / rho_sq;
                    const float sinphi = (p_m1 * p0_m3 - p_m3 * p0_m1) / rho_sq;
                    Coord p = m1 * p_m1 + m2 * p_m2 + m3 * p_m3; // p0 vector "rotated" to diffracting condition
                    Coord S = S0 + p;

                    float phi = -1.0f * std::atan2(sinphi, cosphi);

                    const Coord e1 = (S % S0).Normalize();

                    const float zeta_abs = std::fabs(m2 * e1);

                    if (zeta_abs < settings.min_zeta)
                        continue;

                    float epsilon3 = std::fabs(phi * zeta_abs);

                    if (epsilon3 > settings.mosaicity_multiplier * mos_eff_rad)
                        continue;

                    // Reciprocal Lorentz (Kabsch 2010): L^-1 = |m2 . (S x S0)| / (|S| |S0|) =
                    // |zeta * sin angle(S,S0)|. The original divided by the scalar product
                    // S.S0 = |S||S0|cos(2theta), adding a spurious 1/cos(2theta) (1.8x at 1 A) that
                    // corrupts the absolute/Wilson scale (it cancels within a resolution shell, so
                    // CC1/2 / CCref / R-meas are neutral).
                    const float lorentz_reciprocal = std::fabs(m2 * (S % S0)) / (S.Length() * S0.Length());
                    const float c1 = zeta_abs / (std::sqrt(2.0f) * mos_eff_rad);

                    const float partiality = (std::erf((phi + half_wedge_angle_rad) * c1)
                        - std::erf((phi - half_wedge_angle_rad) * c1)) / 2.0f;
                    // Inlined RecipToDector with rot1 and rot2 (rot3 = 0)
                    // Apply rotation matrix transpose
                    float S_rot_x = rot[0] * S.x + rot[1] * S.y + rot[2] * S.z;
                    float S_rot_y = rot[3] * S.x + rot[4] * S.y + rot[5] * S.z;
                    float S_rot_z = rot[6] * S.x + rot[7] * S.y + rot[8] * S.z;

                    if (S_rot_z <= 0)
                        continue;

                    float x = beam_x + F * S_rot_x / S_rot_z;
                    float y = beam_y + F * S_rot_y / S_rot_z;

                    if ((x < 0) || (x >= det_width_pxl) || (y < 0) || (y >= det_height_pxl))
                        continue;

                    float dist_ewald_sphere = std::fabs(S.Length() - one_over_wavelength);

                    float d = 1.0f / sqrtf(p0_sq);
                    reflections[i] = Reflection{
                        .h = h,
                        .k = k,
                        .l = l,
                        .delta_phi_deg = phi * 180.0f / static_cast<float>(PI),
                        .predicted_x = x,
                        .predicted_y = y,
                        .observed_x = NAN,
                        .observed_y = NAN,
                        .d = d,
                        .dist_ewald = dist_ewald_sphere,
                        .rlp = lorentz_reciprocal,
                        .partiality = partiality,
                        .zeta = zeta_abs,
                        .image_scale_corr = lorentz_reciprocal / partiality,
                    };
                    i++;
                }
            }
        }
    }
    return TruncateToOutput(i);
}