Jungfraujoch/image_analysis/bragg_prediction/BraggPrediction.cpp

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

#include <algorithm>

#include "../../common/JFJochMath.h"
#include "../../common/Logger.h"
#include "BraggPrediction.h"
#include "../bragg_integration/SystematicAbsence.h"

void BraggPrediction::ReportOverflow(int predicted) {
    if (overflow_reported)
        return;
    overflow_reported = true;
    Logger("BraggPrediction").Warning(
        "A frame predicted {} reflections but the buffer holds {} - the ones kept are whichever the GPU "
        "wrote first, so this dataset does NOT process reproducibly. Results will differ between runs. "
        "The cell is large enough that the prediction cap no longer fits it.",
        predicted, kPredictionCapacity);
}

int BraggPrediction::TruncateToOutput(int count) {
    if (count <= kPredictionOutput)
        return count;

    std::partial_sort(reflections.begin(), reflections.begin() + kPredictionOutput,
                      reflections.begin() + count,
                      [](const Reflection &a, const Reflection &b) {
                          if (a.dist_ewald != b.dist_ewald) return a.dist_ewald < b.dist_ewald;
                          if (a.h != b.h) return a.h < b.h;
                          if (a.k != b.k) return a.k < b.k;
                          return a.l < b.l;
                      });
    return kPredictionOutput;
}

namespace {
    // Number of bandwidth sigmas included in the (radially thickened) Ewald-shell
    // acceptance window. 3σ captures essentially the whole pink-beam smear; matches
    // the conservative end of the mosaicity cutoff used by callers.
    constexpr float kBandwidthCutoffSigmas = 3.0f;
}

BraggPrediction::BraggPrediction(int max_reflections)
    : max_reflections(max_reflections), reflections(max_reflections) {}

const std::vector<Reflection> &BraggPrediction::GetReflections() const {
    return reflections;
}

int BraggPrediction::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 float epsilon = 1e-5f;
    const float s0_sq = S0 * S0;
    const float rad_to_deg = 180.0f / static_cast<float>(PI);

    int i = 0;

    for (int h = -settings.max_hkl; h <= settings.max_hkl; h++) {
        // Precompute A* h contribution
        const float Ah_x = Astar.x * h;
        const float Ah_y = Astar.y * h;
        const float Ah_z = Astar.z * h;

        for (int k = -settings.max_hkl; k <= settings.max_hkl; k++) {
            // Accumulate B* k contribution
            const float AhBk_x = Ah_x + Bstar.x * k;
            const float AhBk_y = Ah_y + Bstar.y * k;
            const float AhBk_z = Ah_z + Bstar.z * k;

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

                if (i >= max_reflections)
                    continue;

                float recip_x = AhBk_x + Cstar.x * l;
                float recip_y = AhBk_y + Cstar.y * l;
                float recip_z = AhBk_z + Cstar.z * l;

                float recip_sq = recip_x * recip_x + recip_y * recip_y + recip_z * recip_z;
                if (recip_sq > one_over_dmax_sq)
                    continue;

                float S_x = recip_x + S0.x;
                float S_y = recip_y + S0.y;
                float S_z = recip_z + S0.z;
                float S_len = sqrtf(S_x * S_x + S_y * S_y + S_z * S_z);
                float dist_ewald_sphere = std::fabs(S_len - one_over_wavelength);

                // Energy bandwidth thickens the Ewald shell radially: at the
                // diffraction condition |S|-1/λ shifts by recip_z·(Δλ/λ), i.e.
                // σ_bw = |recip_z|·bandwidth_sigma (= bλ/2d²). Broaden the acceptance
                // window in quadrature so high-resolution shells (smeared most, ∝1/d²)
                // are not clipped.
                float radial_cutoff = settings.ewald_dist_cutoff;
                if (settings.bandwidth_sigma > 0.0f) {
                    const float bw_tol = kBandwidthCutoffSigmas * settings.bandwidth_sigma * std::fabs(recip_z);
                    radial_cutoff = std::sqrt(radial_cutoff * radial_cutoff + bw_tol * bw_tol);
                }

                if (dist_ewald_sphere <= radial_cutoff ) {
                    const float s0_p0 = S0.x * recip_x + S0.y * recip_y + S0.z * recip_z;
                    const float val = s0_sq * recip_sq - s0_p0 * s0_p0;

                    float delta_phi_deg = NAN;
                    if (std::fabs(val) >= epsilon && s0_sq > epsilon) {
                        const float a_num = (s0_sq - 0.25f * recip_sq) * recip_sq;
                        if (a_num >= 0.0f) {
                            const float A = std::sqrt(a_num / val);
                            const float B = (A * s0_p0 + 0.5f * recip_sq) / s0_sq;

                            const float p_star_x = A * recip_x - B * S0.x;
                            const float p_star_y = A * recip_y - B * S0.y;
                            const float p_star_z = A * recip_z - B * S0.z;

                            const float p_star_sq = p_star_x * p_star_x + p_star_y * p_star_y + p_star_z * p_star_z;
                            const float denom = std::sqrt(p_star_sq * recip_sq);

                            if (denom >= epsilon) {
                                float c = (p_star_x * recip_x + p_star_y * recip_y + p_star_z * recip_z) / denom;
                                c = std::fmax(-1.0f, std::fmin(1.0f, c));
                                delta_phi_deg = std::acos(c) * rad_to_deg;
                            }
                        }
                    }

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

                    // Project to detector coordinates
                    // Assume detector is along x,y,z coordinates after rotation
                    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 d = 1.0f / sqrtf(recip_sq);
                    reflections[i] = Reflection{
                        .h = h,
                        .k = k,
                        .l = l,
                        .delta_phi_deg = delta_phi_deg,
                        .predicted_x = x,
                        .predicted_y = y,
                        .observed_x = NAN,
                        .observed_y = NAN,
                        .d = d,
                        .dist_ewald = dist_ewald_sphere,
                        .rlp = 1.0,
                        .partiality = 1.0f,
                        .zeta = 1.0,
                        .image_scale_corr = 1.0
                    };
                    ++i;
                }
            }
        }
    }
    return TruncateToOutput(i);
}