Files
Jungfraujoch/image_analysis/geom_refinement/XtalOptimizer.cpp
2025-10-20 20:43:44 +02:00

552 lines
20 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 beam_x, double beam_y,
double lambda,
double pixel_size, double distance_mm,
double rot1, double rot2,
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),
distance(distance_mm),
rot1(rot1),
rot2(rot2),
beam_x(beam_x),
beam_y(beam_y),
symmetry(symmetry) {
}
template<typename T>
bool operator()(const T *const corr_x,
const T *const corr_y,
const T *const p0,
const T *const p1,
const T *const p2,
T *residual) const {
T c1 = ceres::cos(T(rot1));
T c2 = ceres::cos(T(rot2));
T s1 = ceres::sin(T(rot1));
T s2 = ceres::sin(T(rot2));
// x_lab in mm
T x_lab = (T(obs_x) - beam_x - corr_x[0]) * T(pixel_size);
T y_lab = (T(obs_y) - beam_y - corr_y[0]) * T(pixel_size);
T z_lab = T(distance);
// 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(recip);
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 cz = ceres::sqrt(T(1) - cx * cx - cy * cy);
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 double distance;
const double rot1, rot2;
const double beam_x, beam_y;
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 corr_x = 0;
double corr_y = 0;
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) {
Coord recip = data.geom.DetectorToRecip(pt.x, pt.y);
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, 3, 3, 3>(
new XtalResidual(pt.x, pt.y,
data.geom.GetBeamX_pxl(),
data.geom.GetBeamY_pxl(),
data.geom.GetWavelength_A(),
data.geom.GetPixelSize_mm(),
data.geom.GetDetectorDistance_mm(),
data.geom.GetPoniRot1_rad(),
data.geom.GetPoniRot2_rad(),
h, k, l,
data.crystal_system)),
nullptr,
&corr_x,
&corr_y,
latt_vec0,
latt_vec1,
latt_vec2
);
}
if (problem.NumResidualBlocks() < data.min_spots)
return false;
if (!data.refine_beam_center) {
problem.SetParameterBlockConstant(&corr_x);
problem.SetParameterBlockConstant(&corr_y);
}
// 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.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 = corr_x;
data.beam_corr_y = corr_y;
data.geom.BeamX_pxl(data.geom.GetBeamX_pxl() + corr_x)
.BeamY_pxl(data.geom.GetBeamY_pxl() + corr_y);
}
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);
}