diff --git a/image_analysis/geom_refinement/GeometryRefiner.cpp b/image_analysis/geom_refinement/GeometryRefiner.cpp index bebb5a41f..c1835b121 100644 --- a/image_analysis/geom_refinement/GeometryRefiner.cpp +++ b/image_analysis/geom_refinement/GeometryRefiner.cpp @@ -151,8 +151,8 @@ struct BeamBlock { double m00 = 1, m01 = 0, m10 = 0, m11 = 1; // DetectorOrientation::Matrix(), as XtalResidual takes it double pixel = 1; - template - void Poni(const T *const beam, const T &distance_mm, const T *const detector_rot, T *poni) const { + template + void Poni(const T *const beam, const T &distance_mm, const R *const detector_rot, T *poni) const { if (!direct) { poni[0] = beam[0]; poni[1] = beam[1]; @@ -167,42 +167,98 @@ struct BeamBlock { } }; -// XtalResidual with the beam block read through BeamBlock. -struct BeamXtalResidual { +// The residual of one spot (XtalResidual) on the blocks this refinement can move: the beam (read through +// BeamBlock), the distance, the frame's orientation and the cell lengths, plus the module's move (dx, dy, rot) +// where the spot's module is placed. The rotation axis and the cell angles are never refined here, and the +// tilt only on request (TiltedSpotResidual): they enter as plain numbers, not as blocks held constant, +// because Ceres sizes its dual numbers from the declared blocks - 17 to 20 derivatives per number, 6 to 8 of +// them always zero. The arithmetic on the values is the same either way. +// +// A placed module: moving the spot by R(rot) u - u + (dx, dy), with (ux, uy) its offset from the module +// centre, shifts its offset from the beam exactly as moving the beam the other way does, so the beam handed +// to XtalResidual carries the move and everything else is the residual of an unmoved spot. +struct SpotResidual { XtalResidual residual; BeamBlock beam_block; + const double *detector_rot; // the tilt where it is held + const double *rotation_axis; + const double *cell_ang; + double ux = 0, uy = 0; + + template + bool Evaluate(const T *const beam, const T *const distance_mm, const R *const tilt, const T *const p0, + const T *const p1, const T *const move, T *out) const { + using std::cos; + using std::sin; + T b[2]; + beam_block.Poni(beam, distance_mm[0], tilt, b); + if (move) { + const T c = cos(move[2]), s = sin(move[2]); + b[0] = b[0] - ((c - T(1)) * T(ux) - s * T(uy) + move[0]); + b[1] = b[1] - (s * T(ux) + (c - T(1)) * T(uy) + move[1]); + } + const R axis[3] = {R(rotation_axis[0]), R(rotation_axis[1]), R(rotation_axis[2])}; + T recip_obs[3]; + residual.ObservedRecip(b, distance_mm[0], tilt, axis, recip_obs); + + const T p2[3] = {T(cell_ang[0]), T(cell_ang[1]), T(cell_ang[2])}; + T recip_unrot[3]; + residual.PredictedRecipUnrot(p1, p2, recip_unrot); + + T e_pred_recip[3]; + const AngleAxisRotator rot_p0(p0); + rot_p0.Rotate(recip_unrot, e_pred_recip); + + out[0] = recip_obs[0] - e_pred_recip[0]; + out[1] = recip_obs[1] - e_pred_recip[1]; + out[2] = recip_obs[2] - e_pred_recip[2]; + return true; + } + + // beam_distance: the beam (x, y) and the distance, as one block. + template + bool operator()(const T *const beam_distance, const T *const p0, const T *const p1, T *out) const { + return Evaluate(beam_distance, beam_distance + 2, detector_rot, p0, p1, static_cast(nullptr), out); + } template - bool operator()(const T *const beam, const T *const distance_mm, const T *const detector_rot, - const T *const rotation_axis, const T *const p0, const T *const p1, const T *const p2, + bool operator()(const T *const beam_distance, const T *const p0, const T *const p1, const T *const move, T *out) const { - T poni[2]; - beam_block.Poni(beam, distance_mm[0], detector_rot, poni); - return residual(poni, distance_mm, detector_rot, rotation_axis, p0, p1, p2, out); + return Evaluate(beam_distance, beam_distance + 2, detector_rot, p0, p1, move, out); } }; -// XtalResidual for a spot on a module that is being placed: block `move` = (dx, dy, rot) of the module, -// (ux, uy) the spot's offset from the module centre. Moving the spot by R(rot) u - u + (dx, dy) shifts its -// offset from the beam exactly as moving the beam the other way does, so the beam handed to XtalResidual -// carries the move and everything else is the residual of an unmoved spot. -struct ModuleXtalResidual { - XtalResidual residual; - BeamBlock beam_block; - double ux, uy; +// The same with the tilt refined, as a block after the beam and distance. +struct TiltedSpotResidual : SpotResidual { + template + bool operator()(const T *const beam_distance, const T *const tilt, const T *const p0, const T *const p1, + T *out) const { + return Evaluate(beam_distance, beam_distance + 2, tilt, p0, p1, static_cast(nullptr), out); + } template - bool operator()(const T *const beam, const T *const distance_mm, const T *const detector_rot, - const T *const rotation_axis, const T *const p0, const T *const p1, const T *const p2, + bool operator()(const T *const beam_distance, const T *const tilt, const T *const p0, const T *const p1, const T *const move, T *out) const { - using std::cos; - using std::sin; - T poni[2]; - beam_block.Poni(beam, distance_mm[0], detector_rot, poni); - const T c = cos(move[2]), s = sin(move[2]); - const T moved_beam[2] = {poni[0] - ((c - T(1)) * T(ux) - s * T(uy) + move[0]), - poni[1] - (s * T(ux) + (c - T(1)) * T(uy) + move[1])}; - return residual(moved_beam, distance_mm, detector_rot, rotation_axis, p0, p1, p2, out); + return Evaluate(beam_distance, beam_distance + 2, tilt, p0, p1, move, out); + } +}; + +// The residual of a fit that moves only the crystal orientations (refine_beam_distance off, no module +// placed): the observed reciprocal vector and the unrotated prediction are then constants of the spot, taken +// once in plain numbers, and only the orientation's three derivatives are carried. +struct OrientationResidual { + double recip_obs[3]; + double recip_unrot[3]; + + template + bool operator()(const T *const p0, T *out) const { + T e_pred_recip[3]; + const AngleAxisRotator rot_p0(p0); + rot_p0.Rotate(recip_unrot, e_pred_recip); + out[0] = recip_obs[0] - e_pred_recip[0]; + out[1] = recip_obs[1] - e_pred_recip[1]; + out[2] = recip_obs[2] - e_pred_recip[2]; + return true; } }; @@ -252,9 +308,13 @@ GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_ge const gemmi::CrystalSystem system = MapSystem(settings.crystal_system); // Shared geometry parameter blocks (refined). The rotation axis is constant, and so is the detector tilt - // unless settings.refine_tilt. - double beam[2] = {nominal_geom.GetBeamX_pxl(), nominal_geom.GetBeamY_pxl()}; - double distance_mm = nominal_geom.GetDetectorDistance_mm(); + // unless settings.refine_tilt. The beam and the distance are one block of three: every block refined + // beside it has three parameters too (orientations, cell lengths, module moves), so Ceres eliminates the + // orientations with its fixed-size Schur code rather than the general one (1.7x faster on one thread). + double beam_distance[3] = {nominal_geom.GetBeamX_pxl(), nominal_geom.GetBeamY_pxl(), + nominal_geom.GetDetectorDistance_mm()}; + double *const beam = beam_distance; + double &distance_mm = beam_distance[2]; double detector_rot[2] = {nominal_geom.GetPoniRot1_rad(), nominal_geom.GetPoniRot2_rad()}; double rot_vec[3] = {1.0, 0.0, 0.0}; const double rot3 = nominal_geom.GetPoniRot3_rad(); @@ -344,6 +404,11 @@ GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_ge if (!refined[m]) move[m] = {0, 0, 0}; + // Where nothing but the crystal orientations moves, each spot's residual is reduced to them + // (OrientationResidual). + const bool orientation_only = !settings.refine_beam_distance + && std::none_of(refined.begin(), refined.end(), [](bool r) { return r; }); + auto place = [&](const GeomRefineSpot &s, double &x, double &y) { if (s.module >= 0 && static_cast(s.module) < modules.size()) PlaceOnModule(modules[s.module], move[s.module].data(), s.x, s.y, x, y); @@ -360,6 +425,8 @@ GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_ge return 0.3 - (0.3 - 0.1) * t; }; + ceres::CauchyLoss loss(loss_scale); + // Whether the last solve that ran produced a usable solution; folded into result.ok below, // which otherwise reports success on spot and frame counts alone. bool solve_usable = false; @@ -376,7 +443,10 @@ GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_ge const double tol = round_tol(round); const double tol_sq = tol * tol; - ceres::Problem problem; + // One robust loss for every spot, kept here rather than one allocated (and freed) per spot. + ceres::Problem::Options problem_options; + problem_options.loss_function_ownership = ceres::DO_NOT_TAKE_OWNERSHIP; + ceres::Problem problem(problem_options); auto ordering = std::make_shared(); int spots_used = 0, frames_used = 0; @@ -396,30 +466,44 @@ GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_ge if (dn > tol_sq) continue; - if (s.module >= 0 && static_cast(s.module) < modules.size() && refined[s.module]) { + const XtalResidual xr(x, y, lambda, pixel, cos_rot3, sin_rot3, 0.0, h, k, l, system, orientation); + if (orientation_only) { + auto *r = new OrientationResidual; + double b[2]; + beam_block.Poni(beam, distance_mm, detector_rot, b); + xr.ObservedRecip(b, distance_mm, detector_rot, rot_vec, r->recip_obs); + xr.PredictedRecipUnrot(cell_len, cell_ang, r->recip_unrot); + problem.AddResidualBlock(new ceres::AutoDiffCostFunction(r), + &loss, orient[i].data()); + } else if (s.module >= 0 && static_cast(s.module) < modules.size() && refined[s.module]) { + // The spot as found: the module's move is a block of the residual. const auto &m = modules[s.module]; - problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new ModuleXtalResidual{XtalResidual(s.x, s.y, lambda, pixel, cos_rot3, sin_rot3, - 0.0, h, k, l, system, orientation), - beam_block, s.x - m.x0, s.y - m.y0}), - new ceres::CauchyLoss(loss_scale), - beam, &distance_mm, detector_rot, rot_vec, orient[i].data(), cell_len, cell_ang, - move[s.module].data()); + const SpotResidual sr{XtalResidual(s.x, s.y, lambda, pixel, cos_rot3, sin_rot3, 0.0, + h, k, l, system, orientation), + beam_block, detector_rot, rot_vec, cell_ang, s.x - m.x0, s.y - m.y0}; + if (refine_tilt) + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction( + new TiltedSpotResidual{sr}), + &loss, + beam_distance, detector_rot, orient[i].data(), cell_len, move[s.module].data()); + else + problem.AddResidualBlock( + new ceres::AutoDiffCostFunction(new SpotResidual{sr}), + &loss, + beam_distance, orient[i].data(), cell_len, move[s.module].data()); } else if (refine_tilt) problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new BeamXtalResidual{XtalResidual(x, y, lambda, pixel, cos_rot3, sin_rot3, 0.0, - h, k, l, system, orientation), beam_block}), - new ceres::CauchyLoss(loss_scale), - beam, &distance_mm, detector_rot, rot_vec, orient[i].data(), cell_len, cell_ang); + new ceres::AutoDiffCostFunction( + new TiltedSpotResidual{{xr, beam_block, detector_rot, rot_vec, cell_ang}}), + &loss, + beam_distance, detector_rot, orient[i].data(), cell_len); else problem.AddResidualBlock( - new ceres::AutoDiffCostFunction( - new XtalResidual(x, y, lambda, pixel, cos_rot3, sin_rot3, 0.0, - h, k, l, system, orientation)), - new ceres::CauchyLoss(loss_scale), - beam, &distance_mm, detector_rot, rot_vec, orient[i].data(), cell_len, cell_ang); + new ceres::AutoDiffCostFunction( + new SpotResidual{xr, beam_block, detector_rot, rot_vec, cell_ang}), + &loss, + beam_distance, orient[i].data(), cell_len); ++used_in_frame; ++spots_used; } @@ -447,9 +531,11 @@ GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_ge new ceres::AutoDiffCostFunction(reg), nullptr, cell_len); } - // Shared blocks go in the second elimination group; constant ones are stripped by Ceres. - for (double *b : {beam, &distance_mm, detector_rot, rot_vec, cell_len, cell_ang}) - ordering->AddElementToGroup(b, 1); + // Shared blocks go in the second elimination group (those the problem has; constant ones are + // stripped by Ceres). + for (double *b : {beam_distance, detector_rot, cell_len}) + if (problem.HasParameterBlock(b)) + ordering->AddElementToGroup(b, 1); // The refined modules' moves: bounded to what a misplaced module plausibly is, and held to // zero in sum by the gauge (ModuleGauge). A 1 px mean shift costs as much there as moving @@ -494,21 +580,26 @@ GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_ge // Beam + distance bounds around the nominal value; cell lengths loosely bounded (the // regularizer does the anchoring). - problem.SetParameterLowerBound(beam, 0, beam0[0] - settings.beam_range_px); - problem.SetParameterUpperBound(beam, 0, beam0[0] + settings.beam_range_px); - problem.SetParameterLowerBound(beam, 1, beam0[1] - settings.beam_range_px); - problem.SetParameterUpperBound(beam, 1, beam0[1] + settings.beam_range_px); - problem.SetParameterLowerBound(&distance_mm, 0, - std::max(1.0, nominal_geom.GetDetectorDistance_mm() - settings.distance_range_mm)); - problem.SetParameterUpperBound(&distance_mm, 0, - nominal_geom.GetDetectorDistance_mm() + settings.distance_range_mm); - for (int j = 0; j < 3; ++j) { - problem.SetParameterLowerBound(cell_len, j, 0.7 * cell_len0[j]); - problem.SetParameterUpperBound(cell_len, j, 1.3 * cell_len0[j]); - } + if (!orientation_only) { + problem.SetParameterLowerBound(beam_distance, 0, beam0[0] - settings.beam_range_px); + problem.SetParameterUpperBound(beam_distance, 0, beam0[0] + settings.beam_range_px); + problem.SetParameterLowerBound(beam_distance, 1, beam0[1] - settings.beam_range_px); + problem.SetParameterUpperBound(beam_distance, 1, beam0[1] + settings.beam_range_px); + problem.SetParameterLowerBound(beam_distance, 2, + std::max(1.0, nominal_geom.GetDetectorDistance_mm() - settings.distance_range_mm)); + problem.SetParameterUpperBound(beam_distance, 2, + nominal_geom.GetDetectorDistance_mm() + settings.distance_range_mm); + for (int j = 0; j < 3; ++j) { + problem.SetParameterLowerBound(cell_len, j, 0.7 * cell_len0[j]); + problem.SetParameterUpperBound(cell_len, j, 1.3 * cell_len0[j]); + } - if (!settings.refine_distance) - problem.SetParameterBlockConstant(&distance_mm); + if (!settings.refine_beam_distance) { + problem.SetParameterBlockConstant(beam_distance); + problem.SetParameterBlockConstant(cell_len); + } else if (!settings.refine_distance) + problem.SetManifold(beam_distance, new ceres::SubsetManifold(3, {2})); + } if (refine_tilt) { const double tilt_range = 2.0 * PI / 180.0; for (int j = 0; j < 2; ++j) { @@ -517,14 +608,6 @@ GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_ge problem.SetParameterUpperBound(detector_rot, j, (j == 0 ? nominal_geom.GetPoniRot1_rad() : nominal_geom.GetPoniRot2_rad()) + tilt_range); } - } else - problem.SetParameterBlockConstant(detector_rot); - problem.SetParameterBlockConstant(rot_vec); - problem.SetParameterBlockConstant(cell_ang); // angles not refined in v1 - if (!settings.refine_beam_distance) { - problem.SetParameterBlockConstant(beam); - problem.SetParameterBlockConstant(&distance_mm); - problem.SetParameterBlockConstant(cell_len); } ceres::Solver::Options options; diff --git a/tests/ModuleRefinementTest.cpp b/tests/ModuleRefinementTest.cpp index 08494e13c..4a4f6e234 100644 --- a/tests/ModuleRefinementTest.cpp +++ b/tests/ModuleRefinementTest.cpp @@ -103,6 +103,24 @@ TEST_CASE("RefineGlobalGeometry_ModulePlacement") { CHECK(placed.beam_x_px == Catch::Approx(1000).margin(0.1)); CHECK(placed.beam_y_px == Catch::Approx(1000).margin(0.1)); CHECK(placed.distance_mm == Catch::Approx(150).margin(0.05)); + + // The distance held (it shares a block with the beam): it stays exactly where it was given. + GeometryRefinerSettings held_distance = st; + held_distance.refine_distance = false; + const auto fixed = RefineGlobalGeometry(geom, latt.GetUnitCell(), frames, held_distance); + REQUIRE(fixed.ok); + CHECK(fixed.distance_mm == geom.GetDetectorDistance_mm()); + CHECK(fixed.modules[0].dx == Catch::Approx(truth[0].dx).margin(0.1)); + + // Only the orientations moved, on the modules as placed: the spots fit as well as in the placement. + GeometryRefinerSettings ev = st; + ev.refine_modules = false; + ev.refine_beam_distance = false; + ev.modules = placed.modules; + const auto scored = RefineGlobalGeometry(geom, latt.GetUnitCell(), frames, ev); + REQUIRE(scored.ok); + CHECK(scored.median_residual_px < 0.1); + CHECK(scored.distance_mm == geom.GetDetectorDistance_mm()); } // The coarse search ahead of the placement: a module off by several pixels is found near where it is, and the