mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 81 Wiki Activity

Pre-scan on the GPU: background beam-centre walk and beam-stop mask

Browse Source 
The two pre-scan steps that were still CPU-bound in a GPU build now run where the
projection already is.

- FindBeamCenterFromBackground: the per-iteration binning pass and the two clip rounds
  run on the device (BeamCenterBackgroundGPU); the fit itself stays on the host. Each
  cell is summed in the host's order (pixel order within the host's row blocks, blocks
  in order), and the per-pixel cell/derivative formula is shared (BackgroundBand.h).
  The angles come from BackgroundAtan2 (IEEE ops only) instead of atan2f, and both
  translation units are compiled without FMA contraction, so host and device give the
  same bits: 0 of 6.5 M pixels in a different cell, identical walks on the three
  in-house rotation sets. With glibc/CUDA atan2f and default contraction ~30 pixels per
  16 Mpx sweep changed cell and the fitted centre moved by up to 0.05 px.
- ShadowFinder::GetMask: the whole mask (pooling, ring medians, components, morphology,
  hole fill, arm search) runs on the device from ShadowAccumulatorGPU's projection
  (ShadowMaskGPU), so the 360 MB projection no longer comes back; the mean projection is
  divided on the device too (same bits). The two small fits over rings and sectors
  (BlockedOutTo, HarmonicFit) are shared with the host path in ShadowFinderInternal.h.
  Integers, comparisons, sorts and components are exact; the polarization trig, the
  Poisson log and the arm-search azimuth are not, so a pixel at a threshold can differ.

The one-time change against the previous CPU arithmetic (BackgroundAtan2, no
contraction), measured on the myoglobin, cytochrome C and thaumatin rotation sets:
ring centre moves 0.002-0.045 px (fit sigma 0.75-1.2 px), beam-centre capture
0.01-0.04 px; beam-stop mask differs on 31 / 144 / 53 pixels of 259k / 144k / 198k
(25 of the myoglobin ones are GPU-vs-CPU arithmetic in the mask, the rest follow the
centre); hot-pixel mask identical. Spot width, integration radii, bandwidth, beam-centre
arbitration, indexing, space group, cell, resolution and the merged statistics table
are identical; only the error model moves in its 4th digit. CPU build: the same
centres and decisions.

Timing (GPU, box at load 30-38): ring walk 0.54 -> 0.23-0.27 s, mask 1.24-1.44 ->
0.18-0.22 s, beam-centre capture walk 1.1-1.3 -> 0.31-0.35 s.

