Nothing kept a neighbour's flux out of a reflection's own signal disk. The union mask keeps neighbour cores out of the BACKGROUND ring, but the r1 disk was read whole, so on a dense pattern a crowded reflection measures part of its neighbour as its own. Ownership is decided once per image into a per-pixel (quantised distance, reflection) key written with an atomic minimum, so the nearest predicted centre wins whatever order the writes arrive in and the lowest index breaks a tie. `--overlap exclude`, now the default, drops the pixels a nearer neighbour owns from the profile fit. A profile fit is the amplitude of a normalised profile, so leaving pixels out renormalises the estimator by construction and the reflection stays unbiased rather than being discarded; the summation-fallback guard is scaled back to the disk the box-sum seed actually read, so it still compares like with like. `--overlap reject` is the XDS MINPK alternative - drop the reflection when less than `--overlap-minpk` of its expected profile is cleanly its own. A box sum has no profile to renormalise with, so `exclude` is a no-op there and only `reject` acts on it. Widening the split - keeping a pixel only where no other centre is within its distance PLUS a margin - was built and measured, and it is worse monotonically: the residual bias of the pixels that were kept grows from +0.072 to +0.209 in ln intensity at 0 to 3 px of margin. What the margin removes is the reflection's own profile, not the neighbour's tail, so the plain nearest-centre split is the rule. Measured on the full 38-crystal rotation battery against the same binary with the treatment off: ISa better 15 / worse 8, summed shortfall against XDS 39.7 -> 28.1. Three of the losses are the two-pass loop taking its other branch - their median mosaicity moves between the two known attractors - rather than the change under test; excluding those it is better 15 / worse 5 and the shortfall goes 31.3 -> 14.4. The two crowded crystals gain 38% and 52% of their ISa, one of them passing XDS. High-shell CC1/2 over the 35 crystals that neither flipped branch nor carry a collapsed error model is better 7 / worse 7. Space groups unchanged at 35/38. The owner map is built only when a treatment is asked for and costs 1.1% of the battery's wall clock - 23% on a genuinely crowded crystal, nothing where no two predictions touch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
796 lines
41 KiB
Plaintext
796 lines
41 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "BraggIntegrationEngineGPU.h"
|
|
|
|
using namespace bragg_engine;
|
|
|
|
namespace {
|
|
|
|
inline void cuda_err(cudaError_t val) {
|
|
if (val != cudaSuccess)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
|
|
}
|
|
|
|
// Fixed scalars passed by value to every kernel (mirrors BraggIntegrationEngine's members).
|
|
struct BraggGpuParams {
|
|
int W, H;
|
|
float r1_sq, r2, r2_sq, r3, r3_sq;
|
|
int R, G, GG;
|
|
float bkg_clip_nsigma; // high-side background sigma-clip multiplier (0 = no clip)
|
|
int empirical; // ProfileEmpirical vs ProfileGaussian
|
|
int use_ellipse;
|
|
float bw_sigma;
|
|
float c_radial;
|
|
float F_px;
|
|
float beam_x, beam_y;
|
|
float bkg_trim; // idea 1: symmetric trimmed-mean background fraction (0 = plain ring mean)
|
|
BraggStencilParams stencil; // per-reflection signal/background geometry (BraggStencil.h)
|
|
int rad_w; // radial-background window held in shared memory, in bins of one pixel
|
|
int n_kern; // radial-background kernels in the table
|
|
int overlap; // 0 = off, 1 = reject, 2 = exclude (OverlapMode)
|
|
int boxsum_reject; // BoxSum mode under Reject: drop on the disk-AREA fraction (see the CPU engine)
|
|
int exclude_px; // drop a neighbour's pixels from the disk: Exclude, and not a box sum
|
|
float claim_sq, inv_claim; // how far a reflection claims pixels in the owner map
|
|
float minpk; // Reject: least clean profile fraction that is kept
|
|
};
|
|
|
|
__device__ inline bool valid(int32_t v) { return v != INT32_MIN && v != INT32_MAX; }
|
|
|
|
// Learned second moments, (sum v*rad^2, sum v*tan^2, sum v) per shell plus one global slot at N_SHELL.
|
|
constexpr int MOM_STRIDE = 3;
|
|
|
|
// --- Mark the r2 signal region of every predicted reflection (race-free: all writes are 1). The
|
|
// region is the INNER stencil ellipse in the neighbour's own frame; see the CPU engine.
|
|
// The same sweep fills the owner map when an overlap treatment is on: a pixel two signal regions
|
|
// share belongs to the nearer centre, which an atomicMin over (distance, index) settles without
|
|
// any ordering or sorting (BraggStencil.h). The claim disk sits inside the ellipse this loop
|
|
// already walks, so ownership costs one atomic on the pixels already visited. ---
|
|
__global__ void mark_mask(const float *px_x, const float *px_y, uint8_t *mask, uint32_t *owner,
|
|
BraggGpuParams p, int n) {
|
|
const int i = blockIdx.x;
|
|
if (i >= n) return;
|
|
const float cx = px_x[i], cy = px_y[i];
|
|
const BraggStencil st = MakeBraggStencil(cx, cy, p.stencil);
|
|
const int x0 = max(0, (int) floorf(cx - st.ex_in - 1.0f));
|
|
const int x1 = min(p.W - 1, (int) ceilf(cx + st.ex_in + 1.0f));
|
|
const int y0 = max(0, (int) floorf(cy - st.ey_in - 1.0f));
|
|
const int y1 = min(p.H - 1, (int) ceilf(cy + st.ey_in + 1.0f));
|
|
const int bw = x1 - x0 + 1, bh = y1 - y0 + 1;
|
|
if (bw <= 0 || bh <= 0) return;
|
|
for (int t = threadIdx.x; t < bw * bh; t += blockDim.x) {
|
|
const int x = x0 + t % bw, y = y0 + t / bw;
|
|
const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
|
|
if (d.inner < p.r2_sq) mask[y * p.W + x] = 1;
|
|
if (p.overlap && d.signal < p.claim_sq)
|
|
atomicMin(&owner[y * p.W + x], BraggOwnerKey(sqrtf(d.signal), p.inv_claim, i));
|
|
}
|
|
}
|
|
|
|
// --- Pass A box-sum: rough I / background / centroid / strong flag, one block per reflection. ---
|
|
__global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
|
|
const int32_t *img, const uint8_t *mask, const uint32_t *owner,
|
|
BraggGpuParams p, int n,
|
|
int *cx_o, int *cy_o, float *I_o, float *sigma_o, float *bkg_o,
|
|
float *bkgvar_o, float *varbkg_o, float *obsx_o, float *obsy_o, uint8_t *ok_o, uint8_t *strong_o,
|
|
uint8_t *hasobs_o, unsigned long long *invd2mm,
|
|
float *isum_o, int *ninner_o, int *rbin_o, int *kbin_o,
|
|
float *rad_sum, int *rad_cnt, int n_rad) {
|
|
const int i = blockIdx.x;
|
|
if (i >= n) return;
|
|
|
|
__shared__ unsigned long long s_Isum, s_Ix, s_Iy;
|
|
__shared__ int s_ninner, s_ninner_valid, s_nbkg, s_ndisk, s_nown;
|
|
__shared__ double s_bkgsum;
|
|
__shared__ int s_accept;
|
|
__shared__ double s_bkg, s_thr, s_clipsum;
|
|
__shared__ int s_clipn;
|
|
__shared__ float s_r0;
|
|
// The stencil spans r0 +- the radial semi-axis of the outer ellipse, so the window has to be
|
|
// sized from the widest aperture on the detector - which the host does, into p.rad_w. A fixed
|
|
// 32-bin window silently dropped the outer bins as soon as anything was elongated.
|
|
extern __shared__ float s_rad[];
|
|
float *s_radv = s_rad;
|
|
int *s_radn = (int *) (s_rad + p.rad_w);
|
|
__shared__ int s_radbase;
|
|
if (threadIdx.x == 0) {
|
|
const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y;
|
|
s_r0 = sqrtf(rx * rx + ry * ry);
|
|
s_Isum = 0; s_Ix = 0; s_Iy = 0;
|
|
s_ninner = 0; s_ninner_valid = 0; s_nbkg = 0; s_bkgsum = 0.0;
|
|
s_ndisk = 0; s_nown = 0;
|
|
}
|
|
__syncthreads();
|
|
|
|
const float cx = px_x[i], cy = px_y[i];
|
|
// Every thread builds the same stencil from the same inputs; it is a few flops against a whole
|
|
// bounding box of pixel reads, so it costs less than staging it through shared memory.
|
|
const BraggStencil st = MakeBraggStencil(cx, cy, p.stencil);
|
|
const int x0 = max(0, (int) floorf(cx - st.ex_out - 1.0f));
|
|
const int x1 = min(p.W - 1, (int) ceilf(cx + st.ex_out + 1.0f));
|
|
const int y0 = max(0, (int) floorf(cy - st.ey_out - 1.0f));
|
|
const int y1 = min(p.H - 1, (int) ceilf(cy + st.ey_out + 1.0f));
|
|
const int bw = x1 - x0 + 1, bh = y1 - y0 + 1;
|
|
const int area = (bw > 0 && bh > 0) ? bw * bh : 0;
|
|
|
|
long long l_Isum = 0, l_Ix = 0, l_Iy = 0;
|
|
int l_ni = 0, l_niv = 0, l_nb = 0, l_nd = 0, l_no = 0;
|
|
double l_bkg = 0.0;
|
|
for (int t = threadIdx.x; t < area; t += blockDim.x) {
|
|
const int x = x0 + t % bw, y = y0 + t / bw;
|
|
const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
|
|
const int32_t px = img[y * p.W + x];
|
|
if (d.signal < p.r1_sq) {
|
|
// A pixel a nearer neighbour owns carries that neighbour's flux; see the CPU engine.
|
|
++l_nd;
|
|
if (p.overlap) {
|
|
if (BraggOwnedBy(owner[y * p.W + x], i)) ++l_no;
|
|
else if (p.exclude_px) continue;
|
|
}
|
|
++l_ni;
|
|
if (valid(px)) { l_Isum += px; l_Ix += (long long) x * px; l_Iy += (long long) y * px; ++l_niv; }
|
|
} else if (d.inner >= p.r2_sq && d.outer < p.r3_sq) {
|
|
if (mask[y * p.W + x]) continue;
|
|
if (!valid(px)) continue;
|
|
l_bkg += (double) px; ++l_nb;
|
|
}
|
|
}
|
|
atomicAdd(&s_Isum, (unsigned long long) l_Isum);
|
|
atomicAdd(&s_Ix, (unsigned long long) l_Ix);
|
|
atomicAdd(&s_Iy, (unsigned long long) l_Iy);
|
|
atomicAdd(&s_ninner, l_ni);
|
|
atomicAdd(&s_ninner_valid, l_niv);
|
|
atomicAdd(&s_nbkg, l_nb);
|
|
atomicAdd(&s_bkgsum, l_bkg);
|
|
atomicAdd(&s_ndisk, l_nd);
|
|
atomicAdd(&s_nown, l_no);
|
|
__syncthreads();
|
|
|
|
for (int t = threadIdx.x; t < p.rad_w; t += blockDim.x) { s_radv[t] = 0.0f; s_radn[t] = 0; }
|
|
if (threadIdx.x == 0) s_radbase = (int) lroundf(s_r0) - p.rad_w / 2;
|
|
if (threadIdx.x == 0) {
|
|
s_accept = (s_ninner_valid == s_ninner && s_nbkg > 5) ? 1 : 0;
|
|
s_bkg = s_accept ? (s_bkgsum / (double) s_nbkg) : 0.0;
|
|
s_thr = s_bkg + (double) p.bkg_clip_nsigma * sqrt(fmax(s_bkg, 1.0));
|
|
s_clipsum = 0.0; s_clipn = 0;
|
|
}
|
|
__syncthreads();
|
|
|
|
// Trimmed-mean background (idea 1): collect the annulus into shared memory, bitonic-sort the block, and
|
|
// average the middle (1 - 2*bkg_trim) fraction - robust to the high-side contamination that biases the
|
|
// plain ring mean. Falls back to the flat mean when the ring exceeds the shared buffer.
|
|
__shared__ int s_bvals[BKG_TRIM_MAX];
|
|
__shared__ int s_bn;
|
|
if (threadIdx.x == 0) s_bn = 0;
|
|
__syncthreads();
|
|
const bool do_trim = s_accept && p.bkg_trim > 0.0f && s_nbkg > 5 && s_nbkg <= BKG_TRIM_MAX;
|
|
if (do_trim) {
|
|
for (int t = threadIdx.x; t < area; t += blockDim.x) {
|
|
const int x = x0 + t % bw, y = y0 + t / bw;
|
|
const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
|
|
if (!(d.inner >= p.r2_sq && d.outer < p.r3_sq)) continue;
|
|
if (mask[y * p.W + x]) continue;
|
|
const int32_t px = img[y * p.W + x];
|
|
if (!valid(px)) continue;
|
|
const int slot = atomicAdd(&s_bn, 1);
|
|
if (slot < BKG_TRIM_MAX) s_bvals[slot] = px;
|
|
}
|
|
__syncthreads();
|
|
const int nb = min(s_bn, BKG_TRIM_MAX);
|
|
int n2 = 1;
|
|
while (n2 < nb) n2 <<= 1;
|
|
for (int t = threadIdx.x + nb; t < n2; t += blockDim.x) s_bvals[t] = INT32_MAX; // pad to power of 2
|
|
__syncthreads();
|
|
for (int k = 2; k <= n2; k <<= 1) // ascending bitonic sort of s_bvals[0..n2)
|
|
for (int j = k >> 1; j > 0; j >>= 1) {
|
|
for (int idx = threadIdx.x; idx < n2; idx += blockDim.x) {
|
|
const int ixj = idx ^ j;
|
|
if (ixj > idx) {
|
|
const bool up = ((idx & k) == 0);
|
|
const int a = s_bvals[idx], b = s_bvals[ixj];
|
|
if ((up && a > b) || (!up && a < b)) { s_bvals[idx] = b; s_bvals[ixj] = a; }
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
if (threadIdx.x == 0) {
|
|
const int lo = (int) (nb * p.bkg_trim), hi = nb - lo;
|
|
if (hi > lo) {
|
|
double s = 0.0;
|
|
for (int t = lo; t < hi; ++t) s += (double) s_bvals[t];
|
|
s_bkg = s / (double) (hi - lo);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
// Second ring pass for the high-side sigma-clip (re-reads the annulus; avoids storing bkg values).
|
|
if (s_accept && p.bkg_clip_nsigma > 0.0f && !do_trim) {
|
|
double c_l = 0.0; int cn_l = 0;
|
|
for (int t = threadIdx.x; t < area; t += blockDim.x) {
|
|
const int x = x0 + t % bw, y = y0 + t / bw;
|
|
const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
|
|
if (!(d.inner >= p.r2_sq && d.outer < p.r3_sq)) continue;
|
|
if (mask[y * p.W + x]) continue;
|
|
const int32_t px = img[y * p.W + x];
|
|
if (!valid(px)) continue;
|
|
if ((double) px <= s_thr) {
|
|
c_l += px; ++cn_l;
|
|
if (n_rad > 0) {
|
|
// The radial curve is binned on the TRUE detector radius, so this is the
|
|
// unshrunk radial projection - no per-pixel sqrt.
|
|
const int idx = (int) lroundf(s_r0 + d.rad) - s_radbase;
|
|
if (idx >= 0 && idx < p.rad_w) {
|
|
atomicAdd(&s_radv[idx], (float) px); // shared, not global
|
|
atomicAdd(&s_radn[idx], 1);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
atomicAdd(&s_clipsum, c_l); atomicAdd(&s_clipn, cn_l);
|
|
}
|
|
__syncthreads();
|
|
|
|
if (n_rad > 0) {
|
|
for (int t = threadIdx.x; t < p.rad_w; t += blockDim.x) {
|
|
if (s_radn[t] == 0) continue;
|
|
const int b = min(max(s_radbase + t, 0), n_rad - 1);
|
|
atomicAdd(&rad_sum[b], s_radv[t]);
|
|
atomicAdd(&rad_cnt[b], s_radn[t]);
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
if (threadIdx.x != 0) return;
|
|
if (!s_accept) { ok_o[i] = 0; strong_o[i] = 0; hasobs_o[i] = 0; return; }
|
|
// A box sum has no profile to renormalise a disk it has taken pixels out of, so dropping the
|
|
// reflection is the only overlap treatment it has (see the CPU engine).
|
|
if (p.boxsum_reject && (float) s_nown < p.minpk * (float) s_ndisk) {
|
|
ok_o[i] = 0; strong_o[i] = 0; hasobs_o[i] = 0; return;
|
|
}
|
|
|
|
double bkg = s_bkg;
|
|
int n_bkg_used = s_nbkg; // pixels behind the FINAL background value (trim/clip shrink it)
|
|
if (do_trim) {
|
|
const int nb = min(s_bn, BKG_TRIM_MAX);
|
|
const int lo = (int) (nb * p.bkg_trim), hi = nb - lo;
|
|
if (hi > lo) n_bkg_used = hi - lo;
|
|
}
|
|
if (p.bkg_clip_nsigma > 0.0f && s_clipn > 5) { bkg = s_clipsum / (double) s_clipn; n_bkg_used = s_clipn; }
|
|
const long long Isum = (long long) s_Isum;
|
|
const double I = (double) Isum - (double) s_ninner * bkg;
|
|
// See the CPU engine: bkg is estimated from n_bkg_used pixels and subtracted n_inner times, so
|
|
// var(I) = Isum + n_inner^2 * bkg/n_bkg_used. Both engines must agree.
|
|
const double bkg_var = bkg / (double) n_bkg_used;
|
|
const double var_bkg_term = (double) s_ninner * (double) s_ninner * bkg_var;
|
|
double sigma = 1.0;
|
|
uint8_t hasobs = 0; double ox = 0.0, oy = 0.0;
|
|
if (Isum > 0) {
|
|
sigma = fmax(sigma, sqrt((double) Isum + var_bkg_term));
|
|
ox = (double) (long long) s_Ix / (double) Isum;
|
|
oy = (double) (long long) s_Iy / (double) Isum;
|
|
hasobs = 1;
|
|
}
|
|
cx_o[i] = (int) lroundf(cx);
|
|
cy_o[i] = (int) lroundf(cy);
|
|
I_o[i] = (float) I; sigma_o[i] = (float) sigma; bkg_o[i] = (float) bkg;
|
|
bkgvar_o[i] = (float) bkg_var;
|
|
varbkg_o[i] = (float) ((double) s_ninner * bkg + var_bkg_term);
|
|
isum_o[i] = (float) Isum;
|
|
ninner_o[i] = s_ninner;
|
|
rbin_o[i] = min(max((int) lroundf(s_r0), 0), n_rad > 0 ? n_rad - 1 : 0);
|
|
kbin_o[i] = BraggStencilKernelIndex(st, p.n_kern);
|
|
obsx_o[i] = (float) ox; obsy_o[i] = (float) oy; hasobs_o[i] = hasobs;
|
|
ok_o[i] = 1;
|
|
strong_o[i] = (sigma > 0.0 && I / sigma >= STRONG_I_OVER_SIGMA) ? 1 : 0;
|
|
|
|
const float d = dd[i];
|
|
if (d > 0.0f) {
|
|
// Positive doubles keep IEEE bit-pattern ordering, so atomicMin/Max on the ull view works.
|
|
const unsigned long long b = (unsigned long long) __double_as_longlong(1.0 / ((double) d * d));
|
|
atomicMin(&invd2mm[0], b);
|
|
atomicMax(&invd2mm[1], b);
|
|
}
|
|
}
|
|
|
|
// Resolution shell of one reflection from the global inv-d^2 range (mirrors CPU shell_of). Computed
|
|
// inline in learn_profile and fit so no separate shell array/kernel is needed.
|
|
__device__ inline int compute_shell(float d, const unsigned long long *invd2mm) {
|
|
const unsigned long long mn = invd2mm[0], mx = invd2mm[1];
|
|
if (!(d > 0.0f) || mx <= mn) return 0;
|
|
const double invd2 = 1.0 / ((double) d * d);
|
|
const double dmn = __longlong_as_double((long long) mn), dmx = __longlong_as_double((long long) mx);
|
|
const int s = (int) ((invd2 - dmn) / (dmx - dmn) * N_SHELL);
|
|
return s < 0 ? 0 : (s >= N_SHELL ? N_SHELL - 1 : s);
|
|
}
|
|
|
|
// --- Zero the profile accumulators and seed the inv-d^2 range, in one launch (replaces a handful
|
|
// of small cudaMemsetAsync calls, which matter when kernel-launch latency is high). ---
|
|
__global__ void reset(float *shell_grid, float *global_grid, float *mom, int *shell_n, int *global_n,
|
|
unsigned long long *invd2mm, int GG) {
|
|
for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < N_SHELL * GG; k += blockDim.x * gridDim.x)
|
|
shell_grid[k] = 0.0f;
|
|
for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < GG; k += blockDim.x * gridDim.x)
|
|
global_grid[k] = 0.0f;
|
|
if (blockIdx.x == 0 && threadIdx.x < MOM_STRIDE * (N_SHELL + 1)) mom[threadIdx.x] = 0.0f;
|
|
if (blockIdx.x == 0 && threadIdx.x < N_SHELL) shell_n[threadIdx.x] = 0;
|
|
if (blockIdx.x == 0 && threadIdx.x == 0) {
|
|
*global_n = 0;
|
|
invd2mm[0] = ~0ull; // min seed
|
|
invd2mm[1] = 0ull; // max seed
|
|
}
|
|
}
|
|
|
|
// --- Learn the profile: each strong spot adds its bkg-subtracted, I-normalised grid to its shell
|
|
// (and the global grid), and its second moments to the same shell's moment slot. The grid stays
|
|
// in the DETECTOR frame (that is where the empirical profile is applied); the moments are taken in
|
|
// the spot's own radial/tangential frame, because a detector-frame stack is azimuthally averaged
|
|
// and cannot tell a radially smeared spot from a tangentially wide one (see the CPU engine).
|
|
// One block per reflection. ---
|
|
__global__ void learn_profile(const int32_t *img, const uint32_t *owner,
|
|
const float *px_x, const float *px_y,
|
|
const int *cx_a, const int *cy_a, const float *dd,
|
|
const unsigned long long *invd2mm,
|
|
const float *I_a, const float *bkg_a, const uint8_t *ok_a, const uint8_t *strong_a,
|
|
float *shell_grid, float *global_grid, float *mom, int *shell_n, int *global_n,
|
|
BraggGpuParams p, int n) {
|
|
const int i = blockIdx.x;
|
|
if (i >= n || !ok_a[i] || !strong_a[i]) return;
|
|
const float I = I_a[i];
|
|
if (!(I > 0.0f)) return;
|
|
const int cx = cx_a[i], cy = cy_a[i], sh = compute_shell(dd[i], invd2mm);
|
|
const float bkg = bkg_a[i];
|
|
const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y;
|
|
const float Rpx = sqrtf(rx * rx + ry * ry);
|
|
const float ux = Rpx > 1e-6f ? rx / Rpx : 1.0f, uy = Rpx > 1e-6f ? ry / Rpx : 0.0f;
|
|
float *sg = shell_grid + (size_t) sh * p.GG;
|
|
__shared__ float s_rad, s_tan, s_w;
|
|
if (threadIdx.x == 0) { s_rad = 0.0f; s_tan = 0.0f; s_w = 0.0f; }
|
|
__syncthreads();
|
|
float l_rad = 0.0f, l_tan = 0.0f, l_w = 0.0f;
|
|
for (int k = threadIdx.x; k < p.GG; k += blockDim.x) {
|
|
const int dx = k % p.G - p.R, dy = k / p.G - p.R;
|
|
const int x = cx + dx, y = cy + dy;
|
|
if (x < 0 || y < 0 || x >= p.W || y >= p.H) continue;
|
|
const int32_t px = img[y * p.W + x];
|
|
if (!valid(px)) continue;
|
|
if (p.overlap == 2 && !BraggOwnedBy(owner[y * p.W + x], i)) continue;
|
|
const float v = ((float) px - bkg) / I;
|
|
atomicAdd(&sg[k], v);
|
|
atomicAdd(&global_grid[k], v);
|
|
if ((float) (dx * dx + dy * dy) >= p.r1_sq) continue;
|
|
const float rad = dx * ux + dy * uy, tn = -dx * uy + dy * ux;
|
|
l_rad += v * rad * rad; l_tan += v * tn * tn; l_w += v;
|
|
}
|
|
// Reduce in shared memory first: every block adds to the same handful of global words, so one
|
|
// atomic per block instead of one per thread.
|
|
atomicAdd(&s_rad, l_rad); atomicAdd(&s_tan, l_tan); atomicAdd(&s_w, l_w);
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) {
|
|
atomicAdd(&mom[MOM_STRIDE * sh + 0], s_rad);
|
|
atomicAdd(&mom[MOM_STRIDE * sh + 1], s_tan);
|
|
atomicAdd(&mom[MOM_STRIDE * sh + 2], s_w);
|
|
atomicAdd(&mom[MOM_STRIDE * N_SHELL + 0], s_rad);
|
|
atomicAdd(&mom[MOM_STRIDE * N_SHELL + 1], s_tan);
|
|
atomicAdd(&mom[MOM_STRIDE * N_SHELL + 2], s_w);
|
|
atomicAdd(&shell_n[sh], 1);
|
|
atomicAdd(global_n, 1);
|
|
}
|
|
}
|
|
|
|
// --- Turn the learned moments into radial/tangential variances and, for empirical, normalise the
|
|
// learned grid into a profile. One block per grid: blocks [0,N_SHELL) are the shells, block
|
|
// N_SHELL is the global one. ---
|
|
__global__ void build_profiles(const float *shell_grid, const float *global_grid, const float *mom,
|
|
const int *shell_n, const int *global_n,
|
|
float *shell_P, float *global_P, float *sigma2_r, float *sigma2_t,
|
|
BraggGpuParams p) {
|
|
const int b = blockIdx.x;
|
|
const float *grid; int nstrong; float *P;
|
|
if (b < N_SHELL) { grid = shell_grid + (size_t) b * p.GG; nstrong = shell_n[b]; P = shell_P + (size_t) b * p.GG; }
|
|
else { grid = global_grid; nstrong = *global_n; P = global_P; }
|
|
|
|
if (threadIdx.x == 0) {
|
|
const float w = mom[MOM_STRIDE * b + 2];
|
|
sigma2_r[b] = w > 0.0f ? fmaxf(0.25f, mom[MOM_STRIDE * b + 0] / w) : 1.0f;
|
|
sigma2_t[b] = w > 0.0f ? fmaxf(0.25f, mom[MOM_STRIDE * b + 1] / w) : 1.0f;
|
|
}
|
|
if (!p.empirical) return;
|
|
|
|
__shared__ float s_sum;
|
|
if (threadIdx.x == 0) s_sum = 0.0f;
|
|
__syncthreads();
|
|
float l_sum = 0.0f;
|
|
for (int k = threadIdx.x; k < p.GG; k += blockDim.x) {
|
|
const float g = fmaxf(0.0f, grid[k]); // a profile has to be non-negative
|
|
P[k] = g;
|
|
l_sum += g;
|
|
}
|
|
atomicAdd(&s_sum, l_sum);
|
|
__syncthreads();
|
|
const bool normalise = nstrong > 0 && s_sum > 0.0f;
|
|
for (int k = threadIdx.x; k < p.GG; k += blockDim.x) P[k] = normalise ? P[k] / s_sum : 0.0f;
|
|
}
|
|
|
|
// --- Radial background curvature correction: one thread per reflection, no pixel reads.
|
|
// The disk and the annulus are concentric, so a background linear in position cancels between
|
|
// them; what is left is the curvature of the radial background. k_diff (annulus minus disk
|
|
// histogram over radial offset) turns that into one short dot product. An empty radial bin
|
|
// contributes the reflection's own background, so an empty neighbourhood gives exactly zero
|
|
// correction because the kernel weights sum to zero. The kernel depends on how far this
|
|
// reflection's ring was elongated, so k_diff is a table and kbin_a picks the row.
|
|
// Mirrors BraggIntegrationEngineCPU. ---
|
|
__global__ void radial_correct(const float *rad_sum, const int *rad_cnt, int n_rad,
|
|
const float *k_diff, int k_len, int k_off,
|
|
const float *isum_a, const int *ninner_a, const int *rbin_a,
|
|
const int *kbin_a,
|
|
const uint8_t *ok_a, float *bkg_o, float *I_o, int n) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n || !ok_a[i]) return;
|
|
const float bkg = bkg_o[i];
|
|
const float *kern = k_diff + (size_t) kbin_a[i] * k_len;
|
|
float corr = 0.0f;
|
|
for (int k = 0; k < k_len; ++k) {
|
|
int b = rbin_a[i] + k - k_off;
|
|
b = min(max(b, 0), n_rad - 1);
|
|
const float v = rad_cnt[b] > 0 ? rad_sum[b] / (float) rad_cnt[b] : bkg;
|
|
corr += kern[k] * v;
|
|
}
|
|
const float bkg_new = bkg - corr;
|
|
bkg_o[i] = bkg_new;
|
|
I_o[i] = isum_a[i] - (float) ninner_a[i] * bkg_new;
|
|
}
|
|
|
|
// --- Pass B Kabsch profile fit: I = sum P(c-B)/v over sum P^2/v, v = B + max(I,0)P (iterate).
|
|
// One block per reflection; the (possibly elongated) profile is built in shared memory. ---
|
|
__global__ void fit(const int32_t *img, const uint32_t *owner, const float *px_x, const float *px_y,
|
|
const int *cx_a, const int *cy_a, const float *dd, const unsigned long long *invd2mm,
|
|
const float *I_seed, const float *sigma_seed, const float *bkg_a,
|
|
const float *bkgvar_a, const float *varbkg_seed, const uint8_t *ok_a,
|
|
const float *shell_P, const float *global_P,
|
|
const float *sigma2_r, const float *sigma2_t, const int *shell_n,
|
|
float *I_o, float *sigma_o, float *varbkg_o, uint8_t *ok_o, BraggGpuParams p, int n) {
|
|
const int i = blockIdx.x;
|
|
if (i >= n) return;
|
|
extern __shared__ float Pbuf[];
|
|
__shared__ float s_gs, s_num, s_den, s_I, s_wsum;
|
|
__shared__ float s_pown, s_mall, s_mown;
|
|
__shared__ int s_Rf, s_Gf;
|
|
|
|
if (!ok_a[i]) { if (threadIdx.x == 0) ok_o[i] = 0; return; }
|
|
|
|
const int cx = cx_a[i], cy = cy_a[i];
|
|
const int sh = compute_shell(dd[i], invd2mm);
|
|
const bool use_shell = shell_n[sh] >= MIN_STRONG_PER_SHELL; // else fall back to the global profile
|
|
const float bkg = bkg_a[i];
|
|
|
|
if (p.empirical) {
|
|
const float *Psrc = use_shell ? (shell_P + (size_t) sh * p.GG) : global_P;
|
|
if (threadIdx.x == 0) { s_Rf = p.R; s_Gf = p.G; }
|
|
__syncthreads();
|
|
for (int k = threadIdx.x; k < p.GG; k += blockDim.x) Pbuf[k] = Psrc[k];
|
|
__syncthreads();
|
|
} else {
|
|
const int si = use_shell ? sh : N_SHELL; // N_SHELL is the global slot
|
|
const float s2t = sigma2_t[si];
|
|
const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y;
|
|
const float Rpx = sqrtf(rx * rx + ry * ry);
|
|
const float tan2t = Rpx / p.F_px;
|
|
float s2r = s2t, ux = 1.0f, uy = 0.0f;
|
|
bool elong = false;
|
|
if (p.use_ellipse) {
|
|
// Measured radial excess over the tangential width, floored by the analytic bandwidth +
|
|
// parallax/capture term (see the CPU engine).
|
|
const float sbw = p.bw_sigma * Rpx;
|
|
const float radial_extra = fmaxf(sigma2_r[si] - s2t, sbw * sbw + p.c_radial * tan2t * tan2t);
|
|
if (Rpx > 1e-6f && radial_extra > 0.25f) { ux = rx / Rpx; uy = ry / Rpx; s2r = s2t + radial_extra; elong = true; }
|
|
}
|
|
const int Rf = elong ? min(3 * p.R, (int) ceilf(p.r2 + 2.0f * sqrtf(s2r))) : p.R;
|
|
const int Gf = 2 * Rf + 1;
|
|
if (threadIdx.x == 0) { s_Rf = Rf; s_Gf = Gf; s_gs = 0.0f; }
|
|
__syncthreads();
|
|
const float fx = px_x[i] - cx, fy = px_y[i] - cy;
|
|
float l_gs = 0.0f;
|
|
for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) {
|
|
const float ex = (k % Gf - Rf) - fx, ey = (k / Gf - Rf) - fy;
|
|
const float rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux;
|
|
const float g = expf(-rad * rad / (2.0f * s2r) - tn * tn / (2.0f * s2t));
|
|
Pbuf[k] = g; l_gs += g;
|
|
}
|
|
atomicAdd(&s_gs, l_gs);
|
|
__syncthreads();
|
|
const float gs = s_gs;
|
|
for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) Pbuf[k] /= gs;
|
|
__syncthreads();
|
|
}
|
|
|
|
const int Rf = s_Rf, Gf = s_Gf, GfGf = Gf * Gf;
|
|
const float B = fmaxf(bkg, (float) PIXEL_VARIANCE_FLOOR);
|
|
|
|
// How much of the expected profile is cleanly this reflection's own: s_pown over the whole grid
|
|
// (XDS's MINPK quantity, what Reject cuts on) and s_mown / s_mall over the r1 disk alone, which
|
|
// is what the summation seed the runaway guard compares against actually saw. See the CPU engine.
|
|
if (p.overlap) {
|
|
if (threadIdx.x == 0) { s_pown = 0.0f; s_mall = 0.0f; s_mown = 0.0f; }
|
|
__syncthreads();
|
|
float l_pown = 0.0f, l_mall = 0.0f, l_mown = 0.0f;
|
|
for (int k = threadIdx.x; k < GfGf; k += blockDim.x) {
|
|
const float Pp = Pbuf[k];
|
|
if (Pp <= 0.0f) continue;
|
|
const int dx = k % Gf - Rf, dy = k / Gf - Rf;
|
|
const int x = cx + dx, y = cy + dy;
|
|
if (x < 0 || y < 0 || x >= p.W || y >= p.H) continue;
|
|
if (!valid(img[y * p.W + x])) continue;
|
|
const bool own = BraggOwnedBy(owner[y * p.W + x], i);
|
|
// Zeroing the profile here is how Exclude drops the pixel: the fit skips any cell with
|
|
// P <= 0 already, and P is not renormalised, so the fitted amplitude comes out on the
|
|
// scale of the WHOLE profile - the renormalisation is the estimator's own doing.
|
|
if (!own && p.overlap == 2) Pbuf[k] = 0.0f;
|
|
if (own) l_pown += Pp;
|
|
if ((float) (dx * dx + dy * dy) < p.r1_sq) {
|
|
l_mall += Pp;
|
|
if (own) l_mown += Pp;
|
|
}
|
|
}
|
|
atomicAdd(&s_pown, l_pown); atomicAdd(&s_mall, l_mall); atomicAdd(&s_mown, l_mown);
|
|
__syncthreads();
|
|
if (p.overlap == 1 && s_pown < p.minpk) {
|
|
if (threadIdx.x == 0) ok_o[i] = 0;
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (threadIdx.x == 0) s_I = I_seed[i];
|
|
__syncthreads();
|
|
|
|
for (int iter = 0; iter < 4; ++iter) {
|
|
if (threadIdx.x == 0) { s_num = 0.0f; s_den = 0.0f; s_wsum = 0.0f; }
|
|
__syncthreads();
|
|
const float Ihere = s_I;
|
|
float l_num = 0.0f, l_den = 0.0f, l_wsum = 0.0f;
|
|
for (int k = threadIdx.x; k < GfGf; k += blockDim.x) {
|
|
const float Pp = Pbuf[k];
|
|
if (Pp <= 0.0f) continue;
|
|
const int x = cx + (k % Gf - Rf), y = cy + (k / Gf - Rf);
|
|
if (x < 0 || y < 0 || x >= p.W || y >= p.H) continue;
|
|
const int32_t px = img[y * p.W + x];
|
|
if (!valid(px)) continue;
|
|
const float v = fmaxf(B + Ihere * Pp, (float) WEIGHT_VARIANCE_MIN_FRACTION * B);
|
|
l_num += Pp * ((float) px - bkg) / v;
|
|
l_den += Pp * Pp / v;
|
|
l_wsum += Pp / v;
|
|
}
|
|
atomicAdd(&s_num, l_num); atomicAdd(&s_den, l_den); atomicAdd(&s_wsum, l_wsum);
|
|
__syncthreads();
|
|
if (threadIdx.x == 0 && s_den > 0.0f) s_I = s_num / s_den;
|
|
__syncthreads();
|
|
}
|
|
|
|
if (threadIdx.x == 0) {
|
|
if (s_den > 0.0f) {
|
|
// 1/s_den takes the background as exact; dI/dbkg = -s_wsum/s_den adds the background
|
|
// estimate's own error (see the CPU engine).
|
|
const float wr = s_wsum / s_den;
|
|
float I = s_I, sigma = sqrtf(1.0f / s_den + wr * wr * bkgvar_a[i]);
|
|
float var_bkg = fmaxf(0.0f, 1.0f / s_den - fmaxf(0.0f, I) + wr * wr * bkgvar_a[i]);
|
|
// Guard against profile-fit runaways (see the CPU engine): fall back to the summation seed
|
|
// when the profile result diverges from it. Exclude has taken pixels out of both, so the
|
|
// full-profile intensity is scaled back to the disk the seed read; nothing excluded gives
|
|
// exactly 1.
|
|
const float gs = (p.overlap == 2 && s_mall > 0.0f) ? s_mown / s_mall : 1.0f;
|
|
if (fabsf(I * gs - I_seed[i]) > (float) PROFILE_SUMMATION_MAX_NSIGMA * sigma_seed[i]) {
|
|
I = I_seed[i];
|
|
sigma = sigma_seed[i];
|
|
var_bkg = varbkg_seed[i];
|
|
}
|
|
I_o[i] = I; sigma_o[i] = sigma; varbkg_o[i] = var_bkg; ok_o[i] = 1;
|
|
} else ok_o[i] = 0;
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
BraggIntegrationEngineGPU::BraggIntegrationEngineGPU(const DiffractionExperiment &experiment,
|
|
std::shared_ptr<CudaStream> stream)
|
|
: BraggIntegrationEngine(experiment),
|
|
stream(std::move(stream)),
|
|
d_mask(npixel),
|
|
d_shell_grid(static_cast<size_t>(bragg_engine::N_SHELL) * GG),
|
|
d_global_grid(GG),
|
|
d_shell_P(static_cast<size_t>(bragg_engine::N_SHELL) * GG),
|
|
d_global_P(GG),
|
|
d_mom(MOM_STRIDE * (bragg_engine::N_SHELL + 1)),
|
|
d_sigma2_r(bragg_engine::N_SHELL + 1),
|
|
d_sigma2_t(bragg_engine::N_SHELL + 1),
|
|
d_shell_n(bragg_engine::N_SHELL),
|
|
d_global_n(1),
|
|
d_invd2(2) {
|
|
threads = 128;
|
|
|
|
// Fit profile grid: R for empirical / box, up to 3R (radially elongated) for the Gaussian.
|
|
const int max_Rf = empirical ? R : 3 * R;
|
|
const int max_Gf = 2 * max_Rf + 1;
|
|
fit_shared_bytes = static_cast<size_t>(max_Gf) * max_Gf * sizeof(float);
|
|
|
|
// Shared radial window of the background curve in boxsum. The stencil spans r0 +- the radial
|
|
// semi-axis of the outer ellipse, widest at the far corner of the detector, and the window
|
|
// covers twice that so a reflection anywhere on the detector fits.
|
|
rad_w = 2 * (static_cast<int>(std::ceil(r3 + BraggStencilGrow_px(static_cast<float>(r_max), stencil))) + 1) + 1;
|
|
boxsum_shared_bytes = static_cast<size_t>(rad_w) * (sizeof(float) + sizeof(int));
|
|
|
|
// The current device, not device 0: workers are pinned round-robin across GPUs, so device 0's
|
|
// shared-memory size can belong to a different card than the one these kernels launch on.
|
|
int device = 0;
|
|
cuda_err(cudaGetDevice(&device));
|
|
cudaDeviceProp prop{};
|
|
cuda_err(cudaGetDeviceProperties(&prop, device));
|
|
if (fit_shared_bytes > prop.sharedMemPerBlock)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
|
|
"BraggIntegrationEngineGPU: profile grid exceeds shared memory (r2 too large)");
|
|
// boxsum's dynamic window sits alongside its static shared arrays (the trimmed-mean buffer
|
|
// dominates them), so leave room for those rather than budgeting the whole block.
|
|
if (boxsum_shared_bytes + sizeof(int) * BKG_TRIM_MAX + 256 > prop.sharedMemPerBlock)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
|
|
"BraggIntegrationEngineGPU: background ring exceeds shared memory");
|
|
|
|
// Radial background curve: one bin per pixel of distance from the beam, out to the far corner.
|
|
// Allocated whenever the correction COULD run, so the auto mode can turn it on for an individual
|
|
// image; whether it runs for a given image is decided in Run() from the per-image bkg_radial.
|
|
if (overlap != OverlapMode::Off)
|
|
d_owner = CudaDevicePtr<uint32_t>(npixel);
|
|
|
|
if (bkg_radial || bkg_radial_auto) {
|
|
n_rad = static_cast<int>(std::ceil(r_max)) + 2;
|
|
d_rad_sum = CudaDevicePtr<float>(n_rad);
|
|
d_rad_cnt = CudaDevicePtr<int>(n_rad);
|
|
d_k_diff = CudaDevicePtr<float>(k_diff.size());
|
|
cuda_err(cudaMemcpy(d_k_diff, k_diff.data(), sizeof(float) * k_diff.size(),
|
|
cudaMemcpyHostToDevice));
|
|
}
|
|
}
|
|
|
|
void BraggIntegrationEngineGPU::EnsureCapacity(size_t n) {
|
|
if (n <= capacity)
|
|
return;
|
|
// Grow with slack. cudaMalloc/cudaFree take a device-wide lock in the driver, so growing to
|
|
// exactly n makes every image that sets a new reflection-count record stall all other workers.
|
|
const size_t new_capacity = std::max(n, capacity + capacity / 2);
|
|
|
|
d_px_x = CudaDevicePtr<float>(new_capacity);
|
|
d_px_y = CudaDevicePtr<float>(new_capacity);
|
|
d_d = CudaDevicePtr<float>(new_capacity);
|
|
d_cx = CudaDevicePtr<int>(new_capacity);
|
|
d_cy = CudaDevicePtr<int>(new_capacity);
|
|
d_I = CudaDevicePtr<float>(new_capacity);
|
|
d_sigma = CudaDevicePtr<float>(new_capacity);
|
|
d_bkg = CudaDevicePtr<float>(new_capacity);
|
|
d_bkg_var = CudaDevicePtr<float>(new_capacity);
|
|
d_var_bkg = CudaDevicePtr<float>(new_capacity);
|
|
d_isum = CudaDevicePtr<float>(new_capacity);
|
|
d_ninner = CudaDevicePtr<int>(new_capacity);
|
|
d_rbin = CudaDevicePtr<int>(new_capacity);
|
|
d_kbin = CudaDevicePtr<int>(new_capacity);
|
|
d_obs_x = CudaDevicePtr<float>(new_capacity);
|
|
d_obs_y = CudaDevicePtr<float>(new_capacity);
|
|
d_ok = CudaDevicePtr<uint8_t>(new_capacity);
|
|
d_strong = CudaDevicePtr<uint8_t>(new_capacity);
|
|
d_has_obs = CudaDevicePtr<uint8_t>(new_capacity);
|
|
|
|
h_px_x.resize(new_capacity); h_px_y.resize(new_capacity); h_d.resize(new_capacity);
|
|
h_I.resize(new_capacity); h_sigma.resize(new_capacity); h_bkg.resize(new_capacity);
|
|
h_var_bkg.resize(new_capacity);
|
|
h_obs_x.resize(new_capacity); h_obs_y.resize(new_capacity);
|
|
h_ok.resize(new_capacity); h_has_obs.resize(new_capacity);
|
|
capacity = new_capacity;
|
|
}
|
|
|
|
std::vector<Reflection> BraggIntegrationEngineGPU::Run(const ImagePreprocessorBuffer &image,
|
|
const std::vector<Reflection> &predicted,
|
|
size_t npredicted, int64_t image_number) {
|
|
std::vector<BraggFitResult> results(npredicted);
|
|
if (image.size() != npixel || npredicted == 0)
|
|
return Finalize(predicted, npredicted, results, image_number);
|
|
|
|
const int32_t *img = image.getGPUBuffer();
|
|
if (img == nullptr)
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"BraggIntegrationEngineGPU: image buffer is not on the GPU");
|
|
|
|
EnsureCapacity(npredicted);
|
|
const int n = static_cast<int>(npredicted);
|
|
for (size_t i = 0; i < npredicted; ++i) {
|
|
h_px_x[i] = predicted[i].predicted_x;
|
|
h_px_y[i] = predicted[i].predicted_y;
|
|
h_d[i] = predicted[i].d;
|
|
}
|
|
cuda_err(cudaMemcpyAsync(d_px_x, h_px_x.data(), sizeof(float) * npredicted, cudaMemcpyHostToDevice, *stream));
|
|
cuda_err(cudaMemcpyAsync(d_px_y, h_px_y.data(), sizeof(float) * npredicted, cudaMemcpyHostToDevice, *stream));
|
|
cuda_err(cudaMemcpyAsync(d_d, h_d.data(), sizeof(float) * npredicted, cudaMemcpyHostToDevice, *stream));
|
|
|
|
BraggGpuParams p{
|
|
.W = static_cast<int>(xpixel), .H = static_cast<int>(ypixel),
|
|
.r1_sq = r1_sq, .r2 = r2, .r2_sq = r2_sq, .r3 = r3, .r3_sq = r3_sq,
|
|
.R = R, .G = G, .GG = GG,
|
|
.bkg_clip_nsigma = mode != IntegratorMode::BoxSum ? bkg_clip_nsigma : 0.0f,
|
|
.empirical = empirical ? 1 : 0,
|
|
.use_ellipse = use_ellipse ? 1 : 0,
|
|
.bw_sigma = static_cast<float>(bw_sigma), .c_radial = static_cast<float>(c_radial),
|
|
.F_px = static_cast<float>(F_px),
|
|
.beam_x = beam_x, .beam_y = beam_y,
|
|
.bkg_trim = bkg_trim, // effective trim fraction (0 for stills), set by the base ctor from settings
|
|
.stencil = stencil, .rad_w = rad_w, .n_kern = n_kern,
|
|
.overlap = overlap == OverlapMode::Reject ? 1 : (overlap == OverlapMode::Exclude ? 2 : 0),
|
|
.boxsum_reject = (overlap == OverlapMode::Reject && mode == IntegratorMode::BoxSum) ? 1 : 0,
|
|
.exclude_px = (overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum) ? 1 : 0,
|
|
.claim_sq = claim * claim, .inv_claim = inv_claim, .minpk = overlap_min_peak,
|
|
};
|
|
|
|
// Whether the radial correction runs for THIS image. n_rad only says the buffers exist - under
|
|
// the auto mode they are allocated for every image and used for the ones the ice score selects.
|
|
const int rad_n = bkg_radial ? n_rad : 0;
|
|
|
|
// Pass A: reset accumulators, mask, then box-sum.
|
|
cuda_err(cudaMemsetAsync(d_mask, 0, npixel, *stream));
|
|
if (p.overlap)
|
|
cuda_err(cudaMemsetAsync(d_owner, 0xff, sizeof(uint32_t) * npixel, *stream)); // BRAGG_OWNER_NONE
|
|
if (rad_n > 0) {
|
|
cuda_err(cudaMemsetAsync(d_rad_sum, 0, sizeof(float) * n_rad, *stream));
|
|
cuda_err(cudaMemsetAsync(d_rad_cnt, 0, sizeof(int) * n_rad, *stream));
|
|
}
|
|
reset<<<32, 256, 0, *stream>>>(d_shell_grid, d_global_grid, d_mom, d_shell_n, d_global_n, d_invd2, GG);
|
|
mark_mask<<<n, threads, 0, *stream>>>(d_px_x, d_px_y, d_mask, d_owner, p, n);
|
|
boxsum<<<n, threads, boxsum_shared_bytes, *stream>>>(d_px_x, d_px_y, d_d, img, d_mask, d_owner, p, n,
|
|
d_cx, d_cy, d_I, d_sigma, d_bkg, d_bkg_var, d_var_bkg, d_obs_x, d_obs_y,
|
|
d_ok, d_strong, d_has_obs, d_invd2,
|
|
d_isum, d_ninner, d_rbin, d_kbin,
|
|
d_rad_sum, d_rad_cnt, rad_n);
|
|
|
|
// Correct the flat annulus background for the curvature of the radial background before anything
|
|
// downstream (profile fit, variance) reads it.
|
|
if (rad_n > 0)
|
|
radial_correct<<<(n + threads - 1) / threads, threads, 0, *stream>>>(
|
|
d_rad_sum, d_rad_cnt, rad_n, d_k_diff, k_len, k_off,
|
|
d_isum, d_ninner, d_rbin, d_kbin, d_ok, d_bkg, d_I, n);
|
|
|
|
if (mode != IntegratorMode::BoxSum) {
|
|
// Pass B: learn (shell computed inline) -> build -> fit.
|
|
learn_profile<<<n, threads, 0, *stream>>>(img, d_owner, d_px_x, d_px_y, d_cx, d_cy, d_d, d_invd2,
|
|
d_I, d_bkg, d_ok, d_strong,
|
|
d_shell_grid, d_global_grid, d_mom, d_shell_n, d_global_n, p, n);
|
|
build_profiles<<<bragg_engine::N_SHELL + 1, threads, 0, *stream>>>(
|
|
d_shell_grid, d_global_grid, d_mom, d_shell_n, d_global_n,
|
|
d_shell_P, d_global_P, d_sigma2_r, d_sigma2_t, p);
|
|
fit<<<n, threads, fit_shared_bytes, *stream>>>(img, d_owner, d_px_x, d_px_y, d_cx, d_cy, d_d, d_invd2,
|
|
d_I, d_sigma, d_bkg, d_bkg_var, d_var_bkg, d_ok, d_shell_P, d_global_P,
|
|
d_sigma2_r, d_sigma2_t, d_shell_n,
|
|
d_I, d_sigma, d_var_bkg, d_ok, p, n);
|
|
}
|
|
|
|
cuda_err(cudaMemcpyAsync(h_I.data(), d_I, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(h_sigma.data(), d_sigma, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(h_bkg.data(), d_bkg, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(h_var_bkg.data(), d_var_bkg, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(h_ok.data(), d_ok, sizeof(uint8_t) * npredicted, cudaMemcpyDeviceToHost, *stream));
|
|
// Pass A always fills the box-sum centroid (observed spot position), so copy it back in all modes -
|
|
// post-refinement uses it as the observed position (beam-centre / distance).
|
|
cuda_err(cudaMemcpyAsync(h_obs_x.data(), d_obs_x, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(h_obs_y.data(), d_obs_y, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(h_has_obs.data(), d_has_obs, sizeof(uint8_t) * npredicted, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
|
|
for (size_t i = 0; i < npredicted; ++i) {
|
|
if (!h_ok[i]) continue;
|
|
results[i].I = h_I[i];
|
|
results[i].sigma = h_sigma[i];
|
|
results[i].bkg = h_bkg[i];
|
|
results[i].var_bkg = h_var_bkg[i];
|
|
results[i].ok = true;
|
|
if (h_has_obs[i]) {
|
|
results[i].observed_x = h_obs_x[i];
|
|
results[i].observed_y = h_obs_y[i];
|
|
results[i].has_observed = true;
|
|
}
|
|
}
|
|
return Finalize(predicted, npredicted, results, image_number);
|
|
}
|