Metal integration engine: the clip follows a handed-over radial curve and saturated pixels are counted, as on the CPU

The two pixel-level differences left between the Metal kernels and the CPU engine are closed:
- clip_sums follows the radial curve handed over by RadialBackgroundFromRings (RadialClipOffset, written
  out line for line as radial_clip_offset; the curve in buffer 11, BraggPassAFrame::n_curve), so the
  CPU fallback for such frames (3467c2191) is removed;
- box_sums counts the saturated signal-disk pixels (BraggBoxSums::n_saturated), so a still's overload
  kept as a lower bound counts them at the saturation limit, as on the CPU and CUDA.
Everything else - the gain (clip floor, sigma floor, fit floor, Finalize), the radial correction and D14 -
is the CPU engine's own host code on the Metal path already.

The structs the host hands the kernels carry static_asserts on their size in the shared headers, so the
host and the Metal compiler check the same numbers.

The [Metal] bit-identity test now covers the handed-over curve under every integrator, stills keeping
overloads (overloads_kept compared and required non-zero), gains 4 and 1.7, and a footprint still with a
curve; `clipped` is compared too. A [portable] CPU-only case checks everywhere that these cases reach
their modes. Not compiled here (no macOS): the Mac CI job is the check.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit 285d40c99b6356a52d224e610bc7b28d69973a06)

image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal

