leonarski_fandClaude Opus 5.5
f6c948732b
The stills geometry refinement (module placement fits, the held-out fits that decide them, the tilt decision) spent most of its time differentiating parameters it never moves and in Ceres' general Schur elimination: - Each spot's residual declared the rotation axis, the cell angles and the (usually held) tilt as blocks. Ceres sizes its dual numbers from the declared blocks, so a placement residual carried 20 derivatives to use 12. They now enter as plain numbers (SpotResidual / TiltedSpotResidual). - A fit that moves only the orientations (the held-out scoring) gets a residual over the orientation alone, with the observed and predicted vectors taken once per spot (OrientationResidual): 3 derivatives instead of 17. - The beam and the distance are one block of three, so every refined block has three parameters and Ceres uses its fixed-size Schur eliminator (1.7x faster linear solve on one thread). A held distance is a SubsetManifold on that block. - One robust loss shared by every spot instead of one allocated per spot. Same arithmetic on the same values: replaying the dumped placement inputs of four multi-panel stills runs (AGIPD 128 ASICs, MPCCD 8, JF4M 8, JF16M 32 modules) on one thread gives identical held-out residuals, beam, distance, cell and every module shift at the printed precision. Single- thread CPU of the nine placement fits: 397 -> 112 s (MPCCD), 179 -> 80 s (JF4M), 131 -> 68 s (JF16M). On AGIPD the concurrent step goes 8.5 -> 6.4 s wall (quiet replay, 32 threads). Test: RefineGlobalGeometry_ModulePlacement now also covers a held distance and an orientation-only fit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
780 lines
39 KiB
C++
780 lines
39 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "GeometryRefiner.h"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <cmath>
|
|
#include <vector>
|
|
|
|
#include "../../common/JFJochMath.h" // PI
|
|
#include "../../common/ParallelFor.h"
|
|
#include "LatticeReduction.h"
|
|
#include "XtalResidual.h"
|
|
#include "ceres/ceres.h"
|
|
|
|
namespace {
|
|
|
|
// Trigonal is described with the hexagonal cell + B, same as the per-image XtalOptimizer.
|
|
gemmi::CrystalSystem MapSystem(gemmi::CrystalSystem s) {
|
|
return (s == gemmi::CrystalSystem::Trigonal) ? gemmi::CrystalSystem::Hexagonal : s;
|
|
}
|
|
|
|
// The cell angles (radians) that XtalResidual bakes into its B matrix for a given system. Only
|
|
// monoclinic (beta) and triclinic (all three) read p2; the rest use fixed symmetry angles. We keep the
|
|
// angles constant in v1, so these are set once from the input cell.
|
|
void InitCellAngles(gemmi::CrystalSystem system, const UnitCell &cell, double ang[3]) {
|
|
const double d2r = PI / 180.0;
|
|
switch (system) {
|
|
case gemmi::CrystalSystem::Hexagonal:
|
|
ang[0] = PI / 2.0; ang[1] = PI / 2.0; ang[2] = 2.0 * PI / 3.0;
|
|
break;
|
|
case gemmi::CrystalSystem::Monoclinic:
|
|
// XtalResidual reads p2[0] as beta.
|
|
ang[0] = cell.beta * d2r; ang[1] = PI / 2.0; ang[2] = PI / 2.0;
|
|
break;
|
|
case gemmi::CrystalSystem::Triclinic:
|
|
ang[0] = cell.alpha * d2r; ang[1] = cell.beta * d2r; ang[2] = cell.gamma * d2r;
|
|
break;
|
|
default: // orthorhombic / tetragonal / cubic
|
|
ang[0] = PI / 2.0; ang[1] = PI / 2.0; ang[2] = PI / 2.0;
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Effective (lengths, alpha,beta,gamma) that reproduce the lattice XtalResidual builds from (len, ang),
|
|
// so a frame's lattice can be reconstructed with AngleAxisAndCellToLattice for HKL re-assignment.
|
|
void EffectiveCell(gemmi::CrystalSystem system, const double len[3], const double ang[3],
|
|
double lengths[3], double &alpha, double &beta, double &gamma) {
|
|
switch (system) {
|
|
case gemmi::CrystalSystem::Hexagonal:
|
|
lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[2];
|
|
alpha = PI / 2.0; beta = PI / 2.0; gamma = 2.0 * PI / 3.0;
|
|
break;
|
|
case gemmi::CrystalSystem::Tetragonal:
|
|
lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[2];
|
|
alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0;
|
|
break;
|
|
case gemmi::CrystalSystem::Cubic:
|
|
lengths[0] = len[0]; lengths[1] = len[0]; lengths[2] = len[0];
|
|
alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0;
|
|
break;
|
|
case gemmi::CrystalSystem::Monoclinic:
|
|
lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2];
|
|
alpha = PI / 2.0; beta = ang[0]; gamma = PI / 2.0;
|
|
break;
|
|
case gemmi::CrystalSystem::Triclinic:
|
|
lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2];
|
|
alpha = ang[0]; beta = ang[1]; gamma = ang[2];
|
|
break;
|
|
default: // orthorhombic
|
|
lengths[0] = len[0]; lengths[1] = len[1]; lengths[2] = len[2];
|
|
alpha = PI / 2.0; beta = PI / 2.0; gamma = PI / 2.0;
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Orientation seed (angle-axis) for a frame's lattice, using the per-system decomposition the per-image
|
|
// optimizer uses so it matches XtalResidual's B convention.
|
|
void OrientationSeed(gemmi::CrystalSystem system, const CrystalLattice &latt, double rod[3]) {
|
|
double scratch[3];
|
|
if (system == gemmi::CrystalSystem::Hexagonal) {
|
|
LatticeToRodriguesAndLengths_Hex(latt, rod, scratch);
|
|
} else if (system == gemmi::CrystalSystem::Monoclinic) {
|
|
double beta;
|
|
LatticeToRodriguesLengthsBeta_Mono(latt, rod, scratch, beta);
|
|
} else {
|
|
LatticeToRodriguesAndLengths_GS(latt, rod, scratch);
|
|
}
|
|
}
|
|
|
|
UnitCell BuildUnitCell(gemmi::CrystalSystem system, const double len[3], const double ang[3]) {
|
|
const double r2d = 180.0 / PI;
|
|
UnitCell uc{};
|
|
switch (system) {
|
|
case gemmi::CrystalSystem::Hexagonal:
|
|
uc = {(float)len[0], (float)len[0], (float)len[2], 90.0f, 90.0f, 120.0f};
|
|
break;
|
|
case gemmi::CrystalSystem::Tetragonal:
|
|
uc = {(float)len[0], (float)len[0], (float)len[2], 90.0f, 90.0f, 90.0f};
|
|
break;
|
|
case gemmi::CrystalSystem::Cubic:
|
|
uc = {(float)len[0], (float)len[0], (float)len[0], 90.0f, 90.0f, 90.0f};
|
|
break;
|
|
case gemmi::CrystalSystem::Monoclinic:
|
|
uc = {(float)len[0], (float)len[1], (float)len[2], 90.0f, (float)(ang[0] * r2d), 90.0f};
|
|
break;
|
|
case gemmi::CrystalSystem::Triclinic:
|
|
uc = {(float)len[0], (float)len[1], (float)len[2],
|
|
(float)(ang[0] * r2d), (float)(ang[1] * r2d), (float)(ang[2] * r2d)};
|
|
break;
|
|
default: // orthorhombic
|
|
uc = {(float)len[0], (float)len[1], (float)len[2], 90.0f, 90.0f, 90.0f};
|
|
break;
|
|
}
|
|
return uc;
|
|
}
|
|
|
|
// Anchors the cell lengths to the input cell. Penalises each length's deviation from its input value
|
|
// with a large weight, breaking the low-resolution distance<->cell-scale degeneracy (position ~
|
|
// distance / cell_scale) so distance is determined by the data with the cell pinned.
|
|
struct CellLengthRegularizer {
|
|
double w0, w1, w2;
|
|
double l0, l1, l2;
|
|
template <typename T>
|
|
bool operator()(const T *const len, T *r) const {
|
|
r[0] = T(w0) * (len[0] - T(l0));
|
|
r[1] = T(w1) * (len[1] - T(l1));
|
|
r[2] = T(w2) * (len[2] - T(l2));
|
|
return true;
|
|
}
|
|
};
|
|
|
|
// Where a spot found at (x, y) on module m really is (ModuleShift).
|
|
void PlaceOnModule(const ModuleShift &m, const double move[3], double x, double y, double &px, double &py) {
|
|
const double ux = x - m.x0, uy = y - m.y0;
|
|
const double c = std::cos(move[2]), s = std::sin(move[2]);
|
|
px = m.x0 + c * ux - s * uy + move[0];
|
|
py = m.y0 + s * ux + c * uy + move[1];
|
|
}
|
|
|
|
// How a residual reads the beam block. Without the tilt refined it is the PONI, as everywhere else. With the
|
|
// tilt refined it is the DIRECT BEAM, the point the undeflected beam lands on, and the PONI is worked out from
|
|
// it: a tilt moves the whole pattern by D tan(tilt)/pixel, exactly as a moved PONI does, so in PONI terms
|
|
// the two trade tens of pixels for each other (and the beam's box cannot hold the trade), while the direct
|
|
// beam does not move with the tilt and the tilt is left with only what no beam can imitate - the keystone.
|
|
// In XtalResidual's frame the direct beam is where M (u, v) = (-D tan rot1, -D tan rot2 / cos rot1), with
|
|
// (u, v) its offset from the PONI in mm and M the image orientation (a mirror and quarter turns, so M^-1 = M^T).
|
|
struct BeamBlock {
|
|
bool direct = false;
|
|
double m00 = 1, m01 = 0, m10 = 0, m11 = 1; // DetectorOrientation::Matrix(), as XtalResidual takes it
|
|
double pixel = 1;
|
|
|
|
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];
|
|
return;
|
|
}
|
|
using std::cos;
|
|
using std::tan;
|
|
const T dx = -distance_mm * tan(detector_rot[0]);
|
|
const T dy = -distance_mm * tan(detector_rot[1]) / cos(detector_rot[0]);
|
|
poni[0] = beam[0] - (T(m00) * dx + T(m10) * dy) / T(pixel);
|
|
poni[1] = beam[1] - (T(m01) * dx + T(m11) * dy) / T(pixel);
|
|
}
|
|
};
|
|
|
|
// 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 <typename T, typename R>
|
|
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<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;
|
|
}
|
|
|
|
// beam_distance: the beam (x, y) and the distance, as one block.
|
|
template <typename T>
|
|
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<const T *>(nullptr), out);
|
|
}
|
|
|
|
template <typename T>
|
|
bool operator()(const T *const beam_distance, const T *const p0, const T *const p1, const T *const move,
|
|
T *out) const {
|
|
return Evaluate(beam_distance, beam_distance + 2, detector_rot, p0, p1, move, out);
|
|
}
|
|
};
|
|
|
|
// The same with the tilt refined, as a block after the beam and distance.
|
|
struct TiltedSpotResidual : SpotResidual {
|
|
template <typename T>
|
|
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<const T *>(nullptr), out);
|
|
}
|
|
|
|
template <typename T>
|
|
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 {
|
|
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 <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;
|
|
}
|
|
};
|
|
|
|
// What module moves cannot be told apart from: all modules shifted alike is the beam centre, all turned
|
|
// alike (about their own centres) is the crystal orientations, and all pushed outwards alike is the
|
|
// distance. The sums of those three over the moved modules are held at zero - the beam, the orientations
|
|
// and the distance then carry them - so the modules are left with what only they can explain. The sums are
|
|
// weighted by the spots each module holds: a module the frames barely reach is placed poorly, and an
|
|
// unweighted sum would hand its error to all the others.
|
|
//
|
|
// The detector tilt is not among them. Within one module a tilt is a gradient (a scale and a shear across
|
|
// it), which no rigid move of the module makes, and with the tilt refined the placement gains no new
|
|
// direction it cannot tell apart; holding the modules' moves to no tilt-like pattern as well only cost
|
|
// accuracy on frames the fit had not seen.
|
|
struct ModuleGauge {
|
|
std::vector<std::array<double, 2>> outward; // unit vector from the beam to each module centre
|
|
std::vector<double> share; // each module's fraction of the spots
|
|
double lever_px; // turns are weighed at this distance from the centre
|
|
double weight;
|
|
|
|
template <typename T>
|
|
bool operator()(T const *const *move, T *out) const {
|
|
T sx(0), sy(0), sr(0), so(0);
|
|
for (size_t i = 0; i < outward.size(); ++i) {
|
|
const T w(share[i]);
|
|
sx += w * move[i][0];
|
|
sy += w * move[i][1];
|
|
sr += w * move[i][2] * T(lever_px);
|
|
so += w * (move[i][0] * T(outward[i][0]) + move[i][1] * T(outward[i][1]));
|
|
}
|
|
out[0] = T(weight) * sx;
|
|
out[1] = T(weight) * sy;
|
|
out[2] = T(weight) * sr;
|
|
out[3] = T(weight) * so;
|
|
return true;
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_geom,
|
|
const UnitCell &input_cell,
|
|
const std::vector<GeomRefineFrame> &frames,
|
|
const GeometryRefinerSettings &settings) {
|
|
GeometryRefinerResult result;
|
|
try {
|
|
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. 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();
|
|
// Same for every spot of every frame, so taken once here rather than per residual.
|
|
const double cos_rot3 = std::cos(rot3), sin_rot3 = std::sin(rot3);
|
|
const DetectorOrientation orientation = nominal_geom.GetOrientation();
|
|
const double lambda = nominal_geom.GetWavelength_A();
|
|
const double pixel = nominal_geom.GetPixelSize_mm();
|
|
|
|
// With the tilt refined, the beam block is the direct beam (BeamBlock), started and bounded where the
|
|
// nominal geometry puts the direct beam.
|
|
const bool refine_tilt = settings.refine_tilt && settings.refine_beam_distance;
|
|
BeamBlock beam_block;
|
|
if (refine_tilt) {
|
|
const RotMatrix m = orientation.Matrix();
|
|
beam_block = {true, m.Column(0).x, m.Column(1).x, m.Column(0).y, m.Column(1).y, pixel};
|
|
double poni_at_zero[2];
|
|
const double zero[2] = {0, 0};
|
|
beam_block.Poni(zero, distance_mm, detector_rot, poni_at_zero);
|
|
beam[0] -= poni_at_zero[0];
|
|
beam[1] -= poni_at_zero[1];
|
|
}
|
|
const double beam0[2] = {beam[0], beam[1]};
|
|
// The PONI the beam block stands for, at the current tilt and distance.
|
|
auto poni = [&](double &x, double &y) {
|
|
double p[2];
|
|
beam_block.Poni(beam, distance_mm, detector_rot, p);
|
|
x = p[0];
|
|
y = p[1];
|
|
};
|
|
|
|
// Shared cell blocks. Lengths refined (regularized); angles constant.
|
|
double cell_len[3] = {input_cell.a, input_cell.b, input_cell.c};
|
|
const double cell_len0[3] = {input_cell.a, input_cell.b, input_cell.c};
|
|
double cell_ang[3];
|
|
InitCellAngles(system, input_cell, cell_ang);
|
|
|
|
double eff_len[3], eff_alpha, eff_beta, eff_gamma; // recomputed each round from cell_len/ang
|
|
|
|
// Per-frame orientation seeds.
|
|
std::vector<std::array<double, 3>> orient(frames.size());
|
|
for (size_t i = 0; i < frames.size(); ++i)
|
|
OrientationSeed(system, frames[i].lattice, orient[i].data());
|
|
|
|
// Reciprocal length of one detector pixel at the nominal geometry, used to scale the robust
|
|
// loss (~3 px) into the reciprocal-space residual units.
|
|
const double recip_per_px = pixel / (distance_mm * lambda);
|
|
const double loss_scale = 3.0 * recip_per_px;
|
|
|
|
// The modules' moves (dx, dy, rot), starting where settings put them. Which of them are refined is
|
|
// decided once, on the indexed spots each holds, so the problem keeps its shape from round to round.
|
|
const auto &modules = settings.modules;
|
|
std::vector<std::array<double, 3>> move(modules.size());
|
|
for (size_t m = 0; m < modules.size(); ++m)
|
|
move[m] = {modules[m].dx, modules[m].dy, modules[m].rot_rad};
|
|
// A module counts the spots on it that its frame's lattice indexes where the module starts - not
|
|
// every spot found on it: a frame's other spots (noise, a second lattice, ice) would otherwise make a
|
|
// module the frames barely reach look well determined.
|
|
std::vector<int> module_spots(modules.size(), 0);
|
|
for (const auto &f : frames) {
|
|
const Coord v0 = f.lattice.Vec0(), v1 = f.lattice.Vec1(), v2 = f.lattice.Vec2();
|
|
for (const auto &s : f.spots)
|
|
if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size()) {
|
|
double x, y;
|
|
PlaceOnModule(modules[s.module], move[s.module].data(), s.x, s.y, x, y);
|
|
const Coord recip = nominal_geom.DetectorToRecip(x, y);
|
|
const double hf = recip * v0, kf = recip * v1, lf = recip * v2;
|
|
if (std::fabs(hf - std::round(hf)) < 0.15 && std::fabs(kf - std::round(kf)) < 0.15
|
|
&& std::fabs(lf - std::round(lf)) < 0.15)
|
|
module_spots[s.module]++;
|
|
}
|
|
}
|
|
std::vector<bool> refined(modules.size(), false);
|
|
int refined_count = 0;
|
|
for (size_t m = 0; m < modules.size(); ++m)
|
|
if (settings.refine_modules && module_spots[m] >= settings.min_module_spots) {
|
|
refined[m] = true;
|
|
refined_count++;
|
|
}
|
|
// One module alone has nothing to be placed against: its move is the beam and the orientation.
|
|
if (refined_count < 2)
|
|
std::fill(refined.begin(), refined.end(), false);
|
|
// A module too thin to be placed stays where it was given, whatever its start: a search on a few dozen
|
|
// spots finds noise.
|
|
if (settings.refine_modules)
|
|
for (size_t m = 0; m < modules.size(); ++m)
|
|
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);
|
|
else {
|
|
x = s.x;
|
|
y = s.y;
|
|
}
|
|
};
|
|
|
|
// Tolerance schedule (fractional-hkl norm), tightening each round like the per-image optimizer.
|
|
auto round_tol = [&](int r) {
|
|
if (settings.rounds <= 1) return 0.2;
|
|
const double t = static_cast<double>(r) / (settings.rounds - 1);
|
|
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;
|
|
|
|
for (int round = 0; round < settings.rounds; ++round) {
|
|
EffectiveCell(system, cell_len, cell_ang, eff_len, eff_alpha, eff_beta, eff_gamma);
|
|
|
|
double poni_x, poni_y;
|
|
poni(poni_x, poni_y);
|
|
DiffractionGeometry cur = nominal_geom;
|
|
cur.BeamX_pxl(poni_x).BeamY_pxl(poni_y).DetectorDistance_mm(distance_mm)
|
|
.PoniRot1_rad(static_cast<float>(detector_rot[0])).PoniRot2_rad(static_cast<float>(detector_rot[1]));
|
|
|
|
const double tol = round_tol(round);
|
|
const double tol_sq = tol * tol;
|
|
|
|
// 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;
|
|
for (size_t i = 0; i < frames.size(); ++i) {
|
|
const CrystalLattice latt = AngleAxisAndCellToLattice(
|
|
orient[i].data(), eff_len, eff_alpha, eff_beta, eff_gamma);
|
|
const Coord v0 = latt.Vec0(), v1 = latt.Vec1(), v2 = latt.Vec2();
|
|
|
|
int used_in_frame = 0;
|
|
for (const auto &s : frames[i].spots) {
|
|
double x, y;
|
|
place(s, x, y);
|
|
const Coord recip = cur.DetectorToRecip(x, y);
|
|
const double hf = recip * v0, kf = recip * v1, lf = recip * v2;
|
|
const double h = std::round(hf), k = std::round(kf), l = std::round(lf);
|
|
const double dn = (hf - h) * (hf - h) + (kf - k) * (kf - k) + (lf - l) * (lf - l);
|
|
if (dn > tol_sq)
|
|
continue;
|
|
|
|
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];
|
|
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<TiltedSpotResidual, 3, 3, 2, 3, 3, 3>(
|
|
new TiltedSpotResidual{sr}),
|
|
&loss,
|
|
beam_distance, detector_rot, orient[i].data(), cell_len, move[s.module].data());
|
|
else
|
|
problem.AddResidualBlock(
|
|
new ceres::AutoDiffCostFunction<SpotResidual, 3, 3, 3, 3, 3>(new SpotResidual{sr}),
|
|
&loss,
|
|
beam_distance, orient[i].data(), cell_len, move[s.module].data());
|
|
} else if (refine_tilt)
|
|
problem.AddResidualBlock(
|
|
new ceres::AutoDiffCostFunction<TiltedSpotResidual, 3, 3, 2, 3, 3>(
|
|
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<SpotResidual, 3, 3, 3, 3>(
|
|
new SpotResidual{xr, beam_block, detector_rot, rot_vec, cell_ang}),
|
|
&loss,
|
|
beam_distance, orient[i].data(), cell_len);
|
|
++used_in_frame;
|
|
++spots_used;
|
|
}
|
|
if (used_in_frame > 0) {
|
|
ordering->AddElementToGroup(orient[i].data(), 0); // eliminate orientations first
|
|
++frames_used;
|
|
}
|
|
}
|
|
|
|
if (spots_used < settings.min_spots_total) {
|
|
if (round == 0)
|
|
return result; // ok stays false
|
|
break; // keep the previous round's solution
|
|
}
|
|
|
|
// Cell-length regularizer (weight ~ coeff * sqrt(N_obs) / L0^2 -> effectively pins the cell).
|
|
const double sqrtN = std::sqrt(static_cast<double>(spots_used));
|
|
if (settings.refine_beam_distance) {
|
|
CellLengthRegularizer *reg = new CellLengthRegularizer{
|
|
settings.cell_reg_coeff * sqrtN / (cell_len0[0] * cell_len0[0]),
|
|
settings.cell_reg_coeff * sqrtN / (cell_len0[1] * cell_len0[1]),
|
|
settings.cell_reg_coeff * sqrtN / (cell_len0[2] * cell_len0[2]),
|
|
cell_len0[0], cell_len0[1], cell_len0[2]};
|
|
problem.AddResidualBlock(
|
|
new ceres::AutoDiffCostFunction<CellLengthRegularizer, 3, 3>(reg), nullptr, cell_len);
|
|
}
|
|
|
|
// 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
|
|
// every spot by 1 px would, ten times over.
|
|
std::vector<size_t> moved;
|
|
for (size_t m = 0; m < modules.size(); ++m)
|
|
if (refined[m] && problem.HasParameterBlock(move[m].data()))
|
|
moved.push_back(m);
|
|
if (moved.size() >= 2) {
|
|
auto *gauge = new ModuleGauge;
|
|
double lever = 0, total = 0;
|
|
for (size_t m : moved)
|
|
total += module_spots[m];
|
|
for (size_t m : moved) {
|
|
const double ox = modules[m].x0 - nominal_geom.GetBeamX_pxl();
|
|
const double oy = modules[m].y0 - nominal_geom.GetBeamY_pxl();
|
|
const double r = std::max(1.0, std::hypot(ox, oy));
|
|
gauge->outward.push_back({ox / r, oy / r});
|
|
gauge->share.push_back(module_spots[m] / total);
|
|
lever += r;
|
|
}
|
|
gauge->lever_px = lever / static_cast<double>(moved.size());
|
|
gauge->weight = 10.0 * recip_per_px * sqrtN;
|
|
auto *cost = new ceres::DynamicAutoDiffCostFunction<ModuleGauge, 4>(gauge);
|
|
std::vector<double *> blocks;
|
|
for (size_t m : moved) {
|
|
cost->AddParameterBlock(3);
|
|
blocks.push_back(move[m].data());
|
|
}
|
|
cost->SetNumResiduals(4);
|
|
problem.AddResidualBlock(cost, nullptr, blocks);
|
|
}
|
|
for (size_t m : moved) {
|
|
ordering->AddElementToGroup(move[m].data(), 1);
|
|
const double start[3] = {modules[m].dx, modules[m].dy, modules[m].rot_rad};
|
|
const double range[3] = {12.0, 12.0, 1.5 * PI / 180.0};
|
|
for (int j = 0; j < 3; ++j) {
|
|
problem.SetParameterLowerBound(move[m].data(), j, start[j] - range[j]);
|
|
problem.SetParameterUpperBound(move[m].data(), j, start[j] + range[j]);
|
|
}
|
|
}
|
|
|
|
// Beam + distance bounds around the nominal value; cell lengths loosely bounded (the
|
|
// regularizer does the anchoring).
|
|
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_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) {
|
|
problem.SetParameterLowerBound(detector_rot, j, (j == 0 ? nominal_geom.GetPoniRot1_rad()
|
|
: nominal_geom.GetPoniRot2_rad()) - tilt_range);
|
|
problem.SetParameterUpperBound(detector_rot, j, (j == 0 ? nominal_geom.GetPoniRot1_rad()
|
|
: nominal_geom.GetPoniRot2_rad()) + tilt_range);
|
|
}
|
|
}
|
|
|
|
ceres::Solver::Options options;
|
|
options.linear_solver_type = ceres::DENSE_SCHUR;
|
|
options.linear_solver_ordering = ordering;
|
|
options.num_threads = std::max(1, settings.num_threads);
|
|
options.max_num_iterations = 100;
|
|
options.max_solver_time_in_seconds = 60.0;
|
|
options.logging_type = ceres::LoggingType::SILENT;
|
|
options.minimizer_progress_to_stdout = false;
|
|
|
|
ceres::Solver::Summary summary;
|
|
ceres::Solve(options, &problem, &summary);
|
|
solve_usable = summary.IsSolutionUsable();
|
|
|
|
result.frames_used = frames_used;
|
|
result.spots_used = spots_used;
|
|
}
|
|
|
|
// Final fit quality: median detector residual of the retained spots at the refined geometry.
|
|
EffectiveCell(system, cell_len, cell_ang, eff_len, eff_alpha, eff_beta, eff_gamma);
|
|
double poni_x, poni_y;
|
|
poni(poni_x, poni_y);
|
|
DiffractionGeometry cur = nominal_geom;
|
|
cur.BeamX_pxl(poni_x).BeamY_pxl(poni_y).DetectorDistance_mm(distance_mm)
|
|
.PoniRot1_rad(static_cast<float>(detector_rot[0])).PoniRot2_rad(static_cast<float>(detector_rot[1]));
|
|
std::vector<double> residuals;
|
|
constexpr double CAP_PX = 2.0;
|
|
double capped_sum = 0;
|
|
size_t capped_n = 0;
|
|
std::vector<double> module_sq(modules.size(), 0), module_d(modules.size(), 0);
|
|
std::vector<size_t> module_n(modules.size(), 0);
|
|
result.predicted_px.resize(frames.size());
|
|
for (size_t i = 0; i < frames.size(); ++i) {
|
|
const CrystalLattice latt = AngleAxisAndCellToLattice(
|
|
orient[i].data(), eff_len, eff_alpha, eff_beta, eff_gamma);
|
|
const Coord a = latt.Astar(), b = latt.Bstar(), c = latt.Cstar();
|
|
const Coord v0 = latt.Vec0(), v1 = latt.Vec1(), v2 = latt.Vec2();
|
|
for (const auto &s : frames[i].spots) {
|
|
double x, y;
|
|
place(s, x, y);
|
|
const Coord recip = cur.DetectorToRecip(x, y);
|
|
const double hf = recip * v0, kf = recip * v1, lf = recip * v2;
|
|
const double h = std::round(hf), k = std::round(kf), l = std::round(lf);
|
|
const double dn = (hf - h) * (hf - h) + (kf - k) * (kf - k) + (lf - l) * (lf - l);
|
|
const Coord recip_pred = a * (float)h + b * (float)k + c * (float)l;
|
|
const auto [px, py] = cur.RecipToDetector(recip_pred);
|
|
const double res = (std::isfinite(px) && std::isfinite(py)) ? std::hypot(px - x, py - y) : CAP_PX;
|
|
capped_sum += std::min(res, CAP_PX) * std::min(res, CAP_PX);
|
|
capped_n++;
|
|
// The prediction where the spot would be found: the spot's own module's move taken back.
|
|
float fx = NAN, fy = NAN;
|
|
if (std::isfinite(px) && std::isfinite(py)) {
|
|
fx = static_cast<float>(s.x + (px - x));
|
|
fy = static_cast<float>(s.y + (py - y));
|
|
}
|
|
result.predicted_px[i].push_back({fx, fy});
|
|
if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size()) {
|
|
module_sq[s.module] += std::min(res, CAP_PX) * std::min(res, CAP_PX);
|
|
module_d[s.module] += 1.0 / std::max(1e-6f, recip.Length());
|
|
module_n[s.module]++;
|
|
}
|
|
if (dn > 0.1 * 0.1)
|
|
continue;
|
|
if (std::isfinite(px) && std::isfinite(py))
|
|
residuals.push_back(res);
|
|
}
|
|
}
|
|
if (capped_n > 0)
|
|
result.capped_rms_px = std::sqrt(capped_sum / static_cast<double>(capped_n));
|
|
if (!residuals.empty()) {
|
|
std::nth_element(residuals.begin(), residuals.begin() + residuals.size() / 2, residuals.end());
|
|
result.median_residual_px = residuals[residuals.size() / 2];
|
|
}
|
|
|
|
result.ok = solve_usable && result.spots_used >= settings.min_spots_total && result.frames_used > 0;
|
|
result.beam_x_px = poni_x;
|
|
result.beam_y_px = poni_y;
|
|
result.distance_mm = distance_mm;
|
|
result.cell = BuildUnitCell(system, cell_len, cell_ang);
|
|
result.rot1_rad = detector_rot[0];
|
|
result.rot2_rad = detector_rot[1];
|
|
result.module_capped_rms_px.resize(modules.size(), 0);
|
|
result.module_mean_d_A.resize(modules.size(), 0);
|
|
for (size_t m = 0; m < modules.size(); ++m)
|
|
if (module_n[m] > 0) {
|
|
result.module_capped_rms_px[m] = std::sqrt(module_sq[m] / static_cast<double>(module_n[m]));
|
|
result.module_mean_d_A[m] = module_d[m] / static_cast<double>(module_n[m]);
|
|
}
|
|
result.modules = modules;
|
|
for (size_t m = 0; m < modules.size(); ++m) {
|
|
result.modules[m].dx = move[m][0];
|
|
result.modules[m].dy = move[m][1];
|
|
result.modules[m].rot_rad = move[m][2];
|
|
result.modules[m].spots = module_spots[m];
|
|
}
|
|
return result;
|
|
} catch (...) {
|
|
result.ok = false;
|
|
return result;
|
|
}
|
|
}
|
|
|
|
std::vector<ModuleScan> ScanModuleShifts(const DiffractionGeometry &geom,
|
|
const std::vector<GeomRefineFrame> &frames,
|
|
const std::vector<ModuleShift> &modules,
|
|
double range_px, size_t nthreads) {
|
|
// A spot is on its lattice when every fractional index is within this of a whole number.
|
|
constexpr double ON_LATTICE = 0.15;
|
|
|
|
struct Spot { float x, y; int frame; };
|
|
std::vector<std::vector<Spot>> on_module(modules.size());
|
|
for (size_t i = 0; i < frames.size(); i++)
|
|
for (const auto *list : {&frames[i].spots, &frames[i].unindexed})
|
|
for (const auto &s : *list)
|
|
if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size())
|
|
on_module[s.module].push_back({s.x, s.y, static_cast<int>(i)});
|
|
|
|
std::vector<std::array<Coord, 3>> axes(frames.size());
|
|
for (size_t i = 0; i < frames.size(); i++)
|
|
axes[i] = {frames[i].lattice.Vec0(), frames[i].lattice.Vec1(), frames[i].lattice.Vec2()};
|
|
|
|
auto score = [&](const std::vector<Spot> &spots, const ModuleShift &m, double dx, double dy) {
|
|
int n = 0;
|
|
for (const auto &s : spots) {
|
|
double x, y;
|
|
const double move[3] = {m.dx + dx, m.dy + dy, m.rot_rad};
|
|
PlaceOnModule(m, move, s.x, s.y, x, y);
|
|
const Coord recip = geom.DetectorToRecip(x, y);
|
|
bool on = true;
|
|
for (const auto &v : axes[s.frame]) {
|
|
const double f = recip * v;
|
|
on &= std::fabs(f - std::round(f)) < ON_LATTICE;
|
|
}
|
|
n += on;
|
|
}
|
|
return n;
|
|
};
|
|
|
|
std::vector<ModuleScan> out(modules.size());
|
|
for (size_t m = 0; m < modules.size(); m++) {
|
|
auto &r = out[m];
|
|
r.on_lattice_at_zero = r.on_lattice_at_best = score(on_module[m], modules[m], 0, 0);
|
|
// Coarse: whole-pixel steps over the range; then half-pixel steps around the best of them.
|
|
for (int pass = 0; pass < 2; pass++) {
|
|
const double step = pass == 0 ? 1.0 : 0.5;
|
|
const double half = pass == 0 ? range_px : 1.0;
|
|
const double cx = r.dx, cy = r.dy;
|
|
std::vector<double> offsets;
|
|
for (double d = -half; d <= half; d += step)
|
|
offsets.push_back(d);
|
|
// Every shift of the grid is scored on the workers, then the grid is walked in order: the first
|
|
// shift with the most spots wins, as it would on one thread.
|
|
const int n = static_cast<int>(offsets.size());
|
|
std::vector<int> on_lattice(n * n);
|
|
ParallelFor(n * n, nthreads, [&](int k) {
|
|
on_lattice[k] = score(on_module[m], modules[m], cx + offsets[k % n], cy + offsets[k / n]);
|
|
});
|
|
for (int k = 0; k < n * n; k++)
|
|
if (on_lattice[k] > r.on_lattice_at_best) {
|
|
r.on_lattice_at_best = on_lattice[k];
|
|
r.dx = cx + offsets[k % n];
|
|
r.dy = cy + offsets[k / n];
|
|
}
|
|
}
|
|
}
|
|
return out;
|
|
}
|
|
|