Jungfraujoch/image_analysis/structure_refinement/ModelGrid.cpp

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "ModelGrid.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <iterator>
#include <utility>

#include "gemmi/solmask.hpp"   // SolventMasker

#include "../../common/ParallelFor.h"

namespace {

using Table = gemmi::IT92<float>;

// The box of grid points gemmi's Grid::use_points_in_box<true>() walks around one atom.
struct AtomBox {
    gemmi::Fractional fpos;
    int du = 0, dv = 0, dw = 0;
    double radius = 0;
};

// The box of use_points_around<true>(fpos, radius, ..., false), sized and clamped as gemmi does it.
AtomBox MakeAtomBox(const gemmi::Grid<float> &grid, const gemmi::Fractional &fpos, double radius) {
    AtomBox box;
    box.fpos = fpos;
    box.radius = radius;
    box.du = std::min(static_cast<int>(std::ceil(radius / grid.spacing[0])), grid.nu - 1);
    box.dv = std::min(static_cast<int>(std::ceil(radius / grid.spacing[1])), grid.nv - 1);
    box.dw = std::min(static_cast<int>(std::ceil(radius / grid.spacing[2])), grid.nw - 1);
    return box;
}

// The points of one atom's box on one plane that lie 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 BoxPoints {
    int n = 0;
    std::vector<size_t> idx;
    std::vector<double> d2, dx, dy, dz;
    // Room for `capacity` points, none of them set.
    void Clear(size_t capacity) {
        n = 0;
        if (idx.size() < capacity) {
            idx.resize(capacity);
            d2.resize(capacity);
            dx.resize(capacity);
            dy.resize(capacity);
            dz.resize(capacity);
        }
    }
};

