Refinement score: the model's structure factors as REFMAC 5.8 scores them at 0 cycles

Three terms of the score of a model - the same F_model at every stage, the as-given start and the
refined end alike - and a breakdown of it:

- Riding hydrogens from the start: the dictionary's hydrogens are added before the model is first
  scored (REFMAC's 0-cycle restores them; they move R by 0.010-0.013 at 1.1 A). The engine is made
  once, with them.
- The bulk-solvent mask on a grid of its own, min(d_min/3, 0.4 A), at most 2^26 points. On the
  density grid (1.1 A at 3.3 A) gemmi's 0.8 A shrink found no grid offset and was skipped, which put
  R 0.015-0.03 high at >= 2.5 A. 0.3 and 0.4 A agree to 0.001.
- The overall scale and anisotropic B by robust least squares (pseudo-Huber, iteratively
  reweighted, c = median |residual|) in the fine pass of the k_sol/b_sol grid (ModelScaleBox::robust,
  refinement only; MR and validation keep the plain fit). Plain least squares on |F| is the exact LS
  optimum but lets a few badly-fitted strong reflections set the scale (1.1 A set: R 0.1706 against
  REFMAC's 0.1600). The scale fit of refinement now runs on the CPU (no ModelScaleGPUPool).

RefinementSettings::score_only scores the model as given and stops; RefinementResult::score_start
carries every term (k, B, anisotropic B, k_sol, b_sol, mask spacing, R in 20 shells, -lnL work/free,
the likelihood FOM, hydrogens; per-reflection terms with score_only). GeometryStatistics gains rmsZ.
rugnux_refine_test: --score-only, --dump.

Measured against REFMAC 5.8.0431 ncyc 0 (refi bref iso) on the same model, data and flags, 14 open-arm
sets 1.1-3.6 A: R-work/R-free within 0.002 on 12 (was +0.003..+0.03); 7n2s and 9yl4 score 0.003/0.005
below REFMAC, which a 0.6-0.8 A mask grid reproduces. REFMAC with ANISOU stripped scores identically,
so isotropic B is not a parity term. Open: the likelihood's D/FOM (ours 0.01-0.11 higher at high
resolution); geometry differs only by REFMAC's metal LINK restraints and residues outside the
embedded dictionary.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit 4f94b21611874ad9dcf89e6771fe311a0e10bb06)
(cherry picked from commit 00db6a6f539ac2da515c0a649474335f4c4b07b8)

image_analysis/structure_refinement/Refinement.cpp

@@ -27,6 +27,7 @@
 #include "../../common/Logger.h"
 #include "ModelFFT.h"
 #include "ModelScaling.h"
+#include "RefinementMath.h"
 #include "RefinementXray.h"
 #include "SigmaA.h"
 #ifdef JFJOCH_USE_CUDA
