Merge branch 'rb-gpu-scale' into rb-gpu-core

rugnux/CMakeLists.txt

@@ -35,6 +35,7 @@ ADD_LIBRARY(JFJochRugnux STATIC
         HotPixels.h
         $<$<BOOL:${JFJOCH_CUDA_AVAILABLE}>:HotPixelsGPU.cu>
         HotPixelsGPU.h
+        $<$<BOOL:${JFJOCH_CUDA_AVAILABLE}>:ModelScaleGPU.cu>
         DiagnosticOutput.cpp
         DiagnosticOutput.h
         SpindleCuspLoss.cpp

rugnux/ModelScaleGPU.cu

@@ -0,0 +1,478 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#include "ModelScaleGPU.h"
+
+#include <algorithm>
+#include <cmath>
+#include <stdexcept>
+#include <string>
+#include <utility>
+
+#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<float>'s arithmetic for one point: the solvent term and |F| in float, as a
+// complex<float> is, the anisotropic scale in double.
+template <int NA>
+__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<float>(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<float>(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<float>(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<size_t>(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<size_t>(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<double> initial_a, best_a, trial;
+    std::vector<double> 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<double> at;
+
+    explicit LevMarRun(const std::vector<double> &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<double> 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<int>(na));
+        for (size_t i = 0; i < na; i++)
+            trial[i] += best_a[i];
+        mode = MODE_WSSR;
+        at = trial;
+    }
+
+    const std::vector<double> &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<size_t>(MAX_FITS) * BLOCKS * SLOTS),
+      sums_(static_cast<size_t>(MAX_FITS) * SLOTS),
+      host_fits_(MAX_FITS), host_sums_(static_cast<size_t>(MAX_FITS) * SLOTS) {}
+
+void ModelScaleGPU::SetPoints(const std::vector<std::array<int, 3>> &hkl,
+                              const std::vector<double> &stol2,
+                              const std::vector<float> &fobs,
+                              const std::vector<float> &sigma,
+                              const std::vector<std::array<double, 6>> &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<int>(hkl.size());
+    n_params_ = 1 + static_cast<int>(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<unsigned char> 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<ModelScaleFitState> &fits) {
+    const int nfits = static_cast<int>(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><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
+        case 3: scale_sums<3><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
+        case 4: scale_sums<4><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
+        case 5: scale_sums<5><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
+        case 6: scale_sums<6><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, d_fcmol, d_fmask, n_, fits_, c, partial_); break;
+        case 7: scale_sums<7><<<grid, THREADS, 0, stream_>>>(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<<<nfits, SLOTS, 0, stream_>>>(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<ModelScaleFitState> &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<double> 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<double> &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<LevMarRun> runs;
+    for (const auto &f : fits)
+        runs.emplace_back(get_parameters(f));
+    while (true) {
+        std::vector<size_t> live;
+        std::vector<ModelScaleFitState> 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<ModelScaleFitState> 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<ModelScaleFitState> &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<ModelScaleFitState> 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<ModelScaleFitState> 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;
+}

rugnux/ModelScaleGPU.h

@@ -0,0 +1,87 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#pragma once
+
+#include <array>
+#include <cstddef>
+#include <vector>
+
+#include <cuda_runtime.h>
+
+#include "../image_analysis/indexing/CUDAMemHelpers.h"
+
+// The rigid body's scale fit on the GPU: what RigidBodyTarget::Residuals does with a gemmi
+// Scaling<float> (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<std::array<int, 3>> &hkl,
+                   const std::vector<double> &stol2,
+                   const std::vector<float> &fobs,
+                   const std::vector<float> &sigma,
+                   const std::vector<std::array<double, 6>> &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<ModelScaleFitState> &fits);
+    void FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std::vector<ModelScaleFitState> &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<int> hkl_;                  // 3 per point
+    CudaDevicePtr<double> stol2_;
+    CudaDevicePtr<float> fobs_;
+    CudaDevicePtr<unsigned char> strong_;
+    CudaDevicePtr<ModelScaleFitState> fits_;
+    CudaDevicePtr<double> partial_;           // per fit, per block, per slot
+    CudaDevicePtr<double> sums_;              // per fit, per slot
+    CudaHostPtr<ModelScaleFitState> host_fits_;
+    CudaHostPtr<double> host_sums_;
+};

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

tests/ModelScaleGPUTest.cpp

@@ -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