PixelRefine: R1 multiplier (fix tight-Term2 regression) + env-gated orientation refine

Full-jet R-free showed the factored model regressed vs the traditional pipeline
(R-free 0.34 vs 0.26, CCref 58 vs 76), traced to Term 2 using the raw measured
profile width: it is too tight (~0.002 A^-1), so the template sits off the ~0.4 px
centroid-floor scatter and underestimates the intensity (box-sum is immune). Fix,
XDS-style (integrate over ~6 sigma): multiply the measured R1 by r1_multiplier
(default 6) before use. On the jet at 1.3 A this recovers CCref 55.2 -> 59.3
(~ traditional's 60.5); crystal 2 -0.7. Tunable via env PR_RMULT for R-free.

Also reinstates the pre-factored per-image orientation refinement (per-pixel
ShoeboxResidual against the shoebox, regularised to the spot-centroid orientation,
re-predict), behind env PR_ORIENT (off by default), to A/B its effect on serial-data
R-free. On the jet CCref it is a no-op (XtalOptimizer orientation already at the
optimum), but CCref is a weak R-free proxy here, so it is left for R-free validation.

NOTE: PR_RMULT / PR_ORIENT are temporary diagnostic env knobs; the final multiplier
value and the orientation-refine decision are to be fixed against R-free, then baked.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

image_analysis/IndexAndRefine.cpp

@@ -448,6 +448,8 @@ bool IndexAndRefine::PixelRefineIntegrate(DataMessage &msg,
 
     // Signal-box radius from the shared integration setting (same knob as BraggIntegrate2D).
     prd.shoebox_radius = static_cast<int>(std::lround(experiment.GetBraggIntegrationSettings().GetR1()));
+    if (const char *m = std::getenv("PR_RMULT")) prd.r1_multiplier = std::stod(m); // TEMP: Term-2 R1 multiplier sweep
+    if (std::getenv("PR_ORIENT")) prd.refine_orientation = true; // A/B: per-image orientation refinement
 
     std::vector<uint8_t> buffer;
     const uint8_t *ptr = image.GetUncompressedPtr(buffer);

image_analysis/pixel_refinement/PixelRefine.cpp

@@ -24,6 +24,7 @@ struct PixelObs {
     double x, y;          // detector pixel coordinate
     double Iobs;          // raw pixel value (signal + background)
     double Ibkg;          // local background estimate (per-shoebox level, raw counts)
+    double weight;        // 1/sigma_pixel (used only by the optional orientation refinement)
 };
 
 // One reflection together with the pixels of its shoebox.
@@ -336,6 +337,69 @@ struct IntensityResidual {
     double J, inv_sigma, partiality, pol, I_ref, inv_4d2;
 };
 
+// Per-shoebox per-pixel forward-model cost (raw counts), used ONLY by the optional
+// orientation refinement: I_pred = G*Itrue*B_term*P_radial*P_tang*pol + Ibkg. The
+// expensive node geometry is computed once per reflection; the cheap ObservedRecip +
+// Gaussian profile run per pixel. Shares ObservedRecip/PredictedNode with GeometryProbe.
+struct ShoeboxResidual {
+    ShoeboxResidual(const ReflGroup &g, double lambda, double pixel_size,
+                    gemmi::CrystalSystem symmetry)
+        : pixels(g.pixels), Itrue(g.Itrue), R_bw_sq(g.R_bw_sq), pol(g.pol),
+          exp_h(g.h), exp_k(g.k), exp_l(g.l),
+          inv_lambda(1.0 / lambda), pixel_size(pixel_size), symmetry(symmetry) {}
+    template<typename T>
+    bool operator()(const T *const *params, T *residual) const {
+        // 0 beam 1 dist 2 detector_rot 3 p0 4 p1 5 p2 6 scale 7 B 8 R
+        const T *beam = params[0]; const T *distance_mm = params[1]; const T *detector_rot = params[2];
+        const T *p0 = params[3]; const T *p1 = params[4]; const T *p2 = params[5];
+        const T *scale_factor = params[6]; const T *B = params[7]; const T *R = params[8];
+        if (R[0] < T(1e-10) || R[1] < T(1e-10))
+            return false;
+        Eigen::Matrix<T, 3, 1> e_pred_recip, n_radial;
+        T q_sq;
+        if (!PredictedNode(p0, p1, p2, exp_h, exp_k, exp_l, symmetry, inv_lambda,
+                           e_pred_recip, n_radial, q_sq))
+            return false;
+        const T B_term = ceres::exp(-B[0] * q_sq / T(4.0));
+        const T R0_eff_sq = R[0] * R[0] + T(R_bw_sq);
+        for (size_t i = 0; i < pixels.size(); ++i) {
+            const PixelObs &obs = pixels[i];
+            Eigen::Matrix<T, 3, 1> e_obs_recip;
+            ObservedRecip(beam, distance_mm, detector_rot, obs.x, obs.y, pixel_size, inv_lambda, e_obs_recip);
+            const Eigen::Matrix<T, 3, 1> delta_q = e_obs_recip - e_pred_recip;
+            const T eps_radial = delta_q.dot(n_radial);
+            const T eps_tang_sq = (delta_q - eps_radial * n_radial).squaredNorm();
+            const T P_radial = ceres::exp(-eps_radial * eps_radial / R0_eff_sq);
+            const T P_tang = ceres::exp(-eps_tang_sq / (R[1] * R[1])) / (T(M_PI) * R[1] * R[1]);
+            const T signal = scale_factor[0] * T(Itrue) * B_term * P_radial * P_tang * T(pol);
+            residual[i] = (signal + T(obs.Ibkg) - T(obs.Iobs)) * T(obs.weight);
+        }
+        return true;
+    }
+    std::vector<PixelObs> pixels;
+    const double Itrue, R_bw_sq, pol;
+    const double exp_h, exp_k, exp_l;
+    const double inv_lambda, pixel_size;
+    gemmi::CrystalSystem symmetry;
+};
+
+// Anchors the orientation (angle-axis vector) to its prior with a data-scaled weight, so
+// the per-image fit can only make a small, signal-supported sub-spot correction.
+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
 
 PixelRefine::PixelRefine(const DiffractionExperiment &experiment,
@@ -443,69 +507,154 @@ void PixelRefine::Run(const T *image,
     double latt_vec0[3], latt_vec1[3], latt_vec2[3];
     BuildParameterBlocks(data, beam, dist_mm, detector_rot, latt_vec0, latt_vec1, latt_vec2);
 
-    // ---- 1. Predict shoeboxes for the current geometry ------------------------
-    DiffractionExperiment exp_iter = experiment;
-    exp_iter.BeamX_pxl(data.geom.GetBeamX_pxl())
-            .BeamY_pxl(data.geom.GetBeamY_pxl())
-            .DetectorDistance_mm(data.geom.GetDetectorDistance_mm())
-            .PoniRot1_rad(data.geom.GetPoniRot1_rad())
-            .PoniRot2_rad(data.geom.GetPoniRot2_rad());
-    const int nrefl = prediction.Calc(exp_iter, data.latt, settings_prediction);
-
-    // ---- 2. Collect per-reflection shoebox pixels + local background ----------
-    // GetReflections() returns the full pre-sized buffer; only the first nrefl
-    // entries are valid for this image. A spot-core mask over ALL predictions keeps
-    // each reflection's background ring from picking up a neighbour's signal.
-    const auto &predicted = prediction.GetReflections();
-    const auto spot_mask = BuildSpotMask(predicted, nrefl, xpixel, ypixel, radius);
-
+    // ---- 1-2. Predict shoeboxes + collect pixels (a lambda so it can be re-run after
+    //           an orientation refinement moves the predicted positions). ----------
+    // A spot-core mask over ALL predictions keeps each reflection's background ring from
+    // picking up a neighbour's signal. Pixels carry a fit weight (background-limited
+    // variance, signal-weighted toward the predicted centre) used only by the optional
+    // orientation refinement - the factored integration weights by v = Ibkg directly.
     std::vector<ReflGroup> groups;
-    for (int ri = 0; ri < nrefl; ++ri) {
-        const auto &refl = predicted[ri];
-        const auto hkl = hkl_key_generator(refl);
-        if (!reference_data.contains(hkl))
-            continue;
-
-        // Local flat background from the ring around the shoebox (raw counts). If we
-        // cannot estimate a clean local background the reflection is dropped, exactly
-        // as BraggIntegrate2D marks it unobserved when too few background pixels survive.
-        double Ibkg = 0.0;
-        if (!EstimateLocalBackground(image, spot_mask, xpixel, ypixel,
-                                     refl.predicted_x, refl.predicted_y,
-                                     radius, bkg_outer_radius, Ibkg))
-            continue;
-
-        ReflGroup g;
-        g.h = refl.h;
-        g.k = refl.k;
-        g.l = refl.l;
-        g.d = refl.d;
-        g.Itrue = reference_data[hkl];
-        g.R_bw_sq = bandwidth_radial_sq(refl.d);
-        g.pol = polarization(refl.predicted_x, refl.predicted_y);
-        g.Ibkg = Ibkg;
-        g.predicted_x = refl.predicted_x;
-        g.predicted_y = refl.predicted_y;
-
-        const auto box = ShoeboxBounds(refl.predicted_x, refl.predicted_y, radius, xpixel, ypixel);
-        for (int y = box.min_y; y <= box.max_y; ++y) {
-            for (int x = box.min_x; x <= box.max_x; ++x) {
-                const size_t npixel = xpixel * y + x;
-                // Skip sentinel (masked / saturated) pixels.
-                if (image[npixel] == std::numeric_limits<T>::max())
-                    continue;
-                if (std::is_signed_v<T> && (image[npixel] == std::numeric_limits<T>::min()))
-                    continue;
-                g.pixels.push_back({static_cast<double>(x), static_cast<double>(y),
-                                    static_cast<double>(image[npixel]), Ibkg});
+    auto build_groups = [&]() {
+        groups.clear();
+        DiffractionExperiment exp_iter = experiment;
+        exp_iter.BeamX_pxl(data.geom.GetBeamX_pxl())
+                .BeamY_pxl(data.geom.GetBeamY_pxl())
+                .DetectorDistance_mm(data.geom.GetDetectorDistance_mm())
+                .PoniRot1_rad(data.geom.GetPoniRot1_rad())
+                .PoniRot2_rad(data.geom.GetPoniRot2_rad());
+        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);
+        for (int ri = 0; ri < nrefl; ++ri) {
+            const auto &refl = predicted[ri];
+            const auto hkl = hkl_key_generator(refl);
+            if (!reference_data.contains(hkl))
+                continue;
+            double Ibkg = 0.0;
+            if (!EstimateLocalBackground(image, spot_mask, xpixel, ypixel,
+                                         refl.predicted_x, refl.predicted_y,
+                                         radius, bkg_outer_radius, Ibkg))
+                continue;
+            ReflGroup g;
+            g.h = refl.h;
+            g.k = refl.k;
+            g.l = refl.l;
+            g.d = refl.d;
+            g.Itrue = reference_data[hkl];
+            g.R_bw_sq = bandwidth_radial_sq(refl.d);
+            g.pol = polarization(refl.predicted_x, refl.predicted_y);
+            g.Ibkg = Ibkg;
+            g.predicted_x = refl.predicted_x;
+            g.predicted_y = refl.predicted_y;
+            const auto box = ShoeboxBounds(refl.predicted_x, refl.predicted_y, radius, xpixel, ypixel);
+            for (int y = box.min_y; y <= box.max_y; ++y) {
+                for (int x = box.min_x; x <= box.max_x; ++x) {
+                    const size_t npixel = xpixel * y + x;
+                    if (image[npixel] == std::numeric_limits<T>::max())
+                        continue;
+                    if (std::is_signed_v<T> && (image[npixel] == std::numeric_limits<T>::min()))
+                        continue;
+                    double weight = 1.0 / std::sqrt(std::max(Ibkg, 1.0));
+                    if (data.fit_signal_sigma_pix > 0.0) {
+                        const double dx = x - refl.predicted_x, dy = y - refl.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);
+                    }
+                    g.pixels.push_back({static_cast<double>(x), static_cast<double>(y),
+                                        static_cast<double>(image[npixel]), Ibkg, weight});
+                }
             }
+            if (!g.pixels.empty())
+                groups.push_back(std::move(g));
         }
-        if (!g.pixels.empty())
-            groups.push_back(std::move(g));
-    }
+    };
+    build_groups();
     if (groups.empty())
         return;
 
+    // ---- Optional per-image orientation refinement (off by default) -----------------
+    // The pre-factored per-pixel step: refine the orientation (and a nuisance scale)
+    // against the shoebox pixels via ShoeboxResidual, regularised to the spot-centroid
+    // orientation, then re-predict. Dropped from the factored model (geometry fixed to
+    // XtalOptimizer); behind PixelRefineData::refine_orientation to A/B its R-free effect.
+    if (data.refine_orientation) {
+        const double orient_prior[3] = {latt_vec0[0], latt_vec0[1], latt_vec0[2]};
+        ceres::Problem oprob;
+        size_t npix = 0;
+        for (const auto &g : groups) {
+            auto *cost = new ceres::DynamicAutoDiffCostFunction<ShoeboxResidual>(
+                new ShoeboxResidual(g, lambda, pixel_size, data.crystal_system));
+            for (int b : {2, 1, 2, 3, 3, 3, 1, 1, 2})
+                cost->AddParameterBlock(b);
+            cost->SetNumResiduals(static_cast<int>(g.pixels.size()));
+            oprob.AddResidualBlock(cost, nullptr, beam, &dist_mm, detector_rot,
+                                   latt_vec0, latt_vec1, latt_vec2,
+                                   &data.scale_factor, &data.B_factor, data.R);
+            npix += g.pixels.size();
+        }
+        // Refine only orientation (latt_vec0) + the nuisance scale G; everything else fixed.
+        oprob.SetParameterBlockConstant(beam);
+        oprob.SetParameterBlockConstant(&dist_mm);
+        oprob.SetParameterBlockConstant(detector_rot);
+        oprob.SetParameterBlockConstant(latt_vec1);
+        oprob.SetParameterBlockConstant(latt_vec2);
+        oprob.SetParameterBlockConstant(&data.B_factor);
+        oprob.SetParameterBlockConstant(data.R);
+        oprob.SetParameterLowerBound(&data.scale_factor, 0, 0.0);
+        if (data.orient_reg_sigma_deg > 0.0 && npix > 0) {
+            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>(npix)) / sigma_rad;
+            oprob.AddResidualBlock(new ceres::AutoDiffCostFunction<OrientationRegularizer, 3, 3>(
+                new OrientationRegularizer(w, orient_prior)), nullptr, latt_vec0);
+        }
+        if (data.scale_reg_sigma > 0.0) {
+            const double w = std::sqrt(static_cast<double>(groups.size()) / data.scale_reg_sigma);
+            oprob.AddResidualBlock(new ceres::AutoDiffCostFunction<ScalarRegularizer, 1, 1>(
+                new ScalarRegularizer(w, 1.0)), nullptr, &data.scale_factor);
+        }
+        ceres::Solver::Options oopt;
+        oopt.linear_solver_type = ceres::DENSE_QR;
+        oopt.logging_type = ceres::LoggingType::SILENT;
+        oopt.minimizer_progress_to_stdout = false;
+        oopt.max_solver_time_in_seconds = data.max_time_s;
+        oopt.num_threads = 1;
+        ceres::Solver::Summary osum;
+        ceres::Solve(oopt, &oprob, &osum);
+
+        // Write the refined orientation back into data.latt (cell held at latt_vec1/2),
+        // then re-predict + rebuild groups at the new orientation. Scale is reset; Term 1
+        // re-fits it properly below.
+        switch (data.crystal_system) {
+            case gemmi::CrystalSystem::Tetragonal:
+                latt_vec1[1] = latt_vec1[0];
+                data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, M_PI/2, M_PI/2);
+                break;
+            case gemmi::CrystalSystem::Cubic:
+                latt_vec1[1] = latt_vec1[0];
+                latt_vec1[2] = latt_vec1[0];
+                data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, M_PI/2, M_PI/2);
+                break;
+            case gemmi::CrystalSystem::Hexagonal:
+                latt_vec1[1] = latt_vec1[0];
+                data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, M_PI/2, 2.0*M_PI/3.0);
+                break;
+            case gemmi::CrystalSystem::Monoclinic:
+                data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, latt_vec2[0], M_PI/2);
+                break;
+            case gemmi::CrystalSystem::Orthorhombic:
+                data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1, M_PI/2, M_PI/2, M_PI/2);
+                break;
+            default:
+                data.latt = AngleAxisAndCellToLattice(latt_vec0, latt_vec1,
+                                                      latt_vec2[0], latt_vec2[1], latt_vec2[2]);
+                break;
+        }
+        data.scale_factor = 1.0;
+        build_groups();
+        if (groups.empty())
+            return;
+    }
+
     // ---- 3. Term 2: per-resolution tangential profile width R1 ----------------
     // R1 = sqrt(2*<eps_t^2>) from the intensity-weighted tangential second moment of
     // the strong spots, binned by resolution (low res small spots, high res larger).
