From 6326c8b30156fcc1091f49d1dba66d93ec02dc5f Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 12 Sep 2026 22:07:04 +0200 Subject: [PATCH] beam centre: the capture's transforms run on the GPU, and the CPU stays the fallback BeamCenterFFTScore is split the way the FFT indexer is: an engine interface with a cuFFT implementation and an fftw3f one, chosen by whether CUDA is compiled in, a device is visible and the card has room for the image. Everything that decides anything - the preparation, the masked Pearson, the shortlist and the margins - is shared, so the engines can differ only in how the four convolutions are computed, and the parity test compares two shortlists rather than two answers. The whole-detector transform was the entire added cost of the capture (4-15 s a run on CPU, all of it the transforms). On the device it is milliseconds, so the composition is now cheaper than the walk it replaced rather than dearer - which is what makes 2x2 binning, the other way out, unnecessary: full resolution is affordable and the binned capture was measured to pick a neighbouring peak on one dataset of 51. Device discipline, because the card is shared with the run's own analysis workers: the four convolution surfaces are brought back to the host and combined there, so the device holds only one real buffer and three spectra; the spectrum of the image is reused for its square; plans and buffers are created inside the call that needs them and freed when it returns; and an image that would not fit is scored on the CPU instead. Tests, none of which need a GPU or a dataset: the autoconvolution identity (an exactly symmetric image is recovered at integer and half-pixel centres, scoring exactly 1), shift equivariance, the four convolutions against brute force, the no-variance overlap that VARIANCE_FLOOR exists for, the smooth pad, GPU against CPU, and the composition - a walk seeded at the capture where the walk alone declines, the fallback to the capture alone at BEAM_CENTER_CAPTURE_SIGMA_PXL, and the beam stop being blanked out of the scored image. --- .../geom_refinement/BeamCenterFFT.cpp | 270 ++++------------ .../geom_refinement/BeamCenterFFT.h | 5 + .../geom_refinement/BeamCenterFFTCPU.cpp | 198 ++++++++++++ .../geom_refinement/BeamCenterFFTCPU.h | 16 + .../geom_refinement/BeamCenterFFTEngine.h | 58 ++++ .../geom_refinement/BeamCenterFFTGPU.cu | 298 ++++++++++++++++++ .../geom_refinement/BeamCenterFFTGPU.h | 30 ++ image_analysis/geom_refinement/CMakeLists.txt | 16 +- tests/BeamCenterFFTGPUTest.cpp | 148 +++++++++ tests/BeamCenterFFTTest.cpp | 224 +++++++++++++ tests/BeamCenterFromBackgroundTest.cpp | 125 ++++++++ tests/CMakeLists.txt | 1 + 12 files changed, 1185 insertions(+), 204 deletions(-) create mode 100644 image_analysis/geom_refinement/BeamCenterFFTCPU.cpp create mode 100644 image_analysis/geom_refinement/BeamCenterFFTCPU.h create mode 100644 image_analysis/geom_refinement/BeamCenterFFTEngine.h create mode 100644 image_analysis/geom_refinement/BeamCenterFFTGPU.cu create mode 100644 image_analysis/geom_refinement/BeamCenterFFTGPU.h create mode 100644 tests/BeamCenterFFTGPUTest.cpp diff --git a/image_analysis/geom_refinement/BeamCenterFFT.cpp b/image_analysis/geom_refinement/BeamCenterFFT.cpp index f17474384..044593b92 100644 --- a/image_analysis/geom_refinement/BeamCenterFFT.cpp +++ b/image_analysis/geom_refinement/BeamCenterFFT.cpp @@ -5,35 +5,19 @@ #include #include -#include -#include #include -#include #include -#include +#include "BeamCenterFFTCPU.h" +#include "BeamCenterFFTEngine.h" +#include "../../common/CUDAWrapper.h" +#ifdef JFJOCH_USE_CUDA +#include "BeamCenterFFTGPU.h" +#include "../../common/JFJochException.h" +#endif namespace { -// FFTW planning is not thread-safe (execution is); same idiom as FFTIndexerCPU. -std::mutex &FftwPlanMutex() { - static std::mutex m; - return m; -} - -// Smallest 2-3-5-7-smooth length >= n. A smooth transform length matters a lot here: the naive -// doubled Eiger2 16M size costs 3x the time and 4x the plan memory of the padded one (measured). -int64_t NextSmooth(int64_t n) { - for (int64_t k = n;; ++k) { - int64_t v = k; - for (int p : {2, 3, 5, 7}) - while (v % p == 0) - v /= p; - if (v == 1) - return k; - } -} - // The image prepared for scoring: hot pixels clipped at the 99.9th percentile, the median // background level removed, negatives clamped, masked pixels zero. The final subtraction of the // valid-pixel mean changes nothing mathematically - the masked Pearson is exactly invariant under @@ -88,67 +72,6 @@ Prepared PrepareImage(const std::vector &mean) { return out; } -// One real-to-complex 2D FFT workspace of fixed padded size, reused for every forward transform. -class Rfft2 { - int64_t ny, nx, nxc; - std::vector real; - fftwf_plan fwd = nullptr; - fftwf_plan bwd = nullptr; - -public: - std::vector> spectrum; - - Rfft2(int64_t ny_, int64_t nx_) : ny(ny_), nx(nx_), nxc(nx_ / 2 + 1) { - real.resize(static_cast(ny * nx)); - spectrum.resize(static_cast(ny * nxc)); - std::unique_lock lock(FftwPlanMutex()); - fwd = fftwf_plan_dft_r2c_2d(static_cast(ny), static_cast(nx), real.data(), - reinterpret_cast(spectrum.data()), - FFTW_ESTIMATE); - bwd = fftwf_plan_dft_c2r_2d(static_cast(ny), static_cast(nx), - reinterpret_cast(spectrum.data()), - real.data(), FFTW_ESTIMATE); - if (!fwd || !bwd) - throw std::runtime_error("BeamCenterFFT: fftwf plan failed"); - } - Rfft2(const Rfft2 &) = delete; - Rfft2 &operator=(const Rfft2 &) = delete; - ~Rfft2() { - std::unique_lock lock(FftwPlanMutex()); - if (fwd) - fftwf_destroy_plan(fwd); - if (bwd) - fftwf_destroy_plan(bwd); - } - - // Zero-pad `src` (height x width) into the workspace and return its spectrum. - std::vector> Forward(const std::vector &src, int64_t height, - int64_t width) { - std::fill(real.begin(), real.end(), 0.0f); - for (int64_t y = 0; y < height; y++) - std::memcpy(&real[static_cast(y * nx)], &src[static_cast(y * width)], - sizeof(float) * static_cast(width)); - fftwf_execute(fwd); - return spectrum; // copy: the workspace is reused - } - - // Inverse-transform the elementwise product u*v, normalised, cropped to out_h x out_w. - std::vector InverseProduct(const std::vector> &u, - const std::vector> &v, int64_t out_h, - int64_t out_w) { - const float norm = 1.0f / (static_cast(ny) * static_cast(nx)); - for (size_t i = 0; i < spectrum.size(); i++) - spectrum[i] = u[i] * v[i]; - fftwf_execute(bwd); - std::vector out(static_cast(out_h * out_w)); - for (int64_t y = 0; y < out_h; y++) - for (int64_t x = 0; x < out_w; x++) - out[static_cast(y * out_w + x)] = - real[static_cast(y * nx + x)] * norm; - return out; - } -}; - // Greedy shortlist of the 2D surface: strongest peak first, a square of half-width `nms` around // each taken peak suppressed. The surface is indexed by t = 2c, so the centre grid is half-pixel. std::vector Shortlist2D(std::vector &surface, int64_t h, int64_t w, @@ -217,128 +140,63 @@ float Margin(const std::vector &peaks) { // images every overlap that passes the pair-count gate is orders of magnitude above it. constexpr double VARIANCE_FLOOR = 1e-4; -std::vector PearsonSurface(const std::vector &C, const std::vector &S, - const std::vector &Q, const std::vector &D, - float min_pair_fraction, double global_variance) { - const size_t n = C.size(); +std::vector PearsonSurface(const BeamCenterConvSurfaces2D &conv, float min_pair_fraction, + double global_variance) { + const size_t n = conv.C.size(); float d_max = 0.0f; for (size_t i = 0; i < n; i++) - d_max = std::max(d_max, D[i]); + d_max = std::max(d_max, conv.D[i]); const double lim = min_pair_fraction * d_max; std::vector r(n, -std::numeric_limits::infinity()); for (size_t i = 0; i < n; i++) { - const double d = D[i]; + const double d = conv.D[i]; if (d <= lim) continue; - const double s = S[i]; - const double num = d * C[i] - s * s; - const double den = d * Q[i] - s * s; + const double s = conv.S[i]; + const double num = d * conv.C[i] - s * s; + const double den = d * conv.Q[i] - s * s; if (den > d * d * VARIANCE_FLOOR * global_variance) r[i] = static_cast(num / den); } return r; } -// 1D line-mirror score for the row coordinate of `a` (h x w row-major): each column correlated -// with its own mirror, the four convolution accumulators summed over columns in double complex -// (the same cancellation argument as above), one inverse transform each at the end. -std::vector LineScore(const std::vector &a, const std::vector &m, int64_t h, - int64_t w, float min_pair_fraction, double global_variance) { - const int64_t nfft = NextSmooth(2 * h); - const int64_t nc = nfft / 2 + 1; - - std::vector in(static_cast(nfft)); - std::vector> A(static_cast(nc)), M(static_cast(nc)), - A2(static_cast(nc)); - std::vector> accC(static_cast(nc)), accS(static_cast(nc)), - accQ(static_cast(nc)), accD(static_cast(nc)); - - fftwf_plan fwd, bwd; - { - std::unique_lock lock(FftwPlanMutex()); - fwd = fftwf_plan_dft_r2c_1d(static_cast(nfft), in.data(), - reinterpret_cast(A.data()), FFTW_ESTIMATE); - bwd = fftwf_plan_dft_c2r_1d(static_cast(nfft), - reinterpret_cast(A.data()), in.data(), - FFTW_ESTIMATE); - } - if (!fwd || !bwd) - throw std::runtime_error("BeamCenterFFT: fftwf 1D plan failed"); - - auto forward_into = [&](auto value_of, std::vector> &dst) { - for (int64_t y = 0; y < h; y++) - in[static_cast(y)] = value_of(y); - std::fill(in.begin() + h, in.end(), 0.0f); - fftwf_execute_dft_r2c(fwd, in.data(), reinterpret_cast(A.data())); - dst = A; - }; - - for (int64_t x = 0; x < w; x++) { - forward_into([&](int64_t y) { return a[static_cast(y * w + x)]; }, A2); - std::swap(A, A2); // A = spectrum of the column of a - std::vector> Acol = A; - forward_into([&](int64_t y) { return m[static_cast(y * w + x)]; }, M); - forward_into( - [&](int64_t y) { - const float v = a[static_cast(y * w + x)]; - return v * v; - }, - A2); - for (int64_t i = 0; i < nc; i++) { - const std::complex ca(Acol[i]), cm(M[i]), ca2(A2[i]); - accC[i] += ca * ca; - accS[i] += ca * cm; - accQ[i] += ca2 * cm; - accD[i] += cm * cm; - } - } - - const double norm = 1.0 / static_cast(nfft); - auto inverse = [&](const std::vector> &acc) { - for (int64_t i = 0; i < nc; i++) - A[i] = std::complex(acc[i]); - fftwf_execute_dft_c2r(bwd, reinterpret_cast(A.data()), in.data()); - std::vector out(static_cast(2 * h)); - for (int64_t i = 0; i < 2 * h; i++) - out[static_cast(i)] = in[static_cast(i)] * norm; - return out; - }; - const std::vector C = inverse(accC), S = inverse(accS), Q = inverse(accQ), - D = inverse(accD); - { - std::unique_lock lock(FftwPlanMutex()); - fftwf_destroy_plan(fwd); - fftwf_destroy_plan(bwd); - } - +// The same score for a 1D line mirror, whose accumulators were summed over the other coordinate. +std::vector PearsonLine(const BeamCenterConvSurfaces1D &conv, float min_pair_fraction, + double global_variance) { + const size_t n = conv.C.size(); double d_max = 0.0; - for (double d : D) + for (double d : conv.D) d_max = std::max(d_max, d); const double lim = min_pair_fraction * d_max; - std::vector r(static_cast(2 * h), -std::numeric_limits::infinity()); - for (int64_t i = 0; i < 2 * h; i++) { - const double num = D[i] * C[i] - S[i] * S[i]; - const double den = D[i] * Q[i] - S[i] * S[i]; + std::vector r(n, -std::numeric_limits::infinity()); + for (size_t i = 0; i < n; i++) { + const double num = conv.D[i] * conv.C[i] - conv.S[i] * conv.S[i]; + const double den = conv.D[i] * conv.Q[i] - conv.S[i] * conv.S[i]; // Same no-evidence floor as PearsonSurface: see VARIANCE_FLOOR. - if (D[i] > lim && den > D[i] * D[i] * VARIANCE_FLOOR * global_variance) - r[static_cast(i)] = static_cast(num / den); + if (conv.D[i] > lim && den > conv.D[i] * conv.D[i] * VARIANCE_FLOOR * global_variance) + r[i] = static_cast(num / den); } return r; } -std::vector Transpose(const std::vector &src, int64_t h, int64_t w) { - std::vector out(src.size()); - for (int64_t y = 0; y < h; y++) - for (int64_t x = 0; x < w; x++) - out[static_cast(x * h + y)] = src[static_cast(y * w + x)]; - return out; -} - } // namespace +int64_t BeamCenterFFTPadSize(int64_t n) { + for (int64_t k = n;; ++k) { + int64_t v = k; + for (int p : {2, 3, 5, 7}) + while (v % p == 0) + v /= p; + if (v == 1) + return k; + } +} + BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height, const std::vector &mean, - const BeamCenterFFTSettings &settings) { + const BeamCenterFFTSettings &settings, + BeamCenterFFTEngine &engine) { if (width <= 0 || height <= 0 || mean.size() != static_cast(width) * static_cast(height)) throw std::runtime_error("BeamCenterFFT: image dimensions do not match the projection"); @@ -346,35 +204,22 @@ BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height, BeamCenterFFTResult result; const Prepared prep = PrepareImage(mean); - // 2D point-inversion surface. { - Rfft2 fft(NextSmooth(2 * height), NextSmooth(2 * width)); - const auto A = fft.Forward(prep.a, height, width); - const auto M = fft.Forward(prep.m, height, width); - std::vector a2(prep.a.size()); - for (size_t i = 0; i < a2.size(); i++) - a2[i] = prep.a[i] * prep.a[i]; - const auto A2 = fft.Forward(a2, height, width); - - const auto C = fft.InverseProduct(A, A, 2 * height, 2 * width); - const auto S = fft.InverseProduct(A, M, 2 * height, 2 * width); - const auto Q = fft.InverseProduct(A2, M, 2 * height, 2 * width); - const auto D = fft.InverseProduct(M, M, 2 * height, 2 * width); - - auto surface = PearsonSurface(C, S, Q, D, settings.min_pair_fraction, prep.variance); + const auto conv = engine.PointSurfaces(prep.a, prep.m, height, width); + auto surface = PearsonSurface(conv, settings.min_pair_fraction, prep.variance); result.point = Shortlist2D(surface, 2 * height, 2 * width, settings.nms_radius_pxl, settings.candidates_point); } - // 1D line mirrors: the y score mirrors the row coordinate directly; the x score is the same - // computation on the transposed image. { - auto ry = LineScore(prep.a, prep.m, height, width, settings.min_pair_fraction, - prep.variance); + const auto conv_y = engine.LineSurfaces(prep.a, prep.m, height, width, + BeamCenterMirror::Rows); + auto ry = PearsonLine(conv_y, settings.min_pair_fraction, prep.variance); result.line_y = Shortlist1D(ry, settings.nms_radius_pxl, settings.candidates_line, false); - const auto at = Transpose(prep.a, height, width); - const auto mt = Transpose(prep.m, height, width); - auto rx = LineScore(at, mt, width, height, settings.min_pair_fraction, prep.variance); + + const auto conv_x = engine.LineSurfaces(prep.a, prep.m, height, width, + BeamCenterMirror::Columns); + auto rx = PearsonLine(conv_x, settings.min_pair_fraction, prep.variance); result.line_x = Shortlist1D(rx, settings.nms_radius_pxl, settings.candidates_line, true); } @@ -384,6 +229,25 @@ BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height, return result; } +BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height, + const std::vector &mean, + const BeamCenterFFTSettings &settings) { +#ifdef JFJOCH_USE_CUDA + if (get_gpu_count() > 0 && BeamCenterFFTGPU::FitsInDeviceMemory(width, height)) { + BeamCenterFFTGPU gpu; + try { + return BeamCenterFFTScore(width, height, mean, settings, gpu); + } catch (const JFJochException &) { + // The card is shared with this run's analysis workers, so the memory the check above + // saw can be gone by the time it is asked for. A capture is not worth failing a run + // over: score it on the CPU instead. + } + } +#endif + BeamCenterFFTCPU cpu; + return BeamCenterFFTScore(width, height, mean, settings, cpu); +} + BeamCenterFFTResult FindBeamCenterFFT(const DiffractionExperiment &experiment, const std::vector &mean, const BeamCenterFFTSettings &settings) { diff --git a/image_analysis/geom_refinement/BeamCenterFFT.h b/image_analysis/geom_refinement/BeamCenterFFT.h index 349987a4b..ff7cd188a 100644 --- a/image_analysis/geom_refinement/BeamCenterFFT.h +++ b/image_analysis/geom_refinement/BeamCenterFFT.h @@ -35,6 +35,11 @@ // // The input is the mean-projection the pre-scan already accumulates for the beam-stop search // (ShadowFinder::GetMeanProjection), NAN where nothing was counted - so this costs no new I/O. +// +// The transforms run on the GPU where there is one with room for them and on fftw3f otherwise +// (BeamCenterFFTGPU / BeamCenterFFTCPU); everything that decides anything is shared between the +// two. Whole-detector reach is what this buys, and on a 16 Mpixel detector the transforms are the +// whole of its cost, so which engine runs them is the difference between seconds and milliseconds. struct BeamCenterFFTCandidate { float beam_x_pxl = 0.0f; // NAN in a line_y candidate diff --git a/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp b/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp new file mode 100644 index 000000000..6c59f77d6 --- /dev/null +++ b/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp @@ -0,0 +1,198 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "BeamCenterFFTCPU.h" + +#include +#include +#include +#include +#include + +#include + +namespace { + +// FFTW planning is not thread-safe (execution is); same idiom as FFTIndexerCPU. +std::mutex &FftwPlanMutex() { + static std::mutex m; + return m; +} + +// One real-to-complex 2D FFT workspace of fixed padded size, reused for every forward transform. +class Rfft2 { + int64_t ny, nx, nxc; + std::vector real; + fftwf_plan fwd = nullptr; + fftwf_plan bwd = nullptr; + +public: + std::vector> spectrum; + + Rfft2(int64_t ny_, int64_t nx_) : ny(ny_), nx(nx_), nxc(nx_ / 2 + 1) { + real.resize(static_cast(ny * nx)); + spectrum.resize(static_cast(ny * nxc)); + std::unique_lock lock(FftwPlanMutex()); + fwd = fftwf_plan_dft_r2c_2d(static_cast(ny), static_cast(nx), real.data(), + reinterpret_cast(spectrum.data()), + FFTW_ESTIMATE); + bwd = fftwf_plan_dft_c2r_2d(static_cast(ny), static_cast(nx), + reinterpret_cast(spectrum.data()), + real.data(), FFTW_ESTIMATE); + if (!fwd || !bwd) + throw std::runtime_error("BeamCenterFFT: fftwf plan failed"); + } + Rfft2(const Rfft2 &) = delete; + Rfft2 &operator=(const Rfft2 &) = delete; + ~Rfft2() { + std::unique_lock lock(FftwPlanMutex()); + if (fwd) + fftwf_destroy_plan(fwd); + if (bwd) + fftwf_destroy_plan(bwd); + } + + // Zero-pad `src` (height x width) into the workspace and return its spectrum. + std::vector> Forward(const std::vector &src, int64_t height, + int64_t width) { + std::fill(real.begin(), real.end(), 0.0f); + for (int64_t y = 0; y < height; y++) + std::memcpy(&real[static_cast(y * nx)], &src[static_cast(y * width)], + sizeof(float) * static_cast(width)); + fftwf_execute(fwd); + return spectrum; // copy: the workspace is reused + } + + // Inverse-transform the elementwise product u*v, normalised, cropped to out_h x out_w. + std::vector InverseProduct(const std::vector> &u, + const std::vector> &v, int64_t out_h, + int64_t out_w) { + const float norm = 1.0f / (static_cast(ny) * static_cast(nx)); + for (size_t i = 0; i < spectrum.size(); i++) + spectrum[i] = u[i] * v[i]; + fftwf_execute(bwd); + std::vector out(static_cast(out_h * out_w)); + for (int64_t y = 0; y < out_h; y++) + for (int64_t x = 0; x < out_w; x++) + out[static_cast(y * out_w + x)] = + real[static_cast(y * nx + x)] * norm; + return out; + } +}; + +std::vector Transpose(const std::vector &src, int64_t h, int64_t w) { + std::vector out(src.size()); + for (int64_t y = 0; y < h; y++) + for (int64_t x = 0; x < w; x++) + out[static_cast(x * h + y)] = src[static_cast(y * w + x)]; + return out; +} + +} // namespace + +BeamCenterConvSurfaces2D BeamCenterFFTCPU::PointSurfaces(const std::vector &a, + const std::vector &m, int64_t h, + int64_t w) { + Rfft2 fft(BeamCenterFFTPadSize(2 * h), BeamCenterFFTPadSize(2 * w)); + const auto A = fft.Forward(a, h, w); + const auto M = fft.Forward(m, h, w); + std::vector a2(a.size()); + for (size_t i = 0; i < a2.size(); i++) + a2[i] = a[i] * a[i]; + const auto A2 = fft.Forward(a2, h, w); + + BeamCenterConvSurfaces2D out; + out.C = fft.InverseProduct(A, A, 2 * h, 2 * w); + out.S = fft.InverseProduct(A, M, 2 * h, 2 * w); + out.Q = fft.InverseProduct(A2, M, 2 * h, 2 * w); + out.D = fft.InverseProduct(M, M, 2 * h, 2 * w); + return out; +} + +// Each sequence is correlated with its own mirror and the four accumulators are summed over the +// other coordinate in double complex (the same cancellation argument as the 2D surface), so only +// four inverse transforms are needed however many sequences there are. +BeamCenterConvSurfaces1D BeamCenterFFTCPU::LineSurfaces(const std::vector &a_in, + const std::vector &m_in, int64_t h, + int64_t w, BeamCenterMirror mirror) { + // Mirroring the column coordinate is the same computation on the transposed image. + const std::vector at = mirror == BeamCenterMirror::Rows ? std::vector() + : Transpose(a_in, h, w); + const std::vector mt = mirror == BeamCenterMirror::Rows ? std::vector() + : Transpose(m_in, h, w); + const std::vector &a = mirror == BeamCenterMirror::Rows ? a_in : at; + const std::vector &m = mirror == BeamCenterMirror::Rows ? m_in : mt; + const int64_t nseq = mirror == BeamCenterMirror::Rows ? h : w; + const int64_t nbatch = mirror == BeamCenterMirror::Rows ? w : h; + + const int64_t nfft = BeamCenterFFTPadSize(2 * nseq); + const int64_t nc = nfft / 2 + 1; + + std::vector in(static_cast(nfft)); + std::vector> A(static_cast(nc)), M(static_cast(nc)), + A2(static_cast(nc)); + std::vector> accC(static_cast(nc)), accS(static_cast(nc)), + accQ(static_cast(nc)), accD(static_cast(nc)); + + fftwf_plan fwd, bwd; + { + std::unique_lock lock(FftwPlanMutex()); + fwd = fftwf_plan_dft_r2c_1d(static_cast(nfft), in.data(), + reinterpret_cast(A.data()), FFTW_ESTIMATE); + bwd = fftwf_plan_dft_c2r_1d(static_cast(nfft), + reinterpret_cast(A.data()), in.data(), + FFTW_ESTIMATE); + } + if (!fwd || !bwd) + throw std::runtime_error("BeamCenterFFT: fftwf 1D plan failed"); + + auto forward_into = [&](auto value_of, std::vector> &dst) { + for (int64_t y = 0; y < nseq; y++) + in[static_cast(y)] = value_of(y); + std::fill(in.begin() + nseq, in.end(), 0.0f); + fftwf_execute_dft_r2c(fwd, in.data(), reinterpret_cast(A.data())); + dst = A; + }; + + for (int64_t x = 0; x < nbatch; x++) { + forward_into([&](int64_t y) { return a[static_cast(y * nbatch + x)]; }, A2); + std::swap(A, A2); // A = spectrum of the sequence of a + std::vector> Acol = A; + forward_into([&](int64_t y) { return m[static_cast(y * nbatch + x)]; }, M); + forward_into( + [&](int64_t y) { + const float v = a[static_cast(y * nbatch + x)]; + return v * v; + }, + A2); + for (int64_t i = 0; i < nc; i++) { + const std::complex ca(Acol[i]), cm(M[i]), ca2(A2[i]); + accC[i] += ca * ca; + accS[i] += ca * cm; + accQ[i] += ca2 * cm; + accD[i] += cm * cm; + } + } + + const double norm = 1.0 / static_cast(nfft); + auto inverse = [&](const std::vector> &acc) { + for (int64_t i = 0; i < nc; i++) + A[i] = std::complex(acc[i]); + fftwf_execute_dft_c2r(bwd, reinterpret_cast(A.data()), in.data()); + std::vector out(static_cast(2 * nseq)); + for (int64_t i = 0; i < 2 * nseq; i++) + out[static_cast(i)] = in[static_cast(i)] * norm; + return out; + }; + BeamCenterConvSurfaces1D out; + out.C = inverse(accC); + out.S = inverse(accS); + out.Q = inverse(accQ); + out.D = inverse(accD); + { + std::unique_lock lock(FftwPlanMutex()); + fftwf_destroy_plan(fwd); + fftwf_destroy_plan(bwd); + } + return out; +} diff --git a/image_analysis/geom_refinement/BeamCenterFFTCPU.h b/image_analysis/geom_refinement/BeamCenterFFTCPU.h new file mode 100644 index 000000000..0385ba084 --- /dev/null +++ b/image_analysis/geom_refinement/BeamCenterFFTCPU.h @@ -0,0 +1,16 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include "BeamCenterFFTEngine.h" + +// The beam-centre transforms on fftw3f. Always built (FFTW is fetched unconditionally), and the +// fallback wherever there is no GPU to run them on. +class BeamCenterFFTCPU : public BeamCenterFFTEngine { +public: + BeamCenterConvSurfaces2D PointSurfaces(const std::vector &a, const std::vector &m, + int64_t h, int64_t w) override; + BeamCenterConvSurfaces1D LineSurfaces(const std::vector &a, const std::vector &m, + int64_t h, int64_t w, BeamCenterMirror mirror) override; +}; diff --git a/image_analysis/geom_refinement/BeamCenterFFTEngine.h b/image_analysis/geom_refinement/BeamCenterFFTEngine.h new file mode 100644 index 000000000..b711645b2 --- /dev/null +++ b/image_analysis/geom_refinement/BeamCenterFFTEngine.h @@ -0,0 +1,58 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include + +#include "BeamCenterFFT.h" + +// The transform half of BeamCenterFFTScore, with a CPU (fftw3f) and a GPU (cuFFT) implementation +// chosen the way the FFT indexer's are. Everything that decides anything - the preparation of the +// image, the masked Pearson combination, the shortlist and the margins - is shared, so the two +// engines can differ only in how the four convolutions are computed. That is what makes them +// comparable: the parity test scores one image with each and compares the shortlists. +// +// The masked Pearson needs four convolutions of the prepared image `a` and its valid-pixel mask +// `m`: C = a (*) a, S = a (*) m, Q = a^2 (*) m and D = m (*) m. + +// The 2D point-inversion score: each surface cropped to 2h x 2w, row-major. +struct BeamCenterConvSurfaces2D { + std::vector C, S, Q, D; +}; + +// The 1D line-mirror score: each sequence correlated with its own mirror and the four accumulators +// summed over the other coordinate, so each surface is 2 x (the mirrored extent) long. +struct BeamCenterConvSurfaces1D { + std::vector C, S, Q, D; +}; + +// Which coordinate the 1D score mirrors: Rows mirrors y (each column is one sequence) and gives +// the beam's y; Columns mirrors x and gives its x. +enum class BeamCenterMirror { Rows, Columns }; + +class BeamCenterFFTEngine { +public: + virtual ~BeamCenterFFTEngine() = default; + + // `a` and `m` are h x w row-major, as PrepareImage leaves them. + virtual BeamCenterConvSurfaces2D PointSurfaces(const std::vector &a, + const std::vector &m, int64_t h, + int64_t w) = 0; + virtual BeamCenterConvSurfaces1D LineSurfaces(const std::vector &a, + const std::vector &m, int64_t h, int64_t w, + BeamCenterMirror mirror) = 0; +}; + +// Smallest 2-3-5-7-smooth transform length >= n. Both engines pad to it, so both score the same +// grid. A smooth length is not a micro-optimisation here: the naive next-power-of-two pad of a +// 16 Mpixel detector costs 3x the time of the smooth one on both FFTW and cuFFT, and 2.5 GB of +// cuFFT plan memory against a few hundred MB (measured). +[[nodiscard]] int64_t BeamCenterFFTPadSize(int64_t n); + +// BeamCenterFFTScore with the engine named rather than chosen, for the GPU/CPU parity test. +[[nodiscard]] BeamCenterFFTResult BeamCenterFFTScore(int64_t width, int64_t height, + const std::vector &mean, + const BeamCenterFFTSettings &settings, + BeamCenterFFTEngine &engine); diff --git a/image_analysis/geom_refinement/BeamCenterFFTGPU.cu b/image_analysis/geom_refinement/BeamCenterFFTGPU.cu new file mode 100644 index 000000000..797591056 --- /dev/null +++ b/image_analysis/geom_refinement/BeamCenterFFTGPU.cu @@ -0,0 +1,298 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "BeamCenterFFTGPU.h" + +#include + +#include +#include + +#include "../indexing/CUDAMemHelpers.h" + +namespace { + +constexpr int BLOCK = 256; +// Threads per frequency in the 1D accumulation. A power of two, for the shared-memory reduction. +constexpr int ACCUMULATE_THREADS = 128; + +int Blocks(int64_t n) { return static_cast((n + BLOCK - 1) / BLOCK); } + +void Check(cudaError_t err, const char *what) { + if (err != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, what); +} + +void Check(cufftResult res, const char *what) { + if (res != CUFFT_SUCCESS) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, what); +} + +std::vector Dims(int64_t a) { return {static_cast(a)}; } +std::vector Dims(int64_t a, int64_t b) { return {static_cast(a), static_cast(b)}; } + +__global__ void SquareInPlace(float *data, int64_t n) { + const int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i < n) + data[i] *= data[i]; +} + +__global__ void MultiplySpectra(const cufftComplex *u, const cufftComplex *v, cufftComplex *out, + int64_t n) { + const int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i < n) + out[i] = cuCmulf(u[i], v[i]); +} + +// The inverse transform leaves an unnormalised ny x nx result; only its 2h x 2w corner is the +// convolution surface, so scale that in place and copy it out with one strided copy. +__global__ void ScaleCorner(float *data, int64_t nx, int64_t out_h, int64_t out_w, float norm) { + const int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= out_h * out_w) + return; + const int64_t y = i / out_w; + data[y * nx + (i - y * out_w)] *= norm; +} + +// Lay each sequence out contiguously, zero-padded to nfft, so the 1D transforms are a plain batch: +// out[b * nfft + s] = img[b * stride_batch + s * stride_seq], squared where asked. +__global__ void GatherSequences(const float *img, int64_t nseq, int64_t nbatch, int64_t stride_seq, + int64_t stride_batch, int64_t nfft, int square, float *out) { + const int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= nbatch * nfft) + return; + const int64_t b = i / nfft; + const int64_t s = i - b * nfft; + float v = 0.0f; + if (s < nseq) { + v = img[b * stride_batch + s * stride_seq]; + if (square) + v *= v; + } + out[i] = v; +} + +// The four accumulators of the 1D score, summed over the batch in double - the same widening the +// host path does, and for the same reason: the numerator and denominator below are cancellations +// of large near-equal terms. One block per frequency. +__global__ void AccumulateSpectra(const cufftComplex *A, const cufftComplex *M, + const cufftComplex *A2, int64_t nbatch, int64_t nc, double2 *accC, + double2 *accS, double2 *accQ, double2 *accD) { + __shared__ double2 red[4][ACCUMULATE_THREADS]; + const int64_t i = blockIdx.x; + double2 c = {0.0, 0.0}, s = {0.0, 0.0}, q = {0.0, 0.0}, d = {0.0, 0.0}; + for (int64_t b = threadIdx.x; b < nbatch; b += blockDim.x) { + const cufftComplex fa = A[b * nc + i], fm = M[b * nc + i], fa2 = A2[b * nc + i]; + const double ar = fa.x, ai = fa.y, mr = fm.x, mi = fm.y, qr = fa2.x, qi = fa2.y; + c.x += ar * ar - ai * ai; + c.y += ar * ai + ai * ar; + s.x += ar * mr - ai * mi; + s.y += ar * mi + ai * mr; + q.x += qr * mr - qi * mi; + q.y += qr * mi + qi * mr; + d.x += mr * mr - mi * mi; + d.y += mr * mi + mi * mr; + } + red[0][threadIdx.x] = c; + red[1][threadIdx.x] = s; + red[2][threadIdx.x] = q; + red[3][threadIdx.x] = d; + for (unsigned stride = blockDim.x / 2; stride > 0; stride /= 2) { + __syncthreads(); + if (threadIdx.x < stride) + for (int k = 0; k < 4; k++) { + red[k][threadIdx.x].x += red[k][threadIdx.x + stride].x; + red[k][threadIdx.x].y += red[k][threadIdx.x + stride].y; + } + } + if (threadIdx.x == 0) { + accC[i] = red[0][0]; + accS[i] = red[1][0]; + accQ[i] = red[2][0]; + accD[i] = red[3][0]; + } +} + +// The four accumulators narrowed back to single precision and packed as one batch of four, ready +// for the inverse transform - exactly what the host path does before its four inverses. +__global__ void PackAccumulators(const double2 *accC, const double2 *accS, const double2 *accQ, + const double2 *accD, int64_t nc, cufftComplex *out) { + const int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= nc) + return; + const double2 *acc[4] = {accC, accS, accQ, accD}; + for (int k = 0; k < 4; k++) + out[k * nc + i] = make_cuFloatComplex(static_cast(acc[k][i].x), + static_cast(acc[k][i].y)); +} + +void CheckLastKernel(const char *what) { + Check(cudaGetLastError(), what); +} + +} // namespace + +BeamCenterConvSurfaces2D BeamCenterFFTGPU::PointSurfaces(const std::vector &a, + const std::vector &m, int64_t h, + int64_t w) { + const int64_t ny = BeamCenterFFTPadSize(2 * h); + const int64_t nx = BeamCenterFFTPadSize(2 * w); + const int64_t nxc = nx / 2 + 1; + const int64_t nreal = ny * nx; + const int64_t ncomplex = ny * nxc; + + CudaDevicePtr d_real(nreal); + CudaDevicePtr d_a(ncomplex), d_m(ncomplex), d_prod(ncomplex); + + CudaFFTPlan fwd(2, Dims(ny, nx), Dims(ny, nx), 1, static_cast(nreal), Dims(ny, nxc), 1, + static_cast(ncomplex), CUFFT_R2C, 1); + CudaFFTPlan inv(2, Dims(ny, nx), Dims(ny, nxc), 1, static_cast(ncomplex), Dims(ny, nx), 1, + static_cast(nreal), CUFFT_C2R, 1); + + auto upload = [&](const std::vector &src) { + Check(cudaMemset(d_real.get(), 0, nreal * sizeof(float)), "BeamCenterFFT: memset failed"); + Check(cudaMemcpy2D(d_real.get(), nx * sizeof(float), src.data(), w * sizeof(float), + w * sizeof(float), h, cudaMemcpyHostToDevice), + "BeamCenterFFT: image upload failed"); + }; + + upload(a); + Check(cufftExecR2C(fwd, d_real, d_a), "BeamCenterFFT: forward transform failed"); + upload(m); + Check(cufftExecR2C(fwd, d_real, d_m), "BeamCenterFFT: forward transform failed"); + + const float norm = 1.0f / (static_cast(ny) * static_cast(nx)); + auto inverse_product = [&](const cufftComplex *u, const cufftComplex *v) { + MultiplySpectra<<>>(u, v, d_prod.get(), ncomplex); + CheckLastKernel("BeamCenterFFT: spectrum product failed"); + Check(cufftExecC2R(inv, d_prod, d_real), "BeamCenterFFT: inverse transform failed"); + ScaleCorner<<>>(d_real.get(), nx, 2 * h, 2 * w, norm); + CheckLastKernel("BeamCenterFFT: surface scaling failed"); + std::vector out(static_cast(4 * h * w)); + Check(cudaMemcpy2D(out.data(), 2 * w * sizeof(float), d_real.get(), nx * sizeof(float), + 2 * w * sizeof(float), 2 * h, cudaMemcpyDeviceToHost), + "BeamCenterFFT: surface download failed"); + return out; + }; + + BeamCenterConvSurfaces2D out; + out.C = inverse_product(d_a, d_a); + out.S = inverse_product(d_a, d_m); + out.D = inverse_product(d_m, d_m); + // The spectrum of a is finished with, so a^2 goes in its place rather than in a fourth buffer. + upload(a); + SquareInPlace<<>>(d_real.get(), nreal); + CheckLastKernel("BeamCenterFFT: squaring failed"); + Check(cufftExecR2C(fwd, d_real, d_a), "BeamCenterFFT: forward transform failed"); + out.Q = inverse_product(d_a, d_m); + return out; +} + +BeamCenterConvSurfaces1D BeamCenterFFTGPU::LineSurfaces(const std::vector &a, + const std::vector &m, int64_t h, + int64_t w, BeamCenterMirror mirror) { + const bool rows = mirror == BeamCenterMirror::Rows; + const int64_t nseq = rows ? h : w; + const int64_t nbatch = rows ? w : h; + const int64_t stride_seq = rows ? w : 1; // step between two samples of one sequence + const int64_t stride_batch = rows ? 1 : w; // step between two sequences + const int64_t nfft = BeamCenterFFTPadSize(2 * nseq); + const int64_t nc = nfft / 2 + 1; + + BeamCenterConvSurfaces1D out; + CudaDevicePtr accC(nc), accS(nc), accQ(nc), accD(nc); + { + CudaDevicePtr d_img(h * w), d_in(nfft * nbatch); + CudaDevicePtr d_a(nbatch * nc), d_m(nbatch * nc), d_a2(nbatch * nc); + CudaFFTPlan fwd(1, Dims(nfft), Dims(nfft), 1, static_cast(nfft), Dims(nc), 1, + static_cast(nc), CUFFT_R2C, static_cast(nbatch)); + + auto transform = [&](int square, cufftComplex *dst) { + GatherSequences<<>>(d_img.get(), nseq, nbatch, stride_seq, + stride_batch, nfft, square, + d_in.get()); + CheckLastKernel("BeamCenterFFT: sequence gather failed"); + Check(cufftExecR2C(fwd, d_in, dst), "BeamCenterFFT: forward transform failed"); + }; + auto upload = [&](const std::vector &src) { + Check(cudaMemcpy(d_img.get(), src.data(), h * w * sizeof(float), + cudaMemcpyHostToDevice), + "BeamCenterFFT: image upload failed"); + }; + + upload(a); + transform(0, d_a); + transform(1, d_a2); + upload(m); + transform(0, d_m); + + AccumulateSpectra<<(nc), ACCUMULATE_THREADS>>>(d_a, d_m, d_a2, nbatch, nc, + accC, accS, accQ, accD); + CheckLastKernel("BeamCenterFFT: spectrum accumulation failed"); + } + + CudaDevicePtr d_pack(4 * nc); + CudaDevicePtr d_out(4 * nfft); + PackAccumulators<<>>(accC, accS, accQ, accD, nc, d_pack.get()); + CheckLastKernel("BeamCenterFFT: accumulator packing failed"); + CudaFFTPlan inv(1, Dims(nfft), Dims(nc), 1, static_cast(nc), Dims(nfft), 1, + static_cast(nfft), CUFFT_C2R, 4); + Check(cufftExecC2R(inv, d_pack, d_out), "BeamCenterFFT: inverse transform failed"); + + std::vector host(static_cast(4 * nfft)); + Check(cudaMemcpy(host.data(), d_out.get(), host.size() * sizeof(float), + cudaMemcpyDeviceToHost), + "BeamCenterFFT: line surface download failed"); + + const double norm = 1.0 / static_cast(nfft); + auto take = [&](int k) { + std::vector v(static_cast(2 * nseq)); + for (int64_t i = 0; i < 2 * nseq; i++) + v[static_cast(i)] = host[static_cast(k * nfft + i)] * norm; + return v; + }; + out.C = take(0); + out.S = take(1); + out.Q = take(2); + out.D = take(3); + return out; +} + +size_t BeamCenterFFTGPU::DeviceMemoryNeeded(int64_t width, int64_t height) { + const int64_t ny = BeamCenterFFTPadSize(2 * height); + const int64_t nx = BeamCenterFFTPadSize(2 * width); + const int64_t nxc = nx / 2 + 1; + size_t work_fwd = 0, work_inv = 0; + cufftEstimate2d(static_cast(ny), static_cast(nx), CUFFT_R2C, &work_fwd); + cufftEstimate2d(static_cast(ny), static_cast(nx), CUFFT_C2R, &work_inv); + size_t needed = ny * nx * sizeof(float) + 3 * ny * nxc * sizeof(cufftComplex) + work_fwd + + work_inv; + + // The two 1D phases, the larger of them: they run after the point surfaces, with its buffers + // already freed. + for (int pass = 0; pass < 2; pass++) { + const int64_t nseq = pass == 0 ? height : width; + const int64_t nbatch = pass == 0 ? width : height; + const int64_t nfft = BeamCenterFFTPadSize(2 * nseq); + const int64_t nc = nfft / 2 + 1; + std::vector n = Dims(nfft), embed_real = Dims(nfft), embed_complex = Dims(nc); + size_t work = 0; + cufftEstimateMany(1, n.data(), embed_real.data(), 1, static_cast(nfft), + embed_complex.data(), 1, static_cast(nc), CUFFT_R2C, + static_cast(nbatch), &work); + const size_t line = (height * width + nfft * nbatch) * sizeof(float) + + 3 * nbatch * nc * sizeof(cufftComplex) + work; + needed = std::max(needed, line); + } + return needed; +} + +bool BeamCenterFFTGPU::FitsInDeviceMemory(int64_t width, int64_t height) { + size_t free_bytes = 0, total_bytes = 0; + if (cudaMemGetInfo(&free_bytes, &total_bytes) != cudaSuccess) + return false; + // Headroom for the allocator's own rounding, and for whatever else lands on a shared card + // between this question and the allocations it is asked about. + constexpr size_t HEADROOM = 512ull << 20; + return DeviceMemoryNeeded(width, height) + HEADROOM <= free_bytes; +} diff --git a/image_analysis/geom_refinement/BeamCenterFFTGPU.h b/image_analysis/geom_refinement/BeamCenterFFTGPU.h new file mode 100644 index 000000000..a3f2c2ce3 --- /dev/null +++ b/image_analysis/geom_refinement/BeamCenterFFTGPU.h @@ -0,0 +1,30 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Included only where CUDA is known to be present, i.e. under JFJOCH_USE_CUDA - the same rule as +// FFTIndexerGPU.h. Deliberately free of CUDA headers so a host translation unit (the parity test) +// can construct the engine. + +#include "BeamCenterFFTEngine.h" + +// The beam-centre transforms on cuFFT. The capture is one call per run on an image of tens of +// megapixels, so plans and device buffers are created inside the call that needs them and freed +// when it returns: the card is shared with the analysis workers, and a plan held across the run +// would be several hundred megabytes doing nothing. +class BeamCenterFFTGPU : public BeamCenterFFTEngine { +public: + BeamCenterConvSurfaces2D PointSurfaces(const std::vector &a, const std::vector &m, + int64_t h, int64_t w) override; + BeamCenterConvSurfaces1D LineSurfaces(const std::vector &a, const std::vector &m, + int64_t h, int64_t w, BeamCenterMirror mirror) override; + + // Peak device memory the two phases need for an image of this size: buffers plus the cuFFT + // plans' own workspace. + [[nodiscard]] static size_t DeviceMemoryNeeded(int64_t width, int64_t height); + // Whether the visible device has that much free right now. The card is shared - with the + // analysis workers of this run and, on a development box, with whatever else is on it - so a + // capture that would not fit runs on the CPU instead of taking the card down with it. + [[nodiscard]] static bool FitsInDeviceMemory(int64_t width, int64_t height); +}; diff --git a/image_analysis/geom_refinement/CMakeLists.txt b/image_analysis/geom_refinement/CMakeLists.txt index 56510285f..b004e8d4e 100644 --- a/image_analysis/geom_refinement/CMakeLists.txt +++ b/image_analysis/geom_refinement/CMakeLists.txt @@ -6,6 +6,9 @@ ADD_LIBRARY(JFJochGeomRefinement STATIC BeamCenterFromSpots.h BeamCenterFFT.cpp BeamCenterFFT.h + BeamCenterFFTEngine.h + BeamCenterFFTCPU.cpp + BeamCenterFFTCPU.h RingOptimizer.cpp RingOptimizer.h AssignSpotsToRings.cpp @@ -30,5 +33,16 @@ ADD_LIBRARY(JFJochGeomRefinement STATIC ) # fftw3f: BeamCenterFFT scores every candidate centre with one transform set (always fetched, -# same target the CPU FFT indexer links). +# same target the CPU FFT indexer links), and is the fallback wherever there is no GPU. TARGET_LINK_LIBRARIES(JFJochGeomRefinement Ceres::ceres Eigen3::Eigen JFJochCommon fftw3f) + +IF (JFJOCH_CUDA_AVAILABLE) + TARGET_SOURCES(JFJochGeomRefinement PRIVATE BeamCenterFFTGPU.cu BeamCenterFFTGPU.h) + # Same static/dynamic cuFFT choice as the FFT indexer, and for the same reasons - see the long + # note in image_analysis/indexing/CMakeLists.txt. + IF (JFJOCH_PORTABLE_ONLY AND TARGET CUDA::cufft_static) + TARGET_LINK_LIBRARIES(JFJochGeomRefinement CUDA::cufft_static) + ELSE() + TARGET_LINK_LIBRARIES(JFJochGeomRefinement CUDA::cufft) + ENDIF() +ENDIF() diff --git a/tests/BeamCenterFFTGPUTest.cpp b/tests/BeamCenterFFTGPUTest.cpp new file mode 100644 index 000000000..3c82a002c --- /dev/null +++ b/tests/BeamCenterFFTGPUTest.cpp @@ -0,0 +1,148 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include +#include "../common/CUDAWrapper.h" + +#ifdef JFJOCH_USE_CUDA + +#include +#include +#include +#include + +#include "../image_analysis/geom_refinement/BeamCenterFFTCPU.h" +#include "../image_analysis/geom_refinement/BeamCenterFFTGPU.h" + +namespace { + +// An isotropic scattered background about (cx, cy) with a solvent ring, a module gap and a darker +// sextant - the same shape of image the corpus measurement runs on, at a size a test can afford. +std::vector SynthesiseBackground(int64_t w, int64_t h, float cx, float cy) { + std::vector mean(static_cast(w * h), NAN); + std::mt19937 rng(20260913); + std::normal_distribution gauss(0.0f, 1.0f); + for (int64_t y = 0; y < h; y++) { + if (y >= 500 && y < 520) + continue; // module gap + for (int64_t x = 0; x < w; x++) { + const float r = std::hypot(static_cast(x) - cx, static_cast(y) - cy); + const float t = (r - 330.0f) / 40.0f; + float value = 130.0f * std::exp(-r / 260.0f) + 50.0f * std::exp(-0.5f * t * t) + 6.0f; + const float phi = std::atan2(static_cast(y) - cy, static_cast(x) - cx); + if (phi > 0.0f && phi < 1.0f) + value *= 0.8f; + mean[static_cast(y * w + x)] = value + gauss(rng) * std::sqrt(value / 60.0f); + } + } + return mean; +} + +void CompareShortlist(const std::vector &gpu, + const std::vector &cpu, const char *what) { + INFO(what); + REQUIRE(gpu.size() == cpu.size()); + for (size_t i = 0; i < cpu.size(); i++) { + INFO("candidate " << i); + // The positions are what a caller consumes, and they are a grid index: they must be the + // same index, not a nearby one. + CHECK(((gpu[i].beam_x_pxl == cpu[i].beam_x_pxl) + || (std::isnan(gpu[i].beam_x_pxl) && std::isnan(cpu[i].beam_x_pxl)))); + CHECK(((gpu[i].beam_y_pxl == cpu[i].beam_y_pxl) + || (std::isnan(gpu[i].beam_y_pxl) && std::isnan(cpu[i].beam_y_pxl)))); + CHECK(gpu[i].score == Catch::Approx(cpu[i].score).epsilon(1e-5)); + } +} + +} // namespace + +// The GPU engine exists to make the capture affordable, not to answer differently: the two engines +// share the preparation, the masked Pearson, the shortlist and the margins, and differ only in how +// the four convolutions are computed. So the shortlists must be the same grid positions and the +// margins must agree - a capture that depended on which engine ran it would be a capture that could +// not be reasoned about. +TEST_CASE("BeamCenterFFTGPU_MatchesTheCPUEngine", "[BeamCenter][BeamCenterFFTGPU]") { + if (get_gpu_count() == 0) { + WARN("No CUDA GPU present. Skipping BeamCenterFFTGPU_MatchesTheCPUEngine"); + return; + } + + const int64_t w = 1100, h = 900; + const float cx = 397.25f, cy = 511.5f; + const auto mean = SynthesiseBackground(w, h, cx, cy); + + BeamCenterFFTCPU cpu_engine; + BeamCenterFFTGPU gpu_engine; + const BeamCenterFFTSettings settings; + const auto cpu = BeamCenterFFTScore(w, h, mean, settings, cpu_engine); + const auto gpu = BeamCenterFFTScore(w, h, mean, settings, gpu_engine); + + REQUIRE(!cpu.point.empty()); + CompareShortlist(gpu.point, cpu.point, "point surface"); + CompareShortlist(gpu.line_x, cpu.line_x, "x line mirror"); + CompareShortlist(gpu.line_y, cpu.line_y, "y line mirror"); + + // The margin is the surface's self-diagnostic, and it is a difference of two scores that sit + // within a fraction of a percent of each other - the quantity most exposed to a change of + // transform. + CHECK(gpu.margin_point == Catch::Approx(cpu.margin_point).margin(1e-5)); + CHECK(gpu.margin_line_x == Catch::Approx(cpu.margin_line_x).margin(1e-5)); + CHECK(gpu.margin_line_y == Catch::Approx(cpu.margin_line_y).margin(1e-5)); +} + +// The same identity test the CPU engine carries, on the device: an exactly centrosymmetric image +// scores exactly 1 at its centre and the capture lands on that grid point. It is the cheapest +// check there is for an indexing or sign error that only the device path has. +TEST_CASE("BeamCenterFFTGPU_RecoversAnExactlySymmetricCentre", "[BeamCenter][BeamCenterFFTGPU]") { + if (get_gpu_count() == 0) { + WARN("No CUDA GPU present. Skipping BeamCenterFFTGPU_RecoversAnExactlySymmetricCentre"); + return; + } + + // Keyed on the canonical member of the +- pair, so the image is exactly its own mirror about + // (cx, cy) - the same construction the CPU identity test uses. + const int64_t w = 220, h = 180; + const float cx = 83.5f, cy = 61.5f; + std::vector mean(static_cast(w * h)); + for (int64_t y = 0; y < h; y++) + for (int64_t x = 0; x < w; x++) { + int64_t dx = std::lround(2.0f * (static_cast(x) - cx)); + int64_t dy = std::lround(2.0f * (static_cast(y) - cy)); + if (dy < 0 || (dy == 0 && dx < 0)) { + dx = -dx; + dy = -dy; + } + uint64_t k = static_cast(dx) * 0x9e3779b97f4a7c15ULL + + static_cast(dy) * 0xc2b2ae3d27d4eb4fULL; + k ^= k >> 31; + k *= 0xbf58476d1ce4e5b9ULL; + k ^= k >> 29; + mean[static_cast(y * w + x)] = 100.0f + static_cast(k % 4096u) * 0.25f; + } + + BeamCenterFFTGPU engine; + const auto r = BeamCenterFFTScore(w, h, mean, BeamCenterFFTSettings{}, engine); + REQUIRE(!r.point.empty()); + CHECK(r.point[0].beam_x_pxl == cx); + CHECK(r.point[0].beam_y_pxl == cy); + CHECK(r.point[0].score == Catch::Approx(1.0f).margin(1e-4)); +} + +// The card is shared, so the engine has to be able to say that an image will not fit before it +// starts allocating - that is what sends the capture to the CPU instead of taking the run down. +TEST_CASE("BeamCenterFFTGPU_SizesItsOwnFootprint", "[BeamCenter][BeamCenterFFTGPU]") { + if (get_gpu_count() == 0) { + WARN("No CUDA GPU present. Skipping BeamCenterFFTGPU_SizesItsOwnFootprint"); + return; + } + + // A 16 Mpixel detector is the largest this runs on, and it has to be a bounded cost. + CHECK(BeamCenterFFTGPU::DeviceMemoryNeeded(4148, 4362) < (3ull << 30)); + // Bigger images need more. + CHECK(BeamCenterFFTGPU::DeviceMemoryNeeded(4148, 4362) + > BeamCenterFFTGPU::DeviceMemoryNeeded(1030, 1064)); + // Nothing on this card can hold a hundred-gigapixel image. + CHECK(!BeamCenterFFTGPU::FitsInDeviceMemory(300000, 300000)); +} + +#endif diff --git a/tests/BeamCenterFFTTest.cpp b/tests/BeamCenterFFTTest.cpp index 70279d323..095d2fec2 100644 --- a/tests/BeamCenterFFTTest.cpp +++ b/tests/BeamCenterFFTTest.cpp @@ -3,10 +3,14 @@ #include +#include #include +#include +#include #include #include "../image_analysis/geom_refinement/BeamCenterFFT.h" +#include "../image_analysis/geom_refinement/BeamCenterFFTCPU.h" namespace { @@ -69,3 +73,223 @@ TEST_CASE("BeamCenterFFT_SyntheticBackground", "[BeamCenter]") { CHECK(best < 12.0f); } } + +namespace { + +// A value keyed on the distance from (cx, cy) in each coordinate, so the image is EXACTLY its own +// mirror about that centre under point inversion AND under either line mirror - the three symmetries +// the three surfaces measure. The preparation the score applies - clipping, median subtraction, +// clamping, the mean shift - is pointwise and monotone, so it preserves all of them, and the masked +// Pearson at the true centre is then exactly 1. Every other centre must score below it, because the +// score is a correlation coefficient and cannot exceed 1. +float SymmetricValue(int64_t x, int64_t y, float cx, float cy) { + const int64_t dx = std::llabs(std::lround(2.0f * (static_cast(x) - cx))); + const int64_t dy = std::llabs(std::lround(2.0f * (static_cast(y) - cy))); + uint64_t k = static_cast(dx) * 0x9e3779b97f4a7c15ULL + + static_cast(dy) * 0xc2b2ae3d27d4eb4fULL; + k ^= k >> 31; + k *= 0xbf58476d1ce4e5b9ULL; + k ^= k >> 29; + return 100.0f + static_cast(k % 4096u) * 0.25f; +} + +// The symmetric pattern in a rectangle at (x0, y0), NAN everywhere else, so the same pattern can be +// placed anywhere on the detector. The true centre is (x0 + cx, y0 + cy). +std::vector SymmetricPatch(int64_t w, int64_t h, int64_t x0, int64_t y0, int64_t rw, + int64_t rh, float cx, float cy) { + std::vector mean(static_cast(w * h), NAN); + for (int64_t y = 0; y < rh; y++) + for (int64_t x = 0; x < rw; x++) + mean[static_cast((y0 + y) * w + x0 + x)] = SymmetricValue(x, y, cx, cy); + return mean; +} + +// The four convolutions straight from their definitions, for an image small enough to afford it: +// C(t) = sum_x a(x) a(t - x), S = a (*) m, Q = a^2 (*) m, D = m (*) m. +BeamCenterConvSurfaces2D BruteForcePoint(const std::vector &a, const std::vector &m, + int64_t h, int64_t w) { + BeamCenterConvSurfaces2D out; + const size_t n = static_cast(4 * h * w); + out.C.assign(n, 0.0f); + out.S.assign(n, 0.0f); + out.Q.assign(n, 0.0f); + out.D.assign(n, 0.0f); + for (int64_t ty = 0; ty < 2 * h; ty++) + for (int64_t tx = 0; tx < 2 * w; tx++) { + double c = 0, s = 0, q = 0, d = 0; + for (int64_t y = std::max(0, ty - h + 1); y <= std::min(h - 1, ty); y++) + for (int64_t x = std::max(0, tx - w + 1); x <= std::min(w - 1, tx); x++) { + const size_t i = static_cast(y * w + x); + const size_t j = static_cast((ty - y) * w + tx - x); + c += static_cast(a[i]) * a[j]; + s += static_cast(a[i]) * m[j]; + q += static_cast(a[i]) * a[i] * m[j]; + d += static_cast(m[i]) * m[j]; + } + const size_t t = static_cast(ty * 2 * w + tx); + out.C[t] = static_cast(c); + out.S[t] = static_cast(s); + out.Q[t] = static_cast(q); + out.D[t] = static_cast(d); + } + return out; +} + +// The same for the 1D mirror of the row coordinate: each column autoconvolved, summed over columns. +BeamCenterConvSurfaces1D BruteForceLine(const std::vector &a, const std::vector &m, + int64_t h, int64_t w) { + BeamCenterConvSurfaces1D out; + out.C.assign(static_cast(2 * h), 0.0); + out.S.assign(static_cast(2 * h), 0.0); + out.Q.assign(static_cast(2 * h), 0.0); + out.D.assign(static_cast(2 * h), 0.0); + for (int64_t t = 0; t < 2 * h; t++) + for (int64_t x = 0; x < w; x++) + for (int64_t y = std::max(0, t - h + 1); y <= std::min(h - 1, t); y++) { + const size_t i = static_cast(y * w + x); + const size_t j = static_cast((t - y) * w + x); + out.C[static_cast(t)] += static_cast(a[i]) * a[j]; + out.S[static_cast(t)] += static_cast(a[i]) * m[j]; + out.Q[static_cast(t)] += static_cast(a[i]) * a[i] * m[j]; + out.D[static_cast(t)] += static_cast(m[i]) * m[j]; + } + return out; +} + +// Largest relative difference of two surfaces, scaled by the largest term of the reference - a +// convolution surface spans orders of magnitude and its small entries are differences of large +// ones, so an elementwise relative test measures roundoff and nothing else. +template double MaxScaledDifference(const std::vector &got, const std::vector &ref) { + double scale = 0.0, worst = 0.0; + for (T v : ref) + scale = std::max(scale, std::abs(static_cast(v))); + for (size_t i = 0; i < ref.size(); i++) + worst = std::max(worst, std::abs(static_cast(got[i]) - static_cast(ref[i]))); + return scale > 0.0 ? worst / scale : worst; +} + +} // namespace + +// The identity the whole component rests on: the centrosymmetry score about c is the +// autoconvolution at 2c, so one transform set scores every centre. On an exactly centrosymmetric +// image the score at the true centre is exactly 1 and nowhere else can reach it - at integer and +// at half-pixel centres alike, because the surface is sampled on the half-pixel grid. +TEST_CASE("BeamCenterFFT_RecoversAnExactlySymmetricCentre", "[BeamCenter]") { + const int64_t w = 220, h = 180; + + const float cx = GENERATE(83.0f, 83.5f); + const float cy = GENERATE(61.0f, 61.5f); + + const auto mean = SymmetricPatch(w, h, 0, 0, w, h, cx, cy); + const auto r = BeamCenterFFTScore(w, h, mean); + + REQUIRE(!r.point.empty()); + CHECK(r.point[0].beam_x_pxl == cx); + CHECK(r.point[0].beam_y_pxl == cy); + CHECK(r.point[0].score == Catch::Approx(1.0f).margin(1e-4)); + // The two line mirrors see the same symmetry one coordinate at a time. + REQUIRE(!r.line_x.empty()); + REQUIRE(!r.line_y.empty()); + CHECK(r.line_x[0].beam_x_pxl == cx); + CHECK(r.line_y[0].beam_y_pxl == cy); +} + +// Move the image and the answer moves with it, exactly. The score is a convolution, so this is a +// property and not a tolerance: an off-by-one in the padding, the crop or the t = 2c indexing shows +// up here as a fixed offset, which is the class of error a synthetic image at one position cannot +// see. +TEST_CASE("BeamCenterFFT_IsShiftEquivariant", "[BeamCenter]") { + const int64_t w = 260, h = 200; + const int64_t rw = 150, rh = 120; + const float cx = 71.5f, cy = 55.0f; + + const auto at_origin = BeamCenterFFTScore(w, h, SymmetricPatch(w, h, 0, 0, rw, rh, cx, cy)); + REQUIRE(!at_origin.point.empty()); + + const int64_t tx = 37, ty = 43; + const auto shifted = BeamCenterFFTScore(w, h, SymmetricPatch(w, h, tx, ty, rw, rh, cx, cy)); + REQUIRE(!shifted.point.empty()); + + CHECK(shifted.point[0].beam_x_pxl == at_origin.point[0].beam_x_pxl + static_cast(tx)); + CHECK(shifted.point[0].beam_y_pxl == at_origin.point[0].beam_y_pxl + static_cast(ty)); + CHECK(shifted.line_x[0].beam_x_pxl == at_origin.line_x[0].beam_x_pxl + static_cast(tx)); + CHECK(shifted.line_y[0].beam_y_pxl == at_origin.line_y[0].beam_y_pxl + static_cast(ty)); +} + +// The four convolutions the score is built from, against their definitions. The masked Pearson is a +// cancellation of large near-equal terms and the shortlist margins ride on differences of ~0.3 %, +// so what matters is not that the transform is approximately right but by how much it is wrong. +TEST_CASE("BeamCenterFFT_ConvolutionsAgreeWithBruteForce", "[BeamCenter]") { + const int64_t w = 23, h = 17; + std::vector a(static_cast(w * h)), m(static_cast(w * h)); + std::mt19937 rng(20260913); + std::uniform_real_distribution uniform(-1.0f, 1.0f); + for (size_t i = 0; i < a.size(); i++) { + // A quarter of the pixels masked, in a pattern with no symmetry of its own. + m[i] = (i * 7919u) % 4u == 0u ? 0.0f : 1.0f; + a[i] = m[i] * (30.0f + 10.0f * uniform(rng)); + } + + BeamCenterFFTCPU engine; + const auto got = engine.PointSurfaces(a, m, h, w); + const auto ref = BruteForcePoint(a, m, h, w); + CHECK(MaxScaledDifference(got.C, ref.C) < 1e-6); + CHECK(MaxScaledDifference(got.S, ref.S) < 1e-6); + CHECK(MaxScaledDifference(got.Q, ref.Q) < 1e-6); + CHECK(MaxScaledDifference(got.D, ref.D) < 1e-6); + + const auto got_line = engine.LineSurfaces(a, m, h, w, BeamCenterMirror::Rows); + const auto ref_line = BruteForceLine(a, m, h, w); + CHECK(MaxScaledDifference(got_line.C, ref_line.C) < 1e-6); + CHECK(MaxScaledDifference(got_line.S, ref_line.S) < 1e-6); + CHECK(MaxScaledDifference(got_line.Q, ref_line.Q) < 1e-6); + CHECK(MaxScaledDifference(got_line.D, ref_line.D) < 1e-6); +} + +// A mirror overlap with no variance in it - a region the preparation has flattened to a constant - +// makes the score 0/0, and in floating point that mints an r of several hundred out of pure +// roundoff. The score is a correlation coefficient: no candidate may ever come back above 1. +TEST_CASE("BeamCenterFFT_ANoVarianceOverlapScoresNothing", "[BeamCenter]") { + const int64_t w = 400, h = 320; + const int64_t pw = 240, ph = 200; + const float cx = 0.5f * static_cast(pw - 1), cy = 0.5f * static_cast(ph - 1); + + // A flat detector-wide background - which the preparation clamps to a constant - with the + // textured, exactly symmetric patch filling one corner of it. Every centre whose overlap lies + // inside the flat part carries no evidence at all. + std::vector mean(static_cast(w * h), 100.0f); + for (int64_t y = 0; y < ph; y++) + for (int64_t x = 0; x < pw; x++) + mean[static_cast(y * w + x)] = SymmetricValue(x, y, cx, cy); + + const auto r = BeamCenterFFTScore(w, h, mean); + REQUIRE(!r.point.empty()); + for (const auto &c : r.point) + CHECK(c.score <= 1.0f + 1e-3f); + CHECK(r.point[0].beam_x_pxl == cx); + CHECK(r.point[0].beam_y_pxl == cy); +} + +// The transform length. A smooth length is what makes the whole-detector transform affordable: the +// naive next power of two for a 16 Mpixel detector costs 3x the time and, on cuFFT, 2.5 GB of plan +// memory against a few hundred MB. +TEST_CASE("BeamCenterFFT_PadsToASmoothLength", "[BeamCenter]") { + CHECK(BeamCenterFFTPadSize(1) == 1); + CHECK(BeamCenterFFTPadSize(1024) == 1024); + CHECK(BeamCenterFFTPadSize(1025) == 1029); // 3 * 7^3 + CHECK(BeamCenterFFTPadSize(2027) == 2048); + CHECK(BeamCenterFFTPadSize(4931) == 5000); // 2^3 * 5^4 + + for (int64_t n = 2000; n < 20000; n += 7) { + const int64_t pad = BeamCenterFFTPadSize(n); + REQUIRE(pad >= n); + int64_t v = pad; + for (int p : {2, 3, 5, 7}) + while (v % p == 0) + v /= p; + REQUIRE(v == 1); + // And it is always close: the padding never costs a whole extra octave the way a + // power-of-two length can. + REQUIRE(pad < n + n / 20); + } +} diff --git a/tests/BeamCenterFromBackgroundTest.cpp b/tests/BeamCenterFromBackgroundTest.cpp index 55129a853..73a547e92 100644 --- a/tests/BeamCenterFromBackgroundTest.cpp +++ b/tests/BeamCenterFromBackgroundTest.cpp @@ -4,7 +4,9 @@ #include #include +#include #include +#include #include "../image_analysis/geom_refinement/BeamCenterFromBackground.h" #include "../common/DetectorSetup.h" @@ -182,3 +184,126 @@ TEST_CASE("BeamCenterFromBackground_DoesNotDependOnTheThreadCount", "[BeamCenter CHECK(one->beam_y_pxl == many->beam_y_pxl); CHECK(one->sigma_pxl == many->sigma_pxl); } + +// FindBeamCenter is the walk with an FFT capture in front of it: the capture chooses the basin and +// the walk finishes inside it and reports what it knows the answer to. The tests below are about +// the composition - that the walk really is seeded where the capture landed, that a decline falls +// back the way it is documented to, and that the beam stop is taken out of the image the capture +// scores. + +namespace { + +// A projection that is NAN everywhere except one exactly centrosymmetric patch. The capture has an +// unambiguous answer on it (a correlation coefficient cannot exceed 1, and the patch reaches 1 at +// its own centre) while the ring fit has nothing at all: no ring of the 12-2.2 A band is covered +// by a patch that small, from any start. +std::vector SymmetricPatchOnly(const DiffractionExperiment &experiment, float cx, float cy, + int half_size) { + const auto W = static_cast(experiment.GetXPixelsNumConv()); + const auto H = static_cast(experiment.GetYPixelsNumConv()); + std::vector mean(static_cast(W) * H, NAN); + for (int y = static_cast(cy) - half_size; y <= static_cast(cy) + half_size; y++) + for (int x = static_cast(cx) - half_size; x <= static_cast(cx) + half_size; x++) { + // Keyed on the canonical member of the +- pair about (cx, cy), so the patch is exactly + // its own mirror there. + int64_t dx = std::lround(2.0f * (static_cast(x) - cx)); + int64_t dy = std::lround(2.0f * (static_cast(y) - cy)); + if (dy < 0 || (dy == 0 && dx < 0)) { + dx = -dx; + dy = -dy; + } + uint64_t k = static_cast(dx) * 0x9e3779b97f4a7c15ULL + + static_cast(dy) * 0xc2b2ae3d27d4eb4fULL; + k ^= k >> 31; + k *= 0xbf58476d1ce4e5b9ULL; + k ^= k >> 29; + mean[static_cast(y) * W + x] = 100.0f + static_cast(k % 4096u) * 0.25f; + } + return mean; +} + +} // namespace + +// The case the composition exists for. With the centre in the file destroyed - the value a +// beamline writes when it has nothing to write - the walk has no ring of the fitted band covered +// from where it starts, so it declines and says nothing at all. The capture does not care how +// wrong the file is: it scores every centre on the detector in one transform set, and the walk +// then finishes from there with a real fit sigma, low enough that a caller would commit it. +TEST_CASE("FindBeamCenter_AnswersWhereTheWalkDeclines", "[BeamCenter]") { + DiffractionExperiment x = TestExperiment(); + PixelMask pixel_mask(x); + + const DiffractionGeometry geom_true = x.GetDiffractionGeometry(); + const auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 1.0f); + x.BeamX_pxl(0.0f).BeamY_pxl(0.0f); // the destroyed header + + CHECK(!FindBeamCenterFromBackground(x, pixel_mask, projection).has_value()); + + const auto composed = FindBeamCenter(x, pixel_mask, projection); + REQUIRE(composed.has_value()); + CHECK(composed->beam_x_pxl == Catch::Approx(geom_true.GetBeamX_pxl()).margin(1.0)); + CHECK(composed->beam_y_pxl == Catch::Approx(geom_true.GetBeamY_pxl()).margin(1.0)); + CHECK(composed->sigma_pxl < 1.0f); +} + +// The end of the fallback chain. Where the walk has nothing to fit from either start - neither from +// the capture nor from the centre in the file - the capture stands alone, and it is disclosed as a +// capture: BEAM_CENTER_CAPTURE_SIGMA_PXL sits above every ceiling a caller commits geometry on, so +// the answer reaches the consumers that only need a hypothesis to test and is refused by the two +// that would change the geometry with it. +TEST_CASE("FindBeamCenter_FallsBackToTheCaptureAlone", "[BeamCenter]") { + DiffractionExperiment x = TestExperiment(); + PixelMask pixel_mask(x); + + const float cx = 700.5f, cy = 640.0f; + const auto projection = SymmetricPatchOnly(x, cx, cy, 90); + + CHECK(!FindBeamCenterFromBackground(x, pixel_mask, projection).has_value()); + + const auto composed = FindBeamCenter(x, pixel_mask, projection); + REQUIRE(composed.has_value()); + CHECK(composed->beam_x_pxl == cx); + CHECK(composed->beam_y_pxl == cy); + CHECK(composed->sigma_pxl == BEAM_CENTER_CAPTURE_SIGMA_PXL); +} + +// The beam stop is blanked out of the image the capture scores. A large one-sided opaque region is +// centrosymmetric about ITS own centre, and where it is big enough that preference beats the +// background's - measured on a real umbra over 9.6 % of the detector, the capture landed 48 px out +// and masking it put it back to 1.1 px. The mask that reaches the capture already has the stop in +// it, so this costs nothing but has to keep working. +TEST_CASE("FindBeamCenter_BlanksTheBeamStopOutOfTheCapture", "[BeamCenter]") { + DiffractionExperiment x = TestExperiment(); + PixelMask pixel_mask(x); + + const DiffractionGeometry geom_true = x.GetDiffractionGeometry(); + auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 1.0f); + + // An umbra on one side of the beam: opaque, so the pixels under it carry no background. + const auto W = static_cast(x.GetXPixelsNumConv()); + const auto H = static_cast(x.GetYPixelsNumConv()); + const float sx = geom_true.GetBeamX_pxl() + 330.0f, sy = geom_true.GetBeamY_pxl() + 70.0f; + std::vector stop(static_cast(W) * H, 0); + for (int y = 0; y < H; y++) + for (int x_pxl = 0; x_pxl < W; x_pxl++) + if (std::hypot(x_pxl - sx, y - sy) < 300.0f) { + stop[static_cast(y) * W + x_pxl] = 1; + projection[static_cast(y) * W + x_pxl] = 0.0f; + } + + const auto unmasked = BeamCenterFFTScore(W, H, projection); + REQUIRE(!unmasked.point.empty()); + + pixel_mask.LoadBeamStopMask(x, stop); + BeamCenterFFTResult capture; + const auto composed = FindBeamCenter(x, pixel_mask, projection, 0, &capture); + REQUIRE(composed.has_value()); + REQUIRE(!capture.point.empty()); + + const float masked_error = std::hypot(capture.point[0].beam_x_pxl - geom_true.GetBeamX_pxl(), + capture.point[0].beam_y_pxl - geom_true.GetBeamY_pxl()); + const float unmasked_error = std::hypot(unmasked.point[0].beam_x_pxl - geom_true.GetBeamX_pxl(), + unmasked.point[0].beam_y_pxl - geom_true.GetBeamY_pxl()); + CHECK(masked_error < 12.0f); + CHECK(masked_error <= unmasked_error); +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index aa5c87021..6d748a52b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -109,6 +109,7 @@ ADD_EXECUTABLE(jfjoch_test LatticeReductionTest.cpp BeamCenterFromBackgroundTest.cpp BeamCenterFFTTest.cpp + BeamCenterFFTGPUTest.cpp BeamCenterFromSpotsTest.cpp )