mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 81 Wiki Activity

rugnux: model validation does the same work in less time

Browse Source 
--model validation (battery-only for users) was 18% of the battery's time. Every
number it produces is unchanged to the bit (p.mtz, maps, placed model and every
model-validation line of the report md5/diff-identical on 11 open sets); only
when and where the work runs changes:

- The bulk-solvent grid fit (FitModelScale, most of the CPU time) fits each
  solvent pair on a copy of gemmi::Scaling's target that takes
  |Fcalc + k_sol exp(-b_sol s^2) Fmask| once per pair instead of at every
  solver evaluation; same expressions, same types (new test checks a grid
  point against gemmi's own Scaling fit with ==).
- Fcalc density and the solvent mask are made on two threads; the model's
  structure factors beside the GPU engine reservation.
- The indexing probe fits the relabellings concurrently.
- The null's replicates run beside the real model's placement (they start
  from a snapshot of the model as read); one GPU engine per replicate plus
  one for the real fit instead of a cap of 4 (engines are interchangeable
  and deterministic).
- The 2mFo-DFc, mFo-DFc and anomalous maps are made and written
  concurrently; the placed model is written beside the reflection files.
- A rigid-body zone whose solvent-mask grid needs gemmi's shrink is sent to
  the CPU when the engines are reserved (ModelMaskGPU::ShrinkIsNoOp), instead
  of failing on the GPU and validating everything again on the CPU - the
  same CPU result, without the wasted first attempt.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
This commit is contained in:
leonarski_f
2026-10-09 00:01:14 +02:00
co-authored by Claude Opus 5.5
parent 42555a55fe
commit f956eb25e8
11 changed files with 374 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/ModelValidation.cpp
+123 -89
View File
@@ -537,20 +537,29 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
dc.rate = 1.5;
dc.set_grid_cell_and_spacegroup(ms.st);
dc.set_refmac_compatible_blur(ms.st.models[0]);
// The mask is made on a thread of its own, beside the density: both only read the model, each
// fills a grid of its own, and between them they are the whole cost of this step.
// Refmac radii give a slightly lower R than the Cctbx set on our test cases, at no cost.
std::future<gemmi::AsuData<std::complex<float>>> fmask = std::async(std::launch::async,
[&model = ms.st.models[0], unit_cell = dc.grid.unit_cell, spacegroup = dc.grid.spacegroup,
spacing = dc.requested_grid_spacing(), to_d] {
gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
gemmi::Grid<float> mask_grid;
mask_grid.unit_cell = unit_cell;
mask_grid.spacegroup = spacegroup;
mask_grid.set_size_from_spacing(spacing, gemmi::GridSizeRounding::Up);
masker.put_mask_on_grid(mask_grid, model);
return MapToFPhi(mask_grid).prepare_asu_data(to_d, 0);
});
dc.put_model_density_on_grid(ms.st.models[0]);
ms.fcalc = MapToFPhi(dc.grid)
.prepare_asu_data(dc.d_min, dc.blur, false, false, false);
// Refmac radii give a slightly lower R than the Cctbx set on our test cases, at no cost.
gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
gemmi::Grid<float> mask_grid;
mask_grid.unit_cell = dc.grid.unit_cell;
mask_grid.spacegroup = dc.grid.spacegroup;
mask_grid.set_size_from_spacing(dc.requested_grid_spacing(), gemmi::GridSizeRounding::Up);
masker.put_mask_on_grid(mask_grid, ms.st.models[0]);
ms.fmask = MapToFPhi(mask_grid).prepare_asu_data(dc.d_min, 0);
ms.fmask = fmask.get();
};
compute_model_factors(mdl, d_min);
// The model's structure factors are made beside the reservation of the GPU engines below, which
// does not need them.
std::future<void> model_factors = std::async(std::launch::async, [&] { compute_model_factors(mdl, d_min); });
// The rigid body's device, decided once for the whole validation - the real fit and every replicate
// of the null on the same one, so that a verdict never depends on how full the card was when a
@@ -564,13 +573,16 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
for (const MergedReflection &r : obs)
if (!std::isnan(r.F) && r.d >= finest_zone)
++zone_observations;
// One engine for the real model and one for each replicate of the null placed beside it.
rigid_body_engines = RigidBodyGPUPool::Create(st.models[0], ucell, *sg, d_min, zone_observations,
std::min<size_t>(std::max<size_t>(nthreads, 1), 4), logger);
std::min<size_t>(std::max<size_t>(nthreads, 1), NULL_REPLICATES) + 1,
logger);
rigid_body_pool = rigid_body_engines.get();
}
#else
(void) rigid_body_gpu;
#endif
model_factors.get();
gemmi::GroupOps gops = sg->operations();
gemmi::ReciprocalAsu asu(sg);
@@ -726,11 +738,18 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
};
auto probe_indexing = [&](ModelState &ms, const std::vector<MergedReflection> &obs_in) {
IndexingProbe out;
// Each relabelling is fitted on a thread of its own: the fits only read the structure factors.
std::vector<std::future<double>> relabelled;
for (const auto &op : reindex_ops)
relabelled.push_back(std::async(std::launch::async, [&, op] {
return fit_model(ms, ReindexReflections(obs_in, op)).r_free;
}));
out.identity_fit = fit_model(ms, obs_in);
std::vector<double> r_free{out.identity_fit->r_free}; // identity first, then the twin laws
double best_r_free = r_free.front();
for (const auto &op : reindex_ops) {
const double cand = fit_model(ms, ReindexReflections(obs_in, op)).r_free;
for (size_t i = 0; i < reindex_ops.size(); i++) {
const gemmi::Op &op = reindex_ops[i];
const double cand = relabelled[i].get();
r_free.push_back(cand);
if (cand < best_r_free) {
best_r_free = cand;
@@ -762,7 +781,6 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
if (obs_model == &obs)
identity_fit = std::move(indexing.identity_fit);
indexing.identity_fit.reset();
Placement real = place_and_fit(mdl, *obs_model, std::move(identity_fit), d_min);
// --- is there anything for the model to decide? ---
// There are three: the space-group label, where the model asserts the other enantiomorph; the
@@ -777,12 +795,11 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
const bool decision_pending = result.model_enantiomorph_candidate
|| !(indexing.op == gemmi::Op::identity())
|| !(result.change_of_basis_op == gemmi::Op::identity());
// A null that cannot be built is a question that could not be put, and this design already has a
// state for that. Without this the run dies here - model validation runs BEFORE the reflection
// files are written, so one failed replicate would take the .mtz, .cif, .hkl and _unmerged.mtz
// with it, after the merge has already been paid for. The replicates are rotated models fed to a
// scaling path that throws on data it cannot pair up, so this is the adversarial input for it.
try {
std::vector<gemmi::Mat33> equivalent; // the orientations the crystal cannot tell from the model's
double reach_deg = 0;
std::vector<gemmi::Mat33> null_rotation;
double null_d_min = d_min;
std::vector<MergedReflection> obs_null;
if (decision_pending) {
// --- the null: this model, this data, in random orientations ---
// An R-factor on its own says nothing about whether a model belongs to a crystal. What a model
@@ -821,12 +838,10 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
std::vector<gemmi::Op> lattice_ops{gemmi::Op::identity()};
for (const gemmi::Op &law : ReindexAmbiguityOperators(cell, ambiguity_sg))
lattice_ops.push_back(law);
std::vector<gemmi::Mat33> equivalent;
for (const gemmi::Op &law : lattice_ops)
for (const gemmi::Op &s : gops.sym_ops)
equivalent.push_back(cartesian_rotation(s.combine(law), ucell));
const double reach_deg = RigidBodyReachDeg(st.models[0], d_min);
std::vector<gemmi::Mat33> null_rotation;
reach_deg = RigidBodyReachDeg(st.models[0], d_min);
{
std::mt19937 rng(NULL_SEED);
int redraws = 0;
@@ -857,15 +872,79 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
if (r.rfree_flag && !std::isnan(r.F))
free_d.push_back(r.d);
std::sort(free_d.begin(), free_d.end(), std::greater<>());
double null_d_min = std::max(d_min, NULL_D_MIN);
null_d_min = std::max(d_min, NULL_D_MIN);
if (free_d.size() < NULL_MIN_FREE)
null_d_min = d_min;
else
null_d_min = std::min<double>(null_d_min, free_d[NULL_MIN_FREE - 1]);
std::vector<MergedReflection> obs_null;
for (const MergedReflection &r : obs)
if (r.d >= null_d_min)
obs_null.push_back(r);
}
const gemmi::Structure null_model = st; // the model as read, which the replicates are made from
// Each replicate works on its own copy of the model as read and of its structure factors, and the
// replicates share nothing else, so they run concurrently, and beside the real model's placement
// below, which moves the real model - the one the maps and the atom-density readout are still
// taken from - and nothing a replicate reads.
//
// The check before the fit is only a filter: the rigid body can walk further than its nominal
// reach, and a replicate it carries onto an orientation equivalent to the model's has been
// refined into the model's own solution (measured: 24 deg, R-free 0.56 -> 0.39, which took the
// null's spread from 0.016 to 0.066 and a fitting model below the gate). So where the
// replicate ENDS is checked as well, and one that ended within reach is drawn again and
// placed again - from a generator of its own, so the draws a replicate takes do not depend on
// the order the concurrent replicates finish in.
std::vector<double> null_r_work(NULL_REPLICATES), null_margin(NULL_REPLICATES);
std::atomic<int> placed_redraws{0};
auto run_replicate = [&](int i) {
std::mt19937 rng(NULL_SEED + 1 + i);
gemmi::Mat33 rot = null_rotation[i];
for (int redraw = 0;; redraw++) {
ModelState rep;
rep.st = null_model;
reorient_about_centroid(rep.st.models[0], rot);
compute_model_factors(rep, null_d_min);
std::optional<Fit> identity_fit;
if (!reindex_ops.empty()) {
IndexingProbe probe = probe_indexing(rep, obs_null);
null_margin[i] = probe.margin;
identity_fit = std::move(probe.identity_fit); // replicates are placed against obs
}
const Placement placed = place_and_fit(rep, obs_null, std::move(identity_fit), null_d_min);
null_r_work[i] = placed.fit.r_work;
const gemmi::Mat33 ended = placed.rb_rotation.multiply(rot);
if (redraw >= NULL_MAX_REDRAWS || angle_to_nearest_deg(ended, equivalent) >= reach_deg)
break;
++placed_redraws;
rot = random_rotation(rng);
for (int k = 0; k < NULL_MAX_REDRAWS && angle_to_nearest_deg(rot, equivalent) < reach_deg; k++)
rot = random_rotation(rng);
}
};
// On threads of their own, not on ParallelFor's pool. A parallel pass reached from a pool
// worker runs inline (ParallelFor.h), so there the scale fits and the rigid-body Jacobian of
// every replicate but the one on the calling thread ran serially, and the null lasted as long
// as its slowest serial replicate. From threads outside the pool those passes spread over the
// pool as the real fit's do. Each pass splits its work the same way wherever it runs, so the
// numbers are the same.
std::atomic<int> next_replicate{0};
std::vector<std::future<void>> replicate_threads;
if (decision_pending)
for (size_t t = 0; t < std::min<size_t>(std::max<size_t>(nthreads, 1), NULL_REPLICATES); t++)
replicate_threads.push_back(std::async(std::launch::async, [&] {
for (int i = next_replicate++; i < NULL_REPLICATES; i = next_replicate++)
run_replicate(i);
}));
// The real model is placed while they run; what is compared with them is read off it below.
Placement real = place_and_fit(mdl, *obs_model, std::move(identity_fit), d_min);
// A null that cannot be built is a question that could not be put, and this design already has a
// state for that. Without this the run dies here - model validation runs BEFORE the reflection
// files are written, so one failed replicate would take the .mtz, .cif, .hkl and _unmerged.mtz
// with it, after the merge has already been paid for. The replicates are rotated models fed to a
// scaling path that throws on data it cannot pair up, so this is the adversarial input for it.
try {
if (decision_pending) {
// The real model's side of it: the indexing margin from the model as read, as a replicate's
// is, and the R-work of the model where its placement left it, against the reflections that
// placement was fitted to.
@@ -889,58 +968,6 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
}
double real_r_work = coarse_r_work(*obs_model);
// Each replicate works on its own copy of the model and of its structure factors, and the
// replicates share nothing else, so they run concurrently and the real model - which the
// maps and the atom-density readout below are still taken from - is never moved.
//
// The check before the fit is only a filter: the rigid body can walk further than its nominal
// reach, and a replicate it carries onto an orientation equivalent to the model's has been
// refined into the model's own solution (measured: 24 deg, R-free 0.56 -> 0.39, which took the
// null's spread from 0.016 to 0.066 and a fitting model below the gate). So where the
// replicate ENDS is checked as well, and one that ended within reach is drawn again and
// placed again - from a generator of its own, so the draws a replicate takes do not depend on
// the order the concurrent replicates finish in.
std::vector<double> null_r_work(NULL_REPLICATES), null_margin(NULL_REPLICATES);
std::atomic<int> placed_redraws{0};
auto run_replicate = [&](int i) {
std::mt19937 rng(NULL_SEED + 1 + i);
gemmi::Mat33 rot = null_rotation[i];
for (int redraw = 0;; redraw++) {
ModelState rep;
rep.st = mdl.st;
SetModelPositions(rep.st.models[0], as_read);
reorient_about_centroid(rep.st.models[0], rot);
compute_model_factors(rep, null_d_min);
std::optional<Fit> identity_fit;
if (!reindex_ops.empty()) {
IndexingProbe probe = probe_indexing(rep, obs_null);
null_margin[i] = probe.margin;
identity_fit = std::move(probe.identity_fit); // replicates are placed against obs
}
const Placement placed = place_and_fit(rep, obs_null, std::move(identity_fit), null_d_min);
null_r_work[i] = placed.fit.r_work;
const gemmi::Mat33 ended = placed.rb_rotation.multiply(rot);
if (redraw >= NULL_MAX_REDRAWS || angle_to_nearest_deg(ended, equivalent) >= reach_deg)
break;
++placed_redraws;
rot = random_rotation(rng);
for (int k = 0; k < NULL_MAX_REDRAWS && angle_to_nearest_deg(rot, equivalent) < reach_deg; k++)
rot = random_rotation(rng);
}
};
// On threads of their own, not on ParallelFor's pool. A parallel pass reached from a pool
// worker runs inline (ParallelFor.h), so there the scale fits and the rigid-body Jacobian of
// every replicate but the one on the calling thread ran serially, and the null lasted as long
// as its slowest serial replicate. From threads outside the pool those passes spread over the
// pool as the real fit's do. Each pass splits its work the same way wherever it runs, so the
// numbers are the same.
std::atomic<int> next_replicate{0};
std::vector<std::future<void>> replicate_threads;
for (size_t t = 0; t < std::min<size_t>(std::max<size_t>(nthreads, 1), NULL_REPLICATES); t++)
replicate_threads.push_back(std::async(std::launch::async, [&] {
for (int i = next_replicate++; i < NULL_REPLICATES; i = next_replicate++)
run_replicate(i);
}));
for (std::future<void> &f : replicate_threads)
f.get();
if (placed_redraws > 0)
@@ -1235,24 +1262,18 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
}
// --- write the maps and score the 2mFo-DFc map at atom centres (a real map peaks there) ---
const gemmi::Grid<float> grid2fofc = map_from_coefficients(map2fofc);
const double rms2 = write_ccp4(grid2fofc, output_prefix + "_2fofc.ccp4");
write_ccp4(map_from_coefficients(mapfofc), output_prefix + "_fofc.ccp4");
{
double s = 0; int n = 0;
for (gemmi::Model &m : st.models)
for (gemmi::Chain &ch : m.chains)
for (gemmi::Residue &r : ch.residues)
for (gemmi::Atom &a : r.atoms) { s += grid2fofc.interpolate_value(a.pos, MAP_INTERPOLATION_ORDER); ++n; }
result.mean_atom_density_sigma = (n > 0 && rms2 > 0) ? (s / n) / rms2 : 0;
}
// The three maps - this one, the difference map and the anomalous map - share nothing but what
// they read, so the other two are made and written on threads of their own beside this one.
std::future<void> fofc_map = std::async(std::launch::async, [&] {
write_ccp4(map_from_coefficients(mapfofc), output_prefix + "_fofc.ccp4");
});
// --- anomalous difference map, where the merge kept the Bijvoet split ---
// Coefficients F(+) - F(-) carried on the model phase turned back by 90 degrees. Its peaks sit on
// the anomalous scatterers, so reading the map at each of the model's own atoms names them,
// rather than leaving a list of coordinates for someone to look up.
// Following ANODE, Thorn & Sheldrick (2011) J. Appl. Cryst. 44, 1285-1287
{
std::future<void> anomalous_map = std::async(std::launch::async, [&] {
// Read from the merged reflections as they came in, and carry each one into the model's frame
// here: the hand each Bijvoet difference belongs to is a property of the frame the merge was
// made in, and both operators can change it. F(+) and F(-) are attached to the + index of the
@@ -1332,7 +1353,20 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
sites.resize(MAX_ANOMALOUS_SITES);
result.anomalous_sites = std::move(sites);
}
});
const gemmi::Grid<float> grid2fofc = map_from_coefficients(map2fofc);
const double rms2 = write_ccp4(grid2fofc, output_prefix + "_2fofc.ccp4");
{
double s = 0; int n = 0;
for (gemmi::Model &m : st.models)
for (gemmi::Chain &ch : m.chains)
for (gemmi::Residue &r : ch.residues)
for (gemmi::Atom &a : r.atoms) { s += grid2fofc.interpolate_value(a.pos, MAP_INTERPOLATION_ORDER); ++n; }
result.mean_atom_density_sigma = (n > 0 && rms2 > 0) ? (s / n) / rms2 : 0;
}
fofc_map.get();
anomalous_map.get();
// --- the enantiomorph, now that there is something to decide it on ---
// Two things have to hold before the model's hand is written on these data. The model has to