Jungfraujoch/common/DiffractionGeometry.cpp

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <cmath>

#include "DiffractionGeometry.h"
#include "RawToConvertedGeometry.h"

Coord DiffractionGeometry::LabCoord(float x, float y) const {
    Coord detectorCoord = {(x - beam_x_pxl) * pixel_size_mm ,
                           (y - beam_y_pxl) * pixel_size_mm ,
                           det_distance_mm};
    return poni_rot * detectorCoord;
}

std::pair<float, float> DiffractionGeometry::GetDirectBeam_pxl() const {
    return RecipToDector({0,0,0});
}

Coord DiffractionGeometry::GetScatteringVector() const {
    return {0, 0, 1.0f / wavelength_A};
}

Coord DiffractionGeometry::DetectorToRecip(float x, float y) const {
    return LabCoord(x, y).Normalize() / wavelength_A - GetScatteringVector();
}

std::pair<float, float> DiffractionGeometry::RecipToDector(const Coord &recip) const {
    auto S_unrotated = recip + GetScatteringVector();
    auto S = poni_rot.transpose() * S_unrotated;
    if (S.z <= 0)
        return {NAN, NAN};

    float coeff = det_distance_mm / (S.z * pixel_size_mm);
    float x = beam_x_pxl + S.x * coeff;
    float y = beam_y_pxl + S.y * coeff;
    return {x, y};
}

float DiffractionGeometry::TwoTheta_rad(float x, float y) const {
    auto lab = LabCoord(x, y);

    float r = sqrtf(lab.x * lab.x + lab.y * lab.y);
    return atan2f(r, lab.z);
}

float DiffractionGeometry::Phi_rad(float x, float y) const {
    auto lab = LabCoord(x, y);
    auto v = atan2f(lab.y, lab.x);
    if (v < 0)
        v += 2.0f * M_PI;
    return v;
}

float DiffractionGeometry::PxlToRes(float x, float y) const {
    float two_theta = TwoTheta_rad(x, y);
    return wavelength_A / (2.0f * sinf(two_theta/2.0f));
}

float DiffractionGeometry::PxlToQ(float x, float y) const {
    return 2.0f * M_PI / PxlToRes(x,y);
}

float DiffractionGeometry::PxlToRes(float dist_pxl) const {
    // This is agnostic to detector rotation!!!
    if (dist_pxl == 0)
        return INFINITY;
    float tan_2theta = dist_pxl * pixel_size_mm / det_distance_mm;
    float theta = atanf(tan_2theta) / 2.0;
    float d_A = wavelength_A / (2.0f * sinf(theta));
    return d_A;
}

float DiffractionGeometry::ResToPxl(float d_A) const {
    if (d_A == 0)
        return INFINITY;

    float sin_theta = wavelength_A / (2 * d_A);
    float theta = asinf(sin_theta);
    float tan_2theta = tanf(2 * theta);
    return tan_2theta * det_distance_mm / pixel_size_mm;
}

float DiffractionGeometry::DistFromEwaldSphere(const Coord &recip) const {
    auto S = recip + GetScatteringVector();
    return S.Length() - (1.0f/wavelength_A);
}

float DiffractionGeometry::CalcAzIntSolidAngleCorr(float q) const {
    float sin_theta = q * wavelength_A / (4 * static_cast<float>(M_PI));
    float cos_2theta = 1.0f - 2.0f * sin_theta * sin_theta; // cos(2*alpha) = 1 - 2 * sin(alpha)^2
    float cos_2theta_3 = cos_2theta * cos_2theta * cos_2theta;
    return cos_2theta_3;
}

float DiffractionGeometry::CalcAzIntSolidAngleCorr(float x, float y) const {
    float cos_2theta = cosf(TwoTheta_rad(x, y));
    float cos_2theta_3 = cos_2theta * cos_2theta * cos_2theta;
    return cos_2theta_3;
}

float DiffractionGeometry::CalcAzIntPolarizationCorr(float x, float y, float coeff) const {
    auto cos_2theta = cosf(TwoTheta_rad(x, y));
    float cos_2theta_2 = cos_2theta * cos_2theta;
    float cos_2phi = cosf(2.0f * Phi_rad(x, y));
    return 0.5f * (1.0f + cos_2theta_2 - coeff * cos_2phi * (1.0f - cos_2theta_2));
}

