leonarski_fandClaude Opus 5.5
0486575dec
Resolved: the ModelGrid.h include in ModelStructureFactorsGPU.cpp (AnisoRadiusB, SolventMaskGrid).
ModelReferenceIntensities keeps its mask on the density's grid: it is a reference at fixed solvent
constants, not a scored mask, and RigidBodyGPU_MatchesCPU and ModelScaleGPU_MatchesGemmiOnAModelsPoints
make their amplitudes from it - with the finer mask those amplitudes moved enough to land the anisotropic
test cluster on a near-tie of gemmi's scale fit (residuals 2.8e-4 against 2e-4, b_star 9.5e-7 against
1.6e-7). ModelStructureFactorsGPU_MatchesCPU at 1.5 A in F 41 3 2 (F_calc 1.28e-4 of the mean |F|
against 1e-4) fails on 89b37f9bc already and is left as it is.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37
527 lines
25 KiB
C++
527 lines
25 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,
|
|
bool ion_charges) {
|
|
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 = ion_charges ? IonFormFactor<float>(atom.element.elem, atom.charge)
|
|
: 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 = AnisoRadiusB(aniso_b);
|
|
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);
|
|
}
|
|
|
|
gemmi::Grid<float> SolventMaskGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
|
|
double spacing) {
|
|
spacing = std::min({spacing, MASK_SPACING_SHRINK, d_min / 3.0});
|
|
gemmi::Grid<float> grid;
|
|
grid.unit_cell = cell;
|
|
grid.spacegroup = &sg;
|
|
// set_size_from_spacing()'s size, without allocating the values
|
|
const std::array<int, 3> n = gemmi::good_grid_size({{cell.a / spacing, cell.b / spacing, cell.c / spacing}},
|
|
gemmi::GridSizeRounding::Up, &sg);
|
|
grid.nu = n[0];
|
|
grid.nv = n[1];
|
|
grid.nw = n[2];
|
|
grid.calculate_spacing();
|
|
grid.axis_order = gemmi::AxisOrder::XYZ;
|
|
return 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]);
|
|
}
|