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

ModelGrid: an atom's points within its radius are gathered, then gridded in one vectorised pass

Browse Source 
UseAtomBoxesByPlane handed its functor one grid point at a time, so the density
(gemmi's ExpSum/ExpAnisoSum with unsafe_expapprox) ran scalar - on the CPU path of
model validation the rigid body's density copies were 48% of all validation
CPU time (8sa8, perf). The walk over an atom's box on a plane is unchanged (same
order, same delta.x stepping) but now only collects the points within the radius
(index, d2, delta); the density functor then evaluates them a Gaussian at a time
over all points - each point still sums its Gaussians in gemmi's order from 0 -
in loops GCC vectorises (AVX2, no intrinsics), and adds them to the grid in walk
order. The mask functor zeroes the gathered points. Same grid, bit for bit, at
any thread count (ModelValidation_ParallelGriddingMatchesGemmi: iso + aniso,
five groups, 1 and 4 threads).

Measured (CPU build, --mode scale --model, shared machine, indicative): 8xtg
validation 28.1 -> 13.6 s wall, run total 665 -> 366 user-s (with the FitModelScale
commit); outputs identical.

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-10 10:47:27 +02:00
co-authored by Claude Opus 5.5
parent b530c2d326
commit 03cd78322a
1 changed files with 84 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/ModelGrid.cpp
+84 -17
View File
@@ -6,6 +6,7 @@
#include <algorithm>
#include <array>
#include <cmath>
#include <iterator>
#include <utility>
#include "gemmi/solmask.hpp" // SolventMasker
@@ -34,12 +35,36 @@ AtomBox MakeAtomBox(const gemmi::Grid<float> &grid, const gemmi::Fractional &fpo
return box;
}
// The points of one atom's box on one plane that lie within its radius, in the order gemmi visits them:
// their index in the grid, squared distance and delta (the atom's centre minus the point, orthogonal).
struct BoxPoints {
int n = 0;
std::vector<size_t> idx;
std::vector<double> d2, dx, dy, dz;
// Room for `capacity` points, none of them set.
void Clear(size_t capacity) {
n = 0;
if (idx.size() < capacity) {
idx.resize(capacity);
d2.resize(capacity);
dx.resize(capacity);
dy.resize(capacity);
dz.resize(capacity);
}
}
};
// gemmi's use_points_in_box<true>() run over a list of atoms (grid.hpp, do_use_points_in_box), in
// parallel over the grid's w planes: a plane is one task, and it walks every atom whose box reaches
// it, in list order, visiting only that plane's points. Every point is so handed exactly the calls
// gemmi's serial loop over the atoms hands it, in the same order, and the grid comes out bit for bit
// the same whatever the thread count. func_for(i) gives atom i's func(point, d2, delta), taken as a
// local copy so that the compiler can hold it in registers while the atom's points are written.
// the same whatever the thread count. func_for(i) gives atom i's func, taken as a local copy so that
// the compiler can hold it in registers while the atom's points are written.
//
// gemmi hands func one point at a time, as it walks the box. Here the walk, the same one, only
// collects the points within the radius of the atom, and func(data, points) is handed them all at
// once, to visit in the order given: a loop over a few hundred points the compiler can vectorise,
// where one over a row of the box is a handful of points on the coarse grids of the rigid body.
template <typename FuncFor>
void UseAtomBoxesByPlane(gemmi::Grid<float> &grid, const std::vector<AtomBox> &boxes, size_t nthreads,
FuncFor func_for) {
@@ -54,6 +79,7 @@ void UseAtomBoxesByPlane(gemmi::Grid<float> &grid, const std::vector<AtomBox> &b
}
}
ParallelFor(nw, nthreads, [&](int plane) {
BoxPoints points;
for (int i : plane_atoms[plane]) {
const AtomBox &b = boxes[i];
const auto func = func_for(i);
@@ -64,6 +90,11 @@ void UseAtomBoxesByPlane(gemmi::Grid<float> &grid, const std::vector<AtomBox> &b
const int w_lo = gemmi::iround(nctr.z) - b.dw, w_hi = gemmi::iround(nctr.z) + b.dw;
const int u_0 = gemmi::modulo(u_lo, nu), v_0 = gemmi::modulo(v_lo, nv);
gemmi::Fractional fdelta(nctr.x - u_lo, 0, 0);
// At most the (2du+1)(2dv+1) points of the box on each of its w that fall on this plane (below).
points.Clear(static_cast<size_t>(2 * b.du + 1) * (2 * b.dv + 1) * (1 + 2 * b.dw / nw));
size_t *idx = points.idx.data();
double *d2 = points.d2.data(), *dx = points.dx.data(), *dy = points.dy.data(), *dz = points.dz.data();
int n = 0;
// The box's w that fall on this plane, ascending as gemmi walks them: one, unless the box is
// wider than the cell.
for (int w = w_lo + gemmi::modulo(plane - w_lo, nw); w <= w_hi; w += nw) {
@@ -74,24 +105,27 @@ void UseAtomBoxesByPlane(gemmi::Grid<float> &grid, const std::vector<AtomBox> &b
const double dist_sq0 = gemmi::sq(delta.y) + gemmi::sq(delta.z);
if (dist_sq0 > max_dist_sq)
continue;
float *t = &grid.data[grid.index_q(u_0, v_, plane)];
const size_t row = grid.index_q(0, v_, plane);
for (int u = u_lo, u_ = u_0;;) {
const double dist_sq = dist_sq0 + gemmi::sq(delta.x);
if (!(dist_sq > max_dist_sq))
func(*t, dist_sq, delta);
if (!(dist_sq > max_dist_sq)) {
idx[n] = row + u_;
d2[n] = dist_sq;
dx[n] = delta.x;
dy[n] = delta.y;
dz[n] = delta.z;
++n;
}
if (u >= u_hi)
break;
++u;
++u_;
++t;
if (u_ == nu) {
u_ = 0;
t -= nu;
}
u_ = (u_ + 1 == nu ? 0 : u_ + 1);
delta.x -= grid.orth_n.a11;
}
}
}
points.n = n;
func(grid.data.data(), points);
}
});
}
@@ -145,11 +179,41 @@ struct AtomDensity {
IsoSum iso{};
AnisoSum aniso{};
void operator()(float &point, double r2, const gemmi::Position &delta) const {
if (!is_aniso)
point += float(occ * iso.calculate((CReal)r2));
else
point += float(occ * aniso.calculate(delta));
// gemmi's ExpSum / ExpAnisoSum::calculate() for the points of a box, a Gaussian at a time: every
// point still sums its Gaussians in gemmi's order, starting from 0, and the loops over the points
// vectorise where gemmi's loop over one point's few Gaussians does not. The points are then added
// to the grid in the order given.
void operator()(float *data, const BoxPoints &p) const {
constexpr int BLOCK = 256;
const int n = p.n;
for (int k0 = 0; k0 < n; k0 += BLOCK) {
const int m = std::min(BLOCK, n - k0);
float density[BLOCK], x[BLOCK];
for (int k = 0; k < m; ++k)
density[k] = 0;
if (!is_aniso) {
float r2[BLOCK];
for (int k = 0; k < m; ++k)
r2[k] = (CReal) p.d2[k0 + k];
for (size_t i = 0; i < std::size(iso.a); ++i) {
for (int k = 0; k < m; ++k)
x[k] = std::max(iso.b[i] * r2[k], -88.f);
for (int k = 0; k < m; ++k)
density[k] += iso.a[i] * gemmi::unsafe_expapprox(x[k]);
}
} else {
for (size_t i = 0; i < std::size(aniso.a); ++i) {
for (int k = 0; k < m; ++k) {
const gemmi::Vec3f r((float) p.dx[k0 + k], (float) p.dy[k0 + k], (float) p.dz[k0 + k]);
x[k] = std::max(aniso.b[i].r_u_r(r), -88.f);
}
for (int k = 0; k < m; ++k)
density[k] += aniso.a[i] * gemmi::unsafe_expapprox(x[k]);
}
}
for (int k = 0; k < m; ++k)
data[p.idx[k0 + k]] += float(occ * density[k]);
}
}
};
@@ -240,7 +304,10 @@ void PutMaskOnGrid(gemmi::Grid<float> &grid, const gemmi::Model &model, const st
boxes.push_back(MakeAtomBox(grid, grid.unit_cell.fractionalize(atom->pos), r));
}
UseAtomBoxesByPlane(grid, boxes, nthreads, [](int) {
return [](float &point, double, const gemmi::Position &) { point = 0.f; };
return [](float *data, const BoxPoints &p) {
for (int k = 0; k < p.n; ++k)
data[p.idx[k]] = 0.f;
};
});
SymmetrizeOrbits(grid, orbit_leaders, nthreads, [](float a, float b) { return a < b ? a : b; });
masker.remove_islands(grid);