mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 84 Wiki Activity
Files
4bf41983ed13bd71ed304029a7dbc837249b6bb2
Jungfraujoch/tests/ModelScaleGPUTest.cpp
T
leonarski_fandClaude Opus 5.5 4bf41983ed rugnux: the model is parsed by the caller and handed in as RugnuxModel 
The analysis read the --model file itself, four to five times per run
(the model kind, every validation, molecular replacement, the stills
reference intensities) and located the PAE sidecar from the file name.
The library is to do no file I/O, so the model now arrives in memory:

- RugnuxModel (image_analysis/structure_refinement/RugnuxModel.h): the
  gemmi structure as read, whether it is ModelCIF, the PAE matrix where a
  prediction came with one, and the name the report gives it.
- ParseRugnuxModel / ReadRugnuxModelFile (writer/RugnuxModelFile.h, in
  JFJochRugnuxFiles): bytes in, gzip by magic, PDB/mmCIF by content,
  _ma_qa_metric detection as before; the file reader takes the PAE JSON
  from the same place PaePathFor named.
- ProcessConfig::model replaces model_path; PrepareSearchModel,
  RunMolecularReplacement, ValidateAgainstModel and
  ModelReferenceIntensities take the model; result fields name it by
  model_name.
- rugnux reads the model before the run starts and stops with an error
  where the file holds no model (it used to run and report the model as
  not validated). The viewer reads it when the job is set up.
- The broker parses an uploaded model at once and answers 400 where it is
  not one; it keeps the parsed model in memory, with no temporary
  directory or file. It takes no PAE.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-17 16:11:58 +02:00

387 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"
#include "../writer/RugnuxModelFile.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(ReadRugnuxModelFile(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