PixelRefine: Results seem to be much better

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<typename T>
+    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<typename T>
+    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<class T>
+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<double>::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<float>(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<int>(std::lround(px));
+        const int cy = static_cast<int>(std::lround(py));
+        const int outer = radius + 1;
+        if (cx - outer < 0 || cy - outer < 0 ||
+            cx + outer >= static_cast<int>(xpixel) || cy + outer >= static_cast<int>(ypixel))
+            return qnan;
+        double sig = 0.0;
+        int nsig = 0;
+        std::vector<double> 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<size_t>(xpixel) * y + x];
+                if (raw == std::numeric_limits<T>::max())
+                    return qnan;
+                if (std::is_signed_v<T> && raw == std::numeric_limits<T>::min())
+                    return qnan;
+                const double v = static_cast<double>(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<float>(cutoff),
+        .ewald_dist_cutoff = static_cast<float>(data.ewald_dist_cutoff),
         .max_hkl = 100,
         .centering = data.centering,
-        .bandwidth_sigma = bw_sigma
+        .bandwidth_sigma = static_cast<float>(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> matched;
+    double r_max = 0.0, r_min = std::numeric_limits<double>::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<float>(m.h) + B * static_cast<float>(m.k)
+                            + C * static_cast<float>(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<int>(std::lround(data.sweep_max_deg * M_PI / 180.0 * r_max)), 1, 25);
+    const int n_scale = std::clamp(
+        static_cast<int>(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<float>(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<class T>
@@ -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<float>(data.ewald_dist_cutoff),
+        .max_hkl = 100,
+        .centering = data.centering,
+        .bandwidth_sigma = static_cast<float>(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<double>(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<double>(x),
                         .y = static_cast<double>(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<double>(residual_pixels)) / sigma_rad;
+            auto *reg = new ceres::AutoDiffCostFunction<OrientationRegularizer, 3, 3>(
+                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<double>(residual_pixels) / data.scale_reg_sigma);
+                auto *reg = new ceres::AutoDiffCostFunction<ScalarRegularizer, 1, 1>(
+                    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<double>(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 <I/sig> 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<float> 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<float>(data.ewald_dist_cutoff),
+        .max_hkl = 100,
+        .centering = data.centering,
+        .bandwidth_sigma = static_cast<float>(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<float> 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<float>(data.ewald_dist_cutoff),
+        .max_hkl = 100,
+        .centering = data.centering,
+        .bandwidth_sigma = static_cast<float>(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);