diff --git a/image_analysis/structure_refinement/ModelGrid.cpp b/image_analysis/structure_refinement/ModelGrid.cpp index 95cd5c355..2dde37fe2 100644 --- a/image_analysis/structure_refinement/ModelGrid.cpp +++ b/image_analysis/structure_refinement/ModelGrid.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include "gemmi/solmask.hpp" // SolventMasker @@ -34,12 +35,36 @@ AtomBox MakeAtomBox(const gemmi::Grid &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 idx; + std::vector 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() 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 void UseAtomBoxesByPlane(gemmi::Grid &grid, const std::vector &boxes, size_t nthreads, FuncFor func_for) { @@ -54,6 +79,7 @@ void UseAtomBoxesByPlane(gemmi::Grid &grid, const std::vector &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 &grid, const std::vector &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(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 &grid, const std::vector &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 &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);