// gemmi's use_points_in_box<true>() run over a list of atoms (grid.hpp, do_use_points_in_box), in
// parallel over the grid's w planes: a plane is one task, and it walks every atom whose box reaches
// it, in list order, visiting only that plane's points. Every point is so handed exactly the calls
// gemmi's serial loop over the atoms hands it, in the same order, and the grid comes out bit for bit
// the same whatever the thread count. func_for(i) gives atom i's func, taken as a local copy so that
// the compiler can hold it in registers while the atom's points are written.
//
// gemmi hands func one point at a time, as it walks the box. Here the walk, the same one, only
// collects the points within the radius of the atom, and func(data, points) is handed them all at
// once, to visit in the order given: a loop over a few hundred points the compiler can vectorise,
// where one over a row of the box is a handful of points on the coarse grids of the rigid body.
template <typename FuncFor>
void UseAtomBoxesByPlane(gemmi::Grid<float> &grid, const std::vector<AtomBox> &boxes, size_t nthreads,
                         FuncFor func_for) {
    const int nu = grid.nu, nv = grid.nv, nw = grid.nw;
    std::vector<std::vector<int>> plane_atoms(nw);
    for (int i = 0; i < static_cast<int>(boxes.size()); i++) {
        const int w0 = gemmi::iround(boxes[i].fpos.z * nw);
        for (int w = w0 - boxes[i].dw; w <= w0 + boxes[i].dw; 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) {
        BoxPoints points;
        for (int i : plane_atoms[plane]) {
            const AtomBox &b = boxes[i];
            const auto func = func_for(i);
            const double max_dist_sq = b.radius * b.radius;
            const gemmi::Fractional nctr(b.fpos.x * nu, b.fpos.y * nv, b.fpos.z * nw);
            const int u_lo = gemmi::iround(nctr.x) - b.du, u_hi = gemmi::iround(nctr.x) + b.du;
            const int v_lo = gemmi::iround(nctr.y) - b.dv, v_hi = gemmi::iround(nctr.y) + b.dv;
            const int w_lo = gemmi::iround(nctr.z) - b.dw, w_hi = gemmi::iround(nctr.z) + b.dw;
            const int u_0 = gemmi::modulo(u_lo, nu), v_0 = gemmi::modulo(v_lo, nv);
            gemmi::Fractional fdelta(nctr.x - u_lo, 0, 0);
            // At most the (2du+1)(2dv+1) points of the box on each of its w that fall on this plane (below).
            points.Clear(static_cast<size_t>(2 * b.du + 1) * (2 * b.dv + 1) * (1 + 2 * b.dw / nw));
            size_t *idx = points.idx.data();
            double *d2 = points.d2.data(), *dx = points.dx.data(), *dy = points.dy.data(), *dz = points.dz.data();
            int n = 0;
            // The box's w that fall on this plane, ascending as gemmi walks them: one, unless the box is
            // wider than the cell.
            for (int w = w_lo + gemmi::modulo(plane - w_lo, nw); w <= w_hi; w += nw) {
                fdelta.z = nctr.z - w;
                for (int v = v_lo, v_ = v_0; v <= v_hi; ++v, v_ = (v_ + 1 == nv ? 0 : v_ + 1)) {
                    fdelta.y = nctr.y - v;
                    gemmi::Position delta(grid.orth_n.multiply(fdelta));
                    const double dist_sq0 = gemmi::sq(delta.y) + gemmi::sq(delta.z);
                    if (dist_sq0 > max_dist_sq)
                        continue;
                    const size_t row = grid.index_q(0, v_, plane);
                    for (int u = u_lo, u_ = u_0;;) {
                        const double dist_sq = dist_sq0 + gemmi::sq(delta.x);
                        if (!(dist_sq > max_dist_sq)) {
                            idx[n] = row + u_;
                            d2[n] = dist_sq;
                            dx[n] = delta.x;
                            dy[n] = delta.y;
                            dz[n] = delta.z;
                            ++n;
                        }
                        if (u >= u_hi)
                            break;
                        ++u;
                        u_ = (u_ + 1 == nu ? 0 : u_ + 1);
                        delta.x -= grid.orth_n.a11;
                    }
                }
            }
            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.
template <typename Func>
void SymmetrizeOrbits(gemmi::Grid<float> &grid, const std::vector<size_t> &leaders, size_t nthreads,
                      Func func) {
    const std::vector<gemmi::GridOp> ops = grid.get_scaled_ops_except_id();
    const size_t nu = grid.nu, nv = grid.nv;
    ParallelChunks(static_cast<int>(leaders.size()), nthreads, [&](int lo, int hi) {
        std::vector<size_t> mates(ops.size());
        for (int i = lo; i < hi; i++) {
            const size_t idx = leaders[i];
            const int u = static_cast<int>(idx % nu), v = static_cast<int>(idx / nu % nv),
                      w = static_cast<int>(idx / (nu * nv));
            for (size_t k = 0; k < ops.size(); ++k) {
                const std::array<int, 3> t = ops[k].apply(u, v, w);
                mates[k] = grid.index_n(t[0], t[1], t[2]);
            }
            float value = grid.data[idx];
            for (size_t k : mates)
                value = func(value, grid.data[k]);
            grid.data[idx] = value;
            for (size_t k : mates)
                grid.data[k] = value;
        }
    });
}

std::vector<const gemmi::Atom *> ModelAtoms(const gemmi::Model &model) {
    std::vector<const gemmi::Atom *> atoms;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms)
                atoms.push_back(&a);
    return atoms;
}

// One atom's density, added to a grid point as gemmi's do_add_atom_density_to_grid() adds it.
struct AtomDensity {
    using CReal = Table::Coef::coef_type;
    using IsoSum = decltype(std::declval<const Table::Coef &>().precalculate_density_iso(CReal(), CReal()));
    using AnisoSum = decltype(std::declval<const Table::Coef &>().precalculate_density_aniso_b(
        gemmi::SMat33<CReal>(), CReal()));

    bool is_aniso = false;
    float occ = 0;
    IsoSum iso{};
    AnisoSum aniso{};

    // gemmi's ExpSum / ExpAnisoSum::calculate() for the points of a box, a Gaussian at a time: every
    // 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 {
        constexpr int BLOCK = 256;
        const int n = p.n;
        for (int k0 = 0; k0 < n; k0 += BLOCK) {
            const int m = std::min(BLOCK, n - k0);
            float density[BLOCK], x[BLOCK];
            for (int k = 0; k < m; ++k)
                density[k] = 0;
            if (!is_aniso) {
                float r2[BLOCK];
                for (int k = 0; k < m; ++k)
                    r2[k] = (CReal) p.d2[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);
                    for (int k = 0; k < m; ++k)
                        density[k] += iso.a[i] * gemmi::unsafe_expapprox(x[k]);
                }
            } 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]);
                        x[k] = std::max(aniso.b[i].r_u_r(r), -88.f);
                    }
                    for (int k = 0; k < m; ++k)
                        density[k] += aniso.a[i] * gemmi::unsafe_expapprox(x[k]);
                }
            }
            for (int k = 0; k < m; ++k)
                data[p.idx[k0 + k]] += float(occ * density[k]);
        }
    }
};

}  // namespace

