BraggPrediction: Add partiality calculation

image_analysis/IndexAndRefine.cpp

@@ -167,11 +167,14 @@ void IndexAndRefine::QuickPredictAndIntegrate(DataMessage &msg,
     if (experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA())
         ewald_dist_cutoff = experiment.GetBraggIntegrationSettings().GetFixedProfileRadius_recipA().value() * 3.0f;
 
-    float max_angle_deg = 0.2f;
+    float wedge_deg = 0.0f;
+    float mos_deg = 0.1f;
+
     if (experiment.GetGoniometer().has_value()) {
-        max_angle_deg = experiment.GetGoniometer()->GetWedge_deg() / 2.0;
+        wedge_deg = experiment.GetGoniometer()->GetWedge_deg() / 2.0;
+
         if (msg.mosaicity_deg)
-            max_angle_deg += msg.mosaicity_deg.value();
+            mos_deg = msg.mosaicity_deg.value();
     }
 
     const BraggPredictionSettings settings_prediction{
@@ -179,7 +182,8 @@ void IndexAndRefine::QuickPredictAndIntegrate(DataMessage &msg,
         .ewald_dist_cutoff = ewald_dist_cutoff,
         .max_hkl = 100,
         .centering = outcome.symmetry.centering,
-        .max_angle_deg = max_angle_deg
+        .wedge_deg = std::fabs(wedge_deg),
+        .mosaicity_deg = std::fabs(mos_deg)
     };
 
     auto nrefl = prediction.Calc(outcome.experiment, latt, settings_prediction);

image_analysis/bragg_prediction/BraggPrediction.h

@@ -15,7 +15,8 @@ struct BraggPredictionSettings {
     float ewald_dist_cutoff = 0.0005;
     int max_hkl = 100;
     char centering = 'P';
-    float max_angle_deg = 0.2f;
+    float wedge_deg = 0.1f;
+    float mosaicity_deg = 0.2f;
 };
 
 class BraggPrediction {

image_analysis/bragg_prediction/BraggPredictionRot.cpp

@@ -45,7 +45,10 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys
     const float m3_S0 = m3 * S0;
 
     int i = 0;
-    const float max_angle_rad = settings.max_angle_deg * static_cast<float>(M_PI) / 180.f;
+
+    const float mos_angle_rad = settings.mosaicity_deg * static_cast<float>(M_PI) / 180.f;
+    const float wedge_angle_rad = settings.wedge_deg * static_cast<float>(M_PI) / 180.f;
+    const float max_angle_rad = wedge_angle_rad / 2 + mos_angle_rad;
 
     for (int h = -settings.max_hkl; h <= settings.max_hkl; h++) {
         // Precompute A* h contribution
@@ -93,11 +96,18 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys
 
                     float phi = -1.0f * std::atan2(sinphi, cosphi);
 
-                    if (phi > max_angle_rad || phi < -max_angle_rad)
+                    if (phi > wedge_angle_rad || phi < -max_angle_rad)
                         continue;
 
-                    const float lorentz_reciprocal = std::fabs(m2 * (S % S0))/(S*S0);
+                    const Coord e1 = (S % S0).Normalize();
 
+                    const float zeta_abs = std::fabs(m2 * e1);
+
+                    const float lorentz_reciprocal = std::fabs(m2 * (S % S0))/(S*S0);
+                    const float c1 = std::sqrt(2.0f) * mos_angle_rad / zeta_abs;
+
+                    const float partiality = (std::erf(zeta_abs * (phi + wedge_angle_rad / 2.0f) * c1)
+                        - std::erf(zeta_abs * (phi - wedge_angle_rad / 2.0f) * c1)) / 2.0f;
                     // Inlined RecipToDector with rot1 and rot2 (rot3 = 0)
                     // Apply rotation matrix transpose
                     float S_rot_x = rot[0] * S.x + rot[1] * S.y + rot[2] * S.z;
@@ -127,7 +137,8 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys
                         .rlp = lorentz_reciprocal,
                         .S_x = S.x,
                         .S_y = S.y,
-                        .S_z = S.z
+                        .S_z = S.z,
+                        .partiality = partiality
                     };
                     i++;
                 }

image_analysis/bragg_prediction/BraggPredictionRotGPU.cu

@@ -99,15 +99,33 @@ namespace {
             float Sy = C.S0.y + py;
             float Sz = C.S0.z + pz;
 
-            float phi = -1.0f * atan2f(sinphi, cosphi) * 180.0f / static_cast<float>(M_PI);
-            if (phi > C.max_angle_deg || phi < -C.max_angle_deg)
+            float phi = -1.0f * atan2f(sinphi, cosphi);
+            float phi_deg = phi * 180.0f / static_cast<float>(M_PI);
+            float max_angle_rad = C.mos_angle_rad + C.wedge_angle_rad / 2.0f;
+
+            if (phi_deg > max_angle_rad || phi_deg < -max_angle_rad)
                 continue;
 
+            // e1 = normalize(S x S0) - direction perpendicular to both S and S0
+            float e1x, e1y, e1z;
+            cross3(Sx, Sy, Sz, C.S0.x, C.S0.y, C.S0.z, e1x, e1y, e1z);
+            normalize3(e1x, e1y, e1z);
+
+            // zeta = |m2 · e1| - the "lorentz-like" geometric factor for partiality
+            float zeta_abs = fabsf(dot3(C.m2.x, C.m2.y, C.m2.z, e1x, e1y, e1z));
+
             float cx, cy, cz;
             cross3(Sx, Sy, Sz, C.S0.x, C.S0.y, C.S0.z, cx, cy, cz);
             float lorentz = fabsf(dot3(C.m2.x, C.m2.y, C.m2.z, cx, cy, cz)) /
                             dot3(Sx, Sy, Sz, C.S0.x, C.S0.y, C.S0.z);
 
+            // Partiality calculation (Kabsch et al.)
+            // c1 = sqrt(2) * sigma / zeta, where sigma = mosaicity
+            float c1 = sqrtf(2.0f) * C.mos_angle_rad / zeta_abs;
+            float half_wedge = C.wedge_angle_rad / 2.0f;
+            float partiality = (erff(zeta_abs * (phi + half_wedge) * c1)
+                              - erff(zeta_abs * (phi - half_wedge) * c1)) / 2.0f;
+
             // Use S (rotated) for projection
             float Srx = C.rot[0] * Sx + C.rot[1] * Sy + C.rot[2] * Sz;
             float Sry = C.rot[3] * Sx + C.rot[4] * Sy + C.rot[5] * Sz;
@@ -135,6 +153,7 @@ namespace {
             out[count].S_x = Sx;
             out[count].S_y = Sy;
             out[count].S_z = Sz;
+            out[count].partiality = partiality;
             count++;
         }
         return count;
@@ -167,35 +186,40 @@ namespace {
     }
 
     inline KernelConstsRot BuildKernelConstsRot(const DiffractionExperiment &experiment,
-                                                const CrystalLattice &lattice,
-                                                float high_res_A,
-                                                float max_angle_deg,
-                                                char centering) {
-        KernelConstsRot kc{};
-        auto geom = experiment.GetDiffractionGeometry();
+                                            const CrystalLattice &lattice,
+                                            float high_res_A,
+                                            float mos_angle_deg,
+                                            float wedge_angle_deg,
+                                            char centering) {
+         KernelConstsRot kc{};
+         auto geom = experiment.GetDiffractionGeometry();
 
-        kc.det_width_pxl = static_cast<float>(experiment.GetXPixelsNum());
-        kc.det_height_pxl = static_cast<float>(experiment.GetYPixelsNum());
-        kc.beam_x = geom.GetBeamX_pxl();
-        kc.beam_y = geom.GetBeamY_pxl();
-        kc.coeff_const = geom.GetDetectorDistance_mm() / geom.GetPixelSize_mm();
+         kc.det_width_pxl = static_cast<float>(experiment.GetXPixelsNum());
+         kc.det_height_pxl = static_cast<float>(experiment.GetYPixelsNum());
+         kc.beam_x = geom.GetBeamX_pxl();
+         kc.beam_y = geom.GetBeamY_pxl();
+         kc.coeff_const = geom.GetDetectorDistance_mm() / geom.GetPixelSize_mm();
 
-        float one_over_dmax = 1.0f / high_res_A;
-        kc.one_over_dmax_sq = one_over_dmax * one_over_dmax;
-        kc.one_over_wavelength = 1.0f / geom.GetWavelength_A();
-        kc.max_angle_deg = max_angle_deg;
+         float one_over_dmax = 1.0f / high_res_A;
+         kc.one_over_dmax_sq = one_over_dmax * one_over_dmax;
+         kc.one_over_wavelength = 1.0f / geom.GetWavelength_A();
+         kc.max_angle_deg = wedge_angle_deg / 2.0f + mos_angle_deg;
 
-        kc.Astar = lattice.Astar();
-        kc.Bstar = lattice.Bstar();
-        kc.Cstar = lattice.Cstar();
-        kc.S0 = geom.GetScatteringVector();
+         // Store mosaicity and wedge in radians for partiality calculation
+         kc.mos_angle_rad = mos_angle_deg * static_cast<float>(M_PI) / 180.0f;
+         kc.wedge_angle_rad = wedge_angle_deg * static_cast<float>(M_PI) / 180.0f;
 
-        auto rotT = geom.GetPoniRotMatrix().transpose().arr();
-        for (int i = 0; i < 9; ++i) kc.rot[i] = rotT[i];
+         kc.Astar = lattice.Astar();
+         kc.Bstar = lattice.Bstar();
+         kc.Cstar = lattice.Cstar();
+         kc.S0 = geom.GetScatteringVector();
 
-        kc.centering = centering;
-        return kc;
-    }
+         auto rotT = geom.GetPoniRotMatrix().transpose().arr();
+         for (int i = 0; i < 9; ++i) kc.rot[i] = rotT[i];
+
+         kc.centering = centering;
+         return kc;
+     }
 
     inline void BuildGoniometerBasis(const DiffractionExperiment &experiment, KernelConstsRot &kc) {
         const auto gon_opt = experiment.GetGoniometer();
@@ -237,7 +261,10 @@ BraggPredictionRotGPU::BraggPredictionRotGPU(int max_reflections)
 int BraggPredictionRotGPU::Calc(const DiffractionExperiment &experiment,
                                 const CrystalLattice &lattice,
                                 const BraggPredictionSettings &settings) {
-    KernelConstsRot hK = BuildKernelConstsRot(experiment, lattice, settings.high_res_A, settings.max_angle_deg,
+    KernelConstsRot hK = BuildKernelConstsRot(experiment, lattice,
+                                              settings.high_res_A,
+                                              settings.mosaicity_deg,
+                                              settings.wedge_deg,
                                               settings.centering);
     BuildGoniometerBasis(experiment, hK);
 

image_analysis/bragg_prediction/BraggPredictionRotGPU.h

@@ -16,7 +16,8 @@ struct KernelConstsRot {
     float coeff_const;
     float one_over_wavelength;
     float one_over_dmax_sq;
-    float max_angle_deg;
+    float mos_angle_rad;
+    float wedge_angle_rad;
     Coord Astar, Bstar, Cstar, S0;
     Coord m1, m2, m3;
     float m2_S0;