306 lines
12 KiB
C++
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];
|
|
|
|
|
|
}
|
|
}
|
|
}
|
|
|
|
} |