Jungfraujoch/image_analysis/structure_refinement/ModelDensityMetal.h

// 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();