@@ -98,15 +99,15 @@ struct Context {
     double d_min = 0;
     RefinementReflections refl;
     std::unique_ptr<RefinementEngine> engine;
-#ifdef JFJOCH_USE_CUDA
-    std::unique_ptr<ModelScaleGPUPool> scale_gpu;
-#endif
+    std::unique_ptr<BulkSolventMask> mask;
     size_t nthreads = 1;
     int device = -1;            // the engine's GPU
     size_t engine_bytes = 0;    // and what it holds there
     XrayRowModel rows;
     std::vector<std::complex<float>> fcalc, fmask;
     double k_sol = 0, b_sol = 0;
+    double k_overall = 0;
+    gemmi::SMat33<double> b_star{0, 0, 0, 0, 0, 0};
     double b_floor = B_FLOOR;   // the cycle's lowest allowed B
     double b_ceiling = B_CEILING_MAX;
     // Riding hydrogens: per atom the parent it rides on (-1: none), and where it sits from the parent.
@@ -175,7 +176,7 @@ struct Context {
         }
         {
             const RefinementStage s2("refit.fmask");
-            engine->Fmask(Model(), fmask);
+            mask->Fmask(Model(), refl.hkl, fmask);
         }
         std::optional<RefinementStage> s3("refit.scale");
 
@@ -192,15 +193,13 @@ struct Context {
         gemmi::Scaling<float> scaling(cell, sg);
         scaling.use_solvent = true;
         scaling.prepare_points(fc, fo, &fm);
-#ifdef JFJOCH_USE_CUDA
-        if (scale_gpu)
-            scale_gpu->Fit(scaling, SolventBMax(Model()), nthreads);
-        else
-#endif
-            FitModelScale(scaling, {.b_hi = SolventBMax(Model())}, nthreads);
+        // the scale of the score: by robust least squares (ModelScaleBox::robust), on the CPU
+        FitModelScale(scaling, {.b_hi = SolventBMax(Model()), .robust = true}, nthreads);
         s3.reset();
         k_sol = scaling.k_sol;
         b_sol = scaling.b_sol;
+        k_overall = scaling.k_overall;
+        b_star = scaling.b_star;
         rows.k_total.resize(refl.Size());
         rows.fmask.resize(refl.Size());
         std::vector<std::complex<float>> fmodel(refl.Size());
@@ -217,6 +216,71 @@ struct Context {
         return ComputeRFactors(refl, fmodel);
     }
 
+    // Every term of the score of the atoms as they are, at the scale, the solvent and the likelihood parameters
+    // of the last Refit().
+    RefinementScore Score() {
+        const std::vector<std::complex<float>> fmodel = FModel();
+        RefinementScore sc;
+        sc.k_overall = k_overall;
+        sc.k_sol = k_sol;
+        sc.b_sol = b_sol;
+        sc.mask_spacing = mask->Spacing();
+        // b_star is in the reciprocal basis (exp(-h^T b_star h / 4)); in Cartesian B = O b_star O^T
+        const gemmi::Mat33 b = cell.orth.mat.multiply(b_star.as_mat33()).multiply(cell.orth.mat.transpose());
+        sc.b_iso = (b[0][0] + b[1][1] + b[2][2]) / 3.0;
+        sc.b_aniso = {b[0][0] - sc.b_iso, b[1][1] - sc.b_iso, b[2][2] - sc.b_iso, b[0][1], b[0][2], b[1][2]};
+        const RFactors r = ComputeRFactors(refl, fmodel);
+        sc.r_work = r.r_work;
+        sc.r_free = r.r_free;
+        sc.n_work = r.n_work;
+        sc.n_free = r.n_free;
+        sc.nll_work = sc.nll_free = 0;
+        double fom_sum = 0;
+        for (size_t m = 0; m < refl.Size(); m++) {
+            double dfc = 0;
+            const double fc = std::abs(fmodel[m]);
+            const double nll = refine_math::AmplitudeNegLogLikelihood<double>(refl.fo[m], fc, rows.d[m], rows.var[m],
+                                                                             refl.centric[m], dfc);
+            (refl.free[m] ? sc.nll_free : sc.nll_work) += nll;
+            // the figure of merit of the same likelihood: <cos(dphi)> = I1(X)/I0(X), X = 2 Fo D Fc / var, for an
+            // acentric reflection; tanh(X / 2) for a centric one
+            // Following Murshudov, Vagin & Dodson (1997) Acta Cryst. D53, 240-255
+            const double x = 2.0 * refl.fo[m] * rows.d[m] * fc / rows.var[m];
+            double log_i0 = 0, ratio = 0;
+            refine_math::BesselI0I1(x, log_i0, ratio);
+            fom_sum += refl.centric[m] ? std::tanh(0.5 * x) : ratio;
+        }
+        sc.fom = fom_sum / static_cast<double>(refl.Size());
+        for (int p : rider_parent)
+            sc.hydrogens += p >= 0;
+        constexpr int SHELLS = 20;
+        double s_lo = refl.inv_d2[0], s_hi = refl.inv_d2[0];
+        for (double s2 : refl.inv_d2) {
+            s_lo = std::min(s_lo, s2);
+            s_hi = std::max(s_hi, s2);
+        }
+        std::vector<std::array<double, 4>> sums(SHELLS, {0, 0, 0, 0});
+        sc.shells.resize(SHELLS);
+        for (int k = 0; k < SHELLS; k++) {
+            sc.shells[k].d_max = 1.0 / std::sqrt(s_lo + (s_hi - s_lo) * k / SHELLS);
+            sc.shells[k].d_min = 1.0 / std::sqrt(s_lo + (s_hi - s_lo) * (k + 1) / SHELLS);
+        }
+        for (size_t m = 0; m < refl.Size(); m++) {
+            const int k = std::min(SHELLS - 1, static_cast<int>((refl.inv_d2[m] - s_lo) / (s_hi - s_lo) * SHELLS));
+            const int f = refl.free[m] ? 2 : 0;
+            sums[k][f] += std::fabs(refl.fo[m] - std::abs(fmodel[m]));
+            sums[k][f + 1] += refl.fo[m];
+            (refl.free[m] ? sc.shells[k].n_free : sc.shells[k].n_work)++;
+        }
+        for (int k = 0; k < SHELLS; k++) {
+            if (sums[k][1] > 0)
+                sc.shells[k].r_work = sums[k][0] / sums[k][1];
+            if (sums[k][3] > 0)
+                sc.shells[k].r_free = sums[k][2] / sums[k][3];
+        }
+        return sc;
+    }
+
     // F_model with the cycle's scale and solvent, for the atoms as they are now.
     std::vector<std::complex<float>> FModel() {
         engine->Fcalc(Model(), fcalc);
@@ -604,7 +668,6 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg
     const RefinementStage stage("engine");
     ctx.engine.reset();
 #ifdef JFJOCH_USE_CUDA
-    ctx.scale_gpu.reset();
     if (gpu && get_gpu_count() > 0) {
         try {
             const int device = get_gpu();
@@ -616,15 +679,12 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg
                                                  ctx.b_ceiling, budget, &bytes, settings.nthreads);
             ctx.device = device;
             ctx.engine_bytes = bytes;
-            if (ctx.engine)
-                ctx.scale_gpu = std::make_unique<ModelScaleGPUPool>(device, ctx.refl.Size(), 1);
-            else
+            if (!ctx.engine)
                 logger.Info("Refinement: on the CPU - the GPU cannot take this case or {:.1f} GB is over its budget",
                             bytes / 1e9);
         } catch (const JFJochException &e) {
             logger.Warning("Refinement: the GPU failed ({}), on the CPU", e.what());
             ctx.engine.reset();
-            ctx.scale_gpu.reset();
         }
     }
 #else
@@ -681,6 +741,7 @@ std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, co
     ctx.refl = main.refl;
     ctx.nthreads = nthreads;
     ctx.b_ceiling = main.b_ceiling;
+    ctx.mask = std::make_unique<BulkSolventMask>(ctx.cell, *ctx.sg, ctx.d_min, nthreads);
     slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
     ctx.atoms = AtomList(ctx.Model());
     ctx.rider_parent = slot->topo.rider_parent;
@@ -691,7 +752,6 @@ std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, co
                                                  ctx.b_ceiling, SIZE_MAX, nullptr, nthreads);
             if (!ctx.engine)
                 return nullptr;
-            ctx.scale_gpu = std::make_unique<ModelScaleGPUPool>(device, ctx.refl.Size(), 1);
         } catch (const JFJochException &) {
             return nullptr;
         }
