diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index 479d5ac1..34244a4a 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/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(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 buffer; const uint8_t *ptr = image.GetUncompressedPtr(buffer); diff --git a/image_analysis/pixel_refinement/PixelRefine.cpp b/image_analysis/pixel_refinement/PixelRefine.cpp index 75facde4..843d6360 100644 --- a/image_analysis/pixel_refinement/PixelRefine.cpp +++ b/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 + 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 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 e_obs_recip; + ObservedRecip(beam, distance_mm, detector_rot, obs.x, obs.y, pixel_size, inv_lambda, e_obs_recip); + const Eigen::Matrix 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 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 + 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 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::max()) - continue; - if (std::is_signed_v && (image[npixel] == std::numeric_limits::min())) - continue; - g.pixels.push_back({static_cast(x), static_cast(y), - static_cast(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::max()) + continue; + if (std::is_signed_v && (image[npixel] == std::numeric_limits::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(x), static_cast(y), + static_cast(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( + 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(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(npix)) / sigma_rad; + oprob.AddResidualBlock(new ceres::AutoDiffCostFunction( + new OrientationRegularizer(w, orient_prior)), nullptr, latt_vec0); + } + if (data.scale_reg_sigma > 0.0) { + const double w = std::sqrt(static_cast(groups.size()) / data.scale_reg_sigma); + oprob.AddResidualBlock(new ceres::AutoDiffCostFunction( + 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*) 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 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); } diff --git a/image_analysis/pixel_refinement/PixelRefine.h b/image_analysis/pixel_refinement/PixelRefine.h index dfc560c5..1e7a2dac 100644 --- a/image_analysis/pixel_refinement/PixelRefine.h +++ b/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).