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

Rigid body: backend seam, GPU engine for density/composition/Jacobian, engine pool

Browse Source 
RigidBodyTarget becomes the CPU backend behind RigidBodyTargetBase, numerics
unchanged. RigidBodyTargetGPU evaluates the same target on a GPU engine: atoms
placed on the device, density by a deterministic brick gather (atoms in model
order per point), cuFFT R2C, the SymmetryComposition gather for Fcalc and
dF/dt in double, the numeric rotation columns as a batch-3 transform, the
Jacobian rows and the Kaufman sums by fixed-order reductions. The mask and the
scale fit are still computed on the host (stubs until ModelMaskGPU and
ModelScaleGPU land).

A pool of 1-4 engines is reserved once per validation, capped at min(25% of
the card, free - 1 GB); GPU or CPU is decided once for the whole validation,
and a CUDA failure restarts it on the CPU. JFJOCH_RIGID_BODY_CPU forces the CPU.

PutMaskOnGrid skips gemmi's shrink where its stencil is empty (every rigid-body
zone): exact, it cannot change a point.

GPU vs CPU on the five test groups: residuals within 7e-6 of <F>, Jacobian
columns within 4e-5 relative, fits 4e-7 A apart.

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 16:51:16 +02:00
co-authored by Claude Opus 5.5
parent 9f00fd1f4d
commit 15843eb49d
10 changed files with 1876 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/ModelValidationTest.cpp
+187 -2
View File
@@ -21,6 +21,10 @@
#include "../rugnux/ModelGrid.h"
#include "../rugnux/ModelValidation.h"
#include "../rugnux/RigidBodyRefine.h"
#ifdef JFJOCH_USE_CUDA
#include "../rugnux/RigidBodyGPU.h"
#include "../common/CUDAWrapper.h"
#endif
#include "../rugnux/SigmaA.h"
#include "../rugnux/WriteModel.h"
#include "../image_analysis/scale_merge/ReindexAmbiguity.h"
@@ -1027,8 +1031,34 @@ namespace {
// it. Per column (those in `columns`): the cosine between the two and the ratio of their norms.
// Also logged, not checked: how far the difference moves once the mask is let move with the body,
// which is what holding it costs.
// The model's own amplitudes to `d_min`: "observed" data whose minimum is where the model is.
gemmi::AsuData<gemmi::ValueSigma<float>> OwnAmplitudes(const char *cryst, const gemmi::Structure &st,
double d_min, Logger &logger) {
const auto path = WriteTemp("rigid_body_own_amplitudes.pdb", ClusterPdb(cryst).c_str());
const auto ref = ModelReferenceIntensities(path, {}, {}, 3.0, logger);
std::filesystem::remove(path);
gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
fobs.unit_cell_ = st.cell;
fobs.spacegroup_ = st.find_spacegroup();
for (const auto &r : ref)
if (st.cell.calculate_d({{r.h, r.k, r.l}}) >= d_min)
fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}});
fobs.ensure_sorted();
return fobs;
}
// The rigid body's target on the CPU, or on an engine of `gpu` where it is given.
std::unique_ptr<RigidBodyTargetBase> MakeTarget(gemmi::Model &model, const gemmi::Structure &st,
RigidBodyGPUPool *gpu) {
#ifdef JFJOCH_USE_CUDA
if (gpu != nullptr)
return std::make_unique<RigidBodyTargetGPU>(*gpu, model, st.cell, *st.find_spacegroup(), 4);
#endif
return std::make_unique<RigidBodyTarget>(model, st.cell, *st.find_spacegroup(), 4);
}
void CheckRigidBodyJacobian(const char *cryst, const double q0[6], const std::vector<int> &columns,
double min_cosine, double max_norm_error) {
double min_cosine, double max_norm_error, bool gpu = false) {
Logger logger("CheckRigidBodyJacobian");
const gemmi::Structure st = AnisoCluster(cryst);
const gemmi::SpaceGroup *sg = st.find_spacegroup();
@@ -1048,7 +1078,15 @@ namespace {
fobs.ensure_sorted();
gemmi::Model model = st.models[0];
RigidBodyTarget target(model, st.cell, *sg, 4);
std::unique_ptr<RigidBodyGPUPool> pool;
#ifdef JFJOCH_USE_CUDA
if (gpu) {
pool = RigidBodyGPUPool::Create(model, st.cell, *sg, zone, fobs.v.size(), 1, logger);
REQUIRE(pool);
}
#endif
const std::unique_ptr<RigidBodyTargetBase> target_backend = MakeTarget(model, st, pool.get());
RigidBodyTargetBase &target = *target_backend;
target.SetZone(fobs, zone);
const size_t n = target.NumObservations();
std::vector<double> r(n), jacobian(n * 6);
@@ -1190,3 +1228,150 @@ TEST_CASE("ModelValidation_RigidBodyReportsTheRotationItApplied", "[ModelValidat
CHECK(std::acos(std::clamp((trace - 1) / 2, -1.0, 1.0)) * 180 / PI ==
Catch::Approx(result.angle_deg).margin(1e-6));
}
#ifdef JFJOCH_USE_CUDA
// The GPU target is the CPU target's function, computed on the device: at the same placement the two
// give the same residuals and the same Jacobian to rounding (float distances on the device, cuFFT for
// FFTW), on every group of the composition tests, both zones, anisotropic atoms included.
TEST_CASE("RigidBodyGPU_MatchesCPU", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
SUCCEED("no GPU");
return;
}
Logger logger("RigidBodyGPU_MatchesCPU");
const double q0[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05};
for (const char *cryst : kRigidBodyCrysts) {
const gemmi::Structure st = AnisoCluster(cryst);
const gemmi::SpaceGroup *sg = st.find_spacegroup();
for (double zone : {6.0, 3.5}) {
const auto fobs = OwnAmplitudes(cryst, st, zone, logger);
gemmi::Model cpu_model = st.models[0], gpu_model = st.models[0];
auto pool = RigidBodyGPUPool::Create(gpu_model, st.cell, *sg, zone, fobs.v.size(), 1, logger);
REQUIRE(pool);
RigidBodyTarget cpu(cpu_model, st.cell, *sg, 4);
RigidBodyTargetGPU gpu(*pool, gpu_model, st.cell, *sg, 4);
cpu.SetZone(fobs, zone);
gpu.SetZone(fobs, zone);
const size_t n = cpu.NumObservations();
REQUIRE(gpu.NumObservations() == n);
std::vector<double> rc(n), rg(n), jc(6 * n), jg(6 * n);
REQUIRE(cpu.Residuals(q0, rc.data()));
REQUIRE(gpu.Residuals(q0, rg.data()));
REQUIRE(cpu.Jacobian(q0, jc.data()));
REQUIRE(gpu.Jacobian(q0, jg.data()));
CHECK(gpu.unmatched == cpu.unmatched);
CHECK(gpu.k_sol == cpu.k_sol);
CHECK(gpu.b_sol == cpu.b_sol);
double worst = 0, rms = 0;
for (size_t i = 0; i < n; i++) {
worst = std::max(worst, std::fabs(rc[i] - rg[i]));
rms += rc[i] * rc[i] / static_cast<double>(n);
}
INFO(cryst << " at " << zone << " A: residuals differ by at most " << worst << ", rms residual "
<< std::sqrt(rms));
CHECK(worst <= 1e-4);
for (int j = 0; j < 6; j++) {
double diff = 0, norm = 0;
for (size_t i = 0; i < n; i++) {
diff += gemmi::sq(jc[6 * i + j] - jg[6 * i + j]);
norm += gemmi::sq(jc[6 * i + j]);
}
INFO(cryst << " at " << zone << " A, column " << j << ": relative difference "
<< std::sqrt(diff / norm));
CHECK(std::sqrt(diff) <= 2e-3 * std::sqrt(norm) + 1e-9);
}
}
}
}
// The whole Jacobian on the GPU against the central difference of the GPU's own residuals, with the
// bounds of ModelValidation_RigidBodyJacobianMatchesNumericScaleRefit.
TEST_CASE("RigidBodyGPU_JacobianMatchesNumericScaleRefit", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
SUCCEED("no GPU");
return;
}
const double close[6] = {0.0, 0.0, 0.0, 0.05, -0.04, 0.03};
const double rotated[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05};
for (const char *cryst : kRigidBodyCrysts) {
CheckRigidBodyJacobian(cryst, close, {0, 1, 2, 3, 4, 5}, 0.99, 0.02, true);
CheckRigidBodyJacobian(cryst, rotated, {0, 1, 2}, 0.99, 0.03, true);
}
}
namespace {
// A whole placement on the GPU from the model displaced by 0.54 A and turned 2 deg, against the
// model's own amplitudes; the model is left where the fit put it.
RigidBodyRefineResult DisplacedFit(const char *cryst, gemmi::Structure &st, RigidBodyGPUPool *pool,
Logger &logger) {
const gemmi::SpaceGroup *sg = st.find_spacegroup();
const auto fobs = OwnAmplitudes(cryst, st, 3.0, logger);
std::vector<gemmi::Position> moved = ModelPositions(st.models[0]);
gemmi::Vec3 centre;
for (const gemmi::Position &p : moved)
centre += p;
centre *= 1.0 / static_cast<double>(moved.size());
const double a = 2.0 * PI / 180.0;
const gemmi::Mat33 rz(std::cos(a), -std::sin(a), 0, std::sin(a), std::cos(a), 0, 0, 0, 1);
for (gemmi::Position &p : moved)
p = gemmi::Position(rz.multiply(gemmi::Vec3(p) - centre) + centre + gemmi::Vec3(0.40, -0.30, 0.20));
SetModelPositions(st.models[0], moved);
return RefineRigidBody(st.models[0], st.cell, *sg, fobs, 3.0, logger, 4, pool);
}
}
// The GPU fit walks where the CPU fit walks: from the same displaced start, to the same placement
// within a few thousandths of an angstrom.
TEST_CASE("RigidBodyGPU_FitAgreesWithCPU", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
SUCCEED("no GPU");
return;
}
Logger logger("RigidBodyGPU_FitAgreesWithCPU");
for (const char *cryst : {kCryst, kPolarCryst, kRigidBodyCrysts[3]}) {
gemmi::Structure cpu_st = AnisoCluster(cryst), gpu_st = AnisoCluster(cryst);
auto pool = RigidBodyGPUPool::Create(gpu_st.models[0], gpu_st.cell, *gpu_st.find_spacegroup(), 3.0, 100000, 1,
logger);
REQUIRE(pool);
const RigidBodyRefineResult cpu = DisplacedFit(cryst, cpu_st, nullptr, logger);
const RigidBodyRefineResult gpu = DisplacedFit(cryst, gpu_st, pool.get(), logger);
CHECK(gpu.converged == cpu.converged);
const std::vector<gemmi::Position> pc = ModelPositions(cpu_st.models[0]), pg = ModelPositions(gpu_st.models[0]);
double rmsd = 0;
for (size_t i = 0; i < pc.size(); i++)
rmsd += pc[i].dist_sq(pg[i]) / static_cast<double>(pc.size());
INFO(cryst << ": CPU " << cpu.angle_deg << " deg " << cpu.shift_A << " A, GPU " << gpu.angle_deg << " deg "
<< gpu.shift_A << " A, " << std::sqrt(rmsd) << " A apart");
CHECK(std::sqrt(rmsd) < 2e-3);
}
}
// Deterministic: the same fit twice, and on a pool of one engine and of four, gives the same placement
// bit for bit.
TEST_CASE("RigidBodyGPU_Deterministic", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
SUCCEED("no GPU");
return;
}
Logger logger("RigidBodyGPU_Deterministic");
const char *cryst = kRigidBodyCrysts[4];
std::vector<std::vector<gemmi::Position>> placed;
for (size_t engines : {1, 1, 4}) {
gemmi::Structure st = AnisoCluster(cryst);
auto pool = RigidBodyGPUPool::Create(st.models[0], st.cell, *st.find_spacegroup(), 3.0, 100000, engines, logger);
REQUIRE(pool);
const RigidBodyRefineResult r = DisplacedFit(cryst, st, pool.get(), logger);
CHECK(r.evaluations > 0);
placed.push_back(ModelPositions(st.models[0]));
}
for (size_t k = 1; k < placed.size(); k++)
for (size_t i = 0; i < placed[0].size(); i++) {
CHECK(placed[k][i].x == placed[0][i].x);
CHECK(placed[k][i].y == placed[0][i].y);
CHECK(placed[k][i].z == placed[0][i].z);
}
}
#endif