@@ -549,7 +698,7 @@ void PixelRefine::Run(const T *image,
         std::vector<double> bin_R1(n_bins, data.R[1]);
         for (int b = 0; b < n_bins; ++b)
             if (bin_M2[b].size() >= 5) {
-                const double r1 = std::sqrt(2.0 * std::max(MedianInPlace(bin_M2[b]), 0.0));
+                const double r1 = data.r1_multiplier * std::sqrt(2.0 * std::max(MedianInPlace(bin_M2[b]), 0.0));
                 if (std::isfinite(r1) && r1 > 1e-4)
                     bin_R1[b] = std::clamp(r1, 1e-4, 0.05);
             }

image_analysis/pixel_refinement/PixelRefine.h

@@ -67,6 +67,22 @@ struct PixelRefineData {
                                   // tangential profile width before Term 2 measures it (A^-1)
     bool   refine_B = true;       // refine the per-image B-factor along with G
 
+    // Term 2 measures the physical tangential width R1, but the *integration* profile must
+    // be generous (XDS-style: integrate over ~6-10 sigma), or a tight template centred on
+    // the prediction sits off the ~0.4 px centroid-floor scatter and underestimates the
+    // intensity. The measured R1 is multiplied by this before use. 1.0 = the raw physical
+    // width (too tight - regressed R-free on the serial jet); ~6 recovers it (closes the
+    // CCref gap to box-sum). Tune the final value against R-free.
+    double r1_multiplier = 6.0;
+
+    // Optional per-image orientation refinement (env-gated, off by default): the
+    // pre-factored per-pixel step that refines the crystal orientation against the
+    // shoebox pixels, regularised to the spot-centroid orientation, before integration.
+    // Reinstated to A/B its effect on serial-data R-free vs the fixed-geometry default.
+    bool   refine_orientation = false;
+    double orient_reg_sigma_deg = 1.0;  // orientation-prior strength (deg)
+    double fit_signal_sigma_pix = 1.5;  // signal-weighting sigma for the orientation fit (px)
+
     // Relative X-ray bandwidth (sigma of dlambda/lambda), e.g. ~0.004 for a 1% FWHM
     // DMM, ~1e-4 for Si(111). Adds a resolution-dependent radial broadening to R[0].
     // 0 = monochromatic (the term switches off entirely).