leonarski_fandClaude Opus 5.5
278c488c37
Both sides kept: modelpar's parallel basis scoring, the second validation started on a forecast beside the first (write gate, schedule parameter), the multi-GPU placement and delete-before-rewrite; gpusf's GPU structure factors, maps and null engine, failure-instead-of-restart, and the merge engine released before the validation (now just before modelpar's ValidateAgainstModel call, after the forecast lambda is set up). Placement: each validation's structure-factor engines (d_min and the null's) are made on the card of the thread that runs it - the main thread's for the first validation, card 1 % count for the speculative second, which pins itself there - so two validations on two cards use both, as the rigid-body pools do. One memory rule for both, per card, from total memory, up front: a validation plans at most half of its card - its structure-factor engines a quarter together (was half for the d_min engine alone), its rigid-body engines a quarter (RigidBodyGPUPool) - and the second validation runs beside the first only where twice the first's plan (rigid-body planned bytes + the d_min engine, twice it where a null is coming, the null's engine being no larger) fits half of all cards' memory together (was: twice the rigid-body plan within a quarter). The card's total is read once when the engine is made; a CUDA error there fails the validation like any other. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
255 lines
12 KiB
C++
255 lines
12 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
// The reflection list follows ReciprocalGrid::prepare_asu_data() of GEMMI's recgrid.hpp
|
|
// (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/recgrid.hpp)
|
|
// (c) Global Phasing Ltd., Mozilla Public License Version 2.0
|
|
|
|
#include "ModelStructureFactorsGPU.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
|
|
#include "gemmi/fourier.hpp" // get_size_for_hkl, get_f_phi_on_grid
|
|
#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent
|
|
|
|
#include "RigidBodyRefine.h" // RigidBodyZoneGrid
|
|
|
|
namespace {
|
|
|
|
using Table = gemmi::IT92<float>;
|
|
|
|
// The density calculator model validation grids F_calc with: d_min, rate 1.5 and the blur gemmi's
|
|
// set_refmac_compatible_blur() gives this model.
|
|
gemmi::DensityCalculator<Table, float> Calculator(const gemmi::Model &model, const gemmi::UnitCell &cell,
|
|
const gemmi::SpaceGroup &sg, double d_min) {
|
|
gemmi::DensityCalculator<Table, float> dc;
|
|
dc.d_min = d_min;
|
|
dc.rate = 1.5;
|
|
dc.grid.unit_cell = cell;
|
|
dc.grid.spacegroup = &sg;
|
|
dc.set_refmac_compatible_blur(model);
|
|
return dc;
|
|
}
|
|
|
|
// The reflections prepare_asu_data(d_min) lists for the half-l transform of an (nu, nv, nw) grid, in its
|
|
// order (h, then k, then l ascending): in the group's reciprocal ASU, strictly inside d_min, inside the
|
|
// grid's Nyquist box, not systematically absent and not 000.
|
|
std::vector<gemmi::Miller> AsuReflections(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
|
|
int nu, int nv, int nw) {
|
|
const gemmi::Miller lim = cell.get_hkl_limits(d_min);
|
|
const int max_h = std::min((nu - 1) / 2, lim[0]);
|
|
const int max_k = std::min((nv - 1) / 2, lim[1]);
|
|
const int max_l = std::min(nw / 2, lim[2]);
|
|
const double max_1_d2 = 1.0 / (d_min * d_min);
|
|
const gemmi::ReciprocalAsu asu(&sg);
|
|
const gemmi::GroupOps gops = sg.operations();
|
|
std::vector<gemmi::Miller> rows;
|
|
gemmi::Miller hkl;
|
|
for (hkl[0] = -max_h; hkl[0] <= max_h; ++hkl[0])
|
|
for (hkl[1] = -max_k; hkl[1] <= max_k; ++hkl[1])
|
|
for (hkl[2] = -max_l; hkl[2] <= max_l; ++hkl[2])
|
|
if (asu.is_in(hkl) && cell.calculate_1_d2(hkl) < max_1_d2 && !gops.is_systematically_absent(hkl) &&
|
|
!(hkl[0] == 0 && hkl[1] == 0 && hkl[2] == 0))
|
|
rows.push_back(hkl);
|
|
return rows;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void ModelDensityAtoms(const gemmi::Model &model, const gemmi::DensityCalculator<Table, float> &dc,
|
|
std::vector<ModelDensityAtom> &atoms, std::vector<int> &mask_atom,
|
|
std::vector<float> &mask_radius) {
|
|
atoms.clear();
|
|
mask_atom.clear();
|
|
mask_radius.clear();
|
|
const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
|
|
int index = 0;
|
|
for (const gemmi::Chain &ch : model.chains)
|
|
for (const gemmi::Residue &r : ch.residues)
|
|
for (const gemmi::Atom &atom : r.atoms) {
|
|
using CReal = Table::Coef::coef_type;
|
|
const auto &coef = Table::get(atom.element, atom.charge, atom.serial);
|
|
const float addend = dc.addends.get(atom.element);
|
|
ModelDensityAtom a{};
|
|
a.occ = atom.occ;
|
|
a.aniso = atom.aniso.nonzero();
|
|
if (!a.aniso) {
|
|
const CReal b = static_cast<CReal>(atom.b_iso + dc.blur);
|
|
const auto precal = coef.precalculate_density_iso(b, addend);
|
|
a.radius = dc.estimate_radius(precal, b);
|
|
for (int k = 0; k < 5; k++) {
|
|
a.a[k] = precal.a[k];
|
|
a.b[k][0] = precal.b[k];
|
|
}
|
|
} 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);
|
|
a.radius = static_cast<float>(dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max));
|
|
const auto precal = coef.precalculate_density_aniso_b(aniso_b, addend);
|
|
for (int k = 0; k < 5; k++) {
|
|
a.a[k] = precal.a[k];
|
|
const gemmi::SMat33<float> &m = precal.b[k];
|
|
const float e[6] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23};
|
|
std::copy(e, e + 6, a.b[k]);
|
|
}
|
|
}
|
|
atoms.push_back(a);
|
|
if (!((masker.ignore_hydrogen && atom.is_hydrogen()) ||
|
|
(masker.ignore_zero_occupancy_atoms && atom.occ <= 0))) {
|
|
mask_atom.push_back(index);
|
|
mask_radius.push_back(static_cast<float>(
|
|
masker.constant_r + masker.rprobe + gemmi::refmac_radius_for_bulk_solvent(atom.element.elem)));
|
|
}
|
|
++index;
|
|
}
|
|
}
|
|
|
|
ModelStructureFactorsGPU::ModelStructureFactorsGPU(int device, const gemmi::Model &model, const gemmi::UnitCell &cell,
|
|
const gemmi::SpaceGroup &sg, double d_min)
|
|
: device_(device), device_total_(ModelStructureFactorsGPUEngine::TotalMemory(device)), cell_(cell), sg_(&sg), d_min_(d_min), dc_(Calculator(model, cell, sg, d_min)) {
|
|
const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min); // DensityCalculator's for d_min
|
|
setup_.grid.nu = grid.nu;
|
|
setup_.grid.nv = grid.nv;
|
|
setup_.grid.nw = grid.nw;
|
|
for (int i = 0; i < 3; i++)
|
|
for (int j = 0; j < 3; j++) {
|
|
setup_.grid.orth[3 * i + j] = cell.orth.mat[i][j];
|
|
setup_.grid.frac[3 * i + j] = cell.frac.mat[i][j];
|
|
}
|
|
setup_.volume = cell.volume;
|
|
|
|
const gemmi::GroupOps gops = sg.operations();
|
|
setup_.den = gemmi::Op::DEN;
|
|
for (const gemmi::Op &op : gops.sym_ops) {
|
|
std::array<int, 12> o{};
|
|
for (int i = 0; i < 3; i++) {
|
|
for (int j = 0; j < 3; j++)
|
|
o[3 * i + j] = op.rot[i][j];
|
|
o[9 + i] = op.tran[i];
|
|
}
|
|
setup_.sym_ops.push_back(o);
|
|
}
|
|
for (const gemmi::Op::Tran &cen : gops.cen_ops)
|
|
for (const gemmi::Op &op : gops.sym_ops) {
|
|
ModelMaskOp m{};
|
|
for (int i = 0; i < 3; i++) {
|
|
for (int j = 0; j < 3; j++)
|
|
m.rot[3 * i + j] = static_cast<double>(op.rot[i][j]) / gemmi::Op::DEN;
|
|
m.tran[i] = static_cast<double>(op.tran[i] + cen[i]) / gemmi::Op::DEN;
|
|
}
|
|
setup_.mask_ops.push_back(m);
|
|
}
|
|
|
|
rows_ = AsuReflections(cell, sg, d_min, grid.nu, grid.nv, grid.nw);
|
|
const double n_cen = static_cast<double>(gops.cen_ops.size());
|
|
for (const gemmi::Miller &h : rows_) {
|
|
setup_.rows.push_back({h[0], h[1], h[2]});
|
|
// prepare_asu_data()'s unblur, exp(B_blur |s|^2 / 4)
|
|
setup_.row_scale.push_back(n_cen * std::exp(dc_.blur * 0.25 * cell.calculate_1_d2(h)));
|
|
}
|
|
|
|
std::vector<ModelDensityAtom> atoms;
|
|
std::vector<int> mask_atom;
|
|
std::vector<float> mask_radius;
|
|
ModelDensityAtoms(model, dc_, atoms, mask_atom, mask_radius);
|
|
setup_.max_atoms = atoms.size();
|
|
setup_.max_pairs = ModelDensityGPU::PairBound(setup_.grid, atoms);
|
|
supported_ = !rows_.empty() && ModelDensityGPU::Supports(setup_.grid, atoms) &&
|
|
setup_.mask_ops.size() <= ModelMaskGPU::MAX_OPS;
|
|
if (!supported_)
|
|
return;
|
|
|
|
// The largest map: Map()'s coefficients are on a subset of these reflections, and get_size_for_hkl()
|
|
// only grows with the indices and the resolution it is given.
|
|
gemmi::AsuData<std::complex<float>> all;
|
|
all.unit_cell_ = cell;
|
|
all.spacegroup_ = &sg;
|
|
for (const gemmi::Miller &h : rows_)
|
|
all.v.push_back({h, {1.0f, 0.0f}});
|
|
map_size_ = gemmi::get_size_for_hkl(all, {{0, 0, 0}}, 3.0);
|
|
// MapFromFPhi() transforms a half-l grid of nw / 2 + 1 planes back to 2 (nw / 2) of them.
|
|
const int map_nw = 2 * (map_size_[2] / 2);
|
|
setup_.map_points = static_cast<size_t>(map_size_[0]) * map_size_[1] * map_nw;
|
|
setup_.map_complex_points = static_cast<size_t>(map_size_[0]) * map_size_[1] * (map_nw / 2 + 1);
|
|
fft_work_bytes_ = ModelStructureFactorsGPUEngine::FFTWorkBytes(device_, setup_.grid, map_size_[0], map_size_[1], map_nw);
|
|
device_bytes_ = ModelStructureFactorsGPUEngine::DeviceBytes(setup_, fft_work_bytes_);
|
|
}
|
|
|
|
ModelStructureFactorsGPU::~ModelStructureFactorsGPU() = default;
|
|
|
|
void ModelStructureFactorsGPU::Reserve() {
|
|
std::lock_guard lock(m_);
|
|
if (!engine_)
|
|
engine_ = std::make_unique<ModelStructureFactorsGPUEngine>(device_, setup_, fft_work_bytes_);
|
|
}
|
|
|
|
void ModelStructureFactorsGPU::Compute(const gemmi::Model &model, gemmi::AsuData<std::complex<float>> &fcalc,
|
|
gemmi::AsuData<std::complex<float>> &fmask) {
|
|
std::vector<ModelDensityAtom> atoms;
|
|
std::vector<int> mask_atom;
|
|
std::vector<float> mask_radius;
|
|
ModelDensityAtoms(model, dc_, atoms, mask_atom, mask_radius);
|
|
|
|
// Each atom in grid units and, for the mask, fractional - both wrapped into the cell.
|
|
const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
|
|
std::vector<std::array<float, 4>> pos;
|
|
std::vector<std::array<double, 3>> frac;
|
|
for (const gemmi::Chain &ch : model.chains)
|
|
for (const gemmi::Residue &r : ch.residues)
|
|
for (const gemmi::Atom &a : r.atoms) {
|
|
const gemmi::Fractional f0 = cell_.fractionalize(a.pos);
|
|
std::array<double, 3> f{f0.x, f0.y, f0.z};
|
|
std::array<float, 4> g{};
|
|
for (int k = 0; k < 3; k++) {
|
|
f[k] -= std::floor(f[k]);
|
|
if (f[k] >= 1.0) // a tiny negative coordinate wraps to 1 - epsilon, which rounds to 1
|
|
f[k] = 0.0;
|
|
g[k] = static_cast<float>(f[k] * n3[k]);
|
|
if (g[k] >= n3[k])
|
|
g[k] -= n3[k];
|
|
}
|
|
frac.push_back(f);
|
|
pos.push_back(g);
|
|
}
|
|
std::vector<ModelMaskAtom> mask_atoms;
|
|
for (size_t j = 0; j < mask_atom.size(); j++) {
|
|
const std::array<double, 3> &f = frac[mask_atom[j]];
|
|
mask_atoms.push_back({f[0], f[1], f[2], mask_radius[j]});
|
|
}
|
|
|
|
std::vector<std::array<float, 2>> fc, fm;
|
|
{
|
|
std::lock_guard lock(m_);
|
|
engine_->Compute(atoms, pos, mask_atoms, fc, fm);
|
|
}
|
|
fcalc.v.clear();
|
|
fmask.v.clear();
|
|
fcalc.v.reserve(rows_.size());
|
|
fmask.v.reserve(rows_.size());
|
|
for (size_t i = 0; i < rows_.size(); i++) {
|
|
fcalc.v.push_back({rows_[i], {fc[i][0], fc[i][1]}});
|
|
fmask.v.push_back({rows_[i], {fm[i][0], fm[i][1]}});
|
|
}
|
|
fcalc.unit_cell_ = fmask.unit_cell_ = cell_;
|
|
fcalc.spacegroup_ = fmask.spacegroup_ = sg_;
|
|
}
|
|
|
|
gemmi::Grid<float> ModelStructureFactorsGPU::Map(gemmi::AsuData<std::complex<float>> &coef) {
|
|
coef.ensure_sorted();
|
|
const std::array<int, 3> size = gemmi::get_size_for_hkl(coef, {{0, 0, 0}}, 3.0);
|
|
// ZYX: l fastest and halved, the layout a c2r transform takes, where gemmi's XYZ grid halves the slowest
|
|
// axis. The coefficients written are the same.
|
|
const gemmi::FPhiGrid<float> hkl = gemmi::get_f_phi_on_grid<float>(coef, size, true, gemmi::AxisOrder::ZYX);
|
|
gemmi::Grid<float> map;
|
|
map.spacegroup = coef.spacegroup_;
|
|
map.unit_cell = coef.unit_cell_;
|
|
// As MapFromFPhi(): 2 (nw - 1) planes back from nw stored ones.
|
|
map.set_size(size[0], size[1], 2 * (hkl.nu - 1));
|
|
map.axis_order = gemmi::AxisOrder::XYZ;
|
|
std::lock_guard lock(m_);
|
|
engine_->Map(map.nu, map.nv, map.nw, map.unit_cell.volume, reinterpret_cast<const float *>(hkl.data.data()),
|
|
map.data.data());
|
|
return map;
|
|
}
|