mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 81 Wiki Activity
Files
b5980c9ba0e781122b6c41b2775d98fe03427a64
Jungfraujoch/image_analysis/structure_refinement/RigidBodyGPU.cpp
T
leonarski_fandClaude Opus 5.5 b5980c9ba0 rugnux: start the validation in the model's setting on a forecast, beside the first 
Where a model that fits is written on other axes than the data, the files
take its setting and the validation is made a second time on the
relabelled data. That second validation depends on the first only through
the setting and the indexing the first settles, and both are known long
before the first has finished: the change of basis right after the frame
scoring, the indexing the probe prefers before the null. So ValidateAgainstModel
now reports them (ModelFrameForecast, through ModelValidationSchedule::
on_forecast), and the second validation is started there, on the
relabelling AdoptModelFrame and relabel_output would make, applied to a
copy of the merge - beside the first one's null, real fit and maps.

It is kept only where the first decides exactly what was forecast (the
model fits, same change of basis, same indexing; a model asserting the
other enantiomorph is not forecast, as the label is decided last, on the
anomalous map). Until then its log is held (Logger::Buffered, replayed
where the serial run logged it) and its files wait on a gate
(ModelValidationSchedule::write_gate) placed before the first map is
written; otherwise it is released with false and returns unwritten, and
the serial validation runs as before.

