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

Refinement: the repulsion search on the refinement's threads, hydrogens skipped before the neighbour walk

Browse Source 
BuildRepulsions walked the neighbours of every atom and only then dropped the pairs of a riding hydrogen;
those atoms are now skipped before the walk. The atoms are searched in 64 blocks under ParallelFor and the
blocks' pairs joined in atom order, so the restraint list is the same as before for any thread count. The
model is no longer copied for the NeighborSearch. Same bits.

9yl4: repulsions 3.37 -> 1.09 s (91 calls), refinement 17.7 -> 14.2 s; 5reo 0.38 -> 0.12 s
(RTX 5080, -N 8, loaded box).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit ded11388a18f5a4138a73150c4bb073570152727)
(cherry picked from commit 9c96836a90c9f170008d343af2cd885a4e59b321)
This commit is contained in:
leonarski_f committed 2026-10-26 01:38:23 +01:00
1 parent e3782cabe4
commit 001e61c1a4
3 files changed
+74 -62
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/RefinementRestraints.cpp
+67 -56
View File
@@ -13,6 +13,7 @@
#include "gemmi/riding_h.hpp" // place_hydrogens_on_all_atoms
#include "gemmi/topo.hpp"
#include "../../common/ParallelFor.h"
#include "MonomerLibrary.h"
#include "RefinementMath.h"
@@ -457,78 +458,88 @@ RefinementTopology BuildRefinementTopology(gemmi::Structure &st, bool hydrogens)
return out;
}
RestraintSet BuildRepulsions(const RefinementTopology &topo, const gemmi::Model &model, const gemmi::UnitCell &cell_in,
const gemmi::SpaceGroup &sg, double margin) {
RestraintSet out;
RestraintSet BuildRepulsions(const RefinementTopology &topo, gemmi::Model &model, const gemmi::UnitCell &cell_in,
const gemmi::SpaceGroup &sg, double margin, size_t nthreads) {
gemmi::UnitCell cell = cell_in;
cell.set_cell_images_from_spacegroup(&sg);
gemmi::Model copy = model;
float max_radius = 0;
for (float r : topo.vdw_radius)
max_radius = std::max(max_radius, r);
const double reach = 2 * max_radius + margin;
gemmi::NeighborSearch ns(copy, cell, std::max(reach, 5.0));
gemmi::NeighborSearch ns(model, cell, std::max(reach, 5.0));
ns.populate(false);
std::vector<std::vector<int>> first(copy.chains.size());
std::vector<std::vector<int>> first(model.chains.size());
std::vector<gemmi::Position> pos;
for (size_t c = 0; c < copy.chains.size(); c++)
for (const gemmi::Residue &r : copy.chains[c].residues) {
for (size_t c = 0; c < model.chains.size(); c++)
for (const gemmi::Residue &r : model.chains[c].residues) {
first[c].push_back(static_cast<int>(pos.size()));
for (const gemmi::Atom &a : r.atoms)
pos.push_back(a.pos);
}
for (size_t i = 0; i < pos.size(); i++) {
const gemmi::Fractional fi = cell.fractionalize(pos[i]);
ns.for_each(pos[i], '\0', reach, [&](gemmi::NeighborSearch::Mark &m, double) {
const int j = first[m.chain_idx][m.residue_idx] + m.atom_idx;
if (topo.rider_parent[i] >= 0 || topo.rider_parent[j] >= 0 || topo.element[i] == gemmi::El::H ||
topo.element[j] == gemmi::El::H)
return;
if (!topo.refinable[i] && !topo.refinable[j])
return;
if (m.image_idx == 0 && (static_cast<int>(i) >= j ||
std::binary_search(topo.excluded[i].begin(), topo.excluded[i].end(), j)))
return;
if (m.image_idx != 0 && static_cast<int>(i) > j)
return;
if (topo.altloc[i] && topo.altloc[j] && topo.altloc[i] != topo.altloc[j])
return;
// The image of j nearest to i, as a Cartesian x -> R x + t.
const gemmi::FTransform ft = ns.get_image_transformation(m.image_idx);
const gemmi::Fractional fj = ft.apply(cell.fractionalize(pos[j]));
const gemmi::Vec3 shift(std::round(fi.x - fj.x), std::round(fi.y - fj.y), std::round(fi.z - fj.z));
if (m.image_idx == 0 && (shift.x != 0 || shift.y != 0 || shift.z != 0))
return; // the same copy's own lattice translate: the identity term covers the contact
const gemmi::Mat33 rot = cell.orth.mat.multiply(ft.mat).multiply(cell.frac.mat);
const gemmi::Vec3 tran = cell.orth.mat.multiply(ft.vec + shift);
const gemmi::Position image(rot.multiply(pos[j]) + tran);
const double d = pos[i].dist(image);
// The atoms in blocks searched on the threads, each block's pairs listed in atom order; the blocks are
// joined in order below, so the list is the same for any number of threads.
constexpr int BLOCKS = 64;
std::vector<std::vector<RestraintRecord>> found(BLOCKS);
ParallelFor(BLOCKS, nthreads, [&](int block) {
const size_t lo = pos.size() * block / BLOCKS, hi = pos.size() * (block + 1) / BLOCKS;
for (size_t i = lo; i < hi; i++) {
if (topo.rider_parent[i] >= 0 || topo.element[i] == gemmi::El::H)
continue;
const gemmi::Fractional fi = cell.fractionalize(pos[i]);
ns.for_each(pos[i], '\0', reach, [&](gemmi::NeighborSearch::Mark &m, double) {
const int j = first[m.chain_idx][m.residue_idx] + m.atom_idx;
if (topo.rider_parent[j] >= 0 || topo.element[j] == gemmi::El::H)
return;
if (!topo.refinable[i] && !topo.refinable[j])
return;
if (m.image_idx == 0 && (static_cast<int>(i) >= j ||
std::binary_search(topo.excluded[i].begin(), topo.excluded[i].end(), j)))
return;
if (m.image_idx != 0 && static_cast<int>(i) > j)
return;
if (topo.altloc[i] && topo.altloc[j] && topo.altloc[i] != topo.altloc[j])
return;
// The image of j nearest to i, as a Cartesian x -> R x + t.
const gemmi::FTransform ft = ns.get_image_transformation(m.image_idx);
const gemmi::Fractional fj = ft.apply(cell.fractionalize(pos[j]));
const gemmi::Vec3 shift(std::round(fi.x - fj.x), std::round(fi.y - fj.y), std::round(fi.z - fj.z));
if (m.image_idx == 0 && (shift.x != 0 || shift.y != 0 || shift.z != 0))
return; // the same copy's own lattice translate: the identity term covers the contact
const gemmi::Mat33 rot = cell.orth.mat.multiply(ft.mat).multiply(cell.frac.mat);
const gemmi::Vec3 tran = cell.orth.mat.multiply(ft.vec + shift);
const gemmi::Position image(rot.multiply(pos[j]) + tran);
const double d = pos[i].dist(image);
float d0 = topo.vdw_radius[i] + topo.vdw_radius[j];
if (topo.metal[i] || topo.metal[j]) {
const gemmi::El other = topo.element[topo.metal[i] ? j : static_cast<int>(i)];
if (other == gemmi::El::N || other == gemmi::El::O || other == gemmi::El::S)
d0 = METAL_CONTACT;
} else if (topo.hbond[i] && topo.hbond[j])
d0 = HBOND_CONTACT;
if (!(d < d0 + margin))
return;
RestraintRecord r;
r.kind = static_cast<int>(RestraintKind::Repulsion);
r.atom[0] = static_cast<int>(i);
r.atom[1] = j;
r.ideal = d0;
r.sigma = REPULSION_SIGMA;
for (int a = 0; a < 3; a++) {
for (int b = 0; b < 3; b++)
r.rot[3 * a + b] = static_cast<float>(rot[a][b]);
r.tran[a] = static_cast<float>(tran.at(a));
}
float d0 = topo.vdw_radius[i] + topo.vdw_radius[j];
if (topo.metal[i] || topo.metal[j]) {
const gemmi::El other = topo.element[topo.metal[i] ? j : static_cast<int>(i)];
if (other == gemmi::El::N || other == gemmi::El::O || other == gemmi::El::S)
d0 = METAL_CONTACT;
} else if (topo.hbond[i] && topo.hbond[j])
d0 = HBOND_CONTACT;
if (!(d < d0 + margin))
return;
RestraintRecord r;
r.kind = static_cast<int>(RestraintKind::Repulsion);
r.atom[0] = static_cast<int>(i);
r.atom[1] = j;
r.ideal = d0;
r.sigma = REPULSION_SIGMA;
for (int a = 0; a < 3; a++) {
for (int b = 0; b < 3; b++)
r.rot[3 * a + b] = static_cast<float>(rot[a][b]);
r.tran[a] = static_cast<float>(tran.at(a));
}
found[block].push_back(r);
});
}
});
RestraintSet out;
for (const std::vector<RestraintRecord> &records : found)
for (const RestraintRecord &r : records)
out.Add(r, 2);
});
}
out.Finalize(pos.size());
return out;
}