Tests: ShadowFinder_DeviceMaskMatchesHost, BeamCenterFromBackground_DeviceMatchesHost
(bit-exact), plus [ShadowFinder], [BeamCenter], [HotPixelFinder].

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01K5K8jvPPbmCrbqnWkddTuB
This commit is contained in:
leonarski_f
2026-10-03 12:19:05 +02:00
co-authored by Claude Opus 5.5
parent 24e36ae740
commit f849e2d1be
15 changed files with 1468 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/CMakeLists.txt
+3
View File
@@ -99,7 +99,10 @@ ADD_LIBRARY(JFJochImageAnalysis STATIC
beam_stop/ShadowFinder.cpp
beam_stop/ShadowFinder.h
$<$<BOOL:${JFJOCH_CUDA_AVAILABLE}>:beam_stop/ShadowAccumulatorGPU.cu>
$<$<BOOL:${JFJOCH_CUDA_AVAILABLE}>:beam_stop/ShadowMaskGPU.cu>
beam_stop/ShadowAccumulatorGPU.h
beam_stop/ShadowFinderInternal.h
beam_stop/ShadowMaskGPU.h
rotation_indexer/RotationIndexer.cpp
rotation_indexer/RotationIndexer.h
WriteReflections.cpp
image_analysis/beam_stop/ShadowAccumulatorGPU.cu
+10
View File
@@ -190,3 +190,13 @@ void ShadowAccumulatorGPU::Download(std::vector<int64_t> &max_value, std::vector
cudaMemcpyDeviceToHost, *stream));
cuda_err(cudaStreamSynchronize(*stream));
}
std::vector<uint32_t> ShadowAccumulatorGPU::Mask(const ShadowMaskSetup &setup, const std::vector<uint32_t> &pixel_mask) {
FoldPending();
return ShadowMaskOnDevice(setup, pixel_mask, gpu_max, gpu_sum, gpu_count, frames, *stream);
}
std::vector<float> ShadowAccumulatorGPU::MeanProjection(const std::vector<uint32_t> &pixel_mask) {
FoldPending();
return MeanProjectionOnDevice(pixel_mask, gpu_sum, gpu_count, npixels, *stream);
}
image_analysis/beam_stop/ShadowAccumulatorGPU.h
+6
View File
@@ -10,6 +10,7 @@
#include "../../common/CompressedImage.h"
#include "../image_preprocessing/BSLZ4DecoderGPU.h"
#include "../indexing/CUDAMemHelpers.h"
#include "ShadowMaskGPU.h"
// The beam-stop projection accumulated on the device: only the compressed chunk crosses PCIe, and
// both the decode and the per-pixel maximum / sum / count run on the GPU. The projection comes back
@@ -59,6 +60,11 @@ public:
[[nodiscard]] uint32_t GetFrameCount() const { return frames; }
// The beam-stop mask and the mean projection, made where the projection is (ShadowMaskGPU.h), so
// that it does not have to come back at all.
std::vector<uint32_t> Mask(const ShadowMaskSetup &setup, const std::vector<uint32_t> &pixel_mask);
std::vector<float> MeanProjection(const std::vector<uint32_t> &pixel_mask);
// Bring the projection back to the host, folding in whatever the last batch still holds. Cheap
// to call once; it moves 20 bytes per pixel.
void Download(std::vector<int64_t> &max_value, std::vector<int64_t> &sum_value,
image_analysis/beam_stop/ShadowFinder.cpp
+123 -119
View File
@@ -2,6 +2,7 @@
// SPDX-License-Identifier: GPL-3.0-only
#include "ShadowFinder.h"
#include "ShadowFinderInternal.h"
#include <algorithm>
#include <cmath>
@@ -18,64 +19,9 @@
#include "../../common/ParallelFor.h"
#include "../../common/JFJochException.h"
// A pixel is shadow when its background is below this fraction of the background it is
// compared against.
constexpr float SHADOW_RATIO = 0.50f;
using namespace shadow_finder;
// The boundary grows outward into partially shadowed pixels down to this fraction, but no
// further than PENUMBRA_MAX_PX from the core. A pin or a loop casts a wide half-shadow, so the
// reach is a good deal more than the beam stop's own edge needs.
constexpr float PENUMBRA_RATIO = 0.75f;
constexpr int PENUMBRA_MAX_PX = 30;
// Bridge small breaks along the holder arm. A module gap wider than this is bridged separately for
// the arm search (see bridge_gaps).
constexpr int BRIDGE_PX = 6;
// A pixel whose maximum reaches this recorded a real reflection and is never masked - a
// beam stop cannot block a reflection that was measured.
constexpr int64_t MIN_REFLECTION = 25;
// How far below the background it is compared against a pixel must sit before the dip is
// believed, in standard deviations of the counts that back it. The counts are photons, so their
// scatter is Poisson and the deficit is measured against it rather than against a fixed number:
// on a well-exposed sweep a third of the background missing is overwhelming, and on a handful of
// low-background frames the same third is noise. Without this a six-frame pre-scan of a
// low-background sweep masks three quarters of the detector.
constexpr double MIN_DEFICIT_SIGMA = 6.0;
// Smallest region the per-pixel test may return. A shadow is cast by something physical and is
// correspondingly large; an isolated patch this small is the background wandering, not hardware.
// This is what keeps the test specific now that a shadow no longer has to touch the direct beam.
constexpr int MIN_SHADOW_PIXELS = 2000;
// ... and of those, how many must be deep (below SHADOW_RATIO) for a region of merely DIM pixels -
// hardware that lets part of the beam through - to count. A thin holder arm is dim along most of
// its length and deep in places; the background drifting over a detector's edge is dim everywhere
// and deep nowhere.
constexpr int MIN_CORE_PIXELS = 200;
// The arm search compares a pixel with its ring as the ring actually varies around the beam. A ring
// of background is not flat once divided by the polarization factor when that factor is not the
// beam's: the remainder is a second harmonic in azimuth, cos 2phi, which reaches tens of percent at
// high angle. It is measured over radial bands of HARMONIC_BAND_PX, from the median of each of
// HARMONIC_SECTORS sectors that holds at least MIN_SECTOR_PIXELS pixels.
constexpr int HARMONIC_BAND_PX = 64;
constexpr int HARMONIC_SECTORS = 24;
constexpr int MIN_SECTOR_PIXELS = 200;
// Side of the box the background is pooled over before testing. Its area is how many pixels back
// a ring's countability test, which decides where an azimuthal comparison is possible at all.
constexpr int POOL_PX = 5;
constexpr double MEAN_POOLED_PIXELS = POOL_PX * POOL_PX;
// A ring with fewer valid pixels than this says nothing about whether it was counted.
constexpr int MIN_RING_PIXELS = 32;
// A ring lies wholly inside the stop when its background is below this fraction of the background
// further out. This asks about a whole ring rather than about a pixel, so it keeps a threshold of
// its own and does not follow SHADOW_RATIO.
constexpr float BLOCKED_RING_RATIO = 0.35f;
static_assert(ShadowFinder::SHADOW == MASK_SHADOW && ShadowFinder::TRANSMITTING == MASK_TRANSMITTING);
// Binary-image helpers on a width*height frame stored row-major as char (0/1). All run once,
// at GetMask() time, and all are O(pixels) rather than O(pixels * radius).
@@ -635,6 +581,10 @@ void ShadowFinder::ReleaseProjection() {
std::vector<float> ShadowFinder::GetMeanProjection() const {
std::unique_lock ul(m);
#ifdef JFJOCH_USE_CUDA
if (Gpu() && gpu->GetFrameCount() > 0 && host.frames == 0)
return gpu->MeanProjection(pixel_mask);
#endif
const Projection &p = Reduced();
const auto &sum_value = p.sum_value;
const auto &valid_count = p.valid_count;
@@ -652,6 +602,28 @@ std::vector<uint32_t> ShadowFinder::GetMask(size_t nthreads) const {
std::unique_lock ul(m);
if (nthreads == 0)
nthreads = std::max(1u, std::thread::hardware_concurrency());
#ifdef JFJOCH_USE_CUDA
// Where every frame went to the device the mask is made there too, from the projection as it lies.
if (Gpu() && gpu->GetFrameCount() > 0 && host.frames == 0) {
const float diag = std::hypot(static_cast<float>(width), static_cast<float>(height));
if (!std::isfinite(beam_x) || !std::isfinite(beam_y)
|| std::fabs(beam_x - width * 0.5f) > 4.0f * diag || std::fabs(beam_y - height * 0.5f) > 4.0f * diag)
return std::vector<uint32_t>(static_cast<size_t>(width) * height, 0);
ShadowMaskSetup setup;
setup.width = width;
setup.height = height;
setup.beam_x = beam_x;
setup.beam_y = beam_y;
const auto rot = geometry.GetDetectorMatrix().arr();
for (int k = 0; k < 9; k++)
setup.det_matrix[k] = rot[k];
setup.pixel_size_mm = geometry.GetPixelSize_mm();
setup.distance_mm = geometry.GetDetectorDistance_mm();
setup.has_polarization = polarization.has_value();
setup.polarization = polarization.value_or(0.0f);
return gpu->Mask(setup, pixel_mask);
}
#endif
const Projection &p = Reduced();
const auto &max_value = p.max_value;
const auto &sum_value = p.sum_value;
@@ -804,25 +776,7 @@ std::vector<uint32_t> ShadowFinder::GetMask(size_t nthreads) const {
// disk it masked. The median over the rings this walk is willing to judge is what "typically"
// means, and it is robust from both sides - to a bright ring, and to a corner ring of a handful
// of pixels whose median is one pixel's mean.
std::vector<float> judgeable;
for (int rad = 0; rad <= max_radius; rad++)
if (ring_pixels[rad] >= MIN_RING_PIXELS)
judgeable.push_back(baseline[rad]);
float typical_background = 0.0f;
if (!judgeable.empty()) {
const auto middle = judgeable.begin() + judgeable.size() / 2;
std::nth_element(judgeable.begin(), middle, judgeable.end());
typical_background = *middle;
}
int blocked_out_to = -1;
for (int rad = 0; rad <= max_radius; rad++) {
if (ring_pixels[rad] < MIN_RING_PIXELS)
continue;
if (baseline[rad] >= BLOCKED_RING_RATIO * typical_background)
break;
blocked_out_to = rad;
}
const int blocked_out_to = BlockedOutTo(baseline, ring_pixels);
// The counts a pixel's pooled background is made of, and the counts the ring says it should
// have had. The test is on the deficit between them, in units of its own Poisson scatter.
@@ -975,51 +929,17 @@ std::vector<uint32_t> ShadowFinder::GetMask(size_t nthreads) const {
sector_values[k].insert(sector_values[k].end(), values[k].begin(), values[k].end());
});
block_values.clear();
std::vector<float> harm_c(n_bands, 0.0f), harm_s(n_bands, 0.0f);
ParallelFor(n_bands, nthreads, [&](int band) {
std::vector<double> med(HARMONIC_SECTORS, -1.0), c(HARMONIC_SECTORS), s(HARMONIC_SECTORS);
for (int k = 0; k < HARMONIC_SECTORS; k++) {
auto &v = sector_values[static_cast<size_t>(band) * HARMONIC_SECTORS + k];
if (v.size() >= MIN_SECTOR_PIXELS) {
std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end());
med[k] = v[v.size() / 2];
}
const double phi = (k + 0.5) * 2.0 * std::numbers::pi / HARMONIC_SECTORS - std::numbers::pi;
c[k] = std::cos(2 * phi);
s[k] = std::sin(2 * phi);
}
// Least squares of med = m + p cos 2phi + q sin 2phi over the sectors that are not themselves
// dim against the fit, three times over as the baseline is; the model is relative, 1 + (p cos + q sin)/m.
std::vector<char> use(HARMONIC_SECTORS);
for (int k = 0; k < HARMONIC_SECTORS; k++)
use[k] = med[k] >= 0;
for (int iter = 0; iter < 3; iter++) {
double n = 0, sc = 0, ss = 0, scc = 0, sss = 0, scs = 0, y = 0, yc = 0, ys = 0;
for (int k = 0; k < HARMONIC_SECTORS; k++) {
if (!use[k]) continue;
n++; sc += c[k]; ss += s[k]; scc += c[k] * c[k]; sss += s[k] * s[k]; scs += c[k] * s[k];
y += med[k]; yc += med[k] * c[k]; ys += med[k] * s[k];
}
// Sectors crowded into a narrow arc cannot tell a harmonic from a level; the determinant of
// the normal matrix, per sector cubed, is 1/4 on a full ring.
const double det = n * (scc * sss - scs * scs) - sc * (sc * sss - scs * ss) + ss * (sc * scs - scc * ss);
if (n < 6 || det < 0.01 * n * n * n) {
harm_c[band] = harm_s[band] = 0.0f;
return;
}
const double m = (y * (scc * sss - scs * scs) - sc * (yc * sss - scs * ys) + ss * (yc * scs - scc * ys)) / det;
const double p = (n * (yc * sss - ys * scs) - y * (sc * sss - scs * ss) + ss * (sc * ys - yc * ss)) / det;
const double q = (n * (scc * ys - scs * yc) - sc * (sc * ys - yc * ss) + y * (sc * scs - scc * ss)) / det;
if (m <= 0) {
harm_c[band] = harm_s[band] = 0.0f;
return;
}
harm_c[band] = static_cast<float>(p / m);
harm_s[band] = static_cast<float>(q / m);
for (int k = 0; k < HARMONIC_SECTORS; k++)
use[k] = med[k] >= 0 && med[k] >= PENUMBRA_RATIO * (1.0 + harm_c[band] * c[k] + harm_s[band] * s[k]);
// The median of each sector with enough pixels to have one.
std::vector<double> sector_median(n_sectors, -1.0);
ParallelFor(static_cast<int>(n_sectors), nthreads, [&](int k) {
auto &v = sector_values[k];
if (v.size() >= MIN_SECTOR_PIXELS) {
std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end());
sector_median[k] = v[v.size() / 2];
}
});
std::vector<float> harm_c, harm_s;
HarmonicFit(sector_median, n_bands, harm_c, harm_s);
Plane<char> dim = filled_plane<char>(n_pixels, 0, nthreads);
ParallelChunks(n_pixels, nthreads, [&](int lo, int hi) {
@@ -1062,3 +982,87 @@ std::vector<uint32_t> ShadowFinder::GetMask(size_t nthreads) const {
});
return mask;
}
namespace shadow_finder {
int BlockedOutTo(const std::vector<float> &baseline, const std::vector<int> &ring_pixels) {
const int max_radius = static_cast<int>(baseline.size()) - 1;
// A ring lies inside the stop when its background is a fraction of what this detector's
// background typically is. Counting statistics cannot decide this: on a bright dataset the
// shadow is still well counted. The comparison used to be against the LARGEST background of any
// ring further out, and that reads a sample whose background peaks in a strong ring away from
// the beam - a powder standard, a strong solvent ring - as a beam stop the size of that ring:
// the ordinary background inside it is legitimately below a third of the peak. On one corpus
// dataset it declared 16 % of the detector to be stop, with diffraction rings visible inside the
// disk it masked. The median over the rings this walk is willing to judge is what "typically"
// means, and it is robust from both sides - to a bright ring, and to a corner ring of a handful
// of pixels whose median is one pixel's mean.
std::vector<float> judgeable;
for (int rad = 0; rad <= max_radius; rad++)
if (ring_pixels[rad] >= MIN_RING_PIXELS)
judgeable.push_back(baseline[rad]);
float typical_background = 0.0f;
if (!judgeable.empty()) {
const auto middle = judgeable.begin() + judgeable.size() / 2;
std::nth_element(judgeable.begin(), middle, judgeable.end());
typical_background = *middle;
}
int blocked_out_to = -1;
for (int rad = 0; rad <= max_radius; rad++) {
if (ring_pixels[rad] < MIN_RING_PIXELS)
continue;
if (baseline[rad] >= BLOCKED_RING_RATIO * typical_background)
break;
blocked_out_to = rad;
}
return blocked_out_to;
}
void HarmonicFit(const std::vector<double> &sector_median, int n_bands,
std::vector<float> &harm_c, std::vector<float> &harm_s) {
harm_c.assign(n_bands, 0.0f);
harm_s.assign(n_bands, 0.0f);
for (int band = 0; band < n_bands; band++) {
std::vector<double> med(HARMONIC_SECTORS), c(HARMONIC_SECTORS), s(HARMONIC_SECTORS);
for (int k = 0; k < HARMONIC_SECTORS; k++) {
med[k] = sector_median[static_cast<size_t>(band) * HARMONIC_SECTORS + k];
const double phi = (k + 0.5) * 2.0 * std::numbers::pi / HARMONIC_SECTORS - std::numbers::pi;
c[k] = std::cos(2 * phi);
s[k] = std::sin(2 * phi);
}
// Least squares of med = m + p cos 2phi + q sin 2phi over the sectors that are not themselves
// dim against the fit, three times over as the baseline is; the model is relative, 1 + (p cos + q sin)/m.
std::vector<char> use(HARMONIC_SECTORS);
for (int k = 0; k < HARMONIC_SECTORS; k++)
use[k] = med[k] >= 0;
for (int iter = 0; iter < 3; iter++) {
double n = 0, sc = 0, ss = 0, scc = 0, sss = 0, scs = 0, y = 0, yc = 0, ys = 0;
for (int k = 0; k < HARMONIC_SECTORS; k++) {
if (!use[k]) continue;
n++; sc += c[k]; ss += s[k]; scc += c[k] * c[k]; sss += s[k] * s[k]; scs += c[k] * s[k];
y += med[k]; yc += med[k] * c[k]; ys += med[k] * s[k];
}
// Sectors crowded into a narrow arc cannot tell a harmonic from a level; the determinant of
// the normal matrix, per sector cubed, is 1/4 on a full ring.
const double det = n * (scc * sss - scs * scs) - sc * (sc * sss - scs * ss) + ss * (sc * scs - scc * ss);
if (n < 6 || det < 0.01 * n * n * n) {
harm_c[band] = harm_s[band] = 0.0f;
break;
}
const double m = (y * (scc * sss - scs * scs) - sc * (yc * sss - scs * ys) + ss * (yc * scs - scc * ys)) / det;
const double p = (n * (yc * sss - ys * scs) - y * (sc * sss - scs * ss) + ss * (sc * ys - yc * ss)) / det;
const double q = (n * (scc * ys - scs * yc) - sc * (sc * ys - yc * ss) + y * (sc * scs - scc * ss)) / det;
if (m <= 0) {
harm_c[band] = harm_s[band] = 0.0f;
break;
}
harm_c[band] = static_cast<float>(p / m);
harm_s[band] = static_cast<float>(q / m);
for (int k = 0; k < HARMONIC_SECTORS; k++)
use[k] = med[k] >= 0 && med[k] >= PENUMBRA_RATIO * (1.0 + harm_c[band] * c[k] + harm_s[band] * s[k]);
}
}
}
} // namespace shadow_finder
image_analysis/beam_stop/ShadowFinderInternal.h
+90
View File
@@ -0,0 +1,90 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
// What ShadowFinder::GetMask shares with its device twin (ShadowMaskGPU): the constants it tests
// against, and the two small fits made over whole rings and sectors, which stay on the host on both
// paths.
#include <cstdint>
#include <vector>
namespace shadow_finder {
// The values of ShadowFinder::SHADOW and ShadowFinder::TRANSMITTING, for the device code, which does
// not include ShadowFinder.h.
inline constexpr uint32_t MASK_SHADOW = 1;
inline constexpr uint32_t MASK_TRANSMITTING = 2;
// A pixel is shadow when its background is below this fraction of the background it is
// compared against.
inline constexpr float SHADOW_RATIO = 0.50f;
// The boundary grows outward into partially shadowed pixels down to this fraction, but no
// further than PENUMBRA_MAX_PX from the core. A pin or a loop casts a wide half-shadow, so the
// reach is a good deal more than the beam stop's own edge needs.
inline constexpr float PENUMBRA_RATIO = 0.75f;
inline constexpr int PENUMBRA_MAX_PX = 30;
// Bridge small breaks along the holder arm. A module gap wider than this is bridged separately for
// the arm search (see bridge_gaps).
inline constexpr int BRIDGE_PX = 6;
// A pixel whose maximum reaches this recorded a real reflection and is never masked - a
// beam stop cannot block a reflection that was measured.
inline constexpr int64_t MIN_REFLECTION = 25;
// How far below the background it is compared against a pixel must sit before the dip is
// believed, in standard deviations of the counts that back it. The counts are photons, so their
// scatter is Poisson and the deficit is measured against it rather than against a fixed number:
// on a well-exposed sweep a third of the background missing is overwhelming, and on a handful of
// low-background frames the same third is noise. Without this a six-frame pre-scan of a
// low-background sweep masks three quarters of the detector.
inline constexpr double MIN_DEFICIT_SIGMA = 6.0;
// Smallest region the per-pixel test may return. A shadow is cast by something physical and is
// correspondingly large; an isolated patch this small is the background wandering, not hardware.
// This is what keeps the test specific now that a shadow no longer has to touch the direct beam.
inline constexpr int MIN_SHADOW_PIXELS = 2000;
// ... and of those, how many must be deep (below SHADOW_RATIO) for a region of merely DIM pixels -
// hardware that lets part of the beam through - to count. A thin holder arm is dim along most of
// its length and deep in places; the background drifting over a detector's edge is dim everywhere
// and deep nowhere.
inline constexpr int MIN_CORE_PIXELS = 200;
// The arm search compares a pixel with its ring as the ring actually varies around the beam. A ring
// of background is not flat once divided by the polarization factor when that factor is not the
// beam's: the remainder is a second harmonic in azimuth, cos 2phi, which reaches tens of percent at
// high angle. It is measured over radial bands of HARMONIC_BAND_PX, from the median of each of
// HARMONIC_SECTORS sectors that holds at least MIN_SECTOR_PIXELS pixels.
inline constexpr int HARMONIC_BAND_PX = 64;
inline constexpr int HARMONIC_SECTORS = 24;
inline constexpr int MIN_SECTOR_PIXELS = 200;
// Side of the box the background is pooled over before testing. Its area is how many pixels back
// a ring's countability test, which decides where an azimuthal comparison is possible at all.
inline constexpr int POOL_PX = 5;
inline constexpr double MEAN_POOLED_PIXELS = POOL_PX * POOL_PX;
// A ring with fewer valid pixels than this says nothing about whether it was counted.
inline constexpr int MIN_RING_PIXELS = 32;
// A ring lies wholly inside the stop when its background is below this fraction of the background
// further out. This asks about a whole ring rather than about a pixel, so it keeps a threshold of
// its own and does not follow SHADOW_RATIO.
inline constexpr float BLOCKED_RING_RATIO = 0.35f;
// The rings that lie wholly inside the stop: walking outward, every judgeable ring (at least
// MIN_RING_PIXELS pixels) before the first whose baseline reaches BLOCKED_RING_RATIO of the typical
// background. -1 when there is none. See GetMask.
int BlockedOutTo(const std::vector<float> &baseline, const std::vector<int> &ring_pixels);
// The second harmonic in azimuth of each radial band, relative to its level, fitted to the medians of
// its HARMONIC_SECTORS sectors (sector_median[band * HARMONIC_SECTORS + k], negative where the sector
// has too few pixels). See GetMask.
void HarmonicFit(const std::vector<double> &sector_median, int n_bands,
std::vector<float> &harm_c, std::vector<float> &harm_s);
} // namespace shadow_finder
image_analysis/beam_stop/ShadowMaskGPU.cu
+636
View File
@@ -0,0 +1,636 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "ShadowMaskGPU.h"
#include <cub/device/device_radix_sort.cuh>
#include "ShadowFinderInternal.h"
#include "../../common/JFJochMath.h"
#include "../indexing/CUDAMemHelpers.h"
#include "../../common/JFJochException.h"
using namespace shadow_finder;
namespace {
constexpr int THREADS = 256;
void check(cudaError_t err, const char *what) {
if (err != cudaSuccess)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
std::string("Beam stop mask: ") + what + ": " + cudaGetErrorString(err));
}
unsigned grid(size_t n) {
return static_cast<unsigned>((n + THREADS - 1) / THREADS);
}
// A float as an unsigned integer that sorts the same way, and back.
__device__ uint32_t float_key(float f) {
const uint32_t u = __float_as_uint(f);
return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
}
__device__ float key_float(uint32_t k) {
return __uint_as_float((k & 0x80000000u) ? (k & 0x7fffffffu) : ~k);
}
// The first index in a sorted key array whose key is not below `value`.
__device__ size_t lower_bound(const uint64_t *keys, size_t n, uint64_t value) {
size_t lo = 0, hi = n;
while (lo < hi) {
const size_t mid = (lo + hi) / 2;
if (keys[mid] < value) lo = mid + 1;
else hi = mid;
}
return lo;
}
__device__ double poisson_deficit_sigma(double observed, double expected) {
if (expected <= 0.0 || observed >= expected)
return 0.0;
const double ll = 2.0 * (expected - observed + (observed > 0.0 ? observed * log(observed / expected) : 0.0));
return ll > 0.0 ? sqrt(ll) : 0.0;
}
// DiffractionGeometry::CalcAzIntPolarizationCorr about the centre the rings are drawn about.
__device__ float polarization_factor(const ShadowMaskSetup &s, float x, float y) {
const float u = (x - s.beam_x) * s.pixel_size_mm;
const float v = (y - s.beam_y) * s.pixel_size_mm;
const float *m = s.det_matrix;
const float lx = m[0] * u + m[1] * v + m[2] * s.distance_mm;
const float ly = m[3] * u + m[4] * v + m[5] * s.distance_mm;
const float lz = m[6] * u + m[7] * v + m[8] * s.distance_mm;
const float two_theta = atan2f(sqrtf(lx * lx + ly * ly), lz);
float phi = atan2f(ly, lx);
if (phi < 0)
phi += 2.0f * PI;
const float cos_2theta = cosf(two_theta);
const float cos_2theta_2 = cos_2theta * cos_2theta;
const float cos_2phi = cosf(2.0f * phi);
return 0.5f * (1.0f + cos_2theta_2 - s.polarization * cos_2phi * (1.0f - cos_2theta_2));
}
__global__ void setup_kernel(ShadowMaskSetup s, const uint32_t *__restrict__ pixel_mask,
const int64_t *__restrict__ sum_value, const uint32_t *__restrict__ valid_count,
float *__restrict__ pol, char *__restrict__ valid, int *__restrict__ radius,
double *__restrict__ num, int32_t *__restrict__ den, int *__restrict__ max_radius) {
const size_t n = static_cast<size_t>(s.width) * s.height;
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n)
return;
const int x = static_cast<int>(i % s.width), y = static_cast<int>(i / s.width);
const float dx = x - s.beam_x, dy = y - s.beam_y;
const float p = s.has_polarization ? polarization_factor(s, static_cast<float>(x), static_cast<float>(y)) : 1.0f;
pol[i] = p;
float mean = 0.0f;
char v = 0;
if (valid_count[i] > 0 && pixel_mask[i] == 0 && p > 0.0f) {
mean = static_cast<float>(static_cast<double>(sum_value[i]) / valid_count[i] / p);
v = 1;
}
valid[i] = v;
num[i] = v ? mean : 0.0;
den[i] = v ? 1 : 0;
const int r = static_cast<int>(lroundf(sqrtf(dx * dx + dy * dy)));
radius[i] = r;
atomicMax(max_radius, r);
}
// Sum over the k x k box about each pixel, zero outside the frame: one running sum per row, then one
// per column, each with exactly the terms and order of the host's (box_sum in ShadowFinder.cpp).
template <typename T>
__global__ void box_rows(const T *__restrict__ in, T *__restrict__ out, int W, int H, int half) {
const int y = blockIdx.x * blockDim.x + threadIdx.x;
if (y >= H) return;
const T *src = in + static_cast<size_t>(y) * W;
T *dst = out + static_cast<size_t>(y) * W;
T s = 0;
for (int x = 0; x <= min(half, W - 1); x++)
s += src[x];
for (int x = 0; x < W; x++) {
dst[x] = s;
if (x + half + 1 < W) s += src[x + half + 1];
if (x - half >= 0) s -= src[x - half];
}
}
template <typename T>
__global__ void box_columns(const T *__restrict__ in, T *__restrict__ out, int W, int H, int half) {
const int x = blockIdx.x * blockDim.x + threadIdx.x;
if (x >= W) return;
T s = 0;
for (int y = 0; y <= min(half, H - 1); y++)
s += in[static_cast<size_t>(y) * W + x];
for (int y = 0; y < H; y++) {
out[static_cast<size_t>(y) * W + x] = s;
if (y + half + 1 < H) s += in[static_cast<size_t>(y + half + 1) * W + x];
if (y - half >= 0) s -= in[static_cast<size_t>(y - half) * W + x];
}
}
// Dilation of a 0/1 plane by the (2r+1) square clipped to the frame, as a count over a sliding window
// along rows and then along columns (dilate in ShadowFinder.cpp).
__global__ void dilate_rows(const char *__restrict__ in, char *__restrict__ out, int W, int H, int r) {
const int y = blockIdx.x * blockDim.x + threadIdx.x;
if (y >= H) return;
const char *src = in + static_cast<size_t>(y) * W;
char *dst = out + static_cast<size_t>(y) * W;
int count = 0;
for (int x = 0; x <= min(r, W - 1); x++)
count += src[x];
for (int x = 0; x < W; x++) {
dst[x] = count > 0;
if (x + r + 1 < W) count += src[x + r + 1];
if (x - r >= 0) count -= src[x - r];
}
}
__global__ void dilate_columns(const char *__restrict__ in, char *__restrict__ out, int W, int H, int r) {
const int x = blockIdx.x * blockDim.x + threadIdx.x;
if (x >= W) return;
int count = 0;
for (int y = 0; y <= min(r, H - 1); y++)
count += in[static_cast<size_t>(y) * W + x];
for (int y = 0; y < H; y++) {
out[static_cast<size_t>(y) * W + x] = count > 0;
if (y + r + 1 < H) count += in[static_cast<size_t>(y + r + 1) * W + x];
if (y - r >= 0) count -= in[static_cast<size_t>(y - r) * W + x];
}
}
__global__ void pooled_kernel(size_t n, const double *__restrict__ pooled_sum, const int32_t *__restrict__ pooled_count,
const char *__restrict__ valid, const int *__restrict__ radius,
float *__restrict__ pooled, uint64_t *__restrict__ ring_key) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
const float p = pooled_count[i] > 0 ? static_cast<float>(pooled_sum[i] / pooled_count[i]) : 0.0f;
pooled[i] = p;
ring_key[i] = valid[i] ? (static_cast<uint64_t>(radius[i]) << 32) | float_key(p) : UINT64_MAX;
}
// Where each of the keys' leading 32-bit groups (ring or sector) starts in a sorted key array.
__global__ void group_offsets(const uint64_t *__restrict__ keys, size_t n, int groups, int *__restrict__ offset) {
const int g = blockIdx.x * blockDim.x + threadIdx.x;
if (g > groups) return;
offset[g] = static_cast<int>(lower_bound(keys, n, static_cast<uint64_t>(g) << 32));
}
// The ring's baseline, three iterations of an order statistic over its sorted values (GetMask).
__global__ void baseline_kernel(int rings, const uint64_t *__restrict__ keys, const int *__restrict__ offset,
float *__restrict__ baseline) {
const int r = blockIdx.x * blockDim.x + threadIdx.x;
if (r >= rings) return;
const int lo = offset[r], n = offset[r + 1] - offset[r];
int excluded = 0;
float b = 0.0f;
for (int iter = 0; iter < 3; iter++) {
const int avail = n - excluded;
b = (avail <= 0) ? 0.0f : key_float(static_cast<uint32_t>(keys[lo + excluded + avail / 2]));
const float d = fmaxf(b, 1e-6f);
int excl = 0;
while (excl < n && key_float(static_cast<uint32_t>(keys[lo + excl])) / d < SHADOW_RATIO)
excl++;
excluded = excl;
}
baseline[r] = b;
}
__global__ void low_kernel(size_t n, uint32_t frames, const char *__restrict__ valid, const float *__restrict__ pooled,
const int32_t *__restrict__ pooled_count, const float *__restrict__ pol,
const int *__restrict__ radius, const float *__restrict__ baseline,
float *__restrict__ ratio, float *__restrict__ deficit, char *__restrict__ low) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
const float base = baseline[radius[i]];
const float rt = valid[i] ? pooled[i] / fmaxf(base, 1e-6f) : 1.0f;
ratio[i] = rt;
float df = 0.0f;
char l = 0;
if (valid[i]) {
const double counted = static_cast<double>(frames) * pooled_count[i] * pol[i];
df = static_cast<float>(poisson_deficit_sigma(pooled[i] * counted, base * counted));
l = rt < SHADOW_RATIO && df > MIN_DEFICIT_SIGMA;
}
deficit[i] = df;
low[i] = l;
}
// 8-connected components by union-find: every component ends up named by its smallest pixel index,
// whatever order the unions ran in.
__device__ int find_root(const int *parent, int x) {
while (parent[x] != x)
x = parent[x];
return x;
}
__global__ void cc_init(size_t n, const char *__restrict__ member, int *__restrict__ parent) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
parent[i] = member[i] ? static_cast<int>(i) : -1;
}
__device__ void cc_unite(int *parent, int a, int b) {
while (true) {
a = find_root(parent, a);
b = find_root(parent, b);
if (a == b) return;
if (a < b) { const int t = a; a = b; b = t; }
if (atomicCAS(&parent[a], a, b) == a) return;
}
}
__global__ void cc_union(int W, int H, const char *__restrict__ member, int *parent) {
const size_t n = static_cast<size_t>(W) * H;
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n || !member[i]) return;
const int x = static_cast<int>(i % W), y = static_cast<int>(i / W);
// The four neighbours before this pixel; the other four see it from their side.
if (x > 0 && member[i - 1]) cc_unite(parent, static_cast<int>(i), static_cast<int>(i - 1));
if (y > 0) {
const size_t up = i - W;
if (member[up]) cc_unite(parent, static_cast<int>(i), static_cast<int>(up));
if (x > 0 && member[up - 1]) cc_unite(parent, static_cast<int>(i), static_cast<int>(up - 1));
if (x + 1 < W && member[up + 1]) cc_unite(parent, static_cast<int>(i), static_cast<int>(up + 1));
}
}
__global__ void cc_flatten(size_t n, int *parent) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n || parent[i] < 0) return;
parent[i] = find_root(parent, static_cast<int>(i));
}
// Per component (by root): how many of its pixels are `a`, and how many are both `a` and `b`.
__global__ void cc_count(size_t n, const int *__restrict__ root, const char *__restrict__ a, const char *__restrict__ b,
int *__restrict__ count_a, int *__restrict__ count_ab) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n || root[i] < 0 || !a[i]) return;
atomicAdd(&count_a[root[i]], 1);
if (count_ab && b[i])
atomicAdd(&count_ab[root[i]], 1);
}
__global__ void region_kernel(size_t n, const int *__restrict__ root, const int *__restrict__ n_low,
const char *__restrict__ low, const char *__restrict__ valid, const int *__restrict__ radius,
int blocked_out_to, char *__restrict__ region) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
char r = root[i] >= 0 && n_low[root[i]] >= MIN_SHADOW_PIXELS ? low[i] : 0;
if (valid[i] && radius[i] <= blocked_out_to)
r = 1;
region[i] = r;
}
__global__ void lit_kernel(size_t n, const uint32_t *__restrict__ valid_count, const int64_t *__restrict__ max_value,
char *__restrict__ lit) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
lit[i] = (valid_count[i] > 0) && (max_value[i] >= MIN_REFLECTION);
}
__global__ void reflection_kernel(int W, int H, const char *__restrict__ lit, char *__restrict__ reflection) {
const size_t n = static_cast<size_t>(W) * H;
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
const int x = static_cast<int>(i % W), y = static_cast<int>(i / W);
char r = 0;
if (lit[i]) {
int neighbours = 0;
for (int dy = -1; dy <= 1; dy++)
for (int dx = -1; dx <= 1; dx++) {
const int yy = y + dy, xx = x + dx;
if ((dx || dy) && yy >= 0 && yy < H && xx >= 0 && xx < W && lit[static_cast<size_t>(yy) * W + xx])
neighbours++;
}
r = neighbours >= 2;
}
reflection[i] = r;
}
__global__ void penumbra_kernel(size_t n, const char *__restrict__ penumbra, const char *__restrict__ valid,
const float *__restrict__ ratio, const float *__restrict__ deficit, char *__restrict__ region) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
if (penumbra[i] && valid[i] && ratio[i] < PENUMBRA_RATIO && deficit[i] > MIN_DEFICIT_SIGMA)
region[i] = 1;
}
__global__ void invert_kernel(size_t n, const char *__restrict__ in, char *__restrict__ out) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
out[i] = !in[i];
}
// Background components that touch the frame's edge are outside; the rest are holes.
__global__ void outside_kernel(int W, int H, const int *__restrict__ root, int *__restrict__ outside) {
const int k = blockIdx.x * blockDim.x + threadIdx.x;
const int perimeter = 2 * W + 2 * H;
if (k >= perimeter) return;
int x, y;
if (k < W) { x = k; y = 0; }
else if (k < 2 * W) { x = k - W; y = H - 1; }
else if (k < 2 * W + H) { x = 0; y = k - 2 * W; }
else { x = W - 1; y = k - 2 * W - H; }
const int r = root[static_cast<size_t>(y) * W + x];
if (r >= 0) outside[r] = 1;
}
__global__ void final_kernel(size_t n, const int *__restrict__ background_root, const int *__restrict__ outside,
const char *__restrict__ reflection_grown, char *__restrict__ region,
uint32_t *__restrict__ mask) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
char r = region[i];
if (background_root[i] >= 0 && !outside[background_root[i]])
r = 1; // a hole
if (reflection_grown[i])
r = 0;
region[i] = r;
mask[i] = r ? MASK_SHADOW : 0;
}
__global__ void sector_key_kernel(ShadowMaskSetup s, const char *__restrict__ valid, const char *__restrict__ region,
const int *__restrict__ radius, const float *__restrict__ ratio,
uint64_t *__restrict__ key) {
const size_t n = static_cast<size_t>(s.width) * s.height;
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
if (!valid[i] || region[i]) {
key[i] = UINT64_MAX;
return;
}
const float dx = static_cast<float>(i % s.width) - s.beam_x, dy = static_cast<float>(i / s.width) - s.beam_y;
const double phi = atan2f(dy, dx) + PI;
const int sector = min(HARMONIC_SECTORS - 1, static_cast<int>(phi / (2.0 * PI) * HARMONIC_SECTORS));
const uint64_t k = static_cast<uint64_t>(radius[i] / HARMONIC_BAND_PX) * HARMONIC_SECTORS + sector;
key[i] = (k << 32) | float_key(ratio[i]);
}
__global__ void sector_median_kernel(int sectors, const uint64_t *__restrict__ keys, const int *__restrict__ offset,
double *__restrict__ median) {
const int k = blockIdx.x * blockDim.x + threadIdx.x;
if (k >= sectors) return;
const int n = offset[k + 1] - offset[k];
median[k] = n >= MIN_SECTOR_PIXELS ? key_float(static_cast<uint32_t>(keys[offset[k] + n / 2])) : -1.0;
}
__global__ void dim_kernel(ShadowMaskSetup s, uint32_t frames, const char *__restrict__ valid,
const char *__restrict__ region, const float *__restrict__ ratio,
const float *__restrict__ deficit, const int *__restrict__ radius,
const float *__restrict__ harm_c, const float *__restrict__ harm_s,
const int32_t *__restrict__ pooled_count, const float *__restrict__ pol,
const float *__restrict__ pooled, const float *__restrict__ baseline,
char *__restrict__ dim) {
const size_t n = static_cast<size_t>(s.width) * s.height;
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
char d = 0;
if (valid[i] && !region[i] && ratio[i] < PENUMBRA_RATIO) {
// cos 2phi and sin 2phi from the offset to the beam.
const float dx = static_cast<float>(i % s.width) - s.beam_x, dy = static_cast<float>(i / s.width) - s.beam_y;
const float r2 = fmaxf(dx * dx + dy * dy, 1e-6f);
const int band = radius[i] / HARMONIC_BAND_PX;
const float model = fminf(1.0f, 1.0f + harm_c[band] * (dx * dx - dy * dy) / r2
+ harm_s[band] * 2.0f * dx * dy / r2);
if (model >= 1.0f) {
d = deficit[i] > MIN_DEFICIT_SIGMA;
} else {
const double counted = static_cast<double>(frames) * pooled_count[i] * pol[i];
d = ratio[i] < PENUMBRA_RATIO * model
&& poisson_deficit_sigma(pooled[i] * counted, baseline[radius[i]] * model * counted) > MIN_DEFICIT_SIGMA;
}
}
dim[i] = d;
}
// Join a region across the module gaps it crosses, one line per thread (bridge_gaps in
// ShadowFinder.cpp). Both directions read `region` and only ever set pixels of `out` to 1.
__global__ void bridge_rows(int W, int H, const char *__restrict__ region, const char *__restrict__ valid,
char *out) {
const int y = blockIdx.x * blockDim.x + threadIdx.x;
if (y >= H) return;
const size_t row = static_cast<size_t>(y) * W;
int k = 0;
while (k < W) {
if (valid[row + k]) { k++; continue; }
const int start = k;
while (k < W && !valid[row + k]) k++;
if (start > 0 && k < W && region[row + start - 1] && region[row + k])
for (int j = start; j < k; j++) out[row + j] = 1;
}
}
__global__ void bridge_columns(int W, int H, const char *__restrict__ region, const char *__restrict__ valid,
char *out) {
const int x = blockIdx.x * blockDim.x + threadIdx.x;
if (x >= W) return;
const auto at = [&](int y) { return static_cast<size_t>(y) * W + x; };
int k = 0;
while (k < H) {
if (valid[at(k)]) { k++; continue; }
const int start = k;
while (k < H && !valid[at(k)]) k++;
if (start > 0 && k < H && region[at(start - 1)] && region[at(k)])
for (int j = start; j < k; j++) out[at(j)] = 1;
}
}
__global__ void transmitting_kernel(size_t n, const int *__restrict__ root, const char *__restrict__ dim,
const int *__restrict__ n_dim, const int *__restrict__ n_low,
uint32_t *__restrict__ mask) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
const int c = root[i];
if (c >= 0 && dim[i] && n_dim[c] >= MIN_SHADOW_PIXELS && n_low[c] >= MIN_CORE_PIXELS)
mask[i] = MASK_TRANSMITTING;
}
__global__ void mean_kernel(size_t n, const uint32_t *__restrict__ pixel_mask, const int64_t *__restrict__ sum_value,
const uint32_t *__restrict__ valid_count, float *__restrict__ mean) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= n) return;
mean[i] = valid_count[i] > 0 && pixel_mask[i] == 0
? static_cast<float>(static_cast<double>(sum_value[i]) / valid_count[i]) : NAN;
}
// The device side of one GetMask: planes, sort scratch and the stream to run on.
class MaskEngine {
public:
const ShadowMaskSetup s;
const int W, H;
const size_t n;
cudaStream_t stream;
MaskEngine(const ShadowMaskSetup &setup, cudaStream_t st)
: s(setup), W(setup.width), H(setup.height), n(static_cast<size_t>(setup.width) * setup.height), stream(st) {}
void Check(const char *what) const {
check(cudaGetLastError(), what);
}
void Dilate(const char *in, char *out, char *scratch, int r) const {
dilate_rows<<<grid(H), THREADS, 0, stream>>>(in, scratch, W, H, r);
dilate_columns<<<grid(W), THREADS, 0, stream>>>(scratch, out, W, H, r);
Check("dilate");
}
// Components of `member`, each pixel's root in `root` (-1 outside every component).
void Label(const char *member, int *root) const {
cc_init<<<grid(n), THREADS, 0, stream>>>(n, member, root);
cc_union<<<grid(n), THREADS, 0, stream>>>(W, H, member, root);
cc_flatten<<<grid(n), THREADS, 0, stream>>>(n, root);
Check("components");
}
void SortKeys(uint64_t *keys, uint64_t *sorted) const {
size_t bytes = 0;
check(cub::DeviceRadixSort::SortKeys(nullptr, bytes, keys, sorted, n, 0, 64, stream), "sort size");
CudaDevicePtr<uint8_t> scratch(bytes);
check(cub::DeviceRadixSort::SortKeys(scratch.get(), bytes, keys, sorted, n, 0, 64, stream), "sort");
// The scratch is freed on the allocation stream, which knows nothing of this one.
check(cudaStreamSynchronize(stream), "sort");
}
template <typename T>
std::vector<T> Download(const T *device, size_t count) const {
std::vector<T> host(count);
check(cudaMemcpyAsync(host.data(), device, count * sizeof(T), cudaMemcpyDeviceToHost, stream), "download");
check(cudaStreamSynchronize(stream), "download");
return host;
}
template <typename T>
void Upload(T *device, const std::vector<T> &host) const {
check(cudaMemcpyAsync(device, host.data(), host.size() * sizeof(T), cudaMemcpyHostToDevice, stream), "upload");
}
};
} // namespace
std::vector<uint32_t> ShadowMaskOnDevice(const ShadowMaskSetup &setup, const std::vector<uint32_t> &pixel_mask,
const int64_t *max_value, const int64_t *sum_value,
const uint32_t *valid_count, uint32_t frames, cudaStream_t stream) {
const MaskEngine e(setup, stream);
const size_t n = e.n;
const int W = e.W, H = e.H;
CudaDevicePtr<uint32_t> d_pixel_mask(n), mask(n);
e.Upload(d_pixel_mask.get(), pixel_mask);
// Mean projection over the polarization factor, usable pixels and radius from the beam centre.
CudaDevicePtr<float> pol(n), pooled(n), ratio(n), deficit(n);
CudaDevicePtr<char> valid(n), low(n), region(n), a(n), b(n), c(n);
CudaDevicePtr<int> radius(n), root(n), count_a(n), count_b(n), max_radius(1);
CudaDevicePtr<double> num(n), dsum(n);
CudaDevicePtr<int32_t> den(n), dcount(n);
check(cudaMemsetAsync(max_radius.get(), 0, sizeof(int), stream), "memset");
setup_kernel<<<grid(n), THREADS, 0, stream>>>(setup, d_pixel_mask, sum_value, valid_count, pol, valid, radius,
num, den, max_radius);
e.Check("setup");
const int max_r = e.Download(max_radius.get(), 1)[0];
const int rings = max_r + 1;
// The background pooled over a small box.
box_rows<double><<<grid(H), THREADS, 0, stream>>>(num, dsum, W, H, POOL_PX / 2);
box_columns<double><<<grid(W), THREADS, 0, stream>>>(dsum, num, W, H, POOL_PX / 2);
box_rows<int32_t><<<grid(H), THREADS, 0, stream>>>(den, dcount, W, H, POOL_PX / 2);
box_columns<int32_t><<<grid(W), THREADS, 0, stream>>>(dcount, den, W, H, POOL_PX / 2);
e.Check("pooling");
const double *pooled_sum = num;
const int32_t *pooled_count = den;
// The rings, each sorted once; the baseline is an order statistic of them.
CudaDevicePtr<uint64_t> keys(n), sorted(n);
pooled_kernel<<<grid(n), THREADS, 0, stream>>>(n, pooled_sum, pooled_count, valid, radius, pooled, keys);
e.Check("pooled");
e.SortKeys(keys, sorted);
CudaDevicePtr<int> ring_offset(rings + 1);
group_offsets<<<grid(rings + 1), THREADS, 0, stream>>>(sorted, n, rings, ring_offset);
CudaDevicePtr<float> baseline(rings);
baseline_kernel<<<grid(rings), THREADS, 0, stream>>>(rings, sorted, ring_offset, baseline);
e.Check("baseline");
const auto host_baseline = e.Download(baseline.get(), rings);
const auto offsets = e.Download(ring_offset.get(), rings + 1);
std::vector<int> ring_pixels(rings);
for (int r = 0; r < rings; r++)
ring_pixels[r] = offsets[r + 1] - offsets[r];
const int blocked_out_to = BlockedOutTo(host_baseline, ring_pixels);
// Low pixels, and the regions of them large enough to be hardware.
low_kernel<<<grid(n), THREADS, 0, stream>>>(n, frames, valid, pooled, pooled_count, pol, radius, baseline,
ratio, deficit, low);
e.Check("low");
e.Dilate(low, a, c, BRIDGE_PX);
e.Label(a, root);
check(cudaMemsetAsync(count_a.get(), 0, n * sizeof(int), stream), "memset");
cc_count<<<grid(n), THREADS, 0, stream>>>(n, root, low, low, count_a, nullptr);
region_kernel<<<grid(n), THREADS, 0, stream>>>(n, root, count_a, low, valid, radius, blocked_out_to, region);
e.Check("region");
// Recorded reflections; `b` holds them until they are given back at the end.
lit_kernel<<<grid(n), THREADS, 0, stream>>>(n, valid_count, max_value, a);
reflection_kernel<<<grid(n), THREADS, 0, stream>>>(W, H, a, b);
e.Check("reflections");
// Penumbra, round and fill.
e.Dilate(region, a, c, PENUMBRA_MAX_PX);
penumbra_kernel<<<grid(n), THREADS, 0, stream>>>(n, a, valid, ratio, deficit, region);
e.Dilate(region, a, c, 2);
invert_kernel<<<grid(n), THREADS, 0, stream>>>(n, a, region);
e.Dilate(region, a, c, 2);
invert_kernel<<<grid(n), THREADS, 0, stream>>>(n, a, region); // region = erode(dilate(region))
invert_kernel<<<grid(n), THREADS, 0, stream>>>(n, region, a); // the background
e.Label(a, root);
check(cudaMemsetAsync(count_a.get(), 0, n * sizeof(int), stream), "memset");
outside_kernel<<<grid(2 * W + 2 * H), THREADS, 0, stream>>>(W, H, root, count_a);
e.Dilate(b, c, a, 1); // the reflections, grown by one
final_kernel<<<grid(n), THREADS, 0, stream>>>(n, root, count_a, c, region, mask);
e.Check("fill");
// The arm search: sector medians of the ratio, the harmonic of each band, the dim pixels.
const int n_bands = max_r / HARMONIC_BAND_PX + 1;
const int n_sectors = n_bands * HARMONIC_SECTORS;
sector_key_kernel<<<grid(n), THREADS, 0, stream>>>(setup, valid, region, radius, ratio, keys);
e.Check("sectors");
e.SortKeys(keys, sorted);
CudaDevicePtr<int> sector_offset(n_sectors + 1);
group_offsets<<<grid(n_sectors + 1), THREADS, 0, stream>>>(sorted, n, n_sectors, sector_offset);
CudaDevicePtr<double> sector_median(n_sectors);
sector_median_kernel<<<grid(n_sectors), THREADS, 0, stream>>>(n_sectors, sorted, sector_offset, sector_median);
e.Check("sector medians");
std::vector<float> harm_c, harm_s;
HarmonicFit(e.Download(sector_median.get(), n_sectors), n_bands, harm_c, harm_s);
CudaDevicePtr<float> d_harm_c(n_bands), d_harm_s(n_bands);
e.Upload(d_harm_c.get(), harm_c);
e.Upload(d_harm_s.get(), harm_s);
dim_kernel<<<grid(n), THREADS, 0, stream>>>(setup, frames, valid, region, ratio, deficit, radius, d_harm_c, d_harm_s,
pooled_count, pol, pooled, baseline, a);
e.Check("dim");
e.Dilate(a, b, c, BRIDGE_PX);
check(cudaMemcpyAsync(c.get(), b.get(), n, cudaMemcpyDeviceToDevice, stream), "copy");
bridge_rows<<<grid(H), THREADS, 0, stream>>>(W, H, b, valid, c);
bridge_columns<<<grid(W), THREADS, 0, stream>>>(W, H, b, valid, c);
e.Label(c, root);
check(cudaMemsetAsync(count_a.get(), 0, n * sizeof(int), stream), "memset");
check(cudaMemsetAsync(count_b.get(), 0, n * sizeof(int), stream), "memset");
cc_count<<<grid(n), THREADS, 0, stream>>>(n, root, a, low, count_a, count_b);
transmitting_kernel<<<grid(n), THREADS, 0, stream>>>(n, root, a, count_a, count_b, mask);
e.Check("transmitting");
return e.Download(mask.get(), n);
}
std::vector<float> MeanProjectionOnDevice(const std::vector<uint32_t> &pixel_mask, const int64_t *sum_value,
const uint32_t *valid_count, size_t npixels, cudaStream_t stream) {
CudaDevicePtr<uint32_t> d_pixel_mask(npixels);
CudaDevicePtr<float> mean(npixels);
check(cudaMemcpyAsync(d_pixel_mask.get(), pixel_mask.data(), npixels * sizeof(uint32_t), cudaMemcpyHostToDevice,
stream), "upload");
mean_kernel<<<grid(npixels), THREADS, 0, stream>>>(npixels, d_pixel_mask, sum_value, valid_count, mean);
check(cudaGetLastError(), "mean");
std::vector<float> host(npixels);
check(cudaMemcpyAsync(host.data(), mean.get(), npixels * sizeof(float), cudaMemcpyDeviceToHost, stream), "download");
check(cudaStreamSynchronize(stream), "mean");
return host;
}
image_analysis/beam_stop/ShadowMaskGPU.h
+39
View File
@@ -0,0 +1,39 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
// Included only under JFJOCH_USE_CUDA.
#include <cstdint>
#include <vector>
#include <cuda_runtime.h>
// What the mask is drawn about: the detector, the centre the rings are drawn about, and the geometry
// the polarization factor is read off (ShadowFinder keeps it as a DiffractionGeometry; the device
// takes it as numbers).
struct ShadowMaskSetup {
int width = 0, height = 0;
float beam_x = 0.0f, beam_y = 0.0f;
float det_matrix[9] = {}; // row major
float pixel_size_mm = 0.0f;
float distance_mm = 0.0f;
bool has_polarization = false;
float polarization = 0.0f;
};
// ShadowFinder::GetMask on the device, from the projection ShadowAccumulatorGPU holds there. Step for
// step the host's algorithm - the same pooling, ring medians, components, morphology and arm search -
// and the same answer wherever the arithmetic is exact: every integer, comparison, sort and component
// is. What is not is the floating point the two compilers evaluate differently - the polarization
// factor's trigonometry, the Poisson test's logarithm and the azimuth of the arm search - so a pixel
// within a rounding of one of those thresholds can come out the other way.
std::vector<uint32_t> ShadowMaskOnDevice(const ShadowMaskSetup &setup, const std::vector<uint32_t> &pixel_mask,
const int64_t *max_value, const int64_t *sum_value,
const uint32_t *valid_count, uint32_t frames, cudaStream_t stream);
// The mean projection the host's GetMeanProjection makes, computed where the sums are: the same
// division, so the same bits, and a quarter of the bytes to bring back.
std::vector<float> MeanProjectionOnDevice(const std::vector<uint32_t> &pixel_mask, const int64_t *sum_value,
const uint32_t *valid_count, size_t npixels, cudaStream_t stream);
image_analysis/geom_refinement/BackgroundBand.h
+107
View File
@@ -0,0 +1,107 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
// Where one pixel falls in the background beam-centre fit (FindBeamCenterFromBackground): which
// radial-bin x sector cell of the fitted band it lands in about a trial centre, and the derivative
// of its 2theta with respect to that centre. Written once and compiled both for the host fit and for
// its device twin (BeamCenterBackgroundGPU), so the two read the same formula.
//
// The same formula is the same number only when both sides evaluate it the same way. The library
// atan2f differs between glibc and CUDA in the last bit, and the compilers fuse multiply-adds each
// in their own way; either moves a pixel within a rounding of a cell edge into the neighbouring cell
// (measured on 16 Mpx sweeps: ~30 of 6.5 million pixels, enough to move the fitted centre by up to
// 0.05 px). So the angles come from BackgroundAtan2 below and both translation units are compiled
// without contraction (geom_refinement/CMakeLists.txt), and the host and the device then agree to
// the bit.
#include <cmath>
#include <cstdint>
#include "../../common/JFJochMath.h"
#ifdef __CUDACC__
#define BACKGROUND_BAND_HD __host__ __device__ inline
#else
#define BACKGROUND_BAND_HD inline
#endif
struct BackgroundBand {
static constexpr int SECTORS = 36;
static constexpr int RADIAL_BINS = 120;
static constexpr int CELLS = RADIAL_BINS * SECTORS;
float rot[9]; // detector matrix, row major
float pixel_size; // mm
float distance; // mm
float tt_lo, tt_hi; // the band in 2theta
float d_tt; // width of one radial bin
double tan_lo, tan_hi; // the band in tan(2theta), widened for the quick rejection
};
// atan2(y, x) from IEEE operations alone - add, multiply, divide, square root, all correctly rounded on
// the host and on the device - so that, compiled without contraction (see CMakeLists.txt), the host
// and the device return the same bits. The library atan2f does not: glibc's and CUDA's differ in the
// last place. Two half-angle reductions take the argument below tan(pi/16), where eleven terms of the
// series leave an error under 1e-17.
BACKGROUND_BAND_HD double BackgroundAtan2(double y, double x) {
const double ax = x < 0 ? -x : x, ay = y < 0 ? -y : y;
if (ax == 0.0 && ay == 0.0)
return 0.0;
const bool swap = ay > ax;
double t = swap ? ax / ay : ay / ax;
t = t / (1.0 + sqrt(1.0 + t * t));
t = t / (1.0 + sqrt(1.0 + t * t));
const double t2 = t * t;
double series = 1.0 / 23.0;
for (int k = 21; k >= 1; k -= 2)
series = 1.0 / k - t2 * series;
double a = 4.0 * t * series;
if (swap) a = PI / 2 - a;
if (x < 0) a = PI - a;
return y < 0 ? -a : a;
}
// The cell of pixel (x, y) about (beam_x, beam_y), or -1 when it is outside the band; for a pixel in
// it, also the two components of the derivative of its 2theta with respect to the centre.
BACKGROUND_BAND_HD int BackgroundBandCell(const BackgroundBand &b, int x, int y, float beam_x, float beam_y,
float &jac_x, float &jac_y) {
const float *rot = b.rot;
const float u = (x - beam_x) * b.pixel_size;
const float v = (y - beam_y) * b.pixel_size;
const float lx = rot[0] * u + rot[1] * v + rot[2] * b.distance;
const float ly = rot[3] * u + rot[4] * v + rot[5] * b.distance;
const float lz = rot[6] * u + rot[7] * v + rot[8] * b.distance;
const float rho_sq = lx * lx + ly * ly;
// Most pixels lie outside the band. Those clearly outside it in tan(2theta) = rho / lz - by a
// margin far above float rounding - skip the square root and both atan2 below; every pixel the
// exact test would keep still reaches it.
if (lz > 0.0f) {
const double lz_sq = static_cast<double>(lz) * lz;
if (rho_sq < b.tan_lo * b.tan_lo * lz_sq || rho_sq > b.tan_hi * b.tan_hi * lz_sq)
return -1;
}
const float rho = sqrtf(rho_sq);
const float two_theta = static_cast<float>(BackgroundAtan2(rho, lz));
if (two_theta < b.tt_lo || two_theta >= b.tt_hi || rho == 0.0f)
return -1;
const float phi = static_cast<float>(BackgroundAtan2(ly, lx));
// Both bins are clamped: a pixel one float ulp below the top of the band divides to exactly
// RADIAL_BINS, which is one cell past the end of every accumulator.
int r_bin = static_cast<int>((two_theta - b.tt_lo) / b.d_tt);
r_bin = r_bin < 0 ? 0 : (r_bin > BackgroundBand::RADIAL_BINS - 1 ? BackgroundBand::RADIAL_BINS - 1 : r_bin);
int s_bin = static_cast<int>((phi + PI) / (2 * PI) * BackgroundBand::SECTORS);
s_bin = s_bin < 0 ? 0 : (s_bin > BackgroundBand::SECTORS - 1 ? BackgroundBand::SECTORS - 1 : s_bin);
// d(2theta)/d(beam), through the lab coordinate: the detector coordinate depends on the centre
// only as (x - beam_x), so moving the centre is moving the pixel.
const float denominator = rho * rho + lz * lz;
const float g_x = lz * lx / (rho * denominator);
const float g_y = lz * ly / (rho * denominator);
const float g_z = -rho / denominator;
jac_x = -b.pixel_size * (g_x * rot[0] + g_y * rot[3] + g_z * rot[6]);
jac_y = -b.pixel_size * (g_x * rot[1] + g_y * rot[4] + g_z * rot[7]);
return r_bin * BackgroundBand::SECTORS + s_bin;
}
image_analysis/geom_refinement/BeamCenterBackgroundGPU.cu
+210
View File
@@ -0,0 +1,210 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "BeamCenterBackgroundGPU.h"
#include <cub/device/device_radix_sort.cuh>
#include "../indexing/CUDAMemHelpers.h"
namespace {
constexpr int THREADS = 256;
void check(cudaError_t err, const char *what) {
if (err != cudaSuccess)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
std::string("Beam centre from background: ") + what + ": " + cudaGetErrorString(err));
}
// Every pixel's cell (BackgroundBand::CELLS where it is outside the band or unusable), its index,
// and its derivatives.
__global__ void bin_kernel(BackgroundBand band, int width, size_t npixels, float beam_x, float beam_y,
const char *__restrict__ usable, int32_t *__restrict__ key,
int32_t *__restrict__ index, float *__restrict__ jac_x,
float *__restrict__ jac_y) {
const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
if (i >= npixels)
return;
int cell = BackgroundBand::CELLS;
float jx = 0.0f, jy = 0.0f;
if (usable[i]) {
const int c = BackgroundBandCell(band, static_cast<int>(i % width), static_cast<int>(i / width),
beam_x, beam_y, jx, jy);
if (c >= 0)
cell = c;
}
key[i] = cell;
index[i] = static_cast<int32_t>(i);
jac_x[i] = jx;
jac_y[i] = jy;
}
// Where each cell's pixels start in the sorted keys; offset[CELLS] is the number of binned pixels.
__global__ void offset_kernel(const int32_t *__restrict__ sorted_key, size_t n, int32_t *__restrict__ offset) {
const int c = blockIdx.x * blockDim.x + threadIdx.x;
if (c > BackgroundBand::CELLS)
return;
size_t lo = 0, hi = n;
while (lo < hi) {
const size_t mid = (lo + hi) / 2;
if (sorted_key[mid] < c) lo = mid + 1;
else hi = mid;
}
offset[c] = static_cast<int32_t>(lo);
}
// One thread per cell, walking its pixels in pixel order: a partial sum per host row block, added to
// the cell's total when the block changes - the host's arithmetic, step for step. A cell whose pixels
// are clipped (clip_limit != nullptr) skips the ones above its limit, as the host's clip rounds do.
__global__ void sum_kernel(int width, const int32_t *__restrict__ offset, const int32_t *__restrict__ index,
const int32_t *__restrict__ row_block, const float *__restrict__ mean,
const float *__restrict__ jac_x, const float *__restrict__ jac_y,
const float *__restrict__ clip_limit,
double *__restrict__ sum, double *__restrict__ sum_sq,
double *__restrict__ sum_jx, double *__restrict__ sum_jy,
int32_t *__restrict__ count) {
const int c = blockIdx.x * blockDim.x + threadIdx.x;
if (c >= BackgroundBand::CELLS)
return;
const bool clipped = clip_limit != nullptr;
const float limit = clipped ? clip_limit[c] : 0.0f;
double s = 0, ss = 0, jx = 0, jy = 0;
double bs = 0, bss = 0, bjx = 0, bjy = 0;
int32_t n = 0;
int block = -1;
for (int32_t p = offset[c]; p < offset[c + 1]; p++) {
const int32_t i = index[p];
const float value = mean[i];
if (clipped && (limit < 0.0f || value > limit))
continue;
const int b = row_block[i / width];
if (b != block) {
s += bs; ss += bss; jx += bjx; jy += bjy;
bs = bss = bjx = bjy = 0;
block = b;
}
n++;
bs += value;
bss += static_cast<double>(value) * value;
if (!clipped) {
bjx += jac_x[i];
bjy += jac_y[i];
}
}
s += bs; ss += bss; jx += bjx; jy += bjy;
sum[c] = s;
sum_sq[c] = ss;
count[c] = n;
if (!clipped) {
sum_jx[c] = jx;
sum_jy[c] = jy;
}
}
} // namespace
struct BeamCenterBackgroundGPU::Impl {
int width, height;
size_t npixels;
CudaStream stream;
CudaDevicePtr<char> usable;
CudaDevicePtr<float> mean;
CudaDevicePtr<int32_t> row_block;
CudaDevicePtr<int32_t> key, sorted_key, index, sorted_index;
CudaDevicePtr<float> jac_x, jac_y;
CudaDevicePtr<int32_t> offset;
CudaDevicePtr<float> clip_limit;
CudaDevicePtr<double> sum, sum_sq, sum_jx, sum_jy;
CudaDevicePtr<int32_t> count;
CudaDevicePtr<uint8_t> sort_scratch;
size_t sort_scratch_bytes = 0;
Impl(int w, int h)
: width(w), height(h), npixels(static_cast<size_t>(w) * h),
usable(npixels), mean(npixels), row_block(h),
key(npixels), sorted_key(npixels), index(npixels), sorted_index(npixels),
jac_x(npixels), jac_y(npixels), offset(BackgroundBand::CELLS + 1),
clip_limit(BackgroundBand::CELLS),
sum(BackgroundBand::CELLS), sum_sq(BackgroundBand::CELLS),
sum_jx(BackgroundBand::CELLS), sum_jy(BackgroundBand::CELLS),
count(BackgroundBand::CELLS) {
// Keys run to CELLS inclusive, the bin of everything outside the band.
check(cub::DeviceRadixSort::SortPairs(nullptr, sort_scratch_bytes, key.get(), sorted_key.get(),
index.get(), sorted_index.get(), npixels, 0, end_bit(), stream),
"sort size");
sort_scratch = CudaDevicePtr<uint8_t>(sort_scratch_bytes);
}
static int end_bit() {
int bits = 0;
while ((1 << bits) <= BackgroundBand::CELLS) bits++;
return bits;
}
void Download(std::vector<double> &s, std::vector<double> &ss, std::vector<double> *jx,
std::vector<double> *jy, std::vector<int32_t> &n) {
const size_t cells = BackgroundBand::CELLS;
check(cudaMemcpyAsync(s.data(), sum.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy");
check(cudaMemcpyAsync(ss.data(), sum_sq.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy");
if (jx) check(cudaMemcpyAsync(jx->data(), sum_jx.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy");
if (jy) check(cudaMemcpyAsync(jy->data(), sum_jy.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy");
check(cudaMemcpyAsync(n.data(), count.get(), cells * sizeof(int32_t), cudaMemcpyDeviceToHost, stream), "copy");
check(cudaStreamSynchronize(stream), "sums");
}
};
BeamCenterBackgroundGPU::BeamCenterBackgroundGPU(int width, int height, const std::vector<int> &block_row,
const char *usable, const float *mean)
: impl(std::make_unique<Impl>(width, height)) {
std::vector<int32_t> row_block(height);
for (size_t b = 0; b + 1 < block_row.size(); b++)
for (int y = block_row[b]; y < block_row[b + 1]; y++)
row_block[y] = static_cast<int32_t>(b);
check(cudaMemcpyAsync(impl->row_block.get(), row_block.data(), height * sizeof(int32_t),
cudaMemcpyHostToDevice, impl->stream), "upload");
check(cudaMemcpyAsync(impl->usable.get(), usable, impl->npixels, cudaMemcpyHostToDevice, impl->stream), "upload");
check(cudaMemcpyAsync(impl->mean.get(), mean, impl->npixels * sizeof(float), cudaMemcpyHostToDevice,
impl->stream), "upload");
check(cudaStreamSynchronize(impl->stream), "upload");
}
BeamCenterBackgroundGPU::~BeamCenterBackgroundGPU() = default;
void BeamCenterBackgroundGPU::Bin(const BackgroundBand &band, float beam_x, float beam_y,
std::vector<double> &sum, std::vector<double> &sum_sq,
std::vector<double> &sum_jx, std::vector<double> &sum_jy,
std::vector<int32_t> &count) {
Impl &d = *impl;
const auto blocks = static_cast<unsigned>((d.npixels + THREADS - 1) / THREADS);
bin_kernel<<<blocks, THREADS, 0, d.stream>>>(band, d.width, d.npixels, beam_x, beam_y, d.usable,
d.key, d.index, d.jac_x, d.jac_y);
check(cudaGetLastError(), "bin");
// A radix sort is stable, so each cell's pixels come out in pixel order.
check(cub::DeviceRadixSort::SortPairs(d.sort_scratch.get(), d.sort_scratch_bytes, d.key.get(),
d.sorted_key.get(), d.index.get(), d.sorted_index.get(),
d.npixels, 0, Impl::end_bit(), d.stream), "sort");
constexpr int cell_blocks = (BackgroundBand::CELLS + 1 + THREADS - 1) / THREADS;
offset_kernel<<<cell_blocks, THREADS, 0, d.stream>>>(d.sorted_key, d.npixels, d.offset);
check(cudaGetLastError(), "offsets");
sum_kernel<<<cell_blocks, THREADS, 0, d.stream>>>(d.width, d.offset, d.sorted_index, d.row_block, d.mean,
d.jac_x, d.jac_y, nullptr, d.sum, d.sum_sq,
d.sum_jx, d.sum_jy, d.count);
check(cudaGetLastError(), "sums");
d.Download(sum, sum_sq, &sum_jx, &sum_jy, count);
}
void BeamCenterBackgroundGPU::Clip(const std::vector<float> &clip_limit,
std::vector<double> &sum, std::vector<double> &sum_sq,
std::vector<int32_t> &count) {
Impl &d = *impl;
check(cudaMemcpyAsync(d.clip_limit.get(), clip_limit.data(), BackgroundBand::CELLS * sizeof(float),
cudaMemcpyHostToDevice, d.stream), "upload");
constexpr int cell_blocks = (BackgroundBand::CELLS + THREADS - 1) / THREADS;
sum_kernel<<<cell_blocks, THREADS, 0, d.stream>>>(d.width, d.offset, d.sorted_index, d.row_block, d.mean,
d.jac_x, d.jac_y, d.clip_limit, d.sum, d.sum_sq,
d.sum_jx, d.sum_jy, d.count);
check(cudaGetLastError(), "clip");
d.Download(sum, sum_sq, nullptr, nullptr, count);
}
image_analysis/geom_refinement/BeamCenterBackgroundGPU.h
+40
View File
@@ -0,0 +1,40 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
// Included only under JFJOCH_USE_CUDA. Free of CUDA headers, so the host fit can hold one.
#include <cstdint>
#include <memory>
#include <vector>
#include "BackgroundBand.h"
// The two passes over the pixels of the background beam-centre fit (FindBeamCenterFromBackground),
// on the device: binning every usable pixel into its cell about a trial centre, and summing the
// binned pixels again under a clip. They are all of the fit's cost; the fit itself stays on the host.
//
// Each cell is summed in the order the host sums it - pixel order within each of the host's row
// blocks, the blocks then added in block order - so a pixel that lands in the same cell on both
// sides adds the same rounding on both. Whatever differs comes from BackgroundBandCell (see there).
class BeamCenterBackgroundGPU {
struct Impl;
std::unique_ptr<Impl> impl;
public:
// block_row: the first row of each of the host's row blocks, and one past the last row at the end.
BeamCenterBackgroundGPU(int width, int height, const std::vector<int> &block_row,
const char *usable, const float *mean);
~BeamCenterBackgroundGPU();
// Bin the band about (beam_x, beam_y) and sum each cell: the values, their squares, the two
// derivatives, and the count.
void Bin(const BackgroundBand &band, float beam_x, float beam_y,
std::vector<double> &sum, std::vector<double> &sum_sq,
std::vector<double> &sum_jx, std::vector<double> &sum_jy, std::vector<int32_t> &count);
// Sum the pixels the last Bin put in each cell again, leaving out those above the cell's
// clip_limit and every pixel of a cell whose limit is negative.
void Clip(const std::vector<float> &clip_limit,
std::vector<double> &sum, std::vector<double> &sum_sq, std::vector<int32_t> &count);
};
image_analysis/geom_refinement/BeamCenterFromBackground.cpp
+72 -54
View File
@@ -2,6 +2,7 @@
// SPDX-License-Identifier: GPL-3.0-only
#include "BeamCenterFromBackground.h"
#include "BackgroundBand.h"
#include <algorithm>
#include <cmath>
@@ -10,6 +11,10 @@
#include "../../common/CompressedImage.h"
#include "../../common/JFJochMath.h"
#include "../../common/ParallelFor.h"
#ifdef JFJOCH_USE_CUDA
#include "../../common/CUDAWrapper.h"
#include "BeamCenterBackgroundGPU.h"
#endif
namespace {
@@ -18,8 +23,8 @@ namespace {
constexpr float BAND_LOW_RES_A = 12.0f;
constexpr float BAND_HIGH_RES_A = 2.2f;
constexpr int SECTORS = 36;
constexpr int RADIAL_BINS = 120;
constexpr int SECTORS = BackgroundBand::SECTORS;
constexpr int RADIAL_BINS = BackgroundBand::RADIAL_BINS;
// A cell with fewer pixels than this has no usable mean.
constexpr int MIN_PIXELS_PER_CELL = 20;
@@ -84,7 +89,7 @@ float median_of(std::vector<float> &v) {
std::optional<BeamCenterEstimate>
FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const PixelMask &mask,
const std::vector<float> &mean, size_t nthreads,
std::optional<std::pair<float, float>> start) {
std::optional<std::pair<float, float>> start, bool allow_device) {
if (nthreads == 0)
nthreads = std::max(1u, std::thread::hardware_concurrency());
const auto W = static_cast<int>(experiment.GetXPixelsNumConv());
@@ -165,10 +170,31 @@ FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const Pixe
const double tan_lo = std::tan(static_cast<double>(tt_lo)) * (1.0 - 1e-3);
const double tan_hi = tt_hi < PI / 2 ? std::tan(static_cast<double>(tt_hi)) * (1.0 + 1e-3) : INFINITY;
BackgroundBand band{};
for (int k = 0; k < 9; k++)
band.rot[k] = rot[k];
band.pixel_size = pixel_size;
band.distance = distance;
band.tt_lo = tt_lo;
band.tt_hi = tt_hi;
band.d_tt = d_tt;
band.tan_lo = tan_lo;
band.tan_hi = tan_hi;
#ifndef JFJOCH_USE_CUDA
(void) allow_device;
#else
// With a GPU the two passes over the pixels run there, and only the cells come back.
std::unique_ptr<BeamCenterBackgroundGPU> gpu;
if (allow_device && get_gpu_count() > 0)
gpu = std::make_unique<BeamCenterBackgroundGPU>(W, H, block_row, usable.data(), mean.data());
#endif
float step_x = 0.0f, step_y = 0.0f, sigma_x = 0.0f, sigma_y = 0.0f;
float previous_x = 0.0f, previous_y = 0.0f;
int reversals = 0;
for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++) {
// The binning pass about the current centre, then the clip rounds over the pixels it binned.
const auto bin_cpu = [&] {
ParallelFor(BLOCKS, nthreads, [&](int b) {
double *b_sum = block_sum.data() + static_cast<size_t>(b) * n_cells;
double *b_sum_sq = block_sum_sq.data() + static_cast<size_t>(b) * n_cells;
@@ -188,48 +214,54 @@ FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const Pixe
const size_t i = static_cast<size_t>(y) * W + x;
if (!usable[i])
continue;
const float u = (x - beam_x) * pixel_size;
const float v = (y - beam_y) * pixel_size;
const float lx = rot[0] * u + rot[1] * v + rot[2] * distance;
const float ly = rot[3] * u + rot[4] * v + rot[5] * distance;
const float lz = rot[6] * u + rot[7] * v + rot[8] * distance;
const float rho_sq = lx * lx + ly * ly;
if (lz > 0.0f) {
const double lz_sq = static_cast<double>(lz) * lz;
if (rho_sq < tan_lo * tan_lo * lz_sq || rho_sq > tan_hi * tan_hi * lz_sq)
continue;
}
const float rho = std::sqrt(rho_sq);
const float two_theta = std::atan2(rho, lz);
if (two_theta < tt_lo || two_theta >= tt_hi || rho == 0.0f)
float jac_x, jac_y;
const int cell = BackgroundBandCell(band, x, y, beam_x, beam_y, jac_x, jac_y);
if (cell < 0)
continue;
const float phi = std::atan2(ly, lx);
// Both bins are clamped: a pixel one float ulp below the top of the band divides
// to exactly RADIAL_BINS, which is one cell past the end of every accumulator.
const int r_bin = std::clamp(static_cast<int>((two_theta - tt_lo) / d_tt), 0, RADIAL_BINS - 1);
const int s_bin = std::clamp(static_cast<int>((phi + PI) / (2 * PI) * SECTORS), 0, SECTORS - 1);
const int cell = r_bin * SECTORS + s_bin;
// d(2theta)/d(beam), through the lab coordinate: the detector coordinate depends
// on the centre only as (x - beam_x), so moving the centre is moving the pixel.
const float denominator = rho * rho + lz * lz;
const float g_x = lz * lx / (rho * denominator);
const float g_y = lz * ly / (rho * denominator);
const float g_z = -rho / denominator;
cells[n_band] = cell;
values[n_band] = mean[i];
n_band++;
b_count[cell]++;
b_sum[cell] += mean[i];
b_sum_sq[cell] += static_cast<double>(mean[i]) * mean[i];
b_jx[cell] += -pixel_size * (g_x * rot[0] + g_y * rot[3] + g_z * rot[6]);
b_jy[cell] += -pixel_size * (g_x * rot[1] + g_y * rot[4] + g_z * rot[7]);
b_jx[cell] += jac_x;
b_jy[cell] += jac_y;
}
}
band_pixels[b] = n_band;
});
fold(true);
};
const auto clip_cpu = [&] {
ParallelFor(BLOCKS, nthreads, [&](int b) {
double *b_sum = block_sum.data() + static_cast<size_t>(b) * n_cells;
double *b_sum_sq = block_sum_sq.data() + static_cast<size_t>(b) * n_cells;
int32_t *b_count = block_count.data() + static_cast<size_t>(b) * n_cells;
std::fill(b_sum, b_sum + n_cells, 0.0);
std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0);
std::fill(b_count, b_count + n_cells, 0);
const int32_t *cells = band_cell.data() + static_cast<size_t>(block_row[b]) * W;
const float *values = band_value.data() + static_cast<size_t>(block_row[b]) * W;
for (size_t j = 0; j < band_pixels[b]; j++) {
const int32_t c = cells[j];
const float value = values[j];
if (clip_limit[c] < 0.0f || value > clip_limit[c])
continue;
b_count[c]++;
b_sum[c] += value;
b_sum_sq[c] += static_cast<double>(value) * value;
}
});
fold(false);
};
for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++) {
#ifdef JFJOCH_USE_CUDA
if (gpu)
gpu->Bin(band, beam_x, beam_y, sum, sum_sq, sum_jx, sum_jy, count);
else
#endif
bin_cpu();
count_all = count; // the Jacobian sums belong to the unclipped pixel set
@@ -240,26 +272,12 @@ FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const Pixe
const double variance = std::max(sum_sq[c] / count[c] - m * m, 0.0);
clip_limit[c] = static_cast<float>(m + CLIP_SIGMA * std::sqrt(variance));
}
ParallelFor(BLOCKS, nthreads, [&](int b) {
double *b_sum = block_sum.data() + static_cast<size_t>(b) * n_cells;
double *b_sum_sq = block_sum_sq.data() + static_cast<size_t>(b) * n_cells;
int32_t *b_count = block_count.data() + static_cast<size_t>(b) * n_cells;
std::fill(b_sum, b_sum + n_cells, 0.0);
std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0);
std::fill(b_count, b_count + n_cells, 0);
const int32_t *cells = band_cell.data() + static_cast<size_t>(block_row[b]) * W;
const float *values = band_value.data() + static_cast<size_t>(block_row[b]) * W;
for (size_t j = 0; j < band_pixels[b]; j++) {
const int32_t c = cells[j];
const float value = values[j];
if (clip_limit[c] < 0.0f || value > clip_limit[c])
continue;
b_count[c]++;
b_sum[c] += value;
b_sum_sq[c] += static_cast<double>(value) * value;
}
});
fold(false);
#ifdef JFJOCH_USE_CUDA
if (gpu)
gpu->Clip(clip_limit, sum, sum_sq, count);
else
#endif
clip_cpu();
}
// Radial profile: the median over the sectors that have a mean, on rings that are
image_analysis/geom_refinement/BeamCenterFromBackground.h
+5 -1
View File
@@ -40,10 +40,14 @@ struct BeamCenterEstimate {
// `start` is where the walk begins; the centre in the file when it is not given. The walk advances
// by a bounded distance per iteration, so where it starts decides how much of its budget is spent
// travelling and - on a surface with more than one basin - which fixed point it can reach at all.
//
// With a GPU the passes over the pixels run on it (BeamCenterBackgroundGPU); allow_device = false
// keeps them on the host, which is what the parity test compares against.
std::optional<BeamCenterEstimate>
FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const PixelMask &mask,
const std::vector<float> &mean, size_t nthreads = 0,
std::optional<std::pair<float, float>> start = {});
std::optional<std::pair<float, float>> start = {},
bool allow_device = true);
// The precision of a centre that is the FFT capture alone, with no walk behind it. The capture is
// a half-pixel grid position read off a surface, measured over 75 rotation datasets at a median
image_analysis/geom_refinement/CMakeLists.txt
+13 -1
View File
@@ -24,6 +24,7 @@ ADD_LIBRARY(JFJochGeomRefinement STATIC
XtalOptimizer.cpp
XtalOptimizer.h
XtalResidual.h
BackgroundBand.h
PostRefine.cpp
PostRefine.h
GeometryRefiner.cpp
@@ -37,7 +38,8 @@ ADD_LIBRARY(JFJochGeomRefinement STATIC
TARGET_LINK_LIBRARIES(JFJochGeomRefinement Ceres::ceres Eigen3::Eigen JFJochCommon fftw3f)
IF (JFJOCH_CUDA_AVAILABLE)
TARGET_SOURCES(JFJochGeomRefinement PRIVATE BeamCenterFFTGPU.cu BeamCenterFFTGPU.h)
TARGET_SOURCES(JFJochGeomRefinement PRIVATE BeamCenterFFTGPU.cu BeamCenterFFTGPU.h
BeamCenterBackgroundGPU.cu BeamCenterBackgroundGPU.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)
@@ -46,3 +48,13 @@ IF (JFJOCH_CUDA_AVAILABLE)
TARGET_LINK_LIBRARIES(JFJochGeomRefinement CUDA::cufft)
ENDIF()
ENDIF()
# The background beam-centre fit bins every pixel with BackgroundBand.h on the host and on the device
# and must get the same bits on both (see there). That needs no multiply-add contracted on either side:
# GCC and Clang contract by default and nvcc does too, while MSVC does not without /fp:contract.
IF (JFJOCH_CUDA_AVAILABLE)
SET_SOURCE_FILES_PROPERTIES(BeamCenterBackgroundGPU.cu PROPERTIES COMPILE_OPTIONS "--fmad=false")
ENDIF()
IF (CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang")
SET_SOURCE_FILES_PROPERTIES(BeamCenterFromBackground.cpp PROPERTIES COMPILE_OPTIONS "-ffp-contract=off")
ENDIF()
tests/BeamCenterFromBackgroundTest.cpp
+26
View File
@@ -307,3 +307,29 @@ TEST_CASE("FindBeamCenter_BlanksTheBeamStopOutOfTheCapture", "[BeamCenter]") {
CHECK(masked_error < 12.0f);
CHECK(masked_error <= unmasked_error);
}
#ifdef JFJOCH_USE_CUDA
#include "../common/CUDAWrapper.h"
// With a GPU the passes over the pixels run on it. The cell a pixel lands in is computed from IEEE
// operations alone and without contraction on both sides, and each cell is summed in the host's order,
// so the walk ends on the same bits - a tilted detector and an offset centre included.
TEST_CASE("BeamCenterFromBackground_DeviceMatchesHost", "[BeamCenter]") {
if (get_gpu_count() == 0)
SKIP("no GPU");
DiffractionExperiment x = TestExperiment();
x.PoniRot1_rad(0.005f).PoniRot2_rad(-0.003f);
PixelMask pixel_mask(x);
const DiffractionGeometry geom_true = OffsetBy(x.GetDiffractionGeometry(), 7.0f, -4.5f);
const auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 0.5f);
const auto host = FindBeamCenterFromBackground(x, pixel_mask, projection, 0, {}, /*allow_device=*/false);
const auto device = FindBeamCenterFromBackground(x, pixel_mask, projection, 0, {}, /*allow_device=*/true);
REQUIRE(host.has_value());
REQUIRE(device.has_value());
CHECK(device->beam_x_pxl == host->beam_x_pxl);
CHECK(device->beam_y_pxl == host->beam_y_pxl);
CHECK(device->sigma_pxl == host->sigma_pxl);
}
#endif
tests/ShadowFinderTest.cpp
+88
View File
@@ -471,3 +471,91 @@ TEST_CASE("ShadowFinder_ABrightRingIsNotABeamStop", "[ShadowFinder]") {
// Anything much beyond the stop, its arm and their penumbra means the walk ran away.
CHECK(std::count(mask.begin(), mask.end(), 1u) < 6000);
}
#ifdef JFJOCH_USE_CUDA
#include "../common/CUDAWrapper.h"
#include "../compression/JFJochCompressor.h"
namespace {
// The mask and the mean projection of the same frames, once from the host projection (the frames
// handed over uncompressed) and once from the device's (handed over as bitshuffle+LZ4, which is
// what sends them to the GPU).
struct HostAndDevice {
std::vector<uint32_t> host_mask, device_mask;
std::vector<float> host_mean, device_mean;
};
HostAndDevice MaskBothWays(const DiffractionExperiment &x, const std::vector<std::vector<int32_t>> &frames) {
const PixelMask pixel_mask(x);
HostAndDevice out;
std::vector<uint8_t> buffer;
{
ShadowFinder finder(x, pixel_mask);
for (const auto &frame : frames) {
DataMessage msg{};
msg.image = CompressedImage(frame, W, H);
finder.AddImage(msg, buffer);
}
out.host_mask = finder.GetMask();
out.host_mean = finder.GetMeanProjection();
}
{
ShadowFinder finder(x, pixel_mask);
JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
std::vector<std::vector<uint8_t>> compressed;
for (const auto &frame : frames) {
compressed.push_back(compressor.Compress(frame));
DataMessage msg{};
msg.image = CompressedImage(compressed.back().data(), compressed.back().size(), W, H,
CompressedImageMode::Int32, CompressionAlgorithm::BSHUF_LZ4);
finder.AddImage(msg, buffer);
}
out.device_mask = finder.GetMask();
out.device_mean = finder.GetMeanProjection();
}
return out;
}
}
// The mask is made on the GPU wherever the projection is there. On scenes where no pixel sits within a
// rounding of a threshold the two must agree to the pixel: every step but the polarization factor,
// the Poisson test's logarithm and the arm search's azimuth is exact on both, and these scenes are
// built so that none of the three decides anything at an edge.
TEST_CASE("ShadowFinder_DeviceMaskMatchesHost", "[ShadowFinder]") {
if (get_gpu_count() == 0)
SKIP("no GPU");
SECTION("a beam stop with a reflection behind it") {
std::vector<std::vector<int32_t>> frames;
for (int f = 0; f < NFRAMES; f++)
frames.push_back(Scene(false, f == 0));
const auto r = MaskBothWays(TestExperiment(), frames);
CHECK(std::count(r.host_mask.begin(), r.host_mask.end(), 1u) == 2612);
CHECK(r.device_mask == r.host_mask);
CHECK(std::memcmp(r.device_mean.data(), r.host_mean.data(), r.host_mean.size() * sizeof(float)) == 0);
}
SECTION("an arm that lets part of the beam through, across a module gap") {
constexpr int32_t BRIGHT = 50;
constexpr int ARM_HALF_WIDE = 15, GAP_X0 = 200, GAP_X1 = 216, OPAQUE_FROM_X = 232;
std::vector<std::vector<int32_t>> frames;
for (int f = 0; f < NFRAMES; f++) {
frames.emplace_back(static_cast<size_t>(W) * H, BRIGHT);
auto &frame = frames.back();
for (int y = 0; y < H; y++)
for (int xi = 0; xi < W; xi++) {
const int dx = xi - C, dy = y - C;
if (dx * dx + dy * dy <= STOP_R * STOP_R)
frame[I(xi, y)] = 0;
else if (dx >= 0 && std::abs(dy) <= ARM_HALF_WIDE)
frame[I(xi, y)] = xi >= OPAQUE_FROM_X ? 0 : BRIGHT * 6 / 10;
if (xi >= GAP_X0 && xi <= GAP_X1)
frame[I(xi, y)] = INT32_MIN;
}
}
const auto r = MaskBothWays(TestExperiment(), frames);
CHECK(std::count(r.host_mask.begin(), r.host_mask.end(), ShadowFinder::TRANSMITTING) > 0);
CHECK(r.device_mask == r.host_mask);
}
}
#endif