ScaleOnTheFly: WIP: ScalingPostRefResidual

image_analysis/scale_merge/ScaleOnTheFly.cpp

@@ -104,6 +104,151 @@ namespace {
 
         double partiality;
     };
+
+struct ScalingPostRefResidual : public ScalingResidual {
+    // In this algorithm at the point of post-refinement we don't anymore care for where maximum of
+    // the reflection was located and if it fits the observed position.
+    // This reflections was already integrated and we cannot integrate it better at this point.
+    // But, we could adjust partiality to indicate that this reflection was wrongly predicted.
+    // I.e., integrated position was far away from true reflection, so partiality must be low.
+    // This is an empiric model and need to see if this will work in practice at all.
+    // I hope it will allow the model to find that reflections were misindexed.
+    // We assume at this point that initial indexing was done properly and integration was generally OK => most low resolution reflections fit correctly
+    // Yet we know, that small errors in indexing are inducing misalignment at high resolution - sometimes it is visible that high-resolution reflections
+    // are away from shoe-boxes observed in the image, if we can catch this at post-refinement/scaling step, this would be great.
+    // Next logical step is to do this pixel-wise - for each pixel refine partiality and merge pixels
+    // This should work for per-image scaling, or even, maybe, for full rotation datasets (3600 images)
+    // Then we could properly take into account misalignment of shoe-box center vs. partiality and also remove most pixels
+    // in the shoe-box that don't really contribute to the reflection.
+    // But...it could also drift to downweighting partiality for all high resolution reflections to make loss function "fake happy".
+
+    // We assume rot3 == 0. Rot3 is not really helping much in crystallography (other than fixing polarization correction)
+    ScalingPostRefResidual(const Reflection &r, double Itrue, double sigma,
+                           const DiffractionGeometry &geometry,
+                           const CrystalLattice &lattice,
+                           double exp_h, double exp_k,
+                           double exp_l)
+    : ScalingResidual(r, Itrue, sigma),
+          integration_center_x(r.predicted_x),
+          integration_center_y(r.predicted_y),
+          inv_lambda(SafeInv(geometry.GetWavelength_A(), 0.0)),
+          pixel_size(geometry.GetPixelSize_mm()),
+          det_dist_mm(geometry.GetDetectorDistance_mm()),
+          beam_x(geometry.GetBeamX_pxl()),
+          beam_y(geometry.GetBeamY_pxl()),
+          exp_h(exp_h),
+          exp_k(exp_k),
+          exp_l(exp_l),
+          Astar(lattice.Astar()),
+          Bstar(lattice.Bstar()),
+          Cstar(lattice.Cstar()),
+          c1(std::cos(geometry.GetPoniRot1_rad())),
+          s1(std::sin(geometry.GetPoniRot1_rad())),
+          c2(std::cos(geometry.GetPoniRot2_rad())),
+          s2(std::sin(geometry.GetPoniRot2_rad())) {
+    }
+
+    template <typename T>
+    T CalcPartiality(const T *const R_radial,
+                    const T *const R_tangential,
+                    const T *const beam_corr,
+                    const T *const p0) const {
+        // Detector coordinates in mm
+        const T det_x = (T(integration_center_x) - beam_x - beam_corr[0]) * T(pixel_size);
+        const T det_y = (T(integration_center_y) - beam_y - beam_corr[1]) * T(pixel_size);
+        const T det_z = T(det_dist_mm);
+
+        // Apply Ry(rot1) first: rotate around Y
+        const T t1_x = T(c1) * det_x + T(s1) * det_z;
+        const T t1_y = det_y;
+        const T t1_z = T(-s1) * det_x + T(c1) * det_z;
+
+        // Then apply Rx(-rot2): rotate around X
+        const T x = t1_x;
+        const T y = T(c2) * t1_y + T(s2) * t1_z;
+        const T z = - T(s2) * t1_y + T(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_obs[3];
+        recip_obs[0] = x * inv_norm * inv_lambda;
+        recip_obs[1] = y * inv_norm * inv_lambda;
+        recip_obs[2] = (z * inv_norm - T(1.0)) * inv_lambda;
+
+        const Eigen::Matrix<T, 3, 1> e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]);
+
+        const T astar_unrot[3] = {T(Astar.x), T(Astar.y), T(Astar.z)};
+        const T bstar_unrot[3] = {T(Bstar.x), T(Bstar.y), T(Bstar.z)};
+        const T cstar_unrot[3] = {T(Cstar.x), T(Cstar.y), T(Cstar.z)};
+
+        T astar_rot[3], bstar_rot[3], cstar_rot[3];
+
+        ceres::AngleAxisRotatePoint(p0, astar_unrot, astar_rot);
+        ceres::AngleAxisRotatePoint(p0, bstar_unrot, bstar_rot);
+        ceres::AngleAxisRotatePoint(p0, cstar_unrot, cstar_rot);
+
+        const Eigen::Matrix<T, 3, 1> e_pred_recip(T(exp_h) * astar_rot[0] + T(exp_k) * bstar_rot[0] + T(exp_l) * cstar_rot[0],
+            T(exp_h) * astar_rot[1] + T(exp_k) * bstar_rot[1] + T(exp_l) * cstar_rot[1],
+            T(exp_h) * astar_rot[2] + T(exp_k) * bstar_rot[2] + T(exp_l) * cstar_rot[2]
+        );
+
+        // 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
+        // ─────────────────────────────────────────────────────────────
+
+        return ceres::exp(
+            - eps_radial * eps_radial / (R_radial[0]    * R_radial[0])
+            - eps_tangential_sq       / (R_tangential[0] * R_tangential[0])
+        );
+    }
+
+    template<typename T>
+    bool operator()(const T *const G,
+                    const T *const B,
+                    const T *const R_radial,
+                    const T *const R_tangential,
+                    const T *const beam_corr,
+                    const T *const p0,
+                    T *residual) const {
+        if (R_radial[0] < T(1e-10) || R_tangential[0] < T(1e-10))
+            return false;
+
+        const T B_term = ceres::exp(B[0] * T(b_resolution_coeff));
+
+        const T partiality = CalcPartiality(R_radial, R_tangential, beam_corr, p0);
+        residual[0] = (G[0] * partiality * B_term * T(lp) * T(Itrue)
+                       - T(Iobs)) * T(weight);
+        return true;
+    }
+
+    const double integration_center_x, integration_center_y;
+    const double inv_lambda;
+    const double pixel_size;
+    const double det_dist_mm;
+    const double beam_x, beam_y;
+    const double exp_h;
+    const double exp_k;
+    const double exp_l;
+    const Coord Astar, Bstar, Cstar;
+    const double c1,s1,c2,s2;
+};
 }
 
 ScaleOnTheFly::ScaleOnTheFly(const DiffractionExperiment &x, const std::vector<MergedReflection> &ref)