leonarski_fandClaude Opus 5
1239c49731
Setting a bandwidth flipped three unrelated switches at once: it changed the profile's radial capture term, it moved the width measurement from the signal disk to the whole fit grid, and it silently overrode the background clip and trim, so --background-clip under --bandwidth was ignored - the two runs were bit-identical. The width measurement was the damaging one. The fit grid is an azimuthally averaged stack, so its second moment is sigma_r^2 + sigma_t^2 and the radial smear of a bandwidth leaked into the tangential model - a tangential width of 3.04 px against a 1.06 px truth, inflating the effective background pixel count where the weak signal is. The result was a step rather than a slope: on genuinely monochromatic data, declaring a 0.2% bandwidth cost ISa 28.4 -> 22.2. Measure the two widths separately, accumulated in each spot's own radial/tangential frame over the signal disk, from the signed profile cells - away from the peak a learned cell is background noise centred on zero, so the signed sum is unbiased, while clamping it at zero turns that noise into a pedestal the r^2 weight reads as width. The radial term is then the measured excess or the analytic floor, whichever is larger. With the two widths separated there is nothing left for the broadband switch to select, so it is gone - which is the proof the three were independent. The background clip and trim now come from the settings in every case; the tuned 3-sigma broadband default moves to the rugnux front end, which is the only place that knows whether the user gave a value. Monochromatic data: declaring a 0.2% bandwidth now costs ISa 28.4 -> 27.9 rather than 22.2, and forcing the old 3-sigma clip in the new build reproduces the good result, so none of the step came from the clip. On large-bandwidth data CC1/2 improves in 8 of 10 shells. Across 12 monochromatic crystals the space groups are unchanged and CC1/2 moves by at most 0.2 points. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
704 lines
34 KiB
Plaintext
704 lines
34 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)
|
|
};
|
|
|
|
__device__ inline bool valid(int32_t v) { return v != INT32_MIN && v != INT32_MAX; }
|
|
|
|
// idea 1: max annulus pixels held in shared memory for the per-block trimmed-mean sort (else flat mean).
|
|
constexpr int BKG_TRIM_MAX = 512;
|
|
|
|
// 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 disk of every predicted reflection (race-free: all writes are 1). ---
|
|
__global__ void mark_mask(const float *px_x, const float *px_y, uint8_t *mask, BraggGpuParams p, int n) {
|
|
const int i = blockIdx.x;
|
|
if (i >= n) return;
|
|
const float cx = px_x[i], cy = px_y[i];
|
|
const int x0 = max(0, (int) floorf(cx - p.r2 - 1.0f));
|
|
const int x1 = min(p.W - 1, (int) ceilf(cx + p.r2 + 1.0f));
|
|
const int y0 = max(0, (int) floorf(cy - p.r2 - 1.0f));
|
|
const int y1 = min(p.H - 1, (int) ceilf(cy + p.r2 + 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 float ddx = (float) x - cx, ddy = (float) y - cy;
|
|
if (ddx * ddx + ddy * ddy < p.r2_sq) mask[y * p.W + x] = 1;
|
|
}
|
|
}
|
|
|
|
// --- 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, 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,
|
|
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;
|
|
__shared__ double s_bkgsum;
|
|
__shared__ int s_accept;
|
|
__shared__ double s_bkg, s_thr, s_clipsum;
|
|
__shared__ int s_clipn;
|
|
__shared__ float s_ux, s_uy, s_r0;
|
|
// The whole stencil spans r0 +- r3, so a 32-bin window centred on the reflection always covers it.
|
|
constexpr int RAD_W = 32;
|
|
__shared__ float s_radv[RAD_W];
|
|
__shared__ int s_radn[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_ux = s_r0 > 1e-6f ? rx / s_r0 : 1.0f;
|
|
s_uy = s_r0 > 1e-6f ? ry / s_r0 : 0.0f;
|
|
}
|
|
if (threadIdx.x == 0) {
|
|
s_Isum = 0; s_Ix = 0; s_Iy = 0;
|
|
s_ninner = 0; s_ninner_valid = 0; s_nbkg = 0; s_bkgsum = 0.0;
|
|
}
|
|
__syncthreads();
|
|
|
|
const float cx = px_x[i], cy = px_y[i];
|
|
const int x0 = max(0, (int) floorf(cx - p.r3 - 1.0f));
|
|
const int x1 = min(p.W - 1, (int) ceilf(cx + p.r3 + 1.0f));
|
|
const int y0 = max(0, (int) floorf(cy - p.r3 - 1.0f));
|
|
const int y1 = min(p.H - 1, (int) ceilf(cy + p.r3 + 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;
|
|
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 float ddx = (float) x - cx, ddy = (float) y - cy;
|
|
const float d2 = ddx * ddx + ddy * ddy;
|
|
const int32_t px = img[y * p.W + x];
|
|
if (d2 < p.r1_sq) {
|
|
++l_ni;
|
|
if (valid(px)) { l_Isum += px; l_Ix += (long long) x * px; l_Iy += (long long) y * px; ++l_niv; }
|
|
} else if (d2 >= p.r2_sq && d2 < 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);
|
|
__syncthreads();
|
|
|
|
for (int t = threadIdx.x; t < RAD_W; t += blockDim.x) { s_radv[t] = 0.0f; s_radn[t] = 0; }
|
|
if (threadIdx.x == 0) s_radbase = (int) lroundf(s_r0) - 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 float ddx = (float) x - cx, ddy = (float) y - cy;
|
|
const float d2 = ddx * ddx + ddy * ddy;
|
|
if (!(d2 >= p.r2_sq && d2 < 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 float ddx = (float) x - cx, ddy = (float) y - cy;
|
|
const float d2 = ddx * ddx + ddy * ddy;
|
|
if (!(d2 >= p.r2_sq && d2 < 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) {
|
|
// Radial offset = projection on the beam->reflection direction, no per-pixel sqrt.
|
|
const float off = ddx * s_ux + ddy * s_uy;
|
|
const int idx = (int) lroundf(s_r0 + off) - s_radbase;
|
|
if (idx >= 0 && idx < 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 < 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; }
|
|
|
|
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);
|
|
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 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;
|
|
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. 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 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];
|
|
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 += k_diff[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 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__ 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);
|
|
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.
|
|
if (fabsf(I - 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);
|
|
|
|
// 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)");
|
|
|
|
// 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 (bkg_radial || bkg_radial_auto) {
|
|
const double fx = std::max<double>(beam_x, static_cast<double>(xpixel) - beam_x);
|
|
const double fy = std::max<double>(beam_y, static_cast<double>(ypixel) - beam_y);
|
|
n_rad = static_cast<int>(std::ceil(std::hypot(fx, fy))) + 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_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
|
|
};
|
|
|
|
// 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 (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, p, n);
|
|
boxsum<<<n, threads, 0, *stream>>>(d_px_x, d_px_y, d_d, img, d_mask, 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_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, static_cast<int>(k_diff.size()), k_off,
|
|
d_isum, d_ninner, d_rbin, 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_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_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);
|
|
}
|