Merge mac-mr-metal (density series) into rc176: Metal model density, bit-identical to the CPU; the CPU density walk in single precision (RESULT-CHANGING at the last bit); autorelease-pool leak fix in the Metal rotation function

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4

image_analysis/structure_refinement/CMakeLists.txt

@@ -39,5 +39,18 @@ IF (APPLE)
             MolecularReplacementMetal.metal ${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp)
     SET_SOURCE_FILES_PROPERTIES(MolecularReplacementMetal.mm PROPERTIES COMPILE_OPTIONS -fobjc-arc)
     SET_SOURCE_FILES_PROPERTIES(MolecularReplacementMetal.metal PROPERTIES HEADER_FILE_ONLY ON)
+    # The model density's gather (ModelDensityMetal.h), carried in the same way.
+    FILE(READ ${CMAKE_CURRENT_SOURCE_DIR}/ModelDensityMetal.metal text)
+    SET_PROPERTY(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/ModelDensityMetal.metal)
+    FILE(WRITE ${CMAKE_CURRENT_BINARY_DIR}/ModelDensityMetalSource.cpp.tmp
+         "// Generated from image_analysis/structure_refinement by its CMakeLists.txt - do not edit.\n"
+         "extern const char *const model_density_metal_source;\n"
+         "const char *const model_density_metal_source = R\"DENSITY_METAL(\n${text})DENSITY_METAL\";\n")
+    CONFIGURE_FILE(${CMAKE_CURRENT_BINARY_DIR}/ModelDensityMetalSource.cpp.tmp
+                   ${CMAKE_CURRENT_BINARY_DIR}/ModelDensityMetalSource.cpp COPYONLY)
+    TARGET_SOURCES(JFJochStructureRefinement PRIVATE ModelDensityMetal.mm ModelDensityMetal.h
+            ModelDensityMetal.metal ${CMAKE_CURRENT_BINARY_DIR}/ModelDensityMetalSource.cpp)
+    SET_SOURCE_FILES_PROPERTIES(ModelDensityMetal.mm PROPERTIES COMPILE_OPTIONS -fobjc-arc)
+    SET_SOURCE_FILES_PROPERTIES(ModelDensityMetal.metal PROPERTIES HEADER_FILE_ONLY ON)
     TARGET_LINK_LIBRARIES(JFJochStructureRefinement "-framework Metal" "-framework Foundation")
 ENDIF()

image_analysis/structure_refinement/ModelDensityMetal.h

@@ -0,0 +1,42 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#pragma once
+
+// PutModelDensityOnGrid()'s density (ModelGrid.cpp) on an Apple GPU (macOS only), before the symmetrize: the same
+// grid, bit for bit. A gather, as ModelDensityGPU.cu's under CUDA: the grid is cut into bricks of 4^3 points,
+// one threadgroup each, and every point adds, in model order, each atom whose box holds it and whose radius
+// reaches it - the points UseDensityBoxesByPlane() hands the atom, with the same single-precision arithmetic
+// (ModelDensityMetal.metal spells it out), gemmi's float expapprox and the same order of the sums. Nothing is
+// summed with an atomic. tests/ModelDensityMetalTest.cpp holds it to the CPU's grid.
+
+#include <cstddef>
+#include <vector>
+
+// One atom, as ModelGrid.cpp's DensityBox and AtomDensity hold it: its box about the grid point c nearest its
+// centre, c - d to c + d; gemmi's five Gaussians of the isotropic or the anisotropic density.
+struct ModelDensityMetalAtom {
+    float nctr[3];    // the centre in grid units
+    int c[3], d[3];
+    float max_r2;     // the radius squared
+    float occ;
+    int aniso;        // 0: b[k][0] multiplies r^2; 1: b[k] is u11 u22 u33 u12 u13 u23
+    float a[5];
+    float b[5][6];
+};
+
+// The grid, u fastest, and gemmi's orth_n over its size in single precision (upper triangular).
+struct ModelDensityMetalGrid {
+    int nu = 0, nv = 0, nw = 0;
+    float m00 = 0, m01 = 0, m02 = 0, m11 = 0, m12 = 0, m22 = 0;
+};
+
+// The density of `atoms` into `out` (nu nv nw floats, every point written) on the GPU. False where it is not
+// taken there - no Metal device the process may use (rugnux --no-gpu hides it, see hide_gpus), or an atom whose
+// box is wider than the cell, which reaches a point through two images;
+// `out` is then untouched. `nthreads` CPU threads sort the atoms into the bricks. Throws JFJochException on a
+// Metal failure.
+bool ModelDensityMetal(const ModelDensityMetalGrid &grid, const std::vector<ModelDensityMetalAtom> &atoms, float *out,
+                       size_t nthreads);
+// How many grids the GPU has put the density on in this process.
+size_t ModelDensityMetalGrids();

image_analysis/structure_refinement/ModelDensityMetal.metal

