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 <algorithm>
#include <cmath>

#include "../../common/JFJochMath.h"
#include "BraggPredictionRot.h"
#include "../SensorAbsorption.h"
#include "../bragg_integration/SystematicAbsence.h"
#include "RockingSlice.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.GetDetectorMatrix().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 four_S0_sq = 4 * S0 * S0;
    const float m3_S0 = m3 * S0;

    int i = 0;
    const Coord beam_dir = S0.Normalize();

    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;

    // Angle-dependent sensor efficiency. Per-dataset constants (thickness, material, wavelength)
    // collapse to two numbers here; the per-reflection part is one exponential below. Inert - and
    // bit-identical to not applying it - wherever the sensor is opaque, which is every long
    // wavelength, so it needs no flag and no threshold anyone has to choose.
    const auto &det = experiment.GetDetectorSetup();
    const auto sensor_qe = sensor_absorption::SensorQE::Build(
        det.GetSensorMaterial(), det.GetSensorThickness_um(), geom.GetWavelength_A());
    // The air in the sample-to-pixel flight path carries the same cos(alpha) dependence, with the
    // opposite sign. See sensor_absorption::FlightPathAttenuation; off leaves d_over_L at 0, which is
    // bit-identical to not applying it.
    const auto air = sensor_absorption::FlightPathAttenuation::Build(
        experiment.GetBraggIntegrationSettings().GetFlightPath(), geom.GetDetectorDistance_mm(),
        geom.GetWavelength_A());

    // The phi test after the atan2 below, taken before it: a solution with |phi| above phi_limit fails
    // that test whatever its own bandwidth term, since mos_eff_rad is largest at the resolution limit.
    // |phi| > phi_limit is cos(phi) < cos(phi_limit), which the (cos, sin) pair answers without the
    // angle. phi_limit carries a 1e-3 rad margin, far above float rounding, so every solution rejected
    // here is rejected by that test too. -1 = nothing is rejected here.
    double cos_phi_limit = -1.0;
    double phi_limit = -1.0;
    if (settings.min_zeta > 0.0f) {
        double mos_max_rad = mos_angle_rad;
        if (bandwidth_sigma > 0.0f) {
            const double sin_theta = half_wavelength_A * static_cast<double>(one_over_dmax);
            const double dphi_bw = sin_theta < 1.0
                                   ? bandwidth_sigma * sin_theta / std::sqrt(1.0 - sin_theta * sin_theta)
                                   : INFINITY;
            mos_max_rad = std::sqrt(mos_max_rad * mos_max_rad + dphi_bw * dphi_bw);
        }
        phi_limit = half_wedge_angle_rad
                    + (settings.mosaicity_multiplier * mos_max_rad + 1e-5) / settings.min_zeta
                    + 1e-3;
        if (phi_limit < PI)
            cos_phi_limit = std::cos(phi_limit);
        else
            phi_limit = -1.0;
    }

    // Which l of an (h, k) column can pass that test at all, so the walk does not visit the rest of
    // the box. f(p0) = |p0|^2 + 2 S0.p0 is zero where p0 lies on the Ewald sphere, and rotating p0 by
    // phi about the axis moves S0.p0 by at most |S0_perp| |p0_perp| |phi|. So a reflection whose
    // solution lies within phi_limit has |f| <= 2 |S0_perp| |p0| phi_limit <= f_max, and along a column
    // f is a quadratic in l, so the l that can pass are at most two intervals. Both are widened - f_max
    // by 1%, each end by a whole index - so every l the tests below keep is still walked, in the same
    // order and with the same arithmetic; the rest the phi test would have rejected anyway.
    auto dot = [](const Coord &a, const Coord &b) {
        return static_cast<double>(a.x) * b.x + static_cast<double>(a.y) * b.y + static_cast<double>(a.z) * b.z;
    };
    const double S0_perp = std::sqrt(std::max(0.0, dot(S0, S0) - static_cast<double>(m2_S0) * m2_S0));
    const double f_max = 1.01 * 2.0 * S0_perp * one_over_dmax * phi_limit;
    const double C_C = dot(Cstar, Cstar), S0_C = dot(S0, Cstar);
    // A root of the quadratic as an index, held inside the box first so the conversion cannot overflow.
    auto to_l = [&](double v) {
        return static_cast<int>(std::clamp(v, -settings.max_l - 2.0, settings.max_l + 2.0));
    };

    // p0 = A* h + B* k + C* l, evaluated as ((A* h) + (B* k)) + (C* l) exactly as before, with the
    // terms that do not change in the inner loops taken out of them.
    std::vector<Coord> Cstar_l(2 * settings.max_l + 1);
    for (int l = -settings.max_l; l <= settings.max_l; l++)
        Cstar_l[l + settings.max_l] = Cstar * l;

    for (int h = -settings.max_h; h <= settings.max_h; h++) {
        const Coord Astar_h = Astar * h;

        for (int k = -settings.max_k; k <= settings.max_k; k++) {
            const Coord AB = Astar_h + Bstar * k;

            // The l of this column that can pass (above): f(l) = C_C l^2 + b l + c <= f_max on
            // [l_lo, l_hi], and f(l) < -f_max strictly inside (skip_lo, skip_hi).
            int l_lo = -settings.max_l, l_hi = settings.max_l;
            int skip_lo = settings.max_l + 1, skip_hi = settings.max_l + 1;
            if (phi_limit > 0.0 && !enumerate_full_box) {
                const double b = 2.0 * (dot(AB, Cstar) + S0_C);
                const double c = dot(AB, AB) + 2.0 * dot(S0, AB);
                const double disc_hi = b * b - 4.0 * C_C * (c - f_max);
                if (!(disc_hi >= 0.0))
                    continue;
                const double r_hi = std::sqrt(disc_hi);
                l_lo = std::max(l_lo, to_l(std::ceil((-b - r_hi) / (2.0 * C_C))) - 1);
                l_hi = std::min(l_hi, to_l(std::floor((-b + r_hi) / (2.0 * C_C))) + 1);
                const double disc_lo = b * b - 4.0 * C_C * (c + f_max);
                if (disc_lo > 0.0) {
                    const double r_lo = std::sqrt(disc_lo);
                    skip_lo = to_l(std::floor((-b - r_lo) / (2.0 * C_C))) + 2;
                    skip_hi = to_l(std::ceil((-b + r_lo) / (2.0 * C_C))) - 2;
                }
            }

            for (int l = l_lo; l <= l_hi; l++) {
                if (l >= skip_lo && l <= skip_hi)
                    l = skip_hi + 1;
                if (l > l_hi)
                    break;
                if (systematic_absence(h, k, l, settings.centering))
                    continue;

                const Coord &Cl = Cstar_l[l + settings.max_l];
                const Coord p0(AB.x + Cl.x, AB.y + Cl.y, AB.z + Cl.z);

                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 > four_S0_sq))
                    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) {
                    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;
                    const double c = cosphi, s = sinphi;
                    if (c < cos_phi_limit * std::sqrt(c * c + s * s))
                        continue;
                    float phi = -1.0f * std::atan2(sinphi, cosphi);

                    // Most solutions are far from this frame. zeta >= min_zeta, so the rocking-curve test
                    // below (epsilon3) rejects every solution with min_zeta * (|phi| - half wedge) above the
                    // limit; testing that first, with a margin well above float rounding, skips the rest
                    // of the arithmetic for them and changes nothing about the ones that pass.
                    if (settings.min_zeta > 0.0f
                        && settings.min_zeta * (std::fabs(static_cast<double>(phi)) - half_wedge_angle_rad)
                           > settings.mosaicity_multiplier * static_cast<double>(mos_eff_rad) + 1e-5)
                        continue;

                    Coord p = m1 * p_m1 + m2 * p_m2 + m3 * p_m3; // p0 vector "rotated" to diffracting condition
                    Coord S = S0 + p;

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

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

                    if (zeta_abs < settings.min_zeta)
                        continue;

                    // Is any of this reflection's rocking curve inside THIS image's oscillation range?
                    // phi is the offset from the frame's mid-exposure angle to the exact diffracting
                    // condition, so the curve - a Gaussian of width mos_eff/zeta in phi - has to be
                    // measured against the frame's EDGE, not its centre. Testing |phi|*zeta alone
                    // asks whether the frame centre is within the curve, which is a different and
                    // stricter question: it rejects a reflection whose curve overlaps the exposure
                    // but whose exact condition falls outside it. Since consecutive frame centres are
                    // one wedge apart, the nearest centre can be half a wedge away, so a reflection is
                    // rejected on EVERY frame - lost entirely, not merely clipped - once
                    // mos_eff < zeta * wedge / (2 * multiplier). That is coarse slicing on a sharp
                    // crystal at high zeta, where the reflection is fully recorded and cleanest.
                    float epsilon3 = std::fabs(phi * zeta_abs) - half_wedge_angle_rad * 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 = RockingSlicePartiality(phi, half_wedge_angle_rad, c1);
                    // Place the partial where this frame recorded it: at the centre of the frame's slice of
                    // the rocking curve, which is the exact-condition position walked along its Debye ring by
                    // the rotation between the two (RockingSlice.h). Walking along the ring is turning S
                    // about the beam, by the rotation's component along the ring over the ring's radius.
                    const float c_slice = RockingSliceCentroid_rad(phi, half_wedge_angle_rad, c1, partiality);
                    const Coord S_par = beam_dir * (beam_dir * S), S_perp = S - S_par;
                    const Coord S_turn = beam_dir % S;   // |S_turn| = |S_perp|
                    const float psi = c_slice * ((m2 % p) * S_turn) / (S_perp * S_perp);
                    S = S_par + S_perp * std::cos(psi) + S_turn * std::sin(psi);

                    // Inlined RecipToDetector: the full transposed detector matrix, tilt and discrete orientation
                    // 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);

                    // Sensor quantum efficiency at this reflection's own angle of incidence on the
                    // detector. The angle is taken against the DETECTOR NORMAL - S_rot is the
                    // diffracted direction in the detector's own frame, so its z component over its
                    // length is that cosine already, at no cost. Taking it here rather than from the
                    // resolution is what makes it right on a tilted detector, where the incidence
                    // angle stops being a function of resolution and the correction stops
                    // cancelling within a resolution shell.
                    const float cos_alpha = S_rot_z / S.Length();
                    const float qe_corr = sensor_qe.Factor(cos_alpha);
                    const float flight_corr = air.Factor(cos_alpha);

                    if (i == max_reflections)
                        GrowCapacity(2 * max_reflections);

                    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,
                        .prescaling_corr = lorentz_reciprocal,
                        .qe_corr = qe_corr,
                        .flight_corr = flight_corr,
                        .partiality = partiality,
                        .zeta = zeta_abs,
                        .image_scale_corr = lorentz_reciprocal * qe_corr * flight_corr / partiality,
                    };
                    i++;
                }
            }
        }
    }
    return TruncateToOutput(i);
}