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

RESULT-CHANGING (ulp level): the model density's walk places its points in single precision

Browse Source 
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018My7tqt8a4LyKKjP3FhScT
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-18 17:04:17 +02:00
1 parent f74df573cf
commit fe58c5041c
2 files changed
+133 -10
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/ModelGrid.cpp
+130 -9
View File
@@ -130,6 +130,126 @@ void UseAtomBoxesByPlane(gemmi::Grid<float> &grid, const std::vector<AtomBox> &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<float> &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<float>(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<float>(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<float> &grid) {
const gemmi::UpperTriangularMat33 &m = grid.orth_n;
return {static_cast<float>(m.a11), static_cast<float>(m.a12), static_cast<float>(m.a13),
static_cast<float>(m.a22), static_cast<float>(m.a23), static_cast<float>(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<size_t> idx;
std::vector<float> 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 <typename FuncFor>
void UseDensityBoxesByPlane(gemmi::Grid<float> &grid, const std::vector<DensityBox> &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<std::vector<int>> plane_atoms(nw);
for (int i = 0; i < static_cast<int>(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<int> &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<size_t>(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<float>(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<float>(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<float>(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<size_t> OrbitLeaders(const gemmi::Grid<float> &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<Table, float> &dc, const gemmi::Model &model,
const std::vector<size_t> &orbit_leaders, size_t nthreads) {
using CReal = AtomDensity::CReal;
dc.initialize_grid();
const std::vector<const gemmi::Atom *> atoms = ModelAtoms(model);
const int n = static_cast<int>(atoms.size());
std::vector<AtomBox> boxes(n);
std::vector<DensityBox> boxes(n);
std::vector<AtomDensity> density(n);
ParallelChunks(n, nthreads, [&](int lo, int hi) {
for (int i = lo; i < hi; i++) {
@@ -270,17 +391,17 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator<Table, float> &dc, const gem
if (!d.is_aniso) {
const CReal b = static_cast<CReal>(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; });
}