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

ModelScaleGPU: the rigid body's scale fit on the GPU

Browse Source 
gemmi's Scaling<float> fit as RigidBodyTarget uses it - fit_isotropic_b_approximately() and the
Levenberg-Marquardt of fit_parameters() with k_sol and b_sol fixed - and FitModelScale's k_sol/b_sol
grid, with the sums over the reflections on the device (double, fixed launch shape, shuffle tree per
warp, warps and then blocks summed in order: no float atomics, bit-identical repeats). The LevMar
control is gemmi's, ported line for line to the host and unrolled into its requests, so the 88 coarse
and up to 25 fine grid fits share one launch per step.

Against gemmi on real zone hkl sets with synthetic amplitudes: k_overall and b* within 1e-9..1e-5
relative, the same grid winner every time, R within 1e-7.

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:56:34 +02:00
co-authored by Claude Opus 5.5
parent 9f00fd1f4d
commit e39489dfe0
5 changed files with 777 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/ModelScaleGPUTest.cpp
+210
View File
@@ -0,0 +1,210 @@
// 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 "gemmi/asumask.hpp"
#include "gemmi/scaling.hpp"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"
#include "../rugnux/ModelScaling.h"
#include "../rugnux/ModelScaleGPU.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;
}
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);
cudaMemcpy(fcmol, fc.data(), fc.size() * sizeof(float2), cudaMemcpyHostToDevice);
cudaMemcpy(fmask, fm.data(), fm.size() * sizeof(float2), cudaMemcpyHostToDevice);
}
};
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;
}
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));
for (int i = 0; i < 6; i++)
CHECK(std::fabs(gpu.b_star[i] - b_cpu[i]) <= 1e-4 * MaxAbs(b_cpu) + 1e-12);
}
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)
return;
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()) < 1e-6);
}
}
}
TEST_CASE("ModelScaleGPU_SolventGridMatchesFitModelScale", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
return;
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);
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);
CheckSameScale(fit.scale, cpu);
}
}
TEST_CASE("ModelScaleGPU_Deterministic", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
return;
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);
}
}
#endif