mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 87 Wiki Activity

Merge fc-modules: module placement fits differentiate only what they move, the observed and predicted halves apart; threads split between the fits (decision-identical)

Browse Source 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-23 12:50:54 +02:00
2 parents 3a7bb28ed2 b84d648c05
commit 27a03fbb06
3 files changed
+231 -89
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/geom_refinement/GeometryRefiner.cpp
+205 -85
View File
@@ -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 <typename T>
void Poni(const T *const beam, const T &distance_mm, const T *const detector_rot, T *poni) const {
template <typename T, typename R>
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 <typename T>
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<int32_t> &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<true, true>(parameters, residuals, jacobians)
: EvaluateWith<true, false>(parameters, residuals, jacobians);
return module ? EvaluateWith<false, true>(parameters, residuals, jacobians)
: EvaluateWith<false, false>(parameters, residuals, jacobians);
}
template <typename T>
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 TILT, bool MODULE>
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<double, N_OBS>;
using JP = ceres::Jet<double, 6>;
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<JP> 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 <typename T, typename R>
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 <typename T>
bool operator()(const T *const p0, T *out) const {
T e_pred_recip[3];
const AngleAxisRotator<T> 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<size_t>(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<ceres::ParameterBlockOrdering>();
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<size_t>(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<OrientationResidual, 3, 3>(r),
&loss, orient[i].data());
} else if (s.module >= 0 && static_cast<size_t>(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<ModuleXtalResidual, 3, 2, 1, 2, 3, 3, 3, 3, 3>(
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<BeamXtalResidual, 3, 2, 1, 2, 3, 3, 3, 3>(
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<XtalResidual, 3, 2, 1, 2, 3, 3, 3, 3>(
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<CellLengthRegularizer, 3, 3>(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;
rugnux/RugnuxStills.cpp
+8 -4
View File
@@ -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<ModuleShift> &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);
tests/ModuleRefinementTest.cpp
+18
View File
@@ -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