float DiffractionGeometry::GetBeamX_pxl() const {
    return beam_x_pxl;
}

float DiffractionGeometry::GetBeamY_pxl() const {
    return beam_y_pxl;
}

float DiffractionGeometry::GetDetectorDistance_mm() const {
    return det_distance_mm;
}

float DiffractionGeometry::GetPixelSize_mm() const {
    return pixel_size_mm;
}

float DiffractionGeometry::GetWavelength_A() const {
    return wavelength_A;
}

DiffractionGeometry &DiffractionGeometry::BeamX_pxl(float input) {
    beam_x_pxl = input;
    return *this;
}

DiffractionGeometry &DiffractionGeometry::BeamY_pxl(float input) {
    beam_y_pxl = input;
    return *this;
}

DiffractionGeometry &DiffractionGeometry::DetectorDistance_mm(float input) {
    if (input < 1.0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Det distance must be above 1.0 mm ");
    det_distance_mm = input;
    return *this;
}

DiffractionGeometry &DiffractionGeometry::PixelSize_mm(float input) {
    if (input <= 0.0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Pixel size must be positive number");
    pixel_size_mm = input;
    return *this;
}

DiffractionGeometry &DiffractionGeometry::Wavelength_A(float input) {
    if (input <= 0.0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Wavelength must be positive number");
    wavelength_A = input;
    return *this;
}

float DiffractionGeometry::AngleFromEwaldSphere_deg(const Coord &p0) const {
    // https://journals.iucr.org/d/issues/2014/08/00/dz5332/index.html
    Coord S0 = GetScatteringVector();

    const float epsilon = 1e-5f;

    float S0_sq = S0 * S0;
    float p0_sq = p0 * p0;
    float S0_p0 = S0 * p0;

    float val = S0_sq * p0_sq - S0_p0 * S0_p0;

    if (fabsf(val) < epsilon)
        return NAN;

    float A = std::sqrt((S0_sq - 1.0f/4.0f * p0_sq) * p0_sq / val);
    float B = (A * S0_p0 + p0_sq / 2.0f) / S0_sq;

    Coord p_star = A * p0 - B * S0;

    return angle_deg(p_star, p0);
}

void DiffractionGeometry::UpdatePoniRotMatrix() {
    poni_rot = RotMatrix(-poni_rot_3, {0,0,1})
               * RotMatrix(-poni_rot_2, {1,0,0})
               * RotMatrix(poni_rot_1, {0,1,0});
}

DiffractionGeometry &DiffractionGeometry::PoniRot1_rad(float input) {
    poni_rot_1 = input;
    UpdatePoniRotMatrix();
    return *this;
}

DiffractionGeometry &DiffractionGeometry::PoniRot2_rad(float input) {
    poni_rot_2 = input;
    UpdatePoniRotMatrix();
    return *this;
}

DiffractionGeometry &DiffractionGeometry::PoniRot3_rad(float input) {
    poni_rot_3 = input;
    UpdatePoniRotMatrix();
    return *this;
}

float DiffractionGeometry::GetPoniRot1_rad() const {
    return poni_rot_1;
}

float DiffractionGeometry::GetPoniRot2_rad() const {
    return poni_rot_2;
}

float DiffractionGeometry::GetPoniRot3_rad() const {
    return poni_rot_3;
}

std::pair<float, float> DiffractionGeometry::ResPhiToPxl(float d_A, float phi_rad) const {
    // Guard invalid inputs
    if (wavelength_A <= 0.0f || d_A <= wavelength_A / 2.0f)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Resolution to high for a given wavelength");

    float sin_theta = wavelength_A / (2.0f * d_A);
    float theta = asinf(sin_theta);
    float k = 1.0f / wavelength_A;

    float s2t = sinf(2.0f * theta);
    float c2t = cosf(2.0f * theta);

    float cphi = cosf(phi_rad);
    float sphi = sinf(phi_rad);

    return RecipToDector(Coord{ k * s2t * cphi,k * s2t * sphi,k * (c2t - 1.0f)});
}

Coord DiffractionGeometry::ProjectToEwaldSphere(const Coord &p0) const {
    Coord S0 = GetScatteringVector();
    Coord S = p0 + S0;
    S = S.Normalize() / wavelength_A;
    return S - S0;
}

const RotMatrix &DiffractionGeometry::GetPoniRotMatrix() const {
    return poni_rot;
}