@@ -855,6 +915,23 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
         return result;
     }
 
+    ctx.mask = std::make_unique<BulkSolventMask>(cell, sg, ctx.d_min, settings.nthreads);
+
+    // The dictionary's view of the model, and its hydrogens riding on their parents from the start, so that the
+    // model is scored as the restrained stage refines it.
+    const bool restrained = settings.restrained || settings.score_only;
+    // The weight trials' slots start from a copy of the model as it is before the topology is made.
+    const gemmi::Structure before_topology = ctx.st;
+    RefinementTopology topo;
+    if (restrained) {
+        std::optional<RefinementStage> topo_stage("topology");
+        topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
+        topo_stage.reset();
+        ctx.atoms = AtomList(ctx.Model());
+        ctx.rider_parent = topo.rider_parent;
+        PlaceRiders(ctx, topo);
+    }
+
     MakeEngine(ctx, settings.gpu, settings, logger);
     result.on_gpu = ctx.engine->OnGPU();
 
@@ -863,11 +940,51 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
     result.n_free = r.n_free;
     result.r_work_start = r.r_work;
     result.r_free_start = r.r_free;
-    logger.Info("Refinement: {} atoms, {} reflections to {:.2f} A ({} free), on the {}; as given R-work {:.4f}, R-free {:.4f}",
-                ctx.atoms.size(), ctx.refl.Size(), result.d_min, r.n_free, result.on_gpu ? "GPU" : "CPU", r.r_work,
-                r.r_free);
+    result.score_start = ctx.Score();
+    logger.Info("Refinement: {} atoms ({} riding hydrogens), {} reflections to {:.2f} A ({} free), on the {}; as given "
+                "R-work {:.4f}, R-free {:.4f}; k_sol {:.3f}, b_sol {:.1f}, overall B {:.1f}, mask grid {:.2f} A",
+                ctx.atoms.size(), result.score_start.hydrogens, ctx.refl.Size(), result.d_min, r.n_free,
+                result.on_gpu ? "GPU" : "CPU", r.r_work, r.r_free, ctx.k_sol, ctx.b_sol, result.score_start.b_iso,
+                result.score_start.mask_spacing);
+
+    // The maps, the blobs and the model of where the refinement ended.
+    auto finish = [&]() {
+        MakeMaps(ctx, merged, wavelength_A, result);
+        {
+            const RefinementStage stage("blobs");
+            result.blobs = FindDensityBlobs(*result.difference_map, result.difference_map_rms, ctx.st, result.d_min,
+                                            result.anomalous_map.get(), result.anomalous_map_rms);
+        }
+        result.k_sol = ctx.k_sol;
+        result.b_sol = ctx.b_sol;
+        result.model = std::make_shared<gemmi::Structure>(ctx.st);
+        result.ok = true;
+        result.seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
+    };
+    if (settings.score_only) {
+        result.score_start.hkl = ctx.refl.hkl;
+        result.score_start.fcalc = ctx.fcalc;
+        result.score_start.fmask = ctx.rows.fmask;
+        result.score_start.k_total = ctx.rows.k_total;
+        result.score_start.d = ctx.rows.d;
+        result.score_start.var = ctx.rows.var;
+        result.r_work_placement = result.r_work_start;
+        result.r_free_placement = result.r_free_start;
+        result.restrained_residues = topo.restrained_residues;
+        result.fixed_residues = topo.fixed_residues;
+        result.fixed_names = topo.fixed_names;
+        result.geometry_start = MeasureGeometry(topo.geometry, BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0),
+                                                ctx.XyzB());
+        finish();
+        logger.Info("Refinement: the model scored only; geometry rmsd bonds {:.4f} A (rmsZ {:.2f}), angles {:.2f} deg "
+                    "(rmsZ {:.2f})", result.geometry_start.rmsd_bonds, result.geometry_start.rmsz_bonds,
+                    result.geometry_start.rmsd_angles, result.geometry_start.rmsz_angles);
+        return result;
+    }
 
     RefinePlacement(ctx, settings.group_b, result);
