diff --git a/image_analysis/geom_refinement/GeometryRefiner.cpp b/image_analysis/geom_refinement/GeometryRefiner.cpp index bebb5a41f..581aebfd7 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,143 @@ struct BeamBlock { } }; -// XtalResidual with the beam block read through BeamBlock. -struct BeamXtalResidual { - XtalResidual residual; - BeamBlock beam_block; - - 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, - 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); +// The residual of one spot (XtalResidual) on the blocks this refinement can move, in this order: the beam +// (read through BeamBlock) and the distance as one block, the tilt where it is refined, the frame's orientation, +// the cell lengths, and the module's move (dx, dy, rot) where the spot's module is placed. The rotation axis, the +// cell angles and a held tilt enter as plain numbers. +// +// The residual is the difference of two halves that share no block: the observed reciprocal vector (beam, +// distance, tilt, move) and the predicted one (orientation, cell). Ceres' AutoDiffCostFunction would carry the +// derivatives of every block through both halves; here each half is differentiated on its own blocks only, which +// gives the same values and the same derivatives with half as many of them per number. +// +// 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. +class SpotCost : public ceres::CostFunction { +public: + SpotCost(const XtalResidual &residual, const BeamBlock &beam_block, const double *detector_rot, + const double *rotation_axis, const double *cell_ang, bool tilt, bool module, double ux = 0, + double uy = 0) + : residual(residual), beam_block(beam_block), detector_rot(detector_rot), rotation_axis(rotation_axis), + cell_ang(cell_ang), tilt(tilt), module(module), ux(ux), uy(uy) { + set_num_residuals(3); + std::vector &sizes = *mutable_parameter_block_sizes(); + sizes.push_back(3); + if (tilt) + sizes.push_back(2); + sizes.push_back(3); + sizes.push_back(3); + if (module) + sizes.push_back(3); } -}; -// 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; + bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const override { + if (tilt) + return module ? EvaluateWith(parameters, residuals, jacobians) + : EvaluateWith(parameters, residuals, jacobians); + return module ? EvaluateWith(parameters, residuals, jacobians) + : EvaluateWith(parameters, residuals, jacobians); + } - 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, - const T *const move, T *out) const { +private: + template + bool EvaluateWith(double const *const *p, double *out, double **jac) const { + constexpr int N_OBS = 3 + (TILT ? 2 : 0) + (MODULE ? 3 : 0); + using JO = ceres::Jet; + using JP = ceres::Jet; + constexpr int TILT_BLOCK = 1, P0_BLOCK = TILT ? 2 : 1, P1_BLOCK = P0_BLOCK + 1, MOVE_BLOCK = P1_BLOCK + 1; + + // The observed half, on the beam and distance, the tilt and the move. + JO bd[3], t[2], m[3]; + for (int j = 0; j < 3; ++j) + bd[j] = JO(p[0][j], j); + if constexpr (TILT) + for (int j = 0; j < 2; ++j) + t[j] = JO(p[TILT_BLOCK][j], 3 + j); + if constexpr (MODULE) + for (int j = 0; j < 3; ++j) + m[j] = JO(p[MOVE_BLOCK][j], N_OBS - 3 + j); + JO obs[3]; + if constexpr (TILT) + Observed(bd, t, MODULE ? m : nullptr, obs); + else + Observed(bd, detector_rot, MODULE ? m : nullptr, obs); + + // The predicted half, on the orientation and the cell lengths. + JP p0[3], p1[3]; + for (int j = 0; j < 3; ++j) { + p0[j] = JP(p[P0_BLOCK][j], j); + p1[j] = JP(p[P1_BLOCK][j], 3 + j); + } + const JP p2[3] = {JP(cell_ang[0]), JP(cell_ang[1]), JP(cell_ang[2])}; + JP recip_unrot[3]; + residual.PredictedRecipUnrot(p1, p2, recip_unrot); + JP pred[3]; + const AngleAxisRotator rot_p0(p0); + rot_p0.Rotate(recip_unrot, pred); + + for (int r = 0; r < 3; ++r) + out[r] = obs[r].a - pred[r].a; + if (!jac) + return true; + // Row-major 3 x size blocks; a block held constant has no Jacobian asked for. + auto fill = [](double *block, int size, auto derivative) { + if (block) + for (int r = 0; r < 3; ++r) + for (int c = 0; c < size; ++c) + block[r * size + c] = derivative(r, c); + }; + fill(jac[0], 3, [&](int r, int c) { return obs[r].v[c]; }); + if constexpr (TILT) + fill(jac[TILT_BLOCK], 2, [&](int r, int c) { return obs[r].v[3 + c]; }); + if constexpr (MODULE) + fill(jac[MOVE_BLOCK], 3, [&](int r, int c) { return obs[r].v[N_OBS - 3 + c]; }); + fill(jac[P0_BLOCK], 3, [&](int r, int c) { return -pred[r].v[c]; }); + fill(jac[P1_BLOCK], 3, [&](int r, int c) { return -pred[r].v[3 + c]; }); + return true; + } + + template + void Observed(const T *const beam_distance, const R *const tilt, const T *const move, T *recip_obs) 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); + T b[2]; + beam_block.Poni(beam_distance, beam_distance[2], 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])}; + residual.ObservedRecip(b, beam_distance[2], tilt, axis, recip_obs); + } + + const XtalResidual residual; + const BeamBlock beam_block; + const double *detector_rot; // the tilt where it is held + const double *rotation_axis; + const double *cell_ang; + const bool tilt, module; + const double ux, uy; +}; + +// 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 +353,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 +449,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 +470,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 +488,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 +511,36 @@ 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()); - } 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); - 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); + const XtalResidual found(s.x, s.y, lambda, pixel, cos_rot3, sin_rot3, 0.0, h, k, l, system, + orientation); + auto *cost = new SpotCost(found, beam_block, detector_rot, rot_vec, cell_ang, refine_tilt, true, + s.x - m.x0, s.y - m.y0); + if (refine_tilt) + problem.AddResidualBlock(cost, &loss, beam_distance, detector_rot, orient[i].data(), + cell_len, move[s.module].data()); + else + problem.AddResidualBlock(cost, &loss, beam_distance, orient[i].data(), cell_len, + move[s.module].data()); + } else { + auto *cost = new SpotCost(xr, beam_block, detector_rot, rot_vec, cell_ang, refine_tilt, false); + if (refine_tilt) + problem.AddResidualBlock(cost, &loss, beam_distance, detector_rot, orient[i].data(), + cell_len); + else + problem.AddResidualBlock(cost, &loss, beam_distance, orient[i].data(), cell_len); + } ++used_in_frame; ++spots_used; } @@ -447,9 +568,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 +617,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 +645,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/rugnux/RugnuxStills.cpp b/rugnux/RugnuxStills.cpp index 79394851f..26d4f3d08 100644 --- a/rugnux/RugnuxStills.cpp +++ b/rugnux/RugnuxStills.cpp @@ -466,12 +466,14 @@ void Rugnux::RefineStillsGeometry(int start_image, int end_image, int images_to_ half[i % 2].push_back(placing[i]); // The four held-out fits and the placement on every frame do not depend on each other, and one fit - // keeps only five or six cores busy however many threads it is given (measured on 32): all five run - // side by side, each with every thread - on a busy machine that measured faster than a fifth each - // (JF16M: 10.9 s against 14.6 s, 18.9 s one after the other). The placement is computed whether or - // not the held-out residuals then call for it. + // keeps only a few cores busy however many threads it is given (its Schur elimination contends on the + // shared blocks): all five run side by side. The placement, on twice the frames, is the one they wait + // for, so it gets half the threads and each held-out fit an eighth - measured on 32 threads (AGIPD, + // MPCCD, JF4M, JF16M) 10-15% faster than every fit on every thread, which mostly spins on the locks. + // The placement is computed whether or not the held-out residuals then call for it. TimingMark placement_mark(logger, "module placement fits"); GeometryRefinerSettings st = gr_settings; + st.num_threads = std::max(1, config_.nthreads / 2); st.modules = start; st.refine_modules = true; auto placement = std::async(std::launch::async, [&] { @@ -481,6 +483,7 @@ void Rugnux::RefineStillsGeometry(int start_image, int end_image, int images_to_ // Held-out residual of the geometry fitted on half `fit`, scored on every spot of the other half. auto held_out = [&](const std::vector &modules, bool refine, int fit) { GeometryRefinerSettings sf = gr_settings; + sf.num_threads = std::max(1, config_.nthreads / 8); sf.modules = modules; sf.refine_modules = refine; const auto f = RefineGlobalGeometry(nominal, *cell, half[fit], sf); @@ -489,6 +492,7 @@ void Rugnux::RefineStillsGeometry(int start_image, int end_image, int images_to_ DiffractionGeometry g = nominal; g.BeamX_pxl(f.beam_x_px).BeamY_pxl(f.beam_y_px).DetectorDistance_mm(f.distance_mm); GeometryRefinerSettings ev = gr_settings; + ev.num_threads = sf.num_threads; ev.refine_beam_distance = false; ev.modules = f.modules; const auto e = RefineGlobalGeometry(g, f.cell, half[1 - fit], ev); 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