@@ -0,0 +1,108 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+// The model density's gather (ModelDensityMetal.h), compiled when it is first used from this source, which
+// CMakeLists.txt copies into the binary - with fast math off and no a*b+c contracted, as on the host, so that
+// every operation below rounds as the CPU's does.
+
+#include <metal_stdlib>
+using namespace metal;
+
+#pragma clang fp contract(off)
+
+// A brick's edge in points (ModelDensityMetal.mm says the same): every point of a brick tests each atom whose box
+// reaches the brick, and with 4 rather than CUDA's 8 far fewer of them test atoms that miss them - the kernel
+// took a half to a third of the time on an Apple M6.
+constant int BRICK = 4;
+
+struct Atom {   // ModelDensityMetalAtom
+    float nctr[3];
+    int c[3], d[3];
+    float max_r2;
+    float occ;
+    int aniso;
+    float a[5];
+    float b[5][6];
+};
+
+// Where an atom's box starts along each axis, wrapped into the grid, and its width less one (2 d): the test
+// whether a point is in the box, without a division.
+struct Box {
+    int first[3];
+    int span[3];
+};
+
+struct Grid {   // ModelDensityMetalGrid
+    int nu, nv, nw;
+    float m00, m01, m02, m11, m12, m22;
+};
+
+// gemmi's unsafe_expapprox() (formfact.hpp, MPL-2.0; https://github.com/project-gemmi/gemmi), which is based on
+// expapprox() from https://github.com/jhjourdan/SIMD-math-prims: the same operations, the bits moved with as_type
+// where gemmi uses memcpy.
+static float expapprox(float x) {
+    const float val = 12102203.1615614f * x + 1065353216.f;
+    const int vali = static_cast<int>(val);
+    const int xu1 = vali & 0x7F800000;
+    const int xu2 = (vali & 0x7FFFFF) | 0x3F800000;
+    const float a = as_type<float>(xu1);
+    const float b = as_type<float>(xu2);
+    return a * (0.509871020f + b * (0.312146713f + b * (0.166617139f + b *
+            (-2.190619930e-3f + b * 1.3555747234e-2f))));
+}
+
+// The point p's offset in the box along an axis of n points, 0 to span where it is in it.
+static int box_offset(int p, int first, int n) {
+    const int m = p - first;
+    return m < 0 ? m + n : m;
+}
+
+// One threadgroup per brick of BRICK^3 points: each point adds the brick's atoms in model order, as
+// UseDensityBoxesByPlane() and AtomDensity (ModelGrid.cpp) add them on the CPU.
+kernel void model_density(device const Atom *atoms [[buffer(0)]], constant Grid &g [[buffer(1)]],
+                          device const uint *pair_atom [[buffer(2)]], device const uint *brick_first [[buffer(3)]],
+                          device float *out [[buffer(4)]], device const Box *boxes [[buffer(5)]],
+                          uint3 group [[threadgroup_position_in_grid]], uint3 t [[thread_position_in_threadgroup]]) {
+    const uint brick = group.x;
+    const int nbu = (g.nu + BRICK - 1) / BRICK, nbv = (g.nv + BRICK - 1) / BRICK;
+    const int u = int(brick % uint(nbu)) * BRICK + int(t.x);
+    const int v = int(brick / uint(nbu) % uint(nbv)) * BRICK + int(t.y);
+    const int w = int(brick / uint(nbu * nbv)) * BRICK + int(t.z);
+    if (u >= g.nu || v >= g.nv || w >= g.nw)
+        return;
+    float acc = 0.0f;
+    for (uint k = brick_first[brick]; k < brick_first[brick + 1]; k++) {
+        const uint i = pair_atom[k];
+        device const Box &box = boxes[i];
+        const int mu = box_offset(u, box.first[0], g.nu), mv = box_offset(v, box.first[1], g.nv),
+                  mw = box_offset(w, box.first[2], g.nw);
+        if (mu > box.span[0] || mv > box.span[1] || mw > box.span[2])
+            continue;
+        device const Atom &at = atoms[i];
+        // The point's index in the box's own (unwrapped) range, as the CPU walks it.
+        const int uu = at.c[0] - at.d[0] + mu, vv = at.c[1] - at.d[1] + mv, ww = at.c[2] - at.d[2] + mw;
+        const float fx = at.nctr[0] - float(uu), fy = at.nctr[1] - float(vv), fz = at.nctr[2] - float(ww);
+        const float z = g.m22 * fz;
+        const float y = g.m11 * fy + g.m12 * fz;
+        const float yz = y * y + z * z;
+        const float x = g.m00 * fx + (g.m01 * fy + g.m02 * fz);
+        const float r2 = yz + x * x;
+        if (r2 > at.max_r2)
+            continue;
+        float density = 0.0f;
+        if (at.aniso == 0) {
+            for (int q = 0; q < 5; q++)
+                density += at.a[q] * expapprox(max(at.b[q][0] * r2, -88.f));
+        } else {
+            for (int q = 0; q < 5; q++) {
+                device const float *b = at.b[q];
+                // gemmi's SMat33::r_u_r(), in its order
+                const float rur = x * x * b[0] + y * y * b[1] + z * z * b[2]
+                                  + 2 * (x * y * b[3] + x * z * b[4] + y * z * b[5]);
+                density += at.a[q] * expapprox(max(rur, -88.f));
+            }
+        }
+        acc += at.occ * density;
+    }
+    out[uint(u) + uint(g.nu) * (uint(v) + uint(g.nv) * uint(w))] = acc;
+}

image_analysis/structure_refinement/ModelDensityMetal.mm