+    if (restrained)
+        PlaceRiders(ctx, topo);
     r = ctx.Refit();
     result.r_work_placement = r.r_work;
     result.r_free_placement = r.r_free;
@@ -875,18 +992,6 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
                 result.rigid_bodies, result.group_b ? "kept" : "not kept", r.r_work, r.r_free);
 
     if (settings.restrained) {
-        const gemmi::Structure before_topology = ctx.st;
-        std::optional<RefinementStage> topo_stage("topology");
-        RefinementTopology topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
-        topo_stage.reset();
-        // the hydrogens it added: the atoms and the engine made again for them
-        ctx.atoms = AtomList(ctx.Model());
-        ctx.rider_parent = topo.rider_parent;
-        PlaceRiders(ctx, topo);
-        if (static_cast<int>(ctx.atoms.size()) != result.atoms) {
-            MakeEngine(ctx, result.on_gpu, settings, logger);
-            result.on_gpu = ctx.engine->OnGPU();
-        }
         ctx.Refit();
         result.restrained_residues = topo.restrained_residues;
         result.fixed_residues = topo.fixed_residues;
@@ -998,17 +1103,7 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
         result.geometry = chosen.geometry;
     }
 
-    MakeMaps(ctx, merged, wavelength_A, result);
-    {
-        const RefinementStage stage("blobs");
-        result.blobs = FindDensityBlobs(*result.difference_map, result.difference_map_rms, ctx.st, result.d_min,
-                                        result.anomalous_map.get(), result.anomalous_map_rms);
-    }
-    result.k_sol = ctx.k_sol;
-    result.b_sol = ctx.b_sol;
-    result.model = std::make_shared<gemmi::Structure>(ctx.st);
-    result.ok = true;
-    result.seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
+    finish();
     logger.Info("Refinement: done in {:.1f} s, {} target evaluations{}", result.seconds, ctx.evaluations,
                 result.restrained ? fmt::format("; R-work {:.4f}, R-free {:.4f}, weight {:.3g}", result.r_work,
                                                 result.r_free, result.weight)