Build Packages / Create release (push) Successful in 9s
Build Packages / build:portable:macos-arm64:nocuda (push) Successful in 5m8s
Build Packages / build:rugnux:linux-aarch64:cuda (push) Successful in 11m12s
Build Packages / build:portable:linux-x86_64:nocuda (push) Successful in 16m43s
Build Packages / build:portable:linux-x86_64:cuda (push) Successful in 20m56s
Build Packages / build:portable:windows-x86_64:nocuda (push) Successful in 22m4s
Build Packages / build:jfjoch:rocky8:nocuda (push) Successful in 26m57s
Build Packages / HDF5 consumer tests (DIALS, XDS) (push) Successful in 28m24s
Build Packages / build:jfjoch:rocky9:nocuda (push) Successful in 27m58s
Build Packages / build:portable:windows-x86_64:cuda (push) Successful in 29m29s
Build Packages / build:jfjoch:ubuntu2204:nocuda (push) Successful in 18m23s
Build Packages / build:jfjoch:ubuntu2404:nocuda (push) Successful in 19m55s
Build Packages / Generate python client (push) Successful in 32s
Build Packages / Build documentation (push) Successful in 1m32s
Build Packages / build:jfjoch:rocky8:cuda-sls9 (push) Successful in 27m20s
Build Packages / build:jfjoch:ubuntu2204:cuda (push) Successful in 29m27s
Build Packages / build:jfjoch:ubuntu2404:cuda (push) Successful in 32m41s
Build Packages / build:jfjoch:rocky8:cuda (push) Successful in 34m46s
Build Packages / build:jfjoch:rocky9:cuda-sls9 (push) Successful in 37m25s
Build Packages / build:jfjoch:rocky9:cuda (push) Successful in 35m56s
Build Packages / Unit tests (push) Successful in 2h8m44s
Build Packages / Validation Battery (push) Successful in 33m35s
* 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>
508 lines
24 KiB
C++
508 lines
24 KiB
C++
// 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"
|
|
#ifdef __APPLE__
|
|
#include <atomic>
|
|
|
|
#include "../../common/JFJochException.h"
|
|
#include "../../common/Logger.h"
|
|
#include "ModelDensityMetal.h"
|
|
#endif
|
|
|
|
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);
|
|
}
|
|
});
|
|
}
|
|
|
|
// 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.
|
|
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 DensityPoints &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] = 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);
|
|
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(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)
|
|
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]);
|
|
}
|
|
}
|
|
};
|
|
|
|
#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
|
|
// 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 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, [[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<DensityBox> 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] = 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] = MakeDensityBox(dc.grid, fpos, radius);
|
|
}
|
|
}
|
|
});
|
|
#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; });
|
|
}
|
|
|
|
// 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);
|
|
}
|
|
|
|
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]);
|
|
}
|