mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 81 Wiki Activity
Files
8430d697318cfaf35bbaede09b934b897f5aa6ef
Jungfraujoch/tests/ModelScaleGPUTest.cpp
T
leonarski_fandClaude Opus 5.5 1fe73bab41 Model validation: scale fits on the GPU, change-of-basis probe gridded once per frame on all threads 
The scale fit of every model-validation fit (fit_model: the real model, the indexing probe's relabellings
and every null replicate) was FitModelScale on the host - a long-double Levenberg-Marquardt per solvent grid
point, ~80% of the validation's CPU on a 400 A cell. It now runs on the GPU through ModelScaleGPU::FitSolvent,
the engine the rigid body and molecular replacement already use, from a small pool (ModelScaleGPUPool: up to
four engines with their own streams, sized on the card's total memory, at most an eighth of it). The grid,
the fits and the winner rule are FitModelScale's; the sums are taken in double on the device, so R-factors
agree with the host's to rounding. The CPU build, and a GPU build without a card, keep FitModelScale.

The change-of-basis probe gridded each candidate frame with gemmi's serial put_model_density_on_grid, once
for the data as merged and once more for each alternative indexing. It now grids each frame once, with
ModelGrid's parallel PutModelDensityOnGrid (the same grid, bit for bit), from threads outside the pool so the
gridding spreads over it, and scores every labelling on that F_calc. Exact.

4nwv (122k atoms, I222, 400 A cell): model validation 33.2 s -> 14.0 s; every logged R, sigma and decision
the same.

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

386 lines
17 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 <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"
#ifdef JFJOCH_USE_CUDA
#include <algorithm>
#include <cmath>
#include <complex>
#include <cstring>
#include <random>
#include <vector>
#include <filesystem>
#include <fstream>
#include <future>
#include "gemmi/asumask.hpp"
#include "gemmi/mmread_gz.hpp"
#include "gemmi/scaling.hpp"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"
#include "../common/Logger.h"
#include "../image_analysis/structure_refinement/ModelFFT.h"
#include "../image_analysis/structure_refinement/ModelGrid.h"
#include "../image_analysis/structure_refinement/ModelScaling.h"
#include "../image_analysis/structure_refinement/ModelScaleGPU.h"
#include "../image_analysis/structure_refinement/ModelValidation.h"
#include "../image_analysis/structure_refinement/RigidBodyRefine.h"
namespace {
constexpr double PI_ = 3.14159265358979323846;
// Scaling points in the reciprocal asymmetric unit of `sg` to d_min, sorted as prepare_points() leaves
// them. Fcalc has random phases and a Wilson fall-off, the mask term is strong at low resolution and
// roughly opposite in phase, and |Fobs| comes from a known overall scale, anisotropic B and solvent pair
// with 5% noise - so the fit has a right answer and a realistic shape of residual.
gemmi::Scaling<float> MakeScaling(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
gemmi::Scaling<float> scaling(cell, &sg);
scaling.use_solvent = true;
const gemmi::GroupOps gops = sg.operations();
const gemmi::ReciprocalAsu asu(&sg);
const gemmi::Miller lim = cell.get_hkl_limits(d_min);
std::vector<gemmi::Miller> hkl;
for (int h = -lim[0]; h <= lim[0]; h++)
for (int k = -lim[1]; k <= lim[1]; k++)
for (int l = -lim[2]; l <= lim[2]; l++) {
const gemmi::Miller m{{h, k, l}};
if ((h == 0 && k == 0 && l == 0) || cell.calculate_d(m) < d_min || !asu.is_in(m) ||
gops.is_systematically_absent(m))
continue;
hkl.push_back(m);
}
std::sort(hkl.begin(), hkl.end());
// B_cart of 25/35/30 A^2 with an off-diagonal term, as B*
gemmi::SMat33<double> b_cart{25, 35, 30, 4, -3, 2};
const gemmi::SMat33<double> b_star = b_cart.transformed_by(cell.frac.mat);
std::mt19937 rng(20260928);
std::uniform_real_distribution<double> uniform(0.0, 1.0);
std::normal_distribution<double> gauss(0.0, 1.0);
const double k_overall = 3.0, k_sol = 0.38, b_sol = 52.0;
for (const gemmi::Miller &m : hkl) {
const double stol2 = cell.calculate_stol_sq(m);
const double phase = 2 * PI_ * uniform(rng);
const double amp = 200.0 * std::exp(-10.0 * stol2) * std::sqrt(-std::log(1.0 - 0.999 * uniform(rng)));
const std::complex<double> fc = std::polar(amp, phase);
const std::complex<double> fm = std::polar(1500.0 * std::exp(-20.0 * stol2) * (0.5 + uniform(rng)),
phase + PI_ + 0.5 * gauss(rng));
const std::complex<float> fcf(fc), fmf(fm);
const std::complex<double> total = std::complex<double>(fcf) + k_sol * std::exp(-b_sol * stol2) * std::complex<double>(fmf);
const double fobs = k_overall * std::exp(-0.25 * b_star.r_u_r(m)) * std::abs(total) * (1.0 + 0.05 * gauss(rng));
gemmi::Scaling<float>::Point p{};
p.hkl = m;
p.stol2 = stol2;
p.fcmol = fcf;
p.fmask = fmf;
p.fobs = static_cast<float>(std::fabs(fobs));
p.sigma = static_cast<float>(0.05 * std::fabs(fobs) + 0.5);
scaling.points.push_back(p);
}
return scaling;
}
// A synthetic "protein" of 150 carbons, every third anisotropic, in `cryst` - ModelValidationTest.cpp's
// rigid-body fixture - and the Scaling points RigidBodyTarget::Residuals fits at a displaced
// placement q0: the model's own amplitudes as Fobs, and the Fcalc and bulk-solvent Fmask of the
// displaced model, gridded and composed as the rigid body does.
gemmi::Scaling<float> ModelPoints(const char *cryst, double zone) {
std::string pdb = cryst;
std::mt19937 rng(20260902);
std::uniform_real_distribution<double> x(2, 14), y(2, 16), z(2, 18);
char line[96];
for (int i = 1; i <= 150; i++) {
std::snprintf(line, sizeof line, "ATOM %5d C UNK A 1 %8.3f%8.3f%8.3f 1.00 20.00 C\n",
i, x(rng), y(rng), z(rng));
pdb += line;
}
pdb += "END\n";
const std::string path = "model_scale_gpu_test.pdb";
std::ofstream(path) << pdb;
Logger logger("ModelScaleGPUTest");
const auto reference = ModelReferenceIntensities(path, {}, {}, 3.0, logger);
gemmi::Structure st = gemmi::read_structure_gz(path, gemmi::CoorFormat::Detect);
std::filesystem::remove(path);
st.setup_cell_images();
int i = 0;
for (gemmi::Chain &ch : st.models[0].chains)
for (gemmi::Residue &r : ch.residues)
for (gemmi::Atom &a : r.atoms)
if (i++ % 3 == 0)
a.aniso = {0.30f, 0.25f, 0.20f, 0.02f, -0.01f, 0.03f};
const gemmi::SpaceGroup &sg = *st.find_spacegroup();
gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
fobs.unit_cell_ = st.cell;
fobs.spacegroup_ = &sg;
for (const auto &r : reference)
if (st.cell.calculate_d({{r.h, r.k, r.l}}) >= zone)
fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}});
fobs.ensure_sorted();
gemmi::Model model = st.models[0];
RigidBodyTarget target(model, st.cell, sg, 4);
const double q0[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05};
target.Place(q0, model);
gemmi::Grid<float> grid;
grid.unit_cell = st.cell;
grid.spacegroup = &sg;
grid.set_size_from_spacing(zone / 3.0, gemmi::GridSizeRounding::Up);
std::vector<gemmi::Miller> hkl;
for (const auto &hv : fobs.v)
hkl.push_back(hv.hkl);
const SymmetryComposition composition(grid, zone, hkl);
gemmi::DensityCalculator<gemmi::IT92<float>, float> dc;
dc.d_min = zone;
dc.rate = 1.5;
dc.grid.unit_cell = st.cell;
dc.grid.spacegroup = &sg;
dc.set_refmac_compatible_blur(model);
PutModelDensityOnGrid(dc, model, {}, 4);
std::vector<std::complex<double>> fc;
composition.Compose(MapToFPhi(dc.grid), dc.blur, fc, nullptr, 4);
gemmi::Grid<float> mask = grid;
PutMaskOnGrid(mask, model, OrbitLeaders(mask, 4), 4);
const gemmi::FPhiGrid<float> fm = MapToFPhi(mask);
gemmi::AsuData<std::complex<float>> fcalc, fmask;
for (size_t m = 0; m < composition.Hkl().size(); m++) {
fcalc.v.push_back({composition.Hkl()[m], std::complex<float>(fc[m])});
fmask.v.push_back({composition.Hkl()[m], fm.get_value_by_hkl(composition.Hkl()[m])});
}
gemmi::Scaling<float> scaling(st.cell, &sg);
scaling.use_solvent = true;
scaling.prepare_points(fcalc, fobs, &fmask);
return scaling;
}
struct GpuPoints {
CudaStream stream;
ModelScaleGPU scale;
CudaDevicePtr<float2> fcmol, fmask;
explicit GpuPoints(const gemmi::Scaling<float> &s)
: scale(stream, s.points.size()), fcmol(s.points.size()), fmask(s.points.size()) {
std::vector<std::array<int, 3>> hkl;
std::vector<double> stol2;
std::vector<float> fobs, sigma;
std::vector<float2> fc, fm;
for (const auto &p : s.points) {
hkl.push_back(p.hkl);
stol2.push_back(p.stol2);
fobs.push_back(p.fobs);
sigma.push_back(p.sigma);
fc.push_back(make_float2(p.fcmol.real(), p.fcmol.imag()));
fm.push_back(make_float2(p.fmask.real(), p.fmask.imag()));
}
std::vector<std::array<double, 6>> constraints(s.constraint_matrix.begin(), s.constraint_matrix.end());
double frac[9];
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++)
frac[3 * i + j] = s.cell.frac.mat[i][j];
scale.SetPoints(hkl, stol2, fobs, sigma, constraints, frac);
// On the stream the fits run on: it is non-blocking, so nothing orders it after the legacy stream.
REQUIRE(cudaMemcpyAsync(fcmol, fc.data(), fc.size() * sizeof(float2), cudaMemcpyHostToDevice, stream) ==
cudaSuccess);
REQUIRE(cudaMemcpyAsync(fmask, fm.data(), fm.size() * sizeof(float2), cudaMemcpyHostToDevice, stream) ==
cudaSuccess);
REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
}
};
double MaxAbs(const double b[6]) {
double m = 0;
for (int i = 0; i < 6; i++)
m = std::max(m, std::fabs(b[i]));
return m;
}
// The tolerances are those of gemmi's fit rather than of the arithmetic. Its Levenberg-Marquardt stops
// when the WSSR has changed by less than 1e-5 twice, which leaves the scale short of the minimum by
// about sqrt(1e-5) of the residual, so where a step is accepted or the fit stops on a near-tie the last
// bit decides it: measured, gemmi moves by up to 1e-4 of <Fobs> when its own Fcalc changes by 1e-5.
constexpr double R_TOLERANCE = 1e-5;
void CheckSameScale(const ModelScaleParams &gpu, const gemmi::Scaling<float> &cpu) {
const double b_cpu[6] = {cpu.b_star.u11, cpu.b_star.u22, cpu.b_star.u33,
cpu.b_star.u12, cpu.b_star.u13, cpu.b_star.u23};
CHECK(std::fabs(gpu.k_overall - cpu.k_overall) <= 1e-4 * std::fabs(cpu.k_overall));
// b_star = F B_cart F^T, so a B of 0.001 A^2 is 0.001 / a^2 in it for the longest edge a. Without that
// floor a fit whose B comes out near zero is held to 1e-4 of almost nothing (a cubic case at B 0.6 A^2
// differed by 1e-4 A^2).
const double edge = std::max({cpu.cell.a, cpu.cell.b, cpu.cell.c});
for (int i = 0; i < 6; i++)
CHECK(std::fabs(gpu.b_star[i] - b_cpu[i]) <= 1e-4 * MaxAbs(b_cpu) + 1e-3 / (edge * edge));
}
struct Case {
const char *name;
double a, b, c, alpha, beta, gamma;
const char *hm;
};
const Case CASES[] = {
{"triclinic", 41, 47, 53, 82, 97, 104, "P 1"},
{"monoclinic", 72, 44, 51, 90, 112, 90, "C 1 2 1"},
{"tetragonal", 64, 64, 81, 90, 90, 90, "P 4"},
{"hexagonal", 58, 58, 96, 90, 90, 120, "P 6"},
{"cubic", 92, 92, 92, 90, 90, 90, "P 2 3"},
};
} // namespace
TEST_CASE("ModelScaleGPU_FitMatchesGemmi", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
SKIP("No GPU");
for (const Case &c : CASES) {
INFO(c.name);
gemmi::UnitCell cell(c.a, c.b, c.c, c.alpha, c.beta, c.gamma);
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(c.hm);
gemmi::Scaling<float> cpu = MakeScaling(cell, sg, 2.8);
REQUIRE(cpu.points.size() > 2000);
GpuPoints gpu(cpu);
for (const double k_sol : {0.25, 0.4})
for (const double b_sol : {30.0, 60.0}) {
gemmi::Scaling<float> ref = cpu;
ref.k_sol = k_sol;
ref.b_sol = b_sol;
ref.fix_k_sol = true;
ref.fix_b_sol = true;
ref.fit_isotropic_b_approximately();
ref.fit_parameters();
const ModelScaleParams p = gpu.scale.Fit(gpu.fcmol, gpu.fmask, k_sol, b_sol);
CheckSameScale(p, ref);
gemmi::Scaling<float> with_gpu = ref;
with_gpu.k_overall = p.k_overall;
with_gpu.b_star = {p.b_star[0], p.b_star[1], p.b_star[2], p.b_star[3], p.b_star[4], p.b_star[5]};
CHECK(std::fabs(with_gpu.calculate_r_factor() - ref.calculate_r_factor()) < R_TOLERANCE);
}
}
}
TEST_CASE("ModelScaleGPU_SolventGridMatchesFitModelScale", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
SKIP("No GPU");
for (const Case &c : CASES) {
INFO(c.name);
gemmi::UnitCell cell(c.a, c.b, c.c, c.alpha, c.beta, c.gamma);
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(c.hm);
gemmi::Scaling<float> cpu = MakeScaling(cell, sg, 2.8);
GpuPoints gpu(cpu);
const ModelScaleReport report = FitModelScale(cpu, {}, 8);
const ModelSolventFit fit = gpu.scale.FitSolvent(gpu.fcmol, gpu.fmask);
CHECK(fit.n_grid == report.n_grid);
// The same grid point, so the same doubles - or, on a near-tie another card resolves the other
// way, a neighbour whose R is the CPU winner's (checked either way).
const bool same_point = fit.k_sol == cpu.k_sol && fit.b_sol == cpu.b_sol;
CHECK(std::fabs(fit.r - report.r_work_fit) < R_TOLERANCE);
if (same_point)
CheckSameScale(fit.scale, cpu);
}
}
// The rigid body's own sequence on a model's points: FitModelScale once for the zone, then the
// per-evaluation fit at that solvent pair.
TEST_CASE("ModelScaleGPU_MatchesGemmiOnAModelsPoints", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
SKIP("No GPU");
const char *crysts[] = {
"CRYST1 40.000 50.000 60.000 90.00 90.00 90.00 P 1 1\n",
"CRYST1 40.000 50.000 60.000 90.00 100.00 90.00 C 1 2 1 4\n",
"CRYST1 40.000 50.000 60.000 90.00 90.00 90.00 P 21 21 21 4\n",
"CRYST1 60.000 60.000 60.000 90.00 90.00 90.00 I 2 3 24\n",
"CRYST1 80.000 80.000 80.000 90.00 90.00 90.00 F 41 3 2 96\n",
};
for (const char *cryst : crysts)
for (const double zone : {6.0, 3.5}) {
INFO(cryst << " at " << zone << " A");
gemmi::Scaling<float> cpu = ModelPoints(cryst, zone);
REQUIRE(cpu.points.size() > 50);
GpuPoints gpu(cpu);
const ModelScaleReport report = FitModelScale(cpu, {}, 4);
const ModelSolventFit solvent = gpu.scale.FitSolvent(gpu.fcmol, gpu.fmask);
// As above: the same grid point, or a near-tie's neighbour at the CPU winner's R.
const bool same_point = solvent.k_sol == cpu.k_sol && solvent.b_sol == cpu.b_sol;
CHECK(std::fabs(solvent.r - report.r_work_fit) < R_TOLERANCE);
if (same_point)
CheckSameScale(solvent.scale, cpu);
gemmi::Scaling<float> ref = cpu;
ref.fix_k_sol = true;
ref.fix_b_sol = true;
ref.k_overall = 1;
ref.b_star = {0, 0, 0, 0, 0, 0};
ref.fit_isotropic_b_approximately();
ref.fit_parameters();
const ModelScaleParams p = gpu.scale.Fit(gpu.fcmol, gpu.fmask, cpu.k_sol, cpu.b_sol);
CheckSameScale(p, ref);
}
}
TEST_CASE("ModelScaleGPU_Deterministic", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
SKIP("No GPU");
gemmi::UnitCell cell(72, 44, 51, 90, 112, 90);
gemmi::Scaling<float> cpu = MakeScaling(cell, *gemmi::find_spacegroup_by_name("C 1 2 1"), 2.5);
GpuPoints gpu(cpu);
const ModelSolventFit first = gpu.scale.FitSolvent(gpu.fcmol, gpu.fmask);
const ModelScaleParams first_fit = gpu.scale.Fit(gpu.fcmol, gpu.fmask, first.k_sol, first.b_sol);
for (int repeat = 0; repeat < 3; repeat++) {
const ModelSolventFit again = gpu.scale.FitSolvent(gpu.fcmol, gpu.fmask);
CHECK(again.k_sol == first.k_sol);
CHECK(again.b_sol == first.b_sol);
CHECK(std::memcmp(&again.r, &first.r, sizeof(double)) == 0);
CHECK(std::memcmp(&again.scale, &first.scale, sizeof(ModelScaleParams)) == 0);
const ModelScaleParams fit = gpu.scale.Fit(gpu.fcmol, gpu.fmask, first.k_sol, first.b_sol);
CHECK(std::memcmp(&fit, &first_fit, sizeof(ModelScaleParams)) == 0);
}
}
// What model validation's fits take: FitModelScale's answer, to rounding, from a pool shared by more threads
// than it has engines - and the same answer, to the bit, whichever engine a fit got.
TEST_CASE("ModelScaleGPUPool_MatchesFitModelScale", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
SKIP("No GPU");
for (const Case &c : CASES) {
INFO(c.name);
gemmi::UnitCell cell(c.a, c.b, c.c, c.alpha, c.beta, c.gamma);
const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(c.hm);
const gemmi::Scaling<float> points = MakeScaling(cell, sg, 2.8);
gemmi::Scaling<float> cpu = points;
const ModelScaleReport report = FitModelScale(cpu, {}, 8);
ModelScaleGPUPool pool(0, points.points.size(), 2);
std::vector<gemmi::Scaling<float>> gpu(6, points);
std::vector<ModelScaleReport> gpu_report(gpu.size());
std::vector<std::future<void>> fits;
for (size_t i = 0; i < gpu.size(); i++)
fits.push_back(std::async(std::launch::async, [&, i] { gpu_report[i] = pool.Fit(gpu[i], 80.0, 2); }));
for (auto &f : fits)
f.get();
CHECK(gpu_report[0].n_grid == report.n_grid);
CHECK(std::fabs(gpu_report[0].r_work_fit - report.r_work_fit) < R_TOLERANCE);
CHECK(std::fabs(gpu[0].calculate_r_factor() - cpu.calculate_r_factor()) < R_TOLERANCE);
for (size_t i = 1; i < gpu.size(); i++) {
CHECK(gpu[i].k_sol == gpu[0].k_sol);
CHECK(gpu[i].b_sol == gpu[0].b_sol);
CHECK(gpu[i].k_overall == gpu[0].k_overall);
CHECK(std::memcmp(&gpu[i].b_star, &gpu[0].b_star, sizeof(gpu[0].b_star)) == 0);
}
}
}
#endif
Reference in New Issue View Git Blame Copy Permalink