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

rugnux --af-model: anomalous veto on the enantiomorph, search on the GPU, cheaper rescoring

Browse Source 
- An MR ENANTIOMORPH decision is taken into the answer's label only as --model takes a model's
  enantiomorph: ValidateAgainstModel runs on the placed model (in a scratch directory, removed after)
  and its adopted_model_enantiomorph - fit over its null, anomalous map at the atoms not inverted -
  decides; otherwise ENANTIOMORPH_VETOED. The hand-check numbers go to the MODEL PRIOR section.
- Rotation and translation functions on the GPU (MolecularReplacementGPU.cu): one block per
  orientation for the rotation function; per orientation the copies, the bin weights, a gather of the
  Crowther-Blow coefficients (CSR lists, fixed order, no atomics) and a cuFFT C2R, block-reduced
  statistics finished on the host in a fixed order. The CPU path (FFTW) is unchanged in what it
  computes; the orientation-independent part of the translation function (TranslationSetup) is now
  built once per group on either path.
- Rescoring: a placement clashing with its mates over twice the limit is not refined, one rigid-body
  engine per group, and the null is built only for the candidate the choice falls on.

Measured MR step on the RTX 5080 (before -> after, seconds; hand check extra where it runs):
3r6o 1.77 -> 0.67, neg 2.14 -> 0.37, 11if 1.20 -> 0.51 (+0.6 hand check), 7ris 1.46 -> 0.56 (+0.6),
9zmu 4.71 -> 1.07 (+0.9). Rotation 0.50 -> 0.02 s and translation 1.17 -> 0.13 s on 9zmu.
Mirrored-detector runs of two anomalous lysozyme sets: the hand check vetoes (map inverted at the
model's atoms); the unmirrored runs adopt.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-10-11 17:10:48 +02:00
co-authored by Claude Opus 5.5
parent 54993911f1
commit 2a3d60fb82
9 changed files with 820 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/MolecularReplacement.cpp
+180 -105
View File
@@ -27,7 +27,10 @@
#include "ModelScaling.h"
#include "RigidBodyRefine.h"
#ifdef JFJOCH_USE_CUDA
#include "MolecularReplacementGPU.h"
#include "RigidBodyGPU.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#endif
namespace {
@@ -536,17 +539,24 @@ SearchReflections MakeSearchReflections(const std::vector<MergedReflection> &mer
// --- the rotation function ---
std::vector<gemmi::Vec3> CopyVectors(const SearchReflections &refl, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg) {
const gemmi::GroupOps gops = sg.operations();
const size_t nops = gops.sym_ops.size();
const gemmi::Mat33 frac_t = cell.frac.mat.transpose();
std::vector<gemmi::Vec3> s(refl.hkl.size() * nops);
for (size_t i = 0; i < refl.hkl.size(); i++)
for (size_t j = 0; j < nops; j++)
s[i * nops + j] = frac_t.multiply(miller_vec(rotated_index(gops.sym_ops[j], refl.hkl[i])));
return s;
}
std::vector<double> RotationFunction(const MolecularTransform &ft, const SearchReflections &refl,
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
const std::vector<gemmi::Mat33> &rotations, size_t nthreads) {
const gemmi::GroupOps gops = sg.operations();
const size_t nops = gops.sym_ops.size(), n = refl.hkl.size();
const gemmi::Mat33 frac_t = cell.frac.mat.transpose();
// Copy j of the model contributes |F_model(s)|^2 at s = (h R_j) as a Cartesian reciprocal vector.
std::vector<gemmi::Vec3> s(n * nops);
for (size_t i = 0; i < n; i++)
for (size_t j = 0; j < nops; j++)
s[i * nops + j] = frac_t.multiply(miller_vec(rotated_index(gops.sym_ops[j], refl.hkl[i])));
const size_t n = refl.hkl.size();
const std::vector<gemmi::Vec3> s = CopyVectors(refl, cell, sg);
const size_t nops = s.size() / n;
double w_mean = 0;
for (double e2 : refl.e2)
w_mean += e2;
@@ -567,95 +577,112 @@ std::vector<double> RotationFunction(const MolecularTransform &ft, const SearchR
bin_sum[refl.bin[i]] += ic[i];
bin_n[refl.bin[i]]++;
}
double x_mean = 0;
double x_sum = 0, x_ss = 0, xy = 0;
for (size_t i = 0; i < n; i++) {
ic[i] /= bin_sum[refl.bin[i]] / bin_n[refl.bin[i]];
x_mean += ic[i];
const double x = ic[i] / (bin_sum[refl.bin[i]] / bin_n[refl.bin[i]]);
x_sum += x;
x_ss += x * x;
xy += x * (refl.e2[i] - w_mean);
}
x_mean /= static_cast<double>(n);
double sxy = 0, sxx = 0;
for (size_t i = 0; i < n; i++) {
sxy += (ic[i] - x_mean) * (refl.e2[i] - w_mean);
sxx += (ic[i] - x_mean) * (ic[i] - x_mean);
}
out[r] = sxx > 0 && w_ss > 0 ? sxy / std::sqrt(sxx * w_ss) : 0.0;
const double sxx = x_ss - x_sum * x_sum / static_cast<double>(n);
out[r] = sxx > 0 && w_ss > 0 ? xy / std::sqrt(sxx * w_ss) : 0.0;
});
return out;
}
// --- the translation function ---
//
// The full-symmetry translation function as one Fourier synthesis, following Crowther & Blow (1967)
// Acta Cryst. 23, 544-548: with the model at t, copy j adds G_j(h) exp(2 pi i (h R_j).t) exp(2 pi i h.t_j)
// to F(h), so |F(h)|^2 summed against E^2 - 1 over the sphere is, for each pair of copies j != k, a
// Fourier series in t with coefficient (E^2 - 1) g_j conj(g_k), g_j = G_j(h) exp(2 pi i h.t_j), at the
// index h R_j - h R_k. Every pair's terms go on one grid, and one FFT gives the function at every t.
// The pairs j = k are the rotation function's, constant in t.
TranslationPeak TranslationFunction(const MolecularTransform &ft, const SearchReflections &refl,
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
const gemmi::Mat33 &rotation) {
// The full-symmetry translation function as one Fourier synthesis, following Crowther & Blow
// (1967) Acta Cryst. 23, 544-548: with the model at t, copy j adds G_j(h) exp(2 pi i (h R_j).t)
// exp(2 pi i h.t_j) to F(h), so |F(h)|^2 summed against E^2 - 1 over the sphere is, for each pair
// of copies j != k, a Fourier series in t with coefficient (E^2 - 1) G_j conj(G_k) exp(2 pi i
// h.(t_j - t_k)) at the index h R_j - h R_k. Every pair's terms go on one grid, and one FFT gives
// the function at every t. The pairs j = k are the rotation function's, constant in t.
TranslationSetup MakeTranslationSetup(const SearchReflections &refl, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg) {
const gemmi::GroupOps gops = sg.operations();
const size_t nops = gops.sym_ops.size();
TranslationPeak out;
if (nops < 2)
return out; // P1: every position is the same
const int nops = static_cast<int>(gops.sym_ops.size());
const gemmi::Mat33 frac_t = cell.frac.mat.transpose();
const gemmi::Mat33 rt = rotation.transpose();
// Per copy j: k = h R_j, and g = G_j(h) exp(2 pi i h.t_j), so a pair's coefficient is g_j conj(g_k).
struct Term { gemmi::Miller h; std::vector<gemmi::Miller> k; std::vector<std::complex<double>> g; double w; };
std::vector<Term> terms;
std::vector<double> bin_sum(refl.bins, 0.0);
std::vector<int> bin_n(refl.bins, 0);
std::vector<int> term_bin;
for (size_t i = 0; i < refl.hkl.size(); i++) {
if (gops.is_systematically_absent(refl.hkl[i]))
continue;
Term t{refl.hkl[i], {}, {}, refl.multiplicity[i] * (refl.e2[i] - 1.0)};
double sum = 0;
for (const gemmi::Op &op : gops.sym_ops) {
t.k.push_back(rotated_index(op, refl.hkl[i]));
const double ht = (t.h[0] * op.tran[0] + t.h[1] * op.tran[1] + t.h[2] * op.tran[2])
/ static_cast<double>(gemmi::Op::DEN);
t.g.push_back(ft.At(rt.multiply(frac_t.multiply(miller_vec(t.k.back())))) * std::polar(1.0, 2.0 * PI * ht));
sum += std::norm(t.g.back());
}
bin_sum[refl.bin[i]] += sum;
bin_n[refl.bin[i]]++;
term_bin.push_back(refl.bin[i]);
terms.push_back(std::move(t));
}
if (terms.empty())
return out;
int max_m[3] = {0, 0, 0};
for (const Term &t : terms)
for (size_t j = 0; j < nops; j++)
for (size_t k = 0; k < nops; k++)
for (int a = 0; a < 3; a++)
max_m[a] = std::max(max_m[a], std::abs(t.k[j][a] - t.k[k][a]));
const int nu = good_even_size(2 * max_m[0] + 2), nv = good_even_size(2 * max_m[1] + 2),
nw = good_even_size(2 * max_m[2] + 2);
const int nwh = nw / 2 + 1;
std::vector<std::complex<double>> c(static_cast<size_t>(nu) * nv * nwh, 0.0);
for (size_t i = 0; i < terms.size(); i++) {
const Term &t = terms[i];
const double norm = bin_sum[term_bin[i]] / bin_n[term_bin[i]];
for (size_t j = 0; j < nops; j++)
for (size_t k = 0; k < nops; k++) {
if (j == k)
continue;
const int m[3] = {t.k[j][0] - t.k[k][0], t.k[j][1] - t.k[k][1], t.k[j][2] - t.k[k][2]};
// Only l >= 0 is stored: the pair (k, j) carries the conjugate term at -m.
if (m[2] < 0)
continue;
const std::complex<double> v = t.w / norm * t.g[j] * std::conj(t.g[k]);
const int u = (m[0] % nu + nu) % nu, vv = (m[1] % nv + nv) % nv;
c[u + static_cast<size_t>(nu) * (vv + static_cast<size_t>(nv) * m[2])] += v;
TranslationSetup out;
out.ops = nops;
out.bins = refl.bins;
for (int j = 0; j < nops; j++)
for (int k = 0; k < nops; k++)
if (j != k) {
out.pair_j.push_back(j);
out.pair_k.push_back(k);
}
std::vector<gemmi::Miller> rotated; // terms x ops
for (size_t i = 0; i < refl.hkl.size(); i++) {
const gemmi::Miller &h = refl.hkl[i];
if (gops.is_systematically_absent(h))
continue;
for (const gemmi::Op &op : gops.sym_ops) {
rotated.push_back(rotated_index(op, h));
out.s.push_back(frac_t.multiply(miller_vec(rotated.back())));
const double ht = (h[0] * op.tran[0] + h[1] * op.tran[1] + h[2] * op.tran[2])
/ static_cast<double>(gemmi::Op::DEN);
out.phase.push_back(std::polar(1.0, 2.0 * PI * ht));
}
out.weight.push_back(refl.multiplicity[i] * (refl.e2[i] - 1.0));
out.bin.push_back(refl.bin[i]);
}
int max_m[3] = {0, 0, 0};
for (size_t t = 0; t < out.Terms(); t++)
for (size_t p = 0; p < out.pair_j.size(); p++)
for (int a = 0; a < 3; a++)
max_m[a] = std::max(max_m[a], std::abs(rotated[t * nops + out.pair_j[p]][a]
- rotated[t * nops + out.pair_k[p]][a]));
out.nu = good_even_size(2 * max_m[0] + 2);
out.nv = good_even_size(2 * max_m[1] + 2);
out.nw = good_even_size(2 * max_m[2] + 2);
for (size_t t = 0; t < out.Terms(); t++)
for (size_t p = 0; p < out.pair_j.size(); p++) {
const gemmi::Miller &kj = rotated[t * nops + out.pair_j[p]], &kk = rotated[t * nops + out.pair_k[p]];
const int l = kj[2] - kk[2];
out.target_u.push_back(l < 0 ? -1 : ((kj[0] - kk[0]) % out.nu + out.nu) % out.nu);
out.target_v.push_back(((kj[1] - kk[1]) % out.nv + out.nv) % out.nv);
out.target_l.push_back(l);
}
return out;
}
TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup,
const gemmi::Mat33 &rotation) {
TranslationPeak out;
const size_t terms = setup.Terms(), nops = setup.ops, npairs = setup.pair_j.size();
if (nops < 2 || terms == 0)
return out; // P1: every position is the same
const gemmi::Mat33 rt = rotation.transpose();
std::vector<std::complex<double>> g(terms * nops);
std::vector<double> bin_sum(setup.bins, 0.0);
std::vector<int> bin_n(setup.bins, 0);
for (size_t t = 0; t < terms; t++) {
double sum = 0;
for (size_t j = 0; j < nops; j++) {
g[t * nops + j] = ft.At(rt.multiply(setup.s[t * nops + j])) * setup.phase[t * nops + j];
sum += std::norm(g[t * nops + j]);
}
bin_sum[setup.bin[t]] += sum;
bin_n[setup.bin[t]]++;
}
const int nu = setup.nu, nv = setup.nv, nwh = setup.nw / 2 + 1;
std::vector<std::complex<double>> c(static_cast<size_t>(nu) * nv * nwh, 0.0);
for (size_t t = 0; t < terms; t++) {
const double w = setup.weight[t] / (bin_sum[setup.bin[t]] / bin_n[setup.bin[t]]);
for (size_t p = 0; p < npairs; p++) {
const size_t e = t * npairs + p;
if (setup.target_u[e] < 0)
continue;
c[setup.target_u[e] + static_cast<size_t>(nu) * (setup.target_v[e] + static_cast<size_t>(nv) * setup.target_l[e])]
+= w * g[t * nops + setup.pair_j[p]] * std::conj(g[t * nops + setup.pair_k[p]]);
}
}
// MapFromFPhi sums F exp(-2 pi i m.x): with F = conj(C) that is the conjugate of the series, which
// is real.
gemmi::FPhiGrid<float> hkl;
hkl.unit_cell = cell;
hkl.unit_cell = ft.Grid().unit_cell;
hkl.spacegroup = &gemmi::get_spacegroup_p1();
hkl.axis_order = gemmi::AxisOrder::XYZ;
hkl.half_l = true;
@@ -795,9 +822,24 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
t0 = std::chrono::steady_clock::now();
const MolecularTransform ft(model, result.d_min_search, nthreads);
#ifdef JFJOCH_USE_CUDA
// The rotation and translation functions on the GPU where there is one; the rest of the search is
// the same either way.
std::unique_ptr<MolecularReplacementGPU> gpu;
if (get_gpu_count() > 0) {
try {
gpu = std::make_unique<MolecularReplacementGPU>(ft);
} catch (const JFJochException &e) {
result.failure_reason = std::string("GPU: ") + e.what();
return result;
}
}
result.gpu_search = gpu != nullptr;
#endif
result.t_transform = seconds_since(t0);
const std::vector<gemmi::Mat33> sphere = UniformRotations(ROTATION_STEP_DEG);
std::vector<gemmi::Mat33> orientations; // of each candidate's placement, before the rigid body
// One merge: the rotation function over its point group, then for each of its groups the
// translation function and the best placements through packing, the rigid body and R-free.
const auto search_merge = [&](size_t m) {
@@ -810,7 +852,12 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
result.d_min_search, result.d_max_search);
const std::vector<gemmi::Mat33> pg = CartesianPointGroup(merge.space_group, cell);
const std::vector<gemmi::Mat33> rotations = RotationAsu(sphere, pg);
#ifdef JFJOCH_USE_CUDA
const std::vector<double> rf = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, rotations)
: RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads);
#else
const std::vector<double> rf = RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads);
#endif
double rf_mean = 0, rf_sd = 0;
for (double x : rf)
rf_mean += x;
@@ -840,10 +887,23 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
// --- translation, for every group and every peak orientation ---
t0 = std::chrono::steady_clock::now();
const std::vector<gemmi::SpaceGroup> groups = SameRotationGroups(merge.space_group);
std::vector<TranslationPeak> tf_all(groups.size() * peaks.size());
ParallelFor(static_cast<int>(tf_all.size()), nthreads, [&](int i) {
tf_all[i] = TranslationFunction(ft, refl, cell, groups[i / peaks.size()], rotations[peaks[i % peaks.size()]]);
});
std::vector<gemmi::Mat33> peak_rotations;
for (size_t p : peaks)
peak_rotations.push_back(rotations[p]);
std::vector<TranslationPeak> tf_all;
for (const gemmi::SpaceGroup &sg : groups) {
const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg);
std::vector<TranslationPeak> tf(peaks.size());
#ifdef JFJOCH_USE_CUDA
if (gpu)
tf = gpu->TranslationFunction(setup, peak_rotations);
else
#endif
ParallelFor(static_cast<int>(peaks.size()), nthreads, [&](int i) {
tf[i] = TranslationFunction(ft, setup, peak_rotations[i]);
});
tf_all.insert(tf_all.end(), tf.begin(), tf.end());
}
result.t_translation += seconds_since(t0);
for (size_t g = 0; g < groups.size(); g++) {
@@ -872,14 +932,20 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
const size_t i = by_tf[k];
auto placed = std::make_shared<gemmi::Structure>(
PlaceSearchModel(model, rotations[peaks[i]], tf[i].t, cell, sg));
// A placement clashing with its own mates over twice the limit is not refined: the rigid
// body moves a search placement by a few degrees and a fraction of an angstrom, and the
// clash fraction moves by a percent or two with it.
double clash = ClashFraction(*placed);
PlacementScore score{NAN, NAN};
if (one_copy) {
if (k == 0)
engines = make_engines(*placed, sg, *merge.merged, d_min, NULL_REPLICATES + 1, logger);
if (one_copy && clash <= 2 * CLASH_MAX) {
if (engines.get() == nullptr)
engines = make_engines(*placed, sg, *merge.merged, d_min, 1, logger);
score = refine_placement(*placed, sg, *merge.merged, d_min, engines.get(), nthreads);
clash = ClashFraction(*placed);
}
const double clash = ClashFraction(*placed);
if (!have || score.r_free < cand.r_free) { // NaN (not refined) only ever on k = 0
// The refined placement with the lowest R-free; the best translation-function one where
// none was refined.
if (!have || (std::isfinite(score.r_free) && !(cand.r_free <= score.r_free))) {
have = true;
cand.rf_z = rf_sd > 0 ? (rf[peaks[i]] - rf_mean) / rf_sd : 0.0;
cand.tf_z = tf[i].z;
@@ -905,14 +971,7 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
cand.verdict = "NOT_PLACED";
else
cand.verdict = "PLACED";
if (cand.verdict == "PLACED") {
t0 = std::chrono::steady_clock::now();
cand.null_sigma = null_sigma(*cand.placed, orientation, sg, *merge.merged, d_min, engines.get(),
nthreads);
result.t_null += seconds_since(t0);
if (!(cand.null_sigma >= NULL_SIGMA_MIN))
cand.verdict = "NULL";
}
orientations.push_back(orientation);
result.candidates.push_back(std::move(cand));
}
};
@@ -921,15 +980,31 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
// Within one merge the candidates are scored on the same reflections, so their R-free values can
// be compared directly: the placed candidate with the lowest R-free is chosen where it leads every
// other candidate of that merge that packs by R_FREE_MARGIN.
// The null is built for the candidate the choice would fall on, and only for it: a placed candidate
// that does not beat its own null is set aside, and the choice falls on the next.
const auto choose_in = [&](int merge) -> int {
int best = -1;
for (size_t i = 0; i < result.candidates.size(); i++) {
const MRCandidate &c = result.candidates[i];
if (c.merge == merge && c.verdict == "PLACED" && (best < 0 || c.r_free < result.candidates[best].r_free))
best = static_cast<int>(i);
while (true) {
best = -1;
for (size_t i = 0; i < result.candidates.size(); i++) {
const MRCandidate &c = result.candidates[i];
if (c.merge == merge && c.verdict == "PLACED"
&& (best < 0 || c.r_free < result.candidates[best].r_free))
best = static_cast<int>(i);
}
if (best < 0)
return -1;
MRCandidate &c = result.candidates[best];
const auto t_null = std::chrono::steady_clock::now();
RigidBodyEngines engines = make_engines(*c.placed, c.space_group, *merges[merge].merged,
result.d_min_score, NULL_REPLICATES, logger);
c.null_sigma = null_sigma(*c.placed, orientations[best], c.space_group, *merges[merge].merged,
result.d_min_score, engines.get(), nthreads);
result.t_null += seconds_since(t_null);
if (c.null_sigma >= NULL_SIGMA_MIN)
break;
c.verdict = "NULL";
}
if (best < 0)
return -1;
for (size_t i = 0; i < result.candidates.size(); i++) {
const MRCandidate &c = result.candidates[i];
if (static_cast<int>(i) == best || c.merge != merge || c.clash_fraction > CLASH_MAX)