From e39489dfe0d8a239c5212a6b105a9083db652edf Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 16:56:34 +0200 Subject: [PATCH] ModelScaleGPU: the rigid body's scale fit on the GPU gemmi's Scaling 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) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- rugnux/CMakeLists.txt | 1 + rugnux/ModelScaleGPU.cu | 478 ++++++++++++++++++++++++++++++++++++ rugnux/ModelScaleGPU.h | 87 +++++++ tests/CMakeLists.txt | 1 + tests/ModelScaleGPUTest.cpp | 210 ++++++++++++++++ 5 files changed, 777 insertions(+) create mode 100644 rugnux/ModelScaleGPU.cu create mode 100644 rugnux/ModelScaleGPU.h create mode 100644 tests/ModelScaleGPUTest.cpp diff --git a/rugnux/CMakeLists.txt b/rugnux/CMakeLists.txt index ad785e24a..fcad8ae66 100644 --- a/rugnux/CMakeLists.txt +++ b/rugnux/CMakeLists.txt @@ -31,6 +31,7 @@ ADD_LIBRARY(JFJochRugnux STATIC HotPixels.h $<$:HotPixelsGPU.cu> HotPixelsGPU.h + $<$:ModelScaleGPU.cu> DiagnosticOutput.cpp DiagnosticOutput.h SpindleCuspLoss.cpp diff --git a/rugnux/ModelScaleGPU.cu b/rugnux/ModelScaleGPU.cu new file mode 100644 index 000000000..01b9985d6 --- /dev/null +++ b/rugnux/ModelScaleGPU.cu @@ -0,0 +1,478 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "ModelScaleGPU.h" + +#include +#include +#include +#include +#include + +#include "../common/JFJochException.h" + +namespace { + +void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); +} + +// A fixed launch shape, so every point is summed by the same thread in the same order whatever the +// card or the stream: the result is the same bit for bit run to run. +constexpr int THREADS = 128; +constexpr int BLOCKS = 128; +constexpr int SLOTS = 36; // 1 + 28 + 7: the WSSR, the lower half of J^T J and J^T r for 7 parameters +constexpr int MAX_FITS = 128; // FitModelScale's coarse pass is 88 fits + +// What a launch sums over the points for one fit. +constexpr int MODE_ISOTROPIC = 0; // fit_isotropic_b_approximately(): sx, sy, sxx, sxy, n +constexpr int MODE_WSSR = 1; // compute_wssr() +constexpr int MODE_MATRICES = 2; // compute_lm_matrices(): WSSR, J^T J, J^T r +constexpr int MODE_R = 3; // FitModelScale's R: sum |fobs - value|, sum fobs + +struct Constraints { + double row[6][6]; +}; + +// gemmi Scaling's arithmetic for one point: the solvent term and |F| in float, as a +// complex is, the anisotropic scale in double. +template +__global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict__ stol2, + const float *__restrict__ fobs, const unsigned char *__restrict__ strong, + const float2 *__restrict__ fcmol, const float2 *__restrict__ fmask, int n, + const ModelScaleFitState *__restrict__ fits, Constraints c, double *__restrict__ partial) { + constexpr int NS = 1 + NA * (NA + 1) / 2 + NA; // WSSR, J^T J lower half, J^T r + const ModelScaleFitState f = fits[blockIdx.y]; + double acc[NS]; + for (int s = 0; s < NS; s++) + acc[s] = 0; + + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) { + const float solvent = static_cast(f.k_sol * exp(-f.b_sol * stol2[i])); + const float2 fc = fcmol[i], fm = fmask[i]; + const float f_abs = hypotf(fc.x + solvent * fm.x, fc.y + solvent * fm.y); + const double fo = fobs[i]; + + if (f.mode == MODE_ISOTROPIC) { + if (!strong[i]) + continue; + const double x = stol2[i]; + const double y = logf(static_cast(fo / f_abs)); + acc[0] += x; + acc[1] += y; + acc[2] += x * x; + acc[3] += x * y; + acc[4] += 1; + continue; + } + + const double hx = hkl[3 * i], hy = hkl[3 * i + 1], hz = hkl[3 * i + 2]; + const double *u = f.b_star; + const double r_u_r = hx * hx * u[0] + hy * hy * u[1] + hz * hz * u[2] + + 2 * (hx * hy * u[3] + hx * hz * u[4] + hy * hz * u[5]); + const double k_aniso = exp(-0.25 * r_u_r); + + if (f.mode == MODE_WSSR || f.mode == MODE_R) { + // compute_value(): |F| times the scale cast to float, in float + const float value = f_abs * static_cast(f.k_overall * k_aniso); + const double dy = fo - value; + if (f.mode == MODE_WSSR) { + acc[0] += dy * dy; + } else { + acc[0] += fabs(dy); + acc[1] += fo; + } + continue; + } + + // compute_value_and_derivatives() with k_sol and b_sol fixed + const double fe = f_abs * k_aniso; + const double y = f.k_overall * fe; + double dy_da[NA]; + dy_da[0] = fe; + const double du[6] = {-0.25 * y * (hx * hx), -0.25 * y * (hy * hy), -0.25 * y * (hz * hz), + -0.5 * y * (hx * hy), -0.5 * y * (hx * hz), -0.5 * y * (hy * hz)}; + for (int j = 1; j < NA; j++) { + const double *r = c.row[j - 1]; + dy_da[j] = r[0] * du[0] + r[1] * du[1] + r[2] * du[2] + r[3] * du[3] + r[4] * du[4] + r[5] * du[5]; + } + const double dy = fo - y; + acc[0] += dy * dy; + int s = 1; + for (int j = 0; j < NA; j++) + for (int k = 0; k <= j; k++) + acc[s++] += dy_da[j] * dy_da[k]; + for (int j = 0; j < NA; j++) + acc[s++] += dy * dy_da[j]; + } + + // The block's sum of each slot: a fixed shuffle tree within each warp, then the warps in order. + __shared__ double warp_sum[THREADS / 32][NS]; + const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + for (int s = 0; s < NS; s++) { + double v = acc[s]; + for (int offset = 16; offset > 0; offset /= 2) + v += __shfl_down_sync(0xffffffffu, v, offset); + if (lane == 0) + warp_sum[warp][s] = v; + } + __syncthreads(); + if (threadIdx.x < SLOTS) { + double t = 0; + if (threadIdx.x < NS) + for (int w = 0; w < THREADS / 32; w++) + t += warp_sum[w][threadIdx.x]; + partial[(static_cast(blockIdx.y) * BLOCKS + blockIdx.x) * SLOTS + threadIdx.x] = t; + } +} + +// Each fit's block partials summed in block order. +__global__ void sum_blocks(const double *__restrict__ partial, double *__restrict__ sums) { + const int s = threadIdx.x; + double t = 0; + for (int b = 0; b < BLOCKS; b++) + t += partial[(static_cast(blockIdx.x) * BLOCKS + b) * SLOTS + s]; + sums[blockIdx.x * SLOTS + s] = t; +} + +// gemmi's jordan_solve (levmar.hpp): Gauss-Jordan elimination with partial pivoting, x returned in b. +void JordanSolve(double *a, double *b, int n) { + for (int i = 0; i < n; i++) { + int maxnr = -1; + double amax = 0; + for (int j = i; j < n; j++) { + const double aji = std::fabs(a[n * j + i]); + if (aji > amax) { + maxnr = j; + amax = aji; + } + } + if (maxnr == -1) { + for (int j = i; j < n; j++) + if (a[n * i + j] != 0. || b[i] != 0.) + throw std::runtime_error("Trying to reverse singular matrix. Column " + std::to_string(i) + + " is zeroed."); + continue; + } + if (maxnr != i) { + for (int j = i; j < n; j++) + std::swap(a[n * maxnr + j], a[n * i + j]); + std::swap(b[i], b[maxnr]); + } + const double c = 1.0 / a[i * n + i]; + for (int j = i; j < n; j++) + a[i * n + j] *= c; + b[i] *= c; + for (int k = 0; k < n; k++) + if (k != i) { + const double d = a[k * n + i]; + for (int j = i; j < n; j++) + a[k * n + j] -= a[i * n + j] * d; + b[k] -= b[i] * d; + } + } +} + +} // namespace + +// gemmi's LevMar::fit (levmar.hpp), unrolled into the requests it makes of the target, so that many +// fits can share one launch per step. Each call of compute_lm_matrices() or compute_wssr() there is one +// request here; everything between two of them is the same code in the same order. +struct ModelScaleGPU::LevMarRun { + // gemmi's defaults + const int eval_limit = 100; + const double lambda_limit = 1e+15; + const double stop_rel_change = 1e-5; + const double lambda_up_factor = 10; + const double lambda_down_factor = 0.1; + + double lambda = 0.001; + std::vector initial_a, best_a, trial; + std::vector alpha, beta; + double initial_wssr = NAN, wssr = NAN; + int eval_count = 0; + int small_change_counter = 0; + bool done = false; + + int mode = MODE_MATRICES; // the next request, at `at` + std::vector at; + + explicit LevMarRun(const std::vector &a) : initial_a(a), best_a(a), at(a) { + alpha.resize(a.size() * a.size()); + beta.resize(a.size()); + } + + void Take(const double *sums) { + const size_t na = beta.size(); + if (mode == MODE_MATRICES) { + int s = 1; + for (size_t j = 0; j < na; j++) + for (size_t k = 0; k <= j; k++) + alpha[na * j + k] = alpha[na * k + j] = sums[s++]; + for (size_t j = 0; j < na; j++) + beta[j] = sums[s++]; + if (eval_count == 0) { + initial_wssr = wssr = sums[0]; + eval_count = 1; + } else { + ++eval_count; + lambda *= lambda_down_factor; + } + Propose(); + return; + } + const double new_wssr = sums[0]; + ++eval_count; + if (new_wssr < wssr) { + const double rel_change = (wssr - new_wssr) / wssr; + wssr = new_wssr; + best_a = trial; + if (wssr == 0) { + done = true; + return; + } + if (rel_change < stop_rel_change) { + if (++small_change_counter >= 2) { + done = true; + return; + } + } else { + small_change_counter = 0; + } + mode = MODE_MATRICES; + at = best_a; + } else { + if (lambda > lambda_limit) { + done = true; + return; + } + lambda *= lambda_up_factor; + Propose(); + } + } + + // The damped normal equations J^T J + lambda diag(J^T J), solved for the next trial. + void Propose() { + if (eval_limit > 0 && eval_count >= eval_limit) { + done = true; + return; + } + const size_t na = beta.size(); + std::vector temp_alpha = alpha; + for (size_t j = 0; j < na; j++) + temp_alpha[na * j + j] *= (1.0 + lambda); + trial = beta; + JordanSolve(temp_alpha.data(), trial.data(), static_cast(na)); + for (size_t i = 0; i < na; i++) + trial[i] += best_a[i]; + mode = MODE_WSSR; + at = trial; + } + + const std::vector &Result() const { + return wssr < initial_wssr ? best_a : initial_a; + } +}; + +size_t ModelScaleGPU::DeviceBytes(size_t max_points) { + return max_points * (3 * sizeof(int) + sizeof(double) + sizeof(float) + sizeof(unsigned char)) + + MAX_FITS * (sizeof(ModelScaleFitState) + (BLOCKS + 1) * SLOTS * sizeof(double)); +} + +ModelScaleGPU::ModelScaleGPU(cudaStream_t stream, size_t max_points) + : stream_(stream), max_points_(max_points), + hkl_(3 * max_points), stol2_(max_points), fobs_(max_points), strong_(max_points), + fits_(MAX_FITS), partial_(static_cast(MAX_FITS) * BLOCKS * SLOTS), + sums_(static_cast(MAX_FITS) * SLOTS), + host_fits_(MAX_FITS), host_sums_(static_cast(MAX_FITS) * SLOTS) {} + +void ModelScaleGPU::SetPoints(const std::vector> &hkl, + const std::vector &stol2, + const std::vector &fobs, + const std::vector &sigma, + const std::vector> &constraints, + const double frac[9]) { + if (hkl.size() > max_points_ || stol2.size() != hkl.size() || fobs.size() != hkl.size() || + sigma.size() != hkl.size() || constraints.empty() || constraints.size() > 6) + throw std::invalid_argument("ModelScaleGPU::SetPoints: inconsistent points"); + n_ = static_cast(hkl.size()); + n_params_ = 1 + static_cast(constraints.size()); + for (size_t j = 0; j < constraints.size(); j++) + for (int k = 0; k < 6; k++) + constraints_[j][k] = constraints[j][k]; + std::copy(frac, frac + 9, frac_); + + // fit_isotropic_b_approximately()'s filter: it skips weak reflections + std::vector strong(n_); + n_strong_ = 0; + for (int i = 0; i < n_; i++) { + strong[i] = !(fobs[i] < 1 || fobs[i] < sigma[i]); + n_strong_ += strong[i]; + } + if (n_ == 0) + return; + cuda_err(cudaMemcpyAsync(hkl_, hkl.data(), n_ * 3 * sizeof(int), cudaMemcpyHostToDevice, stream_)); + cuda_err(cudaMemcpyAsync(stol2_, stol2.data(), n_ * sizeof(double), cudaMemcpyHostToDevice, stream_)); + cuda_err(cudaMemcpyAsync(fobs_, fobs.data(), n_ * sizeof(float), cudaMemcpyHostToDevice, stream_)); + cuda_err(cudaMemcpyAsync(strong_, strong.data(), n_, cudaMemcpyHostToDevice, stream_)); + cuda_err(cudaStreamSynchronize(stream_)); // the host vectors go out of scope +} + +const double *ModelScaleGPU::Sums(const float2 *d_fcmol, const float2 *d_fmask, + const std::vector &fits) { + const int nfits = static_cast(fits.size()); + if (nfits > MAX_FITS) + throw std::logic_error("ModelScaleGPU: too many fits in one launch"); + std::copy(fits.begin(), fits.end(), host_fits_.get()); + cuda_err(cudaMemcpyAsync(fits_, host_fits_, nfits * sizeof(ModelScaleFitState), cudaMemcpyHostToDevice, stream_)); + Constraints c{}; + std::copy(&constraints_[0][0], &constraints_[0][0] + 36, &c.row[0][0]); + const dim3 grid(BLOCKS, nfits); + switch (n_params_) { + case 2: scale_sums<2><<>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break; + case 3: scale_sums<3><<>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break; + case 4: scale_sums<4><<>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break; + case 5: scale_sums<5><<>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break; + case 6: scale_sums<6><<>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break; + case 7: scale_sums<7><<>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break; + default: throw std::logic_error("ModelScaleGPU: unexpected number of parameters"); + } + cuda_err(cudaGetLastError()); + sum_blocks<<>>(partial_, sums_); + cuda_err(cudaGetLastError()); + cuda_err(cudaMemcpyAsync(host_sums_, sums_, nfits * SLOTS * sizeof(double), cudaMemcpyDeviceToHost, stream_)); + cuda_err(cudaStreamSynchronize(stream_)); + return host_sums_; +} + +// fit_isotropic_b_approximately() then fit_parameters() for every fit in `fits`, which start from the +// k_overall and b_star they hold, and are left holding the answer. +void ModelScaleGPU::FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std::vector &fits) { + const size_t nfits = fits.size(); + for (auto &f : fits) + f.mode = MODE_ISOTROPIC; + const double *iso = Sums(d_fcmol, d_fmask, fits); + for (size_t i = 0; i < nfits; i++) { + const double *s = iso + i * SLOTS; + const double sx = s[0], sy = s[1], sxx = s[2], sxy = s[3], n = s[4]; + if (n <= 5) + continue; + const double slope = (n * sxy - sx * sy) / (n * sxx - sx * sx); + const double intercept = (sy - slope * sx) / n; + const double b_iso = -slope; + fits[i].k_overall = std::exp(intercept); + // set_b_overall({b_iso, b_iso, b_iso, 0, 0, 0}): b_star = F B F^T (SMat33::transformed_by) + const double *m = frac_; + auto elem = [&](int r, int q) { + return m[3 * r] * (m[3 * q] * b_iso) + m[3 * r + 1] * (m[3 * q + 1] * b_iso) + m[3 * r + 2] * (m[3 * q + 2] * b_iso); + }; + const double b[6] = {elem(0, 0), elem(1, 1), elem(2, 2), elem(0, 1), elem(0, 2), elem(1, 2)}; + std::copy(b, b + 6, fits[i].b_star); + } + + // get_parameters() and set_parameters() with k_sol and b_sol fixed + const int na = n_params_; + auto get_parameters = [&](const ModelScaleFitState &f) { + std::vector a(na); + a[0] = f.k_overall; + for (int j = 1; j < na; j++) { + const double *r = constraints_[j - 1]; + const double *u = f.b_star; + a[j] = r[0] * u[0] + r[1] * u[1] + r[2] * u[2] + r[3] * u[3] + r[4] * u[4] + r[5] * u[5]; + } + return a; + }; + auto set_parameters = [&](ModelScaleFitState &f, const std::vector &a) { + f.k_overall = a[0]; + std::fill(f.b_star, f.b_star + 6, 0.0); + for (int j = 1; j < na; j++) + for (int k = 0; k < 6; k++) + f.b_star[k] += constraints_[j - 1][k] * a[j]; + }; + + std::vector runs; + for (const auto &f : fits) + runs.emplace_back(get_parameters(f)); + while (true) { + std::vector live; + std::vector request; + for (size_t i = 0; i < nfits; i++) + if (!runs[i].done) { + ModelScaleFitState f = fits[i]; + set_parameters(f, runs[i].at); + f.mode = runs[i].mode; + live.push_back(i); + request.push_back(f); + } + if (live.empty()) + break; + const double *sums = Sums(d_fcmol, d_fmask, request); + for (size_t j = 0; j < live.size(); j++) + runs[live[j]].Take(sums + j * SLOTS); + } + for (size_t i = 0; i < nfits; i++) + set_parameters(fits[i], runs[i].Result()); +} + +ModelScaleParams ModelScaleGPU::Fit(const float2 *d_fcmol, const float2 *d_fmask, double k_sol, double b_sol) { + ModelScaleParams out; + if (n_ == 0) + return out; + std::vector fits(1, ModelScaleFitState{1.0, {0, 0, 0, 0, 0, 0}, k_sol, b_sol, 0}); + FitBatch(d_fcmol, d_fmask, fits); + out.k_overall = fits[0].k_overall; + std::copy(fits[0].b_star, fits[0].b_star + 6, out.b_star); + return out; +} + +// FitModelScale's grid (ModelScaling.cpp, default ModelScaleBox): a coarse pass over the box and a +// finer one around its winner, every grid point its own fit, the winner the lowest finite R in grid +// order. The fits of a pass run together, one launch per Levenberg-Marquardt step. +ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d_fmask) { + ModelSolventFit out; + if (n_ < 20) + return out; + if (n_strong_ <= 5) + throw std::logic_error("ModelScaleGPU::FitSolvent: too few reflections for independent grid points"); + + const double k_lo = 0.10, k_hi = 0.60, b_lo = 10.0, b_hi = 80.0; // ModelScaleBox{} + double best_r = -1; + auto try_points = [&](std::vector &fits) { + FitBatch(d_fcmol, d_fmask, fits); + for (auto &f : fits) + f.mode = MODE_R; + const double *r = Sums(d_fcmol, d_fmask, fits); + for (size_t i = 0; i < fits.size(); i++) { + ++out.n_grid; + const double den = r[i * SLOTS + 1]; + const double ri = den > 0 ? r[i * SLOTS] / den : 1.0; + if (std::isfinite(ri) && (best_r < 0 || ri < best_r)) { + best_r = ri; + out.k_sol = fits[i].k_sol; + out.b_sol = fits[i].b_sol; + out.scale.k_overall = fits[i].k_overall; + std::copy(fits[i].b_star, fits[i].b_star + 6, out.scale.b_star); + } + } + }; + auto point = [](double ks, double bs) { + return ModelScaleFitState{1.0, {0, 0, 0, 0, 0, 0}, ks, bs, 0}; + }; + + std::vector coarse; + for (double ks = k_lo; ks <= k_hi + 1e-9; ks += 0.05) + for (double bs = b_lo; bs <= b_hi + 1e-9; bs += 10.0) + coarse.push_back(point(ks, bs)); + try_points(coarse); + const double k0 = out.k_sol, b0 = out.b_sol; + const double k_hi2 = std::min(k_hi, k0 + 0.05); + const double b_hi2 = std::min(b_hi, b0 + 10.0); + std::vector fine; + for (double ks = std::max(k_lo, k0 - 0.05); ks <= k_hi2 + 1e-9; ks += 0.025) + for (double bs = std::max(b_lo, b0 - 10.0); bs <= b_hi2 + 1e-9; bs += 5.0) + fine.push_back(point(ks, bs)); + try_points(fine); + out.r = best_r; + return out; +} diff --git a/rugnux/ModelScaleGPU.h b/rugnux/ModelScaleGPU.h new file mode 100644 index 000000000..68edd44fc --- /dev/null +++ b/rugnux/ModelScaleGPU.h @@ -0,0 +1,87 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include + +#include + +#include "../image_analysis/indexing/CUDAMemHelpers.h" + +// The rigid body's scale fit on the GPU: what RigidBodyTarget::Residuals does with a gemmi +// Scaling (use_solvent, k_sol and b_sol fixed) - fit_isotropic_b_approximately() followed by +// gemmi's Levenberg-Marquardt, and FitModelScale's k_sol/b_sol grid. The sums over the reflections run +// on the device, in double, in a fixed order; the Levenberg-Marquardt control (the damping, the 7x7 +// solve, the stop rules) is gemmi's own, ported line for line and run on the host. + +struct ModelScaleParams { + double k_overall = 1.0; + double b_star[6] = {0, 0, 0, 0, 0, 0}; // gemmi SMat33 order u11 u22 u33 u12 u13 u23 +}; + +struct ModelSolventFit { + double k_sol = 0.35, b_sol = 46.0; // gemmi's Scaling defaults when the fit does not run + double r = 1.0; // ModelScaleReport::r_work_fit + int n_grid = 0; + ModelScaleParams scale; // the winner's k_overall and b_star +}; + +// One fit's parameters for one launch, and which sums to take over the points. +struct ModelScaleFitState { + double k_overall; + double b_star[6]; + double k_sol, b_sol; + int mode; +}; + +class ModelScaleGPU { +public: + static size_t DeviceBytes(size_t max_points); + ModelScaleGPU(cudaStream_t stream, size_t max_points); + + // Per zone: the Scaling points in gemmi prepare_points() order, adp_symmetry_constraints(sg) rows + // and the cell's fractionalization matrix (UnitCell::frac.mat, row-major). n <= max_points. + void SetPoints(const std::vector> &hkl, + const std::vector &stol2, + const std::vector &fobs, + const std::vector &sigma, + const std::vector> &constraints, + const double frac[9]); + + // fit_isotropic_b_approximately() + fit_parameters() at a fixed solvent pair, from k_overall = 1 and + // b_star = 0. d_fcmol, d_fmask: one float2 per point on the device. + ModelScaleParams Fit(const float2 *d_fcmol, const float2 *d_fmask, double k_sol, double b_sol); + + // FitModelScale (ModelScaling.cpp) with the default box: the same grid, the same winner. Throws + // std::logic_error where FitModelScale would chain the grid points (five or fewer reflections for the + // isotropic fit). + ModelSolventFit FitSolvent(const float2 *d_fcmol, const float2 *d_fmask); + +private: + struct LevMarRun; + + // The sums of every fit in `fits` over all points, one row of SLOTS doubles per fit, on the host. + const double *Sums(const float2 *d_fcmol, const float2 *d_fmask, const std::vector &fits); + void FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std::vector &fits); + + cudaStream_t stream_; + size_t max_points_; + int n_ = 0; + int n_strong_ = 0; // points fit_isotropic_b_approximately() fits on + int n_params_ = 1; // k_overall + one per constraint row + double constraints_[6][6] = {}; + double frac_[9] = {}; + + CudaDevicePtr hkl_; // 3 per point + CudaDevicePtr stol2_; + CudaDevicePtr fobs_; + CudaDevicePtr strong_; + CudaDevicePtr fits_; + CudaDevicePtr partial_; // per fit, per block, per slot + CudaDevicePtr sums_; // per fit, per slot + CudaHostPtr host_fits_; + CudaHostPtr host_sums_; +}; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7ce420525..9eae5711b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -142,6 +142,7 @@ ADD_EXECUTABLE(jfjoch_test MergeScaleTest.cpp AnisotropyAnalysisTest.cpp ModelScalingTest.cpp + ModelScaleGPUTest.cpp TwinningAnalysisTest.cpp TranslationalNCSTest.cpp RfreeFlagsTest.cpp diff --git a/tests/ModelScaleGPUTest.cpp b/tests/ModelScaleGPUTest.cpp new file mode 100644 index 000000000..4f2bb412a --- /dev/null +++ b/tests/ModelScaleGPUTest.cpp @@ -0,0 +1,210 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include +#include "../common/CUDAWrapper.h" + +#ifdef JFJOCH_USE_CUDA + +#include +#include +#include +#include +#include +#include + +#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 MakeScaling(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) { + gemmi::Scaling 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 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 b_cart{25, 35, 30, 4, -3, 2}; + const gemmi::SMat33 b_star = b_cart.transformed_by(cell.frac.mat); + std::mt19937 rng(20260928); + std::uniform_real_distribution uniform(0.0, 1.0); + std::normal_distribution 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 fc = std::polar(amp, phase); + const std::complex fm = std::polar(1500.0 * std::exp(-20.0 * stol2) * (0.5 + uniform(rng)), + phase + PI_ + 0.5 * gauss(rng)); + const std::complex fcf(fc), fmf(fm); + const std::complex total = std::complex(fcf) + k_sol * std::exp(-b_sol * stol2) * std::complex(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::Point p{}; + p.hkl = m; + p.stol2 = stol2; + p.fcmol = fcf; + p.fmask = fmf; + p.fobs = static_cast(std::fabs(fobs)); + p.sigma = static_cast(0.05 * std::fabs(fobs) + 0.5); + scaling.points.push_back(p); + } + return scaling; +} + +struct GpuPoints { + CudaStream stream; + ModelScaleGPU scale; + CudaDevicePtr fcmol, fmask; + + explicit GpuPoints(const gemmi::Scaling &s) + : scale(stream, s.points.size()), fcmol(s.points.size()), fmask(s.points.size()) { + std::vector> hkl; + std::vector stol2; + std::vector fobs, sigma; + std::vector 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> 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 &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 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 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 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 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 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