Files
Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp
T
leonarski_fandClaude Opus 5.5 81b81ec683 CPU Bragg integration: reflections walked in detector order, output unchanged
The prediction lists a frame's reflections by Miller index, which scatters consecutive ones over the
whole detector, so pass A's box walk, the reflection-mask and owner-map builds and the CPU profile fit
each waited on memory for every box's image rows and mask tiles (the pixel load alone was 18% of pass
A's samples). These loops now visit the reflections sorted by 16-row band, then column. Each of them
gives the same result in any order - integer sums, min/max, per-reflection outputs indexed by the
reflection - so the output is unchanged; the profile learning, whose double sums do depend on the
order, still walks the original one.

cytochrome C --no-gpu: user cycles 10.71e12 -> 10.20e12 (-4.8%); p.hkl and MTZ data identical.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-15 14:52:22 +02:00

885 lines
49 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "BraggIntegrationEngineCPU.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <vector>
#if defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86))
#include <xmmintrin.h>
#endif
#include "../../common/CompressedImage.h"
#include "../../common/JFJochException.h"
using namespace bragg_engine;
namespace {
// The engine reads pixels in the INT32_MIN(masked)/INT32_MAX(saturated) convention.
inline bool valid(int32_t v) { return v != INT32_MIN && v != INT32_MAX; }
// Ask for the cache lines from first to last ahead of their use. Only a hint: it changes no result.
inline void prefetch(const void *p) {
#if defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86))
_mm_prefetch(static_cast<const char *>(p), _MM_HINT_T0);
#elif defined(__GNUC__) || defined(__clang__)
__builtin_prefetch(p);
#endif
}
inline void prefetch_row(const void *first, const void *last) {
for (const char *a = static_cast<const char *>(first); a < last; a += 64) prefetch(a);
prefetch(last);
}
// Sampler over a raw detector image of pixel type T: masked pixels carry the type minimum, saturated
// the type maximum (the FPGA image has no lossy-codec +/-1 band). Only the pixels actually read - the
// reflection disks - are converted, so there is no whole-image pass.
template <class T>
struct ImageSampler {
const T *p;
int64_t special_value;
int64_t saturation;
int32_t operator[](size_t i) const {
const int64_t v = p[i];
if (v == special_value) return INT32_MIN;
if (v == saturation) return INT32_MAX;
return static_cast<int32_t>(v);
}
const void *address(size_t i) const { return p + i; }
};
// The reflections in detector order: by band of 16 rows (a mask tile's height), then by column. The
// prediction lists them by Miller index, which scatters consecutive reflections over the whole frame, so
// each box waited on memory for its image rows and mask tiles; in this order neighbours share them in
// cache. Only for the loops whose results do not depend on the order the reflections are visited in.
std::vector<uint32_t> FrameOrder(const std::vector<Reflection> &predicted, size_t npredicted) {
std::vector<std::pair<uint64_t, uint32_t>> keyed(npredicted);
for (size_t i = 0; i < npredicted; ++i) {
const auto band = static_cast<uint64_t>(std::max(0.0f, predicted[i].predicted_y)) / 16;
const auto col = static_cast<uint64_t>(std::max(0.0f, predicted[i].predicted_x));
keyed[i] = {(band << 32) | col, static_cast<uint32_t>(i)};
}
std::sort(keyed.begin(), keyed.end());
std::vector<uint32_t> order(npredicted);
for (size_t k = 0; k < npredicted; ++k)
order[k] = keyed[k].second;
return order;
}
} // namespace
BraggIntegrationEngineCPU::BraggIntegrationEngineCPU(const DiffractionExperiment &experiment,
const PixelMask &mask)
: BraggIntegrationEngine(experiment),
refl_mask(static_cast<int>(xpixel), static_cast<int>(ypixel), 0),
owner(static_cast<int>(xpixel), static_cast<int>(ypixel), BRAGG_OWNER_NONE),
static_mask(mask.GetPackedMask()) {}
void BraggIntegrationEngineCPU::MarkNeighbours(const std::vector<Reflection> &predicted,
const std::vector<uint32_t> &order) {
const int W = static_cast<int>(xpixel), H = static_cast<int>(ypixel);
// --- Reflection mask: mark the r2 signal region of every predicted reflection, so a neighbour's
// signal is excluded from this reflection's r2..r3 background ring. MASK_FLUX where the
// neighbour puts flux on this frame (NEIGHBOUR_MASK_MIN_PARTIALITY) - only those are excluded
// - and MASK_TAIL where only the tail of a rocking curve does, which is counted for the density
// guard but read as background. The region is the INNER stencil ellipse, taken in each
// neighbour's OWN frame - an elongated neighbour whose streak is still marked as a disk would
// leak its tails into this reflection's ring. ---
refl_mask.Clear();
for (const uint32_t i : order) {
const auto &r = predicted[i];
const uint8_t level = r.partiality >= NEIGHBOUR_MASK_MIN_PARTIALITY ? MASK_FLUX : MASK_TAIL;
const BraggStencil st = MakeBraggStencil(r.predicted_x, r.predicted_y, stencil);
const int x0 = std::max(0, static_cast<int>(std::floor(r.predicted_x - st.ex_in - 1.0f)));
const int x1 = std::min(W - 1, static_cast<int>(std::ceil(r.predicted_x + st.ex_in + 1.0f)));
const int y0 = std::max(0, static_cast<int>(std::floor(r.predicted_y - st.ey_in - 1.0f)));
const int y1 = std::min(H - 1, static_cast<int>(std::ceil(r.predicted_y + st.ey_in + 1.0f)));
for (int y = y0; y <= y1; ++y)
for (int x = x0; x <= x1; ++x) {
const auto d = BraggStencilDistances(st, x - r.predicted_x, y - r.predicted_y);
if (d.inner < r2_sq) {
uint8_t &m = refl_mask.At(x, y);
m = std::max(m, level);
}
}
}
}
void BraggIntegrationEngineCPU::CountNeighbourCrowding(const std::vector<Reflection> &predicted,
size_t npredicted) {
const int W = static_cast<int>(xpixel), H = static_cast<int>(ypixel);
counts.predicted += npredicted;
const std::vector<uint32_t> order = FrameOrder(predicted, npredicted);
MarkNeighbours(predicted, order);
for (const uint32_t i : order) {
const auto &r = predicted[i];
const BraggStencil st = MakeBraggStencil(r.predicted_x, r.predicted_y, stencil);
const int x0 = std::max(0, static_cast<int>(std::floor(r.predicted_x - st.ex_out - 1.0f)));
const int x1 = std::min(W - 1, static_cast<int>(std::ceil(r.predicted_x + st.ex_out + 1.0f)));
const int y0 = std::max(0, static_cast<int>(std::floor(r.predicted_y - st.ey_out - 1.0f)));
const int y1 = std::min(H - 1, static_cast<int>(std::ceil(r.predicted_y + st.ey_out + 1.0f)));
int n_ring = 0, n_clean = 0;
for (int y = y0; y <= y1; ++y)
for (int x = x0; x <= x1; ++x) {
const auto d = BraggStencilDistances(st, x - r.predicted_x, y - r.predicted_y);
if (BraggStencilInSignal(st, x - r.predicted_x, y - r.predicted_y, d.signal, r1_sq)
|| d.inner < r2_sq || d.outer >= r3_sq)
continue;
const size_t idx = static_cast<size_t>(y) * W + x;
if (idx / 32 < static_mask.size() && ((static_mask[idx / 32] >> (idx % 32)) & 1U))
continue;
++n_ring;
if (refl_mask.Get(x, y) == 0)
++n_clean;
}
if (n_clean <= 5 && n_ring > 5)
++counts.bkg_starved_by_neighbour;
}
}
template <class Sampler>
std::vector<Reflection> BraggIntegrationEngineCPU::RunImpl(const Sampler &img,
const std::vector<Reflection> &predicted,
size_t npredicted, int64_t image_number) {
if (npredicted == 0)
return Finalize(predicted, npredicted, std::vector<BraggFitResult>(npredicted), image_number);
counts.predicted += npredicted;
const int W = static_cast<int>(xpixel), H = static_cast<int>(ypixel);
const bool do_clip = bkg_clip_nsigma > 0.0f && mode != IntegratorMode::BoxSum;
// Every loop below whose results do not depend on the order walks the reflections in this one.
const std::vector<uint32_t> order = FrameOrder(predicted, npredicted);
MarkNeighbours(predicted, order);
owner.Clear();
// --- Signal-region ownership: a pixel inside two reflections' signal regions belongs to the
// NEARER predicted centre. The union mask above cannot answer that - it also marks a
// reflection's own core - so ownership gets its own map, one (distance, reflection) key per
// pixel (BraggStencil.h). Built only when an overlap treatment is asked for. ---
// A box sum has no profile to renormalise a disk it has taken pixels out of, so it never excludes.
const bool exclude = overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum;
const bool use_owner = overlap != OverlapMode::Off;
if (use_owner) {
const float claim_sq = claim * claim;
for (const uint32_t i : order) {
const auto &r = predicted[i];
const int x0 = std::max(0, static_cast<int>(std::floor(r.predicted_x - claim)));
const int x1 = std::min(W - 1, static_cast<int>(std::ceil(r.predicted_x + claim)));
const int y0 = std::max(0, static_cast<int>(std::floor(r.predicted_y - claim)));
const int y1 = std::min(H - 1, static_cast<int>(std::ceil(r.predicted_y + claim)));
for (int y = y0; y <= y1; ++y)
for (int x = x0; x <= x1; ++x) {
const float dx = x - r.predicted_x, dy = y - r.predicted_y;
const float d2 = dx * dx + dy * dy;
if (d2 >= claim_sq) continue;
uint32_t &o = owner.At(x, y);
o = std::min(o, BraggOwnerKey(std::sqrt(d2), inv_claim, static_cast<int>(i)));
}
}
}
auto clean = [&](int x, int y, size_t i) {
return !use_owner
|| BraggOwnedBy(owner.Get(x, y), static_cast<int>(i));
};
// --- Pass A: box-sum every reflection (rough I, background, centroid, strong flag). ---
// A saturated pixel, or one unreadable on this frame alone - the overload marker of a detector
// that writes its error value for a pixel it could not count (EIGER), which the preprocessor turns
// into a masked pixel like a gap's. The run's mask holds the gaps, so whatever is unreadable beyond
// it was lost to the flux.
auto static_masked = [&](size_t idx) {
return idx / 32 < static_mask.size() && ((static_mask[idx / 32] >> (idx % 32)) & 1U);
};
PassA a = StartPassA(npredicted);
std::vector<int32_t> bkg_vals; // reused per reflection for the trimmed-mean background (idea 1)
// The ring pixels the background clip below reads, with their radial offsets, gathered while pass
// A walks the box anyway.
struct ClipPixel { int32_t px; float rad; };
std::vector<ClipPixel> clip_ring;
// The box pass A reads: the OUTER ellipse, tightly - taking the largest semi-axis in both
// directions instead would read up to 60% more pixels for nothing.
struct Box { int x0, x1, y0, y1; };
auto box_of = [&](const Reflection &r, const BraggStencil &st) {
return Box{std::max(0, static_cast<int>(std::floor(r.predicted_x - st.ex_out - 1.0f))),
std::min(W - 1, static_cast<int>(std::ceil(r.predicted_x + st.ex_out + 1.0f))),
std::max(0, static_cast<int>(std::floor(r.predicted_y - st.ey_out - 1.0f))),
std::min(H - 1, static_cast<int>(std::ceil(r.predicted_y + st.ey_out + 1.0f)))};
};
for (size_t k = 0; k < npredicted; ++k) {
const size_t i = order[k];
// The next reflection's image rows, asked for while this one is worked on: the reflections
// are spread over a frame far larger than the caches, so each box would otherwise wait on
// memory row by row. (Written out here: GCC drops a call to a function that does nothing but
// prefetch.)
if (k + 1 < npredicted) {
const Reflection &n = predicted[order[k + 1]];
const Box b = box_of(n, MakeBraggStencil(n.predicted_x, n.predicted_y, stencil));
for (int y = b.y0; y <= b.y1 && b.x0 <= b.x1; ++y)
prefetch_row(img.address(static_cast<size_t>(y) * W + b.x0),
img.address(static_cast<size_t>(y) * W + b.x1));
}
const auto &r = predicted[i];
// This reflection's stencil: the r1 signal disk, and the r2..r3 ring elongated radially by
// the analytic smear.
const BraggStencil st = MakeBraggStencil(r.predicted_x, r.predicted_y, stencil);
const Box box = box_of(r, st);
const int x0 = box.x0, x1 = box.x1, y0 = box.y0, y1 = box.y1;
// n_ring / n_clean: readable ring pixels, and those of them no predicted neighbour's signal
// region covers. refl_mask marks d.inner < r2_sq and the ring is d.inner >= r2_sq, so the two
// are complementary and a masked ring pixel is always someone else's - never this reflection's
// own core. That makes the difference an exact count of what the pattern's density took,
// separable from what the detector took.
BraggBoxSums s{};
bkg_vals.clear();
clip_ring.clear();
for (int y = y0; y <= y1; ++y)
for (int x = x0; x <= x1; ++x) {
const auto d = BraggStencilDistances(st, x - r.predicted_x, y - r.predicted_y);
const int32_t px = img[y * W + x];
if (do_clip && d.inner >= r2_sq && d.outer < r3_sq && valid(px) && refl_mask.Get(x, y) != MASK_FLUX)
clip_ring.push_back({px, d.rad});
if (BraggStencilInSignal(st, x - r.predicted_x, y - r.predicted_y, d.signal, r1_sq)) {
// A pixel a nearer neighbour owns carries that neighbour's flux, so in Exclude
// mode it leaves the disk entirely - the sum, the pixel count the background is
// subtracted with, and the all-or-nothing validity gate alike. The box sum only
// counts how many were lost, which is all it can act on.
++s.n_disk;
if (overlap != OverlapMode::Off) {
if (clean(x, y, i)) ++s.n_own;
else if (exclude) continue;
}
++s.n_inner;
if (px == INT32_MAX
|| (px == INT32_MIN && !static_masked(static_cast<size_t>(y) * W + x)))
s.overloaded = 1;
if (!valid(px)) continue;
s.I_sum += px;
s.I_sum_x += static_cast<int64_t>(x) * px;
s.I_sum_y += static_cast<int64_t>(y) * px;
// Position sums over the very same pixels, so the background can be taken out of
// the centroid below. It is not known yet - the ring is read in this same loop.
s.x_sum += x;
s.y_sum += y;
++s.n_inner_valid;
} else if (d.inner >= r2_sq && d.outer < r3_sq) {
if (!valid(px)) continue;
++s.n_ring;
const uint8_t m = refl_mask.Get(x, y);
if (m == 0) ++s.n_clean;
if (m == MASK_FLUX) continue;
s.bkg_sum += px;
if (bkg_trim > 0.0f) bkg_vals.push_back(px);
++s.n_bkg;
}
}
// One high-outlier sigma-clip pass on the background ring (see ClipThreshold), and the radial
// background curve from the pixels it keeps.
BraggClipSums clip{};
if (ClipsRing(s)) {
const int32_t thr = ClipThreshold(s);
for (const auto &[px, rad] : clip_ring)
if (px <= thr) {
clip.sum += px;
++clip.n;
if (bkg_radial) {
// The radial curve is binned on the TRUE detector radius, so the
// offset here is the unshrunk radial projection.
const int b = std::clamp(static_cast<int>(std::lround(st.r0 + rad)), 0, a.n_rad - 1);
a.rad_sum[b] += px;
++a.rad_cnt[b];
}
}
}
FinishPassA(a, i, r, st, s, clip, bkg_vals);
}
return ProfilePasses(img, owner, predicted, npredicted, a, image_number);
}
BraggIntegrationEngineCPU::PassA BraggIntegrationEngineCPU::StartPassA(size_t npredicted) const {
const int W = static_cast<int>(xpixel), H = static_cast<int>(ypixel);
PassA a;
a.rough.resize(npredicted);
// A pixel's radius is the reflection's radius plus the pixel's projection on the beam->reflection
// direction, so no per-pixel sqrt is needed.
a.n_rad = static_cast<int>(std::ceil(std::hypot(std::max<double>(beam_x, W - beam_x),
std::max<double>(beam_y, H - beam_y)))) + 2;
if (bkg_radial) {
a.rad_sum.assign(a.n_rad, 0);
a.rad_cnt.assign(a.n_rad, 0);
}
a.inv_d2_min = std::numeric_limits<double>::max();
a.inv_d2_max = 0.0;
return a;
}
bool BraggIntegrationEngineCPU::ClipsRing(const BraggBoxSums &s) const {
// The clip is the alternative to the trimmed mean, so where it is on the trim is off, and the
// profile modes keep every reflection with a background (FinishPassA).
return bkg_clip_nsigma > 0.0f && mode != IntegratorMode::BoxSum && s.n_bkg > 5;
}
int32_t BraggIntegrationEngineCPU::ClipThreshold(const BraggBoxSums &s) const {
// Reject pixels above mean + n*sqrt(gain*mean) to strip a neighbour core or a zinger that biases
// the mean. The gain is 1 for counted photons, so this is then exactly the Poisson clip. A pixel
// value is an integer, so it is at or below the threshold exactly when it is at or below its floor.
const double bkg = static_cast<double>(s.bkg_sum) / s.n_bkg;
const double thr = bkg + bkg_clip_nsigma * std::sqrt(static_cast<double>(bkg_pixel_gain) * std::max(bkg, 1.0));
return static_cast<int32_t>(std::clamp(std::floor(thr), static_cast<double>(INT32_MIN),
static_cast<double>(INT32_MAX)));
}
void BraggIntegrationEngineCPU::FinishPassA(PassA &a, size_t i, const Reflection &r, const BraggStencil &st,
const BraggBoxSums &s, const BraggClipSums &clip,
std::vector<int32_t> &bkg_vals) {
Rough out;
// Both from the stencil's own radius, so nothing downstream is derived from a second one.
out.r_bin = std::clamp(static_cast<int>(std::lround(st.r0)), 0, a.n_rad - 1);
out.k_bin = BraggStencilKernelIndex(st, n_kern);
out.overloaded = s.overloaded != 0;
if (overlap != OverlapMode::Off && s.n_own < s.n_disk) {
++counts.overlapped;
if (s.n_own < overlap_min_peak * s.n_disk) ++counts.overlapped_major;
}
int n_bkg_used = s.n_bkg; // pixels behind the FINAL background value (trim/clip shrink it)
// A masked, untrusted, gapped or overloaded pixel inside the signal disk used to discard the
// reflection outright. A profile fit does not need it to: the fit is the amplitude of a
// NORMALISED profile, so leaving pixels out renormalises the estimator by construction and
// widens sigma by the information they carried - the same argument Exclude makes for a
// neighbour's pixels. Pass B keeps the reflection only while enough of the expected profile
// survived to constrain the amplitude (XDS's MINPK, dials' valid_foreground_threshold). A box
// sum has no profile to renormalise with, so there it stays all or nothing.
const bool full = s.n_inner_valid == s.n_inner;
// A ring left with five or fewer readable pixels cannot estimate a background, so the
// reflection is dropped whole - the one thing the stencil geometry does to the DATA rather
// than to a measurement. Counted here because it is the only direct evidence of a radius that has
// outgrown the pattern it is integrating (BraggIntegrationCounts).
const bool keep_partial = full || mode != IntegratorMode::BoxSum || out.overloaded;
if (keep_partial && s.n_bkg <= 5) ++counts.bkg_starved;
// Would every predicted neighbour, the tails included, leave the ring starved that the
// detector alone would not? That is the pattern's density, what the guard on a widened
// radius reads (BraggIntegrationCounts).
if (keep_partial && s.n_clean <= 5 && s.n_ring > 5) ++counts.bkg_starved_by_neighbour;
if (keep_partial && s.n_bkg > 5) {
out.bkg = static_cast<double>(s.bkg_sum) / s.n_bkg;
// Symmetric trimmed-mean background fraction (BraggIntegrationSettings, rugnux --background-trim):
// replaces the r2..r3 ring MEAN with an f-trimmed mean, robust to the high-side contamination
// (neighbour wings, tails) that biases the plain mean up and makes it over-subtract weak
// high-angle reflections. The base ctor has already forced it to 0 whenever the clip is in
// force, which is the default - the two are alternatives.
const double bkg_trim_frac = bkg_trim;
if (bkg_trim_frac > 0.0 && bkg_vals.size() > 5
&& bkg_vals.size() <= static_cast<size_t>(bragg_engine::BKG_TRIM_MAX)) {
// Symmetric trimmed mean over the background ring (idea 1): drop the lowest and highest
// bkg_trim_frac of the pixels, average the rest. Robust to the high-side contamination
// that biases the plain ring mean and makes it over-subtract at high resolution.
std::sort(bkg_vals.begin(), bkg_vals.end());
const size_t lo = static_cast<size_t>(bkg_vals.size() * bkg_trim_frac);
const size_t hi = bkg_vals.size() - lo;
if (hi > lo) {
double sum = 0.0;
for (size_t t = lo; t < hi; ++t) sum += bkg_vals[t];
out.bkg = sum / static_cast<double>(hi - lo);
n_bkg_used = static_cast<int>(hi - lo);
}
} else if (ClipsRing(s)) {
if (clip.n > 5) { out.bkg = static_cast<double>(clip.sum) / clip.n; n_bkg_used = clip.n; }
}
// The sum is over the pixels actually READ, so that is the count the background is
// subtracted with; with nothing missing it is the whole disk, exactly as before.
out.I = static_cast<double>(s.I_sum) - static_cast<double>(s.n_inner_valid) * out.bkg;
// I = I_sum - n_inner*bkg, and bkg is itself estimated from n_bkg_used pixels, so its
// error enters n_inner times over: var(I) = I_sum + n_inner^2 * bkg/n_bkg_used. Leaving
// the second term out understates sigma by sqrt(1 + n_inner/n_bkg) - 1.109x at the
// default r1=4/r2=6/r3=10 stencil, on every reflection of every dataset.
out.bkg_var = out.bkg / n_bkg_used;
out.var_bkg = static_cast<double>(s.n_inner_valid) * out.bkg
+ static_cast<double>(s.n_inner_valid) * s.n_inner_valid * out.bkg_var;
out.I_sum = s.I_sum;
out.n_inner = static_cast<int>(s.n_inner_valid);
out.n_disk = s.n_disk;
out.n_own = s.n_own;
const double var_bkg_term = static_cast<double>(s.n_inner_valid) * s.n_inner_valid * out.bkg_var;
out.sigma = 1.0;
if (s.I_sum > 0) {
out.sigma = std::max(out.sigma, std::sqrt(static_cast<double>(s.I_sum) + var_bkg_term));
// The centroid of the SIGNAL, not of the disk. Weighting by the raw counts weights by
// signal plus background, and the background is flat over a disk centred on the
// PREDICTION - so its own centroid is the prediction exactly, and it pulls the answer
// there. The measured displacement comes out shrunk by I/(I + n*bkg), which is worst
// where the background dominates, i.e. at high resolution. Post-refinement fits the
// beam and the distance on these centroids, so it under-corrects by that factor and
// does so resolution-dependently. Subtracting a flat pedestal from a first moment is
// exact: sum(x*(px-bkg)) = I_sum_x - bkg*x_sum over the same pixels.
const double net = static_cast<double>(s.I_sum) - s.n_inner_valid * out.bkg;
if (net > 0.0) {
out.obs_x = (static_cast<double>(s.I_sum_x) - out.bkg * s.x_sum) / net;
out.obs_y = (static_cast<double>(s.I_sum_y) - out.bkg * s.y_sum) / net;
} else {
// Nothing above background to take a centroid of; the raw one is all there is.
out.obs_x = static_cast<double>(s.I_sum_x) / static_cast<double>(s.I_sum);
out.obs_y = static_cast<double>(s.I_sum_y) / static_cast<double>(s.I_sum);
}
// A disk with a hole in it gives a centroid pulled away from the hole, and the hole
// sits at a fixed place on the detector - post-refinement would read that as geometry.
out.has_obs = full;
}
out.cx = static_cast<int>(std::lround(r.predicted_x));
out.cy = static_cast<int>(std::lround(r.predicted_y));
out.ok = true;
out.full = full;
// The profile, its resolution shells and their widths are learned from COMPLETE
// reflections: a partial one has no intensity to normalise its grid by and a hole to
// stack into it. So the learning sees exactly what it saw before this rescue existed.
out.strong = full && out.sigma > 0.0 && out.I / out.sigma >= STRONG_I_OVER_SIGMA;
if (full && r.d > 0.0f) {
const double inv_d2 = 1.0 / (static_cast<double>(r.d) * r.d);
a.inv_d2_min = std::min(a.inv_d2_min, inv_d2);
a.inv_d2_max = std::max(a.inv_d2_max, inv_d2);
}
}
a.rough[i] = out;
}
bool BraggIntegrationEngineCPU::FitKept(const BraggFitOut &o) const {
if (o.p_valid < overlap_min_peak * o.p_grid) return false;
// A hole in the profile's PEAK is a different defect from a hole in its wings, and the mass
// fraction above cannot tell them apart - the peak of a broad spot is a few percent of the
// mass, so MINPK passes a reflection that has lost the one part of the profile its amplitude
// is determined by. See the header.
if (o.p_lost_peak > MINPK_MAX_MISSING_PEAK * o.p_peak) return false;
if (overlap == OverlapMode::Reject && o.p_own < overlap_min_peak) return false;
return true;
}
template <class Sampler, class OwnerMap>
BraggFitOut BraggIntegrationEngineCPU::FitProfile(const Sampler &img, const OwnerMap &owner_map,
const BraggFitSetup &s, size_t i,
const std::vector<float> &shell_profiles,
std::vector<float> &cells,
std::vector<FitPixel> &fit_pixels) const {
const int W = static_cast<int>(xpixel), H = static_cast<int>(ypixel);
const bool exclude = overlap == OverlapMode::Exclude;
const bool use_owner = overlap != OverlapMode::Off;
const int Rf = s.Rf, Gf = 2 * Rf + 1;
BraggFitOut o{};
const float *P;
if (s.shell >= 0)
P = shell_profiles.data() + static_cast<size_t>(s.shell) * GG;
else {
// The grid's image rows, asked for now (as in pass A): building the profile below takes long
// enough for them to arrive before the fit reads them.
const int gx0 = std::max(0, s.cx - Rf), gx1 = std::min(W - 1, s.cx + Rf);
for (int y = std::max(0, s.cy - Rf); y <= std::min(H - 1, s.cy + Rf); ++y)
prefetch_row(img.address(static_cast<size_t>(y) * W + gx0),
img.address(static_cast<size_t>(y) * W + gx1));
cells.resize(static_cast<size_t>(Gf) * Gf);
float gs = 0.0f;
for (int dy = -Rf; dy <= Rf; ++dy)
for (int dx = -Rf; dx <= Rf; ++dx) {
const float g = BraggGaussian(dx, dy, s.fx, s.fy, s.ux, s.uy, s.inv_2s2r, s.inv_2s2t);
cells[(dy + Rf) * Gf + (dx + Rf)] = g;
gs += g;
}
const float inv_gs = 1.0f / gs;
for (float &p : cells) p *= inv_gs;
P = cells.data();
}
// --- How much of the expected profile the fit can actually see. p_valid is the readable
// fraction of it - XDS's MINPK, dials' valid_foreground_threshold - measured against the
// mass that falls on the detector at all, so a reflection is judged on the pixels that
// exist and not on where the sensor ends. p_own is the same quantity over neighbour-owned
// pixels, what Reject cuts on. m_read / m_all is the fraction of the r1 disk the
// summation seed the runaway guard compares against actually summed; with nothing missing
// and nothing excluded it is 1 and the guard is untouched. ---
const BraggStencil st = MakeBraggStencil(s.px, s.py, stencil);
fit_pixels.clear();
for (int dy = -Rf; dy <= Rf; ++dy)
for (int dx = -Rf; dx <= Rf; ++dx) {
const float Pp = P[(dy + Rf) * Gf + (dx + Rf)];
if (Pp <= 0.0f) continue;
const int x = s.cx + dx, y = s.cy + dy;
if (x < 0 || y < 0 || x >= W || y >= H) continue;
// The region the summation seed was taken over (pass A), so the guard compares the two
// over the same pixels.
const bool in_disk = s.footprint
? BraggStencilInSignal(st, x - s.px, y - s.py,
(x - s.px) * (x - s.px) + (y - s.py) * (y - s.py), r1_sq)
: dx * dx + dy * dy < r1_sq;
o.p_grid += Pp;
o.p_peak = std::max(o.p_peak, Pp);
if (in_disk) o.m_all += Pp;
const int32_t v = img[y * W + x];
if (!valid(v)) {
o.p_lost_peak = std::max(o.p_lost_peak, Pp);
continue;
}
o.p_valid += Pp;
const bool own = !use_owner || BraggOwnedBy(owner_map.Get(x, y), static_cast<int>(i));
if (own) o.p_own += Pp;
if (in_disk && (own || !exclude)) o.m_read += Pp;
if (own || !exclude)
fit_pixels.push_back({Pp, static_cast<float>(v) - s.bkg});
}
if (!FitKept(o)) return o;
float If = s.I0;
for (int iter = 0; iter < 4; ++iter) {
float num = 0.0f;
o.den = 0.0f;
o.wsum = 0.0f;
for (const auto &[Pp, signal] : fit_pixels) {
const float iv = 1.0f / std::max(s.B + If * Pp, s.v_min);
num += Pp * signal * iv;
o.den += Pp * Pp * iv;
o.wsum += Pp * iv;
}
if (o.den > 0.0f) If = num / o.den; else break;
}
o.I = If;
return o;
}
template <class Sampler, class OwnerMap>
std::vector<Reflection> BraggIntegrationEngineCPU::ProfilePasses(const Sampler &img, const OwnerMap &owner_map,
const std::vector<Reflection> &predicted,
size_t npredicted, PassA &a,
int64_t image_number) {
const int W = static_cast<int>(xpixel), H = static_cast<int>(ypixel);
std::vector<BraggFitResult> results(npredicted);
std::vector<Rough> &rough = a.rough;
const bool exclude = overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum;
const bool use_owner = overlap != OverlapMode::Off;
auto clean = [&](int x, int y, size_t i) {
return !use_owner
|| BraggOwnedBy(owner_map.Get(x, y), static_cast<int>(i));
};
auto grid_idx = [this](int dx, int dy) { return (dy + R) * G + (dx + R); };
// --- Radial background curvature correction. The annulus mean is blind to the curvature of the
// radial background (a linear background cancels between the concentric disk and annulus), so
// correct it by mean_annulus(B) - mean_disk(B) taken from the curve just accumulated. Reads no
// pixels: one short dot product per reflection. ---
if (bkg_radial) {
for (size_t i = 0; i < npredicted; ++i) {
auto &rh = rough[i];
if (!rh.ok) continue;
// An empty bin contributes the reflection's own background, so a fully empty
// neighbourhood gives corr == 0 exactly (the kernel weights sum to zero).
const float *kern = k_diff.data() + static_cast<size_t>(rh.k_bin) * k_len;
double corr = 0.0;
for (int k = 0; k < k_len; ++k) {
const int b = std::clamp(rh.r_bin + k - k_off, 0, a.n_rad - 1);
const double v = a.rad_cnt[b] > 0 ? static_cast<double>(a.rad_sum[b]) / a.rad_cnt[b] : rh.bkg;
corr += static_cast<double>(kern[k]) * v;
}
rh.bkg -= corr; // annulus mean -> mean over the signal disk
rh.I = static_cast<double>(rh.I_sum) - static_cast<double>(rh.n_inner) * rh.bkg;
}
}
// --- BoxSum mode is BraggIntegrate2D: emit the rough result directly. ---
if (mode == IntegratorMode::BoxSum) {
for (size_t i = 0; i < npredicted; ++i) {
const auto &rh = rough[i];
if (!rh.ok) continue;
// A box sum measures what is in the disk with no model of what should be there, so it can
// neither renormalise nor tell a neighbour's photon from its own: dropping the reflection
// is the only treatment it has, and it is applied only where that is what was ASKED for.
// Exclude means "leave the shared pixels out of the fit", and a box sum has no fit, so it
// is a no-op here rather than a rejection the caller never requested. The fraction is by
// AREA, not by profile mass, so the same threshold cuts harder here than in the profile
// modes.
if (overlap == OverlapMode::Reject && rh.n_own < overlap_min_peak * rh.n_disk) continue;
results[i] = {static_cast<float>(rh.I), static_cast<float>(rh.sigma), static_cast<float>(rh.bkg),
static_cast<float>(rh.obs_x), static_cast<float>(rh.obs_y),
static_cast<float>(rh.var_bkg), true, rh.has_obs, rh.overloaded};
}
return Finalize(predicted, npredicted, results, image_number);
}
auto shell_of = [&](float d) {
if (!(d > 0.0f) || a.inv_d2_max <= a.inv_d2_min) return 0;
const double t = (1.0 / (static_cast<double>(d) * d) - a.inv_d2_min) / (a.inv_d2_max - a.inv_d2_min);
return std::clamp(static_cast<int>(t * N_SHELL), 0, N_SHELL - 1);
};
for (size_t i = 0; i < npredicted; ++i)
if (rough[i].ok) rough[i].shell = shell_of(predicted[i].d);
// --- Learn the profile per shell (+ global) from the strong spots. ---
// Two things are learned. The empirical profile is the average grid in the DETECTOR frame, which is
// where it is applied. The width is a pair of second moments taken in each spot's OWN radial /
// tangential frame: a grid stacked in the detector frame is azimuthally averaged, so its <r^2> is
// sigma_r^2 + sigma_t^2 with no way back, and a radially smeared spot reads as a wide TANGENTIAL
// one. Rotating each contribution into the spot's frame keeps the two apart.
struct Moments { double rad = 0.0, tan = 0.0, w = 0.0; };
struct Sigma2 { double rad = 1.0, tan = 1.0; };
std::vector<std::vector<double>> shell_grid(N_SHELL, std::vector<double>(GG, 0.0));
std::vector<Moments> shell_mom(N_SHELL);
std::vector<int> shell_n(N_SHELL, 0);
std::vector<double> global_grid(GG, 0.0);
Moments global_mom;
int global_n = 0;
for (size_t i = 0; i < npredicted; ++i) {
const auto &rh = rough[i];
if (!rh.ok || !rh.strong || rh.I <= 0.0) continue;
const double rx = predicted[i].predicted_x - beam_x, ry = predicted[i].predicted_y - beam_y;
const double Rpx = std::hypot(rx, ry);
const double ux = Rpx > 1e-6 ? rx / Rpx : 1.0, uy = Rpx > 1e-6 ? ry / Rpx : 0.0;
for (int dy = -R; dy <= R; ++dy)
for (int dx = -R; dx <= R; ++dx) {
const int x = rh.cx + dx, y = rh.cy + dy;
if (x < 0 || y < 0 || x >= W || y >= H) continue;
const int32_t px = img[y * W + x];
if (!valid(px)) continue;
if (exclude && !clean(x, y, i)) continue;
const double v = (static_cast<double>(px) - rh.bkg) / rh.I;
shell_grid[rh.shell][grid_idx(dx, dy)] += v;
global_grid[grid_idx(dx, dy)] += v;
if (dx * dx + dy * dy >= r1_sq) continue;
const double rad = dx * ux + dy * uy, tn = -dx * uy + dy * ux;
shell_mom[rh.shell].rad += v * rad * rad;
shell_mom[rh.shell].tan += v * tn * tn;
shell_mom[rh.shell].w += v;
global_mom.rad += v * rad * rad;
global_mom.tan += v * tn * tn;
global_mom.w += v;
}
++shell_n[rh.shell];
++global_n;
}
// Radial and tangential variances from the moments. The domain is the r1 disk, which is
// azimuthally symmetric and so adds no anisotropy of its own. The cells are SIGNED: away from the
// peak a learned cell is background noise centred on zero, and clamping it at zero turns that
// noise into a positive pedestal that the rad^2 / tan^2 weights read as extra width.
auto widths = [](const Moments &m) {
Sigma2 s;
if (m.w > 0.0) {
s.rad = std::max(0.25, m.rad / m.w);
s.tan = std::max(0.25, m.tan / m.w);
}
return s;
};
// Normalised empirical profile (sum = 1), the average grid over the strong spots of a shell.
// ProfileGaussian does not use it - it rebuilds a per-reflection Gaussian in Pass B.
// Stored in float, as the fit below reads it (and as the GPU engine keeps it).
auto build_profile = [&](const std::vector<double> &grid) {
std::vector<float> P(GG, 0.0f);
double sum = 0.0;
for (int k = 0; k < GG; ++k)
sum += std::max(0.0, grid[k]);
if (sum > 0.0)
for (int k = 0; k < GG; ++k)
P[k] = static_cast<float>(std::max(0.0, grid[k]) / sum);
return P;
};
const std::vector<float> global_P = empirical && global_n > 0 ? build_profile(global_grid)
: std::vector<float>(GG, 0.0f);
const Sigma2 global_sigma2 = widths(global_mom);
std::vector<std::vector<float>> shell_P(N_SHELL, global_P);
std::vector<Sigma2> shell_sigma2(N_SHELL, global_sigma2);
for (int s = 0; s < N_SHELL; ++s) {
if (shell_n[s] < MIN_STRONG_PER_SHELL) continue;
if (empirical) shell_P[s] = build_profile(shell_grid[s]);
shell_sigma2[s] = widths(shell_mom[s]);
}
// --- Pass B: profile-fit each reflection (Kabsch, de-biased variance v = B + I*P; iterate). The
// reweighting is the Kabsch/Otwinowski iteration: Kabsch, Acta Cryst D66, 133-144 (2010);
// Otwinowski & Minor, Methods Enzymol 276, 307-326 (1997).
// Per reflection in three steps: what the fit needs, worked out in double (BraggFitSetup); the
// fit in single precision (FitProfile, or the Metal engine's kernel), with the GPU engine's
// arithmetic (BraggIntegrationEngineGPU.cu, fit) - the profile from exp with the widths'
// reciprocals taken once, normalised by a reciprocal, and one reciprocal of the variance for the
// three weights; and its error, the summation guard and the result, in double. ---
std::vector<float> shell_profiles; // the learned profiles, shell by shell (ProfileEmpirical)
if (empirical)
for (int s = 0; s < N_SHELL; ++s)
shell_profiles.insert(shell_profiles.end(), shell_P[s].begin(), shell_P[s].end());
std::vector<BraggFitSetup> setups(npredicted);
uint32_t cells = 0;
for (size_t i = 0; i < npredicted; ++i) {
const auto &rh = rough[i];
BraggFitSetup &fs = setups[i];
fs = {};
if (!rh.ok) continue;
// An overloaded reflection is not measured by any fit - its brightest pixels are what is
// missing. It is passed on with its box sum and the flag, for the merge to drop its rocking
// event (Reflection::overloaded); the missing-peak cut below would otherwise drop just this
// frame and leave the event to be extrapolated from its flanks.
if (rh.overloaded) {
results[i] = {static_cast<float>(rh.I), static_cast<float>(rh.sigma), static_cast<float>(rh.bkg),
static_cast<float>(rh.obs_x), static_cast<float>(rh.obs_y),
static_cast<float>(rh.var_bkg), true, false, true};
continue;
}
const int sh = rh.shell < 0 ? 0 : rh.shell;
const BraggStencil st = MakeBraggStencil(predicted[i].predicted_x, predicted[i].predicted_y, stencil);
fs.fit = 1;
fs.px = predicted[i].predicted_x;
fs.py = predicted[i].predicted_y;
fs.cx = rh.cx;
fs.cy = rh.cy;
fs.Rf = R;
fs.shell = sh;
fs.footprint = st.fp_s2t > 0.0f;
if (!empirical) {
const double rx = predicted[i].predicted_x - beam_x, ry = predicted[i].predicted_y - beam_y;
const double Rpx = std::hypot(rx, ry);
const double tan2t = Rpx / F_px;
double s2t = shell_sigma2[sh].tan;
double s2r = s2t, ux = 1.0, uy = 0.0;
bool elong = false;
if (use_ellipse) {
// Radial excess over the tangential width: measured where the peak is resolved inside
// the r1 disk, with the analytic bandwidth + parallax/capture term as the floor. The
// analytic term is what carries a streak the disk is too small to measure.
const double sbw = bw_sigma * Rpx;
const double radial_extra = std::max(shell_sigma2[sh].rad - s2t,
sbw * sbw + c_radial * tan2t * tan2t);
if (Rpx > 1e-6 && radial_extra > 0.25) {
ux = rx / Rpx; uy = ry / Rpx;
s2r = s2t + radial_extra;
elong = true;
}
}
// A spot that outgrows the r1 disk has widths the disk cannot measure - the moments above
// saturate near r1^2/4 - so the measured footprint sets them instead (SpotFootprint.h).
if (st.fp_s2t > 0.0f) {
ux = st.ux; uy = st.uy;
s2t = std::max<double>(s2t, st.fp_s2t);
s2r = std::max<double>(s2r, st.fp_s2r);
elong = true;
}
// The Gaussian is built per reflection, centred on the sub-pixel predicted position and (when
// needed) elongated, on a grid grown to hold the wider of its two axes.
fs.Rf = elong ? std::min(3 * R, static_cast<int>(std::ceil(r2 + 2.0 * std::sqrt(std::max(s2r, s2t))))) : R;
fs.shell = -1;
fs.cells = cells;
cells += (2 * fs.Rf + 1) * (2 * fs.Rf + 1);
fs.fx = static_cast<float>(predicted[i].predicted_x - rh.cx);
fs.fy = static_cast<float>(predicted[i].predicted_y - rh.cy);
fs.ux = static_cast<float>(ux);
fs.uy = static_cast<float>(uy);
fs.inv_2s2r = 0.5f / static_cast<float>(s2r);
fs.inv_2s2t = 0.5f / static_cast<float>(s2t);
}
fs.bkg = static_cast<float>(rh.bkg);
fs.B = static_cast<float>(std::max(rh.bkg, PIXEL_VARIANCE_FLOOR));
fs.v_min = static_cast<float>(WEIGHT_VARIANCE_MIN_FRACTION) * fs.B;
fs.I0 = static_cast<float>(rh.I);
}
std::vector<BraggFitOut> fits(npredicted);
if (!FitOnDevice(setups, shell_profiles, fits)) {
std::vector<float> cells_buf;
std::vector<FitPixel> fit_pixels;
for (const uint32_t i : FrameOrder(predicted, npredicted))
if (setups[i].fit)
fits[i] = FitProfile(img, owner_map, setups[i], i, shell_profiles, cells_buf, fit_pixels);
}
for (size_t i = 0; i < npredicted; ++i) {
if (!setups[i].fit) continue;
const auto &rh = rough[i];
const BraggFitOut &fo = fits[i];
if (!FitKept(fo) || !(fo.den > 0.0f)) continue;
const float den = fo.den, wsum = fo.wsum;
double I = fo.I;
// Guard against profile-fit runaways: on a weak / near-zero reflection the reweighted Kabsch
// iteration has no real peak to lock onto and can manufacture intensity the box sum never sees.
// Keep the profile intensity only if it agrees with the summation seed within the margin;
// otherwise fall back to the summation, which is robust there.
// 1/den is the profile-fit variance with the background taken as exact. The fit is
// I = sum(P*(px-bkg)/v) / sum(P^2/v), so dI/dbkg = -wsum/den and the background estimate's
// own error adds (wsum/den)^2 * var(bkg) - the same term the box sum was missing.
double sigma = std::sqrt(1.0 / den + (wsum / den) * (wsum / den) * rh.bkg_var);
// var_bkg is the NON-SIGNAL part of that variance, and 1/den is the fit's variance taken at the
// fitted intensity, so the signal part to remove is I itself - not max(0, I). Clamping it leaves
// a down-fluctuated reflection's own (deflated) variance standing as its background variance,
// which is 2-3x too small; the merge then weights exactly the down-fluctuated observations up.
// The whole point of a separate var_bkg is that it does not move with the observation's own
// fluctuation, and 1/den - I is what does not (1/den grows with I one for one).
double var_bkg = std::max(0.0, 1.0 / den - I
+ (wsum / den) * (wsum / den) * rh.bkg_var);
// The seed is a sum over the disk the box sum actually read, so when pixels are missing from
// both - excluded to a neighbour, or unreadable - the fit's full-profile intensity has to be
// scaled down to that same disk before the two are comparable. Nothing dropped gives 1.
const double guard_scale = fo.m_all > 0.0f ? fo.m_read / fo.m_all : 1.0;
if (std::abs(I * guard_scale - rh.I) > PROFILE_SUMMATION_MAX_NSIGMA * rh.sigma) {
I = rh.I;
sigma = rh.sigma;
var_bkg = rh.var_bkg;
++counts.profile_fallback;
}
// Carry the Pass-A box-sum intensity-weighted centroid (observed spot position) through the
// profile path too - post-refinement uses it as the observed position (beam-centre / distance).
results[i] = {static_cast<float>(I), static_cast<float>(sigma),
static_cast<float>(rh.bkg),
static_cast<float>(rh.obs_x), static_cast<float>(rh.obs_y),
static_cast<float>(var_bkg), true, rh.has_obs, rh.overloaded};
}
return Finalize(predicted, npredicted, results, image_number);
}
template std::vector<Reflection> BraggIntegrationEngineCPU::ProfilePasses(
const BraggIntegrationEngineCPU::BufferSampler &, const BraggIntegrationEngineCPU::FlatOwnerMap &,
const std::vector<Reflection> &, size_t, BraggIntegrationEngineCPU::PassA &, int64_t);
std::vector<Reflection> BraggIntegrationEngineCPU::Run(const ImagePreprocessorBuffer &image,
const std::vector<Reflection> &predicted,
size_t npredicted, int64_t image_number) {
if (image.size() != npixel)
return Finalize(predicted, npredicted, std::vector<BraggFitResult>(npredicted), image_number);
return RunImpl(BufferSampler{image.data()}, predicted, npredicted, image_number);
}
std::vector<Reflection> BraggIntegrationEngineCPU::Run(const CompressedImage &image,
const std::vector<Reflection> &predicted,
size_t npredicted, int64_t image_number) {
if (image.GetWidth() * image.GetHeight() != npixel)
return Finalize(predicted, npredicted, std::vector<BraggFitResult>(npredicted), image_number);
std::vector<uint8_t> scratch;
const auto *ptr = image.GetUncompressedPtr(scratch);
switch (image.GetMode()) {
case CompressedImageMode::Int8:
return RunImpl(ImageSampler<int8_t>{reinterpret_cast<const int8_t *>(ptr), INT8_MIN, INT8_MAX},
predicted, npredicted, image_number);
case CompressedImageMode::Int16:
return RunImpl(ImageSampler<int16_t>{reinterpret_cast<const int16_t *>(ptr), INT16_MIN, INT16_MAX},
predicted, npredicted, image_number);
case CompressedImageMode::Int32:
return RunImpl(ImageSampler<int32_t>{reinterpret_cast<const int32_t *>(ptr), INT32_MIN, INT32_MAX},
predicted, npredicted, image_number);
case CompressedImageMode::Uint8:
return RunImpl(ImageSampler<uint8_t>{reinterpret_cast<const uint8_t *>(ptr), UINT8_MAX, UINT8_MAX},
predicted, npredicted, image_number);
case CompressedImageMode::Uint16:
return RunImpl(ImageSampler<uint16_t>{reinterpret_cast<const uint16_t *>(ptr), UINT16_MAX, UINT16_MAX},
predicted, npredicted, image_number);
case CompressedImageMode::Uint32:
return RunImpl(ImageSampler<uint32_t>{reinterpret_cast<const uint32_t *>(ptr), UINT32_MAX, UINT32_MAX},
predicted, npredicted, image_number);
default:
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Image mode not supported");
}
}