@@ -0,0 +1,196 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#import <Metal/Metal.h>
+
+#include "ModelDensityMetal.h"
+
+#include <algorithm>
+#include <atomic>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <string>
+#include <unistd.h>
+
+#include "../../common/CUDAWrapper.h"   // gpus_are_hidden
+#include "../../common/JFJochException.h"
+#include "../../common/ParallelFor.h"
+
+// The kernel's source, ModelDensityMetal.metal, copied in at configure time (CMakeLists.txt) so the binary needs
+// no file beside it.
+extern const char *const model_density_metal_source;
+
+namespace {
+
+constexpr int BRICK = 4;   // as in the kernel
+
+std::atomic<size_t> grids_on_gpu{0};
+
+static_assert(sizeof(ModelDensityMetalAtom) == 47 * 4, "the kernel's Atom is 47 words, none of them padding");
+
+struct Kernel {
+    id<MTLDevice> device = nil;   // nil: no Metal device
+    id<MTLComputePipelineState> pipeline = nil;
+    id<MTLCommandQueue> queue = nil;
+};
+
+// Compiled once per process, the first time it is asked for.
+const Kernel &GetKernel() {
+    static const Kernel kernel = [] {
+        Kernel k;
+        k.device = MTLCreateSystemDefaultDevice();
+        if (k.device == nil)
+            return k;
+        MTLCompileOptions *options = [MTLCompileOptions new];
+        // Fast math off: nothing reassociated or approximated, as on the CPU (MTLCompileOptions.mathMode says
+        // the same from macOS 15; the build targets 13). The language of macOS 13 on every macOS.
+        options.fastMathEnabled = NO;
+        options.languageVersion = MTLLanguageVersion3_0;
+        NSError *error = nil;
+        id<MTLLibrary> library = [k.device newLibraryWithSource:@(model_density_metal_source) options:options
+                                                          error:&error];
+        if (library == nil)
+            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
+                                  std::string("model density kernel: ") + error.localizedDescription.UTF8String);
+        id<MTLFunction> function = [library newFunctionWithName:@"model_density"];
+        k.pipeline = [k.device newComputePipelineStateWithFunction:function error:&error];
+        if (k.pipeline == nil)
+            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
+                                  std::string("model_density: ") + error.localizedDescription.UTF8String);
+        k.queue = [k.device newCommandQueue];
+        return k;
+    }();
+    return kernel;
+}
+
+id<MTLBuffer> NewBuffer(id<MTLDevice> device, size_t bytes) {
+    id<MTLBuffer> b = [device newBufferWithLength:std::max<size_t>(bytes, 4) options:MTLResourceStorageModeShared];
+    if (b == nil)
+        throw JFJochException(JFJochExceptionCategory::GPUMetalError, "model density: out of GPU memory");
+    return b;
+}
+
+// The bricks along an axis of n points that a box's points first .. first + span (wrapped, span < n, first in
+// [0, n)) fall in, each once, into out; how many.
+int AxisBricks(int first, int span, int n, int *out) {
+    int count = 0;
+    for (int p = first; p <= first + span;) {
+        const int q = p < n ? p : p - n;
+        const int b = q / BRICK;
+        if (std::find(out, out + count, b) == out + count)
+            out[count++] = b;
+        p += std::min(BRICK - q % BRICK, n - q);   // to the next brick, or past the axis' last point
+    }
+    return count;
+}
+
+}  // namespace
+
+bool ModelDensityMetal(const ModelDensityMetalGrid &grid, const std::vector<ModelDensityMetalAtom> &atoms,
+                       float *out, size_t nthreads) {
+    // In a pool of its own: the command buffer and its encoder are autoreleased, and the threads this runs on have
+    // none, so they - and every buffer they hold - would otherwise live as long as the thread.
+    @autoreleasepool {
+        const size_t points = static_cast<size_t>(grid.nu) * grid.nv * grid.nw;
+        if (gpus_are_hidden() || points > std::numeric_limits<uint32_t>::max())
+            return false;
+        const int n3[3] = {grid.nu, grid.nv, grid.nw};
+        for (const ModelDensityMetalAtom &a : atoms)
+            for (int k = 0; k < 3; k++)
+                if (2 * a.d[k] + 1 > n3[k])
+                    return false;   // a box wider than the cell: gemmi's walk adds both images of the atom
+        const Kernel &k = GetKernel();
+        if (k.device == nil)
+            return false;
+
+        // Where each atom's box starts along each axis, wrapped into the grid, and its width less one.
+        id<MTLBuffer> d_boxes = NewBuffer(k.device, atoms.size() * 6 * sizeof(int32_t));
+        auto *boxes = static_cast<int32_t *>(d_boxes.contents);
+        for (size_t i = 0; i < atoms.size(); i++)
+            for (int a = 0; a < 3; a++) {
+                const int first_point = (atoms[i].c[a] - atoms[i].d[a]) % n3[a];
+                boxes[6 * i + a] = first_point < 0 ? first_point + n3[a] : first_point;
+                boxes[6 * i + 3 + a] = 2 * atoms[i].d[a];
+            }
+        // Each brick's atoms in model order: the (brick, atom) pairs counted, then filled, each worker taking the
+        // bricks of a slab of w - a contiguous range of them - and the atoms in order.
+        const int nb3[3] = {(grid.nu + BRICK - 1) / BRICK, (grid.nv + BRICK - 1) / BRICK,
+                            (grid.nw + BRICK - 1) / BRICK};
+        const size_t bricks = static_cast<size_t>(nb3[0]) * nb3[1] * nb3[2];
+        std::vector<uint32_t> first(bricks + 1, 0);
+        const auto pairs_of_slab = [&](int z0, int z1, const auto &pair) {
+            std::vector<int> bu(nb3[0]), bv(nb3[1]), bw(nb3[2]);
+            for (size_t i = 0; i < atoms.size(); i++) {
+                const int32_t *box = boxes + 6 * i;
+                const int nw = AxisBricks(box[2], box[5], grid.nw, bw.data());
+                int nv = -1, nu = 0;
+                for (int kw = 0; kw < nw; kw++) {
+                    if (bw[kw] < z0 || bw[kw] >= z1)
+                        continue;
+                    if (nv < 0) {
+                        nv = AxisBricks(box[1], box[4], grid.nv, bv.data());
+                        nu = AxisBricks(box[0], box[3], grid.nu, bu.data());
+                    }
+                    for (int kv = 0; kv < nv; kv++)
+                        for (int ku = 0; ku < nu; ku++)
+                            pair(bu[ku] + static_cast<size_t>(nb3[0]) * (bv[kv] + static_cast<size_t>(nb3[1]) * bw[kw]),
+                                 i);
+                }
+            }
+        };
+        ParallelChunks(nb3[2], nthreads, [&](int z0, int z1) {
+            pairs_of_slab(z0, z1, [&](size_t brick, size_t) { first[brick + 1]++; });
+        });
+        for (size_t b = 0; b < bricks; b++)
+            first[b + 1] += first[b];
+        id<MTLBuffer> d_atoms = NewBuffer(k.device, atoms.size() * sizeof(ModelDensityMetalAtom));
+        id<MTLBuffer> d_pairs = NewBuffer(k.device, first[bricks] * sizeof(uint32_t));
+        id<MTLBuffer> d_first = NewBuffer(k.device, first.size() * sizeof(uint32_t));
+        std::memcpy(d_atoms.contents, atoms.data(), atoms.size() * sizeof(ModelDensityMetalAtom));
+        std::memcpy(d_first.contents, first.data(), first.size() * sizeof(uint32_t));
+        auto *pairs = static_cast<uint32_t *>(d_pairs.contents);
+        ParallelChunks(nb3[2], nthreads, [&](int z0, int z1) {
+            pairs_of_slab(z0, z1, [&](size_t brick, size_t i) { pairs[first[brick]++] = static_cast<uint32_t>(i); });
+        });
+        // The fill moved every brick's start to the next one's: d_first keeps the starts.
+
+        // The kernel writes into `out` itself where its pages can be handed to the GPU as they are; into a buffer
+        // of its own, copied out, otherwise.
+        const size_t page = static_cast<size_t>(getpagesize());
+        const size_t bytes = points * sizeof(float);
+        id<MTLBuffer> d_out = nil;
+        if (reinterpret_cast<uintptr_t>(out) % page == 0)
+            d_out = [k.device newBufferWithBytesNoCopy:out length:(bytes + page - 1) / page * page
+                                               options:MTLResourceStorageModeShared deallocator:nil];
+        const bool in_place = d_out != nil;
+        if (!in_place)
+            d_out = NewBuffer(k.device, bytes);
+
+        id<MTLCommandBuffer> cb = [k.queue commandBuffer];
+        id<MTLComputeCommandEncoder> e = [cb computeCommandEncoder];
+        [e setComputePipelineState:k.pipeline];
+        [e setBuffer:d_atoms offset:0 atIndex:0];
+        [e setBytes:&grid length:sizeof(grid) atIndex:1];
+        [e setBuffer:d_pairs offset:0 atIndex:2];
+        [e setBuffer:d_first offset:0 atIndex:3];
+        [e setBuffer:d_out offset:0 atIndex:4];
+        [e setBuffer:d_boxes offset:0 atIndex:5];
+        [e dispatchThreadgroups:MTLSizeMake(bricks, 1, 1) threadsPerThreadgroup:MTLSizeMake(BRICK, BRICK, BRICK)];
+        [e endEncoding];
+        [cb commit];
+        [cb waitUntilCompleted];
+        if (cb.status != MTLCommandBufferStatusCompleted)
+            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
+                                  cb.error ? cb.error.localizedDescription.UTF8String
+                                           : "model density: command buffer failed");
+        if (!in_place)
+            std::memcpy(out, d_out.contents, bytes);
+        grids_on_gpu++;
+        return true;
+    }
+}
+
+size_t ModelDensityMetalGrids() {
+    return grids_on_gpu;
+}

