From 07d7174eff02d2b84dec94bfda7c04f91c72b2e4 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Fri, 29 May 2026 14:19:47 +0200 Subject: [PATCH] XtalOptimizer: Accept std::vector> for spots input, to distinguish spots between different images --- image_analysis/IndexAndRefine.cpp | 2 +- image_analysis/RotationIndexer.cpp | 40 ++------- image_analysis/RotationIndexer.h | 2 +- .../geom_refinement/XtalOptimizer.cpp | 82 +++++++++++-------- .../geom_refinement/XtalOptimizer.h | 2 +- tests/XtalOptimizerTest.cpp | 35 ++++---- 6 files changed, 74 insertions(+), 89 deletions(-) diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index f888528e..bffaf3c9 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/image_analysis/IndexAndRefine.cpp @@ -134,7 +134,7 @@ void IndexAndRefine::RefineGeometryIfNeeded(DataMessage &msg, IndexAndRefine::In XtalOptimizerRotationOnly(data, msg.spots, 0.05); break; case GeomRefinementAlgorithmEnum::BeamCenter: - if (XtalOptimizer(data, msg.spots)) { + if (XtalOptimizer(data, {msg.spots})) { outcome.experiment.BeamX_pxl(data.geom.GetBeamX_pxl()) .BeamY_pxl(data.geom.GetBeamY_pxl()); outcome.beam_center_updated = true; diff --git a/image_analysis/RotationIndexer.cpp b/image_analysis/RotationIndexer.cpp index 8ed00d83..38e9a10f 100644 --- a/image_analysis/RotationIndexer.cpp +++ b/image_analysis/RotationIndexer.cpp @@ -11,6 +11,7 @@ RotationIndexer::RotationIndexer(const DiffractionExperiment &x, IndexerThreadPool &indexer) : experiment(x), index_ice_rings(x.GetIndexingSettings().GetIndexIceRings()), + v_(experiment.GetImageNum()), axis_(x.GetGoniometer()), geom_(x.GetDiffractionGeometry()), updated_geom_(geom_), @@ -40,39 +41,8 @@ RotationIndexer::RotationIndexer(const DiffractionExperiment &x, IndexerThreadPo void RotationIndexer::TryIndex() { // Index - std::vector v_sel; - std::vector coords_sel; - if (coords_.size() > max_spots) { - - // Indices into v_ / coords_ - std::vector idx(coords_.size()); - std::iota(idx.begin(), idx.end(), std::size_t{0}); - - // Sort indices by descending intensity - std::partial_sort( - idx.begin(), - idx.begin() + max_spots, - idx.end(), - [this](std::size_t a, std::size_t b) { - return v_[a].intensity > v_[b].intensity; - } - ); - - v_sel.reserve(max_spots); - coords_sel.reserve(max_spots); - - for (std::size_t i = 0; i < max_spots; ++i) { - const std::size_t k = idx[i]; - v_sel.emplace_back(v_[k]); - coords_sel.emplace_back(coords_[k]); - } - } else { - v_sel = v_; - coords_sel = coords_; - } - - auto indexer_result = indexer_.Run(experiment, coords_sel); + auto indexer_result = indexer_.Run(experiment, coords_); if (!indexer_result.lattice.empty() && indexer_result.lattice[0].CalcVolume() > 1.0) { auto sg = experiment.GetGemmiSpaceGroup(); if (sg) { @@ -108,7 +78,7 @@ void RotationIndexer::TryIndex() { if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) data.latt.ReorderMonoclinic(); - if (XtalOptimizer(data, v_sel)) { + if (XtalOptimizer(data, v_)) { indexed_lattice = data.latt; updated_geom_ = data.geom; axis_ = data.axis; @@ -135,12 +105,12 @@ std::optional RotationIndexer::ProcessImage(int64_t image const auto rot = axis_->GetTransformationAngle(angle_deg); if (!indexed_lattice && image >= last_accumulated_image + image_stride) { - v_.reserve(v_.size() + spots.size()); + v_[image].reserve(spots.size()); coords_.reserve(coords_.size() + spots.size()); for (const auto &s: spots) { if (index_ice_rings || !s.ice_ring) { - v_.emplace_back(s); + v_[image].emplace_back(s); coords_.emplace_back(rot * s.ReciprocalCoord(geom_)); } } diff --git a/image_analysis/RotationIndexer.h b/image_analysis/RotationIndexer.h index facda160..2acb2792 100644 --- a/image_analysis/RotationIndexer.h +++ b/image_analysis/RotationIndexer.h @@ -36,7 +36,7 @@ class RotationIndexer { const bool index_ice_rings; float angle_norm_deg = 1.0f; - std::vector v_; + std::vector> v_; std::vector coords_; std::optional axis_; const DiffractionGeometry geom_; diff --git a/image_analysis/geom_refinement/XtalOptimizer.cpp b/image_analysis/geom_refinement/XtalOptimizer.cpp index 4df7a6ae..cefb24cc 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.cpp +++ b/image_analysis/geom_refinement/XtalOptimizer.cpp @@ -504,7 +504,7 @@ CrystalLattice AngleAxisAndLengthsToLattice(const double rod[3], const double le } bool XtalOptimizerInternal(XtalOptimizerData &data, - const std::vector &spots, + const std::vector> &spots, const float tolerance) { try { Coord vec0 = data.latt.Vec0(); @@ -565,50 +565,60 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, const float tolerance_sq = tolerance * tolerance; - // Add residuals for each point - for (const auto &pt: spots) { - if (!data.index_ice_rings && pt.ice_ring) + for (int i = 0; i < spots.size(); i++) { + if (spots[i].empty()) continue; - float angle_rad = 0.0; - - Coord recip = pt.ReciprocalCoord(data.geom); + double angle_rad = 0.0; + std::optional rot_matr; if (data.axis) { - recip = data.axis->GetTransformationAngle(pt.phi) * recip; - angle_rad = pt.phi * M_PI / 180.0; + const float angle_deg = data.axis->GetAngle_deg(i) + data.axis->GetWedge_deg() / 2.0; + angle_rad = angle_deg * M_PI / 180.0; + rot_matr = data.axis->GetTransformationAngle(angle_deg); } - double h_fp = recip * vec0; - double k_fp = recip * vec1; - double l_fp = recip * vec2; + // Add residuals for each point + for (const auto &pt: spots[i]) { + if (!data.index_ice_rings && pt.ice_ring) + continue; - double h = std::round(h_fp); - double k = std::round(k_fp); - double l = std::round(l_fp); + Coord recip = pt.ReciprocalCoord(data.geom); - double norm_sq = (h - h_fp) * (h - h_fp) + (k - k_fp) * (k - k_fp) + (l - l_fp) * (l - l_fp); + if (rot_matr) + recip = rot_matr.value() * recip; - if (norm_sq > tolerance_sq) - continue; + double h_fp = recip * vec0; + double k_fp = recip * vec1; + double l_fp = recip * vec2; - problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new XtalResidual(pt.x, pt.y, - data.geom.GetWavelength_A(), - data.geom.GetPixelSize_mm(), - angle_rad, - h, k, l, - data.crystal_system)), - nullptr, - beam, - &distance_mm, - detector_rot, - rot_vec, - latt_vec0, - latt_vec1, - latt_vec2 - ); + double h = std::round(h_fp); + double k = std::round(k_fp); + double l = std::round(l_fp); + + double norm_sq = (h - h_fp) * (h - h_fp) + (k - k_fp) * (k - k_fp) + (l - l_fp) * (l - l_fp); + + if (norm_sq > tolerance_sq) + continue; + + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction( + new XtalResidual(pt.x, pt.y, + data.geom.GetWavelength_A(), + data.geom.GetPixelSize_mm(), + angle_rad, + h, k, l, + data.crystal_system)), + nullptr, + beam, + &distance_mm, + detector_rot, + rot_vec, + latt_vec0, + latt_vec1, + latt_vec2 + ); + } } if (problem.NumResidualBlocks() < data.min_spots) @@ -722,7 +732,7 @@ bool XtalOptimizerInternal(XtalOptimizerData &data, } } -bool XtalOptimizer(XtalOptimizerData &data, const std::vector &spots) { +bool XtalOptimizer(XtalOptimizerData &data, const std::vector> &spots) { if (!XtalOptimizerInternal(data, spots, 0.3)) return false; XtalOptimizerInternal(data, spots, 0.2); diff --git a/image_analysis/geom_refinement/XtalOptimizer.h b/image_analysis/geom_refinement/XtalOptimizer.h index 9ca306dd..1c0c67dd 100644 --- a/image_analysis/geom_refinement/XtalOptimizer.h +++ b/image_analysis/geom_refinement/XtalOptimizer.h @@ -57,7 +57,7 @@ CrystalLattice AngleAxisAndCellToLattice(const double rod[3], double beta_rad, double gamma_rad); -bool XtalOptimizer(XtalOptimizerData &data, const std::vector &spots); +bool XtalOptimizer(XtalOptimizerData &data, const std::vector> &spots); bool XtalOptimizerRotationOnly(XtalOptimizerData &data, const std::vector &spots, float tolerance); #endif //JFJOCH_XTALOPTIMIZER_H diff --git a/tests/XtalOptimizerTest.cpp b/tests/XtalOptimizerTest.cpp index 1399f0c0..8b90983b 100644 --- a/tests/XtalOptimizerTest.cpp +++ b/tests/XtalOptimizerTest.cpp @@ -39,7 +39,7 @@ TEST_CASE("XtalOptimizer") { xtal_opt.crystal_system = gemmi::CrystalSystem::Triclinic; auto start = std::chrono::high_resolution_clock::now(); - REQUIRE(XtalOptimizer(xtal_opt, spots)); + REQUIRE(XtalOptimizer(xtal_opt, {spots})); auto end = std::chrono::high_resolution_clock::now(); std::cout << "XtalOptimizer took " << std::chrono::duration_cast(end - start).count() << " microseconds" << std::endl; @@ -94,7 +94,7 @@ TEST_CASE("XtalOptimizer_NoBeamCenter") { xtal_opt.max_time = 30.0; auto start = std::chrono::high_resolution_clock::now(); - REQUIRE(XtalOptimizer(xtal_opt, spots)); + REQUIRE(XtalOptimizer(xtal_opt, {spots})); auto end = std::chrono::high_resolution_clock::now(); std::cout << "XtalOptimizer took " << std::chrono::duration_cast(end - start).count() << " microseconds" << std::endl; @@ -145,7 +145,7 @@ TEST_CASE("XtalOptimizer_orthorombic") { xtal_opt.max_time = 30.0; xtal_opt.crystal_system = gemmi::CrystalSystem::Orthorhombic; auto start = std::chrono::high_resolution_clock::now(); - REQUIRE(XtalOptimizer(xtal_opt, spots)); + REQUIRE(XtalOptimizer(xtal_opt, {spots})); auto end = std::chrono::high_resolution_clock::now(); std::cout << "XtalOptimizer took " << std::chrono::duration_cast(end - start).count() << " microseconds" << std::endl; @@ -197,7 +197,7 @@ TEST_CASE("XtalOptimizer_triclinic") { xtal_opt.crystal_system = gemmi::CrystalSystem::Triclinic; xtal_opt.max_time = 36.0; auto start = std::chrono::high_resolution_clock::now(); - REQUIRE(XtalOptimizer(xtal_opt, spots)); + REQUIRE(XtalOptimizer(xtal_opt, {spots})); auto end = std::chrono::high_resolution_clock::now(); std::cout << "XtalOptimizer took " << std::chrono::duration_cast(end - start).count() << " microseconds" << std::endl; @@ -250,7 +250,7 @@ TEST_CASE("XtalOptimizer_tetragonal") { xtal_opt.crystal_system = gemmi::CrystalSystem::Tetragonal; xtal_opt.max_time = 30.0; auto start = std::chrono::high_resolution_clock::now(); - REQUIRE(XtalOptimizer(xtal_opt, spots)); + REQUIRE(XtalOptimizer(xtal_opt, {spots})); auto end = std::chrono::high_resolution_clock::now(); std::cout << "XtalOptimizer took " << std::chrono::duration_cast(end - start).count() << " microseconds" << std::endl; @@ -302,7 +302,7 @@ TEST_CASE("XtalOptimizer_hexagonal") { xtal_opt.crystal_system = gemmi::CrystalSystem::Hexagonal; xtal_opt.max_time = 60.0; auto start = std::chrono::high_resolution_clock::now(); - bool ret = XtalOptimizer(xtal_opt, spots); + bool ret = XtalOptimizer(xtal_opt, {spots}); auto end = std::chrono::high_resolution_clock::now(); std::cout << "XtalOptimizer took " << std::chrono::duration_cast(end - start).count() << " microseconds" << std::endl; @@ -356,7 +356,7 @@ TEST_CASE("XtalOptimizer_hexagonal_unconstrained") { xtal_opt.crystal_system = gemmi::CrystalSystem::Triclinic; xtal_opt.max_time = 30.0; auto start = std::chrono::high_resolution_clock::now(); - REQUIRE(XtalOptimizer(xtal_opt, spots)); + REQUIRE(XtalOptimizer(xtal_opt, {spots})); auto end = std::chrono::high_resolution_clock::now(); std::cout << "XtalOptimizer took " << std::chrono::duration_cast(end - start).count() << " microseconds" << std::endl; @@ -415,7 +415,7 @@ TEST_CASE("XtalOptimizer_cubic") { xtal_opt.crystal_system = gemmi::CrystalSystem::Cubic; xtal_opt.max_time = 30.0; auto start = std::chrono::high_resolution_clock::now(); - REQUIRE(XtalOptimizer(xtal_opt, spots)); + REQUIRE(XtalOptimizer(xtal_opt, {spots})); auto end = std::chrono::high_resolution_clock::now(); std::cout << "XtalOptimizer took " << std::chrono::duration_cast(end - start).count() << " microseconds" << std::endl; @@ -469,7 +469,7 @@ TEST_CASE("XtalOptimizer_monoclinic") { xtal_opt.crystal_system = gemmi::CrystalSystem::Monoclinic; xtal_opt.max_time = 30.0; auto start = std::chrono::high_resolution_clock::now(); - REQUIRE(XtalOptimizer(xtal_opt, spots)); + REQUIRE(XtalOptimizer(xtal_opt, {spots})); auto end = std::chrono::high_resolution_clock::now(); std::cout << "XtalOptimizer took " << std::chrono::duration_cast(end - start).count() << " microseconds" << std::endl; @@ -578,11 +578,13 @@ TEST_CASE("XtalOptimizer_rotation") { .ewald_dist_cutoff = 0.002 }; - std::vector spots; + size_t nimages = 10; + + std::vector> spots(nimages); BraggPrediction prediction; // Predict reflections for images at 0-30 deg. - for (int img = 0; img < 10; ++img) { + for (int img = 0; img < nimages; ++img) { // For a rotated image, per-image lattice is obtained as Multiply(rot.transpose()) const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f; const RotMatrix rot = axis.GetTransformationAngle(angle_deg); @@ -599,7 +601,7 @@ TEST_CASE("XtalOptimizer_rotation") { s.intensity = 1.0f; // minimal positive value s.ice_ring = false; s.indexed = true; - spots.push_back(s); + spots[img].push_back(s); } } @@ -663,11 +665,14 @@ TEST_CASE("XtalOptimizer_refine_rotation_axis") { .ewald_dist_cutoff = 0.002 }; - std::vector spots; + BraggPrediction prediction; + const size_t nimages = 10; + std::vector> spots(nimages); + // Predict reflections for images at 0-30 deg. - for (int img = 0; img < 10; ++img) { + for (int img = 0; img < nimages; ++img) { // For a rotated image, per-image lattice is obtained as Multiply(rot.transpose()) const float angle_deg = axis.GetAngle_deg(img) + axis.GetWedge_deg() / 2.0f; const RotMatrix rot = axis.GetTransformationAngle(angle_deg); @@ -684,7 +689,7 @@ TEST_CASE("XtalOptimizer_refine_rotation_axis") { s.phi = angle_deg; s.ice_ring = false; s.indexed = true; - spots.push_back(s); + spots.at(img).push_back(s); } }