@@ -125,6 +125,8 @@ kernel void box_sums(device const float2 *xy [[buffer(0)]],
                                     && ((static_mask[(uint) idx / 32] >> ((uint) idx % 32)) & 1u);
                 if (v == PX_SATURATED || (v == PX_MASKED && on && !masked))
                     s.overloaded = 1;
+                if (v == PX_SATURATED)
+                    ++s.n_saturated;
                 if (!valid_px(v)) continue;
                 s.I_sum += v;
                 s.I_sum_x += (int64_t) x * v;
@@ -154,8 +156,17 @@ inline void atomic_add_i64(device atomic_uint *w, int64_t v) {
     atomic_fetch_add_explicit(&w[1], hi + (old + lo < old ? 1u : 0u), memory_order_relaxed);
 }
 
-// The background clip: the ring pixels at or below the threshold the CPU computed from box_sums, and
-// the radial background curve from the same pixels.
+// BraggIntegrationEngine::RadialClipOffset: how far the handed-over radial curve rises from the
+// reflection's radius r0 to a ring pixel at radial offset rad, in whole counts.
+inline int32_t radial_clip_offset(device const float *curve, int n, float r0, float rad) {
+    const int b = clamp((int) round(r0 + rad), 0, n - 1);
+    const int b0 = clamp((int) round(r0), 0, n - 1);
+    return (int32_t) round(curve[b] - curve[b0]);
+}
+
+// The background clip: the ring pixels at or below the threshold the CPU computed from box_sums - plus
+// the rise of the handed-over radial curve where one is in use (n_curve > 0) - and the radial
+// background curve from the same pixels where it is accumulated instead (n_rad > 0).
 kernel void clip_sums(device const float2 *xy [[buffer(0)]],
                       device const int32_t *img [[buffer(1)]],
                       device const uint8_t *mask [[buffer(2)]],
@@ -167,6 +178,7 @@ kernel void clip_sums(device const float2 *xy [[buffer(0)]],
                       device atomic_uint *rad_sum [[buffer(8)]],
                       device atomic_int *rad_cnt [[buffer(9)]],
                       device const BraggPanel *panels [[buffer(10)]],
+                      device const float *curve [[buffer(11)]],
                       uint i [[thread_position_in_grid]]) {
     if (i >= f.n || !clips[i])
         return;
@@ -185,7 +197,10 @@ kernel void clip_sums(device const float2 *xy [[buffer(0)]],
             const BraggStencilDist d = BraggStencilDistances(st, x - px, y - py);
             const int idx = y * f.width + x;
             const int32_t v = BraggPanelHolds(pn, x, y) ? img[idx] : PX_MASKED;
-            if (!(d.inner >= f.r2_sq && d.outer < f.r3_sq && valid_px(v) && mask[idx] != MASK_FLUX && v <= t))
+            if (!(d.inner >= f.r2_sq && d.outer < f.r3_sq && valid_px(v) && mask[idx] != MASK_FLUX))
+                continue;
+            const int32_t off = f.n_curve > 0 ? radial_clip_offset(curve, f.n_curve, st.r0, d.rad) : 0;
+            if (!((int64_t) v <= (int64_t) t + off))
                 continue;
             c.sum += v;
             ++c.n;

image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm

@@ -107,6 +107,7 @@ struct BraggIntegrationEngineMetal::Device {
     id<MTLCommandBuffer> unmark = nil;
     bool clear = false;
     id<MTLBuffer> rad_sum = nil, rad_cnt = nil;   // the radial background curve, as the clip kernel adds it up
+    id<MTLBuffer> curve = nil;                    // the curve handed over, which the clip follows
     // Per reflection, grown with the frames.
     size_t capacity = 0;
     id<MTLBuffer> xy = nil, panels = nil, level = nil, sums = nil, clips = nil, thr = nil, clip_out = nil;
@@ -147,9 +148,7 @@ std::vector<Reflection> BraggIntegrationEngineMetal::Run(const ImagePreprocessor
                                                          const std::vector<Reflection> &predicted,
                                                          size_t npredicted, int64_t image_number) {
     // An image of another size, and a frame with nothing predicted, the CPU engine answers as it would.
-    // So does a frame whose radial curve was handed over and is in use: the ring's clip then follows the
-    // curve (RadialClipOffset), which the clip kernel here does not, and the CPU gives that result exactly.
-    if (image.size() != npixel || npredicted == 0 || (bkg_radial && !radial_curve.empty()))
+    if (image.size() != npixel || npredicted == 0)
         return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number);
 
     dev->img = nil;
@@ -223,9 +222,16 @@ std::vector<Reflection> BraggIntegrationEngineMetal::Run(const ImagePreprocessor
         f.use_owner = use_owner;
         f.exclude = overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum;
         // The curve is added up on the GPU only where none was handed over (RadialBackgroundFromRings):
-        // StartPassA sizes its sums for that case alone. A handed-over curve in use never gets here (above).
+        // StartPassA sizes its sums for that case alone.
         const bool accumulate_radial = bkg_radial && radial_curve.empty();
         f.n_rad = accumulate_radial ? a.n_rad : 0;
+        // Where a handed-over curve is in use the clip follows it (RadialClipOffset).
+        f.n_curve = bkg_radial && !radial_curve.empty() ? static_cast<int32_t>(radial_curve.size()) : 0;
+        if (f.n_curve > 0) {
+            if (dev->curve == nil || dev->curve.length < radial_curve.size() * sizeof(float))
+                dev->curve = buffer(radial_curve.size() * sizeof(float));
+            std::copy(radial_curve.begin(), radial_curve.end(), static_cast<float *>(dev->curve.contents));
+        }
         f.static_mask_words = static_cast<uint32_t>(static_mask.size());
         f.n = static_cast<uint32_t>(npredicted);
 
@@ -314,6 +320,7 @@ std::vector<Reflection> BraggIntegrationEngineMetal::Run(const ImagePreprocessor
             [enc setBuffer:(accumulate_radial ? dev->rad_sum : dev->clips) offset:0 atIndex:8];   // not written without the curve
             [enc setBuffer:(accumulate_radial ? dev->rad_cnt : dev->clips) offset:0 atIndex:9];
             [enc setBuffer:dev->panels offset:0 atIndex:10];
+            [enc setBuffer:(f.n_curve > 0 ? dev->curve : dev->thr) offset:0 atIndex:11];   // not read without a curve
             [enc dispatchThreads:grid threadsPerThreadgroup:group];
             [enc endEncoding];
             Wait(cb);

image_analysis/bragg_integration/BraggPassA.h

@@ -46,3 +46,8 @@ struct BraggPassAFrame {
     uint32_t static_mask_words;  // length of the run's packed pixel mask
     uint32_t n;                  // reflections
 };
+
+// Compiled on both sides, so a layout the host and the Metal compiler disagree on fails to compile.
+static_assert(sizeof(BraggBoxSums) == 96, "BraggBoxSums layout shared with the Metal kernels");
+static_assert(sizeof(BraggClipSums) == 16, "BraggClipSums layout shared with the Metal kernels");
+static_assert(sizeof(BraggPassAFrame) == 52, "BraggPassAFrame layout shared with the Metal kernels");

image_analysis/bragg_integration/BraggPassB.h

@@ -46,6 +46,11 @@ struct BraggPassBFrame {
     uint32_t n;                    // reflections
 };
 
+// Compiled on both sides, so a layout the host and the Metal compiler disagree on fails to compile.
+static_assert(sizeof(BraggFitSetup) == 76, "BraggFitSetup layout shared with the Metal kernels");
+static_assert(sizeof(BraggFitOut) == 40, "BraggFitOut layout shared with the Metal kernels");
+static_assert(sizeof(BraggPassBFrame) == 28, "BraggPassBFrame layout shared with the Metal kernels");
+
 // The Gaussian profile's (unnormalised) value at grid offset (dx, dy).
 inline float BraggGaussian(int dx, int dy, float fx, float fy, float ux, float uy, float inv_2s2r,
                            float inv_2s2t) {

image_analysis/bragg_integration/BraggStencil.h

@@ -185,6 +185,10 @@ struct BraggPanel {
     float fx, fy, sx, sy;       // the plane step of one image pixel along x, and of one along y
     int32_t x0, y0, x1, y1;     // the panel's pixels in the image, inclusive
 };
+// The Metal engine hands both to its kernels as they lie in host memory, and this header is compiled on
+// both sides: a layout the two compilers disagree on fails here, on whichever side it is.
+static_assert(sizeof(BraggStencilParams) == 168, "BraggStencilParams layout shared with the Metal kernels");
+static_assert(sizeof(BraggPanel) == 40, "BraggPanel layout shared with the Metal kernels");
 
 BRAGG_STENCIL_HD bool BraggPanelHolds(const BRAGG_STENCIL_THREAD BraggPanel &pn, int x, int y) {
     return x >= pn.x0 && x <= pn.x1 && y >= pn.y0 && y <= pn.y1;

tests/BraggIntegrationEngineMetalTest.cpp

@@ -3,8 +3,6 @@
 
 #include <catch2/catch_all.hpp>
 
-#ifdef __APPLE__
-
 #include <bit>
 #include <cmath>
 #include <cstdint>
@@ -17,12 +15,16 @@
 #include "../common/PixelMask.h"
 #include "../common/Reflection.h"
 #include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h"
+#ifdef __APPLE__
 #include "../image_analysis/bragg_integration/BraggIntegrationEngineMetal.h"
+#endif
 
 // The Metal engine's whole claim is that it IS the CPU engine, bit for bit (BraggIntegrationEngineMetal.h).
 // So nothing here has a tolerance: every float of every reflection is compared as a bit pattern, and
 // every count as an integer. A Mac without a Metal device fails rather than skips - the macOS CI job is
-// the only place this can be checked, and a skip there would check nothing.
+// the only place this can be checked, and a skip there would check nothing. The frames and the cases
+// are built everywhere, and one test case checks on the CPU engine alone that they reach what they are
+// there for.
 
 namespace {
 
@@ -106,20 +108,34 @@ struct Case {
     float stencil_k = 0.0f;
     bool footprint = false;
     float clip_nsigma = 4.0f;
+    bool still = false;       // a still keeps its overloads as lower bounds (saturated pixels at the limit)
+    float gain = 1.0f;        // ADU per photon
+    bool curve = false;       // a radial curve handed over (RadialBackgroundFromRings), which the clip follows
 };
 
+// The radial curve a case hands over: rising steeply enough with the radius that the clip it makes the
+// ring follow (RadialClipOffset) moves by tens of counts across a ring.
+void HandOverCurve(BraggIntegrationEngine &engine) {
+    std::vector<float> rings(400);
+    for (size_t i = 0; i < rings.size(); ++i)
+        rings[i] = 1.0f + 0.01f * static_cast<float>(i * i);
+    engine.RadialBackgroundFromRings(rings, static_cast<int32_t>(rings.size()), 0.05f, 0.01f);
+}
+
 DiffractionExperiment MakeExperiment(const DetectorSetup &det, const Case &c) {
     DiffractionExperiment x(det);
     x.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV)
      .BeamX_pxl(0.4f * static_cast<float>(x.GetXPixelsNum()))
      .BeamY_pxl(0.55f * static_cast<float>(x.GetYPixelsNum()));
-    // Rotation: stills keep overloads as lower bounds on CPU/CUDA, but not yet on Metal (platform difference, to be fixed).
-    x.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}));
-    x.ImportIndexingSettings(IndexingSettings().RotationIndexing(true));
+    // A rotation drops an overloaded observation, a still keeps it as a lower bound (Finalize).
+    if (!c.still) {
+        x.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(-1, 0, 0), {}));
+        x.ImportIndexingSettings(IndexingSettings().RotationIndexing(true));
+    }
     x.BandwidthFWHM(c.bandwidth);
     BraggIntegrationSettings s;
     s.Integrator(c.mode).Overlap(c.overlap).BackgroundRadialCorrection(c.radial).StencilKSigma(c.stencil_k)
-     .BackgroundClipNSigma(c.clip_nsigma);
+     .BackgroundClipNSigma(c.clip_nsigma).PixelGain(c.gain);
     if (c.footprint) {
         // Wider than the r1 disk can hold from the middle out, so both stencils occur on one frame.
         SpotFootprint fp;
@@ -134,6 +150,38 @@ DiffractionExperiment MakeExperiment(const DetectorSetup &det, const Case &c) {
     return x;
 }
 
+}   // namespace
+
+// Runs everywhere, on the CPU engine alone: the cases the Metal comparison below adds for the stills, the
+// gain and the handed-over curve do reach what they are there for, so that comparison cannot pass by
+// comparing nothing.
+TEST_CASE("BraggIntegrationEngineMetal_CasesReachTheirModes", "[portable][Integration]") {
+    Case c;
+    c.still = true;
+    c.gain = 4.0f;
+    c.radial = true;
+    c.curve = true;
+    const DiffractionExperiment x = MakeExperiment(DetJF(2), c);
+    const PixelMask mask(x);
+    BraggIntegrationEngineCPU cpu(x, mask);
+    HandOverCurve(cpu);
+    cpu.BackgroundRadial(true);
+    // The clip follows the curve: it moves by whole counts across a ring.
+    CHECK(cpu.RadialClipOffset(500.0f, 10.0f) > 10);
+    CHECK(cpu.RadialClipOffset(500.0f, -10.0f) < -10);
+    ImagePreprocessorBuffer image(x.GetPixelsNum());
+    const Frame f = RandomFrame(x, 7, {});
+    std::copy(f.image.begin(), f.image.end(), image.data());
+    const auto out = cpu.Run(image, f.predicted, f.predicted.size(), 0);
+    CHECK(out.size() > 1000);
+    // Saturated pixels counted at the limit: overloads kept as lower bounds.
+    CHECK(cpu.Counts().overloads_kept > 0);
+}
+
+#ifdef __APPLE__
+
+namespace {
+
 uint32_t Bits(float f) { return std::bit_cast<uint32_t>(f); }
 
 void RequireSame(const std::vector<Reflection> &cpu, const std::vector<Reflection> &metal) {
@@ -156,6 +204,7 @@ void RequireSame(const std::vector<Reflection> &cpu, const std::vector<Reflectio
         REQUIRE(Bits(a.prescaling_corr) == Bits(b.prescaling_corr));
         REQUIRE(a.observed == b.observed);
         REQUIRE(a.overloaded == b.overloaded);
+        REQUIRE(a.clipped == b.clipped);
     }
 }
 
@@ -166,6 +215,7 @@ void RequireSame(const BraggIntegrationCounts &a, const BraggIntegrationCounts &
     CHECK(a.profile_fallback == b.profile_fallback);
     CHECK(a.overlapped == b.overlapped);
     CHECK(a.overlapped_major == b.overlapped_major);
+    CHECK(a.overloads_kept == b.overloads_kept);
 }
 
 // Three frames through one pair of engines - so whatever the Metal engine keeps between frames (its
@@ -188,6 +238,10 @@ size_t CompareFrames(const DetectorSetup &det, const Case &c) {
 
     BraggIntegrationEngineCPU cpu(x, mask);
     BraggIntegrationEngineMetal metal(x, mask);
+    if (c.curve) {
+        HandOverCurve(cpu);
+        HandOverCurve(metal);
+    }
     ImagePreprocessorBuffer image(x.GetPixelsNum());
 
     size_t n_out = 0;
@@ -204,6 +258,9 @@ size_t CompareFrames(const DetectorSetup &det, const Case &c) {
         n_out += out_cpu.size();
     }
     RequireSame(cpu.Counts(), metal.Counts());
+    // Overloads kept as lower bounds count saturated pixels: the case is not vacuous there either.
+    if (c.still)
+        CHECK(cpu.Counts().overloads_kept > 0);
     // Not the CPU engine compared with itself: every frame went through the kernels.
     REQUIRE(metal.FramesOnGPU() == 3);
     return n_out;
@@ -244,6 +301,39 @@ TEST_CASE("BraggIntegrationEngineMetal_BitIdenticalToCPU", "[portable][Integrati
     Case no_clip;
     no_clip.clip_nsigma = 0.0f;
     cases.emplace_back("no clip", no_clip);
+    // A radial curve handed over, which the ring's clip follows (RadialClipOffset), under every overlap
+    // mode and every integrator.
+    for (IntegratorMode mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical})
+        for (OverlapMode overlap : {OverlapMode::Off, OverlapMode::Exclude}) {
+            Case curve;
+            curve.mode = mode;
+            curve.overlap = overlap;
+            curve.radial = true;
+            curve.curve = true;
+            cases.emplace_back("curve mode " + std::to_string(static_cast<int>(mode)) + " overlap "
+                               + std::to_string(static_cast<int>(overlap)), curve);
+        }
+    // Stills: overloads kept as lower bounds, every saturated pixel counted. At a gain of one, and in
+    // detector units at two gains (the clip, the floors and the output's units all move with it), with
+    // and without the handed-over curve.
+    for (IntegratorMode mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical})
+        for (float gain : {1.0f, 4.0f, 1.7f}) {
+            Case still;
+            still.mode = mode;
+            still.still = true;
+            still.gain = gain;
+            still.radial = gain != 1.0f;
+            still.curve = gain == 4.0f;
+            cases.emplace_back("still mode " + std::to_string(static_cast<int>(mode)) + " gain "
+                               + std::to_string(gain) + (still.radial ? " radial" : "")
+                               + (still.curve ? " curve" : ""), still);
+        }
+    // Footprint stencils on a still with a curve: the clip follows it on a per-reflection stencil.
+    Case footprint_still = footprint;
+    footprint_still.still = true;
+    footprint_still.curve = true;
+    footprint_still.gain = 4.0f;
+    cases.emplace_back("footprint still gain 4 curve", footprint_still);
 
     for (const auto &[name, c] : cases) {
         DYNAMIC_SECTION(name) {
@@ -253,30 +343,6 @@ TEST_CASE("BraggIntegrationEngineMetal_BitIdenticalToCPU", "[portable][Integrati
     }
 }
 
-// A radial curve handed over from the whole frame (RadialBackgroundFromRings) makes the ring's clip follow
-// it, which the clip kernel does not: such a frame is integrated by the CPU engine, so the result is the
-// CPU's exactly and no frame goes to the GPU.
-TEST_CASE("BraggIntegrationEngineMetal_HandedOverRadialCurveStaysOnCPU", "[portable][Integration][Metal]") {
-    Case c;
-    c.radial = true;
-    const DiffractionExperiment x = MakeExperiment(DetJF(2), c);
-    const PixelMask mask(x);
-    BraggIntegrationEngineCPU cpu(x, mask);
-    BraggIntegrationEngineMetal metal(x, mask);
-    std::vector<float> rings(400);
-    for (size_t i = 0; i < rings.size(); ++i)
-        rings[i] = 1.0f + 0.002f * static_cast<float>(i * i);
-    cpu.RadialBackgroundFromRings(rings, static_cast<int32_t>(rings.size()), 0.05f, 0.01f);
-    metal.RadialBackgroundFromRings(rings, static_cast<int32_t>(rings.size()), 0.05f, 0.01f);
-    ImagePreprocessorBuffer image(x.GetPixelsNum());
-    const Frame f = RandomFrame(x, 7, {});
-    std::copy(f.image.begin(), f.image.end(), image.data());
-    const auto out_cpu = cpu.Run(image, f.predicted, f.predicted.size(), 0);
-    RequireSame(out_cpu, metal.Run(image, f.predicted, f.predicted.size(), 0));
-    CHECK(out_cpu.size() > 1000);
-    CHECK(metal.FramesOnGPU() == 0);
-}
-
 // The trimmed-mean background sorts each ring, which only the CPU engine does: such a run keeps the CPU
 // engine (MXAnalysisWithoutFPGA asks Available).
 TEST_CASE("BraggIntegrationEngineMetal_TrimmedMeanStaysOnCPU", "[portable][Integration][Metal]") {