Merge branch 'rb-gpu-polish' into rc173 (no environment variables; memo verification as a test; GPU/CPU parity tests)

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C
This commit is contained in:
2026-09-28 23:02:09 +02:00
co-authored by Claude Opus 5.5
9 changed files with 112 additions and 52 deletions
@@ -9,7 +9,6 @@
#include "../lattice_search/LatticeSearch.h"
#include "../indexing/MultiLatticeSearch.h"
#include <cstdlib>
#include <future>
#include <mutex>
#include <stdexcept>
@@ -255,7 +254,7 @@ void RotationIndexer::RunIndexing() {
if (k == key)
kept = outcome;
}
if (kept && !std::getenv("RUGNUX_VERIFY_FIRST_PASS_MEMO")) {
if (kept && !verify_memo) {
axis_ = kept->axis;
updated_geom_ = kept->updated_geom;
search_result_ = kept->search_result;
@@ -272,7 +271,7 @@ void RotationIndexer::RunIndexing() {
const IndexingOutcome outcome = Outcome();
if (kept) {
if (!SameOutcome(*kept, outcome))
throw std::runtime_error("RotationIndexer: the indexing memo differs from a recomputation");
throw MemoMismatch("RotationIndexer: the indexing memo differs from a recomputation");
return;
}
auto &memo = Memo();
@@ -6,6 +6,7 @@
#include <memory>
#include <vector>
#include <mutex>
#include <stdexcept>
#include "../../common/DiffractionSpot.h"
#include "../../common/DiffractionExperiment.h"
@@ -38,6 +39,12 @@ struct RotationIndexerResult {
std::optional<TiltWalk> tilt_walk;
};
// What a memo verification throws where a reused result differs from its recomputation - its own
// type, so that no pass that catches a failure and carries on can take it for an ordinary one.
struct MemoMismatch : std::runtime_error {
using std::runtime_error::runtime_error;
};
class RotationIndexer {
public:
// RunIndexing is deterministic in what it reads - the accumulated spots and their angles, the
@@ -45,8 +52,8 @@ public:
// of the pool it indexes with - and a rotation run asks the same question again and again (the
// canonical pass repeats its rotation-scale probe, a probe repeats its pass's rescue ladder).
// So what it computes is kept, process-wide, under exactly those inputs, and a second
// RotationIndexer asking with the same inputs takes the answer. RUGNUX_VERIFY_FIRST_PASS_MEMO
// recomputes anyway and throws on any difference.
// RotationIndexer asking with the same inputs takes the answer. VerifyMemo(true) recomputes anyway
// and throws on any difference.
struct IndexingOutcome {
std::optional<GoniometerAxis> axis;
DiffractionGeometry updated_geom;
@@ -68,6 +75,7 @@ private:
// caller lowers it to seed the search on the strongest few, exactly as the per-frame indexer
// escalates 30 -> 80 -> all, for a pattern where the deep list is mostly not this crystal.
size_t max_spots_per_image = DEFAULT_MAX_SPOTS_PER_IMAGE;
bool verify_memo = false;
const bool index_ice_rings;
const bool real_time; // see the constructor
@@ -108,6 +116,8 @@ public:
RotationIndexer(const DiffractionExperiment& x, IndexerThreadPool& indexer, bool real_time = false);
// Keep only this many spots per image in the accumulated cloud (strongest, non-ring first).
void MaxSpotsPerImage(size_t n) { max_spots_per_image = n > 0 ? n : DEFAULT_MAX_SPOTS_PER_IMAGE; }
// Recompute on a memo hit and throw if the two differ - the tests' check that the memo is exact.
void VerifyMemo(bool v) { verify_memo = v; }
// angle_deg is the image's mid-exposure rotation angle; if omitted, the goniometer angle at
// `image` is used (only valid when `image` is the goniometer's own image index).
void ProcessImage(int64_t image, const std::vector<SpotToSave>& spots,
+1 -2
View File
@@ -38,7 +38,6 @@
#include "RigidBodyRefine.h"
#include "SigmaA.h"
#ifdef JFJOCH_USE_CUDA
#include <cstdlib>
#include "RigidBodyGPU.h"
#include "../common/CUDAWrapper.h"
#endif
@@ -555,7 +554,7 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
RigidBodyGPUPool *rigid_body_pool = nullptr;
#ifdef JFJOCH_USE_CUDA
std::unique_ptr<RigidBodyGPUPool> rigid_body_engines;
if (rigid_body_gpu && std::getenv("JFJOCH_RIGID_BODY_CPU") == nullptr) {
if (rigid_body_gpu) {
const double finest_zone = RigidBodyLadder(d_min).back();
size_t zone_observations = 0;
for (const MergedReflection &r : obs)
+13 -11
View File
@@ -148,6 +148,9 @@ namespace {
// has quietly changed. Fail the run instead, and let the operator free the card and repeat it.
bool IsFatalResourceError(const std::exception &e) {
if (dynamic_cast<const std::bad_alloc *>(&e)) return true;
// Not a resource, but just as fatal: a memo verification (verify_first_pass_memo) that found a
// difference must end the run, not pass for a probe that scored nothing.
if (dynamic_cast<const MemoMismatch *>(&e)) return true;
const auto *jf = dynamic_cast<const JFJochException *>(&e);
return jf != nullptr && (jf->Category() == JFJochExceptionCategory::GPUCUDAError
|| jf->Category() == JFJochExceptionCategory::MemAllocFailed);
@@ -3397,7 +3400,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
// detector. Nothing between here and the first pass changes what the key reads.
if (full && indexing_probe_only_ && !force_rotation_result_.has_value()
&& !config_.forced_rotation_lattice.has_value() && config_.rotation_indexing && config_.two_pass_rotation
&& !std::getenv("RUGNUX_VERIFY_FIRST_PASS_MEMO")) {
&& !config_.verify_first_pass_memo) {
const std::vector<double> first_pass_key = FirstPassInputKey();
const auto memo = std::find_if(first_pass_memo_.begin(), first_pass_memo_.end(),
[&](const FirstPassMemo &m) { return m.key == first_pass_key; });
@@ -3579,7 +3582,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
logger.Info("Rotation indexer lattice forced externally - skipping first pass");
} else if (full && config_.rotation_indexing && config_.two_pass_rotation) {
// An indexing probe over inputs a first pass has already been run on returned above, except
// under RUGNUX_VERIFY_FIRST_PASS_MEMO, which runs the pass and compares.
// under verify_first_pass_memo, which runs the pass and compares.
const std::vector<double> first_pass_key = FirstPassInputKey();
std::optional<ValidationSpotEvidence> memo_hit;
const auto memo = std::find_if(first_pass_memo_.begin(), first_pass_memo_.end(),
@@ -3668,7 +3671,6 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
};
int prefetch_batches = 0;
double prefetch_time_s = 0.0;
const bool verify_spots = std::getenv("RUGNUX_VERIFY_FIRST_PASS_MEMO") != nullptr;
const auto prefetch_spots = [&](const std::vector<int> &ordinals) {
// What an earlier pass already found under the same key (first_pass_spots_) is taken from
// there.
@@ -3685,7 +3687,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
std::vector<int> wanted;
for (const int ordinal : ordinals)
if (!spot_cache.contains(ordinal)) {
if (const auto it = kept.find(ordinal); it != kept.end() && !verify_spots)
if (const auto it = kept.find(ordinal); it != kept.end() && !config_.verify_first_pass_memo)
spot_cache.emplace(ordinal, it->second);
else
wanted.push_back(ordinal);
@@ -3731,9 +3733,8 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
}
for (size_t i = 0; i < wanted.size(); i++) {
if (const auto it = kept.find(wanted[i]); it != kept.end() && !SameSpots(it->second, found[i]))
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
fmt::format("First-pass spot memo differs from a recomputation "
"for image {}", wanted[i]));
throw MemoMismatch(fmt::format("First-pass spot memo differs from a recomputation "
"for image {}", wanted[i]));
entry->second.emplace(wanted[i], found[i]);
}
}
@@ -3963,6 +3964,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
const DiffractionExperiment *x = nullptr) {
auto ri = std::make_unique<RotationIndexer>(x ? *x : experiment_, pool);
ri->MaxSpotsPerImage(first_pass_spots_per_image);
ri->VerifyMemo(config_.verify_first_pass_memo);
for (size_t i = 0; i < ordinals.size(); i++) {
if (cancelled_ || ri->AccumulationFull())
break;
@@ -5466,10 +5468,9 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
result.validation_evidence = evidence;
if (memo_hit && (memo_hit->spots != evidence.spots || memo_hit->on_lattice != evidence.on_lattice
|| memo_hit->by_chance != evidence.by_chance))
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
fmt::format("First-pass memo mismatch: stored {}/{}/{}, recomputed {}/{}/{}",
memo_hit->on_lattice, memo_hit->spots, memo_hit->by_chance,
evidence.on_lattice, evidence.spots, evidence.by_chance));
throw MemoMismatch(fmt::format("First-pass memo mismatch: stored {}/{}/{}, recomputed {}/{}/{}",
memo_hit->on_lattice, memo_hit->spots, memo_hit->by_chance,
evidence.on_lattice, evidence.spots, evidence.by_chance));
// Kept only where the pass ended on the inputs it started from, and did not take the
// beam-centre ladder, which a probe (beam_center_searched_) never would.
if (!geometry_prepass && !beam_center_ladder_ran && FirstPassInputKey() == first_pass_key)
@@ -5657,6 +5658,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
const CrystalLattice primary_reduced = reduced(lattice);
for (int id = 1; id <= 3; id++) {
RotationIndexer ri(x, *indexer_pool);
ri.VerifyMemo(config_.verify_first_pass_memo);
for (const int f : spread_ordinals) {
std::vector<SpotToSave> left;
const auto &sp = spot_cache.at(f);
+8 -3
View File
@@ -202,6 +202,11 @@ struct ProcessConfig {
// a user-fixed -S writes nothing, because its centring absences were never integrated.
bool write_p1_crosscheck = true;
bool finalist_ledger = false; // --finalist-ledger; report-only symmetry evidence table
// Where a rotation run would reuse its first-pass memos - a probe's validation evidence, a frame's
// spot list, an indexing outcome - recompute instead and throw if the two differ. Set by the tests
// that prove the reuse exact; not on the command line.
bool verify_first_pass_memo = false;
};
// A rotation lattice's score on the validation frames spread over the whole sweep: their spots (off
@@ -731,7 +736,7 @@ class Rugnux {
// pass read (FirstPassInputKey). The first pass is deterministic in them, so an indexing probe
// whose key matches takes the stored evidence instead of indexing again - the geometry walk's
// probe at the geometry in hand is the canonical pass's own first pass run over. With
// RUGNUX_VERIFY_FIRST_PASS_MEMO set the probe runs anyway and the run fails if the two differ.
// ProcessConfig::verify_first_pass_memo the probe runs anyway and the run fails if the two differ.
struct FirstPassMemo {
std::vector<double> key;
ValidationSpotEvidence evidence;
@@ -740,8 +745,8 @@ class Rugnux {
// The first pass's spot lists, kept across passes under everything a frame's list depends on
// (SpotFindingKey): a probe and the canonical pass after it, or a pass repeating an earlier
// pass's rescue, find the spots of the same frames under the same settings again. Shared with the
// copies the run makes of itself. RUGNUX_VERIFY_FIRST_PASS_MEMO finds them again and throws on a
// difference.
// copies the run makes of itself. ProcessConfig::verify_first_pass_memo finds them again and
// throws on a difference.
struct FirstPassSpots {
std::mutex m;
std::vector<std::pair<std::vector<double>, std::map<int, std::vector<SpotToSave>>>> entries;
+12 -7
View File
@@ -16,6 +16,7 @@
#include <gemmi/solmask.hpp>
#include <gemmi/symmetry.hpp>
#include "../rugnux/ModelGrid.h"
#include "../rugnux/ModelMaskGPU.h"
namespace {
@@ -127,13 +128,13 @@ double DistanceToMaskEdge(const gemmi::Grid<float> &grid, const std::vector<Mode
} // namespace
// The whole mask against gemmi's put_mask_on_grid(), the island step alone against gemmi's on gemmi's
// own pre-island mask (exact), and three repeats bit for bit. The masking itself may differ from gemmi
// only at points that sit at an atom's radius, where the float distance rounds the other way.
// The whole mask against gemmi's put_mask_on_grid() - and so against the CPU path's PutMaskOnGrid() -
// the island step alone against gemmi's on gemmi's own pre-island mask (exact), and three repeats bit
// for bit. The masking itself may differ from gemmi only at points that sit at an atom's radius, where
// the float distance rounds the other way.
TEST_CASE("ModelMaskGPU_MatchesGemmi", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
if (get_gpu_count() == 0)
SKIP("No GPU");
}
CudaStream stream;
bool any_islands = false, any_odd_size = false;
for (const char *cryst : kCrysts) {
@@ -155,6 +156,11 @@ TEST_CASE("ModelMaskGPU_MatchesGemmi", "[ModelValidation][gpu]") {
const int islands = masker.remove_islands(ref);
masker.shrink(ref);
const size_t n = ref.data.size();
// The rigid body's CPU path masks with PutMaskOnGrid, which is gemmi's put_mask_on_grid bit
// for bit - checked here on this grid, so what the GPU is held to below is the CPU path too.
gemmi::Grid<float> cpu = MaskGrid(st, d_min);
PutMaskOnGrid(cpu, st.models[0], OrbitLeaders(cpu, 4), 4);
REQUIRE(std::memcmp(cpu.data.data(), ref.data.data(), n * sizeof(float)) == 0);
any_islands |= islands > 0;
any_odd_size |= ref.nu % 8 != 0 || ref.nv % 8 != 0 || ref.nw % 8 != 0;
@@ -198,9 +204,8 @@ TEST_CASE("ModelMaskGPU_MatchesGemmi", "[ModelValidation][gpu]") {
// The shrink is not implemented: a grid fine enough for it to change anything is refused.
TEST_CASE("ModelMaskGPU_RefusesAGridTheShrinkWouldChange", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
if (get_gpu_count() == 0)
SKIP("No GPU");
}
CudaStream stream;
const gemmi::Structure st = ClusterModel(kCrysts[0]);
const std::vector<ModelMaskOp> ops = MaskOps(*st.find_spacegroup());
+4 -4
View File
@@ -230,7 +230,7 @@ const Case CASES[] = {
TEST_CASE("ModelScaleGPU_FitMatchesGemmi", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
return;
SKIP("No GPU");
for (const Case &c : CASES) {
INFO(c.name);
gemmi::UnitCell cell(c.a, c.b, c.c, c.alpha, c.beta, c.gamma);
@@ -261,7 +261,7 @@ TEST_CASE("ModelScaleGPU_FitMatchesGemmi", "[ModelValidation][gpu]") {
TEST_CASE("ModelScaleGPU_SolventGridMatchesFitModelScale", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
return;
SKIP("No GPU");
for (const Case &c : CASES) {
INFO(c.name);
gemmi::UnitCell cell(c.a, c.b, c.c, c.alpha, c.beta, c.gamma);
@@ -283,7 +283,7 @@ TEST_CASE("ModelScaleGPU_SolventGridMatchesFitModelScale", "[ModelValidation][gp
// per-evaluation fit at that solvent pair.
TEST_CASE("ModelScaleGPU_MatchesGemmiOnAModelsPoints", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
return;
SKIP("No GPU");
const char *crysts[] = {
"CRYST1 40.000 50.000 60.000 90.00 90.00 90.00 P 1 1\n",
"CRYST1 40.000 50.000 60.000 90.00 100.00 90.00 C 1 2 1 4\n",
@@ -319,7 +319,7 @@ TEST_CASE("ModelScaleGPU_MatchesGemmiOnAModelsPoints", "[ModelValidation][gpu]")
TEST_CASE("ModelScaleGPU_Deterministic", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0)
return;
SKIP("No GPU");
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);
+18 -20
View File
@@ -1239,10 +1239,8 @@ TEST_CASE("ModelValidation_RigidBodyReportsTheRotationItApplied", "[ModelValidat
// give the same residuals and the same Jacobian to rounding (float distances on the device, cuFFT for
// FFTW), on every group of the composition tests, both zones, anisotropic atoms included.
TEST_CASE("RigidBodyGPU_MatchesCPU", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
SUCCEED("no GPU");
return;
}
if (get_gpu_count() == 0)
SKIP("No GPU");
Logger logger("RigidBodyGPU_MatchesCPU");
const double q0[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05};
for (const char *cryst : kRigidBodyCrysts) {
@@ -1288,7 +1286,11 @@ TEST_CASE("RigidBodyGPU_MatchesCPU", "[ModelValidation][gpu]") {
}
INFO(cryst << " at " << zone << " A, column " << j << ": relative difference "
<< std::sqrt(diff / norm));
CHECK(std::sqrt(diff) <= 2e-3 * std::sqrt(norm) + 1e-9);
// Measured at up to 5e-5 (the rotation columns, a forward difference of float-gridded
// density; the translation columns agree to about 1e-6). Every column is proportional to
// the overall scale, which the near-tie above can move by 1e-4, so that is the floor of
// an honest bound; 5e-4 leaves room for the difference's own rounding on another card.
CHECK(std::sqrt(diff) <= 5e-4 * std::sqrt(norm) + 1e-9);
}
}
}
@@ -1297,10 +1299,8 @@ TEST_CASE("RigidBodyGPU_MatchesCPU", "[ModelValidation][gpu]") {
// The whole Jacobian on the GPU against the central difference of the GPU's own residuals, with the
// bounds of ModelValidation_RigidBodyJacobianMatchesNumericScaleRefit.
TEST_CASE("RigidBodyGPU_JacobianMatchesNumericScaleRefit", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
SUCCEED("no GPU");
return;
}
if (get_gpu_count() == 0)
SKIP("No GPU");
const double close[6] = {0.0, 0.0, 0.0, 0.05, -0.04, 0.03};
const double rotated[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05};
for (const char *cryst : kRigidBodyCrysts) {
@@ -1330,13 +1330,10 @@ namespace {
}
}
// The GPU fit walks where the CPU fit walks: from the same displaced start, to the same placement
// within a few thousandths of an angstrom.
// The GPU fit walks where the CPU fit walks: from the same displaced start, to the same placement.
TEST_CASE("RigidBodyGPU_FitAgreesWithCPU", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
SUCCEED("no GPU");
return;
}
if (get_gpu_count() == 0)
SKIP("No GPU");
Logger logger("RigidBodyGPU_FitAgreesWithCPU");
for (const char *cryst : {kCryst, kPolarCryst, kRigidBodyCrysts[3]}) {
gemmi::Structure cpu_st = AnisoCluster(cryst), gpu_st = AnisoCluster(cryst);
@@ -1352,17 +1349,18 @@ TEST_CASE("RigidBodyGPU_FitAgreesWithCPU", "[ModelValidation][gpu]") {
rmsd += pc[i].dist_sq(pg[i]) / static_cast<double>(pc.size());
INFO(cryst << ": CPU " << cpu.angle_deg << " deg " << cpu.shift_A << " A, GPU " << gpu.angle_deg << " deg "
<< gpu.shift_A << " A, " << std::sqrt(rmsd) << " A apart");
CHECK(std::sqrt(rmsd) < 2e-3);
// Measured at up to 5e-7 A: the endpoint moves only by what the residuals differ by (a few
// 1e-6 of <Fobs>) over the target's curvature. The bound leaves a factor 20 for another card's
// cuFFT and float rounding, and is still four orders under anything a placement is judged at.
CHECK(std::sqrt(rmsd) < 1e-5);
}
}
// Deterministic: the same fit twice, and on a pool of one engine and of four, gives the same placement
// bit for bit.
TEST_CASE("RigidBodyGPU_Deterministic", "[ModelValidation][gpu]") {
if (get_gpu_count() == 0) {
SUCCEED("no GPU");
return;
}
if (get_gpu_count() == 0)
SKIP("No GPU");
Logger logger("RigidBodyGPU_Deterministic");
const char *cryst = kRigidBodyCrysts[4];
std::vector<std::vector<gemmi::Position>> placed;
+42
View File
@@ -79,3 +79,45 @@ TEST_CASE("Rugnux_Rotation", "[large]") {
reader.Close();
REQUIRE(H5Fget_obj_count(H5F_OBJ_ALL, H5F_OBJ_ALL) == 0);
}
// A rotation run reuses what its first pass found - an indexing probe's validation evidence, each
// frame's spot list, the rotation indexer's outcome - wherever it asks the same question again: in the
// rotation-scale walk's probes and in the canonical pass after them. verify_first_pass_memo recomputes
// at every one of those hits instead and throws where the recomputation differs, so this run passing
// is the check that the reuse is exact. The geometry post-refinement and the scaling it runs inside
// are on because they are what starts the walk.
TEST_CASE("Rugnux_FirstPassMemoMatchesRecomputation", "[large]") {
const auto master = jfjoch_test::LargeDataFile("rotation_master.h5");
if (!master)
SKIP("rotation_master.h5 not available (git-lfs data not pulled)");
RegisterHDF5Filter();
JFJochHDF5Reader reader;
REQUIRE_NOTHROW(reader.ReadFile(*master));
auto dataset = reader.GetDataset();
REQUIRE(dataset);
DiffractionExperiment experiment(dataset->experiment);
IndexingSettings indexing;
indexing.Algorithm(IndexingAlgorithmEnum::Auto);
indexing.RotationIndexing(true);
experiment.ImportIndexingSettings(indexing);
ProcessConfig config;
config.mode = ProcessMode::FullAnalysis;
config.nthreads = default_threads();
config.spot_finding = DiffractionExperiment::DefaultDataProcessingSettings();
config.spot_finding.indexing = true;
config.rotation_indexing = true;
config.two_pass_rotation = true;
config.rotation_postrefine_geometry = true;
config.run_scaling = true;
config.verify_first_pass_memo = true;
Rugnux process(reader, experiment, *dataset->pixel_mask, config);
ProcessResult result;
REQUIRE_NOTHROW(result = process.Run());
CHECK(result.consensus_cell.has_value());
reader.Close();
}