image_analysis/structure_refinement/ModelGrid.cpp

@@ -12,6 +12,13 @@
 #include "gemmi/solmask.hpp"   // SolventMasker
 
 #include "../../common/ParallelFor.h"
+#ifdef __APPLE__
+#include <atomic>
+
+#include "../../common/JFJochException.h"
+#include "../../common/Logger.h"
+#include "ModelDensityMetal.h"
+#endif
 
 namespace {
 
@@ -130,6 +137,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 +310,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 +321,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 +331,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)
@@ -217,6 +344,56 @@ struct AtomDensity {
     }
 };
 
+#ifdef __APPLE__
+// UseDensityBoxesByPlane() with AtomDensity on the Apple GPU (ModelDensityMetal.h): the same grid, bit for bit.
+// False where the GPU does not take it - and from the first time it fails on, with one warning, as the CPU
+// then does the same work.
+std::atomic<bool> density_metal_failed{false};
+
+bool PutDensityOnGridMetal(gemmi::Grid<float> &grid, const std::vector<DensityBox> &boxes,
+                           const std::vector<AtomDensity> &density, size_t nthreads) {
+    if (density_metal_failed)
+        return false;
+    static_assert(std::size(AtomDensity::IsoSum{}.a) == 5 && std::size(AtomDensity::AnisoSum{}.a) == 5);
+    const DensityGeometry m = ModelDensityGeometry(grid);
+    const ModelDensityMetalGrid g{grid.nu, grid.nv, grid.nw, m.m00, m.m01, m.m02, m.m11, m.m12, m.m22};
+    std::vector<ModelDensityMetalAtom> atoms(boxes.size());
+    for (size_t i = 0; i < boxes.size(); i++) {
+        ModelDensityMetalAtom &a = atoms[i];
+        const DensityBox &b = boxes[i];
+        const AtomDensity &d = density[i];
+        for (int k = 0; k < 3; k++) {
+            a.nctr[k] = b.nctr[k];
+            a.c[k] = b.c[k];
+            a.d[k] = b.d[k];
+        }
+        a.max_r2 = b.max_r2;
+        a.occ = d.occ;
+        a.aniso = d.is_aniso ? 1 : 0;
+        for (int q = 0; q < 5; q++) {
+            std::fill(std::begin(a.b[q]), std::end(a.b[q]), 0.0f);
+            if (!d.is_aniso) {
+                a.a[q] = d.iso.a[q];
+                a.b[q][0] = d.iso.b[q];
+            } else {
+                const auto &u = d.aniso.b[q];
+                a.a[q] = d.aniso.a[q];
+                const float row[6] = {u.u11, u.u22, u.u33, u.u12, u.u13, u.u23};
+                std::copy(std::begin(row), std::end(row), a.b[q]);
+            }
+        }
+    }
+    try {
+        return ModelDensityMetal(g, atoms, grid.data.data(), nthreads);
+    } catch (const JFJochException &e) {
+        if (!density_metal_failed.exchange(true))
+            Logger("Rugnux").Warning("Model density: the Apple GPU failed, on the CPU from here on: {}", e.what());
+        std::fill(grid.data.begin(), grid.data.end(), 0.0f);   // the CPU adds to a zero grid; the GPU may have written
+        return false;
+    }
+}
+#endif
+
 }  // namespace
 
 // The points gemmi's Grid::symmetrize() (grid.hpp) reduces the grid's orbits into, ascending. gemmi
