v1.0.0-rc.176 (#86)

* jfjoch_broker: Optional live Rugnux analysis of DECTRIS rotation collections inside the broker, with merged data, maps and the placed model written beside the master file, the rotation lattice in the scan result, and a "Rugnux live" tab and dashboard row in the web frontend.
* Rugnux: Faster processing, with molecular replacement and model validation spread over several GPUs and run on the Apple GPU on macOS.
* Rugnux: The polarization stated in miniCBF and d*TREK headers is used, twinned crystals' space groups are decided more reliably, and a run that finds no lattice still reports every image.

Reviewed-on: #86
Co-authored-by: Filip Leonarski <filip.leonarski@psi.ch>

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; });
 }
 
@@ -317,3 +498,10 @@ void PutMaskOnGrid(gemmi::Grid<float> &grid, const gemmi::Model &model, const st
     if (shrink_has_offsets)
         masker.shrink(grid);
 }
+
+size_t ModelGridPoints(const gemmi::UnitCell &cell, const gemmi::SpaceGroup *sg, double d_min, double rate) {
+    const double spacing = d_min / (2 * rate);
+    const std::array<int, 3> m = gemmi::good_grid_size({{cell.a / spacing, cell.b / spacing, cell.c / spacing}},
+                                                       gemmi::GridSizeRounding::Up, sg);
+    return static_cast<size_t>(m[0]) * static_cast<size_t>(m[1]) * static_cast<size_t>(m[2]);
+}