// The points gemmi's Grid::symmetrize() (grid.hpp) reduces the grid's orbits into, ascending. gemmi
// walks the points by index and takes the orbit of every point it has not visited yet, so it reduces
// each orbit into the orbit's lowest index - the point none of whose mates has a lower one. They depend
// on the grid's size and space group only, so a zone finds them once for all its evaluations.
std::vector<size_t> OrbitLeaders(const gemmi::Grid<float> &grid, size_t nthreads) {
    const std::vector<gemmi::GridOp> ops = grid.get_scaled_ops_except_id();
    if (ops.empty())
        return {};   // P1: nothing to symmetrize
    std::vector<std::vector<size_t>> plane_leaders(grid.nw);
    ParallelFor(grid.nw, nthreads, [&](int w) {
        for (int v = 0; v < grid.nv; ++v)
            for (int u = 0; u < grid.nu; ++u) {
                const size_t idx = grid.index_q(u, v, w);
                bool lowest = true;
                for (size_t k = 0; k < ops.size() && lowest; ++k) {
                    const std::array<int, 3> t = ops[k].apply(u, v, w);
                    lowest = grid.index_n(t[0], t[1], t[2]) >= idx;
                }
                if (lowest)
                    plane_leaders[w].push_back(idx);
            }
    });
    std::vector<size_t> leaders;
    for (const std::vector<size_t> &l : plane_leaders)
        leaders.insert(leaders.end(), l.begin(), l.end());
    return leaders;
}

// 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.
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<AtomDensity> density(n);
    ParallelChunks(n, nthreads, [&](int lo, int hi) {
        for (int i = lo; i < hi; i++) {
            const gemmi::Atom &atom = *atoms[i];
            const auto &coef = Table::get(atom.element, atom.charge, atom.serial);
            const float addend = dc.addends.get(atom.element);
            const gemmi::Fractional fpos = dc.grid.unit_cell.fractionalize(atom.pos);
            AtomDensity &d = density[i];
            d.occ = atom.occ;
            d.is_aniso = atom.aniso.nonzero();
            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));
            } 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);
            }
        }
    });
    UseAtomBoxesByPlane(dc.grid, boxes, nthreads, [&](int i) { return density[i]; });
    SymmetrizeOrbits(dc.grid, orbit_leaders, nthreads, [](float a, float b) { return a + b; });
}

// SolventMasker(AtomicRadiiSet::Refmac).put_mask_on_grid() (gemmi solmask.hpp) with the atoms and the
// symmetry spread over the threads as above - the same mask, bit for bit. The island removal and the
// shrink that follow are gemmi's own, except that the shrink is skipped where it cannot change a point:
// it looks at the grid offsets within rshrink of each point (set_margin_around), and on a grid whose
// spacing is coarser than rshrink along all three axes there are none - which is most rigid-body zones -
// yet it still walks the whole grid.
void PutMaskOnGrid(gemmi::Grid<float> &grid, const gemmi::Model &model, const std::vector<size_t> &orbit_leaders,
                   size_t nthreads) {
    const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
    masker.clear(grid);
    std::vector<AtomBox> boxes;
    for (const gemmi::Atom *atom : ModelAtoms(model)) {
        if ((masker.ignore_hydrogen && atom->is_hydrogen()) ||
            (masker.ignore_zero_occupancy_atoms && atom->occ <= 0))
            continue;
        double r = masker.constant_r + masker.rprobe;
        r += gemmi::refmac_radius_for_bulk_solvent(atom->element.elem);
        boxes.push_back(MakeAtomBox(grid, grid.unit_cell.fractionalize(atom->pos), r));
    }
    UseAtomBoxesByPlane(grid, boxes, nthreads, [](int) {
        return [](float *data, const BoxPoints &p) {
            for (int k = 0; k < p.n; ++k)
                data[p.idx[k]] = 0.f;
        };
    });
    SymmetrizeOrbits(grid, orbit_leaders, nthreads, [](float a, float b) { return a < b ? a : b; });
    masker.remove_islands(grid);
    bool shrink_has_offsets = false;
    for (int i = 0; i < 3; i++)
        shrink_has_offsets = shrink_has_offsets || static_cast<int>(std::floor(masker.rshrink / grid.spacing[i])) > 0;
    if (shrink_has_offsets)
        masker.shrink(grid);
}