@@ -248,15 +425,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) {
+                           const std::vector<size_t> &orbit_leaders, size_t nthreads, [[maybe_unused]] bool gpu) {
     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 +448,20 @@ 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]; });
+#ifdef __APPLE__
+    if (!gpu || !PutDensityOnGridMetal(dc.grid, boxes, density, nthreads))
+#endif
+        UseDensityBoxesByPlane(dc.grid, boxes, nthreads, [&](int i) { return density[i]; });
     SymmetrizeOrbits(dc.grid, orbit_leaders, nthreads, [](float a, float b) { return a + b; });
 }
 

image_analysis/structure_refinement/ModelGrid.h

@@ -14,15 +14,19 @@
 // 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.
 std::vector<size_t> OrbitLeaders(const gemmi::Grid<float> &grid, size_t nthreads);
 
-// `dc` set up as for put_model_density_on_grid(): d_min, rate, blur, the grid's cell and group.
+// `dc` set up as for put_model_density_on_grid(): d_min, rate, blur, the grid's cell and group. On macOS the
+// Apple GPU puts the density on the grid (ModelDensityMetal.h), to the same bits; `gpu` false keeps it on the
+// CPU.
 void PutModelDensityOnGrid(gemmi::DensityCalculator<gemmi::IT92<float>, float> &dc, const gemmi::Model &model,
-                           const std::vector<size_t> &orbit_leaders, size_t nthreads);
+                           const std::vector<size_t> &orbit_leaders, size_t nthreads, bool gpu = true);
 
 // `grid` sized, with its cell and group set.
 void PutMaskOnGrid(gemmi::Grid<float> &grid, const gemmi::Model &model, const std::vector<size_t> &orbit_leaders,

image_analysis/structure_refinement/MolecularReplacementMetal.mm

@@ -109,7 +109,8 @@ MolecularReplacementMetal::MolecularReplacementMetal(const MolecularTransform &f
         for (int l = 2 * bl; l < std::min(2 * bl + 2, g.nw); l++)
             for (int j = 0; j < g.nv; j++)
                 for (int i = 0; i < g.nu; i++) {
-                    const size_t brick = (i >> 1) + static_cast<size_t>(bu) * ((j >> 1) + static_cast<size_t>(bv) * (l >> 1));
+                    const size_t brick =
+                            (i >> 1) + static_cast<size_t>(bu) * ((j >> 1) + static_cast<size_t>(bv) * (l >> 1));
                     out[brick * 8 + (i & 1) + 2 * (j & 1) + 4 * (l & 1)] =
                             g.data[i + static_cast<size_t>(g.nu) * (j + static_cast<size_t>(g.nv) * l)];
                 }
@@ -122,94 +123,99 @@ std::vector<double> MolecularReplacementMetal::RotationFunction(const SearchRefl
                                                                 const gemmi::UnitCell &cell,
                                                                 const gemmi::SpaceGroup &sg,
                                                                 const std::vector<gemmi::Mat33> &rotations) const {
-    const Kernel &k = GetKernel();
-    const size_t n = refl.hkl.size();
-    const std::vector<gemmi::Vec3> s = CopyVectors(refl, cell, sg);
-    const size_t nops = s.size() / n;
-    double w_mean = 0;
-    for (double e2 : refl.e2)
-        w_mean += e2;
-    w_mean /= static_cast<double>(n);
-    double w_ss = 0;
-    for (double e2 : refl.e2)
-        w_ss += (e2 - w_mean) * (e2 - w_mean);
+    // In a pool of its own: the command buffer and its encoder are autoreleased, and the threads this runs on have
+    // none, so they - and every buffer they hold - would otherwise live as long as the thread.
+    @autoreleasepool {
+        const Kernel &k = GetKernel();
+        const size_t n = refl.hkl.size();
+        const std::vector<gemmi::Vec3> s = CopyVectors(refl, cell, sg);
+        const size_t nops = s.size() / n;
+        double w_mean = 0;
+        for (double e2 : refl.e2)
+            w_mean += e2;
+        w_mean /= static_cast<double>(n);
+        double w_ss = 0;
+        for (double e2 : refl.e2)
+            w_ss += (e2 - w_mean) * (e2 - w_mean);
 
-    // The reflections in the order of their bins, each bin's from bin_first[b] to bin_first[b + 1].
-    std::vector<size_t> order(n);
-    std::iota(order.begin(), order.end(), 0);
-    std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { return refl.bin[a] < refl.bin[b]; });
-    std::vector<uint32_t> bin_first(refl.bins + 1, 0);
-    for (size_t i = 0; i < n; i++)
-        bin_first[refl.bin[i] + 1]++;
-    for (int b = 0; b < refl.bins; b++)
-        bin_first[b + 1] += bin_first[b];
-    std::vector<float> sf(3 * n * nops), w(n);
-    for (size_t i = 0; i < n; i++) {
-        for (size_t j = 0; j < nops; j++) {
-            const gemmi::Vec3 &v = s[order[i] * nops + j];
-            sf[3 * (i * nops + j)] = static_cast<float>(v.x);
-            sf[3 * (i * nops + j) + 1] = static_cast<float>(v.y);
-            sf[3 * (i * nops + j) + 2] = static_cast<float>(v.z);
+        // The reflections in the order of their bins, each bin's from bin_first[b] to bin_first[b + 1].
+        std::vector<size_t> order(n);
+        std::iota(order.begin(), order.end(), 0);
+        std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { return refl.bin[a] < refl.bin[b]; });
+        std::vector<uint32_t> bin_first(refl.bins + 1, 0);
+        for (size_t i = 0; i < n; i++)
+            bin_first[refl.bin[i] + 1]++;
+        for (int b = 0; b < refl.bins; b++)
+            bin_first[b + 1] += bin_first[b];
+        std::vector<float> sf(3 * n * nops), w(n);
+        for (size_t i = 0; i < n; i++) {
+            for (size_t j = 0; j < nops; j++) {
+                const gemmi::Vec3 &v = s[order[i] * nops + j];
+                sf[3 * (i * nops + j)] = static_cast<float>(v.x);
+                sf[3 * (i * nops + j) + 1] = static_cast<float>(v.y);
+                sf[3 * (i * nops + j) + 2] = static_cast<float>(v.z);
+            }
+            w[i] = static_cast<float>(refl.e2[order[i]] - w_mean);
         }
-        w[i] = static_cast<float>(refl.e2[order[i]] - w_mean);
-    }
-    // R^T as Metal's float3x3: three columns of four floats.
-    std::vector<float> rot(12 * rotations.size(), 0.0f);
-    for (size_t r = 0; r < rotations.size(); r++)
-        for (int col = 0; col < 3; col++)
-            for (int row = 0; row < 3; row++)
-                rot[12 * r + 4 * col + row] = static_cast<float>(rotations[r].a[col][row]);
+        // R^T as Metal's float3x3: three columns of four floats.
+        std::vector<float> rot(12 * rotations.size(), 0.0f);
+        for (size_t r = 0; r < rotations.size(); r++)
+            for (int col = 0; col < 3; col++)
+                for (int row = 0; row < 3; row++)
+                    rot[12 * r + 4 * col + row] = static_cast<float>(rotations[r].a[col][row]);
 
-    Params params = state_->params;
-    params.nops = static_cast<uint32_t>(nops);
-    params.bins = static_cast<uint32_t>(refl.bins);
-    id<MTLBuffer> d_s = Upload(k.device, sf.data(), sf.size() * sizeof(float));
-    id<MTLBuffer> d_w = Upload(k.device, w.data(), w.size() * sizeof(float));
-    id<MTLBuffer> d_bin_first = Upload(k.device, bin_first.data(), bin_first.size() * sizeof(uint32_t));
-    id<MTLBuffer> d_rot = Upload(k.device, rot.data(), rot.size() * sizeof(float));
-    id<MTLBuffer> d_out = [k.device newBufferWithLength:std::max<size_t>(rotations.size() * refl.bins * 3 * sizeof(float), 4)
-                                                options:MTLResourceStorageModeShared];
-    if (d_out == nil)
-        throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory");
-    std::vector<id<MTLCommandBuffer>> running;
-    for (size_t first = 0; first < rotations.size(); first += BATCH) {
-        const size_t count = std::min(BATCH, rotations.size() - first);
-        id<MTLCommandBuffer> cb = [state_->queue commandBuffer];
-        id<MTLComputeCommandEncoder> e = [cb computeCommandEncoder];
-        [e setComputePipelineState:k.pipeline];
-        [e setBuffer:state_->grid offset:0 atIndex:0];
-        [e setBytes:&params length:sizeof(params) atIndex:1];
-        [e setBuffer:d_rot offset:first * 12 * sizeof(float) atIndex:2];
-        [e setBuffer:d_s offset:0 atIndex:3];
-        [e setBuffer:d_w offset:0 atIndex:4];
-        [e setBuffer:d_bin_first offset:0 atIndex:5];
-        [e setBuffer:d_out offset:first * refl.bins * 3 * sizeof(float) atIndex:6];
-        [e dispatchThreadgroups:MTLSizeMake(count, 1, 1) threadsPerThreadgroup:MTLSizeMake(THREADS, 1, 1)];
-        [e endEncoding];
-        [cb commit];
-        running.push_back(cb);
-    }
-    for (id<MTLCommandBuffer> cb : running) {
-        [cb waitUntilCompleted];
-        if (cb.status != MTLCommandBufferStatusCompleted)
-            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
-                                  cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed");
-    }
-
-    // The correlation of E^2 - 1 with ic over its bin's mean, from the bins' sums: x = ic / mean_b.
-    const auto *sums = static_cast<const float *>(d_out.contents);
-    std::vector<double> out(rotations.size());
-    for (size_t r = 0; r < rotations.size(); r++) {
-        double x_sum = 0, x_ss = 0, xy = 0;
-        for (int b = 0; b < refl.bins; b++) {
-            const float *o = sums + (r * refl.bins + b) * 3;
-            const double mean = o[0] / static_cast<double>(bin_first[b + 1] - bin_first[b]);
-            x_sum += o[0] / mean;
-            x_ss += o[1] / (mean * mean);
-            xy += o[2] / mean;
+        Params params = state_->params;
+        params.nops = static_cast<uint32_t>(nops);
+        params.bins = static_cast<uint32_t>(refl.bins);
+        id<MTLBuffer> d_s = Upload(k.device, sf.data(), sf.size() * sizeof(float));
+        id<MTLBuffer> d_w = Upload(k.device, w.data(), w.size() * sizeof(float));
+        id<MTLBuffer> d_bin_first = Upload(k.device, bin_first.data(), bin_first.size() * sizeof(uint32_t));
+        id<MTLBuffer> d_rot = Upload(k.device, rot.data(), rot.size() * sizeof(float));
+        const size_t out_bytes = rotations.size() * refl.bins * 3 * sizeof(float);
+        id<MTLBuffer> d_out = [k.device newBufferWithLength:std::max<size_t>(out_bytes, 4)
+                                                    options:MTLResourceStorageModeShared];
+        if (d_out == nil)
+            throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory");
+        std::vector<id<MTLCommandBuffer>> running;
+        for (size_t first = 0; first < rotations.size(); first += BATCH) {
+            const size_t count = std::min(BATCH, rotations.size() - first);
+            id<MTLCommandBuffer> cb = [state_->queue commandBuffer];
+            id<MTLComputeCommandEncoder> e = [cb computeCommandEncoder];
+            [e setComputePipelineState:k.pipeline];
+            [e setBuffer:state_->grid offset:0 atIndex:0];
+            [e setBytes:&params length:sizeof(params) atIndex:1];
+            [e setBuffer:d_rot offset:first * 12 * sizeof(float) atIndex:2];
+            [e setBuffer:d_s offset:0 atIndex:3];
+            [e setBuffer:d_w offset:0 atIndex:4];
+            [e setBuffer:d_bin_first offset:0 atIndex:5];
+            [e setBuffer:d_out offset:first * refl.bins * 3 * sizeof(float) atIndex:6];
+            [e dispatchThreadgroups:MTLSizeMake(count, 1, 1) threadsPerThreadgroup:MTLSizeMake(THREADS, 1, 1)];
+            [e endEncoding];
+            [cb commit];
+            running.push_back(cb);
         }
-        const double sxx = x_ss - x_sum * x_sum / static_cast<double>(n);
-        out[r] = sxx > 0 && w_ss > 0 ? xy / std::sqrt(sxx * w_ss) : 0.0;
+        for (id<MTLCommandBuffer> cb : running) {
+            [cb waitUntilCompleted];
+            if (cb.status != MTLCommandBufferStatusCompleted)
+                throw JFJochException(JFJochExceptionCategory::GPUMetalError,
+                                      cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed");
+        }
+
+        // The correlation of E^2 - 1 with ic over its bin's mean, from the bins' sums: x = ic / mean_b.
+        const auto *sums = static_cast<const float *>(d_out.contents);
+        std::vector<double> out(rotations.size());
+        for (size_t r = 0; r < rotations.size(); r++) {
+            double x_sum = 0, x_ss = 0, xy = 0;
+            for (int b = 0; b < refl.bins; b++) {
+                const float *o = sums + (r * refl.bins + b) * 3;
+                const double mean = o[0] / static_cast<double>(bin_first[b + 1] - bin_first[b]);
+                x_sum += o[0] / mean;
+                x_ss += o[1] / (mean * mean);
+                xy += o[2] / mean;
+            }
+            const double sxx = x_ss - x_sum * x_sum / static_cast<double>(n);
+            out[r] = sxx > 0 && w_ss > 0 ? xy / std::sqrt(sxx * w_ss) : 0.0;
+        }
+        return out;
     }
-    return out;
 }

tests/CMakeLists.txt

@@ -22,6 +22,7 @@ ADD_EXECUTABLE(jfjoch_portable_test EXCLUDE_FROM_ALL
         BraggIntegrationEngineCPUTest.cpp
         BraggIntegrationEngineMetalTest.cpp
         MolecularReplacementMetalTest.cpp
+        ModelDensityMetalTest.cpp
         WriteReflectionsTest.cpp
         LoadReferenceMtzFreeFlagsTest.cpp
         LoadReferenceSfCifTest.cpp

tests/ModelDensityMetalTest.cpp

@@ -0,0 +1,91 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#include <catch2/catch_all.hpp>
+
+#ifdef __APPLE__
+
+#include <cstring>
+#include <random>
+#include <vector>
+
+#include <gemmi/dencalc.hpp>
+#include <gemmi/it92.hpp>
+
+#include "../image_analysis/structure_refinement/ModelDensityMetal.h"
+#include "../image_analysis/structure_refinement/ModelGrid.h"
+
+// The model density on the Apple GPU is the CPU's, bit for bit (ModelDensityMetal.h), so the two grids are compared
+// as bit patterns. A Mac without a Metal device fails rather than skips - the macOS CI job is the only place this
+// can be checked, as in BraggIntegrationEngineMetalTest.cpp.
+
+namespace {
+    // Atoms scattered over the whole cell - so that boxes wrap at every face - of a few elements, a quarter of them
+    // anisotropic, with occupancies below one.
+    gemmi::Structure random_crystal(const gemmi::UnitCell &cell, const char *group) {
+        std::mt19937 rng(11);
+        std::uniform_real_distribution<double> frac(0.0, 1.0), b(8.0, 60.0), occ(0.3, 1.0), u(0.05, 0.4);
+        const gemmi::El elements[] = {gemmi::El::C, gemmi::El::N, gemmi::El::O, gemmi::El::S};
+        gemmi::Structure st;
+        st.cell = cell;
+        st.spacegroup_hm = group;
+        st.setup_cell_images();
+        gemmi::Model model(1);
+        gemmi::Chain chain("A");
+        gemmi::Residue res;
+        res.name = "UNK";
+        res.seqid = gemmi::SeqId(1, ' ');
+        for (int i = 0; i < 3000; i++) {
+            gemmi::Atom a;
+            a.name = "X";
+            a.element = elements[i % 4];
+            a.pos = cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng)));
+            a.b_iso = static_cast<float>(b(rng));
+            a.occ = static_cast<float>(occ(rng));
+            if (i % 4 == 3) {
+                const float d1 = static_cast<float>(u(rng)), d2 = static_cast<float>(u(rng)),
+                            d3 = static_cast<float>(u(rng));
+                a.aniso = {d1, d2, d3, 0.3f * std::min(d1, d2), -0.2f * std::min(d1, d3), 0.1f * std::min(d2, d3)};
+            }
+            res.atoms.push_back(a);
+        }
+        chain.residues.push_back(res);
+        model.chains.push_back(chain);
+        st.models.push_back(model);
+        return st;
+    }
+
+    void check_same_grid(const gemmi::UnitCell &cell, const char *group, double d_min) {
+        const gemmi::Structure st = random_crystal(cell, group);
+        gemmi::DensityCalculator<gemmi::IT92<float>, float> cpu, gpu;
+        for (auto *dc : {&cpu, &gpu}) {
+            dc->d_min = d_min;
+            dc->rate = 1.5;
+            dc->set_grid_cell_and_spacegroup(st);
+            dc->set_refmac_compatible_blur(st.models[0]);
+        }
+        // The orbits of the grid initialize_grid() sizes, as ModelValidation.cpp finds them.
+        gemmi::Grid<float> sized;
+        sized.unit_cell = cpu.grid.unit_cell;
+        sized.spacegroup = cpu.grid.spacegroup;
+        sized.set_size_from_spacing(cpu.requested_grid_spacing(), gemmi::GridSizeRounding::Up);
+        const std::vector<size_t> leaders = OrbitLeaders(sized, 4);
+        PutModelDensityOnGrid(cpu, st.models[0], leaders, 4, false);
+        const size_t before = ModelDensityMetalGrids();
+        PutModelDensityOnGrid(gpu, st.models[0], leaders, 4);
+        REQUIRE(ModelDensityMetalGrids() == before + 1);   // the GPU put it on the grid
+        REQUIRE(gpu.grid.data.size() == cpu.grid.data.size());
+        size_t different = 0;
+        for (size_t i = 0; i < cpu.grid.data.size(); i++)
+            different += std::memcmp(&cpu.grid.data[i], &gpu.grid.data[i], sizeof(float)) != 0;
+        CHECK(different == 0);
+    }
+}
+
+TEST_CASE("ModelDensityMetal_BitIdenticalToCPU", "[portable][Metal]") {
+    // Triclinic, so that every term of the orthogonalization is there; then a group with symmetry to reduce.
+    check_same_grid(gemmi::UnitCell(61, 68, 77, 81, 96, 107), "P 1", 2.0);
+    check_same_grid(gemmi::UnitCell(70, 80, 90, 90, 90, 90), "P 21 21 21", 2.0);
+}
+
+#endif