Files
Jungfraujoch/image_analysis/bragg_prediction/BraggPredictionRot.cpp
T
leonarski_fandClaude Opus 5.5 35f42e7334 BraggPredictionRot: walk only the l of each column that can pass the phi test
Rotation prediction enumerated the whole (2h+1)(2k+1)(2l+1) box every frame and kept the ~1% whose
rotation solution lands near the frame. A solution within phi_limit of the frame needs
|f| = ||p0|^2 + 2 S0.p0| <= 2 |S0_perp| |p0| phi_limit, and along an (h, k) column f is a quadratic
in l, so the l that can pass are at most two intervals. Those are walked - widened by 1% in f and by
a whole index at each end - in the same order and with the same arithmetic; the rest the phi test
rejected anyway. Without a phi limit (min_zeta 0) the whole column is walked as before. The GPU
kernel is unchanged.

A new test predicts 1000 random lattices, geometries, axes and rocking widths both ways and
requires the same reflections in the same order, every value bit for bit (30000 cases passed once
locally). CPU build, two interleaved warm runs: 8a1a output loop 241.5 -> 216.3 s, pre-pass loop
90.7 -> 73.8 s, whole run 427.0 -> 384.2 s; 8qaw (with the gzip change before this) loops
116.6 / 251.3 -> 94.4 / 206.8 s, user CPU 14700 -> 12718 s. GPU build within noise.
p.mtz byte-identical on myob, cytc, thau, kdp, 8a1a and 8qaw, GPU and CPU builds.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
2026-10-08 06:46:28 +02:00

321 lines
17 KiB
C++

// 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 = (std::erf((phi + half_wedge_angle_rad) * c1)
- std::erf((phi - half_wedge_angle_rad) * c1)) / 2.0f;
// 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);
}