Files
Jungfraujoch/image_analysis/pixel_refinement/PixelRefine.cpp
T

306 lines
12 KiB
C++

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "PixelRefine.h"
#include <ceres/ceres.h>
#include <ceres/rotation.h>
struct PixelResidual {
// Assume that Itrue and Ibkg are already corrected with solid angle and polarization correction
PixelResidual(double x, double y,
double Itrue, double Iobs,
double Ibkg, double Ibkg_sigma,
double lambda,
double pixel_size,
double angle_rad,
double exp_h, double exp_k,
double exp_l,
gemmi::CrystalSystem symmetry)
: Itrue(Itrue), Iobs(Iobs),
Ibkg(Ibkg), Ibkg_sigma(Ibkg_sigma), obs_x(x),
obs_y(y),
inv_lambda(1.0 / lambda),
pixel_size(pixel_size),
exp_h(exp_h),
exp_k(exp_k),
exp_l(exp_l),
angle_rad(angle_rad),
symmetry(symmetry) {
if (std::fabs(lambda) < 1e-6)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"Lambda cannot be close to zero");
}
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 scale_factor,
const T *const B,
const T *const R,
T *residual) const {
// PyFAI convention (left-handed for rot1/rot2):
// poni_rot = Rz(-rot3) * Rx(-rot2) * Ry(+rot1)
// detector_rot[0] = rot1, detector_rot[1] = rot2 (rot3 = 0 assumed)
const T rot1 = detector_rot[0];
const T rot2 = detector_rot[1];
// Ry(+rot1): rotation around Y-axis
const T c1 = ceres::cos(rot1);
const T s1 = ceres::sin(rot1);
// Rx(-rot2): rotation around X-axis with inverted sign (PyFAI left-handed)
const T c2 = ceres::cos(rot2);
const T s2 = ceres::sin(rot2);
// Detector coordinates in mm
const T det_x = (T(obs_x) - beam[0]) * T(pixel_size);
const T det_y = (T(obs_y) - beam[1]) * T(pixel_size);
const T det_z = T(distance_mm[0]);
// Apply Ry(rot1) first: rotate around Y
const T t1_x = c1 * det_x + s1 * det_z;
const T t1_y = det_y;
const T t1_z = -s1 * det_x + c1 * det_z;
// Then apply Rx(-rot2): rotate around X
const T x = t1_x;
const T y = c2 * t1_y + s2 * t1_z;
const T z = -s2 * t1_y + c2 * t1_z;
// convert to recip space
const T lab_norm = ceres::sqrt(x * x + y * y + z * z);
const T inv_norm = T(1) / lab_norm;
T recip_raw[3];
recip_raw[0] = x * inv_norm * T(inv_lambda);
recip_raw[1] = y * inv_norm * T(inv_lambda);
recip_raw[2] = (z * inv_norm - T(1.0)) * T(inv_lambda);
// Apply goniometer "back-to-start" rotation:
// brings observed reciprocal from image orientation into reference crystal frame
const T aa_back[3] = {
T(angle_rad) * rotation_axis[0],
T(angle_rad) * rotation_axis[1],
T(angle_rad) * rotation_axis[2]
};
T recip_obs[3];
ceres::AngleAxisRotatePoint(aa_back, recip_raw, recip_obs);
const Eigen::Matrix<T, 3, 1> e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]);
// Build unit cell lengths and B (convention: columns are a, b, c prior to global rotation)
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(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:
const T cx = cb;
const T cy = (ca - cb * cg) / sg;
const T v = T(1) - cx * cx - cy * cy;
const T cz = (v >= T(0)) ? ceres::sqrt(v) : T(0);
B(0, 2) = cx;
B(1, 2) = cy;
B(2, 2) = cz;
}
// Build unrotated direct lattice columns: (B * D), then rotate them by p0.
// This avoids AngleAxisToRotationMatrix + matrix multiplications.
const T L0 = e_uc_len[0];
const T L1 = e_uc_len[1];
const T L2 = e_uc_len[2];
T col0_unrot[3] = {B(0, 0) * L0, B(1, 0) * L0, B(2, 0) * L0};
T col1_unrot[3] = {B(0, 1) * L1, B(1, 1) * L1, B(2, 1) * L1};
T col2_unrot[3] = {B(0, 2) * L2, B(1, 2) * L2, B(2, 2) * L2};
T col0_rot[3], col1_rot[3], col2_rot[3];
ceres::AngleAxisRotatePoint(p0, col0_unrot, col0_rot);
ceres::AngleAxisRotatePoint(p0, col1_unrot, col1_rot);
ceres::AngleAxisRotatePoint(p0, col2_unrot, col2_rot);
const Eigen::Matrix<T, 3, 1> A(col0_rot[0], col0_rot[1], col0_rot[2]);
const Eigen::Matrix<T, 3, 1> Bv(col1_rot[0], col1_rot[1], col1_rot[2]);
const Eigen::Matrix<T, 3, 1> C(col2_rot[0], col2_rot[1], col2_rot[2]);
const Eigen::Matrix<T, 3, 1> BxC = Bv.cross(C);
const Eigen::Matrix<T, 3, 1> CxA = C.cross(A);
const Eigen::Matrix<T, 3, 1> AxB = A.cross(Bv);
const T V = A.dot(BxC);
const T invV = T(1) / V;
const Eigen::Matrix<T, 3, 1> Astar = BxC * invV;
const Eigen::Matrix<T, 3, 1> Bstar = CxA * invV;
const Eigen::Matrix<T, 3, 1> Cstar = AxB * invV;
const T h = T(exp_h);
const T k = T(exp_k);
const T l = T(exp_l);
const Eigen::Matrix<T, 3, 1> e_pred_recip = Astar * h + Bstar * k + Cstar * l;
// Ewald sphere centre is at -k_i = (0, 0, -inv_lambda)
// Radial direction: outward normal at g_hkl
const Eigen::Matrix<T, 3, 1> S_pred(
e_pred_recip[0],
e_pred_recip[1],
e_pred_recip[2] + T(inv_lambda) // g_hkl + k_i
);
const T S_pred_norm = S_pred.norm();
if (S_pred_norm < T(1e-10))
return T(0);
const Eigen::Matrix<T, 3, 1> n_radial = S_pred / S_pred_norm;
const Eigen::Matrix<T, 3, 1> delta_q = e_obs_recip - e_pred_recip;
const T eps_radial = delta_q.dot(n_radial);
const Eigen::Matrix<T, 3, 1> dq_tang = delta_q - eps_radial * n_radial;
const T eps_tangential_sq = dq_tang.squaredNorm(); // guaranteed ≥ 0
// ─────────────────────────────────────────────────────────────
const T B_term = ceres::exp(- B[0] * e_pred_recip.squaredNorm() / 4.0);
// Need to normalize by R[0] and R[1]
const T partiality = ceres::exp(- eps_radial * eps_radial / (R[0] * R[0]) - eps_tangential_sq / (R[1] * R[1]));
const T Ipred = partiality * Itrue * scale_factor[0] * B_term - Ibkg;
// Need to weight by sigma
// I would like to use sigma based on Ipred and Ibkg_sigma - need to come up with a better approach
residual[0] = (Ipred - Iobs) / Ibkg_sigma;
return true;
}
const double Itrue, Iobs, Ibkg, Ibkg_sigma;
const double obs_x, obs_y;
const double inv_lambda;
const double pixel_size;
const double exp_h;
const double exp_k;
const double exp_l;
const double angle_rad;
gemmi::CrystalSystem symmetry;
};
PixelRefine::PixelRefine(const DiffractionExperiment &experiment,
const AzimuthalIntegrationMapping &mapping,
const std::vector<MergedReflection> &reference,
BraggPrediction &prediction)
: prediction(prediction),
mapping(mapping),
xpixel(experiment.GetXPixelsNum()),
ypixel(experiment.GetYPixelsNum()),
experiment(experiment),
hkl_key_generator(experiment.GetScalingSettings().GetMergeFriedel(),
experiment.GetSpaceGroupNumber().value_or(1)) {
for (const auto &ref: reference)
reference_data[hkl_key_generator(ref)] = ref.I;
}
template<class T>
void PixelRefine::Run(const T *image,
const AzimuthalIntegrationProfile &profile,
PixelRefineData &data) {
ceres::Problem problem;
// We predict reflections based on initial geometry and default settings
// To be tuned later
const BraggPredictionSettings settings_prediction{
.high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
.max_hkl = 100,
.centering = data.centering
};
prediction.Calc(experiment, data.latt, settings_prediction);
auto azim_result = profile.GetResult();
auto azim_std = profile.GetStd();
// For each reflection we select some area (3-5 pixels around it)
const int radius = 3;
for (const auto &refl : prediction.GetReflections()) {
auto hkl = hkl_key_generator(refl);
// We only handle reflections that are present in the reference set
if (!reference_data.contains(hkl))
continue;
const double I_true = reference_data[hkl];
int min_y = std::max<int>(refl.predicted_y - radius, 0);
int max_y = std::min<int>(refl.predicted_y + radius, ypixel - 1);
int min_x = std::max<int>(refl.predicted_x - radius, 0);
int max_x = std::min<int>(refl.predicted_x + radius, xpixel - 1);
for (int y = min_y; y <= max_y; ++y) {
for (int x = min_x; x <= max_x; ++x) {
const size_t npixel = xpixel * y + x;
int azim_bin = mapping.GetPixelToBin()[npixel];
// If pixel is not mapped to azimuthal bin
// or pixel has special value (lowest/highest integer)
// it should be ignored for the purpose of this try
// We should check if pixel mask is needed, but for most workflows it is already applied
if (azim_bin >= mapping.GetAzimuthalBinCount())
continue;
if (image[npixel] == std::numeric_limits<T>::max())
continue;
if (std::is_signed_v<T>() && (image[npixel] == std::numeric_limits<T>::min()))
continue;
// Get per-pixel polarization and solid angle correction for the pixel from the AzimuthalIntegrationMapping
// Warning! this is missing Lorentz correction, but we don't worry at the moment about it
// Important is -> this correction is also applied to background, so we must be consistent here
float correction = mapping.Corrections()[npixel];
// Get mean pixel value for background in the azimuthal bin + sigma
float bkg_value = azim_result[azim_bin];
float bkg_sigma = azim_std[azim_bin];
float pixel_value = image[npixel];
}
}
}
}