From fe58c5041ce0e7ffb966cfbc8ce87358b4a9567d Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 18 Oct 2026 17:04:17 +0200 Subject: [PATCH] RESULT-CHANGING (ulp level): the model density's walk places its points in single precision PutModelDensityOnGrid walked each atom's box as gemmi does, with every point's place relative to the atom and its test against the atom's radius in double. The Apple GPU has no double, so a Metal gather could only come to the same grid to rounding. The density now has its own walk, UseDensityBoxesByPlane: the same planes, atoms, box and order as before, with the point's place in single precision in one fixed order of operations (f = centre - (u, v, w); x = m00 fx + (m01 fy + m02 fz), y = m11 fy + m12 fz, z = m22 fz with gemmi's upper-triangular orth_n; r2 = (y y + z z) + x x, taken where not above the radius squared). A gather that spells the same arithmetic out then finds the same points with the same bits, and gemmi's float expapprox and the Gaussian sums are already single precision. The solvent mask keeps the double walk: its points are 0 or 1, and a point at the edge would flip. The density differs from gemmi's in the last bits, so everything computed from it moves at that level. 9s02 on an Apple M6 against the previous commit: p.hkl unchanged; the MR decision (AGREES, P 21 21 2) and every R-work/R-free within 0.001; validation R(all) 0.3719 -> 0.3718; the validation rigid body took 29 evaluations instead of 22 to the same place (rotation 0.003 deg); Z of rejected candidates in other groups by up to 0.7; the model and maps differ in their last digits. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_018My7tqt8a4LyKKjP3FhScT --- .../structure_refinement/ModelGrid.cpp | 139 ++++++++++++++++-- .../structure_refinement/ModelGrid.h | 4 +- 2 files changed, 133 insertions(+), 10 deletions(-) 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.