diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index 4f3865c3..f7cb25c1 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/image_analysis/IndexAndRefine.cpp @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only +#include #include "IndexAndRefine.h" #include "bragg_integration/BraggIntegrate2D.h" diff --git a/image_analysis/pixel_refinement/PixelRefine.cpp b/image_analysis/pixel_refinement/PixelRefine.cpp index 98f0a7af..7fbc57a5 100644 --- a/image_analysis/pixel_refinement/PixelRefine.cpp +++ b/image_analysis/pixel_refinement/PixelRefine.cpp @@ -273,6 +273,42 @@ bool PredictedNode(const T *p0, const T *p1, const T *p2, return true; } +// Pulls a scalar parameter towards an expected value with a fixed weight (the +// data-scaled prior). Identical in spirit to ScaleOnTheFly's regularizer: it is what +// keeps the per-image scale G from wandering on weakly-constrained images and +// scrambling the cross-image merge. +struct ScalarRegularizer { + ScalarRegularizer(double weight, double expected) : weight(weight), expected(expected) {} + template + bool operator()(const T *p, T *residual) const { + residual[0] = T(weight) * (p[0] - T(expected)); + return true; + } + double weight; + double expected; +}; + +// Anchors the orientation (angle-axis vector) to its pre-refinement value with a +// data-scaled weight. Without it the three orientation DOF chase the sparse signal +// (and the few noisy background pixels) and the per-image intensities collapse; +// with it the fit can only make a small, signal-supported sub-spot correction - the +// push that brings slightly-misaligned high-resolution reflections onto their +// shoeboxes. Mirrors the G/B regularizers in ScaleOnTheFly. +struct OrientationRegularizer { + OrientationRegularizer(double weight, const double prior[3]) : weight(weight) { + for (int i = 0; i < 3; ++i) + prior_[i] = prior[i]; + } + template + bool operator()(const T *p0, T *residual) const { + for (int i = 0; i < 3; ++i) + residual[i] = T(weight) * (p0[i] - T(prior_[i])); + return true; + } + double weight; + double prior_[3]; +}; + } // namespace // --------------------------------------------------------------------------- @@ -551,33 +587,161 @@ void PixelRefine::BuildParameterBlocks(const PixelRefineData &data, } } -BraggPredictionSettings PixelRefine::BuildPredictionSettings(const PixelRefineData &data) const { - // Radial Ewald-acceptance band: predict exactly the reflections the merge's - // partiality floor would still keep, and no tighter. A reflection's partiality - // is exp(-s^2 / R0_eff^2) (s = excitation error), so it survives min_partiality - // when |s| <= R0 * sqrt(-ln(min_partiality)). Using that as the cutoff keeps - // prediction and the partiality cut consistent. The struct default (0.0005) is - // far narrower than R0 (~0.005), so previously most keepable reflections were - // never predicted and per-reflection multiplicity collapsed ~4x. - const double min_part = std::clamp(experiment.GetScalingSettings().GetMinPartiality(), 1e-6, 0.999); - const double r0 = std::max(data.R[0], 1e-4); - const double cutoff = r0 * std::sqrt(-std::log(min_part)); +template +void PixelRefine::SweepOrientationCell(const T *image, BraggPrediction &prediction, + PixelRefineData &data) const { + const int radius = data.shoebox_radius; + const double beam_x = data.geom.GetBeamX_pxl(); + const double beam_y = data.geom.GetBeamY_pxl(); + const auto qnan = std::numeric_limits::quiet_NaN(); - // Relative bandwidth (sigma of dlambda/lambda): the explicit PixelRefine model - // value if set, else the experiment's nominal bandwidth (FWHM -> sigma). >0 - // thickens the band radially at high resolution (the 1/d^2 pink-beam smear), - // matching the integrator and preventing the outer shells from being clipped. - float bw_sigma = static_cast(data.bandwidth); - if (bw_sigma <= 0.0f) - bw_sigma = experiment.GetBandwidthFWHM().value_or(0.0f) / 2.3548f; + // Box-sum minus local (perimeter) background, raw counts. NaN if the box runs + // off the detector or hits a masked/saturated pixel. + auto integrate = [&](double px, double py) -> double { + const int cx = static_cast(std::lround(px)); + const int cy = static_cast(std::lround(py)); + const int outer = radius + 1; + if (cx - outer < 0 || cy - outer < 0 || + cx + outer >= static_cast(xpixel) || cy + outer >= static_cast(ypixel)) + return qnan; + double sig = 0.0; + int nsig = 0; + std::vector ring; + ring.reserve((2 * outer + 1) * (2 * outer + 1)); + for (int y = cy - outer; y <= cy + outer; ++y) { + for (int x = cx - outer; x <= cx + outer; ++x) { + const T raw = image[static_cast(xpixel) * y + x]; + if (raw == std::numeric_limits::max()) + return qnan; + if (std::is_signed_v && raw == std::numeric_limits::min()) + return qnan; + const double v = static_cast(raw); + if (std::abs(x - cx) <= radius && std::abs(y - cy) <= radius) { + sig += v; + ++nsig; + } else { + ring.push_back(v); + } + } + } + if (ring.size() < 5) + return qnan; + return sig - nsig * MedianInPlace(ring); + }; - return BraggPredictionSettings{ + // Predict (wide band) and collect every reflection that has a reference value, + // with its detector radius. The full set is scored - the strong low-res spots + // anchor the CC, the weak high-res spots are what "appear" at the right cell. + DiffractionExperiment exp_iter = experiment; + exp_iter.BeamX_pxl(beam_x).BeamY_pxl(beam_y) + .DetectorDistance_mm(data.geom.GetDetectorDistance_mm()) + .PoniRot1_rad(data.geom.GetPoniRot1_rad()) + .PoniRot2_rad(data.geom.GetPoniRot2_rad()); + const BraggPredictionSettings settings{ .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(), - .ewald_dist_cutoff = static_cast(cutoff), + .ewald_dist_cutoff = static_cast(data.ewald_dist_cutoff), .max_hkl = 100, .centering = data.centering, - .bandwidth_sigma = bw_sigma + .bandwidth_sigma = static_cast(data.bandwidth) }; + const int nrefl = prediction.Calc(exp_iter, data.latt, settings); + const auto &predicted = prediction.GetReflections(); + + struct Matched { int h, k, l; double refI; }; + std::vector matched; + double r_max = 0.0, r_min = std::numeric_limits::max(); + for (int i = 0; i < nrefl; ++i) { + const auto &r = predicted[i]; + const auto it = reference_data.find(hkl_key_generator(r)); + if (it == reference_data.end()) + continue; + matched.push_back({r.h, r.k, r.l, it->second}); + const double dx = r.predicted_x - beam_x; + const double dy = r.predicted_y - beam_y; + const double rad = std::sqrt(dx * dx + dy * dy); + r_max = std::max(r_max, rad); + r_min = std::min(r_min, rad); + } + if (matched.size() < 20 || r_min <= 1.0 || r_max <= r_min) + return; // too little to anchor a meaningful sweep + + // CC of the box-summed intensities against the reference, over all matched hkls. + auto score = [&](const CrystalLattice &L) -> double { + const Coord A = L.Astar(), B = L.Bstar(), C = L.Cstar(); + double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0; + int n = 0; + for (const auto &m : matched) { + const Coord g = A * static_cast(m.h) + B * static_cast(m.k) + + C * static_cast(m.l); + const auto [x, y] = data.geom.RecipToDetector(g); + if (!std::isfinite(x) || !std::isfinite(y)) + continue; + const double I = integrate(x, y); + if (!std::isfinite(I)) + continue; + const double yv = m.refI; + sx += I; sy += yv; sxx += I * I; syy += yv * yv; sxy += I * yv; ++n; + } + if (n < 10) + return -2.0; + const double nd = n; + const double cov = sxy - sx * sy / nd; + const double vx = sxx - sx * sx / nd; + const double vy = syy - sy * sy / nd; + if (!(vx > 0.0 && vy > 0.0)) + return -2.0; + return cov / std::sqrt(vx * vy); + }; + + // Step = 1 px at the highest resolution. Range = assumed orientation/cell-scale + // uncertainty (a few px at high res), NOT the low-res 2 px cap: the latter is + // ~2*r_max/r_min px at high res - far too permissive, and lets the per-image CC + // overfit. Here the low-res spots barely move (stay anchored). + const double step = 1.0 / r_max; + const int n_rot = std::clamp( + static_cast(std::lround(data.sweep_max_deg * M_PI / 180.0 * r_max)), 1, 25); + const int n_scale = std::clamp( + static_cast(std::lround(data.sweep_max_cell_frac * r_max)), 1, 25); + const Coord axes[3] = {Coord(1, 0, 0), Coord(0, 1, 0), Coord(0, 0, 1)}; + + CrystalLattice best = data.latt; + double best_cc = score(best); + + for (int round = 0; round < 2; ++round) { + for (const auto &axis : axes) { + CrystalLattice axis_best = best; + double axis_cc = best_cc; + for (int i = -n_rot; i <= n_rot; ++i) { + if (i == 0) + continue; + CrystalLattice cand = best.Multiply(RotMatrix(static_cast(i * step), axis)); + const double cc = score(cand); + if (cc > axis_cc) { + axis_cc = cc; + axis_best = cand; + } + } + best = axis_best; + best_cc = axis_cc; + } + CrystalLattice scale_best = best; + double scale_cc = best_cc; + for (int i = -n_scale; i <= n_scale; ++i) { + if (i == 0) + continue; + const double s = 1.0 / (1.0 + i * step); // cell scale (1+eps) -> recip * 1/(1+eps) + CrystalLattice cand = best.Multiply(gemmi::Mat33(s, 0, 0, 0, s, 0, 0, 0, s)); + const double cc = score(cand); + if (cc > scale_cc) { + scale_cc = cc; + scale_best = cand; + } + } + best = scale_best; + best_cc = scale_cc; + } + + data.latt = best; } template @@ -587,10 +751,21 @@ void PixelRefine::Run(const T *image, data.solved = false; data.reflections.clear(); + // Global orientation + cell-scale sweep before the local LSQ, to recentre the + // high-resolution shoeboxes onto signal that small misalignments hide. + if (data.sweep_orientation) + SweepOrientationCell(image, prediction, data); + const double lambda = data.geom.GetWavelength_A(); const double pixel_size = data.geom.GetPixelSize_mm(); - const BraggPredictionSettings settings_prediction = BuildPredictionSettings(data); + const BraggPredictionSettings settings_prediction{ + .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(), + .ewald_dist_cutoff = static_cast(data.ewald_dist_cutoff), + .max_hkl = 100, + .centering = data.centering, + .bandwidth_sigma = static_cast(data.bandwidth) // relative Δλ/λ sigma + }; const int radius = data.shoebox_radius; const int bkg_outer_radius = std::max(radius + 1, data.bkg_outer_radius); @@ -628,6 +803,7 @@ void PixelRefine::Run(const T *image, double latt_vec0[3] = {0, 0, 0}; // orientation (Rodrigues) double latt_vec1[3] = {0, 0, 0}; // lengths double latt_vec2[3] = {0, 0, 0}; // angles (rad) + double orient_prior[3] = {0, 0, 0}; // pre-refinement orientation (regularization anchor) const bool eval_only = (data.max_iterations <= 0); const int n_iter = std::max(1, data.max_iterations); @@ -694,17 +870,32 @@ void PixelRefine::Run(const T *image, const double Iobs = static_cast(image[npixel]); // raw counts - // Per-pixel variance: Poisson noise of the raw counts. - double var = std::max(Iobs, 0.0); - if (!(var > 1.0)) - var = 1.0; + // Variance for the fit weight. Weighting by the observed count + // (var = Iobs) lets down-fluctuated background pixels carry the + // largest 1/sqrt(var) weight, which biases the fit towards "no + // signal" and drove the per-image scale G to 0 on weak images + // (collapsing the merge). Use the local background as the + // (background-limited) variance, constant over the shoebox - the + // same de-biasing applied to the extraction. + double var = std::max(Ibkg, 1.0); + double weight = 1.0 / std::sqrt(var); + + // Signal-weighting: down-weight pixels far from the predicted spot + // centre so the empty shoebox corners cannot dilute or destabilise + // the fit; the signal-bearing core drives the refined parameters. + if (data.fit_signal_sigma_pix > 0.0) { + const double dx = x - g.predicted_x; + const double dy = y - g.predicted_y; + const double s2 = data.fit_signal_sigma_pix * data.fit_signal_sigma_pix; + weight *= std::exp(-0.5 * (dx * dx + dy * dy) / s2); + } PixelObs obs{ .x = static_cast(x), .y = static_cast(y), .Iobs = Iobs, .Ibkg = Ibkg, - .weight = 1.0 / std::sqrt(var) + .weight = weight }; g.pixels.push_back(obs); } @@ -721,6 +912,12 @@ void PixelRefine::Run(const T *image, BuildParameterBlocks(data, beam, dist_mm, detector_rot, latt_vec0, latt_vec1, latt_vec2); + // Anchor for orientation regularization = the spot-centroid orientation we + // started from (captured before any pixel-level refinement moved it). + if (iter == 0) + for (int i = 0; i < 3; ++i) + orient_prior[i] = latt_vec0[i]; + // ---- 4. Build the problem --------------------------------------------- // One residual block per shoebox (N residuals), so the expensive // per-reflection node geometry is evaluated once per reflection instead @@ -753,8 +950,21 @@ void PixelRefine::Run(const T *image, data.residual_count = residual_pixels; // ---- 5. Constrain / bound parameter blocks ---------------------------- - if (!data.refine_orientation) + if (!data.refine_orientation) { problem.SetParameterBlockConstant(latt_vec0); + } else if (data.orient_reg_sigma_deg > 0.0) { + // Anchor orientation to its spot-centroid prior. The weight is scaled to + // the *pixel* data term (sqrt(n_pixels)/sigma_rad), not the reflection + // count - the data has one residual per shoebox pixel, so a reflection- + // scaled prior (~50x too weak) was simply not felt. At a misorientation of + // orient_reg_sigma_deg the prior matches the data, so the fit only moves + // further when the pixels strongly agree it should. + const double sigma_rad = std::max(data.orient_reg_sigma_deg * M_PI / 180.0, 1e-9); + const double w = std::sqrt(static_cast(residual_pixels)) / sigma_rad; + auto *reg = new ceres::AutoDiffCostFunction( + new OrientationRegularizer(w, orient_prior)); + problem.AddResidualBlock(reg, nullptr, latt_vec0); + } if (!data.refine_unit_cell) { problem.SetParameterBlockConstant(latt_vec1); @@ -789,10 +999,20 @@ void PixelRefine::Run(const T *image, } } - if (data.refine_scale) + if (data.refine_scale) { problem.SetParameterLowerBound(&data.scale_factor, 0, 0.0); - else + // Regularize G towards 1 so weakly-constrained images cannot wander + // (an unconstrained 1/G is what collapsed the cross-image merge). Weight + // scaled to the pixel data term (n_pixels), not the reflection count. + if (data.scale_reg_sigma > 0.0) { + const double w = std::sqrt(static_cast(residual_pixels) / data.scale_reg_sigma); + auto *reg = new ceres::AutoDiffCostFunction( + new ScalarRegularizer(w, 1.0)); + problem.AddResidualBlock(reg, nullptr, &data.scale_factor); + } + } else { problem.SetParameterBlockConstant(&data.scale_factor); + } if (!data.refine_B) problem.SetParameterBlockConstant(&data.B_factor); @@ -939,9 +1159,15 @@ void PixelRefine::Run(const T *image, continue; const double Iobs = static_cast(image[np]); // raw counts - double v = std::max(Iobs, 0.0); // Poisson variance - if (!(v > 1.0)) - v = 1.0; + // Variance for the profile-fit weights. Weighting by the *observed* + // per-pixel count (v = Iobs) biases the amplitude negative: a + // down-fluctuated background pixel gets the smallest v and hence the + // largest 1/v weight, so num = sum P_t (Iobs - Ibkg)/v is pulled below + // zero - worst where the signal is weakest, i.e. the high-resolution + // shells (the negative we see there). For background-limited + // reflections the variance is the local background, constant over the + // shoebox, so use that instead of the observed count. + double v = std::max(g.Ibkg, 1.0); // Peak-normalized radial factor (the partiality), in (0,1]. The // bandwidth-broadened radial width matches the model in Model(). @@ -1069,7 +1295,13 @@ std::vector PixelRefine::PredictImage(const T *image, .PoniRot1_rad(data.geom.GetPoniRot1_rad()) .PoniRot2_rad(data.geom.GetPoniRot2_rad()); - const BraggPredictionSettings settings_prediction = BuildPredictionSettings(data); + const BraggPredictionSettings settings_prediction{ + .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(), + .ewald_dist_cutoff = static_cast(data.ewald_dist_cutoff), + .max_hkl = 100, + .centering = data.centering, + .bandwidth_sigma = static_cast(data.bandwidth) // relative Δλ/λ sigma + }; const int nrefl = prediction.Calc(exp_iter, data.latt, settings_prediction); const auto &predicted = prediction.GetReflections(); const auto spot_mask = BuildSpotMask(predicted, nrefl, xpixel, ypixel, radius); @@ -1158,7 +1390,13 @@ std::vector PixelRefine::ChiSquaredImage(const T *image, .PoniRot1_rad(data.geom.GetPoniRot1_rad()) .PoniRot2_rad(data.geom.GetPoniRot2_rad()); - const BraggPredictionSettings settings_prediction = BuildPredictionSettings(data); + const BraggPredictionSettings settings_prediction{ + .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(), + .ewald_dist_cutoff = static_cast(data.ewald_dist_cutoff), + .max_hkl = 100, + .centering = data.centering, + .bandwidth_sigma = static_cast(data.bandwidth) + }; const int nrefl = prediction.Calc(exp_iter, data.latt, settings_prediction); const auto &predicted = prediction.GetReflections(); const auto spot_mask = BuildSpotMask(predicted, nrefl, xpixel, ypixel, radius); diff --git a/image_analysis/pixel_refinement/PixelRefine.h b/image_analysis/pixel_refinement/PixelRefine.h index ffcc08bd..fc4acabd 100644 --- a/image_analysis/pixel_refinement/PixelRefine.h +++ b/image_analysis/pixel_refinement/PixelRefine.h @@ -129,7 +129,50 @@ struct PixelRefineData { bool refine_detector_angles = false; bool refine_scale = true; bool refine_B = false; - bool refine_R = true; + bool refine_R = false; // per-image R refinement is unstable on sparse + // stills; R is held at its nominal value + + // Orientation refinement is anchored to the pre-refinement (spot-centroid) + // orientation with weight sqrt(n_refl)/sigma, so the pixel fit can only nudge + // the orientation by ~this many degrees before the prior pushes back. This is + // what turns the (otherwise overfitting) 3-DOF orientation refinement into a + // small, signal-supported sub-spot correction. Larger => freer; very large + // approaches the unregularized (collapsing) fit. ~1 deg gives the best CCref here; + // beyond ~2 deg the per-image fit overfits and the merge collapses. + double orient_reg_sigma_deg = 1.0; + + // Signal-weighting of the *fit* residuals: each pixel's weight is multiplied by + // a detector-space Gaussian exp(-r^2/2 sigma^2) centred on the predicted spot, so + // the many empty shoebox-corner pixels stop diluting (and destabilising) the fit + // and the signal-bearing core drives the refined scale/R. <= 0 disables (uniform). + double fit_signal_sigma_pix = 1.5; + + // Per-image scale G is regularized towards 1 with weight sqrt(n_refl/scale_reg_sigma) + // (mirrors ScaleOnTheFly). Without this the unconstrained G wanders on weakly + // measured images and 1/G scrambles the cross-image merge. <= 0 disables. + double scale_reg_sigma = 2.0; + + // Radial Ewald-sphere acceptance band for prediction (A^-1): a reflection is + // given a shoebox when ||S|-1/lambda| <= this. The narrow default predicts only + // reflections already on the Ewald sphere; widening it (towards the integrator's + // ~2-3x profile radius) lets in the slightly-misaligned high-resolution + // reflections - more multiplicity, and something for orientation refinement to + // actually centre. Safe to widen only with the per-image fit kept well-behaved + // (de-biased variance + signal-weighting + regularization, all default here). + double ewald_dist_cutoff = 0.0020; + + // Pre-LSQ global orientation+cell sweep (maximises CC vs reference over the + // strongest reflections). Bounds are derived from the detector geometry: the + // step moves the highest-resolution spot by 1 px, the range moves the lowest- + // resolution spot by ~2 px (rotation) / ~1 px (cell scale), so the low-res + // XtalOptimizer solution is preserved while high-res spots are recentred. + bool sweep_orientation = false; + // Sweep half-range as the *orientation uncertainty* (degrees) and cell-scale + // uncertainty (fraction). These set how far the highest-resolution spot may move + // (a few px); low-res spots barely move and stay anchored. Keep small - a large + // range lets the per-image CC overfit and degrades the merge. + double sweep_max_deg = 0.15; + double sweep_max_cell_frac = 0.003; double max_time_s = 5.0; int shoebox_radius = 3; // half-size of the per-reflection signal box (peak region that enters the fit) @@ -164,12 +207,14 @@ class PixelRefine { double detector_rot[2], double latt_vec0[3], double latt_vec1[3], double latt_vec2[3]) const; - // Bragg-prediction settings shared by Run/PredictImage/ChiSquaredImage so all - // three predict an identical reflection set. Critically it sets the radial - // Ewald-acceptance band (ewald_dist_cutoff) and the bandwidth broadening - see - // the definition for why these must track the partiality model, not the - // BraggPredictionSettings struct defaults. - BraggPredictionSettings BuildPredictionSettings(const PixelRefineData &data) const; + // Global orientation + uniform cell-scale sweep run before the LSQ. Re-projects + // the strongest reference reflections through candidate lattices and keeps the + // one maximising CC vs the reference intensities (coordinate descent over the 3 + // Rodrigues axes + cell scale, within geometry-derived pixel bounds). Writes the + // refined orientation/cell back into data.latt. See PixelRefineData::sweep_*. + template + void SweepOrientationCell(const T *image, BraggPrediction &prediction, + PixelRefineData &data) const; public: PixelRefine(const DiffractionExperiment &experiment, const std::vector &reference); diff --git a/tools/jfjoch_process.cpp b/tools/jfjoch_process.cpp index 5876418b..61ea164d 100644 --- a/tools/jfjoch_process.cpp +++ b/tools/jfjoch_process.cpp @@ -12,6 +12,9 @@ #include #include #include +#include +#include +#include #include "../reader/JFJochHDF5Reader.h" #include "../common/Logger.h" @@ -968,6 +971,27 @@ int main(int argc, char **argv) { logger.Info("Reference provided: per-image live scaling already applied; merging directly"); } + // --- Phase 1 diagnostic: distribution of the per-image scale CC vs the + // reference. High per-image CC with low merged CC1/2 => each image is fine + // and the cross-image merge is the problem; low per-image CC => the + // per-image extraction itself is noise. Logged for any reference run. + { + std::vector ccs; + for (const auto &i : indexer.GetIntegrationOutcome()) + if (i.image_scale_cc && std::isfinite(*i.image_scale_cc)) + ccs.push_back(*i.image_scale_cc); + if (!ccs.empty()) { + std::sort(ccs.begin(), ccs.end()); + const double mean = std::accumulate(ccs.begin(), ccs.end(), 0.0) / ccs.size(); + auto q = [&](double f) { return ccs[std::min(ccs.size() - 1, static_cast(f * ccs.size()))]; }; + logger.Info("Per-image scale CC vs reference: n={} mean={:.3f} median={:.3f} " + "p10={:.3f} p90={:.3f} min={:.3f} max={:.3f}", + ccs.size(), mean, q(0.5), q(0.1), q(0.9), ccs.front(), ccs.back()); + } else { + logger.Info("Per-image scale CC vs reference: none available"); + } + } + auto merge_start = std::chrono::steady_clock::now(); MergeOnTheFly merge_engine(experiment);