leonarski_fandClaude Opus 5.5
c83b949b69
The residual of a spot is the observed reciprocal vector (beam and distance, the tilt where refined, the module's move where placed) minus the predicted one (orientation, cell lengths), and the two halves share no parameter block. AutoDiffCostFunction carried all 12 (up to 14) derivatives through both halves; SpotCost differentiates each half on its own blocks (6-8 and 6 derivatives) and writes the Jacobian blocks from them. Values and derivatives are the same numbers. Replaying the dumped placement inputs of four multi-panel stills runs on one thread: identical held-out residuals, beam, distance, cell and module shifts; single-thread time of the nine fits (after the previous commit -> now): AGIPD 66.8 -> 37.4 s, MPCCD 111.8 -> 92.1 s, JF4M 79.8 -> 52.3 s, JF16M 68.4 -> 43.1 s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
817 lines
41 KiB
C++
817 lines
41 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, 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);
|
|
}
|
|
|
|
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);
|
|
}
|
|
|
|
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 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;
|
|
}
|
|
};
|
|
|
|
// 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 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;
|
|
}
|
|
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;
|
|
}
|
|
|