Files
Jungfraujoch/image_analysis/geom_refinement/GeometryRefiner.cpp
T
leonarski_fandClaude Opus 5.5 a39e18b80f stills: detector tilt refined in the geometry pre-pass, kept where held-out frames agree
The stills geometry bundle (RefineGlobalGeometry) can now refine the detector tilt
together with the beam, the distance, the cell and the modules. With the tilt free the
beam block is the direct beam, from which the PONI is worked out inside the residual: in
PONI terms a tilt moves the whole pattern as a moved beam does (D tan(tilt)/pixel), more
than the beam's 8 px box holds, while the direct beam does not move with it and the tilt
is left with the keystone. The two tilt-like rows of ModuleGauge are gone: within a
module a tilt is a gradient no rigid move makes, and the rows cost held-out accuracy.
The shift, turn and outward gauges stay.

The tilt is a hypothesis (DecideStillsTilt): the strong frames of the last geometry
refinement are split in two halves, each fitted with the tilt held and refined (modules
placed in both where the run placed them), each fit scored on the other half, and the
tilt is kept where the refined fits predict the held-out spots better. No threshold.
It is decided after the stills beam-centre trial and its re-refinement, on frames
indexed at the settled centre: a centre error on frames indexed at it reads as a tilt.
--rot1/--rot2 or --no-refine-tilt hold the tilt on stills.

The stills beam-centre trial now compares direct beams. Its hypotheses (the file's
centre, the background walk, the point-symmetry capture) are taken as the point the
beam lands on and each trial puts the PONI where its distance and the run's tilt land
the direct beam there; the walk, fitted through the geometry's tilt, returns a PONI and
is converted. The pre-scan beam-centre check log compares the walk's direct beam, not
its PONI, with the file's direct beam. Without a tilt nothing changes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
2026-10-22 19:26:33 +02:00

697 lines
35 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>
void Poni(const T *const beam, const T &distance_mm, const T *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);
}
};
// 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);
}
};
// 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;
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 {
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);
}
};
// 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.
double beam[2] = {nominal_geom.GetBeamX_pxl(), nominal_geom.GetBeamY_pxl()};
double distance_mm = nominal_geom.GetDetectorDistance_mm();
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};
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;
};
// 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;
ceres::Problem problem;
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;
if (s.module >= 0 && static_cast<size_t>(s.module) < modules.size() && refined[s.module]) {
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);
++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; constant ones are stripped by Ceres.
for (double *b : {beam, &distance_mm, detector_rot, rot_vec, cell_len, cell_ang})
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).
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 (!settings.refine_distance)
problem.SetParameterBlockConstant(&distance_mm);
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);
}
} 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;
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;
}