Rigid body GPU: the per-evaluation scale fit on the device (ModelScaleGPU)

The solvent grid once per zone and the overall scale and anisotropic B at every
evaluation now run through ModelScaleGPU on the engine's stream, on the points
packed on the device. A zone whose grid fits would chain (five or fewer
reflections for the isotropic start) is scaled on the host exactly as the CPU
target scales it.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C

rugnux/RigidBodyGPU.cpp

@@ -133,6 +133,7 @@ RigidBodyGPUZone MakeZone(const gemmi::Model &model, const gemmi::UnitCell &cell
         const gemmi::ValueSigma<float> &vs = fobs->v[i].value;
         zone.obs_row.push_back(m);
         zone.obs_fobs.push_back(vs.value);
+        zone.obs_sigma.push_back(vs.sigma);
         sum += vs.value;
         // the points gemmi's prepare_points() takes
         if (m >= 0 && !std::isnan(vs.value) && !std::isnan(vs.sigma))
@@ -257,6 +258,7 @@ void RigidBodyTargetGPU::SetZone(const gemmi::AsuData<gemmi::ValueSigma<float>>
     fobs_ = fobs;
     d_min_ = d_min;
     solvent_fitted_ = false;
+    host_scale_ = false;
     have_point_ = false;
     const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell_, sg_, d_min);
     orbit_leaders_ = OrbitLeaders(grid, nthreads_);
@@ -272,6 +274,36 @@ void RigidBodyTargetGPU::SetZone(const gemmi::AsuData<gemmi::ValueSigma<float>>
     }
 }
 
