Files
Jungfraujoch/image_analysis/structure_refinement/ModelScaling.h
T
leonarski_fandClaude Opus 5.5 1fe73bab41 Model validation: scale fits on the GPU, change-of-basis probe gridded once per frame on all threads
The scale fit of every model-validation fit (fit_model: the real model, the indexing probe's relabellings
and every null replicate) was FitModelScale on the host - a long-double Levenberg-Marquardt per solvent grid
point, ~80% of the validation's CPU on a 400 A cell. It now runs on the GPU through ModelScaleGPU::FitSolvent,
the engine the rigid body and molecular replacement already use, from a small pool (ModelScaleGPUPool: up to
four engines with their own streams, sized on the card's total memory, at most an eighth of it). The grid,
the fits and the winner rule are FitModelScale's; the sums are taken in double on the device, so R-factors
agree with the host's to rounding. The CPU build, and a GPU build without a card, keep FitModelScale.

The change-of-basis probe gridded each candidate frame with gemmi's serial put_model_density_on_grid, once
for the data as merged and once more for each alternative indexing. It now grids each frame once, with
ModelGrid's parallel PutModelDensityOnGrid (the same grid, bit for bit), from threads outside the pool so the
gridding spreads over it, and scores every labelling on that F_calc. Exact.

4nwv (122k atoms, I222, 400 A cell): model validation 33.2 s -> 14.0 s; every logged R, sigma and decision
the same.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
2026-10-12 21:40:54 +02:00

85 lines
4.2 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
#include <condition_variable>
#include <cstddef>
#include <memory>
#include <mutex>
#include <vector>
#include "gemmi/model.hpp"
#include "gemmi/scaling.hpp"
// Fit the overall scale, the anisotropic B and the flat bulk solvent of a model to observed
// amplitudes, writing the answer into `scaling` (k_overall, b_star, k_sol, b_sol) so that
// gemmi's own scale_data() can then apply it.
//
// The caller decides which reflections are fitted: FitModelScale reads scaling.points, which
// prepare_points() filled. Pass the WORKING set to keep the free reflections out of the scale
// model - that is a property of the call site, not of this function.
//
// Replaces gemmi's fit_isotropic_b_approximately() + fit_parameters(), which is an unbounded
// Levenberg-Marquardt: nothing there stops k_sol and b_sol reaching values a flat solvent model
// is meaningless at (b_sol of 1707 A^2 was measured on this corpus). Here k_sol and b_sol come
// from a grid search over the physical ranges, so an unphysical answer cannot be represented.
// The box the solvent parameters are searched in. The defaults are the physically reasonable
// range reported by Fokine & Urzhumtsev and adopted as phenix's default grid: k_sol in (0.1, 0.8)
// and b_sol in (10, 80) A^2, of which phenix searches b_sol 10-80 in steps of 5. A model with high
// atomic B gets a higher b_sol ceiling (SolventBMax).
struct ModelScaleBox {
double k_lo = 0.10, k_hi = 0.60;
double b_lo = 10.0, b_hi = 80.0;
};
// The b_sol ceiling for `model`: 80 A^2, or twice the model's mean atomic B where that is higher.
// b_sol is how far the solvent's edge is smeared, and the edge can be no sharper than the atoms it
// lies against: a model at B of 300 A^2 beside a solvent held at 80 gets a solvent term that reaches
// far too deep in resolution, which only a too-small k_sol can then offset. Fokine & Urzhumtsev's
// 10-80 A^2 is for models at the B of ordinary entries (their mean b_sol 46 A^2), where the ceiling
// stays at 80.
double SolventBMax(const gemmi::Model &model);
struct ModelScaleReport {
int n_points = 0; // reflections the fit actually used
int n_grid = 0; // grid points evaluated
double r_work_fit = 1.0; // R on the fitted reflections, at the chosen solution
};
// nthreads: the grid points are fitted in parallel on up to this many threads; the answer does not
// depend on it.
ModelScaleReport FitModelScale(gemmi::Scaling<float> &scaling, ModelScaleBox box = {}, size_t nthreads = 1);
#ifdef JFJOCH_USE_CUDA
// FitModelScale on a GPU (ModelScaleGPU::FitSolvent): the same grid over the same box, the same fits and
// the same winner, with the sums over the reflections taken in double on the device where the host takes
// them in long double - so the answer is FitModelScale's to rounding, not to the bit. A few engines,
// reserved once, each a stream and room for one fit at a time; a fit takes a free one and waits for one
// where all are taken. Every kernel is deterministic, so which engine a fit gets changes no number.
class ModelScaleGPUPool {
public:
// What `engines` engines for fits of up to `max_points` reflections take on the card.
static size_t DeviceBytes(size_t max_points, size_t engines);
// On `device`. Throws JFJochException on a CUDA failure.
ModelScaleGPUPool(int device, size_t max_points, size_t engines);
~ModelScaleGPUPool();
// FitModelScale(scaling, {.b_hi = b_hi}, nthreads). Where FitModelScale would walk the grid point to
// point (five or fewer reflections for the isotropic start), it is FitModelScale on the host. Safe to
// call from several threads. Throws JFJochException on a CUDA failure.
ModelScaleReport Fit(gemmi::Scaling<float> &scaling, double b_hi, size_t nthreads);
// What the pool holds on the card.
size_t Bytes() const { return DeviceBytes(max_points_, engines_.size()); }
private:
struct Engine;
int device_;
size_t max_points_;
std::vector<std::unique_ptr<Engine>> engines_;
std::vector<Engine *> idle_;
std::mutex m_;
std::condition_variable cv_;
};
#endif