diff --git a/.gitea/workflows/build_and_test.yml b/.gitea/workflows/build_and_test.yml index b69ad25e5..516f54555 100644 --- a/.gitea/workflows/build_and_test.yml +++ b/.gitea/workflows/build_and_test.yml @@ -269,7 +269,15 @@ jobs: shell: bash run: | cmake --build build-macos --target jfjoch_portable_test -j "$(sysctl -n hw.ncpu)" - cd build-macos/tests && ./jfjoch_portable_test "[portable]" + cd build-macos/tests && ./jfjoch_portable_test "[portable]~[Metal]" + # The Metal Bragg integration engine has to give the CPU engine's result bit for bit + # (BraggIntegrationEngineMetal.h): every reflection's every float, on random crowded frames + # through every integrator, overlap and background mode. Its own step, so a mismatch is named as + # such; it fails rather than skips on a machine without a Metal device, as this is the only job + # that can check it. + - name: Metal integration engine is bit-identical to the CPU engine + shell: bash + run: cd build-macos/tests && ./jfjoch_portable_test "[Metal]" # One build, two artifacts: the viewer .dmg and rugnux as .tar.gz (CPACK_PROJECT_CONFIG_FILE # in CMakeLists.txt gives each generator its components). The XDS plugin is Linux-only. - name: Build viewer DMG and rugnux tgz diff --git a/CMakeLists.txt b/CMakeLists.txt index 1875640bc..54e09879b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -21,6 +21,14 @@ INCLUDE(GNUInstallDirs) SET(CMAKE_CXX_STANDARD 20) SET(CMAKE_CXX_STANDARD_REQUIRED True) +# The Metal Bragg integration engine (image_analysis/bragg_integration) is Objective-C++. Enabled here, +# above every target that links it, as enable_language asks. +IF (APPLE) + ENABLE_LANGUAGE(OBJCXX) + SET(CMAKE_OBJCXX_STANDARD 20) + SET(CMAKE_OBJCXX_STANDARD_REQUIRED True) +ENDIF() + # MSVC gets C++ exception handling only from /EHsc, and CMake puts that flag in the DEFAULT value of # CMAKE_CXX_FLAGS ("/DWIN32 /D_WINDOWS /GR /EHsc"), which is used only when the configure does not # set CMAKE_CXX_FLAGS itself. CI does (-DCMAKE_CXX_FLAGS="/arch:AVX"), and that drops /EHsc: MSVC then @@ -47,9 +55,10 @@ ENDIF() # came out differently from the same source and data. Off, every x86-64 level computes the same bits # (measured: identical p.mtz on nine battery sets, no measurable time cost), and aarch64, where FMA is # always there and GCC contracts by default, follows the same arithmetic. MSVC does not contract under -# its default /fp:precise. Host code of .cu sources gets it through nvcc. +# its default /fp:precise. Host code of .cu sources gets it through nvcc, and the Metal kernels say so in +# their source (BraggIntegrationEngineMetal.mm). IF (NOT MSVC) - ADD_COMPILE_OPTIONS($<$:-ffp-contract=off>) + ADD_COMPILE_OPTIONS($<$:-ffp-contract=off>) ADD_COMPILE_OPTIONS($<$:-Xcompiler=-ffp-contract=off>) ENDIF() # The same for device code. nvcc contracts by default and does it separately for every architecture diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index bbb72b242..6b5eb1272 100644 --- a/common/CUDAWrapper.cpp +++ b/common/CUDAWrapper.cpp @@ -3,6 +3,7 @@ #include "CUDAWrapper.h" +#include #include #include @@ -93,7 +94,17 @@ int32_t get_gpu_count() { return 0; } -void hide_gpus() {} +namespace { + std::atomic gpus_hidden{false}; +} + +void hide_gpus() { + gpus_hidden = true; +} + +bool gpus_are_hidden() { + return gpus_hidden; +} std::vector get_gpu_names() { return {}; diff --git a/common/CUDAWrapper.cu b/common/CUDAWrapper.cu index e3b1648c7..961e706ce 100644 --- a/common/CUDAWrapper.cu +++ b/common/CUDAWrapper.cu @@ -29,6 +29,10 @@ void hide_gpus() { gpus_hidden = true; } +bool gpus_are_hidden() { + return gpus_hidden; +} + int32_t get_gpu_count() { if (gpus_hidden) return 0; diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index 5f0aba705..b3a4cbf1f 100644 --- a/common/CUDAWrapper.h +++ b/common/CUDAWrapper.h @@ -16,6 +16,9 @@ int32_t get_gpu_count(); // CUDA_VISIBLE_DEVICES= does, asked for by the program (rugnux --no-gpu). Call it before anything uses // CUDA; there is no way back. void hide_gpus(); +// Whether hide_gpus() was called - for a GPU path that is not CUDA's (the Metal integration engine), so +// --no-gpu keeps the process off every GPU. +bool gpus_are_hidden(); // Names of the visible GPUs, in device order and one entry per device, so repeated cards repeat. // Empty without CUDA and on a machine with no device, which is also what get_gpu_count() == 0 says. diff --git a/common/JFJochException.h b/common/JFJochException.h index fa3c64809..50882528c 100644 --- a/common/JFJochException.h +++ b/common/JFJochException.h @@ -45,7 +45,8 @@ UDPError, CommunicationError, SimplonError, CalibrationError, -IndexerError +IndexerError, +GPUMetalError }; class JFJochException : public std::exception { @@ -128,6 +129,8 @@ protected: return "Detector calibration error"; case JFJochExceptionCategory::IndexerError: return "Indexer error"; + case JFJochExceptionCategory::GPUMetalError: + return "Metal (GPU) error"; default: return ""; } diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index be74d7ca6..9a4afd7d4 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -523,6 +523,13 @@ corrected once by the exact remainder, which returns the correctly rounded quoti P. W. Markstein, "Computation of elementary functions on the IBM RISC System/6000 processor" (1990), IBM J. Res. Dev. 34, 111-119. +The single-precision exp the Bragg integration's profile fit shares between the CPU and an Apple GPU +(`BraggExp.h`) is built from textbook error-free floating-point arithmetic: the argument reduction of +W. J. Cody and W. Waite, "Software Manual for the Elementary Functions" (1980), Prentice-Hall; the +exact product (with Veltkamp's split) and Fast2Sum of T. J. Dekker, "A floating-point technique for +extending the available precision" (1971), Numer. Math. 18, 224-242; and the 2Sum of D. E. Knuth, +"The Art of Computer Programming, Vol. 2: Seminumerical Algorithms", section 4.2.2. + This software uses the Viridis, Magma and Inferno colormaps from Matplotlib under its BSD-compatible license. J. D. Hunter, "Matplotlib: A 2D graphics environment" (2007), Comput. Sci. Eng. 9, 90-95 [doi:10.1109/MCSE.2007.55](https://doi.org/10.1109/MCSE.2007.55). diff --git a/docs/GPU_DATA_ANALYSIS.md b/docs/GPU_DATA_ANALYSIS.md index b48b2d260..4573c4a69 100644 --- a/docs/GPU_DATA_ANALYSIS.md +++ b/docs/GPU_DATA_ANALYSIS.md @@ -180,3 +180,16 @@ The placements are then refined, scored and checked against their null with the structure factors above; the likelihood scoring runs on the host. The search is described in [§14.10](CPU_DATA_ANALYSIS_DECISIONS.md#1410-molecular-replacement-placing-a-model-the-rigid-body-cannot), its keys and decision rules in [Molecular replacement](RUGNUX_REPORT.md#molecular-replacement). + +## Apple GPUs (Metal) + +A macOS build has no CUDA, and one engine runs on the Apple GPU instead: + +- **Bragg integration** (`BraggIntegrationEngineMetal.mm`, `.metal`) - the CPU engine with its + pixel work moved to the GPU: the neighbour mask, the ownership of shared pixels, the box + summation's walk over every reflection, the background clip and the single-precision profile fit. + Everything computed in double and the profile learning stay on the CPU, as Apple GPUs have no + double precision. The fit's Gaussian uses an exp of our own (`BraggExp.h`) on both, as no GPU exp + matches the CPU's. The same result as the CPU engine **to the bit**, which a macOS CI step checks on + every push; `rugnux --no-gpu` keeps it off the GPU. + [§9.2](CPU_DATA_ANALYSIS_INTEGRATION.md#92-box-summation-seed-and-fallback) diff --git a/image_analysis/MXAnalysisWithoutFPGA.cpp b/image_analysis/MXAnalysisWithoutFPGA.cpp index 67a29a0ea..61b9884f2 100644 --- a/image_analysis/MXAnalysisWithoutFPGA.cpp +++ b/image_analysis/MXAnalysisWithoutFPGA.cpp @@ -19,6 +19,9 @@ #include "spot_finding/ImageSpotFinderCPU.h" #include "spot_finding/AdaptiveSpotFinderCPU.h" #include "bragg_integration/BraggIntegrationEngineCPU.h" +#ifdef __APPLE__ +#include "bragg_integration/BraggIntegrationEngineMetal.h" +#endif #ifdef JFJOCH_USE_CUDA #include "azint/AzIntEngineGPU.h" #include "roi/ROIIntegrationGPU.h" @@ -53,7 +56,13 @@ MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_exp auto cpu_preprocessor = std::make_unique(in_experiment, in_mask); preprocessor_cpu = cpu_preprocessor.get(); preprocessor = std::move(cpu_preprocessor); - bragg_engine = std::make_unique(in_experiment, in_mask); +#ifdef __APPLE__ + // The CPU engine with its integer pixel work on the Apple GPU: the same result, bit for bit. + if (BraggIntegrationEngineMetal::Available(in_experiment)) + bragg_engine = std::make_unique(in_experiment, in_mask); + else +#endif + bragg_engine = std::make_unique(in_experiment, in_mask); if (experiment.ROI().size() >= 1) roi = std::make_unique(experiment); #ifdef JFJOCH_USE_CUDA diff --git a/image_analysis/bragg_integration/BraggExp.h b/image_analysis/bragg_integration/BraggExp.h new file mode 100644 index 000000000..818afd4dc --- /dev/null +++ b/image_analysis/bragg_integration/BraggExp.h @@ -0,0 +1,63 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// exp in single precision, the same bits wherever it is compiled: the host compiler (no FMA contraction, +// see CMakeLists.txt) and the Metal compiler (fast math off) both see only IEEE +, -, * rounded to +// nearest and exact scaling by a power of two, so the profile fit's Gaussian is the same on the CPU and +// on an Apple GPU, whose exp no host function reproduces. It is closer to correctly rounded than macOS's +// expf: at most 0.53 ulp for a normal result, and away from a correctly rounded float for 0.2% of +// arguments below -1 and 0.004% in [-1, 0]. +// +// exp(x) = 2^k exp(r), k = rint(x / ln2), r = x - k ln2 with ln2 split three ways, the first two parts +// of 16 bits so that k times them is exact (Cody & Waite). r is carried as rr + re, rr^2 is formed exactly +// as zh + zl (Dekker's product with Veltkamp's split) and the sums into 1 + rr + rr^2/2 are error-free +// (2Sum, Fast2Sum), so the only rounding that matters is the last one. The rest of the series is a +// degree-4 minimax fit, relative error 1.3e-10. + +#ifdef __METAL_VERSION__ +inline float BraggBitsToFloat(int i) { return as_type(i); } +#else +#include +#include +inline float BraggBitsToFloat(int32_t i) { return std::bit_cast(i); } +#endif + +// p * 2^k, rounded once where the result is subnormal. +inline float BraggScale2(float p, int k) { + if (k < -125) + return (p * BraggBitsToFloat((k + 100 + 127) << 23)) * BraggBitsToFloat((-100 + 127) << 23); + return p * BraggBitsToFloat((k + 127) << 23); +} + +// Following Cody & Waite (1980) Software Manual for the Elementary Functions; Dekker (1971) Numer. Math. +// 18, 224-242; Knuth, The Art of Computer Programming Vol. 2, 4.2.2 (2Sum). +inline float BraggExp(float x) { + if (x < -104.0f) return 0.0f; // below half the smallest subnormal + if (x > 88.72283f) return BraggBitsToFloat(0x7f800000); // overflows + const float shifter = 12582912.0f; // 1.5 * 2^23: adding it rounds to an integer + const float kf = (x * 1.44269504f + shifter) - shifter; + const int k = static_cast(kf); + const float r1 = x - kf * 0.693145751953125f; // exact + const float t2 = kf * 1.4286197256296873e-06f; // exact + const float rr = r1 - t2; // r1 - t2 = rr + e2 (2Sum) + const float bv0 = rr - r1; + const float e2 = (r1 - (rr - bv0)) + (-t2 - bv0); + const float re = e2 - kf * -1.2905320041778356e-11f; + const float sp = rr * 4097.0f; // rr^2 = zh + zl + const float rh = sp - (sp - rr); + const float rl = rr - rh; + const float zh = rr * rr; + const float zl = ((rh * rh - zh) + (rh * rl + rh * rl)) + rl * rl; + const float h = 0.1666666567325592f + rr * (0.04166644066572189f + rr * (0.008333569392561913f + + rr * (0.0013937620678916574f + rr * 0.0001974606711883098f))); + const float c = 0.5f * zl + ((re + re * rr) + (zh * rr) * h); + const float hz = 0.5f * zh; // rr + zh/2 = s1 + e1 (2Sum) + const float s1 = rr + hz; + const float bv1 = s1 - rr; + const float e1 = (rr - (s1 - bv1)) + (hz - bv1); + const float t1 = 1.0f + s1; // 1 + s1 = t1 + e3 (Fast2Sum, |s1| < 1) + const float e3 = (1.0f - t1) + s1; + return BraggScale2(t1 + (e3 + (e1 + c)), k); +} diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp index b62ab2aa9..aa574ae17 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp @@ -36,13 +36,6 @@ inline void prefetch_row(const void *first, const void *last) { prefetch(last); } -// Identity sampler over the preprocessed int32 buffer (already in that convention). -struct BufferSampler { - const int32_t *p; - int32_t operator[](size_t i) const { return p[i]; } - const void *address(size_t i) const { return p + i; } -}; - // 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. @@ -65,9 +58,9 @@ struct ImageSampler { BraggIntegrationEngineCPU::BraggIntegrationEngineCPU(const DiffractionExperiment &experiment, const PixelMask &mask) : BraggIntegrationEngine(experiment), - static_mask(mask.GetPackedMask()), refl_mask(static_cast(xpixel), static_cast(ypixel), 0), - owner(static_cast(xpixel), static_cast(ypixel), BRAGG_OWNER_NONE) {} + owner(static_cast(xpixel), static_cast(ypixel), BRAGG_OWNER_NONE), + static_mask(mask.GetPackedMask()) {} void BraggIntegrationEngineCPU::MarkNeighbours(const std::vector &predicted, size_t npredicted) { const int W = static_cast(xpixel), H = static_cast(ypixel); @@ -133,23 +126,13 @@ template std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, const std::vector &predicted, size_t npredicted, int64_t image_number) { - std::vector results(npredicted); if (npredicted == 0) - return Finalize(predicted, npredicted, results, image_number); + return Finalize(predicted, npredicted, std::vector(npredicted), image_number); counts.predicted += npredicted; const int W = static_cast(xpixel), H = static_cast(ypixel); const bool do_clip = bkg_clip_nsigma > 0.0f && mode != IntegratorMode::BoxSum; - // 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 below is in force, which - // is the default - the two are alternatives. - const double bkg_trim_frac = bkg_trim; - - auto grid_idx = [this](int dx, int dy) { return (dy + R) * G + (dx + R); }; - MarkNeighbours(predicted, npredicted); owner.Clear(); @@ -184,20 +167,6 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, }; // --- Pass A: box-sum every reflection (rough I, background, centroid, strong flag). --- - struct Rough { - double I = 0.0, sigma = NAN, bkg = 0.0, obs_x = 0.0, obs_y = 0.0; - double bkg_var = 0.0; // variance of the background ESTIMATE itself, bkg / n_bkg - double var_bkg = 0.0; // total non-signal variance carried to the merge - int64_t I_sum = 0; // kept so I can be rebuilt after the radial background correction - int n_inner = 0; - int n_disk = 0, n_own = 0; // signal-disk pixels, and how many of them are this reflection's - int r_bin = 0; // rounded distance from the beam centre, indexes the radial curve - int k_bin = 0; // which radial-background kernel this reflection's stencil needs - int cx = 0, cy = 0, shell = -1; - bool ok = false, strong = false, has_obs = false; - bool full = false; // every pixel of the signal disk was readable - bool overloaded = false; // a pixel of the signal disk was saturated - }; // 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 @@ -205,26 +174,13 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, auto static_masked = [&](size_t idx) { return idx / 32 < static_mask.size() && ((static_mask[idx / 32] >> (idx % 32)) & 1U); }; - std::vector rough(npredicted); - double inv_d2_min = std::numeric_limits::max(), inv_d2_max = 0.0; + PassA a = StartPassA(npredicted); std::vector 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 - in the order the clip sums them, so it sums the same values. + // A walks the box anyway. struct ClipPixel { int32_t px; float rad; }; std::vector clip_ring; - // Radial background curve, accumulated from the annulus pixels this pass already reads. 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. - const int n_rad = static_cast(std::ceil(std::hypot(std::max(beam_x, W - beam_x), - std::max(beam_y, H - beam_y)))) + 2; - std::vector rad_sum; - std::vector rad_cnt; - if (bkg_radial) { - rad_sum.assign(n_rad, 0.0); - rad_cnt.assign(n_rad, 0); - } - // 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; }; @@ -247,28 +203,18 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, img.address(static_cast(y) * W + b.x1)); } const auto &r = predicted[i]; - Rough out; // 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; - // Both from the stencil's own radius, so nothing downstream is derived from a second one. - out.r_bin = std::clamp(static_cast(std::lround(st.r0)), 0, n_rad - 1); - out.k_bin = BraggStencilKernelIndex(st, n_kern); - - int64_t I_sum = 0, I_sum_x = 0, I_sum_y = 0, n_inner = 0, n_inner_valid = 0; - int64_t x_sum = 0, y_sum = 0; // positions of the pixels behind I_sum, for the centroid - int n_disk = 0, n_own = 0; // pixels in the signal disk, and how many are this reflection's - double bkg_sum = 0.0; - int n_bkg = 0; - // 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 + // 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. - int n_ring = 0, n_clean = 0; + BraggBoxSums s{}; bkg_vals.clear(); clip_ring.clear(); for (int y = y0; y <= y1; ++y) @@ -282,154 +228,324 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, // 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. - ++n_disk; + ++s.n_disk; if (overlap != OverlapMode::Off) { - if (clean(x, y, i)) ++n_own; + if (clean(x, y, i)) ++s.n_own; else if (exclude) continue; } - ++n_inner; + ++s.n_inner; if (px == INT32_MAX || (px == INT32_MIN && !static_masked(static_cast(y) * W + x))) - out.overloaded = true; + s.overloaded = 1; if (!valid(px)) continue; - I_sum += px; - I_sum_x += static_cast(x) * px; - I_sum_y += static_cast(y) * px; + s.I_sum += px; + s.I_sum_x += static_cast(x) * px; + s.I_sum_y += static_cast(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. - x_sum += x; - y_sum += y; - ++n_inner_valid; + 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; - ++n_ring; + ++s.n_ring; const uint8_t m = refl_mask.Get(x, y); - if (m == 0) ++n_clean; + if (m == 0) ++s.n_clean; if (m == MASK_FLUX) continue; - bkg_sum += static_cast(px); - if (bkg_trim_frac > 0.0) bkg_vals.push_back(px); - ++n_bkg; + s.bkg_sum += px; + if (bkg_trim > 0.0f) bkg_vals.push_back(px); + ++s.n_bkg; } } - if (use_owner && n_own < n_disk) { - ++counts.overlapped; - if (n_own < overlap_min_peak * n_disk) ++counts.overlapped_major; - } - - int n_bkg_used = 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 = n_inner_valid == 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 && 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 && n_clean <= 5 && n_ring > 5) ++counts.bkg_starved_by_neighbour; - if (keep_partial && n_bkg > 5) { - out.bkg = bkg_sum / n_bkg; - if (bkg_trim_frac > 0.0 && bkg_vals.size() > 5 - && bkg_vals.size() <= static_cast(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(bkg_vals.size() * bkg_trim_frac); - const size_t hi = bkg_vals.size() - lo; - if (hi > lo) { - double s = 0.0; - for (size_t t = lo; t < hi; ++t) s += bkg_vals[t]; - out.bkg = s / static_cast(hi - lo); - n_bkg_used = static_cast(hi - lo); - } - } else if (do_clip) { - // One high-outlier sigma-clip pass on the background ring: 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. - const double thr = out.bkg + bkg_clip_nsigma * std::sqrt(static_cast(bkg_pixel_gain) - * std::max(out.bkg, 1.0)); - double s = 0.0; - int n = 0; - for (const auto &[px, rad] : clip_ring) - if (static_cast(px) <= thr) { - s += px; - ++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(std::lround(st.r0 + rad)), 0, n_rad - 1); - rad_sum[b] += static_cast(px); - ++rad_cnt[b]; - } + // 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(std::lround(st.r0 + rad)), 0, a.n_rad - 1); + a.rad_sum[b] += px; + ++a.rad_cnt[b]; } - if (n > 5) { out.bkg = s / n; n_bkg_used = 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(I_sum) - static_cast(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(n_inner_valid) * out.bkg - + static_cast(n_inner_valid) * n_inner_valid * out.bkg_var; - out.I_sum = I_sum; - out.n_inner = static_cast(n_inner_valid); - out.n_disk = n_disk; - out.n_own = n_own; - const double var_bkg_term = static_cast(n_inner_valid) * n_inner_valid * out.bkg_var; - out.sigma = 1.0; - if (I_sum > 0) { - out.sigma = std::max(out.sigma, std::sqrt(static_cast(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(I_sum) - n_inner_valid * out.bkg; - if (net > 0.0) { - out.obs_x = (static_cast(I_sum_x) - out.bkg * x_sum) / net; - out.obs_y = (static_cast(I_sum_y) - out.bkg * y_sum) / net; - } else { - // Nothing above background to take a centroid of; the raw one is all there is. - out.obs_x = static_cast(I_sum_x) / static_cast(I_sum); - out.obs_y = static_cast(I_sum_y) / static_cast(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(std::lround(r.predicted_x)); - out.cy = static_cast(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(r.d) * r.d); - inv_d2_min = std::min(inv_d2_min, inv_d2); - inv_d2_max = std::max(inv_d2_max, inv_d2); - } } - rough[i] = out; + 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(xpixel), H = static_cast(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(std::ceil(std::hypot(std::max(beam_x, W - beam_x), + std::max(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::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(s.bkg_sum) / s.n_bkg; + const double thr = bkg + bkg_clip_nsigma * std::sqrt(static_cast(bkg_pixel_gain) * std::max(bkg, 1.0)); + return static_cast(std::clamp(std::floor(thr), static_cast(INT32_MIN), + static_cast(INT32_MAX))); +} + +void BraggIntegrationEngineCPU::FinishPassA(PassA &a, size_t i, const Reflection &r, const BraggStencil &st, + const BraggBoxSums &s, const BraggClipSums &clip, + std::vector &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(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(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(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(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(hi - lo); + n_bkg_used = static_cast(hi - lo); + } + } else if (ClipsRing(s)) { + if (clip.n > 5) { out.bkg = static_cast(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(s.I_sum) - static_cast(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(s.n_inner_valid) * out.bkg + + static_cast(s.n_inner_valid) * s.n_inner_valid * out.bkg_var; + out.I_sum = s.I_sum; + out.n_inner = static_cast(s.n_inner_valid); + out.n_disk = s.n_disk; + out.n_own = s.n_own; + const double var_bkg_term = static_cast(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(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(s.I_sum) - s.n_inner_valid * out.bkg; + if (net > 0.0) { + out.obs_x = (static_cast(s.I_sum_x) - out.bkg * s.x_sum) / net; + out.obs_y = (static_cast(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(s.I_sum_x) / static_cast(s.I_sum); + out.obs_y = static_cast(s.I_sum_y) / static_cast(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(std::lround(r.predicted_x)); + out.cy = static_cast(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(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 +BraggFitOut BraggIntegrationEngineCPU::FitProfile(const Sampler &img, const OwnerMap &owner_map, + const BraggFitSetup &s, size_t i, + const std::vector &shell_profiles, + std::vector &cells, + std::vector &fit_pixels) const { + const int W = static_cast(xpixel), H = static_cast(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(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(y) * W + gx0), + img.address(static_cast(y) * W + gx1)); + cells.resize(static_cast(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(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(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 +std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler &img, const OwnerMap &owner_map, + const std::vector &predicted, + size_t npredicted, PassA &a, + int64_t image_number) { + const int W = static_cast(xpixel), H = static_cast(ypixel); + std::vector results(npredicted); + std::vector &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(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 @@ -443,8 +559,8 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, const float *kern = k_diff.data() + static_cast(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, n_rad - 1); - const double v = rad_cnt[b] > 0 ? rad_sum[b] / rad_cnt[b] : rh.bkg; + 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(a.rad_sum[b]) / a.rad_cnt[b] : rh.bkg; corr += static_cast(kern[k]) * v; } rh.bkg -= corr; // annulus mean -> mean over the signal disk @@ -473,8 +589,8 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, } auto shell_of = [&](float d) { - if (!(d > 0.0f) || inv_d2_max <= inv_d2_min) return 0; - const double t = (1.0 / (static_cast(d) * d) - inv_d2_min) / (inv_d2_max - inv_d2_min); + if (!(d > 0.0f) || a.inv_d2_max <= a.inv_d2_min) return 0; + const double t = (1.0 / (static_cast(d) * d) - a.inv_d2_min) / (a.inv_d2_max - a.inv_d2_min); return std::clamp(static_cast(t * N_SHELL), 0, N_SHELL - 1); }; for (size_t i = 0; i < npredicted; ++i) @@ -537,22 +653,22 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, }; // 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 &grid) { - std::vector P(GG, 0.0); + std::vector P(GG, 0.0f); double sum = 0.0; - for (int k = 0; k < GG; ++k) { - P[k] = std::max(0.0, grid[k]); - sum += P[k]; - } + for (int k = 0; k < GG; ++k) + sum += std::max(0.0, grid[k]); if (sum > 0.0) - for (double &p : P) p /= sum; + for (int k = 0; k < GG; ++k) + P[k] = static_cast(std::max(0.0, grid[k]) / sum); return P; }; - const std::vector global_P = empirical && global_n > 0 ? build_profile(global_grid) - : std::vector(GG, 0.0); + const std::vector global_P = empirical && global_n > 0 ? build_profile(global_grid) + : std::vector(GG, 0.0f); const Sigma2 global_sigma2 = widths(global_mom); - std::vector> shell_P(N_SHELL, global_P); + std::vector> shell_P(N_SHELL, global_P); std::vector shell_sigma2(N_SHELL, global_sigma2); for (int s = 0; s < N_SHELL; ++s) { if (shell_n[s] < MIN_STRONG_PER_SHELL) continue; @@ -562,14 +678,22 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, // --- 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). --- - std::vector Pbuf; - // The pixels the fit reads, gathered once while p_valid walks the grid: the profile value and the - // background-subtracted count, in the grid's order. - struct FitPixel { double Pp, signal; }; - std::vector fit_pixels; + // 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 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 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 @@ -582,10 +706,15 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, continue; } const int sh = rh.shell < 0 ? 0 : rh.shell; - - int Rf = R; - const std::vector *Pvec = &shell_P[sh]; 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); @@ -614,94 +743,41 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, s2r = std::max(s2r, st.fp_s2r); elong = true; } - // Build the Gaussian per reflection, centred on the sub-pixel predicted position and (when + // 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. - const double fx = predicted[i].predicted_x - rh.cx, fy = predicted[i].predicted_y - rh.cy; - Rf = elong ? std::min(3 * R, static_cast(std::ceil(r2 + 2.0 * std::sqrt(std::max(s2r, s2t))))) : R; - // 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, rh.cx - Rf), gx1 = std::min(W - 1, rh.cx + Rf); - for (int y = std::max(0, rh.cy - Rf); y <= std::min(H - 1, rh.cy + Rf); ++y) - prefetch_row(img.address(static_cast(y) * W + gx0), - img.address(static_cast(y) * W + gx1)); - const int Gf = 2 * Rf + 1; - Pbuf.assign(static_cast(Gf) * Gf, 0.0); - double gs = 0.0; - for (int dy = -Rf; dy <= Rf; ++dy) - for (int dx = -Rf; dx <= Rf; ++dx) { - const double ex = dx - fx, ey = dy - fy; - const double rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux; - const double g = std::exp(-rad * rad / (2.0 * s2r) - tn * tn / (2.0 * s2t)); - Pbuf[(dy + Rf) * Gf + (dx + Rf)] = g; - gs += g; - } - for (double &p : Pbuf) p /= gs; - Pvec = &Pbuf; + fs.Rf = elong ? std::min(3 * R, static_cast(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(predicted[i].predicted_x - rh.cx); + fs.fy = static_cast(predicted[i].predicted_y - rh.cy); + fs.ux = static_cast(ux); + fs.uy = static_cast(uy); + fs.inv_2s2r = 0.5f / static_cast(s2r); + fs.inv_2s2t = 0.5f / static_cast(s2t); } + fs.bkg = static_cast(rh.bkg); + fs.B = static_cast(std::max(rh.bkg, PIXEL_VARIANCE_FLOOR)); + fs.v_min = static_cast(WEIGHT_VARIANCE_MIN_FRACTION) * fs.B; + fs.I0 = static_cast(rh.I); + } - const int Gf = 2 * Rf + 1; + std::vector fits(npredicted); + if (!FitOnDevice(setups, shell_profiles, fits)) { + std::vector cells_buf; + std::vector fit_pixels; + for (size_t i = 0; i < npredicted; ++i) + if (setups[i].fit) + fits[i] = FitProfile(img, owner_map, setups[i], i, shell_profiles, cells_buf, fit_pixels); + } - // --- 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. --- - double p_grid = 0.0, p_valid = 0.0, p_own = 0.0, m_all = 0.0, m_read = 0.0; - double p_peak = 0.0, p_lost_peak = 0.0; - fit_pixels.clear(); - for (int dy = -Rf; dy <= Rf; ++dy) - for (int dx = -Rf; dx <= Rf; ++dx) { - const double Pp = (*Pvec)[(dy + Rf) * Gf + (dx + Rf)]; - if (Pp <= 0.0) continue; - const int x = rh.cx + dx, y = rh.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 below - // compares the two over the same pixels. - const bool in_disk = st.fp_s2t > 0.0f - ? BraggStencilInSignal(st, x - predicted[i].predicted_x, y - predicted[i].predicted_y, - (x - predicted[i].predicted_x) * (x - predicted[i].predicted_x) - + (y - predicted[i].predicted_y) * (y - predicted[i].predicted_y), - r1_sq) - : dx * dx + dy * dy < r1_sq; - p_grid += Pp; - p_peak = std::max(p_peak, Pp); - if (in_disk) m_all += Pp; - if (!valid(img[y * W + x])) { - p_lost_peak = std::max(p_lost_peak, Pp); - continue; - } - p_valid += Pp; - const bool own = clean(x, y, i); - if (own) p_own += Pp; - if (in_disk && (own || !exclude)) m_read += Pp; - if (own || !exclude) - fit_pixels.push_back({Pp, static_cast(img[y * W + x]) - rh.bkg}); - } - if (p_valid < overlap_min_peak * p_grid) continue; - // 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 (p_lost_peak > MINPK_MAX_MISSING_PEAK * p_peak) continue; - if (overlap == OverlapMode::Reject && p_own < overlap_min_peak) continue; - - const double B = std::max(rh.bkg, PIXEL_VARIANCE_FLOOR); - double I = rh.I, den = 0.0, wsum = 0.0; - for (int iter = 0; iter < 4; ++iter) { - double num = 0.0; - den = 0.0; - wsum = 0.0; - for (const auto &[Pp, signal] : fit_pixels) { - const double v = std::max(B + I * Pp, WEIGHT_VARIANCE_MIN_FRACTION * B); - num += Pp * signal / v; - den += Pp * Pp / v; - wsum += Pp / v; - } - if (den > 0.0) I = num / den; else break; - } - if (!(den > 0.0)) continue; + 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. @@ -722,7 +798,7 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, // 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 = m_all > 0.0 ? m_read / m_all : 1.0; + 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; @@ -740,6 +816,10 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, return Finalize(predicted, npredicted, results, image_number); } +template std::vector BraggIntegrationEngineCPU::ProfilePasses( + const BraggIntegrationEngineCPU::BufferSampler &, const BraggIntegrationEngineCPU::FlatOwnerMap &, + const std::vector &, size_t, BraggIntegrationEngineCPU::PassA &, int64_t); + std::vector BraggIntegrationEngineCPU::Run(const ImagePreprocessorBuffer &image, const std::vector &predicted, size_t npredicted, int64_t image_number) { diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h index 95b0a943a..e72497831 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h @@ -4,13 +4,15 @@ #pragma once #include "BraggIntegrationEngine.h" +#include "BraggPassA.h" +#include "BraggPassB.h" #include "../../common/PixelMask.h" class CompressedImage; // Plain-C++ reference/fallback engine: a faithful serial re-expression of BraggIntegrate2D (box // sum) and ProfileIntegrate2D (Kabsch profile fit) reading the preprocessed int32 image. Also the -// numeric oracle the CUDA engine is checked against. +// numeric oracle the CUDA engine is checked against, and the engine the Metal engine extends. class BraggIntegrationEngineCPU : public BraggIntegrationEngine { // Core integrator, templated on a pixel sampler so it reads either the preprocessed int32 buffer // or a raw CompressedImage of any pixel type - both presented per-pixel in the INT32_MIN(masked)/ @@ -53,9 +55,6 @@ class BraggIntegrationEngineCPU : public BraggIntegrationEngine { }; TiledFrame refl_mask; TiledFrame owner; - // The run's pixel mask, packed 32 pixels to a word (PixelMask::GetPackedMask). An unreadable pixel - // it does not explain was unreadable on this frame only - an overload (see Reflection::overloaded). - std::vector static_mask; // Marks every predicted reflection's r2 region in refl_mask (MASK_FLUX / MASK_TAIL). void MarkNeighbours(const std::vector &predicted, size_t npredicted); @@ -64,6 +63,84 @@ class BraggIntegrationEngineCPU : public BraggIntegrationEngine { std::vector RunImpl(const Sampler &img, const std::vector &predicted, size_t npredicted, int64_t image_number); + // Pass B's single-precision part for reflection i, as the Metal engine's kernel does it: the profile + // (the shell's learned one, or the Gaussian, built in `cells`), how much of it is readable, and the + // fit. fit_pixels is scratch: the pixels the fit reads, gathered once. + struct FitPixel { float Pp, signal; }; + template + BraggFitOut FitProfile(const Sampler &img, const OwnerMap &owner_map, const BraggFitSetup &s, size_t i, + const std::vector &shell_profiles, std::vector &cells, + std::vector &fit_pixels) const; + // Whether enough of a reflection's profile is readable, and its own, for the fit to be kept. + bool FitKept(const BraggFitOut &o) const; + +protected: + // The run's pixel mask, packed 32 pixels to a word (PixelMask::GetPackedMask). An unreadable pixel + // it does not explain was unreadable on this frame only - an overload (see Reflection::overloaded). + std::vector static_mask; + + // Identity sampler over the preprocessed int32 buffer (already in the engine's convention). + struct BufferSampler { + const int32_t *p; + int32_t operator[](size_t i) const { return p[i]; } + const void *address(size_t i) const { return p + i; } + }; + // The owner map as one frame-sized array, as the Metal engine builds it (TiledFrame's Get otherwise). + struct FlatOwnerMap { + const uint32_t *p; + size_t width; + uint32_t Get(int x, int y) const { return p[static_cast(y) * width + x]; } + }; + + // Pass A's result for one reflection: the box sum, its background, centroid and strong flag. + struct Rough { + double I = 0.0, sigma = NAN, bkg = 0.0, obs_x = 0.0, obs_y = 0.0; + double bkg_var = 0.0; // variance of the background ESTIMATE itself, bkg / n_bkg + double var_bkg = 0.0; // total non-signal variance carried to the merge + int64_t I_sum = 0; // kept so I can be rebuilt after the radial background correction + int n_inner = 0; + int n_disk = 0, n_own = 0; // signal-disk pixels, and how many of them are this reflection's + int r_bin = 0; // rounded distance from the beam centre, indexes the radial curve + int k_bin = 0; // which radial-background kernel this reflection's stencil needs + int cx = 0, cy = 0, shell = -1; + bool ok = false, strong = false, has_obs = false; + bool full = false; // every pixel of the signal disk was readable + bool overloaded = false; // a pixel of the signal disk was saturated + }; + // Everything pass A hands the profile passes. + struct PassA { + std::vector rough; + // Radial background curve, accumulated from the ring pixels the background clip keeps. Integer + // sums, so the order the pixels arrive in does not matter. + int n_rad = 0; + std::vector rad_sum; + std::vector rad_cnt; + double inv_d2_min = 0.0, inv_d2_max = 0.0; // resolution range of the complete reflections + }; + PassA StartPassA(size_t npredicted) const; + + // Whether pass A sigma-clips this reflection's background ring, and the clip threshold: the ring + // pixels at or below it are kept. A pixel value is an integer, so the threshold is one too. + bool ClipsRing(const BraggBoxSums &s) const; + int32_t ClipThreshold(const BraggBoxSums &s) const; + // Pass A's result for reflection i from what its box walk counted, and from its clip sums where + // ClipsRing. bkg_vals holds the ring pixels for the trimmed-mean background, and is sorted here. + void FinishPassA(PassA &a, size_t i, const Reflection &r, const BraggStencil &st, const BraggBoxSums &s, + const BraggClipSums &clip, std::vector &bkg_vals); + // Everything after pass A: the radial background correction, then either the box-sum results or + // the profile learning and pass B. Defined for BufferSampler with either owner map, and for the + // CompressedImage samplers with the tiled one. + // Pass B's single-precision part for every reflection with setups[i].fit, on a device - the + // learned profiles side by side in shell_profiles, empty for the Gaussian. False where there is no + // device, and the CPU fits them instead. + virtual bool FitOnDevice(const std::vector &setups, const std::vector &shell_profiles, + std::vector &fits) { return false; } + + template + std::vector ProfilePasses(const Sampler &img, const OwnerMap &owner_map, + const std::vector &predicted, size_t npredicted, + PassA &a, int64_t image_number); + public: BraggIntegrationEngineCPU(const DiffractionExperiment &experiment, const PixelMask &mask); diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h new file mode 100644 index 000000000..ffacb973f --- /dev/null +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h @@ -0,0 +1,50 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include + +#include "BraggIntegrationEngineCPU.h" + +// Apple GPU engine (macOS only): BraggIntegrationEngineCPU with its pixel work - the neighbour mask, the +// owner map, pass A's walk over every box, the background clip and pass B's single-precision profile +// fit - run as Metal kernels, one thread per reflection. Everything computed in double and the profile +// learning stay on the CPU, in the CPU engine's own code: Apple GPUs have no double precision. +// +// The result is the CPU engine's, bit for bit: the integer sums do not depend on the order the GPU +// visits pixels in, the fit's float sums are taken in the CPU's order, one reflection per thread, and +// the arithmetic deciding which pixel is which is the same source (BraggStencil.h, BraggPassB.h, and +// BraggExp.h for the Gaussian, as no GPU exp matches the host's), compiled for the GPU with fast math +// off - sqrt and division then round correctly - and, as on the host, with no a*b+c contracted into a +// fused multiply-add. The CPU and the GPU read the same memory, so the preprocessed image is not +// copied. tests/BraggIntegrationEngineMetalTest.cpp holds the engine to all of this. +class BraggIntegrationEngineMetal : public BraggIntegrationEngineCPU { + struct Device; + std::unique_ptr dev; + uint64_t frames_on_gpu = 0; + + bool FitOnDevice(const std::vector &setups, const std::vector &shell_profiles, + std::vector &fits) override; + +public: + // Whether this engine can integrate this experiment: a Metal device the process may use (rugnux + // --no-gpu hides it, see hide_gpus) and settings the kernels cover - not the trimmed-mean background, + // which only the CPU engine sorts the ring for. + static bool Available(const DiffractionExperiment &experiment); + // The Metal device's name, for the log; empty where there is none or the process may not use it. + static std::string DeviceName(); + + BraggIntegrationEngineMetal(const DiffractionExperiment &experiment, const PixelMask &mask); + ~BraggIntegrationEngineMetal() override; + + using BraggIntegrationEngineCPU::Run; // the CompressedImage overload stays the CPU engine's + + std::vector Run(const ImagePreprocessorBuffer &image, + const std::vector &predicted, size_t npredicted, + int64_t image_number) override; + + // The frames whose pixels the GPU walked, rather than the CPU engine this one falls back to. + [[nodiscard]] uint64_t FramesOnGPU() const { return frames_on_gpu; } +}; diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal new file mode 100644 index 000000000..4d2632fa3 --- /dev/null +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal @@ -0,0 +1,275 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +// Kernels of BraggIntegrationEngineMetal. Compiled at run time with fast math off, after BraggStencil.h +// and BraggPassA.h (BraggIntegrationEngineMetal.mm puts the source together). +// +// Each kernel is one of BraggIntegrationEngineCPU's loops over a reflection, one thread per reflection, +// written as the CPU engine writes it. The stencil arithmetic comes from the same header and rounds as +// the CPU does, every sum is an integer, and the two masks are a maximum and a minimum - so neither the +// order a thread visits its pixels in nor the order the threads run in changes a result. + +constant int32_t PX_MASKED = -2147483647 - 1; // INT32_MIN: the preprocessor's masked pixel +constant int32_t PX_SATURATED = 2147483647; // INT32_MAX: and its saturated one +constant uint8_t MASK_TAIL = 1; // bragg_engine::MASK_TAIL / MASK_FLUX +constant uint8_t MASK_FLUX = 2; + +inline bool valid_px(int32_t v) { return v != PX_MASKED && v != PX_SATURATED; } + +// The stencil parameters arrive as bytes: Metal takes no struct with default member initialisers as a +// buffer type, and BraggStencilParams has them. +inline BraggStencilParams stencil_params(constant char *raw) { + return *(constant const BraggStencilParams *) raw; +} + +// MarkNeighbours: every reflection marks its r2 region. One dispatch per level, TAIL first, so a pixel +// two regions share ends up with the larger mark, as the CPU's max gives. want = 0 has every reflection +// unmark its region, which leaves the mask clear for the next frame. +kernel void mark_neighbours(device const float2 *xy [[buffer(0)]], + device const uint8_t *level [[buffer(1)]], + device uint8_t *mask [[buffer(2)]], + constant char *stencil_raw [[buffer(3)]], + constant BraggPassAFrame &f [[buffer(4)]], + constant uint8_t &want [[buffer(5)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n || (want != 0 && level[i] != want)) + return; + const BraggStencilParams sp = stencil_params(stencil_raw); + const float px = xy[i].x, py = xy[i].y; + const BraggStencil st = MakeBraggStencil(px, py, sp); + const int x0 = max(0, (int) floor(px - st.ex_in - 1.0f)); + const int x1 = min(f.width - 1, (int) ceil(px + st.ex_in + 1.0f)); + const int y0 = max(0, (int) floor(py - st.ey_in - 1.0f)); + const int y1 = min(f.height - 1, (int) ceil(py + st.ey_in + 1.0f)); + for (int y = y0; y <= y1; ++y) + for (int x = x0; x <= x1; ++x) { + const BraggStencilDist d = BraggStencilDistances(st, x - px, y - py); + if (d.inner < f.r2_sq) + mask[y * f.width + x] = want; + } +} + +// The owner map: every reflection claims the pixels within `claim` of its centre with an atomic minimum +// of its (distance, index) key, so the nearest centre wins whatever order the threads run in. release +// gives the same pixels back to nobody, which leaves the map clear for the next frame. +kernel void claim_owner(device const float2 *xy [[buffer(0)]], + device atomic_uint *owner [[buffer(1)]], + constant BraggPassAFrame &f [[buffer(2)]], + constant uint &release [[buffer(3)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n) + return; + const float px = xy[i].x, py = xy[i].y; + const float claim_sq = f.claim * f.claim; + const int x0 = max(0, (int) floor(px - f.claim)); + const int x1 = min(f.width - 1, (int) ceil(px + f.claim)); + const int y0 = max(0, (int) floor(py - f.claim)); + const int y1 = min(f.height - 1, (int) ceil(py + f.claim)); + for (int y = y0; y <= y1; ++y) + for (int x = x0; x <= x1; ++x) { + const float dx = x - px, dy = y - py; + const float d2 = dx * dx + dy * dy; + if (d2 >= claim_sq) + continue; + if (release) + atomic_store_explicit(&owner[y * f.width + x], BRAGG_OWNER_NONE, memory_order_relaxed); + else + atomic_fetch_min_explicit(&owner[y * f.width + x], BraggOwnerKey(sqrt(d2), f.inv_claim, (int) i), + memory_order_relaxed); + } +} + +// Pass A's walk over the box of the outer ellipse: the signal-disk and background-ring counts. +kernel void box_sums(device const float2 *xy [[buffer(0)]], + device const int32_t *img [[buffer(1)]], + device const uint8_t *mask [[buffer(2)]], + device const uint32_t *owner [[buffer(3)]], + device const uint32_t *static_mask [[buffer(4)]], + constant char *stencil_raw [[buffer(5)]], + constant BraggPassAFrame &f [[buffer(6)]], + device BraggBoxSums *sums [[buffer(7)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n) + return; + const BraggStencilParams sp = stencil_params(stencil_raw); + const float px = xy[i].x, py = xy[i].y; + const BraggStencil st = MakeBraggStencil(px, py, sp); + const int x0 = max(0, (int) floor(px - st.ex_out - 1.0f)); + const int x1 = min(f.width - 1, (int) ceil(px + st.ex_out + 1.0f)); + const int y0 = max(0, (int) floor(py - st.ey_out - 1.0f)); + const int y1 = min(f.height - 1, (int) ceil(py + st.ey_out + 1.0f)); + BraggBoxSums s = {}; + for (int y = y0; y <= y1; ++y) + for (int x = x0; x <= x1; ++x) { + const BraggStencilDist d = BraggStencilDistances(st, x - px, y - py); + const int idx = y * f.width + x; + const int32_t v = img[idx]; + if (BraggStencilInSignal(st, x - px, y - py, d.signal, f.r1_sq)) { + ++s.n_disk; + if (f.use_owner) { + if (BraggOwnedBy(owner[idx], (int) i)) ++s.n_own; + else if (f.exclude) continue; + } + ++s.n_inner; + // Unreadable here but not in the run's mask: lost to the flux on this frame. + const bool masked = (uint) idx / 32 < f.static_mask_words + && ((static_mask[(uint) idx / 32] >> ((uint) idx % 32)) & 1u); + if (v == PX_SATURATED || (v == PX_MASKED && !masked)) + s.overloaded = 1; + if (!valid_px(v)) continue; + s.I_sum += v; + s.I_sum_x += (int64_t) x * v; + s.I_sum_y += (int64_t) y * v; + s.x_sum += x; + s.y_sum += y; + ++s.n_inner_valid; + } else if (d.inner >= f.r2_sq && d.outer < f.r3_sq) { + if (!valid_px(v)) continue; + ++s.n_ring; + const uint8_t m = mask[idx]; + if (m == 0) ++s.n_clean; + if (m == MASK_FLUX) continue; + s.bkg_sum += v; + ++s.n_bkg; + } + } + sums[i] = s; +} + +// A 64-bit integer sum kept as two 32-bit words, low first: Metal has no 64-bit atomic add. The carry +// out of the low word is known to the thread that caused it, so the pair adds up exactly. +inline void atomic_add_i64(device atomic_uint *w, int64_t v) { + const uint64_t u = (uint64_t) v; + const uint lo = (uint) u, hi = (uint) (u >> 32); + const uint old = atomic_fetch_add_explicit(&w[0], lo, memory_order_relaxed); + 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. +kernel void clip_sums(device const float2 *xy [[buffer(0)]], + device const int32_t *img [[buffer(1)]], + device const uint8_t *mask [[buffer(2)]], + constant char *stencil_raw [[buffer(3)]], + constant BraggPassAFrame &f [[buffer(4)]], + device const uint8_t *clips [[buffer(5)]], + device const int32_t *thr [[buffer(6)]], + device BraggClipSums *out [[buffer(7)]], + device atomic_uint *rad_sum [[buffer(8)]], + device atomic_int *rad_cnt [[buffer(9)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n || !clips[i]) + return; + const BraggStencilParams sp = stencil_params(stencil_raw); + const float px = xy[i].x, py = xy[i].y; + const BraggStencil st = MakeBraggStencil(px, py, sp); + const int x0 = max(0, (int) floor(px - st.ex_out - 1.0f)); + const int x1 = min(f.width - 1, (int) ceil(px + st.ex_out + 1.0f)); + const int y0 = max(0, (int) floor(py - st.ey_out - 1.0f)); + const int y1 = min(f.height - 1, (int) ceil(py + st.ey_out + 1.0f)); + const int32_t t = thr[i]; + BraggClipSums c = {}; + for (int y = y0; y <= y1; ++y) + for (int x = x0; x <= x1; ++x) { + const BraggStencilDist d = BraggStencilDistances(st, x - px, y - py); + const int idx = y * f.width + x; + const int32_t v = img[idx]; + if (!(d.inner >= f.r2_sq && d.outer < f.r3_sq && valid_px(v) && mask[idx] != MASK_FLUX && v <= t)) + continue; + c.sum += v; + ++c.n; + if (f.n_rad > 0) { + // Binned on the true detector radius: the unshrunk radial projection. + const int b = clamp((int) round(st.r0 + d.rad), 0, f.n_rad - 1); + atomic_add_i64(&rad_sum[2 * b], v); + atomic_fetch_add_explicit(&rad_cnt[b], 1, memory_order_relaxed); + } + } + out[i] = c; +} + +// BraggIntegrationEngineCPU::FitProfile. The CPU gathers the pixels the fit reads once; here every +// iteration reads them again, the same pixels in the same order. +kernel void profile_fit(device const BraggFitSetup *setups [[buffer(0)]], + device const int32_t *img [[buffer(1)]], + device const uint32_t *owner [[buffer(2)]], + device const float *shell_profiles [[buffer(3)]], + device float *cells [[buffer(4)]], + constant char *stencil_raw [[buffer(5)]], + constant BraggPassBFrame &f [[buffer(6)]], + device BraggFitOut *out [[buffer(7)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n || !setups[i].fit) + return; + const BraggFitSetup s = setups[i]; + const int Rf = s.Rf, Gf = 2 * Rf + 1; + BraggFitOut o = {}; + + device const float *P; + if (s.shell >= 0) + P = shell_profiles + s.shell * (2 * f.R + 1) * (2 * f.R + 1); + else { + device float *g = cells + s.cells; + float gs = 0.0f; + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float v = BraggGaussian(dx, dy, s.fx, s.fy, s.ux, s.uy, s.inv_2s2r, s.inv_2s2t); + g[(dy + Rf) * Gf + (dx + Rf)] = v; + gs += v; + } + const float inv_gs = 1.0f / gs; + for (int k = 0; k < Gf * Gf; ++k) g[k] *= inv_gs; + P = g; + } + + const BraggStencil st = MakeBraggStencil(s.px, s.py, stencil_params(stencil_raw)); + 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 >= f.width || y >= f.height) continue; + 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), f.r1_sq) + : dx * dx + dy * dy < f.r1_sq; + o.p_grid += Pp; + o.p_peak = o.p_peak < Pp ? Pp : o.p_peak; // std::max + if (in_disk) o.m_all += Pp; + if (!valid_px(img[y * f.width + x])) { + o.p_lost_peak = o.p_lost_peak < Pp ? Pp : o.p_lost_peak; + continue; + } + o.p_valid += Pp; + const bool own = !f.use_owner || BraggOwnedBy(owner[y * f.width + x], (int) i); + if (own) o.p_own += Pp; + if (in_disk && (own || !f.exclude)) o.m_read += Pp; + } + + // FitKept is double arithmetic: the CPU decides it, and the fit runs for every reflection. + float If = s.I0; + for (int iter = 0; iter < 4; ++iter) { + float num = 0.0f; + o.den = 0.0f; + o.wsum = 0.0f; + 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 >= f.width || y >= f.height) continue; + const int32_t v = img[y * f.width + x]; + if (!valid_px(v)) continue; + if (f.exclude && f.use_owner && !BraggOwnedBy(owner[y * f.width + x], (int) i)) continue; + const float signal = (float) v - s.bkg; + const float d = s.B + If * Pp; + const float iv = 1.0f / (d < s.v_min ? s.v_min : d); // std::max + 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; + out[i] = o; +} diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm new file mode 100644 index 000000000..85ffafbe7 --- /dev/null +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm @@ -0,0 +1,405 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#import + +#include "BraggIntegrationEngineMetal.h" + +#include +#include +#include + +#include "../../common/CUDAWrapper.h" +#include "../../common/JFJochException.h" + +using namespace bragg_engine; + +// The kernels' source: BraggStencil.h, BraggPassA.h and BraggIntegrationEngineMetal.metal, copied in +// at configure time (CMakeLists.txt) so the binary needs no file beside it. +extern const char *const bragg_metal_source_stencil; +extern const char *const bragg_metal_source_pass_a; +extern const char *const bragg_metal_source_exp; +extern const char *const bragg_metal_source_pass_b; +extern const char *const bragg_metal_source_kernels; + +namespace { + +static_assert(MASK_TAIL == 1 && MASK_FLUX == 2, "BraggIntegrationEngineMetal.metal spells the two marks out"); + +struct Kernels { + id device = nil; // nil: no Metal device + id mark = nil, claim = nil, box = nil, clip = nil, fit = nil; +}; + +// Compiled once per process: a worker makes an engine per pass, and compiling takes a fraction of a +// second. The Metal compiler ships with macOS, so nothing has to be installed to build or run this. +const Kernels &GetKernels() { + static const Kernels kernels = [] { + Kernels k; + k.device = MTLCreateSystemDefaultDevice(); + if (k.device == nil) + return k; + // No contraction of a*b+c into a fused multiply-add, as everywhere on the host (CMakeLists.txt). + // Metal would otherwise contract wherever its backend saw fit. BraggStencil.h calls C's math + // functions; the defines are Metal's spellings of them. + const std::string source = std::string("#include \n" + "using namespace metal;\n" + "#pragma clang fp contract(off)\n" + "#define sqrtf(x) sqrt(x)\n" + "#define fmaxf(a, b) fmax(a, b)\n" + "#define floorf(x) floor(x)\n" + "#define lroundf(x) ((long) round(x))\n") + + bragg_metal_source_stencil + bragg_metal_source_pass_a + + bragg_metal_source_exp + bragg_metal_source_pass_b + + bragg_metal_source_kernels; + MTLCompileOptions *options = [MTLCompileOptions new]; + // Fast math off: sqrt and division then round correctly and nothing is reassociated, as on the + // CPU. (MTLCompileOptions.mathMode says the same from macOS 15; the build targets 13.) And the + // language of macOS 13 on every macOS, so the kernels mean the same wherever they compile. + options.fastMathEnabled = NO; + options.languageVersion = MTLLanguageVersion3_0; + NSError *error = nil; + id library = [k.device newLibraryWithSource:@(source.c_str()) options:options error:&error]; + if (library == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + std::string("Bragg integration kernels: ") + error.localizedDescription.UTF8String); + auto pipeline = [&](const char *name) { + id function = [library newFunctionWithName:@(name)]; + id p = [k.device newComputePipelineStateWithFunction:function error:&error]; + if (p == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + std::string(name) + ": " + error.localizedDescription.UTF8String); + return p; + }; + k.mark = pipeline("mark_neighbours"); + k.claim = pipeline("claim_owner"); + k.box = pipeline("box_sums"); + k.clip = pipeline("clip_sums"); + k.fit = pipeline("profile_fit"); + return k; + }(); + return kernels; +} + +void WaitCommitted(id cb) { + [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed"); +} + +void Wait(id cb) { + [cb commit]; + WaitCommitted(cb); +} + +} // namespace + +struct BraggIntegrationEngineMetal::Device { + id queue = nil; + id static_mask = nil; + // Frame-sized, made on the first image: a worker's analysis is built for every pass, and most of + // them never integrate. A frame unmarks what it marked once the CPU has read them (`unmark`, which + // the next frame waits for): clearing all of both took 1.5 ms of a 16M-pixel frame, longer than + // integrating a sparse one. `clear` is true once an unmark completed; the first frame, and one + // after a frame that threw part-way through, clear the whole of both instead. + id mask = nil, owner = nil; + id unmark = nil; + bool clear = false; + id rad_sum = nil, rad_cnt = nil; // the radial background curve, as the clip kernel adds it up + // Per reflection, grown with the frames. + size_t capacity = 0; + id xy = nil, level = nil, sums = nil, clips = nil, thr = nil, clip_out = nil; + id fit_setups = nil, fit_out = nil; + // Pass B's Gaussians, grown with the frames, and the learned profiles. + size_t cells_capacity = 0, profiles_capacity = 0; + id cells = nil, profiles = nil; + // The frame's image while Run hands pass B to FitOnDevice; nil otherwise (the CPU engine's own Run). + id img = nil; +}; + +bool BraggIntegrationEngineMetal::Available(const DiffractionExperiment &experiment) { + const auto settings = experiment.GetBraggIntegrationSettings(); + const bool trimmed_mean = !(settings.GetBackgroundClipNSigma() > 0.0f) + && settings.GetBackgroundTrimFraction() > 0.0f; + return !gpus_are_hidden() && !trimmed_mean && GetKernels().device != nil; +} + +std::string BraggIntegrationEngineMetal::DeviceName() { + if (gpus_are_hidden() || GetKernels().device == nil) + return {}; + return GetKernels().device.name.UTF8String; +} + +BraggIntegrationEngineMetal::BraggIntegrationEngineMetal(const DiffractionExperiment &experiment, + const PixelMask &mask) + : BraggIntegrationEngineCPU(experiment, mask), dev(std::make_unique()) { + const Kernels &k = GetKernels(); + dev->queue = [k.device newCommandQueue]; + dev->static_mask = [k.device newBufferWithBytes:static_mask.data() + length:std::max(static_mask.size() * sizeof(uint32_t), 4) + options:MTLResourceStorageModeShared]; +} + +BraggIntegrationEngineMetal::~BraggIntegrationEngineMetal() = default; + +std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessorBuffer &image, + const std::vector &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. + if (image.size() != npixel || npredicted == 0) + return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number); + + dev->img = nil; + @autoreleasepool { + const Kernels &k = GetKernels(); + // The GPU reads the host's own memory, so the image is read where it lies - which Metal allows + // for memory starting on a page, as macOS's allocator hands out any buffer of a frame's size. + // Should one not be, the CPU engine integrates it: the same result. + const size_t page = getpagesize(), bytes = npixel * sizeof(int32_t); + id img = nil; + if (reinterpret_cast(image.data()) % page == 0) + img = [k.device newBufferWithBytesNoCopy:const_cast(image.data()) + length:(bytes + page - 1) / page * page + options:MTLResourceStorageModeShared + deallocator:nil]; + if (img == nil) + return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number); + + counts.predicted += npredicted; + const bool use_owner = overlap != OverlapMode::Off; + if (dev->unmark != nil) { // it reads xy, which is overwritten below + id unmark = dev->unmark; + dev->unmark = nil; + dev->clear = false; // unless it completed + WaitCommitted(unmark); + dev->clear = true; + } + auto buffer = [&](size_t length) { + return [k.device newBufferWithLength:std::max(length, 4) options:MTLResourceStorageModeShared]; + }; + if (dev->mask == nil) { + dev->mask = buffer(npixel); + if (use_owner) + dev->owner = buffer(npixel * sizeof(uint32_t)); + } + const bool was_clear = dev->clear; + dev->clear = false; + if (dev->capacity < npredicted) { + const size_t n = npredicted + npredicted / 2; + dev->xy = buffer(n * 2 * sizeof(float)); + dev->level = buffer(n); + dev->sums = buffer(n * sizeof(BraggBoxSums)); + dev->clips = buffer(n); + dev->thr = buffer(n * sizeof(int32_t)); + dev->clip_out = buffer(n * sizeof(BraggClipSums)); + dev->fit_setups = buffer(n * sizeof(BraggFitSetup)); + dev->fit_out = buffer(n * sizeof(BraggFitOut)); + dev->capacity = n; + } + + auto *xy = static_cast(dev->xy.contents); + auto *level = static_cast(dev->level.contents); + for (size_t i = 0; i < npredicted; ++i) { + xy[2 * i] = predicted[i].predicted_x; + xy[2 * i + 1] = predicted[i].predicted_y; + level[i] = predicted[i].partiality >= NEIGHBOUR_MASK_MIN_PARTIALITY ? MASK_FLUX : MASK_TAIL; + } + + PassA a = StartPassA(npredicted); + BraggPassAFrame f{}; + f.width = static_cast(xpixel); + f.height = static_cast(ypixel); + f.r1_sq = r1_sq; + f.r2_sq = r2_sq; + f.r3_sq = r3_sq; + f.claim = claim; + f.inv_claim = inv_claim; + f.use_owner = use_owner; + f.exclude = overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum; + f.n_rad = bkg_radial ? a.n_rad : 0; + f.static_mask_words = static_cast(static_mask.size()); + f.n = static_cast(npredicted); + + const MTLSize grid = MTLSizeMake(npredicted, 1, 1), group = MTLSizeMake(64, 1, 1); + + // --- The neighbour mask, the owner map and pass A's box walk --- + id cb = [dev->queue commandBuffer]; + id blit; + if (!was_clear) { + blit = [cb blitCommandEncoder]; + [blit fillBuffer:dev->mask range:NSMakeRange(0, npixel) value:0]; + if (use_owner) // every byte 0xff: BRAGG_OWNER_NONE + [blit fillBuffer:dev->owner range:NSMakeRange(0, npixel * sizeof(uint32_t)) value:0xff]; + [blit endEncoding]; + } + const uint32_t claim_not_release = 0; + id enc = [cb computeCommandEncoder]; + for (const uint8_t want : {MASK_TAIL, MASK_FLUX}) { + [enc setComputePipelineState:k.mark]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:dev->level offset:0 atIndex:1]; + [enc setBuffer:dev->mask offset:0 atIndex:2]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:3]; + [enc setBytes:&f length:sizeof(f) atIndex:4]; + [enc setBytes:&want length:sizeof(want) atIndex:5]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + } + if (use_owner) { + [enc setComputePipelineState:k.claim]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:dev->owner offset:0 atIndex:1]; + [enc setBytes:&f length:sizeof(f) atIndex:2]; + [enc setBytes:&claim_not_release length:sizeof(claim_not_release) atIndex:3]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + } + [enc setComputePipelineState:k.box]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:img offset:0 atIndex:1]; + [enc setBuffer:dev->mask offset:0 atIndex:2]; + [enc setBuffer:(use_owner ? dev->owner : dev->mask) offset:0 atIndex:3]; // not read without an owner map + [enc setBuffer:dev->static_mask offset:0 atIndex:4]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:5]; + [enc setBytes:&f length:sizeof(f) atIndex:6]; + [enc setBuffer:dev->sums offset:0 atIndex:7]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + [enc endEncoding]; + Wait(cb); + + // --- The clip thresholds, which are double arithmetic, on the CPU; the clip on the GPU --- + const auto *sums = static_cast(dev->sums.contents); + auto *clips = static_cast(dev->clips.contents); + auto *thr = static_cast(dev->thr.contents); + bool any_clip = false; + for (size_t i = 0; i < npredicted; ++i) { + clips[i] = ClipsRing(sums[i]); + if (clips[i]) { + thr[i] = ClipThreshold(sums[i]); + any_clip = true; + } + } + if (any_clip) { + if (bkg_radial && dev->rad_sum == nil) { + dev->rad_sum = buffer(a.n_rad * 2 * sizeof(uint32_t)); + dev->rad_cnt = buffer(a.n_rad * sizeof(int32_t)); + } + cb = [dev->queue commandBuffer]; + if (bkg_radial) { + blit = [cb blitCommandEncoder]; + [blit fillBuffer:dev->rad_sum range:NSMakeRange(0, a.n_rad * 2 * sizeof(uint32_t)) value:0]; + [blit fillBuffer:dev->rad_cnt range:NSMakeRange(0, a.n_rad * sizeof(int32_t)) value:0]; + [blit endEncoding]; + } + enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:k.clip]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:img offset:0 atIndex:1]; + [enc setBuffer:dev->mask offset:0 atIndex:2]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:3]; + [enc setBytes:&f length:sizeof(f) atIndex:4]; + [enc setBuffer:dev->clips offset:0 atIndex:5]; + [enc setBuffer:dev->thr offset:0 atIndex:6]; + [enc setBuffer:dev->clip_out offset:0 atIndex:7]; + [enc setBuffer:(bkg_radial ? dev->rad_sum : dev->clips) offset:0 atIndex:8]; // not written without the curve + [enc setBuffer:(bkg_radial ? dev->rad_cnt : dev->clips) offset:0 atIndex:9]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + [enc endEncoding]; + Wait(cb); + if (bkg_radial) { + const auto *w = static_cast(dev->rad_sum.contents); + const auto *c = static_cast(dev->rad_cnt.contents); + for (int b = 0; b < a.n_rad; ++b) { + a.rad_sum[b] = static_cast(w[2 * b] | static_cast(w[2 * b + 1]) << 32); + a.rad_cnt[b] = c[b]; + } + } + } + + ++frames_on_gpu; + + // --- The rest is the CPU engine's --- + const auto *clip_out = static_cast(dev->clip_out.contents); + std::vector no_trim; // the trimmed mean is the CPU engine's alone (Available) + for (size_t i = 0; i < npredicted; ++i) { + const Reflection &r = predicted[i]; + FinishPassA(a, i, r, MakeBraggStencil(r.predicted_x, r.predicted_y, stencil), sums[i], + clips[i] ? clip_out[i] : BraggClipSums{}, no_trim); + } + const FlatOwnerMap owner_map{use_owner ? static_cast(dev->owner.contents) : nullptr, xpixel}; + dev->img = img; + auto ret = ProfilePasses(BufferSampler{image.data()}, owner_map, predicted, npredicted, a, image_number); + dev->img = nil; + + // --- Unmark this frame's regions and give back its claims, while the CPU goes on --- + dev->unmark = [dev->queue commandBuffer]; + enc = [dev->unmark computeCommandEncoder]; + const uint8_t unmark = 0; + [enc setComputePipelineState:k.mark]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:dev->level offset:0 atIndex:1]; + [enc setBuffer:dev->mask offset:0 atIndex:2]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:3]; + [enc setBytes:&f length:sizeof(f) atIndex:4]; + [enc setBytes:&unmark length:sizeof(unmark) atIndex:5]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + if (use_owner) { + const uint32_t release = 1; + [enc setComputePipelineState:k.claim]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:dev->owner offset:0 atIndex:1]; + [enc setBytes:&f length:sizeof(f) atIndex:2]; + [enc setBytes:&release length:sizeof(release) atIndex:3]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + } + [enc endEncoding]; + [dev->unmark commit]; + return ret; + } +} + +bool BraggIntegrationEngineMetal::FitOnDevice(const std::vector &setups, + const std::vector &shell_profiles, + std::vector &fits) { + if (dev->img == nil) + return false; + @autoreleasepool { + const Kernels &k = GetKernels(); + const size_t n = setups.size(); + size_t n_cells = 0; + for (const BraggFitSetup &s : setups) + if (s.fit && s.shell < 0) + n_cells = std::max(n_cells, s.cells + static_cast(2 * s.Rf + 1) * (2 * s.Rf + 1)); + auto grow = [&](__strong id &b, size_t &capacity, size_t bytes) { + if (capacity < bytes) { + capacity = bytes + bytes / 2; + b = [k.device newBufferWithLength:std::max(capacity, 4) options:MTLResourceStorageModeShared]; + } + }; + grow(dev->cells, dev->cells_capacity, n_cells * sizeof(float)); + grow(dev->profiles, dev->profiles_capacity, shell_profiles.size() * sizeof(float)); + std::copy(setups.begin(), setups.end(), static_cast(dev->fit_setups.contents)); + std::copy(shell_profiles.begin(), shell_profiles.end(), static_cast(dev->profiles.contents)); + + const bool use_owner = overlap != OverlapMode::Off; + const BraggPassBFrame f{static_cast(xpixel), static_cast(ypixel), r1_sq, R, use_owner, + overlap == OverlapMode::Exclude, static_cast(n)}; + + id cb = [dev->queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:k.fit]; + [enc setBuffer:dev->fit_setups offset:0 atIndex:0]; + [enc setBuffer:dev->img offset:0 atIndex:1]; + [enc setBuffer:(use_owner ? dev->owner : dev->mask) offset:0 atIndex:2]; // not read without an owner map + [enc setBuffer:dev->profiles offset:0 atIndex:3]; + [enc setBuffer:dev->cells offset:0 atIndex:4]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:5]; + [enc setBytes:&f length:sizeof(f) atIndex:6]; + [enc setBuffer:dev->fit_out offset:0 atIndex:7]; + [enc dispatchThreads:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; + [enc endEncoding]; + Wait(cb); + const auto *out = static_cast(dev->fit_out.contents); + std::copy(out, out + n, fits.begin()); + } + return true; +} diff --git a/image_analysis/bragg_integration/BraggPassA.h b/image_analysis/bragg_integration/BraggPassA.h new file mode 100644 index 000000000..8b1b10979 --- /dev/null +++ b/image_analysis/bragg_integration/BraggPassA.h @@ -0,0 +1,45 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Pass A of BraggIntegrationEngineCPU, as the data it hands on: what the walk over one reflection's box +// counts, and the clip sums of its background ring. Integers only, so the counts do not depend on the +// order the pixels are visited in: the Metal engine walks the boxes on the GPU and gets the CPU's +// counts exactly. Compiled by the host and by the Metal compiler, so the two agree on the layout - +// which is also why there are no default member initialisers (Metal does not take a struct with them +// as a buffer type); zero one with {}. + +#ifndef __METAL_VERSION__ +#include +#endif + +struct BraggBoxSums { + int64_t I_sum, I_sum_x, I_sum_y; // readable signal-disk pixels, and their first moments + int64_t x_sum, y_sum; // the positions behind I_sum, to take the background out of the centroid + int64_t bkg_sum; // ring pixels kept for the background + int64_t n_inner, n_inner_valid; // signal-disk pixels kept, and how many of them were readable + int32_t n_disk, n_own; // signal-disk pixels, and how many of them are this reflection's + int32_t n_bkg; // ring pixels behind bkg_sum + int32_t n_ring, n_clean; // readable ring pixels, and those no predicted neighbour covers + int32_t overloaded; // a signal-disk pixel was saturated, or unreadable on this frame only +}; + +// The ring pixels at or below a reflection's clip threshold (BraggIntegrationEngineCPU::ClipThreshold). +struct BraggClipSums { + int64_t sum; + int32_t n; + int32_t unused; +}; + +// What the Metal kernels are told about the frame, beside the stencil (BraggStencilParams). +struct BraggPassAFrame { + int32_t width, height; + float r1_sq, r2_sq, r3_sq; + float claim, inv_claim; // how far a reflection claims pixels in the owner map + int32_t use_owner; // an overlap treatment is on, so the owner map is built and read + int32_t exclude; // ... and a pixel a neighbour owns leaves the signal disk + int32_t n_rad; // radial background bins; 0 when the curve is not accumulated + uint32_t static_mask_words; // length of the run's packed pixel mask + uint32_t n; // reflections +}; diff --git a/image_analysis/bragg_integration/BraggPassB.h b/image_analysis/bragg_integration/BraggPassB.h new file mode 100644 index 000000000..e03981a5e --- /dev/null +++ b/image_analysis/bragg_integration/BraggPassB.h @@ -0,0 +1,55 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Pass B of BraggIntegrationEngineCPU, the profile fit, as the data it hands on. The CPU works out in +// double what one reflection's fit needs (BraggFitSetup), the fit itself is single precision +// (BraggFitOut), and the CPU finishes it in double. The Metal engine runs the middle on the GPU, so this +// is compiled by the host and by the Metal compiler (no default member initialisers, as in BraggPassA.h). + +#ifndef __METAL_VERSION__ +#include + +#include "BraggExp.h" // the Metal engine puts it in front +#endif + +struct BraggFitSetup { + float px, py; // predicted position + int32_t cx, cy; // profile grid centre + int32_t Rf; // grid half-size + int32_t shell; // the learned profile of this shell, or -1: the Gaussian below + uint32_t cells; // where the Metal engine keeps the Gaussian, Gf^2 floats + int32_t fit; // 0: not fitted (no pass A result, or overloaded) + int32_t footprint; // the stencil has a measured footprint: it decides the signal disk + float fx, fy; // sub-pixel offset of the prediction from the grid centre + float ux, uy; // radial direction + float inv_2s2r, inv_2s2t; // 1 / (2 sigma^2), radially and tangentially + float bkg; // background per pixel + float B, v_min; // the variance model's floor, and the least variance a weight takes + float I0; // the box sum, where the fit starts +}; + +struct BraggFitOut { + float p_grid, p_valid, p_own; // profile mass on the detector, readable, and this reflection's own + float p_peak, p_lost_peak; // the largest profile value, and the largest one unreadable + float m_all, m_read; // mass in the signal disk, and the part the box sum read + float I, den, wsum; // the fit; den = 0 when it did not run or failed +}; + +// What the Metal engine's fit kernel is told about the frame, beside the stencil. +struct BraggPassBFrame { + int32_t width, height; + float r1_sq; + int32_t R; // the learned profiles' grid half-size + int32_t use_owner, exclude; // as in BraggPassAFrame + uint32_t n; // reflections +}; + +// 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) { + const float ex = dx - fx, ey = dy - fy; + const float rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux; + return BraggExp(-rad * rad * inv_2s2r - tn * tn * inv_2s2t); +} diff --git a/image_analysis/bragg_integration/BraggStencil.h b/image_analysis/bragg_integration/BraggStencil.h index c92122fb0..00030799d 100644 --- a/image_analysis/bragg_integration/BraggStencil.h +++ b/image_analysis/bragg_integration/BraggStencil.h @@ -49,15 +49,19 @@ // gives q = 0 and hence d2_inner == d2_outer == d2 EXACTLY, in floating point, not merely to within // rounding. k_sigma = 0 therefore reproduces the old circular stencil bit for bit. // -// This header is compiled by both the host compiler and nvcc, so the geometry has exactly one -// definition: the two engines classify by the same rule, built from the same inputs. They are not -// bit-identical to each other - nvcc contracts a*b+c into an FMA and the host baseline has no FMA -// instruction to contract into - so a pixel within a rounding of a boundary can still fall either -// way, which is what the parity test's tolerances are for. +// This header is compiled by the host compiler, by nvcc and by the Metal compiler, so the geometry has +// exactly one definition: the engines classify by the same rule, built from the same inputs. The CUDA +// engine is not bit-identical to the CPU one - nvcc contracts a*b+c into an FMA and the host baseline +// has no FMA instruction to contract into - so a pixel within a rounding of a boundary can still fall +// either way, which is what its parity test's tolerances are for. The Metal engine is: neither the host +// nor the Metal compiler contracts anything, and with fast math off Metal rounds sqrt and division +// correctly, as the CPU does (BraggIntegrationEngineMetal.mm). // ============================================================================= +#ifndef __METAL_VERSION__ #include #include +#endif #ifdef __CUDACC__ #define BRAGG_STENCIL_HD __host__ __device__ inline @@ -65,14 +69,24 @@ #define BRAGG_STENCIL_HD inline #endif +// Metal keeps a program-scope constant in its constant address space, and a reference names the space +// it points into (here always the thread's own), so both say so. +#ifdef __METAL_VERSION__ +#define BRAGG_STENCIL_CONSTANT constant +#define BRAGG_STENCIL_THREAD thread +#else +#define BRAGG_STENCIL_CONSTANT +#define BRAGG_STENCIL_THREAD +#endif + // The measured spot footprint (SpotFootprint.h): the radial and tangential standard deviations of a // spot [px], tabulated against distance from the beam centre in bins of fp_bin_px. Held here, by value, // because both engines build every stencil from these parameters - the GPU gets the table with the // kernel arguments and needs no buffer of its own. -constexpr int BRAGG_FOOTPRINT_MAX_BINS = 16; +BRAGG_STENCIL_CONSTANT constexpr int BRAGG_FOOTPRINT_MAX_BINS = 16; // A spot whose footprint reaches this many of its standard deviations past r1 has outgrown the disk, and // gets the footprint stencil below. -constexpr float BRAGG_FOOTPRINT_NSIGMA = 3.0f; +BRAGG_STENCIL_CONSTANT constexpr float BRAGG_FOOTPRINT_NSIGMA = 3.0f; // How far that stencil reaches, in the footprint's standard deviations: the summation ellipse ends and the // background ring starts there, and the profile is fitted out to it. Further than the 3 sigma that decides // whether a spot outgrew the disk, because a wide spot is not a Gaussian: mosaic streaks and diffuse halos @@ -80,7 +94,7 @@ constexpr float BRAGG_FOOTPRINT_NSIGMA = 3.0f; // corrections applied to both integrations, so only the integration differs, 4 sigma closed about a third // of the gap to XDS's 3D summation on the haloed small-molecule set and was neutral on the streaked one; // 5 sigma was no better than 4. -constexpr float BRAGG_FOOTPRINT_REACH = 4.0f; +BRAGG_STENCIL_CONSTANT constexpr float BRAGG_FOOTPRINT_REACH = 4.0f; // Fixed per-experiment inputs to the stencil law (mirrors BraggIntegrationEngine's members). struct BraggStencilParams { @@ -116,7 +130,7 @@ struct BraggStencil { // Radial growth of the ring at Rpx, in pixels: k_sigma times the bandwidth streak, capped. Shared // so the stencil and the radial-background kernel table agree on what "grown by this much" means. -BRAGG_STENCIL_HD float BraggStencilGrow_px(float Rpx, const BraggStencilParams &p) { +BRAGG_STENCIL_HD float BraggStencilGrow_px(float Rpx, const BRAGG_STENCIL_THREAD BraggStencilParams &p) { if (!(p.k_sigma > 0.0f)) return 0.0f; // The cap bounds what a mis-declared bandwidth can do to the bounding box. A caller that // forgets to set it gets the uncapped growth rather than silently no growth at all. @@ -125,7 +139,8 @@ BRAGG_STENCIL_HD float BraggStencilGrow_px(float Rpx, const BraggStencilParams & } // The footprint at Rpx: linear between bin centres, flat beyond the first and the last. -BRAGG_STENCIL_HD void BraggFootprintAt(float Rpx, const BraggStencilParams &p, float &s_rad, float &s_tan) { +BRAGG_STENCIL_HD void BraggFootprintAt(float Rpx, const BRAGG_STENCIL_THREAD BraggStencilParams &p, + BRAGG_STENCIL_THREAD float &s_rad, BRAGG_STENCIL_THREAD float &s_tan) { const float t = Rpx / p.fp_bin_px - 0.5f; int i = (int) floorf(t); float f = t - (float) i; @@ -139,7 +154,7 @@ BRAGG_STENCIL_HD void BraggFootprintAt(float Rpx, const BraggStencilParams &p, f // Growth of the ring beyond r2 that the footprint asks for along one axis: none while the spot fits // the r1 disk - there the disk resolves the spot and the stencil stays what it always was - and // otherwise out to BRAGG_FOOTPRINT_REACH of it, capped as the bandwidth growth is. -BRAGG_STENCIL_HD float BraggFootprintGrow_px(float sigma, const BraggStencilParams &p) { +BRAGG_STENCIL_HD float BraggFootprintGrow_px(float sigma, const BRAGG_STENCIL_THREAD BraggStencilParams &p) { const float g = BRAGG_FOOTPRINT_REACH * sigma - p.r2; if (!(g > 0.0f)) return 0.0f; return (p.max_grow > 0.0f && g > p.max_grow) ? p.max_grow : g; @@ -147,7 +162,7 @@ BRAGG_STENCIL_HD float BraggFootprintGrow_px(float sigma, const BraggStencilPara // The largest growth any reflection on a detector reaching r_max can get, radially - what the // buffers that hold a ring are sized by. -BRAGG_STENCIL_HD float BraggStencilMaxGrow_px(float r_max, const BraggStencilParams &p) { +BRAGG_STENCIL_HD float BraggStencilMaxGrow_px(float r_max, const BRAGG_STENCIL_THREAD BraggStencilParams &p) { float g = BraggStencilGrow_px(r_max, p); for (int i = 0; i < p.fp_n; ++i) { const float s = fmaxf(p.fp_sigma_rad[i], p.fp_sigma_tan[i]); @@ -156,7 +171,7 @@ BRAGG_STENCIL_HD float BraggStencilMaxGrow_px(float r_max, const BraggStencilPar return g; } -BRAGG_STENCIL_HD BraggStencil MakeBraggStencil(float px_x, float px_y, const BraggStencilParams &p) { +BRAGG_STENCIL_HD BraggStencil MakeBraggStencil(float px_x, float px_y, const BRAGG_STENCIL_THREAD BraggStencilParams &p) { BraggStencil s; const float rx = px_x - p.beam_x, ry = px_y - p.beam_y; const float r0 = sqrtf(rx * rx + ry * ry); @@ -212,7 +227,7 @@ BRAGG_STENCIL_HD BraggStencil MakeBraggStencil(float px_x, float px_y, const Bra // Taken from the stencil, so the kernel and the ring it corrects are built from the same radius - // deriving it from a separately computed r0 let the two engines round differently and pick // different rows. -BRAGG_STENCIL_HD int BraggStencilKernelIndex(const BraggStencil &s, int n_kern) { +BRAGG_STENCIL_HD int BraggStencilKernelIndex(const BRAGG_STENCIL_THREAD BraggStencil &s, int n_kern) { if (n_kern <= 1) return 0; int idx = (int) lroundf(s.grow); if (idx < 0) idx = 0; @@ -234,8 +249,8 @@ struct BraggStencilDist { // minimum, so the nearest centre wins whatever order the writes arrive in and the lowest index breaks // a tie. Eight bits of distance over the claim radius is 1/40 px at the default r2 = 6 px, far finer // than the prediction itself. A pixel nobody claimed is nobody's neighbour, hence the owner's. -constexpr uint32_t BRAGG_OWNER_NONE = 0xffffffffu; -constexpr uint32_t BRAGG_OWNER_IDX_MASK = 0x00ffffffu; +BRAGG_STENCIL_CONSTANT constexpr uint32_t BRAGG_OWNER_NONE = 0xffffffffu; +BRAGG_STENCIL_CONSTANT constexpr uint32_t BRAGG_OWNER_IDX_MASK = 0x00ffffffu; BRAGG_STENCIL_HD int BraggOwnerQuant(float d, float inv_claim) { int q = (int) (d * inv_claim * 255.0f); @@ -258,7 +273,7 @@ BRAGG_STENCIL_HD bool BraggOwnedBy(uint32_t owner, int index) { // intensity at 1, 2 and 3 px, and the merge fell apart with it. What the margin removes is the // reflection's own profile, not the neighbour's tail, so the plain nearest-centre split is the rule. -BRAGG_STENCIL_HD BraggStencilDist BraggStencilDistances(const BraggStencil &s, float ddx, float ddy) { +BRAGG_STENCIL_HD BraggStencilDist BraggStencilDistances(const BRAGG_STENCIL_THREAD BraggStencil &s, float ddx, float ddy) { BraggStencilDist d; d.rad = ddx * s.ux + ddy * s.uy; d.signal = ddx * ddx + ddy * ddy; @@ -276,7 +291,7 @@ BRAGG_STENCIL_HD BraggStencilDist BraggStencilDistances(const BraggStencil &s, f // The signal region: the r1 disk, and for a spot that outgrows it the whole footprint ellipse too, so // the summation - the seed of the profile fit, and what the fit falls back to - holds the spot rather // than its core. Without a footprint this is the plain disk test, bit for bit. -BRAGG_STENCIL_HD bool BraggStencilInSignal(const BraggStencil &s, float ddx, float ddy, float d2, float r1_sq) { +BRAGG_STENCIL_HD bool BraggStencilInSignal(const BRAGG_STENCIL_THREAD BraggStencil &s, float ddx, float ddy, float d2, float r1_sq) { if (d2 < r1_sq) return true; if (!(s.fp_s2t > 0.0f)) return false; const float rad = ddx * s.ux + ddy * s.uy, tn = -ddx * s.uy + ddy * s.ux; diff --git a/image_analysis/bragg_integration/CMakeLists.txt b/image_analysis/bragg_integration/CMakeLists.txt index ac34468cf..2a74ecd82 100644 --- a/image_analysis/bragg_integration/CMakeLists.txt +++ b/image_analysis/bragg_integration/CMakeLists.txt @@ -3,6 +3,9 @@ ADD_LIBRARY(JFJochBraggIntegration STATIC BraggIntegrationEngine.h BraggIntegrationEngineCPU.cpp BraggIntegrationEngineCPU.h + BraggPassA.h + BraggPassB.h + BraggExp.h SpotFootprint.cpp SpotFootprint.h Regression.h @@ -21,3 +24,30 @@ IF (JFJOCH_CUDA_AVAILABLE) TARGET_SOURCES(JFJochBraggIntegration PRIVATE ../indexing/CUDAMemHelpers.h BraggIntegrationEngineGPU.cu BraggIntegrationEngineGPU.h) ENDIF() + +IF (APPLE) + # The Metal engine compiles its kernels when the program starts, from source carried in the binary: + # macOS has the Metal compiler built in, while the offline one is an Xcode download of its own. The + # source is the stencil, exp and pass headers the CPU engine compiles and the kernels, copied + # into a generated .cpp here; editing one of them re-runs the configure. + SET(METAL_PARTS stencil BraggStencil.h pass_a BraggPassA.h exp BraggExp.h pass_b BraggPassB.h + kernels BraggIntegrationEngineMetal.metal) + SET(METAL_SOURCE "// Generated from image_analysis/bragg_integration by its CMakeLists.txt - do not edit.\n") + WHILE (METAL_PARTS) + LIST(POP_FRONT METAL_PARTS name file) + FILE(READ ${CMAKE_CURRENT_SOURCE_DIR}/${file} text) + STRING(APPEND METAL_SOURCE "extern const char *const bragg_metal_source_${name};\n" + "const char *const bragg_metal_source_${name} = R\"BRAGG_METAL(\n// ${file}\n${text})BRAGG_METAL\";\n") + SET_PROPERTY(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/${file}) + ENDWHILE() + # Written only when it changed, so a configure does not rebuild the library. + FILE(WRITE ${CMAKE_CURRENT_BINARY_DIR}/BraggIntegrationEngineMetalSource.cpp.tmp "${METAL_SOURCE}") + CONFIGURE_FILE(${CMAKE_CURRENT_BINARY_DIR}/BraggIntegrationEngineMetalSource.cpp.tmp + ${CMAKE_CURRENT_BINARY_DIR}/BraggIntegrationEngineMetalSource.cpp COPYONLY) + + TARGET_SOURCES(JFJochBraggIntegration PRIVATE BraggIntegrationEngineMetal.mm BraggIntegrationEngineMetal.h + BraggIntegrationEngineMetal.metal ${CMAKE_CURRENT_BINARY_DIR}/BraggIntegrationEngineMetalSource.cpp) + SET_SOURCE_FILES_PROPERTIES(BraggIntegrationEngineMetal.mm PROPERTIES COMPILE_OPTIONS -fobjc-arc) + SET_SOURCE_FILES_PROPERTIES(BraggIntegrationEngineMetal.metal PROPERTIES HEADER_FILE_ONLY ON) + TARGET_LINK_LIBRARIES(JFJochBraggIntegration "-framework Metal" "-framework Foundation") +ENDIF() diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 88e37bdb3..e066009c9 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -275,6 +275,12 @@ namespace { // the loop shrank every scale by 10-30% per round for thirty rounds under the robust fit and // settled in fourteen under the plain one. With the weights the reference uses, the plain // alternating fit is exact coordinate descent on one objective and cannot climb. + double ScaleFromSums(double num, double den, double &info) { + const double G = den > 0.0 ? num / den : NAN; + info = den; + return std::isfinite(G) ? std::max(0.0, G) : 1.0; + } + double SolveScale(const std::vector &obs, double &info) { double num = 0.0, den = 0.0; for (const auto &o : obs) { @@ -282,9 +288,7 @@ namespace { num += w2 * o.coeff * o.Iobs; den += w2 * o.coeff * o.coeff; } - const double G = den > 0.0 ? num / den : NAN; - info = den; - return std::isfinite(G) ? std::max(0.0, G) : 1.0; + return ScaleFromSums(num, den, info); } double median_of(std::vector &v) { @@ -1520,18 +1524,6 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { return n_groups; } -double RotationScaleMerge::ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, - double min_partiality, float &I_corr) { - if (group < 0) return 0.0; - if (!(corr > 0.0f) || !std::isfinite(corr)) return 0.0; - if (partiality < min_partiality) return 0.0; - const float Ic = I * corr; - const float sigma_corr = sigma * corr; - if (!std::isfinite(Ic) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0; - I_corr = Ic; - return 1.0 / (static_cast(sigma_corr) * sigma_corr); -} - void RotationScaleMerge::ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const { // Inverse-variance per-group mean of I*corr = the merge reference (no cell mask: the scaling // reference is built without a reference cell - only the final merge applies it). A group per task @@ -1573,7 +1565,7 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vector &refs, const std::vector &corr, std::vector &out_mean) const { - out_mean.assign(n_groups, NAN); + out_mean.resize(n_groups); ParallelChunks(n_groups, ThreadsForWork(corr.size(), nthreads), [&](int glo, int ghi) { for (int g = glo; g < ghi; ++g) { double sw = 0.0, swI = 0.0; @@ -1581,13 +1573,12 @@ void RotationScaleMerge::ReduceScalingGroupMeans(int n_groups, const std::vector for (int p = lo; p < hi; ++p) { const ScalingRef &r = refs[p]; float I_corr = 0.0f; - const double w = ReferenceWeight(g, corr[partial_group_perm[p]], r.partiality, r.I, r.sigma, - min_partiality, I_corr); + const double w = ReferenceWeight(g, corr[p], r.partiality, r.I, r.sigma, min_partiality, I_corr); if (w <= 0.0) continue; sw += w; swI += w * I_corr; } - if (sw > 0.0) out_mean[g] = swI / sw; + out_mean[g] = sw > 0.0 ? swI / sw : NAN; } }); } @@ -1718,6 +1709,31 @@ void RotationScaleMerge::FitPerFrameG(const std::vector &obs, const std::vect frame_info_scratch = std::move(info); } +void RotationScaleMerge::FitScalingFrames(const std::vector &fit, const std::vector &fit_start, + const std::vector &group_mean_in, std::vector &g) { + const int n_fr = static_cast(fit_start.size()) - 1; + std::vector scaled(n_fr, 0); + std::vector info(n_fr, 0.0); + ParallelFor(n_fr, nthreads, [&](int f) { + double num = 0.0, den = 0.0; + size_t n = 0; + for (int i = fit_start[f]; i < fit_start[f + 1]; ++i) { + const ScalingFit &o = fit[i]; + const double mean = group_mean_in[o.group]; + if (!std::isfinite(mean)) continue; + const double coeff = o.coeff_per_mean * mean; + num += o.w2 * coeff * o.Iobs; + den += o.w2 * coeff * coeff; + ++n; + } + if (n < MIN_REFLECTIONS) return; + g[f] = ScaleFromSums(num, den, info[f]); + scaled[f] = 1; + }); + frame_scaled_scratch = std::move(scaled); + frame_info_scratch = std::move(info); +} + void RotationScaleMerge::ShrinkToRestrained(const char *what, const std::vector &g_restrained, const std::vector &fitted, const std::vector &frame_obs_count, @@ -6414,28 +6430,49 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool logger->Info("Per-frame scaling (partials): the same grouping and settings as an earlier merge, " "so its {} iterations stand", partial_loop.iterations); } else { - // The loop works on its own compact copy of the partials (ScalingObs in the header) and hands - // corr back to them when it is done. + // The loop works on its own compact copy of the partials (ScalingRef and ScalingFit in the + // header) and hands corr back to them when it is done. scaling_index[q] is the partial at + // position q of the copy: the CSR's group order, then the partials in no group. const int n_partials = static_cast(partials.size()); const size_t partial_threads = ThreadsForWork(partials.size(), nthreads); - std::vector scaling_obs(n_partials); + const int n_grouped = partial_group_start.empty() + ? 0 : partial_group_start.back() + partial_group_count.back(); + std::vector scaling_index(partial_group_perm.get(), partial_group_perm.get() + n_grouped); + for (int i = 0; i < n_partials; ++i) + if (partials[i].group < 0) scaling_index.push_back(i); + std::vector scaling_refs(n_partials); std::vector scaling_corr(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - const Obs &o = partials[i]; - scaling_obs[i] = {o.frame, o.group, o.I, o.sigma, o.partiality, o.prescaling_corr, o.zeta, - o.on_ice}; - scaling_corr[i] = o.corr; + for (int q = lo; q < hi; ++q) { + const Obs &o = partials[scaling_index[q]]; + scaling_refs[q] = {o.frame, o.I, o.sigma, o.partiality, o.prescaling_corr}; + scaling_corr[q] = o.corr; } }); - std::vector scaling_refs(partial_group_start.empty() - ? 0 : partial_group_start.back() + partial_group_count.back()); - ParallelChunks(n_groups, partial_threads, [&](int glo, int ghi) { - for (int g = glo; g < ghi; ++g) - for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p) { - const ScalingObs &o = scaling_obs[partial_group_perm[p]]; - scaling_refs[p] = {o.I, o.sigma, o.partiality}; - } + + // FitPerFrameG's filters, all but the reference mean's, frame by frame. + auto fitted_obs = [&](const Obs &o) { + return o.group >= 0 && !o.on_ice && o.partiality >= min_partiality + && std::isfinite(o.zeta) && o.zeta > 0.0f; + }; + const int n_fit_frames = static_cast(frame_start.size()); + std::vector fit_start(n_fit_frames + 1, 0); + ParallelFor(n_fit_frames, nthreads, [&](int f) { + fit_start[f + 1] = static_cast(std::count_if( + partials.begin() + frame_start[f], partials.begin() + frame_start[f] + frame_count[f], + fitted_obs)); + }); + for (int f = 0; f < n_fit_frames; ++f) fit_start[f + 1] += fit_start[f]; + std::vector scaling_fit(fit_start.back()); + ParallelFor(n_fit_frames, nthreads, [&](int f) { + int32_t k = fit_start[f]; + for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) { + const Obs &o = partials[i]; + if (!fitted_obs(o)) continue; + const double weight = SafeInv(o.sigma, 1.0); + scaling_fit[k++] = {o.group, o.partiality * SafeInv(o.prescaling_corr, 1.0), + static_cast(o.I), weight * weight}; + } }); // RunScalingLoop rescales after every iteration before anything reads corr, so the iteration's @@ -6446,14 +6483,14 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, [&] { ReduceScalingGroupMeans(n_groups, scaling_refs, scaling_corr, partial_mean); - FitPerFrameG(scaling_obs, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial); + FitScalingFrames(scaling_fit, fit_start, partial_mean, g_partial); g_fitted = g_partial; frame_fitted = frame_scaled_scratch; }, [&](const std::vector &apply, const std::vector &ratio) { ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - const ScalingObs &o = scaling_obs[i]; - float &corr = scaling_corr[i]; + for (int q = lo; q < hi; ++q) { + const ScalingRef &o = scaling_refs[q]; + float &corr = scaling_corr[q]; if (frame_fitted[o.frame]) { // UpdateCorr const double denom = static_cast(o.partiality) * g_fitted[o.frame]; if (std::isfinite(o.prescaling_corr) && std::isfinite(denom) && denom > 0.0) @@ -6468,7 +6505,11 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool }); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) partials[i].corr = scaling_corr[i]; + for (int q = lo; q < hi; ++q) partials[scaling_index[q]].corr = scaling_corr[q]; + }); + // The memo keeps corr in the partials' own order. + ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) scaling_corr[i] = partials[i].corr; }); PartialLoopMemo m{key, partial_loop, frame_scaled_scratch, g_partial, std::move(scaling_corr)}; if (partial_loop_memos.size() == MAX_PARTIAL_LOOP_MEMOS) diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 5dd3a0290..0b8025e63 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -619,22 +619,35 @@ private: // The CPU partial scaling loop's working copy of the partials. It passes over them several times // per iteration and reads a handful of fields each time, so it works on these rather than on the - // eighty-byte Obs: ScalingObs holds, in the partials' own order, what FitPerFrameG and the corr - // update read; ScalingRef holds, in the CSR's group order, what the group means read besides corr, - // so that only corr is gathered. The same values, so the same arithmetic in the same order. - struct ScalingObs { - int32_t frame, group; - float I, sigma, partiality, prescaling_corr, zeta; - uint8_t on_ice; - }; + // eighty-byte Obs. ScalingRef holds what the group means and the corr update read, in the CSR's + // group order with the partials in no group after them, and corr is kept beside it in the same + // order - so neither pass gathers. ScalingFit holds, frame by frame, the observations + // FitScalingFrames fits: FitPerFrameG's filters, which do not change over the loop, applied once, + // and the factors of its sums that do not depend on the reference worked out once. The same + // values, so the same arithmetic in the same order. struct ScalingRef { - float I, sigma, partiality; + int32_t frame; + float I, sigma, partiality, prescaling_corr; + }; + struct ScalingFit { + int32_t group; + double coeff_per_mean; // partiality / prescaling_corr: FitPerFrameG's coeff for a mean of 1 + double Iobs, w2; }; // The observation's inverse-variance weight in the scaling reference, and its scaled intensity; // 0 when the observation is not in the reference at all. static double ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, - double min_partiality, float &I_corr); + double min_partiality, float &I_corr) { + if (group < 0) return 0.0; + if (!(corr > 0.0f) || !std::isfinite(corr)) return 0.0; + if (partiality < min_partiality) return 0.0; + const float Ic = I * corr; + const float sigma_corr = sigma * corr; + if (!std::isfinite(Ic) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0; + I_corr = Ic; + return 1.0 / (static_cast(sigma_corr) * sigma_corr); + } static double ReferenceWeight(const Obs &o, double min_partiality, float &I_corr) { return ReferenceWeight(o.group, o.corr, o.partiality, o.I, o.sigma, min_partiality, I_corr); } @@ -643,17 +656,20 @@ private: void ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const; // The same over the partials, using the CSR above - the same sums in the same order. void ReducePartialGroupMeans(int n_groups, std::vector &out_mean) const; - // The same again over the scaling loop's copy: refs in group order, corr in the partials' order. + // The same again over the scaling loop's copy, refs and corr both in group order. void ReduceScalingGroupMeans(int n_groups, const std::vector &refs, const std::vector &corr, std::vector &out_mean) const; // Robust per-frame G fit (IRLS, Cauchy k=3), unity=false uses the rotation partiality, unity=true the // scale-fulls (partiality already folded in). Reads out_mean[group] as the reference intensity. - // O is Obs or ScalingObs. template void FitPerFrameG(const std::vector &obs, const std::vector &fstart, const std::vector &fcount, const std::vector &group_mean_in, bool unity, std::vector &g); + // FitPerFrameG(unity = false) over the scaling loop's ScalingFit, frame f's at + // [fit_start[f], fit_start[f + 1]). + void FitScalingFrames(const std::vector &fit, const std::vector &fit_start, + const std::vector &group_mean_in, std::vector &g); // The alternating per-frame scale fit run to convergence: `iterate` fits every frame once against // the current reference (leaving G in `g` and the frames it fitted in frame_scaled_scratch), diff --git a/rugnux/rugnux_cli.cpp b/rugnux/rugnux_cli.cpp index 656ccde43..03b1a1edb 100644 --- a/rugnux/rugnux_cli.cpp +++ b/rugnux/rugnux_cli.cpp @@ -72,6 +72,9 @@ #include "../writer/WriteModel.h" #include "ResultReport.h" #include "RugnuxConsole.h" +#ifdef __APPLE__ +#include "../image_analysis/bragg_integration/BraggIntegrationEngineMetal.h" +#endif // Spots kept per image (the strongest ones) and handed to indexing. Offline reprocessing is not // bound by the online spot budget, so this is rugnux's own default rather than the 250 the @@ -868,6 +871,12 @@ static int RunRugnux(int argc, char **argv, bool no_gpu) { RunProvenance provenance; provenance.command_line = JoinCommandLine(argc, argv); provenance.gpu_count = get_gpu_count(); +#ifdef __APPLE__ + // Without CUDA, the Bragg integration's integer pixel work runs on the Apple GPU (MXAnalysisWithoutFPGA). + const std::string metal_device = BraggIntegrationEngineMetal::DeviceName(); +#else + const std::string metal_device; +#endif provenance.gpu_description = get_gpu_description(); // The terminal shows the run as steps (RugnuxConsole) and every Logger writes to the log file @@ -909,6 +918,7 @@ static int RunRugnux(int argc, char **argv, bool no_gpu) { logger.Info("Command line: {}", provenance.command_line); logger.Info("GPU: {}", no_gpu ? std::string("not used (--no-gpu) - running on the CPU") : provenance.gpu_count > 0 ? fmt::format("{} ({} visible)", provenance.gpu_description, provenance.gpu_count) + : !metal_device.empty() ? fmt::format("{} through Metal, for the Bragg integration only - the rest runs on the CPU", metal_device) : std::string("none visible - running on the CPU")); RegisterHDF5Filter(); @@ -1664,6 +1674,8 @@ static int RunRugnux(int argc, char **argv, bool no_gpu) { : provenance.gpu_count > 0 ? fmt::format("{}{}, {} threads", provenance.gpu_count > 1 ? fmt::format("{} x ", provenance.gpu_count) : "", provenance.gpu_description, nthreads) + : !metal_device.empty() + ? fmt::format("{} CPU threads, Bragg integration on {} (Metal)", nthreads, metal_device) : fmt::format("no GPU visible - {} CPU threads, which is much slower", nthreads)); console->Rule('-'); diff --git a/tests/BraggIntegrationEngineMetalTest.cpp b/tests/BraggIntegrationEngineMetalTest.cpp new file mode 100644 index 000000000..703dcfdd1 --- /dev/null +++ b/tests/BraggIntegrationEngineMetalTest.cpp @@ -0,0 +1,265 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#ifdef __APPLE__ + +#include +#include +#include +#include +#include + +#include "../common/BraggIntegrationSettings.h" +#include "../common/DetectorSetup.h" +#include "../common/DiffractionExperiment.h" +#include "../common/PixelMask.h" +#include "../common/Reflection.h" +#include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h" +#include "../image_analysis/bragg_integration/BraggIntegrationEngineMetal.h" + +// 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. + +namespace { + +struct Frame { + std::vector image; + std::vector predicted; +}; + +// A crowded random frame: predictions at random sub-pixel positions (so the pixel classification meets +// every kind of rounding near the stencil boundaries), a quarter of them on whole and half pixels (so it +// meets exact ties with the radii too), with random partialities (rocking-curve tails and flux, which +// mask the rings differently), about a third of them with a spot, overlapping where they fall close. On +// a background that rises with the radius (the radial correction has a curve to fit) with noise that +// takes it below zero near the beam, as a pedestal-subtracted detector does (a negative pixel is what +// carries into the high word of the GPU's 64-bit radial sums), hot pixels in the rings (so the +// background clip rejects something), and unreadable pixels: some in the run's mask, some saturated, +// some unreadable on this frame alone (both read as an overload), inside spots as well as between them. +Frame RandomFrame(const DiffractionExperiment &x, uint32_t seed, const std::vector &masked) { + const int W = static_cast(x.GetXPixelsNum()), H = static_cast(x.GetYPixelsNum()); + const float beam_x = x.GetBeamX_pxl(), beam_y = x.GetBeamY_pxl(); + std::mt19937 rng(seed); + std::uniform_real_distribution U(0.0f, 1.0f); + Frame f; + f.image.resize(static_cast(W) * H); + for (int y = 0; y < H; ++y) + for (int x_ = 0; x_ < W; ++x_) { + const float r = std::hypot(x_ - beam_x, y - beam_y); + const float mean = 0.5f + 0.01f * r; + f.image[static_cast(y) * W + x_] = static_cast(std::lround(mean + (U(rng) - 0.5f) * 8.0f)); + } + const int nrefl = W * H / 400; + int hkl = 1; + for (int n = 0; n < nrefl; ++n) { + Reflection r{}; + r.h = hkl; r.k = -hkl; r.l = hkl % 7; ++hkl; + r.predicted_x = 2.0f + U(rng) * (W - 4); + r.predicted_y = 2.0f + U(rng) * (H - 4); + if (n % 4 == 1) { + r.predicted_x = std::round(r.predicted_x * 2.0f) / 2.0f; + r.predicted_y = std::round(r.predicted_y); + } + const float u = U(rng); + r.partiality = u < 0.3f ? 0.01f * u : u; // a third of them rocking-curve tails + r.d = 1.2f + 3.0f * U(rng); + r.prescaling_corr = 1.0f; + f.predicted.push_back(r); + if (U(rng) < 0.35f) { + const float amp = 20.0f + 3000.0f * U(rng) * U(rng); + const float s = 0.8f + 1.4f * U(rng); + for (int dy = -8; dy <= 8; ++dy) + for (int dx = -8; dx <= 8; ++dx) { + const int px = static_cast(std::lround(r.predicted_x)) + dx; + const int py = static_cast(std::lround(r.predicted_y)) + dy; + if (px < 0 || py < 0 || px >= W || py >= H) continue; + const float ex = px - r.predicted_x, ey = py - r.predicted_y; + f.image[static_cast(py) * W + px] += static_cast( + std::lround(amp * std::exp(-(ex * ex + ey * ey) / (2.0f * s * s)))); + } + } + } + std::uniform_int_distribution P(0, f.image.size() - 1); + for (int k = 0; k < W * H / 500; ++k) + f.image[P(rng)] += 200 + static_cast(5000 * U(rng)); // hot pixels + for (int k = 0; k < W * H / 3000; ++k) + f.image[P(rng)] = U(rng) < 0.5f ? INT32_MIN : INT32_MAX; // unreadable on this frame + for (size_t i = 0; i < f.predicted.size(); i += 13) { // ... inside spots too + const auto &r = f.predicted[i]; + const size_t idx = static_cast(std::lround(r.predicted_y)) * W + std::lround(r.predicted_x) + i % 3; + if (idx < f.image.size()) f.image[idx] = i % 2 ? INT32_MIN : INT32_MAX; + } + for (size_t idx : masked) + f.image[idx] = INT32_MIN; // what the run's mask holds + return f; +} + +struct Case { + IntegratorMode mode = IntegratorMode::ProfileGaussian; + OverlapMode overlap = OverlapMode::Exclude; + bool radial = false; + std::optional bandwidth; + float stencil_k = 0.0f; + bool footprint = false; + float clip_nsigma = 4.0f; +}; + +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(x.GetXPixelsNum())) + .BeamY_pxl(0.55f * static_cast(x.GetYPixelsNum())); + x.BandwidthFWHM(c.bandwidth); + BraggIntegrationSettings s; + s.Integrator(c.mode).Overlap(c.overlap).BackgroundRadialCorrection(c.radial).StencilKSigma(c.stencil_k) + .BackgroundClipNSigma(c.clip_nsigma); + if (c.footprint) { + // Wider than the r1 disk can hold from the middle out, so both stencils occur on one frame. + SpotFootprint fp; + fp.bin_px = 100.0f; + for (int b = 0; b < 8; ++b) { + fp.sigma_rad.push_back(0.9f + 0.35f * b); + fp.sigma_tan.push_back(0.8f + 0.3f * b); + } + s.Footprint(fp); + } + x.ImportBraggIntegrationSettings(s); + return x; +} + +uint32_t Bits(float f) { return std::bit_cast(f); } + +void RequireSame(const std::vector &cpu, const std::vector &metal) { + REQUIRE(metal.size() == cpu.size()); + for (size_t i = 0; i < cpu.size(); ++i) { + const Reflection &a = cpu[i], &b = metal[i]; + INFO("reflection " << i << " hkl " << a.h << " " << a.k << " " << a.l + << " I " << a.I << " / " << b.I << " bkg " << a.bkg << " / " << b.bkg); + REQUIRE(a.h == b.h); + REQUIRE(a.k == b.k); + REQUIRE(a.l == b.l); + REQUIRE(Bits(a.image_number) == Bits(b.image_number)); + REQUIRE(Bits(a.I) == Bits(b.I)); + REQUIRE(Bits(a.sigma) == Bits(b.sigma)); + REQUIRE(Bits(a.bkg) == Bits(b.bkg)); + REQUIRE(Bits(a.var_bkg) == Bits(b.var_bkg)); + REQUIRE(Bits(a.observed_x) == Bits(b.observed_x)); + REQUIRE(Bits(a.observed_y) == Bits(b.observed_y)); + REQUIRE(Bits(a.image_scale_corr) == Bits(b.image_scale_corr)); + REQUIRE(Bits(a.prescaling_corr) == Bits(b.prescaling_corr)); + REQUIRE(a.observed == b.observed); + REQUIRE(a.overloaded == b.overloaded); + } +} + +void RequireSame(const BraggIntegrationCounts &a, const BraggIntegrationCounts &b) { + CHECK(a.predicted == b.predicted); + CHECK(a.bkg_starved == b.bkg_starved); + CHECK(a.bkg_starved_by_neighbour == b.bkg_starved_by_neighbour); + CHECK(a.profile_fallback == b.profile_fallback); + CHECK(a.overlapped == b.overlapped); + CHECK(a.overlapped_major == b.overlapped_major); +} + +// Three frames through one pair of engines - so whatever the Metal engine keeps between frames (its +// masks, its buffers) has to be left as the CPU engine leaves its own - each compared exactly. +// Returns how many reflections came out, so a case can show it was not vacuous. +size_t CompareFrames(const DetectorSetup &det, const Case &c) { + const DiffractionExperiment x = MakeExperiment(det, c); + REQUIRE(BraggIntegrationEngineMetal::Available(x)); + const size_t W = x.GetXPixelsNum(); + + // A few pixels of the run's mask, under spots and between them: unreadable there, but not overloads. + PixelMask mask(x); + std::vector user(x.GetPixelsNum(), 0); + std::vector masked; + for (size_t k = 0; k < x.GetPixelsNum(); k += 7919) { + user[k] = 1; + masked.push_back(k); + } + mask.LoadUserMask(x, user); + + BraggIntegrationEngineCPU cpu(x, mask); + BraggIntegrationEngineMetal metal(x, mask); + ImagePreprocessorBuffer image(x.GetPixelsNum()); + + size_t n_out = 0; + for (uint32_t frame = 0; frame < 3; ++frame) { + INFO("frame " << frame); + const Frame f = RandomFrame(x, 1000 * frame + static_cast(W), masked); + std::copy(f.image.begin(), f.image.end(), image.data()); + // The auto radial mode switches per image; here it is on for the radial cases, every image. + cpu.BackgroundRadial(c.radial); + metal.BackgroundRadial(c.radial); + const auto out_cpu = cpu.Run(image, f.predicted, f.predicted.size(), frame); + const auto out_metal = metal.Run(image, f.predicted, f.predicted.size(), frame); + RequireSame(out_cpu, out_metal); + n_out += out_cpu.size(); + } + RequireSame(cpu.Counts(), metal.Counts()); + // Not the CPU engine compared with itself: every frame went through the kernels. + REQUIRE(metal.FramesOnGPU() == 3); + return n_out; +} + +} // namespace + +TEST_CASE("BraggIntegrationEngineMetal_BitIdenticalToCPU", "[portable][Integration][Metal]") { + // A JUNGFRAU 1M, and a small odd-sized detector, where most boxes meet an edge. + const DetectorSetup large = DetJF(2); + const DetectorSetup small = DetDECTRIS(211, 197, "small", {}); + + std::vector> cases; + for (IntegratorMode mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical}) + for (OverlapMode overlap : {OverlapMode::Off, OverlapMode::Reject, OverlapMode::Exclude}) + for (bool radial : {false, true}) { + Case c; + c.mode = mode; + c.overlap = overlap; + c.radial = radial; + cases.emplace_back("mode " + std::to_string(static_cast(mode)) + " overlap " + + std::to_string(static_cast(overlap)) + (radial ? " radial" : ""), c); + } + // Rings elongated by the bandwidth, and the measured footprint: per-reflection stencils. + Case broadband; + broadband.bandwidth = 0.01f; + broadband.stencil_k = 3.0f; + broadband.radial = true; + cases.emplace_back("broadband stencil radial", broadband); + Case footprint; + footprint.footprint = true; + footprint.radial = true; + cases.emplace_back("footprint radial", footprint); + Case footprint_box = footprint; + footprint_box.mode = IntegratorMode::BoxSum; + cases.emplace_back("footprint box sum", footprint_box); + // No clip: the plain ring mean, and no clip kernel at all. + Case no_clip; + no_clip.clip_nsigma = 0.0f; + cases.emplace_back("no clip", no_clip); + + for (const auto &[name, c] : cases) { + DYNAMIC_SECTION(name) { + CHECK(CompareFrames(large, c) > 1000); + CHECK(CompareFrames(small, c) > 50); + } + } +} + +// 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]") { + Case trim; + trim.clip_nsigma = 0.0f; + DiffractionExperiment x = MakeExperiment(DetJF(2), trim); + BraggIntegrationSettings s = x.GetBraggIntegrationSettings(); + s.BackgroundTrimFraction(0.1f); + x.ImportBraggIntegrationSettings(s); + CHECK(!BraggIntegrationEngineMetal::Available(x)); +} + +#endif diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6f102faa2..c6dd46391 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -20,6 +20,7 @@ ADD_EXECUTABLE(jfjoch_portable_test EXCLUDE_FROM_ALL SpotFootprintTest.cpp BraggIntegrationEngineCompressedImageTest.cpp BraggIntegrationEngineCPUTest.cpp + BraggIntegrationEngineMetalTest.cpp WriteReflectionsTest.cpp LoadReferenceMtzFreeFlagsTest.cpp LoadReferenceSfCifTest.cpp