+// The scale fitted on the host, exactly as RigidBodyTarget::Residuals() fits it, for a zone the device
+// does not reproduce: one with five or fewer reflections for fit_isotropic_b_approximately(), where
+// FitModelScale()'s grid fits chain from one to the next and every fit starts from the last.
+void RigidBodyTargetGPU::HostScale() {
+    std::vector<std::array<float, 2>> fc, fm;
+    engine_.DownloadPoints(fc, fm);
+    gemmi::Scaling<float> scaling(cell_, &sg_);
+    scaling.use_solvent = true;
+    for (size_t p = 0; p < zone_->point_obs.size(); p++) {
+        const auto &o = fobs_.v[zone_->point_obs[p]];
+        const int m = zone_->obs_row[zone_->point_obs[p]];
+        scaling.points.push_back({o.hkl, zone_->row_stol2[m], std::complex<float>(fc[p][0], fc[p][1]),
+                                  std::complex<float>(fm[p][0], fm[p][1]), o.value.value, o.value.sigma});
+    }
+    if (!solvent_fitted_) {
+        FitModelScale(scaling, {}, nthreads_);
+        k_sol = scaling.k_sol;
+        b_sol = scaling.b_sol;
+        solvent_fitted_ = true;
+    }
+    scaling.k_sol = k_sol;
+    scaling.b_sol = b_sol;
+    scaling.fix_k_sol = true;
+    scaling.fix_b_sol = true;
+    scaling.fit_isotropic_b_approximately();
+    scaling.fit_parameters();
+    k_overall_ = scaling.k_overall;
+    b_star_ = scaling.b_star;
+}
+
 // RigidBodyTarget::Residuals() with the density, the composition, the transforms and the residuals on
 // the device.
 bool RigidBodyTargetGPU::Residuals(const double q[6], double *residuals) {
@@ -289,42 +321,31 @@ bool RigidBodyTargetGPU::Residuals(const double q[6], double *residuals) {
         if (zone_->row_hkl.empty())
             return false;
 
-        std::vector<std::array<float, 2>> fc, fm;
-        engine_.DownloadPoints(fc, fm);
-        gemmi::Scaling<float> scaling(cell_, &sg_);
-        scaling.use_solvent = true;
-        for (size_t p = 0; p < zone_->point_obs.size(); p++) {
-            const auto &o = fobs_.v[zone_->point_obs[p]];
-            const int m = zone_->obs_row[zone_->point_obs[p]];
-            scaling.points.push_back({o.hkl, zone_->row_stol2[m], std::complex<float>(fc[p][0], fc[p][1]),
-                                      std::complex<float>(fm[p][0], fm[p][1]), o.value.value, o.value.sigma});
-        }
-        if (scaling.points.empty())
+        if (zone_->point_obs.empty())
             return false;
         // The scale as RigidBodyTarget::Residuals() fits it: the solvent once per zone, then the
         // overall scale and the anisotropic B at every evaluation.
-        if (!solvent_fitted_) {
-            FitModelScale(scaling, {}, nthreads_);
-            k_sol = scaling.k_sol;
-            b_sol = scaling.b_sol;
-            solvent_fitted_ = true;
+        if (!solvent_fitted_ && !host_scale_) {
+            try {
+                engine_.FitSolvent(k_sol, b_sol);
+                solvent_fitted_ = true;
+            } catch (const std::logic_error &) {
+                host_scale_ = true;
+            }
         }
-        scaling.k_sol = k_sol;
-        scaling.b_sol = b_sol;
-        scaling.fix_k_sol = true;
-        scaling.fix_b_sol = true;
-        scaling.fit_isotropic_b_approximately();
-        scaling.fit_parameters();
-
-        const gemmi::SMat33<double> &b = scaling.b_star;
-        const double b_star[6] = {b.u11, b.u22, b.u33, b.u12, b.u13, b.u23};
-        engine_.Residuals(scaling.k_overall, b_star, k_sol, b_sol, residuals);
+        if (host_scale_) {
+            HostScale();
+        } else {
+            double b[6];
+            engine_.FitScale(k_sol, b_sol, k_overall_, b);
+            b_star_ = {b[0], b[1], b[2], b[3], b[4], b[5]};
+        }
+        const double b_star[6] = {b_star_.u11, b_star_.u22, b_star_.u33, b_star_.u12, b_star_.u13, b_star_.u23};
+        engine_.Residuals(k_overall_, b_star, k_sol, b_sol, residuals);
         unmatched = static_cast<int>(std::count(zone_->obs_row.begin(), zone_->obs_row.end(), -1));
 
         have_point_ = true;
         std::copy(q, q + 6, q_.begin());
-        k_overall_ = scaling.k_overall;
-        b_star_ = scaling.b_star;
         return true;
     } catch (const JFJochException &e) {
         throw RigidBodyGPUFailure(e.what());

rugnux/RigidBodyGPU.cu

@@ -11,6 +11,7 @@
 #include <cub/cub.cuh>
 #include <cufft.h>
 
+#include "ModelScaleGPU.h"
 #include "../common/JFJochException.h"
 #include "../image_analysis/indexing/CUDAMemHelpers.h"
 
@@ -502,6 +503,7 @@ struct RigidBodyGPUEngineImpl {
     CudaDevicePtr<double> products, x;
 
     std::map<std::array<int, 4>, cufftHandle> plans;
+    std::unique_ptr<ModelScaleGPU> scale;
 
     // The zone.
     int n_atoms = 0;
@@ -556,6 +558,16 @@ struct RigidBodyGPUEngineImpl {
         point_fmask = CudaDevicePtr<float2>(no, sync);
         products = CudaDevicePtr<double>(MAX_SCALE_PARAMS * (MAX_SCALE_PARAMS + 6), sync);
         x = CudaDevicePtr<double>(MAX_SCALE_PARAMS * 6, sync);
+        scale = std::make_unique<ModelScaleGPU>(stream, no);
+    }
+
+    // The points' fcmol and fmask at this evaluation.
+    void PackPoints() {
+        if (n_points == 0)
+            return;
+        pack_points_kernel<<<blocks(n_points), THREADS, 0, stream>>>(point_obs, obs_row, static_cast<int>(n_points),
+                                                                    f, fmask, point_fcmol, point_fmask);
+        cuda_err(cudaGetLastError());
     }
 
     ~RigidBodyGPUEngineImpl() {
@@ -706,6 +718,7 @@ size_t RigidBodyGPUEngine::DeviceBytes(const RigidBodyGPUCapacity &c) {
     bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 7 * sizeof(double2) + sizeof(float2));
     bytes += no * (2 * sizeof(int) + sizeof(float) + (1 + 6 + MAX_SCALE_PARAMS) * sizeof(double) + 2 * sizeof(float2));
     bytes += CubBytes(na, c.pairs + 1);
+    bytes += ModelScaleGPU::DeviceBytes(no);
     return bytes;
 }
 
@@ -795,6 +808,17 @@ void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) {
     up(e.obs_row, zone.obs_row.data(), zone.obs_row.size() * sizeof(int));
     up(e.obs_fobs, zone.obs_fobs.data(), zone.obs_fobs.size() * sizeof(float));
     up(e.point_obs, zone.point_obs.data(), zone.point_obs.size() * sizeof(int));
+    std::vector<std::array<int, 3>> point_hkl;
+    std::vector<double> point_stol2;
+    std::vector<float> point_fobs, point_sigma;
+    for (int i : zone.point_obs) {
+        const int m = zone.obs_row[i];
+        point_hkl.push_back(zone.row_hkl[m]);
+        point_stol2.push_back(zone.row_stol2[m]);
+        point_fobs.push_back(zone.obs_fobs[i]);
+        point_sigma.push_back(zone.obs_sigma[i]);
+    }
+    e.scale->SetPoints(point_hkl, point_stol2, point_fobs, point_sigma, zone.constraints, zone.frac);
     e.Plan(1);
     e.Plan(3);
     cuda_err(cudaStreamSynchronize(e.stream));
@@ -830,15 +854,30 @@ void RigidBodyGPUEngine::DownloadPoints(std::vector<std::array<float, 2>> &fcmol
     fmask.resize(e.n_points);
     if (e.n_points == 0)
         return;
-    pack_points_kernel<<<blocks(e.n_points), THREADS, 0, e.stream>>>(e.point_obs, e.obs_row,
-                                                                    static_cast<int>(e.n_points), e.f, e.fmask,
-                                                                    e.point_fcmol, e.point_fmask);
-    cuda_err(cudaGetLastError());
+    e.PackPoints();
     cuda_err(cudaMemcpyAsync(fcmol.data(), e.point_fcmol, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
     cuda_err(cudaMemcpyAsync(fmask.data(), e.point_fmask, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
     cuda_err(cudaStreamSynchronize(e.stream));
 }
 
+void RigidBodyGPUEngine::FitSolvent(double &k_sol, double &b_sol) {
+    RigidBodyGPUEngineImpl &e = *impl_;
+    cuda_err(cudaSetDevice(e.device));
+    e.PackPoints();
+    const ModelSolventFit fit = e.scale->FitSolvent(e.point_fcmol, e.point_fmask);
+    k_sol = fit.k_sol;
+    b_sol = fit.b_sol;
+}
+
+void RigidBodyGPUEngine::FitScale(double k_sol, double b_sol, double &k_overall, double b_star[6]) {
+    RigidBodyGPUEngineImpl &e = *impl_;
+    cuda_err(cudaSetDevice(e.device));
+    e.PackPoints();
+    const ModelScaleParams p = e.scale->Fit(e.point_fcmol, e.point_fmask, k_sol, b_sol);
+    k_overall = p.k_overall;
+    std::copy(p.b_star, p.b_star + 6, b_star);
+}
+
 void RigidBodyGPUEngine::Residuals(double k_overall, const double b_star[6], double k_sol, double b_sol,
                                    double *residuals) {
     RigidBodyGPUEngineImpl &e = *impl_;

rugnux/RigidBodyGPU.h

@@ -75,6 +75,7 @@ public:
 private:
     // The placement at q as x -> R (x - centroid) + centroid + t.
     void Placement(const double q[6], double rotation[9], double translation[3]) const;
+    void HostScale();
 
     RigidBodyGPUPool &pool_;
     RigidBodyGPUEngine &engine_;
@@ -87,6 +88,7 @@ private:
     double d_min_ = 0;
     std::unique_ptr<RigidBodyGPUZone> zone_;
     bool solvent_fitted_ = false;
+    bool host_scale_ = false;             // the zone's scale is fitted on the host (HostScale)
     std::vector<size_t> orbit_leaders_;   // for the host mask
 
     bool have_point_ = false;

rugnux/RigidBodyGPUEngine.h

@@ -53,7 +53,7 @@ struct RigidBodyGPUZone {
 
     // The observations the residuals are over: each one's row (-1 without one) and amplitude.
     std::vector<int> obs_row;
-    std::vector<float> obs_fobs;
+    std::vector<float> obs_fobs, obs_sigma;
     double f_mean = 1;
     // The scale's points: the observations gemmi's prepare_points() would take, in order.
     std::vector<int> point_obs;
@@ -101,7 +101,12 @@ public:
     // instead (the stub until ModelMaskGPU is in).
     void Fcalc(const double rotation[9], const double translation[3]);
     void Fmask(const float *host_mask);
-    // The scale's points: fcmol (Fcalc as complex<float>) and fmask, downloaded (the stub).
+    // The scale at this evaluation's Fcalc and Fmask (ModelScaleGPU): FitModelScale()'s solvent grid, and
+    // the overall scale and anisotropic B at a fixed solvent. FitSolvent() throws std::logic_error where
+    // the grid's fits would chain (too few reflections for the isotropic start), which the host then does.
+    void FitSolvent(double &k_sol, double &b_sol);
+    void FitScale(double k_sol, double b_sol, double &k_overall, double b_star[6]);
+    // The scale's points: fcmol (Fcalc as complex<float>) and fmask, downloaded for the host's fit.
     void DownloadPoints(std::vector<std::array<float, 2>> &fcmol, std::vector<std::array<float, 2>> &fmask);
     // The residuals at the scale given, NumObservations() of them, into `residuals` (host).
     void Residuals(double k_overall, const double b_star[6], double k_sol, double b_sol, double *residuals);