A weighted mean is only unbiased while the weights are independent of the values
being averaged. The IUCr's own nomenclature report (Schwarzenbach et al., Acta
Cryst A45 (1989) 63-75) puts it directly: weights in averaging "should not be
based on the counting statistics of the individual observations whose estimated
variances are biased and result in larger weights for accidentally low
intensities". Two places in the rotation pipeline were doing exactly that, and
between them they drove whole resolution shells of merged intensity negative.
1. The profile fit computed its non-signal variance as
var_bkg = max(0, 1/den - max(0, I) + bkg-estimate term)
The point of a separate var_bkg is that it does NOT move with the
reflection's own fluctuation, and 1/den - I is the quantity that does not:
1/den is the fit variance taken at the fitted intensity and grows with it
roughly one for one. Clamping the subtrahend at zero left a down-fluctuated
reflection's own deflated variance standing as its background variance.
Measured over 6.9 M partials of one weak rotation dataset, var_bkg/bkg came
out at 3.7-5.4 for observations with I < 0 against 11.4-13.7 for I > 0 - the
down-fluctuated half of every reflection carried a variance ~2.7x too small
and was weighted up by the same factor, first in the 3D combine and then
again in the merge. Removing the clamp makes var_bkg flat in I (~13 x bkg
across the whole range).
2. The merge then weighted each combined full by 1/sigma_full^2, and sigma_full
is by construction a function of the full's own answer: the combine's
variance carries a corr*max(0, F) signal term, so every full with F <= 0 got
the smallest variance the model allows while the strongest quartile got
2.26x more. The merge now rebuilds that variance at the reflection's mean
instead, from a linear model var(I) = var_bkg + var_per_I * I that the
combine measures and stores on the full. This mirrors
MergeOnTheFly::CorrectedSigma, whose comment already claimed to mirror the
rotation combine.
Verified against an estimator that cannot see the fluctuation - summing the
partials and dividing by the summed partiality, the classical construction every
other program uses (Greenhough & Suddath, J. Appl. Cryst. 19 (1986) 400-409, via
Leslie, Acta Cryst D55 (1999) 1696-1702: profile fitting biases the individual
partials but not their sum). Reproducing the merge on dumped observations, the
shipped weighting sat ~1.9 sigma below that reference in the noise shells; the
two changes recover most of it, and every intensity-independent weighting
scheme agrees with the reference once (1) is in.
Four-crystal probe, XDS resolution limits, branch fingerprint identical on all
four (so none of these is a two-pass branch flip):
weak cubic case last shell <I/sig> -1.6 -> +0.2 (XDS +0.10), last shell
R_meas 478% -> 250% (XDS 246%), overall <I/sig> 6.1 -> 7.5
(XDS 7.18), R_meas 18.3% -> 18.1%, CC1/2_hi 38.2% -> 43.7%
tetragonal case outer shells <I/sig> -0.4/-0.8/-0.9/-1.0 -> +1.8/+1.2/
+0.9/+0.4, R_meas 184%/595%/7614%/nan -> 95%/119%/135%/232%
(the nan was the shell mean crossing zero), R_meas 33.3% ->
32.9%, CC1/2_hi 38.3% -> 56.5%
trigonal case R_meas 13.0% -> 12.5%, CC1/2_hi 14.4% -> 16.5%
strong control unchanged to every printed digit but ISa
Cost: ISa falls (17.2 -> 14.0 and 16.7 -> 14.9 on the two mid-strength cases,
28.3 -> 27.8 on the control). Strong reflections are untouched by (1) - their
partials are all positive, so var_bkg is bit-identical - but the joint a/b fit
redistributes: honest weak sigmas lower a, and b rises to keep the strong bins
fitted. The median reduced chi^2 improves (1.25 -> 1.14, 1.35 -> 1.28) so the
new split describes the scatter better, but ISa is the one headline metric that
moves the wrong way and it should be watched over the full battery.
The integrator change is shared, so the stills merge sees it too; there it feeds
GetExpectedVarianceMerge, which had been handed the same contaminated var_bkg.
That path is untested here.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
799 lines
41 KiB
Plaintext
799 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). The
|
|
// reweighting is the Kabsch/Otwinowski iteration: Kabsch, Acta Cryst D66, 133-144 (2010);
|
|
// Otwinowski & Minor, Methods Enzymol 276, 307-326 (1997).
|
|
// 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]);
|
|
// The signal part to remove is I itself, not max(0, I) - see the CPU engine.
|
|
float var_bkg = fmaxf(0.0f, 1.0f / s_den - 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);
|
|
}
|