mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 63 Wiki Activity
Files
cdbd0bed5de8a5af816b205559031f079df20b55
Jungfraujoch/image_analysis/IndexAndRefine.cpp
T
leonarski_fandClaude Fable 5 cdbd0bed5d partiality: fix the rotation-prediction wedge, use the stored value in scaling 
BraggPredictionRot halves settings.wedge_deg to the +/- half-wedge of the
partiality erf pair, but IndexAndRefine already passed GetWedge_deg()/2, so the
two /2 compounded to a half-wedge of increment/4 - half the correct Kabsch value
(increment/2, which ScaleOnTheFly's RotationPartiality already used). Pass the
full increment so prediction's partiality matches scaling.

With prediction correct, ScaleOnTheFly now uses the stored r.partiality directly
(the value the reflection was integrated with) rather than recomputing the erf
pair per reflection - recomputing only when scaling overrides the geometry
(-w wedge refinement, --mosaicity, or a scaling wedge override). Output-neutral
on the /data/rotation_test battery (SG/cell/completeness identical, ISa/CC1/2
within run noise on the stable crystals).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-02 14:57:16 +02:00

539 lines
22 KiB
C++
Raw Blame History

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include <cstdlib>
#include "IndexAndRefine.h"
#include "bragg_integration/BraggIntegrate2D.h"
#include "bragg_integration/ProfileIntegrate2D.h"
#include "bragg_integration/CalcISigma.h"
#include "geom_refinement/XtalOptimizer.h"
#include "indexing/AnalyzeIndexing.h"
#include "indexing/FFTIndexer.h"
#include "indexing/MultiLatticeSearch.h"
#include "lattice_search/LatticeSearch.h"
#include "scale_merge/ScaleOnTheFly.h"
IndexAndRefine::IndexAndRefine(const DiffractionExperiment &x, IndexerThreadPool *indexer)
: index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()),
experiment(x),
geom_(x.GetDiffractionGeometry()),
indexer_(indexer),
rotation_indexer_counter(x) {
if (indexer && x.IsRotationIndexing())
rotation_indexer = std::make_unique<RotationIndexer>(x, *indexer);
integration_outcome.resize(x.GetImageNum());
mosaicity.resize(x.GetImageNum(), NAN);
scale_cc.resize(x.GetImageNum(), 0);
unit_cells.resize(x.GetImageNum());
}
std::optional<float> IndexAndRefine::RotationAngle(int64_t image) const {
// Mid-exposure rotation angle for image index `image`, matching the angle used for prediction.
if (const auto g = experiment.GetGoniometer())
return g->GetAngle_deg(static_cast<float>(image)) + g->GetWedge_deg() / 2.0f;
return std::nullopt;
}
void IndexAndRefine::AddImageToRotationIndexer(DataMessage &msg) {
if (rotation_indexer)
rotation_indexer->ProcessImage(msg.number, msg.spots, RotationAngle(msg.number));
}
IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetryRotation(DataMessage &msg) {
IndexingOutcome outcome(experiment);
if (!rotation_indexer)
return outcome;
auto result = rotation_indexer->GetLattice();
if (!result.has_value()) {
auto rot_cnt = rotation_indexer_counter.Process(msg.number);
if (rot_cnt.first)
rotation_indexer->ProcessImage(msg.number, msg.spots, RotationAngle(msg.number));
if (rot_cnt.second)
rotation_indexer->RunIndexing();
result = rotation_indexer->GetLattice();
}
if (result.has_value()) {
// For rotation indexing, indexing rate is calculated only for frames, where "global" rotation indexing solution was found
msg.indexing_result = false;
// get rotated lattice
auto gon = result->axis;
if (gon) {
const float angle_deg = gon->GetAngle_deg(msg.number) + gon->GetWedge_deg() / 2.0f;
const auto rot_to_image = gon->GetTransformationAngle(-angle_deg);
outcome.lattice_candidate = result->lattice.Multiply(rot_to_image);
outcome.extra_lattice_candidates.reserve(result->extra_lattices.size());
for (const auto &el : result->extra_lattices)
outcome.extra_lattice_candidates.push_back(el.Multiply(rot_to_image));
}
outcome.experiment.BeamX_pxl(result->geom.GetBeamX_pxl())
.BeamY_pxl(result->geom.GetBeamY_pxl())
.DetectorDistance_mm(result->geom.GetDetectorDistance_mm())
.PoniRot1_rad(result->geom.GetPoniRot1_rad())
.PoniRot2_rad(result->geom.GetPoniRot2_rad())
.Goniometer(result->axis);
outcome.symmetry.centering = result->search_result.centering;
outcome.symmetry.niggli_class = result->search_result.niggli_class;
outcome.symmetry.crystal_system = result->search_result.system;
}
return outcome;
}
IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetry(DataMessage &msg) {
auto indexing_start_time = std::chrono::steady_clock::now();
IndexingOutcome outcome(experiment);
// Convert input spots to reciprocal space
std::vector<Coord> recip;
recip.reserve(msg.spots.size());
for (const auto &i: msg.spots) {
if (index_ice_rings || !i.ice_ring)
recip.push_back(i.ReciprocalCoord(geom_));
}
auto indexer_result = indexer_->Run(experiment, recip);
if (indexer_result.executed)
msg.indexing_result = false;
if (!indexer_result.lattice.empty()) {
auto latt = indexer_result.lattice[0];
if (latt.CalcVolume() > 1.0) {
auto sg = experiment.GetGemmiSpaceGroup();
const auto algorithm = experiment.GetIndexingAlgorithm();
const bool de_novo = (algorithm == IndexingAlgorithmEnum::FFT
|| algorithm == IndexingAlgorithmEnum::FFTW);
// If space group and cell provided => enforce that symmetry in refinement.
// If not => detect the symmetry from the lattice.
if (sg && experiment.GetUnitCell()) {
outcome.symmetry = LatticeMessage{
.centering = sg->centring_type(),
.niggli_class = 0,
.crystal_system = sg->crystal_system()
};
// Place every frame's cell in ONE consistent setting for the whole dataset:
// mixed axis orders (e.g. [78,78,38] vs [38,78,78]) index the same reflection
// as different HKLs and cannot be merged. LatticeSearch gives the conventional
// setting when its detected symmetry agrees with the user's space group. On
// noisy frames it can instead pick an alternative Bravais setting (e.g. the
// sqrt2 C-centred description of a primitive tetragonal cell,
// [78,78,38]->[110,111,38]); there:
// - FFBIDX already returns the reference setting (c-last), consistent with the
// conventional frames, so its raw lattice is safe -> use it (FFBIDX neutral);
// - de-novo indexers (FFT/FFTW) return a Niggli-primitive cell with a DIFFERENT
// axis order (c-first) that would corrupt the merge -> reject the frame.
// niggli_class is left unassigned (0): it needs the primitive cell incl.
// centering, which LatticeSearch cannot recover from a (possibly centred, e.g.
// C2) user cell. A proper primitive-cell indexing path (CrystFEL-style) is deferred.
auto sym_result = LatticeSearch(latt);
if (sym_result.system == sg->crystal_system())
outcome.lattice_candidate = sym_result.conventional;
else if (!de_novo)
outcome.lattice_candidate = latt;
// else: de-novo + symmetry mismatch -> leave unset, frame is not indexed
} else {
auto sym_result = LatticeSearch(latt);
outcome.symmetry = LatticeMessage{
.centering = sym_result.centering,
.niggli_class = sym_result.niggli_class,
.crystal_system = sym_result.system
};
outcome.lattice_candidate = sym_result.conventional;
}
// Multi-lattice search for stills: store rotations that map the reference
// lattice to each accepted extra lattice. Candidates are materialized later
// in RefineGeometryIfNeeded so they're rooted in the refined (and, for
// monoclinic, reordered) main lattice.
if (outcome.lattice_candidate && indexer_result.lattice.size() > 1) {
auto ml_latt = MultiLatticeSearch(indexer_result.lattice);
for (auto &ml : ml_latt) {
if (outcome.extra_lattice_rotations.size() >= experiment.GetIndexingSettings().GetMaxExtraLattices())
break;
outcome.extra_lattice_rotations.push_back(ml.rotation_vector);
RotMatrix rot(ml.rotation_vector.Length(), ml.rotation_vector.Normalize());
outcome.extra_lattice_candidates.push_back(outcome.lattice_candidate->Multiply(rot));
}
}
}
}
auto indexing_end_time = std::chrono::steady_clock::now();
msg.indexing_time_s = std::chrono::duration<float>(indexing_end_time - indexing_start_time).count();
return outcome;
}
void IndexAndRefine::RefineGeometryIfNeeded(DataMessage &msg, IndexAndRefine::IndexingOutcome &outcome) {
if (!outcome.lattice_candidate)
return;
auto start_time = std::chrono::steady_clock::now();
XtalOptimizerData data{
.geom = outcome.experiment.GetDiffractionGeometry(),
.latt = *outcome.lattice_candidate,
.crystal_system = outcome.symmetry.crystal_system,
.min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
.refine_beam_center = true,
.refine_distance_mm = false,
.refine_detector_angles = false,
.refine_unit_cell = !experiment.IsRotationIndexing(),
.max_time = 0.04 // 40 ms is max allowed time for the operation
};
if (outcome.symmetry.crystal_system == gemmi::CrystalSystem::Trigonal)
data.crystal_system = gemmi::CrystalSystem::Hexagonal;
switch (experiment.GetIndexingSettings().GetGeomRefinementAlgorithm()) {
case GeomRefinementAlgorithmEnum::None:
break;
case GeomRefinementAlgorithmEnum::OrientationOnly:
XtalOptimizerRotationOnly(data, msg.spots, 0.2);
XtalOptimizerRotationOnly(data, msg.spots, 0.1);
XtalOptimizerRotationOnly(data, msg.spots, 0.05);
break;
case GeomRefinementAlgorithmEnum::BeamCenter:
if (XtalOptimizer(data, {msg.spots})) {
outcome.experiment.BeamX_pxl(data.geom.GetBeamX_pxl())
.BeamY_pxl(data.geom.GetBeamY_pxl());
outcome.beam_center_updated = true;
}
break;
}
outcome.lattice_candidate = data.latt;
if (outcome.symmetry.crystal_system == gemmi::CrystalSystem::Monoclinic)
outcome.lattice_candidate->ReorderMonoclinic();
// Rebuild extra-lattice candidates from the refined (and possibly reordered) main
// lattice so they share its cell and obtuse-beta convention.
if (!outcome.extra_lattice_rotations.empty()) {
outcome.extra_lattice_candidates.clear();
outcome.extra_lattice_candidates.reserve(outcome.extra_lattice_rotations.size());
for (const auto &rv : outcome.extra_lattice_rotations) {
RotMatrix rot(rv.Length(), rv.Normalize());
outcome.extra_lattice_candidates.push_back(outcome.lattice_candidate->Multiply(rot));
}
}
// Quick orientation-only refinement of extra lattices (stills path).
// Cell, beam center, detector geometry are taken from the first lattice.
if (!experiment.IsRotationIndexing() && !outcome.extra_lattice_candidates.empty()) {
for (auto &el : outcome.extra_lattice_candidates) {
XtalOptimizerData data_extra{
.geom = data.geom,
.latt = el,
.crystal_system = data.crystal_system,
.min_spots = experiment.GetIndexingSettings().GetViableCellMinSpots(),
.refine_beam_center = false,
.refine_distance_mm = false,
.refine_detector_angles = false,
.refine_unit_cell = false,
.refine_rotation_axis = false,
.index_ice_rings = experiment.GetIndexingSettings().GetIndexIceRings(),
.max_time = 0.02
};
XtalOptimizerRotationOnly(data_extra, msg.spots, 0.1);
el = data_extra.latt;
}
}
if (outcome.beam_center_updated) {
msg.beam_corr_x = data.beam_corr_x;
msg.beam_corr_y = data.beam_corr_y;
}
auto end_time = std::chrono::steady_clock::now();
msg.refinement_time_s = std::chrono::duration_cast<std::chrono::duration<double>>(end_time - start_time).count();
}
void IndexAndRefine::QuickPredictAndIntegrate(DataMessage &msg,
const SpotFindingSettings &spot_finding_settings,
const CompressedImage &image,
BraggPrediction &prediction,
const IndexAndRefine::IndexingOutcome &outcome) {
if (!outcome.lattice_candidate)
return;
CrystalLattice latt = outcome.lattice_candidate.value();
if (rotation_indexer) {
// Use moving average for mosaicity and profile_radius (also add beam center later)
if (msg.mosaicity_deg)
msg.mosaicity_deg = rotation_parameters.Mosaicity(msg.mosaicity_deg.value());
if (msg.profile_radius) {
msg.profile_radius = rotation_parameters.ProfileRadius(msg.profile_radius.value());
}
}
float ewald_dist_cutoff = 0.001f;
if (msg.profile_radius)
ewald_dist_cutoff = msg.profile_radius.value() * 2.0f;
if (experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA())
ewald_dist_cutoff = experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA().value() * 3.0f;
float wedge_deg = 0.0f;
float mos_deg = 0.1f;
if (experiment.GetGoniometer().has_value()) {
// Full oscillation wedge of one frame; BraggPredictionRot halves it to the +/- half-wedge of the
// partiality erf pair (Kabsch). Passing the full increment gives a half-wedge of increment/2 -
// matching ScaleOnTheFly's RotationPartiality, so the predicted partiality is used directly there.
wedge_deg = experiment.GetGoniometer()->GetWedge_deg();
if (msg.mosaicity_deg) {
mos_deg = msg.mosaicity_deg.value();
mosaicity[msg.number] = mos_deg;
}
}
IntegrationOutcome i_outcome{
.geom = outcome.experiment.GetDiffractionGeometry(),
.latt = latt,
.mosaicity_deg = mos_deg,
.image_scale_b_factor_Ang2 = msg.image_scale_b_factor,
.image_scale_cc = msg.image_scale_cc,
};
const BraggPredictionSettings settings_prediction{
.high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
.ewald_dist_cutoff = ewald_dist_cutoff,
.max_hkl = 100,
// Centering is a hypothesis to confirm, not assume: with no user-fixed space group, predict
// in P so the centering-absent reflections are integrated and the space-group search can
// confirm or disprove centering (and catch a missed superstructure). A user-fixed space
// group is trusted, so reject its absences here.
.centering = experiment.GetGemmiSpaceGroup().has_value() ? outcome.symmetry.centering : 'P',
.wedge_deg = std::fabs(wedge_deg),
.mosaicity_deg = std::fabs(mos_deg),
// FWHM -> sigma; 0 when monochromatic, leaving the prediction unchanged.
.bandwidth_sigma = experiment.GetBandwidthFWHM().value_or(0.0f) / 2.3548f
};
// Predict, then integrate with the selected integrator (box-sum or profile-fit).
auto pred_start_time = std::chrono::steady_clock::now();
auto nrefl = prediction.Calc(outcome.experiment, latt, settings_prediction);
auto pred_end_time = std::chrono::steady_clock::now();
msg.bragg_prediction_time_s = std::chrono::duration<float>(pred_end_time - pred_start_time).count();
auto integration_start_time = std::chrono::steady_clock::now();
i_outcome.reflections =
outcome.experiment.GetBraggIntegrationSettings().GetIntegrator() == IntegratorMode::BoxSum
? BraggIntegrate2D(outcome.experiment, image, prediction.GetReflections(), nrefl, msg.number)
: ProfileIntegrate2D(outcome.experiment, image, prediction.GetReflections(), nrefl, msg.number);
msg.integrated_reflections = i_outcome.reflections.size();
auto integration_end_time = std::chrono::steady_clock::now();
msg.integration_time_s = std::chrono::duration<float>(integration_end_time - integration_start_time).count();
constexpr size_t kMaxReflections = 10000;
if (i_outcome.reflections.size() > kMaxReflections) {
// Keep only smallest d (highest resolution)
std::nth_element(i_outcome.reflections.begin(),
i_outcome.reflections.begin() + static_cast<long>(kMaxReflections),
i_outcome.reflections.end(),
[](const Reflection& a, const Reflection& b) {
return a.d < b.d;
});
i_outcome.reflections.resize(kMaxReflections);
// Optional: make output ordered by d (nice for downstream / debugging)
std::sort(i_outcome.reflections.begin(), i_outcome.reflections.end(),
[](const Reflection& a, const Reflection& b) { return a.d < b.d; });
}
CalcISigma(msg, i_outcome.reflections);
CalcWilsonBFactor(msg, i_outcome.reflections);
ScaleImage(msg, i_outcome);
// Copy reflections to outgoing message
msg.reflections = i_outcome.reflections;
{
const std::unique_lock ul(reflections_mutex);
integration_outcome[msg.number] = std::move(i_outcome);
}
}
void IndexAndRefine::ProcessImage(DataMessage &msg,
const SpotFindingSettings &spot_finding_settings,
const CompressedImage &image,
BraggPrediction &prediction) {
if (!indexer_ || !spot_finding_settings.indexing)
return;
IndexingOutcome outcome(experiment);
if (rotation_indexer)
outcome = DetermineLatticeAndSymmetryRotation(msg);
else
outcome = DetermineLatticeAndSymmetry(msg);
if (!outcome.lattice_candidate)
return;
if (experiment.GetIndexingSettings().GetGeomRefinementAlgorithm() != GeomRefinementAlgorithmEnum::None)
RefineGeometryIfNeeded(msg, outcome);
if (!outcome.lattice_candidate.has_value())
return;
if (!AnalyzeIndexing(msg, outcome.experiment, *outcome.lattice_candidate, outcome.extra_lattice_candidates))
return;
{
std::unique_lock ul(reflections_mutex);
unit_cells[msg.number] = outcome.lattice_candidate->GetUnitCell();
}
msg.lattice_type = outcome.symmetry;
if (spot_finding_settings.quick_integration)
QuickPredictAndIntegrate(msg, spot_finding_settings, image, prediction, outcome);
}
std::optional<RotationIndexerResult> IndexAndRefine::FinalizeRotationIndexing() {
if (rotation_indexer) {
if (const auto latt = rotation_indexer->GetLattice())
return latt;
rotation_indexer->RunIndexing();
return rotation_indexer->GetLattice();
}
return {};
}
IndexAndRefine &IndexAndRefine::ReferenceIntensities(std::vector<MergedReflection> &reference) {
scaling_engine = std::make_unique<ScaleOnTheFly>(experiment, reference);
return *this;
}
void IndexAndRefine::ScaleImage(DataMessage &msg, IntegrationOutcome& outcome) {
if (!scaling_engine)
return;
auto scaling_start_time = std::chrono::steady_clock::now();
scaling_engine->Scale(outcome);
auto scaling_end_time = std::chrono::steady_clock::now();
scale_cc[msg.number] = outcome.image_scale_cc.value_or(NAN);
msg.image_scale_b_factor = outcome.image_scale_b_factor_Ang2;
msg.image_scale_factor = outcome.image_scale_g;
msg.image_scale_mosaicity = outcome.mosaicity_deg;
msg.image_scale_cc = outcome.image_scale_cc;
msg.image_scale_time_s = std::chrono::duration<float>(scaling_end_time - scaling_start_time).count();
}
ScalingResult IndexAndRefine::ScaleAllImages(const std::vector<MergedReflection> &reference, size_t nthreads) {
ScaleOnTheFly scaling(experiment, reference);
scaling.Scale(integration_outcome, nthreads);
scale_cc.resize(integration_outcome.size());
for (int i = 0; i < integration_outcome.size(); i++)
scale_cc.at(i) = integration_outcome[i].image_scale_cc.value_or(NAN);
return ScalingResult(integration_outcome);
}
const std::vector<float> &IndexAndRefine::GetImageCC() const {
return scale_cc;
}
const std::vector<std::optional<UnitCell> > & IndexAndRefine::GetUnitCells() const {
return unit_cells;
}
std::optional<UnitCell> IndexAndRefine::GetConsensusUnitCell() const {
const auto dist_tolerance = experiment.GetIndexingSettings().GetUnitCellDistTolerance();
const auto angle_tolerance = experiment.GetIndexingSettings().GetUnitCellAngleTolerance_deg();
if (rotation_indexer) {
auto result = rotation_indexer->GetLattice();
if (!result)
return {};
return result->lattice.GetUnitCell();
}
std::vector<UnitCell> cells;
{
std::unique_lock ul(reflections_mutex);
cells.reserve(unit_cells.size());
for (const auto &cell: unit_cells) {
if (cell && cell->is_finite())
cells.emplace_back(*cell);
}
}
if (cells.empty())
return {};
if (experiment.GetUnitCell()) {
std::vector<UnitCell> accepted;
accepted.reserve(cells.size());
for (const auto &cell: cells) {
if (cell.is_close(*experiment.GetUnitCell(), dist_tolerance, angle_tolerance))
accepted.emplace_back(cell);
}
return MeanUnitCell(accepted);
}
size_t best_count = 0;
UnitCell best_reference{};
for (const auto &ref: cells) {
size_t count = 0;
for (const auto &cell: cells) {
if (cell.is_close(ref, dist_tolerance, angle_tolerance))
++count;
}
if (count > best_count) {
best_count = count;
best_reference = ref;
}
}
if (best_count == 0)
return {};
std::vector<UnitCell> accepted;
accepted.reserve(best_count);
for (const auto &cell: cells) {
if (cell.is_close(best_reference, dist_tolerance, angle_tolerance))
accepted.emplace_back(cell);
}
return MeanUnitCell(accepted);
}
std::vector<IntegrationOutcome> &IndexAndRefine::GetIntegrationOutcome() {
return integration_outcome;
}
const std::vector<IntegrationOutcome> &IndexAndRefine::GetIntegrationOutcome() const {
return integration_outcome;
}
void IndexAndRefine::ForceRotationIndexerLattice(const CrystalLattice &lattice) {
if (rotation_indexer)
rotation_indexer->ForceLattice(lattice);
}
Reference in New Issue View Git Blame Copy Permalink