398 lines
21 KiB
C++
398 lines
21 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "BeamCenterFromBackground.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <thread>
|
|
|
|
#include "../../common/JFJochMath.h"
|
|
#include "../../common/ParallelFor.h"
|
|
|
|
namespace {
|
|
|
|
// The band the background is fitted over. The low-resolution end sits outside the beam stop
|
|
// and its penumbra, the high-resolution end where the solvent ring has died away.
|
|
constexpr float BAND_LOW_RES_A = 12.0f;
|
|
constexpr float BAND_HIGH_RES_A = 2.2f;
|
|
|
|
constexpr int SECTORS = 36;
|
|
constexpr int RADIAL_BINS = 120;
|
|
|
|
// A cell with fewer pixels than this has no usable mean.
|
|
constexpr int MIN_PIXELS_PER_CELL = 20;
|
|
|
|
// A radial bin missing more azimuth than this is a partial ring - it leaves the detector, or a
|
|
// module gap eats it - and a partial ring biases the profile it is compared against.
|
|
constexpr float MIN_SECTOR_COVERAGE = 0.85f;
|
|
|
|
// Fractions of a bin's pixels are Bragg peaks. Two rounds of clipping at the upper 2 sigma take
|
|
// the mean back to the background without needing the pixel values a second time.
|
|
constexpr float CLIP_SIGMA = 2.0f;
|
|
constexpr int CLIP_ROUNDS = 2;
|
|
|
|
// The profile is rebuilt at the trial centre every iteration, so a centre that is off smears the
|
|
// solvent ring and flattens g', which over-estimates the shift. Half steps damp that; the fixed
|
|
// point is unchanged, only the path to it.
|
|
constexpr float DAMPING = 0.5f;
|
|
constexpr float CONVERGED_PXL = 0.02f;
|
|
|
|
// The step is compared with the one before it, and a walk whose step CANCELS its predecessor twice
|
|
// running has crossed its fixed point rather than walked toward it: with DAMPING = 0.5 that happens
|
|
// where the local gain exceeds 2, which is a period-2 limit cycle the walk will not leave.
|
|
// Instrumented, a sweep that spent 96 of its 100 iterations in one oscillated by 0.08 px about a
|
|
// centre it knew to 1.7 px. Cancelling is the whole of it, and it is not the same as merely turning
|
|
// by more than a right angle: a travelling walk wobbles, and on one sweep a single 0.56 px step
|
|
// between 2 px ones turned twice in passing and ended the fit 118 px short of the centre the same
|
|
// fit reaches from the other side - in one build and not another, because the sign of a dot product
|
|
// between two near-orthogonal steps is set by the last ulp and not by the data.
|
|
constexpr int REVERSALS_AT_THE_FIXED_POINT = 2;
|
|
|
|
// A travel budget, not a convergence criterion. The shift a sector's regression can report is
|
|
// bounded by the width of the features it reads - a sector whose profile has moved a long way is
|
|
// not g + d*g' for any d, and the least-squares projection onto g' returns far less than the true
|
|
// d - so the walk advances by a bounded distance per iteration, of the order of twenty pixels,
|
|
// however far it still has to go. Ten iterations therefore cap the fit at about two hundred
|
|
// pixels of travel and a centre further out than that is left part way there, still walking, with
|
|
// the per-iteration precision reported as though it had arrived. The count is set to cross a
|
|
// detector and then converge on it; a centre that is already close leaves on CONVERGED_PXL, or on
|
|
// the reversal test above, after a handful. Measured on a sweep whose centre is 170 px out: the
|
|
// walk crawls at 1-2 px an iteration for its first forty, accelerates across the gap, and reaches
|
|
// CONVERGED_PXL at 107-115 - so a budget of 100 stopped it in the last few pixels of its approach
|
|
// and left where it stopped to the build.
|
|
constexpr int MAX_ITERATIONS = 300;
|
|
|
|
// The two passes over the pixels are split into this many row blocks, each accumulating into cells
|
|
// of its own, and the blocks are folded in block order. The split is a property of the image and
|
|
// not of the machine, so the sums are grouped the same way however many threads ran them.
|
|
constexpr int BLOCKS = 64;
|
|
|
|
// Below these the fit has not seen enough of the detector to be believed at all.
|
|
constexpr int MIN_USABLE_SECTORS = SECTORS * 3 / 5;
|
|
constexpr int MIN_USABLE_RADIAL_BINS = 15;
|
|
|
|
float median_of(std::vector<float> &v) {
|
|
const size_t half = v.size() / 2;
|
|
std::nth_element(v.begin(), v.begin() + half, v.end());
|
|
return v[half];
|
|
}
|
|
|
|
} // namespace
|
|
|
|
std::optional<BeamCenterEstimate>
|
|
FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const PixelMask &mask,
|
|
const std::vector<float> &mean, size_t nthreads,
|
|
std::optional<std::pair<float, float>> start) {
|
|
if (nthreads == 0)
|
|
nthreads = std::max(1u, std::thread::hardware_concurrency());
|
|
const auto W = static_cast<int>(experiment.GetXPixelsNumConv());
|
|
const auto H = static_cast<int>(experiment.GetYPixelsNumConv());
|
|
const size_t n_pixels = static_cast<size_t>(W) * H;
|
|
if (mean.size() != n_pixels)
|
|
return {};
|
|
|
|
const auto &pixel_mask = mask.GetMask(experiment);
|
|
auto geom = experiment.GetDiffractionGeometry();
|
|
|
|
const float wavelength = geom.GetWavelength_A();
|
|
const float sin_high = wavelength / (2.0f * BAND_HIGH_RES_A);
|
|
if (sin_high >= 1.0f)
|
|
return {};
|
|
const float tt_lo = 2.0f * std::asin(wavelength / (2.0f * BAND_LOW_RES_A));
|
|
const float tt_hi = 2.0f * std::asin(sin_high);
|
|
const float d_tt = (tt_hi - tt_lo) / RADIAL_BINS;
|
|
|
|
const auto rot = geom.GetDetectorMatrix().arr(); // row major
|
|
const float pixel_size = geom.GetPixelSize_mm();
|
|
const float distance = geom.GetDetectorDistance_mm();
|
|
|
|
float beam_x = start ? start->first : geom.GetBeamX_pxl();
|
|
float beam_y = start ? start->second : geom.GetBeamY_pxl();
|
|
|
|
constexpr int n_cells = RADIAL_BINS * SECTORS;
|
|
std::vector<int32_t> cell_of(n_pixels);
|
|
std::vector<double> sum(n_cells), sum_sq(n_cells), sum_jx(n_cells), sum_jy(n_cells);
|
|
std::vector<int32_t> count(n_cells), count_all(n_cells);
|
|
std::vector<float> profile(RADIAL_BINS), d_profile(RADIAL_BINS), clip_limit(n_cells);
|
|
std::vector<char> radial_ok(RADIAL_BINS);
|
|
|
|
// One set of cells per block, allocated once and reused by every iteration.
|
|
std::vector<int> block_row(BLOCKS + 1);
|
|
for (int b = 0; b <= BLOCKS; b++)
|
|
block_row[b] = static_cast<int>(static_cast<int64_t>(b) * H / BLOCKS);
|
|
std::vector<double> block_sum(static_cast<size_t>(BLOCKS) * n_cells);
|
|
std::vector<double> block_sum_sq(static_cast<size_t>(BLOCKS) * n_cells);
|
|
std::vector<double> block_jx(static_cast<size_t>(BLOCKS) * n_cells);
|
|
std::vector<double> block_jy(static_cast<size_t>(BLOCKS) * n_cells);
|
|
std::vector<int32_t> block_count(static_cast<size_t>(BLOCKS) * n_cells);
|
|
|
|
float step_x = 0.0f, step_y = 0.0f, sigma_x = 0.0f, sigma_y = 0.0f;
|
|
float previous_x = 0.0f, previous_y = 0.0f;
|
|
int reversals = 0;
|
|
for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++) {
|
|
ParallelFor(BLOCKS, nthreads, [&](int b) {
|
|
double *b_sum = block_sum.data() + static_cast<size_t>(b) * n_cells;
|
|
double *b_sum_sq = block_sum_sq.data() + static_cast<size_t>(b) * n_cells;
|
|
double *b_jx = block_jx.data() + static_cast<size_t>(b) * n_cells;
|
|
double *b_jy = block_jy.data() + static_cast<size_t>(b) * n_cells;
|
|
int32_t *b_count = block_count.data() + static_cast<size_t>(b) * n_cells;
|
|
std::fill(b_sum, b_sum + n_cells, 0.0);
|
|
std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0);
|
|
std::fill(b_jx, b_jx + n_cells, 0.0);
|
|
std::fill(b_jy, b_jy + n_cells, 0.0);
|
|
std::fill(b_count, b_count + n_cells, 0);
|
|
for (int y = block_row[b]; y < block_row[b + 1]; y++) {
|
|
for (int x = 0; x < W; x++) {
|
|
const size_t i = static_cast<size_t>(y) * W + x;
|
|
cell_of[i] = -1;
|
|
if (pixel_mask[i] != 0 || !std::isfinite(mean[i]))
|
|
continue;
|
|
const float u = (x - beam_x) * pixel_size;
|
|
const float v = (y - beam_y) * pixel_size;
|
|
const float lx = rot[0] * u + rot[1] * v + rot[2] * distance;
|
|
const float ly = rot[3] * u + rot[4] * v + rot[5] * distance;
|
|
const float lz = rot[6] * u + rot[7] * v + rot[8] * distance;
|
|
const float rho = std::sqrt(lx * lx + ly * ly);
|
|
const float two_theta = std::atan2(rho, lz);
|
|
if (two_theta < tt_lo || two_theta >= tt_hi || rho == 0.0f)
|
|
continue;
|
|
|
|
const float phi = std::atan2(ly, lx);
|
|
// Both bins are clamped: a pixel one float ulp below the top of the band divides
|
|
// to exactly RADIAL_BINS, which is one cell past the end of every accumulator.
|
|
const int r_bin = std::clamp(static_cast<int>((two_theta - tt_lo) / d_tt), 0, RADIAL_BINS - 1);
|
|
const int s_bin = std::clamp(static_cast<int>((phi + PI) / (2 * PI) * SECTORS), 0, SECTORS - 1);
|
|
const int cell = r_bin * SECTORS + s_bin;
|
|
|
|
// d(2theta)/d(beam), through the lab coordinate: the detector coordinate depends
|
|
// on the centre only as (x - beam_x), so moving the centre is moving the pixel.
|
|
const float denominator = rho * rho + lz * lz;
|
|
const float g_x = lz * lx / (rho * denominator);
|
|
const float g_y = lz * ly / (rho * denominator);
|
|
const float g_z = -rho / denominator;
|
|
cell_of[i] = cell;
|
|
b_count[cell]++;
|
|
b_sum[cell] += mean[i];
|
|
b_sum_sq[cell] += static_cast<double>(mean[i]) * mean[i];
|
|
b_jx[cell] += -pixel_size * (g_x * rot[0] + g_y * rot[3] + g_z * rot[6]);
|
|
b_jy[cell] += -pixel_size * (g_x * rot[1] + g_y * rot[4] + g_z * rot[7]);
|
|
}
|
|
}
|
|
});
|
|
for (int c = 0; c < n_cells; c++) {
|
|
double s = 0, ss = 0, jx = 0, jy = 0;
|
|
int32_t n = 0;
|
|
for (int b = 0; b < BLOCKS; b++) {
|
|
const size_t k = static_cast<size_t>(b) * n_cells + c;
|
|
s += block_sum[k]; ss += block_sum_sq[k];
|
|
jx += block_jx[k]; jy += block_jy[k];
|
|
n += block_count[k];
|
|
}
|
|
sum[c] = s; sum_sq[c] = ss; sum_jx[c] = jx; sum_jy[c] = jy; count[c] = n;
|
|
}
|
|
|
|
count_all = count; // the Jacobian sums belong to the unclipped pixel set
|
|
|
|
for (int round = 0; round < CLIP_ROUNDS; round++) {
|
|
for (int c = 0; c < n_cells; c++) {
|
|
if (count[c] < MIN_PIXELS_PER_CELL) { clip_limit[c] = -1.0f; continue; }
|
|
const double m = sum[c] / count[c];
|
|
const double variance = std::max(sum_sq[c] / count[c] - m * m, 0.0);
|
|
clip_limit[c] = static_cast<float>(m + CLIP_SIGMA * std::sqrt(variance));
|
|
}
|
|
ParallelFor(BLOCKS, nthreads, [&](int b) {
|
|
double *b_sum = block_sum.data() + static_cast<size_t>(b) * n_cells;
|
|
double *b_sum_sq = block_sum_sq.data() + static_cast<size_t>(b) * n_cells;
|
|
int32_t *b_count = block_count.data() + static_cast<size_t>(b) * n_cells;
|
|
std::fill(b_sum, b_sum + n_cells, 0.0);
|
|
std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0);
|
|
std::fill(b_count, b_count + n_cells, 0);
|
|
const size_t lo = static_cast<size_t>(block_row[b]) * W;
|
|
const size_t hi = static_cast<size_t>(block_row[b + 1]) * W;
|
|
for (size_t i = lo; i < hi; i++) {
|
|
const int32_t c = cell_of[i];
|
|
if (c < 0 || clip_limit[c] < 0.0f || mean[i] > clip_limit[c])
|
|
continue;
|
|
b_count[c]++;
|
|
b_sum[c] += mean[i];
|
|
b_sum_sq[c] += static_cast<double>(mean[i]) * mean[i];
|
|
}
|
|
});
|
|
for (int c = 0; c < n_cells; c++) {
|
|
double s = 0, ss = 0;
|
|
int32_t n = 0;
|
|
for (int b = 0; b < BLOCKS; b++) {
|
|
const size_t k = static_cast<size_t>(b) * n_cells + c;
|
|
s += block_sum[k]; ss += block_sum_sq[k]; n += block_count[k];
|
|
}
|
|
sum[c] = s; sum_sq[c] = ss; count[c] = n;
|
|
}
|
|
}
|
|
|
|
// Radial profile: the median over the sectors that have a mean, on rings that are
|
|
// almost fully covered.
|
|
int usable_radial = 0;
|
|
for (int r = 0; r < RADIAL_BINS; r++) {
|
|
std::vector<float> present;
|
|
for (int s = 0; s < SECTORS; s++)
|
|
if (count[r * SECTORS + s] >= MIN_PIXELS_PER_CELL)
|
|
present.push_back(static_cast<float>(sum[r * SECTORS + s] / count[r * SECTORS + s]));
|
|
radial_ok[r] = static_cast<float>(present.size()) >= MIN_SECTOR_COVERAGE * SECTORS;
|
|
profile[r] = radial_ok[r] ? median_of(present) : 0.0f;
|
|
usable_radial += radial_ok[r];
|
|
}
|
|
if (usable_radial < MIN_USABLE_RADIAL_BINS)
|
|
return {};
|
|
// Central difference, so a bin next to a gap in the profile drops out with it. The test
|
|
// reads the ring BEFORE it, so it has to read the covered/not-covered flags as they were,
|
|
// not as this same loop has already rewritten them.
|
|
const std::vector<char> covered = radial_ok;
|
|
for (int r = 0; r < RADIAL_BINS; r++) {
|
|
const bool have = r > 0 && r + 1 < RADIAL_BINS && covered[r - 1] && covered[r] && covered[r + 1];
|
|
d_profile[r] = have ? (profile[r + 1] - profile[r - 1]) / (2 * d_tt) : 0.0f;
|
|
radial_ok[r] = have;
|
|
}
|
|
|
|
// Per sector: regress (profile of the sector - common profile) on {g, g'}. The first
|
|
// coefficient is the sector's amplitude, the second its radial shift; only the shift
|
|
// is carried on.
|
|
std::vector<float> shift, weight, jacobian_x, jacobian_y;
|
|
for (int s = 0; s < SECTORS; s++) {
|
|
double a11 = 0, a12 = 0, a22 = 0, b1 = 0, b2 = 0;
|
|
double jx = 0, jy = 0;
|
|
int n = 0;
|
|
for (int r = 0; r < RADIAL_BINS; r++) {
|
|
const int c = r * SECTORS + s;
|
|
if (!radial_ok[r] || count[c] < MIN_PIXELS_PER_CELL)
|
|
continue;
|
|
const double g = profile[r], dg = d_profile[r];
|
|
const double y = sum[c] / count[c] - profile[r];
|
|
a11 += g * g; a12 += g * dg; a22 += dg * dg;
|
|
b1 += g * y; b2 += dg * y;
|
|
jx += sum_jx[c] / count_all[c];
|
|
jy += sum_jy[c] / count_all[c];
|
|
n++;
|
|
}
|
|
const double det = a11 * a22 - a12 * a12;
|
|
if (n < MIN_USABLE_RADIAL_BINS || det <= 0)
|
|
continue;
|
|
const double amplitude = (a22 * b1 - a12 * b2) / det;
|
|
const double this_shift = (a11 * b2 - a12 * b1) / det;
|
|
// Residual sum of squares from the normal equations, without a second pass.
|
|
double residual = 0;
|
|
for (int r = 0; r < RADIAL_BINS; r++) {
|
|
const int c = r * SECTORS + s;
|
|
if (!radial_ok[r] || count[c] < MIN_PIXELS_PER_CELL)
|
|
continue;
|
|
const double e = sum[c] / count[c] - profile[r] - amplitude * profile[r] - this_shift * d_profile[r];
|
|
residual += e * e;
|
|
}
|
|
const double variance = residual / (n - 2) * (a11 / det);
|
|
if (!(variance > 0))
|
|
continue;
|
|
shift.push_back(static_cast<float>(this_shift));
|
|
weight.push_back(static_cast<float>(1.0 / variance));
|
|
jacobian_x.push_back(static_cast<float>(jx / n));
|
|
jacobian_y.push_back(static_cast<float>(jy / n));
|
|
}
|
|
if (static_cast<int>(shift.size()) < MIN_USABLE_SECTORS)
|
|
return {};
|
|
|
|
// shift_k = Jx_k dx + Jy_k dy, robustified so one bad sector cannot carry the answer.
|
|
std::vector<float> w = weight;
|
|
double c11 = 0, c12 = 0, c22 = 0;
|
|
for (int round = 0; round < 3; round++) {
|
|
c11 = c12 = c22 = 0;
|
|
double r1 = 0, r2 = 0;
|
|
for (size_t k = 0; k < shift.size(); k++) {
|
|
c11 += w[k] * jacobian_x[k] * jacobian_x[k];
|
|
c12 += w[k] * jacobian_x[k] * jacobian_y[k];
|
|
c22 += w[k] * jacobian_y[k] * jacobian_y[k];
|
|
r1 += w[k] * jacobian_x[k] * shift[k];
|
|
r2 += w[k] * jacobian_y[k] * shift[k];
|
|
}
|
|
const double det = c11 * c22 - c12 * c12;
|
|
if (det <= 0)
|
|
return {};
|
|
step_x = static_cast<float>((c22 * r1 - c12 * r2) / det);
|
|
step_y = static_cast<float>((c11 * r2 - c12 * r1) / det);
|
|
std::vector<float> residual(shift.size());
|
|
for (size_t k = 0; k < shift.size(); k++)
|
|
residual[k] = shift[k] - jacobian_x[k] * step_x - jacobian_y[k] * step_y;
|
|
std::vector<float> absolute(residual.size());
|
|
for (size_t k = 0; k < residual.size(); k++) absolute[k] = std::abs(residual[k]);
|
|
const float scale = 1.4826f * median_of(absolute) + 1e-30f;
|
|
for (size_t k = 0; k < shift.size(); k++) {
|
|
const float t = residual[k] / (3 * scale);
|
|
w[k] = weight[k] / (1.0f + t * t);
|
|
}
|
|
}
|
|
|
|
double chi2 = 0;
|
|
for (size_t k = 0; k < shift.size(); k++) {
|
|
const double e = shift[k] - jacobian_x[k] * step_x - jacobian_y[k] * step_y;
|
|
chi2 += w[k] * e * e;
|
|
}
|
|
chi2 = std::max(chi2 / (shift.size() - 2), 1.0);
|
|
const double det = c11 * c22 - c12 * c12;
|
|
sigma_x = static_cast<float>(std::sqrt(c22 / det * chi2));
|
|
sigma_y = static_cast<float>(std::sqrt(c11 / det * chi2));
|
|
|
|
const bool cancels = std::hypot(step_x + previous_x, step_y + previous_y)
|
|
< std::min(std::hypot(step_x, step_y), std::hypot(previous_x, previous_y));
|
|
reversals = cancels ? reversals + 1 : 0;
|
|
previous_x = step_x;
|
|
previous_y = step_y;
|
|
|
|
beam_x += DAMPING * step_x;
|
|
beam_y += DAMPING * step_y;
|
|
if (std::hypot(step_x, step_y) < CONVERGED_PXL || reversals >= REVERSALS_AT_THE_FIXED_POINT)
|
|
break;
|
|
}
|
|
|
|
// A fit that leaves on the iteration cap has not converged - it was still walking when it ran
|
|
// out - and the precision of its last step is not what it knows the centre to. The step it
|
|
// still wanted to take is a floor under what is left, so report that instead: it turns a
|
|
// confidently wrong answer into one the caller's sigma gate refuses. A fit that stopped on
|
|
// CONVERGED_PXL or at its fixed point takes a step far under any sigma worth reporting, so
|
|
// this never touches those.
|
|
return BeamCenterEstimate{beam_x, beam_y,
|
|
std::max({sigma_x, sigma_y, std::hypot(step_x, step_y)})};
|
|
}
|
|
|
|
std::optional<BeamCenterEstimate>
|
|
FindBeamCenter(const DiffractionExperiment &experiment, const PixelMask &mask,
|
|
const std::vector<float> &mean, size_t nthreads, BeamCenterFFTResult *capture) {
|
|
// The image the capture scores: the projection with everything the mask excludes taken out.
|
|
// The beam stop is in that mask by the time this runs, and it is the part that matters - an
|
|
// opaque region on one side of the beam is centrosymmetric about ITS own centre, and where it
|
|
// is large enough that preference beats the background's.
|
|
std::vector<float> scored = mean;
|
|
const auto &pixel_mask = mask.GetMask(experiment);
|
|
if (pixel_mask.size() == scored.size())
|
|
for (size_t i = 0; i < scored.size(); i++)
|
|
if (pixel_mask[i] != 0)
|
|
scored[i] = NAN;
|
|
|
|
const BeamCenterFFTResult fft = FindBeamCenterFFT(experiment, scored);
|
|
if (capture)
|
|
*capture = fft;
|
|
if (fft.point.empty())
|
|
return FindBeamCenterFromBackground(experiment, mask, mean, nthreads);
|
|
|
|
const std::pair<float, float> seed{fft.point[0].beam_x_pxl, fft.point[0].beam_y_pxl};
|
|
if (auto refined = FindBeamCenterFromBackground(experiment, mask, mean, nthreads, seed))
|
|
return refined;
|
|
|
|
// The walk has nothing to fit AT THE CAPTURE. That is a statement about the capture as much as
|
|
// about the background: a centre far enough out leaves too few fully covered rings inside the
|
|
// fitted band for the profile to exist at all. Measured on one small-molecule set whose capture
|
|
// lands 536 px out - the walk from the file's centre answers there to 2.9 px, and taking the
|
|
// capture instead would have been the composition's only regression on the corpus. So ask the
|
|
// walk again where it used to be asked, and keep the capture only where neither start works.
|
|
if (auto from_file = FindBeamCenterFromBackground(experiment, mask, mean, nthreads))
|
|
return from_file;
|
|
return BeamCenterEstimate{seed.first, seed.second, BEAM_CENTER_CAPTURE_SIGMA_PXL};
|
|
}
|