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