leonarski_f
b4b68d368f
RefineModel (Refinement.h) refines a placed model against the merge, reading and writing nothing: - placement stage (R1): every chain of 50+ atoms as a rigid body at once (L-BFGS on six parameters per chain, their gradient from the per-atom one through Ceres jets), then one B shift per residue; - restrained stage (R2): positions and isotropic B of every atom of a residue the embedded dictionary describes, the rest held fixed but kept in Fcalc and the mask; riding hydrogens added by gemmi's topology, moved with their parents, re-placed every macro-cycle. Macro-cycles refit the mask, the scale and bulk solvent (FitModelScale / ModelScaleGPUPool) and the likelihood's D and beta (ML on the free set in resolution bins), rebuild the repulsion pairs (symmetry mates included), then L-BFGS on positions and, where chosen, on B. - weights: a ladder of full refinements at multiples of the gradient-balanced X-ray weight, chosen by R-free within bond/angle bounds (ties to the tighter geometry), then the B refinement tried at its own weights and kept only where it beats B held by more than the tie. The target and its gradient (RefinementXray, RefinementGPU): the Rice/centric amplitude likelihood (RefinementMath.h, shared host/device code), the gradient by Agarwal's FFT method - the adjoint of the reciprocal-space symmetry composition into a half-l (CPU) or half-u (cuFFT) spectrum, one inverse real transform, then each atom's density derivatives (xyz and B) summed over its box. On the GPU everything per evaluation stays on the device: ModelDensityGPU's gather, r2c, a fused compose + likelihood kernel, a CSR gather for the adjoint (no atomics), c2r, one warp per atom for the gradient; the restraints are one kernel per restraint plus a per-atom gather of gradient slots; cub reductions. Deterministic. Restraints (RefinementRestraints): bonds, angles, torsions (with esd, no phi/psi, no alternative sugar puckers), chiral volumes, least-squares planes from gemmi's Topo; repulsions beyond 1-4 (vdW radii from the dictionary's energy library, 2.6 A for H-bond pairs, 1.9 A metal-N/O/S); B similarity of bonded and angle-related atoms. DensityBlobs (D1): gemmi's flood-fill blobs of mFo-DFc above 3 sigma peaking above 4 sigma, placed beside the model and labelled WATER / ION / LIGAND / MODEL_ERROR by distance, contacts, volume and the anomalous map. tools/rugnux_refine_test refines a model against a finished run's p.mtz (developer utility, not installed). Tests: likelihood derivative, the embedded library, restraint values on ideal residues and every restraint kind's gradient against finite differences (CPU, and the GPU against the CPU), the X-ray gradient against finite differences in P1, P212121, C2 and P6122 (CPU and GPU), and a synthetic refinement that recovers a shaken model. Credits in docs/ACKNOWLEDGEMENT.md and at the algorithms. Measured on the open arm's battery outputs against REFMAC 5.8, 10 cycles, same p.mtz, free set and start model (R-free): 5reo 0.2228 vs 0.2236, 8dyz 0.1733 vs 0.1651, 9chw 0.1995 vs 0.1969 (CPU engine). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 8efe6f2322b3bfdae580bb34f256e67730dfc3c2)
79 lines
4.0 KiB
C++
79 lines
4.0 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#pragma once
|
|
|
|
// The device half of restrained refinement on a GPU (RefinementGPU.h, CUDA builds only). The host half works
|
|
// out with gemmi everything that depends on the cell, the group and the reflections; this half grids the atoms,
|
|
// transforms, composes the symmetry, evaluates the likelihood target, takes its derivative back to every atom
|
|
// by the FFT method, and evaluates the geometry restraints. Every kernel is deterministic - gathers and fixed
|
|
// reduction trees, no floating-point atomics - so the same input gives the same bits on every run.
|
|
|
|
#include <array>
|
|
#include <cstddef>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
#include "ModelDensityGPU.h"
|
|
#include "ModelMaskGPU.h"
|
|
#include "RefinementMath.h"
|
|
|
|
// Everything an engine needs that does not depend on the atoms.
|
|
struct RefinementGPUSetup {
|
|
ModelDensityGrid grid; // the structure-factor grid, u fastest
|
|
double volume = 0;
|
|
std::vector<std::array<int, 12>> sym_ops; // gemmi's Op: rot row-major and tran, in units of 1 / den
|
|
int den = 24;
|
|
double centring = 1; // the number of centring vectors
|
|
std::vector<ModelMaskOp> mask_ops; // every operator with every centring vector
|
|
std::vector<std::array<int, 3>> rows; // the reflections, in the order their numbers come back
|
|
std::vector<double> inv_d2;
|
|
std::vector<float> fo;
|
|
std::vector<char> work, centric;
|
|
// The adjoint of the composition, as a gather: the half-u spectrum points the gradient's coefficients
|
|
// land on (ascending), and for each the terms (row * ops + op) * 2 + conj adding to it, in order.
|
|
std::vector<int> scatter_point, scatter_start, scatter_term;
|
|
size_t max_atoms = 0, max_pairs = 0;
|
|
};
|
|
|
|
struct RefinementGPUEngineImpl;
|
|
|
|
class RefinementGPUEngine {
|
|
public:
|
|
static size_t DeviceBytes(const RefinementGPUSetup &setup, size_t fft_work_bytes, size_t max_records,
|
|
size_t max_slots);
|
|
static size_t FFTWorkBytes(int device, const ModelDensityGrid &grid);
|
|
|
|
RefinementGPUEngine(int device, const RefinementGPUSetup &setup, size_t fft_work_bytes);
|
|
~RefinementGPUEngine();
|
|
|
|
// n_cen exp(blur s^2 / 4) per reflection: prepare_asu_data()'s unblur.
|
|
void SetBlur(double blur);
|
|
// Per macro-cycle: F_model = k_total (F + fmask), and the likelihood's D and variance.
|
|
void SetRowModel(const std::vector<float> &k_total, const std::vector<std::array<float, 2>> &fmask,
|
|
const std::vector<float> &d, const std::vector<float> &var);
|
|
|
|
// The atoms, as ModelStructureFactorsGPU takes them: densities, grid positions (fractional times the grid
|
|
// size, not wrapped) and the boxes of ModelDensityGPU::Boxes().
|
|
void SetAtoms(const std::vector<ModelDensityAtom> &atoms, const std::vector<std::array<float, 4>> &pos);
|
|
// F of the atoms at the rows.
|
|
void Fcalc(std::vector<std::array<float, 2>> &f);
|
|
// The bulk-solvent mask's F at the rows (the mask atoms fractional, wrapped).
|
|
void Fmask(const std::vector<ModelMaskAtom> &mask_atoms, std::vector<std::array<float, 2>> &f);
|
|
// -ln L over the working rows for the atoms last set and, with `grad`, dT/d(x, y, z, B) per atom.
|
|
double Evaluate(std::vector<std::array<float, 4>> *grad);
|
|
|
|
// Restraint set `which` (0 or 1), as RestraintSet holds it.
|
|
void SetRestraints(int which, const std::vector<RestraintRecord> &records, const std::vector<int> &plane_atoms,
|
|
const std::vector<int> &slot_start, const std::vector<int> &slot_index, int slots);
|
|
// The sum of E of set `which` at `xyzb` (x, y, z, B per atom) and dE/d(x, y, z, B) per atom.
|
|
double EvaluateRestraints(int which, const std::vector<std::array<float, 4>> &xyzb,
|
|
std::vector<std::array<float, 4>> *grad);
|
|
|
|
// The most atoms of a plane restraint the kernel takes.
|
|
static constexpr int MAX_PLANE_ATOMS = 32;
|
|
|
|
private:
|
|
std::unique_ptr<RefinementGPUEngineImpl> impl_;
|
|
};
|