mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 80 Wiki Activity

ModelScaleGPU: FitSolvent 3x faster; a test on a model's own points

Browse Source 
- |Fcalc + solvent| is taken once per fit instead of at every step: the solvent pair is fixed for the
  whole of a fit.
- The per-step anisotropic factor and the derivatives are computed per point in float, and summed in
  double as before. Double arithmetic was most of the cost on a card with little double throughput. The
  final R of the grid stays in gemmi's double arithmetic.
- Each mode reduces only the slots it fills; 48 blocks per fit instead of 128.
- The first upload of a batch no longer adds a stream synchronisation.

FitSolvent per zone at 3.5 A (RTX 5080, real zone hkl sets, synthetic amplitudes): 3.2 ms at 3.3k
points, 9.2 ms at 72k and 21.6 ms at 213k, against 9.7 / 29.6 / 62.9 ms before. Fit() is 0.25-0.66 ms.
Against gemmi, k_overall and b* agree to 1e-9 - 1.4e-5 relative, and the grid winner is the same.

New test ModelScaleGPU_MatchesGemmiOnAModelsPoints: the rigid body's own points (the ClusterPdb
fixture with anisotropic atoms, its own amplitudes, a displaced placement, the rigid body's Fcalc and
mask), five groups at 6 and 3.5 A. The R tolerance of the synthetic cases is now 1e-5. That is the
resolution of gemmi's own fit: its Levenberg-Marquardt stops at a WSSR change of 1e-5, and perturbing
gemmi's own Fcalc by 1e-5 moves its answer by up to 1.2e-4 of <Fobs>.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C
This commit is contained in:
leonarski_f
2026-09-28 19:28:56 +02:00
co-authored by Claude Opus 5.5
parent e39489dfe0
commit 274d42a2e7
3 changed files with 237 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/ModelScaleGPUTest.cpp
+131 -2
View File
@@ -13,13 +13,22 @@
#include <random>
#include <vector>
#include <filesystem>
#include <fstream>
#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 "../rugnux/ModelFFT.h"
#include "../rugnux/ModelGrid.h"
#include "../rugnux/ModelScaling.h"
#include "../rugnux/ModelScaleGPU.h"
#include "../rugnux/ModelValidation.h"
#include "../rugnux/RigidBodyRefine.h"
namespace {
@@ -76,6 +85,82 @@ gemmi::Scaling<float> MakeScaling(const gemmi::UnitCell &cell, const gemmi::Spac
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;
@@ -113,6 +198,12 @@ double MaxAbs(const double b[6]) {
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};
@@ -163,7 +254,7 @@ TEST_CASE("ModelScaleGPU_FitMatchesGemmi", "[ModelValidation][gpu]") {
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()) < 1e-6);
CHECK(std::fabs(with_gpu.calculate_r_factor() - ref.calculate_r_factor()) < R_TOLERANCE);
}
}
}
@@ -183,11 +274,49 @@ TEST_CASE("ModelScaleGPU_SolventGridMatchesFitModelScale", "[ModelValidation][gp
CHECK(fit.n_grid == report.n_grid);
CHECK(fit.k_sol == cpu.k_sol); // the same grid point, so the same double
CHECK(fit.b_sol == cpu.b_sol);
CHECK(std::fabs(fit.r - report.r_work_fit) < 1e-6);
CHECK(std::fabs(fit.r - report.r_work_fit) < R_TOLERANCE);
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)
return;
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);
CHECK(solvent.k_sol == cpu.k_sol);
CHECK(solvent.b_sol == cpu.b_sol);
CHECK(std::fabs(solvent.r - report.r_work_fit) < R_TOLERANCE);
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)
return;