diff --git a/image_analysis/scale_merge/Merge.cpp b/image_analysis/scale_merge/Merge.cpp index 54ec8d9eb..4f72271c7 100644 --- a/image_analysis/scale_merge/Merge.cpp +++ b/image_analysis/scale_merge/Merge.cpp @@ -15,6 +15,7 @@ #include "../../common/CorrelationCoefficient.h" #include "../../common/ResolutionShells.h" #include "../../common/Definitions.h" +#include "../../common/ParallelFor.h" #include "HKLKey.h" #include "RfreeFlags.h" #include "FrenchWilson.h" @@ -32,6 +33,27 @@ namespace { z = z ^ (z >> 31); return static_cast(z & 1ULL); } + + // Calls reduce(i, out) for the images in parallel, a block of them at a time, and then consume(i, out) + // for each image of the block in image order - the order that every floating-point sum, every + // half-set choice and every hash-map insertion of the merge depends on. The result is therefore the + // serial loop's, bit for bit, at any thread count. + template + void ReduceImagesInOrder(size_t n_images, size_t nthreads, Reduce reduce, Consume consume) { + constexpr size_t BLOCK_IMAGES = 4096; + std::vector> out(std::min(BLOCK_IMAGES, n_images)); + for (size_t first = 0; first < n_images; first += BLOCK_IMAGES) { + const int n = static_cast(std::min(BLOCK_IMAGES, n_images - first)); + ParallelChunks(n, nthreads, [&](int lo, int hi) { + for (int j = lo; j < hi; ++j) { + out[j].clear(); + reduce(first + j, out[j]); + } + }); + for (int j = 0; j < n; ++j) + consume(first + j, out[j]); + } + } } MergeOnTheFly::MergeOnTheFly(const DiffractionExperiment &x) @@ -57,14 +79,8 @@ MergeOnTheFly &MergeOnTheFly::ReferenceCell(const std::optional &cell) return *this; } -void MergeOnTheFly::AddImage(const IntegrationOutcome &outcome, int64_t image_id) { - std::unique_lock ul(merged_mutex); - - if (Mask(outcome)) - return; - - const int prefer = HalfForImage(image_id); - +size_t MergeOnTheFly::Contributions(const IntegrationOutcome &outcome, std::vector &out) const { + size_t rejected = 0; for (const auto &r: outcome.reflections) { if (generator.IsSystematicallyAbsent(r)) continue; @@ -94,24 +110,36 @@ void MergeOnTheFly::AddImage(const IntegrationOutcome &outcome, int64_t image_id const auto mit = reject_median_I.find(hkl_key); if (mit != reject_median_I.end() && std::fabs(I_corr - mit->second) > reject_nsigma * sigma_corr) { - ++reject_count; + ++rejected; continue; } } - auto it = accumulator.find(hkl_key); - if (it == accumulator.end()) - it = accumulator.emplace(hkl_key, MergeAccum{ - .h = hkl.plus ? hkl.h : -hkl.h, - .k = hkl.plus ? hkl.k : -hkl.k, - .l = hkl.plus ? hkl.l : -hkl.l, - }).first; - const float w = 1.0f / (sigma_corr * sigma_corr); - const float wI = w * I_corr; + out.push_back(Contribution{ + .key = hkl_key, + .h = hkl.plus ? hkl.h : -hkl.h, + .k = hkl.plus ? hkl.k : -hkl.k, + .l = hkl.plus ? hkl.l : -hkl.l, + .w = w, + .wI = w * I_corr, + .d = r.d, + .on_ice_ring = r.on_ice_ring + }); + } + return rejected; +} - it->second.sum_wI += wI; - it->second.sum_w += w; +void MergeOnTheFly::Accumulate(const std::vector &contributions, int64_t image_id) { + const int prefer = HalfForImage(image_id); + + for (const auto &c: contributions) { + auto it = accumulator.find(c.key); + if (it == accumulator.end()) + it = accumulator.emplace(c.key, MergeAccum{.h = c.h, .k = c.k, .l = c.l}).first; + + it->second.sum_wI += c.wI; + it->second.sum_w += c.w; // CC1/2 correlates the two half-set means, so a reflection whose observations all land in one // half has no second mean and drops out of the statistic entirely. Taking the image's own @@ -119,18 +147,44 @@ void MergeOnTheFly::AddImage(const IntegrationOutcome &outcome, int64_t image_id // reflection floor(n/2) against ceil(n/2), which is the split cctbx and XDS use. const int half = it->second.n_half[prefer] > it->second.n_half[1 - prefer] ? 1 - prefer : prefer; - it->second.sum_wI_half[half] += wI; - it->second.sum_w_half[half] += w; + it->second.sum_wI_half[half] += c.wI; + it->second.sum_w_half[half] += c.w; it->second.n_half[half]++; - if (r.on_ice_ring) + if (c.on_ice_ring) it->second.on_ice_ring = true; - if (!std::isfinite(it->second.d) && std::isfinite(r.d) && r.d > 0.0f) - it->second.d = r.d; + if (!std::isfinite(it->second.d) && std::isfinite(c.d) && c.d > 0.0f) + it->second.d = c.d; } } +void MergeOnTheFly::AddImage(const IntegrationOutcome &outcome, int64_t image_id) { + std::unique_lock ul(merged_mutex); + + if (Mask(outcome)) + return; + std::vector contributions; + reject_count += Contributions(outcome, contributions); + Accumulate(contributions, image_id); +} + +void MergeOnTheFly::AddImages(const std::vector &outcomes) { + std::unique_lock ul(merged_mutex); + + std::vector rejected(outcomes.size(), 0); + ReduceImagesInOrder( + outcomes.size(), nthreads, + [&](size_t i, std::vector &out) { + if (!Mask(outcomes[i])) + rejected[i] = Contributions(outcomes[i], out); + }, + [&](size_t i, const std::vector &out) { + reject_count += rejected[i]; + Accumulate(out, static_cast(i)); + }); +} + float MergeOnTheFly::CorrectedSigma(float I_corr, float sigma_corr, float image_scale_corr, float var_bkg, uint64_t hkl_key) const { @@ -188,29 +242,36 @@ void MergeOnTheFly::RefineErrorModel(const std::vector &outc // --- 1. Collect accepted, scaled observations grouped by symmetry-equivalent hkl, // applying exactly the filters AddImage uses. --- struct Obs { float I, sigma; }; + struct KeyedObs { uint64_t key; Obs obs; }; std::unordered_map> groups; - for (const auto &outcome: outcomes) { - if (Mask(outcome)) - continue; - for (const auto &r: outcome.reflections) { - if (generator.IsSystematicallyAbsent(r)) - continue; - if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) - continue; - if (!AcceptReflection(r, high_resolution_limit, low_resolution_limit)) - continue; - if (exclude_ice_rings && r.on_ice_ring) - continue; - if (r.partiality < min_partiality) - continue; - const float I_corr = r.I * r.image_scale_corr; - const float sigma_corr = r.sigma * r.image_scale_corr; - if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) - continue; - groups[generator(r).pack()].push_back({I_corr, sigma_corr}); - } - } + ReduceImagesInOrder( + outcomes.size(), nthreads, + [&](size_t i, std::vector &out) { + if (Mask(outcomes[i])) + return; + for (const auto &r: outcomes[i].reflections) { + if (generator.IsSystematicallyAbsent(r)) + continue; + if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) + continue; + if (!AcceptReflection(r, high_resolution_limit, low_resolution_limit)) + continue; + if (exclude_ice_rings && r.on_ice_ring) + continue; + if (r.partiality < min_partiality) + continue; + const float I_corr = r.I * r.image_scale_corr; + const float sigma_corr = r.sigma * r.image_scale_corr; + if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) + continue; + out.push_back({generator(r).pack(), {I_corr, sigma_corr}}); + } + }, + [&](size_t, const std::vector &out) { + for (const auto &o: out) + groups[o.key].push_back(o.obs); + }); // --- 2. One global pool of (sigma^2, ^2, bias-corrected squared deviation). For an // observation in a group of n, the residual from the inverse-variance mean has @@ -337,7 +398,7 @@ void MergeOnTheFly::RefineErrorModel(const std::vector &outc error_model_chi2 = chi2.empty() ? 0.0 : median(chi2) / CHI2_1_MEDIAN; } -bool MergeOnTheFly::Mask(const IntegrationOutcome &outcome) { +bool MergeOnTheFly::Mask(const IntegrationOutcome &outcome) const { if (reference_cell) { auto cell = outcome.latt.GetUnitCell(); if (!cell.is_close(*reference_cell, @@ -355,12 +416,20 @@ bool MergeOnTheFly::Mask(const IntegrationOutcome &outcome) { return false; } -std::vector MergeOnTheFly::ExportReflections() { +std::vector MergeOnTheFly::ExportIntensities() { std::unique_lock ul(merged_mutex); std::vector out; + std::vector> by_key; + by_key.reserve(accumulator.size()); + for (const auto &[key, accum]: accumulator) + by_key.emplace_back(key, &accum); + std::sort(by_key.begin(), by_key.end(), + [](const auto &a, const auto &b) { return a.first < b.first; }); + out.reserve(accumulator.size()); - for (const auto &accum: accumulator | std::views::values) { + for (const auto &[key, accum_ptr]: by_key) { + const MergeAccum &accum = *accum_ptr; if (accum.sum_w <= 0.0) continue; @@ -388,18 +457,23 @@ std::vector MergeOnTheFly::ExportReflections() { out.emplace_back(mr); } + return out; +} +std::vector MergeOnTheFly::ExportReflections() { + auto out = ExportIntensities(); AssignRfreeFlags(out, space_group, scaling_settings.GetRfreeFraction(), 500, reference_cell); ApplyFrenchWilson(out, space_group); return out; } std::vector MergeAll(const DiffractionExperiment &x, - const std::vector &integration_outcome) { + const std::vector &integration_outcome, + size_t nthreads, + const std::optional &reference_cell) { MergeOnTheFly merge(x); - for (size_t i = 0; i < integration_outcome.size(); ++i) - merge.AddImage(integration_outcome[i], static_cast(i)); - return merge.ExportReflections(); + merge.Threads(nthreads).ReferenceCell(reference_cell).AddImages(integration_outcome); + return merge.ExportIntensities(); } struct ShellAccum { @@ -630,43 +704,51 @@ MergeStatistics MergeOnTheFly::MergeStats(const std::vector &m std::unordered_map rmeas_obs; rmeas_obs.reserve(merged.size()); - for (int i = 0; i < integration_outcome.size(); ++i) { - if (Mask(integration_outcome[i])) - continue; + struct ShellObs { uint64_t key; float I_corr; float merged_I; int shell; bool merged; }; + ReduceImagesInOrder( + integration_outcome.size(), nthreads, + [&](size_t i, std::vector &out) { + if (Mask(integration_outcome[i])) + return; + for (const auto &r: integration_outcome[i].reflections) { + if (generator.IsSystematicallyAbsent(r)) + continue; + if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) + continue; + if (!AcceptReflection(r, d_min_limit_A, d_max_limit_A)) + continue; + if (r.partiality < min_partiality) + continue; - for (const auto &r: integration_outcome[i].reflections) { - if (generator.IsSystematicallyAbsent(r)) - continue; - if (r.image_scale_corr <= 0.0 || !std::isfinite(r.image_scale_corr)) - continue; - if (!AcceptReflection(r, d_min_limit_A, d_max_limit_A)) - continue; - if (r.partiality < min_partiality) - continue; + const float I_corr = r.I * r.image_scale_corr; + const float sigma_corr = r.sigma * r.image_scale_corr; + if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) + continue; - const float I_corr = r.I * r.image_scale_corr; - const float sigma_corr = r.sigma * r.image_scale_corr; - if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) - continue; - - const auto shell = shells.GetShell(r.d); - if (!shell.has_value()) - continue; - const int s = *shell; - if (s >= 0 && s < n_shells) { - acc[s].total_obs++; - const auto key = generator(r).pack(); - const auto mit = merged_I.find(key); - if (mit != merged_I.end()) { - auto &ra = rmeas_obs[key]; - ra.sum_abs_dev += std::abs(static_cast(I_corr) - mit->second); - ra.sum_I += I_corr; - ra.n++; - ra.shell = s; + const auto shell = shells.GetShell(r.d); + if (!shell.has_value()) + continue; + const int s = *shell; + if (s >= 0 && s < n_shells) { + const auto key = generator(r).pack(); + const auto mit = merged_I.find(key); + out.push_back({key, I_corr, mit != merged_I.end() ? mit->second : NAN, s, + mit != merged_I.end()}); } } - } - } + }, + [&](size_t, const std::vector &out) { + for (const auto &o: out) { + acc[o.shell].total_obs++; + if (o.merged) { + auto &ra = rmeas_obs[o.key]; + ra.sum_abs_dev += std::abs(static_cast(o.I_corr) - o.merged_I); + ra.sum_I += o.I_corr; + ra.n++; + ra.shell = o.shell; + } + } + }); // R_meas per shell: sum over reflections of sqrt(n/(n-1)) * sum_i|I_i - |, over sum of I_i. // R_merge and R_pim beside it, with the factor 1 and sqrt(1/(n-1)) - R_pim following Weiss (2001) diff --git a/image_analysis/scale_merge/Merge.h b/image_analysis/scale_merge/Merge.h index b8f2b797e..cfb5cd615 100644 --- a/image_analysis/scale_merge/Merge.h +++ b/image_analysis/scale_merge/Merge.h @@ -283,10 +283,11 @@ class MergeOnTheFly { // merged intensities feed the space-group search and the error model - those model fits must not // see the ice-contaminated intensities. The final in-symmetry merge keeps them (for completeness). bool exclude_ice_rings = false; + size_t nthreads = 1; HKLKeyGenerator generator; - std::map accumulator; + std::unordered_map accumulator; // exported in key order // Global error model (XDS form): sigma_corr^2 = a*sigma^2 + (b*)^2. a rescales the // (under-estimated) counting variance; the (b*)^2 term adds the intensity- @@ -315,12 +316,26 @@ class MergeOnTheFly { std::unordered_map reject_median_I; size_t reject_count = 0; - bool Mask(const IntegrationOutcome &outcome); + bool Mask(const IntegrationOutcome &outcome) const; + + // One accepted observation, reduced to what the accumulator adds for it. + struct Contribution { + uint64_t key; + int32_t h, k, l; // the reflection's own index, as the accumulator records it + float w, wI, d; + bool on_ice_ring; + }; + // The observations of one image that reach the merge; returns how many the outlier test dropped. + size_t Contributions(const IntegrationOutcome &outcome, std::vector &out) const; + void Accumulate(const std::vector &contributions, int64_t image_id); public: MergeOnTheFly(const DiffractionExperiment &x); MergeOnTheFly& ReferenceCell(const std::optional &cell); MergeOnTheFly& ExcludeIceRings(bool input) { exclude_ice_rings = input; return *this; } MergeOnTheFly& FilterByImageCC(bool input) { filter_by_image_cc = input; return *this; } + // Workers for the passes over every observation (AddImages, RefineErrorModel, MergeStats); the + // result does not depend on the number. + MergeOnTheFly& Threads(size_t input) { nthreads = std::max(1, input); return *this; } // Fit the global error model from the spread of symmetry-equivalent observations. // Call once before merging; AddImage then applies it. @@ -341,6 +356,9 @@ public: // because it holds every observation at once; a streaming accumulator cannot, and there is no // second, device-side walk of the stills observations for this to disagree with. void AddImage(const IntegrationOutcome& outcome, int64_t image_id); + // AddImage for every outcome, the image id being its index: the same merge, bit for bit, with the + // per-observation work spread over the workers and the accumulation kept in image order. + void AddImages(const std::vector &outcomes); // d_min_override, when set, is the effective high-resolution limit for the shell table (used for // the automatic resolution cutoff computed by the caller); otherwise the manual @@ -355,10 +373,17 @@ public: std::optional declared_range = std::nullopt); std::vector ExportReflections(); + // The merged intensities alone - no free-R flags, no French-Wilson amplitudes - which is all a + // scaling reference reads. + std::vector ExportIntensities(); }; +// The merged intensities of the images whose cell is close to reference_cell (all of them without one): +// the reference a per-image scale is fitted against. std::vector MergeAll(const DiffractionExperiment &x, - const std::vector &reflections); + const std::vector &reflections, + size_t nthreads, + const std::optional &reference_cell); // Pearson CC between one image's corrected intensities (I * image_scale_corr) and a reference set of // full intensities, over the reflections that would enter the merge (non-ice, within the resolution diff --git a/image_analysis/scale_merge/StillsPartialityRefine.cpp b/image_analysis/scale_merge/StillsPartialityRefine.cpp index 44910f6cb..21785a263 100644 --- a/image_analysis/scale_merge/StillsPartialityRefine.cpp +++ b/image_analysis/scale_merge/StillsPartialityRefine.cpp @@ -3,13 +3,14 @@ #include "StillsPartialityRefine.h" +#include +#include #include #include #include #include #include -#include #include #include "Merge.h" @@ -101,44 +102,114 @@ namespace { return G; } - // Ceres residual: refine the orientation tilt (psi_x, psi_y) holding the scale G fixed. The tilt - // rotates the base reciprocal vector q; partiality follows analytically. Residual is the intensity - // mismatch weighted by 1/sigma, exactly matching ScaleOnTheFly's intensity-space objective. - struct PsiResidual { - double qx, qy, qz; + // The rotation by the angle-axis vector (psi_x, psi_y, 0) as a matrix (Rodrigues), so that one tilt + // is applied to all the reflections of a crystal by a matrix product each. + std::array TiltMatrix(double psi_x, double psi_y) { + const double theta = std::sqrt(psi_x * psi_x + psi_y * psi_y); + if (theta < 1e-15) + return {1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0}; + const double kx = psi_x / theta, ky = psi_y / theta; + const double c = std::cos(theta), s = std::sin(theta), t = 1.0 - c; + return {c + t * kx * kx, t * kx * ky, s * ky, + t * kx * ky, c + t * ky * ky, -s * kx, + -s * ky, s * kx, c}; + } + + // ComputeP with the tilt given as its rotation matrix. + struct TiltModel { double s0x, s0y, s0z, inv_lambda; double gamma0, gamma_e, bw; - double G, inv_corr, Iref, Iobs, weight; - template - bool operator()(const T *const psi, T *residual) const { - T q[3] = {T(qx), T(qy), T(qz)}; - T aa[3] = {psi[0], psi[1], T(0.0)}; - T qr[3]; - ceres::AngleAxisRotatePoint(aa, q, qr); - const T Sx = qr[0] + T(s0x), Sy = qr[1] + T(s0y), Sz = qr[2] + T(s0z); - const T de = ceres::sqrt(Sx * Sx + Sy * Sy + Sz * Sz) - T(inv_lambda); - const T dstar = ceres::sqrt(qr[0] * qr[0] + qr[1] * qr[1] + qr[2] * qr[2]); - const T sig_ang = T(gamma_e) * dstar; - const T sbw = T(bw) * ceres::abs(qr[2]); - const T sig2 = T(gamma0) * T(gamma0) + sig_ang * sig_ang + sbw * sbw; - const T p = ceres::exp(T(-0.5) * de * de / sig2); - residual[0] = T(weight) * (T(G) * p * T(inv_corr) * T(Iref) - T(Iobs)); - return true; + double P(const std::array &R, const FitObs &o) const { + const double qx = R[0] * o.qx + R[1] * o.qy + R[2] * o.qz; + const double qy = R[3] * o.qx + R[4] * o.qy + R[5] * o.qz; + const double qz = R[6] * o.qx + R[7] * o.qy + R[8] * o.qz; + const double de = VecLen(qx + s0x, qy + s0y, qz + s0z) - inv_lambda; + const double sig_ang = gamma_e * VecLen(qx, qy, qz); + const double sbw = bw * std::fabs(qz); + const double sig2 = gamma0 * gamma0 + sig_ang * sig_ang + sbw * sbw; + if (!(sig2 > 0.0)) + return 1.0; + return std::exp(-0.5 * de * de / sig2); } }; - // Gaussian prior N(0, sigma_prior^2) on the tilt. The data residuals above are (model-obs)/sigma, a - // proper chi^2, so the MAP prior residual is simply dpsi/sigma_prior - no scale calibration needed. - struct PsiPrior { - double inv_sigma; - template - bool operator()(const T *const psi, T *residual) const { - residual[0] = T(inv_sigma) * psi[0]; - residual[1] = T(inv_sigma) * psi[1]; - return true; + // Refine the orientation tilt (psi_x, psi_y) holding the scale G fixed: minimise + // sum_i rho(r_i^2) / 2 + |psi / sigma_prior|^2 / 2, r_i = weight_i (G p_i(psi) Iref_i / corr_i - Iobs_i), + // with the Cauchy rho(s) = k^2 log(1 + s / k^2) - the intensity-space objective of ScaleOnTheFly, and + // the Gaussian prior a proper MAP term since r_i is already in sigma units. Levenberg-Marquardt on the + // normal equations reweighted by rho', the Jacobian by central differences, every trial tilt held inside + // the bound. Returns false when the cost cannot be evaluated, leaving psi as it came in. + bool RefineTilt(const std::vector &obs, const TiltModel &model, double G, double robust_k, + double inv_sigma_prior, double max_tilt, double psi[2]) { + constexpr int MAX_ITERATIONS = 25; + constexpr double STEP = 1e-7; // rad, for the central differences + const double k2 = robust_k * robust_k; + const double prior2 = inv_sigma_prior * inv_sigma_prior; + + auto cost = [&](const double p[2]) { + const auto R = TiltMatrix(p[0], p[1]); + double c = 0.5 * prior2 * (p[0] * p[0] + p[1] * p[1]); + for (const auto &o: obs) { + const double r = o.weight * (G * model.P(R, o) * o.inv_corr * o.Iref - o.Iobs); + c += 0.5 * k2 * std::log1p(r * r / k2); + } + return c; + }; + + double current = cost(psi); + if (!std::isfinite(current)) + return false; + double lambda = 1e-4; + for (int iter = 0; iter < MAX_ITERATIONS; ++iter) { + const auto R = TiltMatrix(psi[0], psi[1]); + const auto Rxp = TiltMatrix(psi[0] + STEP, psi[1]), Rxm = TiltMatrix(psi[0] - STEP, psi[1]); + const auto Ryp = TiltMatrix(psi[0], psi[1] + STEP), Rym = TiltMatrix(psi[0], psi[1] - STEP); + double A00 = prior2, A01 = 0.0, A11 = prior2; + double b0 = prior2 * psi[0], b1 = prior2 * psi[1]; + for (const auto &o: obs) { + const double a = o.weight * G * o.inv_corr * o.Iref; + const double r = a * model.P(R, o) - o.weight * o.Iobs; + const double jx = a * (model.P(Rxp, o) - model.P(Rxm, o)) / (2.0 * STEP); + const double jy = a * (model.P(Ryp, o) - model.P(Rym, o)) / (2.0 * STEP); + const double w = 1.0 / (1.0 + r * r / k2); // rho' + A00 += w * jx * jx; + A01 += w * jx * jy; + A11 += w * jy * jy; + b0 += w * jx * r; + b1 += w * jy * r; + } + + bool accepted = false; + double trial[2] = {psi[0], psi[1]}; + double trial_cost = current; + while (lambda < 1e10) { + const double M00 = A00 * (1.0 + lambda), M11 = A11 * (1.0 + lambda); + const double det = M00 * M11 - A01 * A01; + if (!(det > 0.0)) + break; + trial[0] = std::clamp(psi[0] - (M11 * b0 - A01 * b1) / det, -max_tilt, max_tilt); + trial[1] = std::clamp(psi[1] - (M00 * b1 - A01 * b0) / det, -max_tilt, max_tilt); + trial_cost = cost(trial); + if (std::isfinite(trial_cost) && trial_cost < current) { + accepted = true; + lambda = std::max(lambda / 10.0, 1e-12); + break; + } + lambda *= 10.0; + } + if (!accepted) + break; + const double step = std::hypot(trial[0] - psi[0], trial[1] - psi[1]); + const double decrease = current - trial_cost; + psi[0] = trial[0]; + psi[1] = trial[1]; + current = trial_cost; + if (decrease <= 1e-6 * current || step <= 1e-10) + break; } - }; + return true; + } } StillsPartialityRefine::StillsPartialityRefine(const DiffractionExperiment &x) @@ -238,6 +309,8 @@ double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome, double psi[2] = {0.0, 0.0}; double G = 1.0; + const TiltModel model{.s0x = S0.x, .s0y = S0.y, .s0z = S0.z, .inv_lambda = inv_lambda, + .gamma0 = gamma0, .gamma_e = gamma_e, .bw = bw}; const bool refine_tilt = obs.size() >= settings_.min_reflections; const int inner = refine_tilt ? settings_.inner_iterations : 1; @@ -258,42 +331,11 @@ double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome, if (!refine_tilt) break; - // (2) Refine the tilt holding G fixed. - ceres::Problem problem; - for (const auto &o: obs) { - auto *cost = new ceres::AutoDiffCostFunction(new PsiResidual{ - .qx = o.qx, .qy = o.qy, .qz = o.qz, - .s0x = S0.x, .s0y = S0.y, .s0z = S0.z, .inv_lambda = inv_lambda, - .gamma0 = gamma0, .gamma_e = gamma_e, .bw = bw, - .G = G, .inv_corr = o.inv_corr, .Iref = o.Iref, .Iobs = o.Iobs, .weight = o.weight}); - problem.AddResidualBlock(cost, new ceres::CauchyLoss(settings_.robust_k), psi); - } - if (settings_.prior_sigma_deg > 0.0) { - const double inv_sigma = kRadToDeg / settings_.prior_sigma_deg; // 1 / sigma_prior (rad) - problem.AddResidualBlock(new ceres::AutoDiffCostFunction( - new PsiPrior{inv_sigma}), nullptr, psi); - } - problem.SetParameterLowerBound(psi, 0, -settings_.max_tilt_rad); - problem.SetParameterUpperBound(psi, 0, settings_.max_tilt_rad); - problem.SetParameterLowerBound(psi, 1, -settings_.max_tilt_rad); - problem.SetParameterUpperBound(psi, 1, settings_.max_tilt_rad); - - ceres::Solver::Options options; - options.linear_solver_type = ceres::DENSE_QR; - options.minimizer_progress_to_stdout = false; - options.num_threads = 1; - options.max_num_iterations = 25; - ceres::Solver::Summary summary; - const double psi_before[2] = {psi[0], psi[1]}; - ceres::Solve(options, &problem, &summary); - // A solve that failed leaves whatever the minimiser last wrote in psi, and that tilt would go - // straight onto every reflection's partiality below. Keep the tilt this crystal came in with - // and stop refining it instead - G alone is still a usable model. - if (!summary.IsSolutionUsable()) { - psi[0] = psi_before[0]; - psi[1] = psi_before[1]; + // (2) Refine the tilt holding G fixed. A cost that cannot be evaluated keeps the tilt this crystal + // came in with and stops refining it - G alone is still a usable model. + const double inv_sigma_prior = settings_.prior_sigma_deg > 0.0 ? kRadToDeg / settings_.prior_sigma_deg : 0.0; + if (!RefineTilt(obs, model, G, settings_.robust_k, inv_sigma_prior, settings_.max_tilt_rad, psi)) break; - } } // Keep what the crystal came in with. This refinement is on by default, so a crystal the tilt model @@ -349,7 +391,14 @@ double StillsPartialityRefine::RefineOne(IntegrationOutcome &outcome, return tilt_deg; } -double StillsPartialityRefine::Run(std::vector &outcomes, size_t nthreads) const { +double StillsPartialityRefine::Run(std::vector &outcomes, size_t nthreads, + const std::optional &reference_cell) const { + const auto &indexing = experiment_.GetIndexingSettings(); + auto merged_crystal = [&](const IntegrationOutcome &o) { + return !reference_cell || o.latt.GetUnitCell().is_close(*reference_cell, + indexing.GetUnitCellDistTolerance(), + indexing.GetUnitCellAngleTolerance_deg()); + }; if (nthreads == 0) nthreads = AvailableCpus(); nthreads = std::max(1, nthreads); @@ -358,7 +407,7 @@ double StillsPartialityRefine::Run(std::vector &outcomes, si for (int outer = 0; outer < settings_.outer_iterations; ++outer) { // Reference full intensities from the current corrections. - const std::vector merged = MergeAll(experiment_, outcomes); + const std::vector merged = MergeAll(experiment_, outcomes, nthreads, reference_cell); std::map reference; for (const auto &m: merged) reference[hkl_key_generator_(m)] = m.I; @@ -370,7 +419,7 @@ double StillsPartialityRefine::Run(std::vector &outcomes, si auto worker = [&]() { size_t i = next.fetch_add(1); while (i < outcomes.size()) { - const double t = RefineOne(outcomes[i], reference); + const double t = merged_crystal(outcomes[i]) ? RefineOne(outcomes[i], reference) : 0.0; if (t > 0.0) { double prev = tilt_sum.load(); while (!tilt_sum.compare_exchange_weak(prev, prev + t)) {} diff --git a/image_analysis/scale_merge/StillsPartialityRefine.h b/image_analysis/scale_merge/StillsPartialityRefine.h index fbce8c1dd..f74568b65 100644 --- a/image_analysis/scale_merge/StillsPartialityRefine.h +++ b/image_analysis/scale_merge/StillsPartialityRefine.h @@ -43,7 +43,10 @@ public: explicit StillsPartialityRefine(const DiffractionExperiment &x); // Refine all crystals in place. Returns the mean |dpsi| applied (degrees), for diagnostics. - double Run(std::vector &outcomes, size_t nthreads = 0) const; + // Only the crystals whose cell is close to reference_cell (all without one) are refined and make the + // reference: the others are not merged. + double Run(std::vector &outcomes, size_t nthreads, + const std::optional &reference_cell) const; private: const DiffractionExperiment experiment_; diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index a2b5f44a4..f390ebc3f 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -1250,7 +1250,9 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // degrade weak stills (measured CC1/2 collapse at the default 3 iters). { phase("Scaling images (" + label + ")"); - auto merge_result = MergeAll(experiment_, indexer->GetIntegrationOutcome()); + TimingMark scale_mark(logger, "stills scaling images " + label); + auto merge_result = MergeAll(experiment_, indexer->GetIntegrationOutcome(), config_.nthreads, + result.consensus_cell); indexer->ScaleAllImages(merge_result); } // Physical partiality post-refinement (default on; --simple-stills disables): refine a per-crystal @@ -1258,8 +1260,9 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // correction (no re-integration). Never for the P1 search pass. if (!for_search && experiment_.GetScalingSettings().GetStillsPartialityRefine()) { phase("Partiality post-refine (" + label + ")"); + TimingMark refine_mark(logger, "stills partiality post-refine " + label); StillsPartialityRefine refiner(experiment_); - const double mean_tilt = refiner.Run(indexer->GetIntegrationOutcome(), config_.nthreads); + const double mean_tilt = refiner.Run(indexer->GetIntegrationOutcome(), config_.nthreads, result.consensus_cell); logger.Info("Stills partiality post-refine: mean |dpsi| = {:.3f} deg", mean_tilt); // The refine fits its OWN per-image scale, with no floor, and adopts it unconditionally // when the image has no prior CC - which is exactly the state the guard above leaves. So @@ -1272,7 +1275,9 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { const std::vector &merge_input = indexer->GetIntegrationOutcome(); phase("Merging"); + TimingMark merge_mark(logger, "stills error model and merge " + label); MergeOnTheFly merge_engine(experiment_); + merge_engine.Threads(config_.nthreads); // For the de-novo P1 pass (for_search) drop the ice-ring reflections from the merged // intensities and the error model, so the space-group search sees clean data; the final // in-symmetry merge keeps them so completeness is not lost. @@ -1296,11 +1301,11 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { logger.Info("Error model (XDS convention): a={:.3f} b={:.3e} ISa={:.1f} chi2={:.2f}", em.a, em.b, em.isa, merge_engine.ErrorModelChi2()); } - for (size_t i = 0; i < merge_input.size(); ++i) - merge_engine.AddImage(merge_input[i], static_cast(i)); + merge_engine.AddImages(merge_input); ScaleMergeResult out; out.merged = merge_engine.ExportReflections(); + merge_mark.Stop(); // Automatic high-resolution cutoff (post-merge): trim the written reflections + reported // shells to the CC1/2 fall-off. The merge, scaling and error model above ran over the full @@ -1313,6 +1318,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { cut_ss.GetResolutionCCTarget(), for_search, logger, &result.resolution_fit_A); phase("Computing statistics"); + TimingMark stats_mark(logger, "stills statistics " + label); out.statistics = merge_engine.MergeStats(out.merged, merge_input, config_.reference_data, effective_d_min); // The same kept reflections binned over the reference range (--report-resolution); the // stills error model is fitted over the whole range already, so no ISa of its own. diff --git a/rugnux/rugnux_cli.cpp b/rugnux/rugnux_cli.cpp index 714bd73d1..6b29580ca 100644 --- a/rugnux/rugnux_cli.cpp +++ b/rugnux/rugnux_cli.cpp @@ -2103,16 +2103,17 @@ static int RunRugnux(int argc, char **argv) { // against a foreign dataset injects cross-dataset systematics and is a worse reference than // the data's own merge. The per-image scale G is the exact one-pass solution, so one pass // (iterating a self-rebuilt reference only re-fits the freshly-scaled noise on weak stills). - ScaleOnTheFly(experiment, MergeAll(experiment, reflections)).Scale(reflections, nthreads); + ScaleOnTheFly(experiment, MergeAll(experiment, reflections, nthreads, experiment.GetUnitCell())).Scale(reflections, nthreads); // Physical partiality post-refinement (default on; --simple-stills disables): refine a per-crystal // orientation tilt against the merge and recompute each reflection's partiality + scale correction // (no re-integration), then merge with the improved corrections. if (experiment.GetScalingSettings().GetStillsPartialityRefine()) { StillsPartialityRefine refiner(experiment); - const double mean_tilt = refiner.Run(reflections, nthreads); + const double mean_tilt = refiner.Run(reflections, nthreads, experiment.GetUnitCell()); logger.Info("Stills partiality post-refine: mean |dpsi| = {:.3f} deg", mean_tilt); } MergeOnTheFly merge_engine(experiment); + merge_engine.Threads(nthreads); merge_engine.ReferenceCell(experiment.GetUnitCell()); // --min-image-cc has to hold for the merge itself, not only for the reported statistics. merge_engine.FilterByImageCC(experiment.GetScalingSettings().GetMinCCForImage() > 0.0); @@ -2128,8 +2129,7 @@ static int RunRugnux(int argc, char **argv) { logger.Info("Error model (XDS convention): a={:.3f} b={:.3e} ISa={:.1f} chi2={:.2f}", em.a, em.b, em.isa, merge_engine.ErrorModelChi2()); } - for (size_t i = 0; i < reflections.size(); ++i) - merge_engine.AddImage(reflections[i], static_cast(i)); + merge_engine.AddImages(reflections); merged_reflections = merge_engine.ExportReflections(); // Automatic high-resolution cutoff (post-merge), matching the full-analysis path: a manual diff --git a/tests/MergeScaleTest.cpp b/tests/MergeScaleTest.cpp index aab3d5e9d..d862fd213 100644 --- a/tests/MergeScaleTest.cpp +++ b/tests/MergeScaleTest.cpp @@ -3,12 +3,15 @@ #include +#include #include #include "../image_analysis/bragg_integration/CalcISigma.h" #include "../image_analysis/scale_merge/ErrorModel.h" #include "../image_analysis/scale_merge/HKLKey.h" #include "../image_analysis/scale_merge/Merge.h" #include "../image_analysis/scale_merge/ResolutionCutoff.h" +#include "../image_analysis/scale_merge/ScaleOnTheFly.h" +#include "../image_analysis/scale_merge/StillsPartialityRefine.h" #include "gemmi/reciproc.hpp" TEST_CASE("HKLKey_NoSG_noMergeFriedel") { @@ -585,3 +588,161 @@ TEST_CASE("Global Wilson B over a widened window, with no upper cap", "[wilson]" CHECK(1.0 / (3.9 * 3.9) - 1.0 / (low_res.d_low_A * low_res.d_low_A) == Catch::Approx(1.0 / 9.0 - 1.0 / 16.0).epsilon(0.01)); } + +// The parallel AddImages is the serial AddImage loop, bit for bit: the per-observation work is spread +// over the workers, but every sum and every half-set choice is still made in image order. More images +// than one of its blocks, so the block boundary is crossed. +TEST_CASE("MergeOnTheFly_AddImagesIsTheSerialMerge") { + DiffractionExperiment x; + x.SetSpaceGroup(TestSpaceGroup()); + + std::mt19937 rng(7); + std::uniform_int_distribution index(-6, 6); + std::uniform_real_distribution unit(0.0f, 1.0f); + std::vector outcomes(5000); + for (auto &o: outcomes) { + for (int j = 0; j < 20; ++j) { + Reflection r{}; + r.h = index(rng); + r.k = index(rng); + r.l = index(rng); + if (r.h == 0 && r.k == 0 && r.l == 0) + continue; + r.d = 20.0f / std::sqrt(static_cast(r.h * r.h + r.k * r.k + r.l * r.l)); + r.I = 100.0f * unit(rng) - 5.0f; + r.sigma = 1.0f + 5.0f * unit(rng); + r.var_bkg = r.sigma * r.sigma * unit(rng); + r.prescaling_corr = 1.0f; + r.partiality = 0.5f + 0.5f * unit(rng); + r.image_scale_corr = 0.5f + unit(rng); + r.on_ice_ring = unit(rng) < 0.05f; + o.reflections.push_back(r); + } + } + + MergeOnTheFly serial(x); + serial.RefineErrorModel(outcomes); + for (size_t i = 0; i < outcomes.size(); ++i) + serial.AddImage(outcomes[i], static_cast(i)); + const auto serial_merged = serial.ExportReflections(); + + MergeOnTheFly parallel(x); + parallel.Threads(8); + parallel.RefineErrorModel(outcomes); + parallel.AddImages(outcomes); + const auto parallel_merged = parallel.ExportReflections(); + + REQUIRE(serial.ErrorModelActive()); + CHECK(serial.ErrorModelA() == parallel.ErrorModelA()); + CHECK(serial.ErrorModelB() == parallel.ErrorModelB()); + REQUIRE(serial_merged.size() == parallel_merged.size()); + for (size_t i = 0; i < serial_merged.size(); ++i) { + CHECK(serial_merged[i].h == parallel_merged[i].h); + CHECK(serial_merged[i].k == parallel_merged[i].k); + CHECK(serial_merged[i].l == parallel_merged[i].l); + CHECK(serial_merged[i].I == parallel_merged[i].I); + CHECK(serial_merged[i].sigma == parallel_merged[i].sigma); + for (int half = 0; half < 2; ++half) // NaN where a half is empty, on both sides + CHECK(std::memcmp(&serial_merged[i].I_half[half], ¶llel_merged[i].I_half[half], sizeof(float)) == 0); + } + + const auto serial_stats = serial.MergeStats(serial_merged, outcomes); + const auto parallel_stats = parallel.MergeStats(parallel_merged, outcomes); + CHECK(serial_stats.overall.total_observations == parallel_stats.overall.total_observations); + CHECK(serial_stats.overall.cc_half == parallel_stats.overall.cc_half); + CHECK(serial_stats.overall.r_meas == parallel_stats.overall.r_meas); +} + +// Stills partiality post-refinement on noiseless synthetic crystals: each crystal's true orientation is +// tilted from the stored one by up to 0.05 deg, its intensities are the true partials of a fixed set of +// full intensities, and its stored partialities those of the stored orientation. The refinement has to +// move the partialities towards the true ones. (The Ceres solve this replaced gave the same numbers: error +// 0.0559 -> 0.0302, mean tilt 0.0206 deg.) +TEST_CASE("StillsPartialityRefine_RecoversTheTilt") { + constexpr float CELL = 50.0f, WAVELENGTH = 1.0f, D_MIN = 2.0f; + constexpr double GAMMA0 = 0.001, ACCEPT = 1.7 * GAMMA0; // RLP width; |dist_ewald| accepted at prediction + DiffractionExperiment x; + x.SetSpaceGroup(gemmi::get_spacegroup_by_name("P 1")); + DiffractionGeometry geom; + geom.Wavelength_A(WAVELENGTH); + const Coord S0 = geom.GetScatteringVector(); + + std::mt19937 rng(11); + std::normal_distribution gauss(0.0, 1.0); + std::uniform_real_distribution unit(0.0, 1.0); + auto full_I = [](int h, int k, int l) { // a fixed, Friedel-symmetric pseudo-random full intensity + const uint64_t key = HKLKeyGenerator(true, gemmi::get_spacegroup_by_name("P 1"))(h, k, l).pack(); + return 10.0 + 990.0 * static_cast((key * 0x9e3779b97f4a7c15ULL) >> 40) / static_cast(1ULL << 24); + }; + auto rotate = [](const double R[9], const Coord &v) { + return Coord(static_cast(R[0] * v.x + R[1] * v.y + R[2] * v.z), + static_cast(R[3] * v.x + R[4] * v.y + R[5] * v.z), + static_cast(R[6] * v.x + R[7] * v.y + R[8] * v.z)); + }; + auto rotation = [](double ax, double ay, double az, double angle, double R[9]) { // Rodrigues + const double n = std::sqrt(ax * ax + ay * ay + az * az); + ax /= n; ay /= n; az /= n; + const double c = std::cos(angle), s = std::sin(angle), t = 1.0 - c; + const double M[9] = {c + t * ax * ax, t * ax * ay - s * az, t * ax * az + s * ay, + t * ax * ay + s * az, c + t * ay * ay, t * ay * az - s * ax, + t * ax * az - s * ay, t * ay * az + s * ax, c + t * az * az}; + std::copy(M, M + 9, R); + }; + auto partiality = [&](const Coord &q) { + const double de = (q + S0).Length() - 1.0 / WAVELENGTH; + return std::exp(-0.5 * de * de / (GAMMA0 * GAMMA0)); + }; + + std::vector outcomes(150); + std::vector> true_p(outcomes.size()); + const int hmax = static_cast(CELL / D_MIN); + for (size_t c = 0; c < outcomes.size(); ++c) { + double R[9], T[9]; + rotation(gauss(rng), gauss(rng), gauss(rng), 2.0 * M_PI * unit(rng), R); + const double tilt = (0.02 + 0.03 * unit(rng)) * M_PI / 180.0, phi = 2.0 * M_PI * unit(rng); + rotation(std::cos(phi), std::sin(phi), 0.0, tilt, T); // about an axis normal to the beam + auto &o = outcomes[c]; + o.geom = geom; + o.latt = CrystalLattice(rotate(R, Coord(CELL, 0, 0)), rotate(R, Coord(0, CELL, 0)), + rotate(R, Coord(0, 0, CELL))); + const double G = 0.5 + unit(rng); + for (int h = -hmax; h <= hmax; ++h) + for (int k = -hmax; k <= hmax; ++k) + for (int l = -hmax; l <= hmax; ++l) { + const Coord q = o.latt.Astar() * h + o.latt.Bstar() * k + o.latt.Cstar() * l; + const float dstar = q.Length(); + if (dstar == 0.0f || dstar > 1.0f / D_MIN) + continue; + if (std::fabs((q + S0).Length() - 1.0 / WAVELENGTH) > ACCEPT) + continue; + const double p_true = partiality(rotate(T, q)); + Reflection r{}; + r.h = h; r.k = k; r.l = l; + r.d = 1.0f / dstar; + r.I = static_cast(G * p_true * full_I(h, k, l)); + r.sigma = static_cast(1.0 + 0.01 * r.I); + r.prescaling_corr = 1.0f; + r.partiality = static_cast(partiality(q)); + r.image_scale_corr = 1.0f / r.partiality; + o.reflections.push_back(r); + true_p[c].push_back(p_true); + } + REQUIRE(o.reflections.size() >= 40); + } + + auto mean_abs_error = [&]() { + double sum = 0.0; + size_t n = 0; + for (size_t c = 0; c < outcomes.size(); ++c) + for (size_t i = 0; i < outcomes[c].reflections.size(); ++i, ++n) + sum += std::fabs(outcomes[c].reflections[i].partiality - true_p[c][i]); + return sum / static_cast(n); + }; + const double before = mean_abs_error(); + ScaleOnTheFly(x, MergeAll(x, outcomes, 4, std::nullopt)).Scale(outcomes, 4); + const double mean_tilt = StillsPartialityRefine(x).Run(outcomes, 4, std::nullopt); + const double after = mean_abs_error(); + CHECK(mean_tilt > 0.015); + CHECK(mean_tilt < 0.06); + CHECK(after < 0.6 * before); // measured 0.54: the prior holds the larger tilts back +}