Files
Jungfraujoch/image_analysis/geom_refinement/XtalOptimizer.cpp
Filip Leonarski dbc09e04f9
All checks were successful
Build Packages / build:rpm (ubuntu2404_nocuda) (push) Successful in 11m43s
Build Packages / build:rpm (rocky8_nocuda) (push) Successful in 12m58s
Build Packages / build:rpm (ubuntu2204_nocuda) (push) Successful in 12m58s
Build Packages / Generate python client (push) Successful in 15s
Build Packages / build:rpm (rocky8) (push) Successful in 13m7s
Build Packages / Create release (push) Has been skipped
Build Packages / build:rpm (ubuntu2204) (push) Successful in 13m12s
Build Packages / build:rpm (rocky8_sls9) (push) Successful in 13m18s
Build Packages / build:rpm (rocky9_nocuda) (push) Successful in 13m33s
Build Packages / Build documentation (push) Successful in 34s
Build Packages / build:rpm (rocky9) (push) Successful in 13m54s
Build Packages / build:rpm (ubuntu2404) (push) Successful in 6m59s
Build Packages / Unit tests (push) Successful in 51m27s
XtalOptimizer: Max time is configurable
2025-12-06 16:40:00 +01:00

604 lines
22 KiB
C++
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "XtalOptimizer.h"
#include "ceres/ceres.h"
#include "ceres/rotation.h"
struct XtalResidual {
XtalResidual(double x, double y,
double lambda,
double pixel_size,
const Eigen::Matrix3d &gonio_back_rot,
double exp_h, double exp_k,
double exp_l,
gemmi::CrystalSystem symmetry)
: obs_x(x), obs_y(y),
lambda(lambda),
pixel_size(pixel_size),
exp_h(exp_h),
exp_k(exp_k),
exp_l(exp_l),
gonio_back_rot_(gonio_back_rot),
symmetry(symmetry) {
}
template<typename T>
bool operator()(const T *const beam_x,
const T *const beam_y,
const T *const distance_mm,
const T *const detector_rot,
const T *const p0,
const T *const p1,
const T *const p2,
T *residual) const {
T c1 = ceres::cos(T(detector_rot[0]));
T c2 = ceres::cos(T(detector_rot[1]));
T s1 = ceres::sin(T(detector_rot[0]));
T s2 = ceres::sin(T(detector_rot[1]));
// x_lab in mm
T x_lab = (T(obs_x) - beam_x[0]) * T(pixel_size);
T y_lab = (T(obs_y) - beam_y[0]) * T(pixel_size);
T z_lab = T(distance_mm[0]);
// apply rotations
T x = x_lab * c1 + z_lab * s1;
T y = y_lab * c2 + (-x_lab * s1 + z_lab * c1) * s2;
T z = -y_lab * s2 + (-x_lab * s1 + z_lab * c1) * c2;
// convert to recip space
T lab_norm = ceres::sqrt(x * x + y * y + z * z);
T recip[3];
recip[0] = x / (lab_norm * T(lambda));
recip[1] = y / (lab_norm * T(lambda));
recip[2] = (z / lab_norm - T(1.0)) / T(lambda);
Eigen::Map<const Eigen::Matrix<T, 3, 1>> e_obs_recip_raw(recip);
// Apply goniometer "back-to-start" rotation:
// brings observed reciprocal from image orientation into reference crystal frame
Eigen::Matrix<T, 3, 3> R_gonio_back = gonio_back_rot_.cast<T>();
Eigen::Matrix<T, 3, 1> e_obs_recip = R_gonio_back * e_obs_recip_raw;
Eigen::Matrix<T, 3, 1> e_pred;
Eigen::Matrix<T, 3, 3> e_latt;
T uc_rot_matrix[9];
ceres::AngleAxisToRotationMatrix(p0, uc_rot_matrix);
Eigen::Map<const Eigen::Matrix<T, 3, 3>> e_uc_rot_matrix(uc_rot_matrix);
Eigen::Matrix<T, 3, 1> e_uc_len = Eigen::Matrix<T, 3, 1>::Zero();
Eigen::Matrix<T, 3, 3> B = Eigen::Matrix<T, 3, 3>::Identity();
if (symmetry == gemmi::CrystalSystem::Hexagonal) {
e_uc_len << p1[0], p1[0], p1[2];
B(0, 1) = T(-0.5); // cos(120)
B(1, 1) = T(sqrt(3.0) / 2.0); // sin(120)
} else if (symmetry == gemmi::CrystalSystem::Orthorhombic) {
e_uc_len << p1[0], p1[1], p1[2];
} else if (symmetry == gemmi::CrystalSystem::Tetragonal) {
e_uc_len << p1[0], p1[0], p1[2];
} else if (symmetry == gemmi::CrystalSystem::Cubic) {
e_uc_len << p1[0], p1[0], p1[0];
} else if (symmetry == gemmi::CrystalSystem::Monoclinic) {
// Unique axis b: alpha = gamma = 90°, beta free (angle between a and c)
e_uc_len << p1[0], p1[1], p1[2];
B(0, 2) = ceres::cos(p2[0]);
B(2, 2) = ceres::sin(p2[0]);
} else {
// Triclinic: p1 = (a,b,c), p2 = (alpha, beta, gamma) in radians
const T ca = ceres::cos(p2[0]);
const T cb = ceres::cos(p2[1]);
const T cg = ceres::cos(p2[2]);
const T sg = ceres::sin(p2[2]);
e_uc_len << p1[0], p1[1], p1[2];
// B_triclinic builds lattice from lengths+angles in crystallographic convention
// columns correspond to a, b, c in Cartesian frame prior to global rotation
B(0, 0) = T(1);
B(1, 0) = T(0);
B(2, 0) = T(0);
B(0, 1) = cg;
B(1, 1) = sg;
B(2, 1) = T(0);
// c vector components:
T cx = cb;
T cy = (ca - cb * cg) / sg;
T v = T(1) - cx * cx - cy * cy;
T cz;
if (v >= T(0))
cz = ceres::sqrt(v);
else
cz = T(0);
B(0, 2) = cx;
B(1, 2) = cy;
B(2, 2) = cz;
}
e_latt = e_uc_rot_matrix * e_uc_len.asDiagonal() * B;
Eigen::Matrix<T, 3, 1> e_hkl;
e_hkl << T(exp_h), T(exp_k), T(exp_l);
auto e_pred_hkl = e_latt.transpose() * e_obs_recip;
residual[0] = exp_h - e_pred_hkl[0];
residual[1] = exp_k - e_pred_hkl[1];
residual[2] = exp_l - e_pred_hkl[2];
return true;
}
const double obs_x, obs_y;
const double lambda;
const double pixel_size;
const double exp_h;
const double exp_k;
const double exp_l;
const Eigen::Matrix3d gonio_back_rot_;
gemmi::CrystalSystem symmetry;
};
inline void LatticeToRodriguesAndLengths_GS(const CrystalLattice &latt,
double rod[3],
double lengths[3]) {
// Load lattice columns
const Coord a = latt.Vec0();
const Coord b = latt.Vec1();
const Coord c = latt.Vec2();
Eigen::Vector3d A(a[0], a[1], a[2]);
Eigen::Vector3d B(b[0], b[1], b[2]);
Eigen::Vector3d C(c[0], c[1], c[2]);
// Lengths = column norms (orthorhombic assumption)
lengths[0] = A.norm();
lengths[1] = B.norm();
lengths[2] = C.norm();
auto safe_unit = [](const Eigen::Vector3d &v, double eps = 1e-15) -> Eigen::Vector3d {
double n = v.norm();
return (n > eps) ? (v / n) : Eigen::Vector3d(1.0, 0.0, 0.0);
};
// GramSchmidt with original order: x from A, y from B orthogonalized vs x
Eigen::Vector3d e1 = safe_unit(A);
Eigen::Vector3d y = B - (e1.dot(B)) * e1;
Eigen::Vector3d e2 = safe_unit(y);
// z from cross to ensure right-handed basis
Eigen::Vector3d e3 = e1.cross(e2);
double n3 = e3.norm();
if (n3 < 1e-15) {
// Degenerate case: B nearly collinear with A → use C instead
y = C - (e1.dot(C)) * e1;
e2 = safe_unit(y);
e3 = e1.cross(e2);
n3 = e3.norm();
if (n3 < 1e-15) {
// Still degenerate: pick any perpendicular to e1
e2 = safe_unit((std::abs(e1.x()) < 0.9)
? Eigen::Vector3d::UnitX().cross(e1)
: Eigen::Vector3d::UnitY().cross(e1));
e3 = e1.cross(e2);
}
} else {
e3 /= n3;
}
Eigen::Matrix3d R;
R.col(0) = e1;
R.col(1) = e2;
R.col(2) = e3;
// Convert rotation to Rodrigues (axis * angle)
Eigen::AngleAxisd aa(R);
Eigen::Vector3d r = aa.angle() * aa.axis();
rod[0] = r.x();
rod[1] = r.y();
rod[2] = r.z();
}
void LatticeToRodriguesAndLengths_Hex(const CrystalLattice &latt, double rod[3], double ac[3]) {
const Coord a = latt.Vec0();
const Coord b = latt.Vec1();
const Coord c = latt.Vec2();
Eigen::Vector3d A(a[0], a[1], a[2]);
Eigen::Vector3d B(b[0], b[1], b[2]);
Eigen::Vector3d C(c[0], c[1], c[2]);
const double a_len = A.norm();
const double b_len = B.norm();
const double c_len = C.norm();
ac[0] = (a_len + b_len) / 2.0;
ac[1] = (a_len + b_len) / 2.0;
ac[2] = c_len;
Eigen::Vector3d e1;
Eigen::Vector3d e3;
if (a_len > 0.0)
e1 = A / a_len;
else
e1 = Eigen::Vector3d::UnitX();
if (c_len > 0.0)
e3 = C / c_len;
else
e3 = Eigen::Vector3d::UnitZ();
Eigen::Vector3d e2 = e3.cross(e1);
if (e2.norm() < 1e-15) {
e2 = (std::abs(e1.x()) < 0.9)
? Eigen::Vector3d::UnitX().cross(e1)
: Eigen::Vector3d::UnitY().cross(e1);
}
e2.normalize();
e3 = e1.cross(e2).normalized();
Eigen::Matrix3d R;
R.col(0) = e1;
R.col(1) = e2;
R.col(2) = e3;
Eigen::AngleAxisd aa(R);
Eigen::Vector3d r = aa.angle() * aa.axis();
rod[0] = r.x();
rod[1] = r.y();
rod[2] = r.z();
}
// Extract rotation (Rodrigues), lengths (a,b,c) and beta (rad) for monoclinic (unique axis b).
// Frame choice: e2 aligned with b; e1 from a projected orthogonal to e2; e3 = e1 x e2.
inline void LatticeToRodriguesLengthsBeta_Mono(const CrystalLattice &latt,
double rod[3],
double lengths[3],
double &beta_rad) {
const Coord a = latt.Vec0();
const Coord b = latt.Vec1();
const Coord c = latt.Vec2();
Eigen::Vector3d A(a[0], a[1], a[2]);
Eigen::Vector3d B(b[0], b[1], b[2]);
Eigen::Vector3d C(c[0], c[1], c[2]);
const double a_len = A.norm();
const double b_len = B.norm();
const double c_len = C.norm();
lengths[0] = a_len;
lengths[1] = b_len;
lengths[2] = c_len;
// beta = angle between a and c
double cos_beta = 0.0;
if (a_len > 0.0 && c_len > 0.0)
cos_beta = std::max(-1.0, std::min(1.0, A.dot(C) / (a_len * c_len)));
beta_rad = std::acos(cos_beta);
Eigen::Vector3d e2, ax;
// e2 along b (unique axis)
if (b_len > 0.0)
e2 = B / b_len;
else
e2 = Eigen::Vector3d::UnitY();
if (a_len > 0.0)
ax = A / a_len;
else
ax = Eigen::Vector3d::UnitX();
Eigen::Vector3d e1 = ax - (ax.dot(e2)) * e2;
double n1 = e1.norm();
if (n1 < 1e-15) {
// Fallback: use any perpendicular to e2
e1 = (std::abs(e2.x()) < 0.9
? Eigen::Vector3d::UnitX().cross(e2)
: Eigen::Vector3d::UnitY().cross(e2));
}
e1.normalize();
// e3 completes the right-handed frame
Eigen::Vector3d e3 = e1.cross(e2).normalized();
Eigen::Matrix3d R;
R.col(0) = e1;
R.col(1) = e2;
R.col(2) = e3;
Eigen::AngleAxisd aa(R);
Eigen::Vector3d r = aa.angle() * aa.axis();
rod[0] = r.x();
rod[1] = r.y();
rod[2] = r.z();
}
CrystalLattice AngleAxisAndLengthsToLattice(const double rod[3], const double lengths[3], bool hex) {
const Eigen::Vector3d r(rod[0], rod[1], rod[2]);
const double angle = r.norm();
Eigen::Matrix3d R = Eigen::Matrix3d::Identity();
if (angle > 0.0)
R = Eigen::AngleAxisd(angle, r / angle).toRotationMatrix();
const Eigen::DiagonalMatrix<double, 3> D(lengths[0], lengths[1], lengths[2]);
Eigen::Matrix3d Bhex = Eigen::Matrix3d::Identity();
if (hex) {
Bhex(0, 1) = -1 / 2.0;
Bhex(1, 1) = sqrt(3) / 2;
}
auto latt = R * D * Bhex;
return CrystalLattice(Coord(latt(0, 0), latt(1, 0), latt(2, 0)),
Coord(latt(0, 1), latt(1, 1), latt(2, 1)),
Coord(latt(0, 2), latt(1, 2), latt(2, 2)));
}
inline CrystalLattice AngleAxisLengthsBetaToLattice_Mono(const double rod[3],
const double lengths[3],
double beta_rad) {
const Eigen::Vector3d r(rod[0], rod[1], rod[2]);
const double angle = r.norm();
Eigen::Matrix3d R = Eigen::Matrix3d::Identity();
if (angle > 0.0)
R = Eigen::AngleAxisd(angle, r / angle).toRotationMatrix();
const Eigen::DiagonalMatrix<double, 3> D(lengths[0], lengths[1], lengths[2]);
Eigen::Matrix3d B = Eigen::Matrix3d::Identity();
// Bmono = [[1,0,cosβ],[0,1,0],[0,0,sinβ]]
B(0, 2) = std::cos(beta_rad);
B(2, 2) = std::sin(beta_rad);
Eigen::Matrix3d latt = R * D * B;
return CrystalLattice(Coord(latt(0, 0), latt(1, 0), latt(2, 0)),
Coord(latt(0, 1), latt(1, 1), latt(2, 1)),
Coord(latt(0, 2), latt(1, 2), latt(2, 2)));
}
bool XtalOptimizerInternal(XtalOptimizerData &data,
const std::vector<SpotToSave> &spots,
const float tolerance) {
try {
data.latt.Regularize(data.crystal_system);
Coord vec0 = data.latt.Vec0();
Coord vec1 = data.latt.Vec1();
Coord vec2 = data.latt.Vec2();
double beta = data.latt.GetUnitCell().beta;
// Initial guess for the parameters
double beam_x = data.geom.GetBeamX_pxl();
double beam_y = data.geom.GetBeamY_pxl();
double distance_mm = data.geom.GetDetectorDistance_mm();
double detector_rot[2] = {data.geom.GetPoniRot1_rad(), data.geom.GetPoniRot2_rad()};
ceres::Problem problem;
double latt_vec0[3], latt_vec1[3], latt_vec2[3];
switch (data.crystal_system) {
case gemmi::CrystalSystem::Orthorhombic:
LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
break;
case gemmi::CrystalSystem::Tetragonal:
LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
latt_vec1[0] = (latt_vec1[0] + latt_vec1[1]) / 2.0;
break;
case gemmi::CrystalSystem::Cubic:
LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
latt_vec1[0] = (latt_vec1[0] + latt_vec1[1] + latt_vec1[2]) / 3.0;
break;
case gemmi::CrystalSystem::Hexagonal:
LatticeToRodriguesAndLengths_Hex(data.latt, latt_vec0, latt_vec1);
break;
case gemmi::CrystalSystem::Monoclinic:
LatticeToRodriguesLengthsBeta_Mono(data.latt, latt_vec0, latt_vec1, beta);
latt_vec2[0] = beta;
latt_vec2[1] = 0.0;
latt_vec2[2] = 0.0;
break;
default:
// Triclinic: initialize a,b,c and α,β,γ from current unit cell
LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
auto uc = data.latt.GetUnitCell();
latt_vec2[0] = uc.alpha * M_PI / 180.0;
latt_vec2[1] = uc.beta * M_PI / 180.0;
latt_vec2[2] = uc.gamma * M_PI / 180.0;
break;
}
// Add residuals for each point
for (const auto &pt: spots) {
Eigen::Matrix3d gonio_back_rot = Eigen::Matrix3d::Identity();
Coord recip = data.geom.DetectorToRecip(pt.x, pt.y);
if (data.axis) {
auto rot = data.axis->GetTransformation(pt.image);
recip = rot * recip;
Eigen::Matrix3d Rg;
for (int i = 0; i < 3; ++i) {
for (int j = 0; j < 3; ++j) {
Rg(i, j) = static_cast<double>(rot.arr()[i + j * 3]);
}
}
gonio_back_rot = Rg.transpose();
}
double h_fp = recip * vec0;
double k_fp = recip * vec1;
double l_fp = recip * vec2;
double h = std::round(h_fp);
double k = std::round(k_fp);
double l = std::round(l_fp);
double norm_sq = (h - h_fp) * (h - h_fp) + (k - k_fp) * (k - k_fp) + (l - l_fp) * (l - l_fp);
if (norm_sq > tolerance * tolerance)
continue;
problem.AddResidualBlock(
new ceres::AutoDiffCostFunction<XtalResidual, 3, 1, 1, 1, 2, 3, 3, 3>(
new XtalResidual(pt.x, pt.y,
data.geom.GetWavelength_A(),
data.geom.GetPixelSize_mm(),
gonio_back_rot,
h, k, l,
data.crystal_system)),
nullptr,
&beam_x,
&beam_y,
&distance_mm,
detector_rot,
latt_vec0,
latt_vec1,
latt_vec2
);
}
if (problem.NumResidualBlocks() < data.min_spots)
return false;
if (!data.refine_distance_mm)
problem.SetParameterBlockConstant(&distance_mm);
else {
const double dist_range = 0.1;
problem.SetParameterLowerBound(&distance_mm, 0, distance_mm * (1.0 - dist_range));
problem.SetParameterLowerBound(&distance_mm, 0, distance_mm * (1.0 + dist_range));
}
if (!data.refine_beam_center) {
problem.SetParameterBlockConstant(&beam_x);
problem.SetParameterBlockConstant(&beam_y);
}
if (!data.refine_detector_angles) {
problem.SetParameterBlockConstant(detector_rot);
} else {
const double rot_range = 3.0 / 180.0 * M_PI;
for (int i = 0; i < 2; ++i) {
problem.SetParameterLowerBound(detector_rot, i, detector_rot[i] - rot_range);
problem.SetParameterUpperBound(detector_rot, i, detector_rot[i] + rot_range);
}
}
// Parameter bounds
// Lengths
for (int i = 0; i < 3; ++i) {
problem.SetParameterLowerBound(latt_vec1, i, data.min_length_A);
problem.SetParameterUpperBound(latt_vec1, i, data.max_length_A);
}
if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) {
const double beta_lo = std::max(1e-6, M_PI * (data.min_angle_deg / 180.0));
const double beta_hi = std::min(M_PI - 1e-6, M_PI * (data.max_angle_deg / 180.0));
problem.SetParameterLowerBound(latt_vec2, 0, beta_lo);
problem.SetParameterUpperBound(latt_vec2, 0, beta_hi);
} else if (data.crystal_system == gemmi::CrystalSystem::Triclinic) {
// α, β, γ bounds (radians)
const double alo = M_PI * (data.min_angle_deg / 180.0);
const double ahi = M_PI * (data.max_angle_deg / 180.0);
for (int i = 0; i < 3; ++i) {
problem.SetParameterLowerBound(latt_vec2, i, alo);
problem.SetParameterUpperBound(latt_vec2, i, ahi);
}
}
// Configure solver
ceres::Solver::Options options;
options.linear_solver_type = ceres::DENSE_QR;
options.minimizer_progress_to_stdout = false;
options.max_solver_time_in_seconds = data.max_time;
options.logging_type = ceres::LoggingType::SILENT;
ceres::Solver::Summary summary;
// Run optimization
ceres::Solve(options, &problem, &summary);
if (data.refine_beam_center) {
data.beam_corr_x = data.geom.GetBeamX_pxl() - beam_x;
data.beam_corr_y = data.geom.GetBeamY_pxl() - beam_y;
data.geom.BeamX_pxl(beam_x).BeamY_pxl(beam_y);
}
if (data.refine_distance_mm)
data.geom.DetectorDistance_mm(distance_mm);
if (data.refine_detector_angles)
data.geom.PoniRot1_rad(detector_rot[0]).PoniRot2_rad(detector_rot[1]);
if (data.crystal_system == gemmi::CrystalSystem::Orthorhombic)
data.latt = AngleAxisAndLengthsToLattice(latt_vec0, latt_vec1, false);
else if (data.crystal_system == gemmi::CrystalSystem::Tetragonal) {
latt_vec1[1] = latt_vec1[0];
data.latt = AngleAxisAndLengthsToLattice(latt_vec0, latt_vec1, false);
} else if (data.crystal_system == gemmi::CrystalSystem::Cubic) {
latt_vec1[1] = latt_vec1[0];
latt_vec1[2] = latt_vec1[0];
data.latt = AngleAxisAndLengthsToLattice(latt_vec0, latt_vec1, false);
} else if (data.crystal_system == gemmi::CrystalSystem::Hexagonal) {
latt_vec1[1] = latt_vec1[0];
data.latt = AngleAxisAndLengthsToLattice(latt_vec0, latt_vec1, true);
} else if (data.crystal_system == gemmi::CrystalSystem::Monoclinic) {
data.latt = AngleAxisLengthsBetaToLattice_Mono(latt_vec0, latt_vec1, latt_vec2[0]);
} else {
// Triclinic: reconstruct with generic B from α,β,γ
const Eigen::Vector3d r(latt_vec0[0], latt_vec0[1], latt_vec0[2]);
const double angle = r.norm();
Eigen::Matrix3d R = Eigen::Matrix3d::Identity();
if (angle > 0.0)
R = Eigen::AngleAxisd(angle, r / angle).toRotationMatrix();
Eigen::Matrix3d B = Eigen::Matrix3d::Identity();
const double ca = std::cos(latt_vec2[0]);
const double cb = std::cos(latt_vec2[1]);
const double cg = std::cos(latt_vec2[2]);
const double sg = std::sin(latt_vec2[2]);
// a along x, b in x-y, c general
B(0, 0) = 1.0;
B(1, 0) = 0.0;
B(2, 0) = 0.0;
B(0, 1) = cg;
B(1, 1) = sg;
B(2, 1) = 0.0;
const double cx = cb;
const double cy = (ca - cb * cg) / sg;
const double cz = std::sqrt(std::max(0.0, 1.0 - cx * cx - cy * cy));
B(0, 2) = cx;
B(1, 2) = cy;
B(2, 2) = cz;
Eigen::DiagonalMatrix<double, 3> D(latt_vec1[0], latt_vec1[1], latt_vec1[2]);
Eigen::Matrix3d latt = R * D * B;
data.latt = CrystalLattice(Coord(latt(0, 0), latt(1, 0), latt(2, 0)),
Coord(latt(0, 1), latt(1, 1), latt(2, 1)),
Coord(latt(0, 2), latt(1, 2), latt(2, 2)));
}
data.latt.Regularize(data.crystal_system);
return true;
} catch (...) {
// Convergence problems, likely not updated
return false;
}
}
bool XtalOptimizer(XtalOptimizerData &data, const std::vector<SpotToSave> &spots) {
if (!XtalOptimizerInternal(data, spots, 0.3))
return false;
return XtalOptimizerInternal(data, spots, 0.1);
}