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

rugnux --af-model: molecular replacement faster - the score's scale fit on the GPU, the expected LLG gate, the serial parts side by side

Browse Source 
The search was CPU-bound: on a 54-chain partial complex ~80% of all CPU samples were the score's
FitModelScale grid (and libm's exp under it), with the GPU busy 13 of 37 s.

- score_placement fits its scale on the GPU where its structure factors are made there
  (ModelScaleGPU's FitSolvent / Fit, as the rigid body already does); equal to rounding, the CPU
  build unchanged.
- A model whose expected log-likelihood gain (sum sigma_A^4 / 2, McCoy et al. 2017; sigma_A^2 the
  part of the scattering one copy is by Matthews, times Luzzati's D^2) is below 10 is not searched
  for. Over the 193-set survey the lowest eLLG of a placed model is 26; every PAE domain below 24
  placed nowhere, so no decision changes.
- The likelihood setups made where first needed, one serving both where the reflections are the same.
- After a candidate fails its null, the other placed candidates' nulls side by side.
- The first copy's placements refined side by side on an engine each; the groups' translation
  functions side by side; CopyFactors and KnownStructureFactors over threads; the further copies'
  CopyFactors and gains for the peaks walked all at once; the known-structure TF setup made while the
  first copy is refined. All exact.

Survey replay (193 sets): no decision changes; two R values differ in the 3rd decimal (GPU scale).
Alternating runs, MODEL_PRIOR_SECONDS base -> new: 54-chain complex 38.3 -> 14.7, 4mtk 18.3 -> 9.7,
9v16 13.7 -> 8.3 (shared machine, indicative).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
This commit is contained in:
leonarski_f
2026-10-12 16:39:03 +02:00
co-authored by Claude Opus 5.5
parent fa79d7587b
commit 6c8ead466a
3 changed files with 316 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/MolecularReplacement.cpp
+312 -96
View File
@@ -40,6 +40,7 @@
#include "RigidBodyGPU.h"
#include "RigidBodyGPUEngine.h"
#include "ModelStructureFactorsGPU.h"
#include "ModelScaleGPU.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#endif
@@ -98,6 +99,9 @@ constexpr double NULL_D_MIN = 3.5;
// a new draw lies t(n-1) sqrt(1 + 1/n) sample sds above their mean at that probability, 4.28 x 1.054 for
// n = 9 - the gate model validation's fit is read at (ModelValidation.cpp, MODEL_FIT_SIGMA).
constexpr double NULL_SIGMA_MIN = 4.51;
// A model whose expected log-likelihood gain (ExpectedLogLikelihoodGain) is below this is not searched for. The
// lowest of a model the search did place, over 193 survey sets: 26; every PAE domain below 24 placed nowhere.
constexpr double ELLG_MIN = 10.0;
double seconds_since(std::chrono::steady_clock::time_point t0) {
return std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
@@ -236,13 +240,23 @@ PlacementScore r_factors(const gemmi::AsuData<std::complex<float>> &fmodel, cons
}
// Model validation's structure-factor engines on the GPU, one per group, copy count and resolution, made
// where the first placement of that kind is scored - or none, where there is no GPU, the cell is too small
// for its gridding or one would take over a quarter of the card. Decided on the card's total memory, so
// the same input always takes the same path.
// where the first placement of that kind is scored, each with the scale fit on the GPU beside it - or none,
// where there is no GPU, the cell is too small for its gridding or the two would take over a quarter of the
// card. Decided on the card's total memory, so the same input always takes the same path.
#ifdef JFJOCH_USE_CUDA
class ScoreEngines {
public:
ModelStructureFactorsGPU *get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) {
struct Engine {
std::unique_ptr<ModelStructureFactorsGPU> sf;
// The scale fit of the reflections the engine's F_calc are matched to, on its card: FitModelScale and
// gemmi's fit at a held solvent as ModelScaleGPU makes them - the same fits, equal to rounding.
std::mutex scale_m;
std::unique_ptr<CudaStream> stream;
std::unique_ptr<ModelScaleGPU> scale;
CudaDevicePtr<std::complex<float>> fcalc, fmask;
};
Engine *get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) {
if (get_gpu_count() == 0)
return nullptr;
size_t atoms = 0;
@@ -256,25 +270,98 @@ public:
try {
auto sf = std::make_unique<ModelStructureFactorsGPU>(RigidBodyGPUEngine::CurrentDevice(), st.models[0],
st.cell, sg, d_min);
if (sf->Supported() && sf->DeviceBytes() <= sf->DeviceTotalMemory() / 4) {
const size_t points = sf->Reflections();
const size_t bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points)
+ 2 * points * sizeof(std::complex<float>);
if (sf->Supported() && bytes <= sf->DeviceTotalMemory() / 4) {
sf->Reserve();
e.sf = std::move(sf);
const int device = RigidBodyGPUEngine::CurrentDevice();
set_gpu(sf->Device());
e.engine = std::make_unique<Engine>();
e.engine->stream = std::make_unique<CudaStream>();
e.engine->scale = std::make_unique<ModelScaleGPU>(*e.engine->stream, points);
e.engine->fcalc = CudaDevicePtr<std::complex<float>>(points);
e.engine->fmask = CudaDevicePtr<std::complex<float>>(points);
e.engine->sf = std::move(sf);
set_gpu(device);
}
} catch (const JFJochException &) {
e.engine.reset();
cuda_clear_error();
}
}
return e.sf.get();
return e.engine.get();
}
private:
struct Entry {
bool tried = false;
std::unique_ptr<ModelStructureFactorsGPU> sf;
std::unique_ptr<Engine> engine;
};
std::map<std::tuple<std::string, size_t, double>, Entry> engines_;
std::mutex m_;
};
// The scale of `scaling`'s points fitted on the engine's card: FitModelScale's grid over the solvent box, or
// with `solvent`, the overall scale and anisotropic B at its k_sol and b_sol held.
void fit_scale_gpu(ScoreEngines::Engine &e, gemmi::Scaling<float> &scaling, const PlacementScore *solvent,
double b_hi, size_t nthreads) {
const size_t n = scaling.points.size();
std::vector<std::array<int, 3>> hkl(n);
std::vector<double> stol2(n);
std::vector<float> fobs(n), sigma(n);
std::vector<std::complex<float>> fc(n), fm(n);
for (size_t i = 0; i < n; i++) {
const auto &p = scaling.points[i];
hkl[i] = {p.hkl[0], p.hkl[1], p.hkl[2]};
stol2[i] = p.stol2;
fobs[i] = p.fobs;
sigma[i] = p.sigma;
fc[i] = p.fcmol;
fm[i] = p.fmask;
}
std::vector<std::array<double, 6>> constraints;
for (const gemmi::Vec6 &c : scaling.constraint_matrix)
constraints.push_back(c);
double frac[9];
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++)
frac[3 * i + j] = scaling.cell.frac.mat[i][j];
std::unique_lock lock(e.scale_m);
const int device = RigidBodyGPUEngine::CurrentDevice();
set_gpu(e.sf->Device());
e.scale->SetPoints(hkl, stol2, fobs, sigma, constraints, frac);
if (cudaMemcpyAsync(e.fcalc.get(), fc.data(), n * sizeof(std::complex<float>), cudaMemcpyHostToDevice,
*e.stream) != cudaSuccess
|| cudaMemcpyAsync(e.fmask.get(), fm.data(), n * sizeof(std::complex<float>), cudaMemcpyHostToDevice,
*e.stream) != cudaSuccess)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Model prior: the scale fit's upload failed");
const auto *d_fc = reinterpret_cast<const float2 *>(e.fcalc.get());
const auto *d_fm = reinterpret_cast<const float2 *>(e.fmask.get());
ModelScaleParams p;
if (solvent != nullptr) {
scaling.k_sol = solvent->k_sol;
scaling.b_sol = solvent->b_sol;
p = e.scale->Fit(d_fc, d_fm, solvent->k_sol, solvent->b_sol);
} else {
try {
const ModelSolventFit fit = e.scale->FitSolvent(d_fc, d_fm, b_hi);
scaling.k_sol = fit.k_sol;
scaling.b_sol = fit.b_sol;
p = fit.scale;
} catch (const ModelScaleGPUTooFewReflections &) {
// FitModelScale walks such a grid in order, each point from where the last ended.
set_gpu(device);
lock.unlock();
FitModelScale(scaling, {.b_hi = b_hi}, nthreads);
return;
}
}
set_gpu(device);
scaling.k_overall = p.k_overall;
scaling.b_star = {p.b_star[0], p.b_star[1], p.b_star[2], p.b_star[3], p.b_star[4], p.b_star[5]};
}
#else
class ScoreEngines {};
#endif
@@ -339,8 +426,9 @@ PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGro
const gemmi::Op *law = nullptr) {
gemmi::AsuData<std::complex<float>> fmask, fmodel;
#ifdef JFJOCH_USE_CUDA
if (ModelStructureFactorsGPU *gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr)
gpu->Compute(st.models[0], fmodel, fmask);
ScoreEngines::Engine *gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr;
if (gpu)
gpu->sf->Compute(st.models[0], fmodel, fmask);
else
#endif
model_factors(st, d_min, fmodel, fmask, nthreads);
@@ -349,6 +437,11 @@ PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGro
gemmi::Scaling<float> scaling(st.cell, &sg);
scaling.use_solvent = true;
scaling.prepare_points(fmodel, fobs, &fmask);
#ifdef JFJOCH_USE_CUDA
if (gpu)
fit_scale_gpu(*gpu, scaling, solvent, SolventBMax(st.models[0]), nthreads);
else
#endif
if (solvent != nullptr) {
// The bulk solvent of an earlier fit of the same crystal, held, as the rigid body holds it.
scaling.k_sol = solvent->k_sol;
@@ -1067,19 +1160,27 @@ TranslationSetup MakeTranslationSetup(const SearchReflections &refl, const gemmi
std::vector<std::complex<double>> KnownStructureFactors(const MolecularTransform &ft, const TranslationSetup &setup,
const gemmi::UnitCell &cell,
const std::vector<gemmi::Mat33> &rotations,
const std::vector<gemmi::Fractional> &positions) {
const std::vector<gemmi::Fractional> &positions,
size_t nthreads) {
const size_t terms = setup.Terms(), nops = setup.ops;
std::vector<std::complex<double>> out(terms, 0.0);
std::vector<gemmi::Mat33> rt;
std::vector<gemmi::Position> x;
for (size_t c = 0; c < rotations.size(); c++) {
const gemmi::Mat33 rt = rotations[c].transpose();
const gemmi::Position x = cell.orthogonalize(positions[c]);
for (size_t t = 0; t < terms; t++)
for (size_t j = 0; j < nops; j++) {
const gemmi::Vec3 &sj = setup.s[t * nops + j];
// (h R_j).t = s_j.x for the Cartesian s_j and x
out[t] += ft.At(rt.multiply(sj)) * setup.phase[t * nops + j] * std::polar(1.0, 2.0 * PI * sj.dot(x));
}
rt.push_back(rotations[c].transpose());
x.push_back(cell.orthogonalize(positions[c]));
}
std::vector<std::complex<double>> out(terms, 0.0);
ParallelChunks(static_cast<int>(terms), ThreadsForWork(terms * nops * rotations.size(), nthreads),
[&](int lo, int hi) {
for (int t = lo; t < hi; t++)
for (size_t c = 0; c < rotations.size(); c++)
for (size_t j = 0; j < nops; j++) {
const gemmi::Vec3 &sj = setup.s[t * nops + j];
// (h R_j).t = s_j.x for the Cartesian s_j and x
out[t] += ft.At(rt[c].multiply(sj)) * setup.phase[t * nops + j]
* std::polar(1.0, 2.0 * PI * sj.dot(x[c]));
}
});
return out;
}
@@ -1218,8 +1319,10 @@ std::vector<std::complex<double>> CopyFactors(const SearchModel &model, const Li
std::vector<std::complex<double>> g(setup.s.size());
if (ft != nullptr && d_min >= ft->DMin()) {
const gemmi::Mat33 rt = rotation.transpose();
for (size_t e = 0; e < g.size(); e++)
g[e] = ft->At(rt.multiply(setup.s[e])) * setup.phase[e];
ParallelChunks(static_cast<int>(g.size()), ThreadsForWork(g.size(), nthreads), [&](int lo, int hi) {
for (int e = lo; e < hi; e++)
g[e] = ft->At(rt.multiply(setup.s[e])) * setup.phase[e];
});
return g;
}
gemmi::Structure st = model.st;
@@ -1832,6 +1935,28 @@ std::vector<gemmi::SpaceGroup> SameRotationGroups(const gemmi::SpaceGroup &sg) {
namespace {
// The expected log-likelihood gain of one copy of the model on the reflections of `merge` the search is made
// on: sum sigma_A^4 / 2, with sigma_A^2 the part of the scattering the copy is - one over the copy number the
// cell most probably holds, by Matthews - times Luzzati's D^2 for the coordinate error the pLDDTs predict.
// Following McCoy et al. (2017) Proc. Natl. Acad. Sci. USA 114, 3637-3641
double ExpectedLogLikelihoodGain(const SearchModel &model, const MRMerge &merge, double d_min, double d_max) {
const gemmi::UnitCell cell(merge.cell);
const SearchReflections refl = MakeSearchReflections(*merge.merged, cell, merge.space_group, d_min, d_max);
const double asu_volume = cell.volume / static_cast<double>(merge.space_group.operations().order());
const std::vector<double> pm = MatthewsProbabilities(asu_volume, model.mass_da);
size_t likely = 0;
for (size_t n = 0; n < pm.size(); n++)
if (pm[n] > pm[likely])
likely = n;
const double fraction = 1.0 / static_cast<double>(likely + 1);
double ellg = 0;
for (const gemmi::Miller &h : refl.hkl) {
const double d = std::exp(-2.0 * PI * PI * model.error_A * model.error_A * cell.calculate_1_d2(h) / 3.0);
ellg += fraction * fraction * d * d * d * d / 2.0;
}
return ellg;
}
// The search and the decision for one search model.
MolecularReplacementResult search_model(const SearchModel &model, const std::vector<MRMerge> &merges, Logger &logger,
size_t nthreads) {
@@ -1854,6 +1979,15 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
result.d_min_score = std::max(SCORE_D_MIN, data_d_min);
logger.Info("Model prior: {}: {} of {} residues kept at pLDDT >= 70 ({:.1f} kDa, radius {:.1f} A)",
model.description, model.residues_kept, model.residues_read, model.mass_da / 1000.0, model.radius_A);
// What a placement of the model can gain at best: a model too small a part of the cell, or too far from
// it, for its placement to stand out is not searched for.
const double ellg = ExpectedLogLikelihoodGain(model, merges[0], result.d_min_search, result.d_max_search);
logger.Info("Model prior: {}: expected LLG {:.1f}", model.description, ellg);
if (ellg < ELLG_MIN) {
result.ok = true;
result.decision = fmt::format("NOT_PLACED: an expected LLG of {:.1f}, below {:.0f}", ellg, ELLG_MIN);
return result;
}
t0 = std::chrono::steady_clock::now();
std::optional<MolecularTransform> ft_made;
@@ -1990,18 +2124,25 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
// --- 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;
for (const gemmi::SpaceGroup &sg : groups) {
const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg);
std::vector<TranslationPeak> tf(peaks.size());
// Each group's on a thread of its own.
std::vector<std::future<std::vector<TranslationPeak>>> tf_running;
for (const gemmi::SpaceGroup &sg : groups)
tf_running.push_back(std::async(std::launch::async, [&, sg] {
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
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]);
});
ParallelFor(static_cast<int>(peaks.size()), nthreads, [&](int i) {
tf[i] = TranslationFunction(ft, setup, peak_rotations[i]);
});
return tf;
}));
std::vector<TranslationPeak> tf_all;
for (auto &r : tf_running) {
const std::vector<TranslationPeak> tf = r.get();
tf_all.insert(tf_all.end(), tf.begin(), tf.end());
}
result.t_translation += seconds_since(t0);
@@ -2062,11 +2203,22 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
by_tf.insert(by_tf.begin(), first.begin(), first.end());
}
const TranslationPeak *tf = tfs.data();
const LikelihoodSetup lsetup = MakeLikelihoodSetup(lrefl, cell, sg);
const std::vector<std::complex<double>> no_known(lsetup.Terms(), 0.0);
// The further copies' peaks are ranked by the likelihood on the search's own reflections, read off
// the molecular transform.
const LikelihoodSetup rsetup = MakeLikelihoodSetup(refl, cell, sg);
// the molecular transform. Each made where it is first needed; where the likelihood's reflections are
// the search's own, one serves both.
std::optional<LikelihoodSetup> lsetup_made, rsetup_made;
const auto lsetup = [&]() -> const LikelihoodSetup & {
if (!lsetup_made)
lsetup_made.emplace(MakeLikelihoodSetup(lrefl, cell, sg));
return *lsetup_made;
};
const auto rsetup = [&]() -> const LikelihoodSetup & {
if (d_llg == result.d_min_search)
return lsetup();
if (!rsetup_made)
rsetup_made.emplace(MakeLikelihoodSetup(refl, cell, sg));
return *rsetup_made;
};
// --- the first copy: the best few through packing, the rigid body and R-free ---
RigidBodyEngines engines;
@@ -2088,27 +2240,53 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
// search placement by a few degrees and a fraction of an angstrom, and the clash fraction moves
// by a percent or two with it - and does not count, so that a group is CLASH only where no
// orientation the rotation function offers packs in it.
int refined = 0;
for (size_t k = 0; k < by_tf.size() && refined < (fits ? RESCORED_PER_GROUP : 1); k++) {
const size_t i = by_tf[k];
auto placed = std::make_shared<gemmi::Structure>(
PlaceSearchModel(model, rots[i], tf[i].t, cell, sg));
double clash = ClashFraction(*placed);
// The placements that go through it are known before any does, so they run through it side by side,
// on an engine each, and are then taken in order.
struct First {
size_t i = 0;
std::shared_ptr<gemmi::Structure> placed;
double clash = 1;
bool refine = false;
PlacementScore score{NAN, NAN, NAN, NAN};
gemmi::Mat33 moved;
if (fits && 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, &scorers,
&moved);
clash = ClashFraction(*placed);
};
std::vector<First> first;
int refined = 0;
for (size_t k = 0; k < by_tf.size() && refined < (fits ? RESCORED_PER_GROUP : 1); k++) {
First f;
f.i = by_tf[k];
f.placed = std::make_shared<gemmi::Structure>(PlaceSearchModel(model, rots[f.i], tf[f.i].t, cell, sg));
f.clash = ClashFraction(*f.placed);
f.refine = fits && f.clash <= 2 * CLASH_MAX;
if (f.refine || !fits)
refined++;
} else if (fits && have) {
first.push_back(std::move(f));
}
// The further copies' translation function, made while the first copy is refined.
std::future<TranslationSetup> setup2_made;
if (fits && cand.copies_max > 1)
setup2_made = std::async(std::launch::async, [&] { return MakeTranslationSetup(refl, cell, sg, true); });
std::vector<First *> to_refine;
for (First &f : first)
if (f.refine)
to_refine.push_back(&f);
if (!to_refine.empty()) {
engines = make_engines(*to_refine[0]->placed, sg, *merge.merged, d_min, to_refine.size(), logger);
std::vector<std::future<void>> running;
for (First *f : to_refine)
running.push_back(std::async(std::launch::async, [&, f] {
f->score = refine_placement(*f->placed, sg, *merge.merged, d_min, engines.get(), nthreads,
&scorers, &f->moved);
f->clash = ClashFraction(*f->placed);
}));
for (auto &r : running)
r.get();
}
for (const First &f : first) {
if (!f.refine && fits && have)
continue;
} else if (!fits) {
refined++;
}
// The refined placement's likelihood, judged against the same orientation elsewhere.
const size_t i = f.i;
const PlacementScore &score = f.score;
// 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))) {
@@ -2116,29 +2294,30 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
cand.rf_z = rf_sd > 0 ? (rfs[i] - rf_mean) / rf_sd : 0.0;
cand.tf_z = tf[i].z;
cand.copy_tf_z = {tf[i].z};
cand.clash_fraction = clash;
cand.clash_fraction = f.clash;
cand.r_work = score.r_work;
cand.r_free = score.r_free;
cand.k_sol = score.k_sol;
cand.b_sol = score.b_sol;
cand.placed = placed;
cand.placed = f.placed;
cand.copies = 1;
orientation = rots[i];
copies = {*placed};
copy_rotation = {moved.multiply(rots[i])};
copy_position = {centroid(*placed)};
copies = {*f.placed};
copy_rotation = {f.moved.multiply(rots[i])};
copy_position = {centroid(*f.placed)};
copy_g = {{}};
copy_rg = {{}};
}
glog.Info("Model prior: {:<12} RF Z {:5.1f} TF Z {:5.1f} clash {:4.1f}% R-work {:.3f} R-free {:.3f}",
sg.xhm(), rf_sd > 0 ? (rfs[i] - rf_mean) / rf_sd : 0.0, tf[i].z, 100 * clash,
sg.xhm(), rf_sd > 0 ? (rfs[i] - rf_mean) / rf_sd : 0.0, tf[i].z, 100 * f.clash,
score.r_work, score.r_free);
}
// The likelihood of the placement kept, judged against the same orientation elsewhere - only where
// its translation function's Z-score does not already stand out.
if (have && std::isfinite(cand.r_free) && cand.tf_z < TF_Z_PLACED) {
copy_g[0] = CopyFactors(model, lsetup, cell, copy_rotation[0], d_llg, nthreads, &ft);
const LikelihoodPeak lk = LikelihoodTranslationZ(lsetup, no_known, 0.0, copy_g[0], copy_position[0],
copy_g[0] = CopyFactors(model, lsetup(), cell, copy_rotation[0], d_llg, nthreads, &ft);
const std::vector<std::complex<double>> no_known(lsetup().Terms(), 0.0);
const LikelihoodPeak lk = LikelihoodTranslationZ(lsetup(), no_known, 0.0, copy_g[0], copy_position[0],
model.error_A, nthreads);
cand.llg = lk.llg;
cand.llg_z = lk.z;
@@ -2186,9 +2365,9 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
std::vector<gemmi::Mat33> rots2 = rots;
rots2.insert(rots2.end(), copy_rotation.begin(), copy_rotation.end());
if (!setup2)
setup2 = MakeTranslationSetup(refl, cell, sg, true);
setup2 = setup2_made.get();
const std::vector<std::complex<double>> known = KnownStructureFactors(ft, *setup2, cell, copy_rotation,
copy_position);
copy_position, nthreads);
std::vector<TranslationPeak> tf2(rots2.size());
#ifdef JFJOCH_USE_CUDA
if (gpu)
@@ -2218,28 +2397,42 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
// likelihood gain of the copy beside the known ones, and the best few of those through the
// rigid body. A copy that only repeats a known one - moved by one of the group's origin shifts,
// it doubles F on half the reflections and cancels it on the other half - gains nothing.
std::vector<std::complex<double>> rknown(rsetup.Terms(), 0.0);
std::vector<std::complex<double>> rknown(rsetup().Terms(), 0.0);
for (size_t c = 0; c < copies.size(); c++) {
if (copy_rg[c].empty())
copy_rg[c] = CopyFactors(model, rsetup, cell, copy_rotation[c], result.d_min_search, nthreads, &ft);
AddCopy(rsetup, copy_rg[c], copy_position[c], rknown);
copy_rg[c] = CopyFactors(model, rsetup(), cell, copy_rotation[c], result.d_min_search, nthreads, &ft);
AddCopy(rsetup(), copy_rg[c], copy_position[c], rknown);
}
double known_scale = 0;
const double rknown_llg = RefinedLogLikelihoodGain(rsetup, rknown, model.error_A, &known_scale);
const double rknown_llg = RefinedLogLikelihoodGain(rsetup(), rknown, model.error_A, &known_scale);
// Down the highest peaks until COPY_TRIALS that gain likelihood - a copy that lowers it is not
// one, and a peak on a copy already placed (or a mate of one) does - and pack. The gain that
// ranks them is taken at the known copies' sigma_A.
std::vector<Trial> trials;
// The CopyFactors the walk down the peaks may need, all made at once.
const auto factors_of = [&](size_t i) -> std::vector<std::complex<double>> & {
return i < rots.size() ? rot_g[i] : copy_rg[i - rots.size()];
};
std::vector<size_t> missing;
for (size_t k = 0; k < by_tf2.size() && k < 2 * COPY_TRIALS; k++)
if (factors_of(by_tf2[k]).empty())
missing.push_back(by_tf2[k]);
rsetup();
ParallelFor(static_cast<int>(missing.size()), nthreads, [&](int m) {
factors_of(missing[m]) = CopyFactors(model, rsetup(), cell, rots2[missing[m]],
result.d_min_search, nthreads, &ft);
});
// Their gains, all taken at once.
std::vector<double> gains(std::min(by_tf2.size(), 2 * COPY_TRIALS));
ParallelFor(static_cast<int>(gains.size()), nthreads, [&](int k) {
std::vector<std::complex<double>> f = rknown;
AddCopy(rsetup(), factors_of(by_tf2[k]), tf2[by_tf2[k]].t, f);
gains[k] = LogLikelihoodGain(rsetup(), f, model.error_A, known_scale) - rknown_llg;
});
for (size_t k = 0; k < by_tf2.size() && k < 2 * COPY_TRIALS && trials.size() < COPY_TRIALS; k++) {
Trial trial;
trial.i = by_tf2[k];
std::vector<std::complex<double>> &g = trial.i < rots.size() ? rot_g[trial.i]
: copy_rg[trial.i - rots.size()];
if (g.empty())
g = CopyFactors(model, rsetup, cell, rots2[trial.i], result.d_min_search, nthreads, &ft);
std::vector<std::complex<double>> f = rknown;
AddCopy(rsetup, g, tf2[trial.i].t, f);
trial.gain = LogLikelihoodGain(rsetup, f, model.error_A, known_scale) - rknown_llg;
trial.gain = gains[k];
if (!(trial.gain > 0))
continue;
trial.copy = PlaceSearchModel(model, rots2[trial.i], tf2[trial.i].t, cell, sg);
@@ -2272,16 +2465,16 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
trial.refined.gain = trial.refined.llg = trial.refined.z = NAN;
if (tf2[trial.i].z < TF_Z_PLACED) {
if (!lknown) {
lknown.emplace(lsetup.Terms(), 0.0);
lknown.emplace(lsetup().Terms(), 0.0);
for (size_t c = 0; c < copies.size(); c++) {
if (copy_g[c].empty())
copy_g[c] = CopyFactors(model, lsetup, cell, copy_rotation[c], d_llg, nthreads, &ft);
AddCopy(lsetup, copy_g[c], copy_position[c], *lknown);
copy_g[c] = CopyFactors(model, lsetup(), cell, copy_rotation[c], d_llg, nthreads, &ft);
AddCopy(lsetup(), copy_g[c], copy_position[c], *lknown);
}
known_llg = RefinedLogLikelihoodGain(lsetup, *lknown, model.error_A);
known_llg = RefinedLogLikelihoodGain(lsetup(), *lknown, model.error_A);
}
trial.g = CopyFactors(model, lsetup, cell, trial.rotation, d_llg, nthreads, &ft);
trial.refined = LikelihoodTranslationZ(lsetup, *lknown, known_llg, trial.g, centroid(trial.copy),
trial.g = CopyFactors(model, lsetup(), cell, trial.rotation, d_llg, nthreads, &ft);
trial.refined = LikelihoodTranslationZ(lsetup(), *lknown, known_llg, trial.g, centroid(trial.copy),
model.error_A, nthreads);
}
glog.Info("Model prior: {:<12} copy {} TF Z {:5.1f} LLG gain {:7.1f}, refined {:7.1f} Z {:5.1f} clash {:4.1f}% R-work {:.3f} R-free {:.3f}",
@@ -2401,19 +2594,9 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
// 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;
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();
// The null of candidate i.
const auto run_null = [&](int i) {
const MRCandidate &c = result.candidates[i];
RigidBodyEngines engines = make_engines(*c.placed, c.space_group, *merges[merge].merged,
result.d_min_score, NULL_REPLICATES, logger);
// One copy: its replicates placed by the translation function, as it was.
@@ -2440,9 +2623,42 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
out.push_back(p.t);
return out;
};
c.null_sigma = null_sigma(*c.placed, orientations[best], c.space_group, *merges[merge].merged,
result.d_min_score, engines.get(), nthreads, &scorers, translate);
result.t_null += seconds_since(t_null);
return null_sigma(*c.placed, orientations[i], c.space_group, *merges[merge].merged, result.d_min_score,
engines.get(), nthreads, &scorers, translate);
};
// The first candidate's null alone; where it fails, the nulls of all the other placed candidates side
// by side, each taken where the choice comes to it.
std::vector<double> nulls(result.candidates.size(), NAN);
bool first = true;
int best = -1;
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];
if (std::isnan(nulls[best])) {
const auto t_null = std::chrono::steady_clock::now();
std::vector<int> todo{best};
if (!first)
for (size_t i = 0; i < result.candidates.size(); i++)
if (static_cast<int>(i) != best && result.candidates[i].merge == merge
&& result.candidates[i].verdict == "PLACED")
todo.push_back(static_cast<int>(i));
std::vector<std::future<double>> running;
for (int i : todo)
running.push_back(std::async(std::launch::async, run_null, i));
for (size_t k = 0; k < todo.size(); k++)
nulls[todo[k]] = running[k].get();
result.t_null += seconds_since(t_null);
first = false;
}
c.null_sigma = nulls[best];
if (c.null_sigma >= NULL_SIGMA_MIN)
break;
c.verdict = "NULL";