diff --git a/image_analysis/structure_refinement/ModelGrid.cpp b/image_analysis/structure_refinement/ModelGrid.cpp index c10ea95fd..82345e07d 100644 --- a/image_analysis/structure_refinement/ModelGrid.cpp +++ b/image_analysis/structure_refinement/ModelGrid.cpp @@ -130,6 +130,126 @@ void UseAtomBoxesByPlane(gemmi::Grid &grid, const std::vector &b }); } +// The density's walk: UseAtomBoxesByPlane()'s box and order, with each point's place relative to the atom and +// the test against its radius taken in single precision, in a fixed order of operations, so that the Apple +// GPU's gather (ModelDensityMetal.metal) - which has no double precision - finds the same points with the same +// bits, spelling the same arithmetic out. gemmi takes them in double; the density differs from its own in the +// last bits. + +// An atom's box: about the grid point nearest its centre, c - d to c + d along each axis. +struct DensityBox { + float nctr[3]; // the centre in grid units, (fractional position) x (grid size) + int c[3], d[3]; + float max_r2; // the radius squared +}; + +DensityBox MakeDensityBox(const gemmi::Grid &grid, const gemmi::Fractional &fpos, double radius) { + const AtomBox b = MakeAtomBox(grid, fpos, radius); + const double n[3] = {fpos.x * grid.nu, fpos.y * grid.nv, fpos.z * grid.nw}; + DensityBox out; + for (int k = 0; k < 3; k++) { + out.nctr[k] = static_cast(n[k]); + out.c[k] = gemmi::iround(n[k]); + } + out.d[0] = b.du; + out.d[1] = b.dv; + out.d[2] = b.dw; + out.max_r2 = static_cast(radius * radius); + return out; +} + +// The orthogonalization over the grid's size (gemmi's orth_n, upper triangular) in single precision: a point's +// place along v and w does not depend on u. +struct DensityGeometry { + float m00, m01, m02, m11, m12, m22; +}; + +DensityGeometry ModelDensityGeometry(const gemmi::Grid &grid) { + const gemmi::UpperTriangularMat33 &m = grid.orth_n; + return {static_cast(m.a11), static_cast(m.a12), static_cast(m.a13), + static_cast(m.a22), static_cast(m.a23), static_cast(m.a33)}; +} + +// The points of one atom's box on one plane 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 DensityPoints { + int n = 0; + std::vector idx; + std::vector r2, dx, dy, dz; + void Clear(size_t capacity) { + n = 0; + if (idx.size() < capacity) { + idx.resize(capacity); + r2.resize(capacity); + dx.resize(capacity); + dy.resize(capacity); + dz.resize(capacity); + } + } +}; + +// UseAtomBoxesByPlane() for the density's boxes: the same planes, atoms and walk. A point at (u, v, w), the +// box's own (unwrapped) indices, is f = centre - (u, v, w) from the atom, and in A +// x = m00 fx + (m01 fy + m02 fz), y = m11 fy + m12 fz, z = m22 fz, r2 = (y y + z z) + x x, +// and is taken where r2 is not above the radius squared. +template +void UseDensityBoxesByPlane(gemmi::Grid &grid, const std::vector &boxes, size_t nthreads, + FuncFor func_for) { + const int nu = grid.nu, nv = grid.nv, nw = grid.nw; + const DensityGeometry m = ModelDensityGeometry(grid); + std::vector> plane_atoms(nw); + for (int i = 0; i < static_cast(boxes.size()); i++) + for (int w = boxes[i].c[2] - boxes[i].d[2]; w <= boxes[i].c[2] + boxes[i].d[2]; w++) { + std::vector &atoms = plane_atoms[gemmi::modulo(w, nw)]; + if (atoms.empty() || atoms.back() != i) + atoms.push_back(i); + } + ParallelFor(nw, nthreads, [&](int plane) { + DensityPoints points; + for (int i : plane_atoms[plane]) { + const DensityBox &b = boxes[i]; + const auto func = func_for(i); + const int u_lo = b.c[0] - b.d[0], u_hi = b.c[0] + b.d[0]; + const int v_lo = b.c[1] - b.d[1], v_hi = b.c[1] + b.d[1]; + const int w_lo = b.c[2] - b.d[2], w_hi = b.c[2] + b.d[2]; + const int u_0 = gemmi::modulo(u_lo, nu), v_0 = gemmi::modulo(v_lo, nv); + points.Clear(static_cast(2 * b.d[0] + 1) * (2 * b.d[1] + 1) * (1 + 2 * b.d[2] / nw)); + int n = 0; + for (int w = w_lo + gemmi::modulo(plane - w_lo, nw); w <= w_hi; w += nw) { + const float fz = b.nctr[2] - static_cast(w); + const float z = m.m22 * fz; + for (int v = v_lo, v_ = v_0; v <= v_hi; ++v, v_ = (v_ + 1 == nv ? 0 : v_ + 1)) { + const float fy = b.nctr[1] - static_cast(v); + const float y = m.m11 * fy + m.m12 * fz; + const float yz = y * y + z * z; + if (yz > b.max_r2) + continue; // so is every point of the row: x x only adds + const float x_yz = m.m01 * fy + m.m02 * fz; + const size_t row = grid.index_q(0, v_, plane); + for (int u = u_lo, u_ = u_0;;) { + const float x = m.m00 * (b.nctr[0] - static_cast(u)) + x_yz; + const float r2 = yz + x * x; + if (!(r2 > b.max_r2)) { + points.idx[n] = row + u_; + points.r2[n] = r2; + points.dx[n] = x; + points.dy[n] = y; + points.dz[n] = z; + ++n; + } + if (u >= u_hi) + break; + ++u; + u_ = (u_ + 1 == nu ? 0 : u_ + 1); + } + } + } + points.n = n; + func(grid.data.data(), points); + } + }); +} + // gemmi's Grid::symmetrize() in parallel over the orbits found by OrbitLeaders(): each orbit is reduced // by one thread, into its leader, with gemmi's operands in gemmi's order, and the value written to its // leader and mates as gemmi writes it. Orbits share no points, so the grid is gemmi's, bit for bit. @@ -183,7 +303,7 @@ struct AtomDensity { // 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 { + void operator()(float *data, const DensityPoints &p) const { constexpr int BLOCK = 256; const int n = p.n; for (int k0 = 0; k0 < n; k0 += BLOCK) { @@ -194,7 +314,7 @@ struct AtomDensity { if (!is_aniso) { float r2[BLOCK]; for (int k = 0; k < m; ++k) - r2[k] = (CReal) p.d2[k0 + k]; + r2[k] = p.r2[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); @@ -204,7 +324,7 @@ struct AtomDensity { } 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]); + const gemmi::Vec3f r(p.dx[k0 + k], p.dy[k0 + k], p.dz[k0 + k]); x[k] = std::max(aniso.b[i].r_u_r(r), -88.f); } for (int k = 0; k < m; ++k) @@ -248,15 +368,16 @@ std::vector OrbitLeaders(const gemmi::Grid &grid, size_t nthreads } // DensityCalculator::put_model_density_on_grid() (gemmi dencalc.hpp, do_add_atom_density_to_grid) -// with the atoms and the symmetry spread over the threads as above - the same grid, bit for bit. Each -// atom's density coefficients and radius are worked out once, up front, rather than once per plane. +// with the atoms and the symmetry spread over the threads as above - the same grid on any number of threads, +// and gemmi's but for the last bits of the single-precision walk (UseDensityBoxesByPlane). Each atom's density +// coefficients and radius are worked out once, up front, rather than once per plane. void PutModelDensityOnGrid(gemmi::DensityCalculator &dc, const gemmi::Model &model, const std::vector &orbit_leaders, size_t nthreads) { using CReal = AtomDensity::CReal; dc.initialize_grid(); const std::vector atoms = ModelAtoms(model); const int n = static_cast(atoms.size()); - std::vector boxes(n); + std::vector boxes(n); std::vector density(n); ParallelChunks(n, nthreads, [&](int lo, int hi) { for (int i = lo; i < hi; i++) { @@ -270,17 +391,17 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator &dc, const gem if (!d.is_aniso) { const CReal b = static_cast(atom.b_iso + dc.blur); d.iso = coef.precalculate_density_iso(b, addend); - boxes[i] = MakeAtomBox(dc.grid, fpos, dc.estimate_radius(d.iso, b)); + boxes[i] = MakeDensityBox(dc.grid, fpos, dc.estimate_radius(d.iso, b)); } else { const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur)); const CReal b_max = std::max(std::max(aniso_b.u11, aniso_b.u22), aniso_b.u33); const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max); d.aniso = coef.precalculate_density_aniso_b(aniso_b, addend); - boxes[i] = MakeAtomBox(dc.grid, fpos, radius); + boxes[i] = MakeDensityBox(dc.grid, fpos, radius); } } }); - UseAtomBoxesByPlane(dc.grid, boxes, nthreads, [&](int i) { return density[i]; }); + UseDensityBoxesByPlane(dc.grid, boxes, nthreads, [&](int i) { return density[i]; }); SymmetrizeOrbits(dc.grid, orbit_leaders, nthreads, [](float a, float b) { return a + b; }); } diff --git a/image_analysis/structure_refinement/ModelGrid.h b/image_analysis/structure_refinement/ModelGrid.h index b6075a999..1ee4e9249 100644 --- a/image_analysis/structure_refinement/ModelGrid.h +++ b/image_analysis/structure_refinement/ModelGrid.h @@ -14,7 +14,9 @@ // gemmi's gridding of a model - the density of DensityCalculator::put_model_density_on_grid() and the // bulk-solvent mask of SolventMasker(AtomicRadiiSet::Refmac).put_mask_on_grid() - spread over // `nthreads` threads of ParallelFor's pool. Every grid point receives the same operations in the same -// order as in gemmi's serial loops, so the grids are gemmi's, bit for bit, on any number of threads. +// order as in gemmi's serial loops, so the grids are the same on any number of threads: the mask is gemmi's +// bit for bit, the density gemmi's but for the last bits (its points are placed in single precision, as the +// Apple GPU places them). // The points gemmi's Grid::symmetrize() reduces the orbits of this grid's size and space group into // (empty in P1). Both functions below take them, for a grid of the same size.