GPU memory: two validations at once take twice the rigid-body engines.
The parallel start is decided up front from sizes, never from what is
free: allowed where twice what the first pool asked for (bytes per engine
times the engines wanted) fits a quarter of the card's TOTAL memory, the
share one validation may take. Threads: both validations submit to the
one ParallelFor pool from threads outside it (the second runs on a
std::async thread, as the null's replicates do), so no pass runs inline
on a pool worker and the pool's size bounds the workers.

Measured on the loaded 16-core workstation, TIMING model validation:
8sa8 30.6 -> 19.7 s, 8xtg 19.9 -> 13.9 s, 9ea5 22.1 -> 16.4 s (this and
the previous commit together). p.mtz, p.hkl, p.cif, the three maps,
p_maps.mtz and p_model.cif/pdb md5-identical to the base on 8sa8, 8xtg,
9ea5, 7qis and myob_x10sa; the validation's log lines identical as a set
(the null's replicate lines were already in completion order).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
2026-10-09 12:10:46 +02:00

445 lines
20 KiB
C++
Raw Blame History

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "RigidBodyGPU.h"
#include <algorithm>
#include <cmath>
#include <Eigen/Dense>
#include <ceres/rotation.h>
#include "gemmi/dencalc.hpp" // DensityCalculator
#include "gemmi/it92.hpp" // IT92 x-ray form factors
#include "gemmi/scaling.hpp" // Scaling
#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent
#include "ModelScaling.h" // FitModelScale
#include "RigidBodyGPUEngine.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#include "../../common/Logger.h"
namespace {
using Table = gemmi::IT92<float>;
// What a zone needs of the model, the cell and the group: the grid, each atom's density and mask
// radius, the images and the scale's constraints. The same for every fit of a validation.
RigidBodyGPUZone ModelZone(const gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
double d_min) {
RigidBodyGPUZone zone;
const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);
zone.nu = grid.nu;
zone.nv = grid.nv;
zone.nw = grid.nw;
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++) {
zone.orth[3 * i + j] = cell.orth.mat[i][j];
zone.frac[3 * i + j] = cell.frac.mat[i][j];
}
zone.volume = cell.volume;
// The density of each atom as PutModelDensityOnGrid() (ModelGrid.cpp) precalculates it.
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);
zone.blur = dc.blur;
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);
RigidBodyGPUAtom 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]);
}
}
zone.atoms.push_back(a);
// The bulk-solvent mask's atoms, as PutMaskOnGrid() takes them.
if (!((masker.ignore_hydrogen && atom.is_hydrogen()) ||
(masker.ignore_zero_occupancy_atoms && atom.occ <= 0))) {
zone.mask_atom.push_back(index);
zone.mask_radius.push_back(static_cast<float>(
masker.constant_r + masker.rprobe + gemmi::refmac_radius_for_bulk_solvent(atom.element.elem)));
}
++index;
}
const gemmi::GroupOps gops = sg.operations();
for (const gemmi::Op::Tran &cen : gops.cen_ops)
for (const gemmi::Op &op : gops.sym_ops) {
std::array<double, 12> image{};
for (int i = 0; i < 3; i++) {
for (int j = 0; j < 3; j++)
image[3 * i + j] = static_cast<double>(op.rot[i][j]) / gemmi::Op::DEN;
image[9 + i] = static_cast<double>(op.tran[i] + cen[i]) / gemmi::Op::DEN;
}
zone.images.push_back(image);
}
for (const gemmi::Vec6 &c : gemmi::adp_symmetry_constraints(&sg))
zone.constraints.push_back(c);
return zone;
}
// The zone's observations and the composition of Fcalc at them.
void AddObservations(RigidBodyGPUZone &zone, const gemmi::UnitCell &cell, const SymmetryComposition &composition,
const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs) {
zone.ops = composition.Ops();
const size_t rows = composition.Hkl().size();
zone.row_hkl.reserve(rows);
zone.row_scale.reserve(rows);
zone.row_stol2.reserve(rows);
for (size_t m = 0; m < rows; m++) {
const gemmi::Miller &h = composition.Hkl()[m];
zone.row_hkl.push_back({h[0], h[1], h[2]});
// prepare_asu_data()'s unblur, exp(B_blur |s|^2 / 4), as Compose() applies it
zone.row_scale.push_back(composition.Centring() * std::exp(zone.blur * 0.25 * composition.InvD2()[m]));
zone.row_stol2.push_back(cell.calculate_stol_sq(h));
}
zone.terms.reserve(composition.Terms().size());
for (const SymmetryComposition::Term &t : composition.Terms()) {
RigidBodyGPUTerm g{};
// The half-u transform holds h >= 0, and F1(-k) = conj F1(k) for a real map.
g.conj = t.k[0] < 0;
const int ku = g.conj ? -t.k[0] : t.k[0], kv = g.conj ? -t.k[1] : t.k[1], kw = g.conj ? -t.k[2] : t.k[2];
g.index = ku + (zone.nu / 2 + 1) * (gemmi::modulo(kv, zone.nv) + zone.nv * gemmi::modulo(kw, zone.nw));
g.phase[0] = t.phase.real();
g.phase[1] = t.phase.imag();
g.s[0] = t.s.x;
g.s[1] = t.s.y;
g.s[2] = t.s.z;
zone.terms.push_back(g);
}
double sum = 0;
for (size_t i = 0; i < fobs.v.size(); i++) {
const int m = composition.Row()[i];
const gemmi::ValueSigma<float> &vs = fobs.v[i].value;
zone.obs_row.push_back(m);
zone.obs_fobs.push_back(vs.value);
zone.obs_sigma.push_back(vs.sigma);
sum += vs.value;
// the points gemmi's prepare_points() takes
if (m >= 0 && !std::isnan(vs.value) && !std::isnan(vs.sigma)) {
zone.point_obs.push_back(static_cast<int>(i));
zone.point_hkl.push_back(zone.row_hkl[m]);
zone.point_stol2.push_back(zone.row_stol2[m]);
zone.point_fobs.push_back(vs.value);
zone.point_sigma.push_back(vs.sigma);
}
}
zone.f_mean = fobs.v.empty() ? 1.0 : sum / static_cast<double>(fobs.v.size());
}
bool SameObservations(const gemmi::AsuData<gemmi::ValueSigma<float>> &a,
const gemmi::AsuData<gemmi::ValueSigma<float>> &b) {
if (a.v.size() != b.v.size())
return false;
for (size_t i = 0; i < a.v.size(); i++)
if (a.v[i].hkl != b.v[i].hkl || !(a.v[i].value.value == b.v[i].value.value) ||
!(a.v[i].value.sigma == b.v[i].value.sigma))
return false;
return true;
}
} // namespace
std::unique_ptr<RigidBodyGPUPool> RigidBodyGPUPool::Create(const gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min,
size_t max_observations, size_t max_engines,
Logger &logger) {
if (get_gpu_count() == 0)
return nullptr;
try {
std::unique_ptr<RigidBodyGPUPool> pool(new RigidBodyGPUPool);
// Sized for the largest zone of the ladder, which every fit in the validation walks a part of.
RigidBodyGPUCapacity cap;
const size_t ops = sg.operations().sym_ops.size();
for (double zone_d : RigidBodyLadder(d_min)) {
const RigidBodyGPUZone &zone = pool->Zone(model, cell, sg, zone_d);
if (!RigidBodyGPUEngine::Supports(zone)) {
logger.Info("Model validation: the rigid body runs on the CPU - the cell is too small for the "
"GPU's gridding at {:.1f} A", zone_d);
return nullptr;
}
if (!RigidBodyGPUEngine::MaskSupports(zone)) {
logger.Info("Model validation: the rigid body runs on the CPU - the solvent mask grid at {:.1f} A "
"is finer than the GPU's mask, which has no shrink step, can take", zone_d);
return nullptr;
}
cap.atoms = std::max(cap.atoms, zone.atoms.size());
cap.grid_points = std::max(cap.grid_points, static_cast<size_t>(zone.nu) * zone.nv * zone.nw);
cap.complex_points = std::max(cap.complex_points, static_cast<size_t>(zone.nu / 2 + 1) * zone.nv * zone.nw);
cap.bricks = std::max(cap.bricks, RigidBodyGPUEngine::Bricks(zone.nu, zone.nv, zone.nw));
cap.pairs = std::max(cap.pairs, RigidBodyGPUEngine::PairBound(zone));
cap.fft_work_bytes = std::max(cap.fft_work_bytes, RigidBodyGPUEngine::FFTWorkBytes(zone.nu, zone.nv, zone.nw));
}
cap.observations = max_observations;
cap.rows = max_observations;
cap.terms = max_observations * ops;
const size_t bytes = RigidBodyGPUEngine::DeviceBytes(cap);
size_t free = 0, total = 0;
RigidBodyGPUEngine::MemoryInfo(free, total);
constexpr size_t HEADROOM = 1ull << 30;
const size_t budget = std::min(total / 4, free > HEADROOM ? free - HEADROOM : 0);
const size_t fit = bytes > 0 ? budget / bytes : 0;
const size_t want = std::min<size_t>(fit, std::max<size_t>(max_engines, 1));
const int device = RigidBodyGPUEngine::CurrentDevice();
pool->planned_bytes_ = bytes * std::max<size_t>(max_engines, 1);
pool->card_bytes_ = total;
for (size_t i = 0; i < want; i++)
pool->engines_.push_back(std::make_unique<RigidBodyGPUEngine>(cap, device));
if (pool->engines_.empty()) {
logger.Info("Model validation: the rigid body runs on the CPU - one GPU engine needs {:.0f} MB and the "
"budget is {:.0f} MB", bytes / 1e6, budget / 1e6);
return nullptr;
}
for (auto &e : pool->engines_)
pool->idle_.push_back(e.get());
logger.Info("Model validation: the rigid body runs on the GPU, {} engine(s) of {:.0f} MB", pool->Engines(),
bytes / 1e6);
return pool;
} catch (const JFJochException &e) {
cuda_clear_error();
logger.Warning("Model validation: the rigid body runs on the CPU - the GPU engines could not be set up ({})",
e.what());
return nullptr;
}
}
RigidBodyGPUPool::~RigidBodyGPUPool() = default;
std::shared_ptr<const RigidBodyGPUZone> RigidBodyGPUPool::ObservedZone(
const gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) {
const RigidBodyGPUZone &model_zone = Zone(model, cell, sg, d_min);
std::lock_guard lock(observed_m_);
for (const ObservedEntry &o : observed_)
if (o.d_min == d_min && SameObservations(o.fobs, fobs))
return o.zone;
const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);
std::vector<gemmi::Miller> hkl;
hkl.reserve(fobs.v.size());
for (const auto &hv : fobs.v)
hkl.push_back(hv.hkl);
const SymmetryComposition composition(grid, d_min, hkl);
auto zone = std::make_shared<RigidBodyGPUZone>(model_zone);
AddObservations(*zone, cell, composition, fobs);
observed_.push_back({d_min, fobs, zone});
return zone;
}
const RigidBodyGPUZone &RigidBodyGPUPool::Zone(const gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min) {
std::lock_guard lock(zones_m_);
auto it = zones_.find(d_min);
if (it == zones_.end())
it = zones_.emplace(d_min, std::make_unique<RigidBodyGPUZone>(ModelZone(model, cell, sg, d_min))).first;
return *it->second;
}
RigidBodyGPUEngine &RigidBodyGPUPool::Acquire() {
std::unique_lock lock(m_);
cv_.wait(lock, [this] { return !idle_.empty(); });
RigidBodyGPUEngine *e = idle_.back();
idle_.pop_back();
return *e;
}
void RigidBodyGPUPool::Release(RigidBodyGPUEngine &engine) {
{
std::lock_guard lock(m_);
idle_.push_back(&engine);
}
cv_.notify_one();
}
RigidBodyTargetGPU::RigidBodyTargetGPU(RigidBodyGPUPool &pool, gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, size_t nthreads)
: RigidBodyTargetBase(model), pool_(pool), lease_(pool), engine_(lease_.Engine()), model_(model), cell_(cell),
sg_(sg), nthreads_(nthreads) {
std::vector<std::array<double, 3>> relative;
for (const gemmi::Position &p : base_)
relative.push_back({p.x - centre_.x, p.y - centre_.y, p.z - centre_.z});
try {
engine_.SetBody(relative);
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
}
RigidBodyTargetGPU::~RigidBodyTargetGPU() = default;
// Place()'s placement as a matrix: the columns are the rotated axes, rotated as Place() rotates a point.
void RigidBodyTargetGPU::Placement(const double q[6], double rotation[9], double translation[3]) const {
const double aa[3] = {q[0] / rms_radius_, q[1] / rms_radius_, q[2] / rms_radius_};
for (int j = 0; j < 3; j++) {
const double e[3] = {j == 0 ? 1.0 : 0.0, j == 1 ? 1.0 : 0.0, j == 2 ? 1.0 : 0.0};
double column[3];
ceres::AngleAxisRotatePoint(aa, e, column);
for (int i = 0; i < 3; i++)
rotation[3 * i + j] = column[i];
}
translation[0] = centre_.x + q[3];
translation[1] = centre_.y + q[4];
translation[2] = centre_.z + q[5];
}
void RigidBodyTargetGPU::SetZone(const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) {
d_min_ = d_min;
solvent_fitted_ = false;
host_scale_ = false;
have_point_ = false;
zone_ = pool_.ObservedZone(model_, cell_, sg_, fobs, d_min);
try {
engine_.SetZone(*zone_);
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
}
// The scale fitted on the host, exactly as RigidBodyTarget::Residuals() fits it, for a zone the device
// does not reproduce: one with five or fewer reflections for fit_isotropic_b_approximately(), where
// FitModelScale()'s grid fits chain from one to the next and every fit starts from the last.
void RigidBodyTargetGPU::HostScale() {
std::vector<std::array<float, 2>> fc, fm;
engine_.DownloadPoints(fc, fm);
gemmi::Scaling<float> scaling(cell_, &sg_);
scaling.use_solvent = true;
for (size_t p = 0; p < zone_->point_obs.size(); p++) {
const std::array<int, 3> &h = zone_->point_hkl[p];
scaling.points.push_back({{h[0], h[1], h[2]}, zone_->point_stol2[p], std::complex<float>(fc[p][0], fc[p][1]),
std::complex<float>(fm[p][0], fm[p][1]), zone_->point_fobs[p], zone_->point_sigma[p]});
}
if (!solvent_fitted_) {
FitModelScale(scaling, {}, nthreads_);
k_sol = scaling.k_sol;
b_sol = scaling.b_sol;
solvent_fitted_ = true;
}
scaling.k_sol = k_sol;
scaling.b_sol = b_sol;
scaling.fix_k_sol = true;
scaling.fix_b_sol = true;
scaling.fit_isotropic_b_approximately();
scaling.fit_parameters();
k_overall_ = scaling.k_overall;
b_star_ = scaling.b_star;
}
// RigidBodyTarget::Residuals() with the density, the composition, the transforms and the residuals on
// the device.
bool RigidBodyTargetGPU::Residuals(const double q[6], double *residuals) {
try {
++evaluations;
double rotation[9], translation[3];
Placement(q, rotation, translation);
if (zone_->row_hkl.empty())
return false;
engine_.Fcalc(rotation, translation);
if (!(hold_mask && have_point_))
engine_.Fmask();
if (zone_->point_obs.empty())
return false;
// The scale as RigidBodyTarget::Residuals() fits it: the solvent once per zone, then the
// overall scale and the anisotropic B at every evaluation.
if (!solvent_fitted_ && !host_scale_) {
try {
engine_.FitSolvent(k_sol, b_sol);
solvent_fitted_ = true;
} catch (const ModelScaleGPUTooFewReflections &) {
host_scale_ = true;
}
}
if (host_scale_) {
HostScale();
} else {
double b[6];
engine_.FitScale(k_sol, b_sol, k_overall_, b);
b_star_ = {b[0], b[1], b[2], b[3], b[4], b[5]};
}
const double b_star[6] = {b_star_.u11, b_star_.u22, b_star_.u33, b_star_.u12, b_star_.u13, b_star_.u23};
engine_.Residuals(k_overall_, b_star, k_sol, b_sol, residuals);
unmatched = static_cast<int>(std::count(zone_->obs_row.begin(), zone_->obs_row.end(), -1));
have_point_ = true;
std::copy(q, q + 6, q_.begin());
return true;
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
}
// RigidBodyTarget::Jacobian(): the rotation columns by forward difference, three copies of the body
// gridded and transformed together, the translation columns exact, the scale re-fit projected out.
bool RigidBodyTargetGPU::Jacobian(const double q[6], double *jacobian) {
if (!have_point_ || !std::equal(q, q + 6, q_.begin())) {
std::vector<double> residuals(NumObservations());
if (!Residuals(q, residuals.data()))
return false;
}
try {
++jacobians;
const double step = RigidBodyJacobianStep(d_min_);
double rotation[3][9], translation[3][3];
for (int j = 0; j < 3; j++) {
double qj[6];
std::copy(q_.begin(), q_.end(), qj);
qj[j] += step;
Placement(qj, rotation[j], translation[j]);
}
const double b_star[6] = {b_star_.u11, b_star_.u22, b_star_.u33, b_star_.u12, b_star_.u13, b_star_.u23};
std::vector<double> jtj, jtq;
engine_.Jacobian(rotation, translation, step, k_overall_, b_star, k_sol, b_sol, jtj, jtq);
// Following Golub & Pereyra (1973) SIAM J. Numer. Anal. 10, 413-432, in the form of Kaufman (1975) BIT 15, 49-57
const int p = static_cast<int>(zone_->constraints.size()) + 1;
Eigen::MatrixXd a(p, p), c(p, 6);
for (int i = 0; i < p; i++) {
for (int j = 0; j < p; j++)
a(i, j) = jtj[i * p + j];
for (int j = 0; j < 6; j++)
c(i, j) = jtq[i * 6 + j];
}
const Eigen::MatrixXd x = a.ldlt().solve(c);
std::vector<double> xv(p * 6);
for (int i = 0; i < p; i++)
for (int j = 0; j < 6; j++)
xv[i * 6 + j] = x(i, j);
engine_.ProjectJacobian(xv, jacobian);
return true;
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
}
size_t RigidBodyTargetGPU::NumObservations() const {
return zone_ ? zone_->obs_row.size() : 0;
}
Reference in New Issue View Git Blame Copy Permalink