jfjoch_process: optional absorption surface for rot3d scaling (--absorption)

Adds an opt-in smooth absorption correction for rotation scaling. After the
rot3d fulls are scaled, --absorption[=num] fits a multiplicative surface
A(s1_crystal) - a degree<=4 monomial basis (real spherical harmonics up to l=4,
as XDS/DIALS) of the diffracted-beam direction in the crystal/goniometer frame,
by ridge-regularized log-linear least-squares of I_scaled/I_ref weighted by
(I/sigma)^2, over num iterations (default 3); the surface divides image_scale_corr
and the fulls are re-merged.

Off by default and a no-op without rot3d. On the test panel (~13 keV, thin
crystals) it is metric-neutral - fitted rms(log A) ~3-4%, ISa/CC1/2 unchanged -
because absorption is negligible there and the per-frame scale G(phi) already
absorbs the angular part. It is kept as a lever for low-energy data (e.g. 6 keV)
where absorption becomes significant. Stored as ScalingSettings::absorption_iter.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

common/ScalingSettings.cpp

@@ -149,6 +149,17 @@ bool ScalingSettings::GetScaleFulls() const {
     return scale_fulls;
 }
 
+ScalingSettings &ScalingSettings::AbsorptionIter(int input) {
+    if (input < 0)
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Absorption iterations must be non-negative");
+    absorption_iter = input;
+    return *this;
+}
+
+int ScalingSettings::GetAbsorptionIter() const {
+    return absorption_iter;
+}
+
 ScalingSettings &ScalingSettings::SmoothGWindow(int input) {
     if (input < 0)
         throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Smooth-G window must be non-negative");

common/ScalingSettings.h

@@ -39,6 +39,11 @@ class ScalingSettings {
     // order). A no-op without rot3d.
     bool scale_fulls = false;
 
+    // Absorption surface: after scaling the rot3d fulls, fit a smooth multiplicative correction
+    // over the diffracted-beam direction in the crystal frame, refined over this many iterations.
+    // 0 = off. Negligible at hard X-rays / thin crystals, but matters at low energy. No-op without rot3d.
+    int absorption_iter = 0;
+
     // Smooth the per-frame scale G across frames (centered moving average of log G over this odd
     // window) before the rot3d combine, so a rocking event's partials share a consistent scale.
     // 0 = off. A no-op without rot3d.
@@ -63,6 +68,7 @@ public:
     ScalingSettings& OutlierRejectNsigma(double input);
     ScalingSettings& Combine3D(bool input);
     ScalingSettings& ScaleFulls(bool input);
+    ScalingSettings& AbsorptionIter(int input);
     ScalingSettings& SmoothGWindow(int input);
 
     ScalingSettings& RfreeFraction(double input);
@@ -95,6 +101,7 @@ public:
     [[nodiscard]] double GetOutlierRejectNsigma() const;
     [[nodiscard]] bool GetCombine3D() const;
     [[nodiscard]] bool GetScaleFulls() const;
+    [[nodiscard]] int GetAbsorptionIter() const;
     [[nodiscard]] int GetSmoothGWindow() const;
 
     [[nodiscard]] double GetRfreeFraction() const;

process/JFJochProcess.cpp

@@ -32,7 +32,11 @@
 #include "../image_analysis/scale_merge/SearchSpaceGroup.h"
 #include "../image_analysis/scale_merge/TwinningAnalysis.h"
 #include "../image_analysis/scale_merge/Combine3D.h"
+#include "../image_analysis/scale_merge/HKLKey.h"
 #include "../image_analysis/WriteReflections.h"
+#include <array>
+#include <map>
+#include <Eigen/Dense>
 
 namespace {
     // Pick up to requested_images ordinals spread evenly across [0, images_to_process) for the
@@ -79,6 +83,85 @@ namespace {
         logger.Info("Scaled fulls (XDS order, Unity model)");
     }
 
+    // Absorption surface. A smooth multiplicative correction A(s1_crystal) shared across the sweep,
+    // fit by ridge-regularized log-linear least-squares of I_scaled/I_ref against a degree<=4
+    // monomial basis of the diffracted-beam direction in the crystal (goniometer) frame - the same
+    // function space as real spherical harmonics up to l=4, as XDS/DIALS use. Applied by dividing
+    // image_scale_corr by A, then re-merged. Negligible at hard X-rays / thin crystals but matters
+    // at low energy (e.g. 6 keV).
+    void AbsorptionSurface(const DiffractionExperiment &exp,
+                           std::vector<IntegrationOutcome> &fulls, int n_iter, Logger &logger) {
+        const auto gon_opt = exp.GetGoniometer();
+        if (!gon_opt) { logger.Warning("Absorption: no goniometer, skipping"); return; }
+        const auto &gon = *gon_opt;
+        const Coord axis = gon.GetAxis().Normalize();
+        const DiffractionGeometry geom = exp.GetDiffractionGeometry();
+        const HKLKeyGenerator keygen(exp.GetScalingSettings().GetMergeFriedel(),
+                                     exp.GetSpaceGroupNumber().value_or(1));
+        constexpr int MAXDEG = 4;
+        constexpr int NB = (MAXDEG + 1) * (MAXDEG + 2) * (MAXDEG + 3) / 6;  // # monomials of degree 0..4
+        auto basis = [&](const Coord &s, std::array<double, NB> &b) {
+            int n = 0;
+            for (int deg = 0; deg <= MAXDEG; ++deg)
+                for (int i = deg; i >= 0; --i)
+                    for (int j = deg - i; j >= 0; --j) {
+                        const int kk = deg - i - j;
+                        double v = 1.0;
+                        for (int a = 0; a < i; ++a) v *= s.x;
+                        for (int a = 0; a < j; ++a) v *= s.y;
+                        for (int a = 0; a < kk; ++a) v *= s.z;
+                        b[n++] = v;
+                    }
+        };
+        auto s1_crystal = [&](const Reflection &r) {
+            const Coord s1lab = geom.LabCoord(r.predicted_x, r.predicted_y).Normalize();
+            const double phi = gon.GetAngle_deg(r.image_number) * 3.14159265358979323846 / 180.0;
+            return RotMatrix(static_cast<float>(-phi), axis) * s1lab;
+        };
+
+        for (int iter = 0; iter < n_iter; ++iter) {
+            const auto ref = MergeAll(exp, fulls, false);
+            std::map<HKLKey, float> refI;
+            for (const auto &m : ref) refI[keygen(m)] = m.I;
+
+            Eigen::MatrixXd AtA = Eigen::MatrixXd::Zero(NB, NB);
+            Eigen::VectorXd Atb = Eigen::VectorXd::Zero(NB);
+            std::array<double, NB> b{};
+            for (const auto &oc : fulls)
+                for (const auto &r : oc.reflections) {
+                    if (!r.observed || !std::isfinite(r.image_scale_corr) || r.image_scale_corr <= 0.0f) continue;
+                    const double Iscaled = static_cast<double>(r.I) * r.image_scale_corr;
+                    const double sig = static_cast<double>(r.sigma) * r.image_scale_corr;
+                    if (!(Iscaled > 0.0) || !(sig > 0.0)) continue;
+                    const auto it = refI.find(keygen(r));
+                    if (it == refI.end() || !(it->second > 0.0f)) continue;
+                    basis(s1_crystal(r), b);
+                    const double y = std::log(Iscaled / it->second);
+                    const double w = std::min(100.0, (Iscaled / sig) * (Iscaled / sig));  // strong reflections define the surface
+                    for (int a = 0; a < NB; ++a) {
+                        Atb(a) += w * b[a] * y;
+                        for (int c = 0; c < NB; ++c) AtA(a, c) += w * b[a] * b[c];
+                    }
+                }
+            const double lambda = 1e-2 * AtA.diagonal().mean();  // ridge: x^2+y^2+z^2=1 makes the basis rank-deficient
+            for (int a = 0; a < NB; ++a) AtA(a, a) += lambda;
+            const Eigen::VectorXd c = AtA.ldlt().solve(Atb);
+
+            double sumsq = 0.0; size_t n_app = 0;
+            for (auto &oc : fulls)
+                for (auto &r : oc.reflections) {
+                    if (!std::isfinite(r.image_scale_corr) || r.image_scale_corr <= 0.0f) continue;
+                    basis(s1_crystal(r), b);
+                    double surf = 0.0;
+                    for (int a = 0; a < NB; ++a) surf += c(a) * b[a];
+                    r.image_scale_corr = static_cast<float>(r.image_scale_corr / std::exp(surf));
+                    sumsq += surf * surf; ++n_app;
+                }
+            logger.Info("Absorption surface iter {}/{}: {} fulls, rms(log A)={:.3f}",
+                        iter + 1, n_iter, n_app, std::sqrt(sumsq / std::max<size_t>(n_app, 1)));
+        }
+    }
+
     // Smooth the per-frame scale G across frames before the rot3d combine. ScaleOnTheFly fits each
     // frame's G independently, so the few partials of one rocking event get inconsistent scales when
     // weight-summed; a centered moving average of log(G) over a small odd frame window removes that
@@ -489,6 +572,10 @@ ProcessResult JFJochProcess::Run(JFJochProcessObserver *observer) {
                 phase("Scaling fulls (XDS order)");
                 ScaleFulls(experiment_, combined, static_cast<int>(config_.scaling_iter), config_.nthreads, logger);
             }
+            if (rot3d && experiment_.GetScalingSettings().GetAbsorptionIter() > 0) {
+                phase("Absorption surface");
+                AbsorptionSurface(experiment_, combined, experiment_.GetScalingSettings().GetAbsorptionIter(), logger);
+            }
             const std::vector<IntegrationOutcome> &merge_input =
                 rot3d ? combined : indexer->GetIntegrationOutcome();
 

tools/jfjoch_process.cpp

@@ -52,6 +52,7 @@ void print_usage() {
     std::cout << "   --scale-fulls                   After -P rot3d combine, refit a per-frame scale on the fulls (XDS order, Unity model); implies -M. Default ON for rot3d" << std::endl;
     std::cout << "   --no-scale-fulls                Disable the rot3d scale-fulls refit (it is on by default for rot3d)" << std::endl;
     std::cout << "   --write-process-h5              Also write the (large) _process.h5 when merging (default: only .mtz/.cif when merging)" << std::endl;
+    std::cout << "   --absorption[=num]              rot3d: fit a smooth absorption surface on the fulls over num iterations (default 3; off by default; matters at low energy)" << std::endl;
     std::cout << "   --smooth-g[=num]                rot3d: smooth per-frame scale G over a num-frame window before the combine (default: 9 for rot3d; 0 = off)" << std::endl;
     std::cout << "   -P, --partiality <txt>          Partiality model fixed|rot|rot3d|unity (default: fixed). rot3d = rot + 3D combine of per-frame partials" << std::endl;
     std::cout << "   -A, --anomalous                 Anomalous mode (don't merge Friedel pairs)" << std::endl;
@@ -103,7 +104,8 @@ enum {
     OPT_SMOOTH_G,
     OPT_RECIPROCAL_PROFILE,
     OPT_NO_SCALE_FULLS,
-    OPT_WRITE_PROCESS_H5
+    OPT_WRITE_PROCESS_H5,
+    OPT_ABSORPTION
 };
 
 static option long_options[] = {
@@ -127,6 +129,7 @@ static option long_options[] = {
     {"scale-fulls", no_argument, nullptr, OPT_SCALE_FULLS},
     {"no-scale-fulls", no_argument, nullptr, OPT_NO_SCALE_FULLS},
     {"write-process-h5", no_argument, nullptr, OPT_WRITE_PROCESS_H5},
+    {"absorption", optional_argument, nullptr, OPT_ABSORPTION},
     {"smooth-g", optional_argument, nullptr, OPT_SMOOTH_G},
     {"refine", required_argument, nullptr, 'r'},
 
@@ -291,6 +294,7 @@ int main(int argc, char **argv) {
     bool run_scaling = false;
     std::optional<bool> scale_fulls_arg;   // --scale-fulls / --no-scale-fulls; default on for rot3d
     bool write_process_h5_flag = false;    // --write-process-h5; also write _process.h5 when merging
+    int absorption_iter = 0;               // --absorption[=n]; rot3d absorption-surface iterations (0 = off)
     std::optional<int> smooth_g_window_arg;   // --smooth-g[=window]; default 9 for rot3d, 0 (off) otherwise
     bool anomalous_mode = false;
     std::optional<int64_t> space_group_number;
@@ -531,6 +535,10 @@ int main(int argc, char **argv) {
         case OPT_WRITE_PROCESS_H5:
             write_process_h5_flag = true;
             break;
+        case OPT_ABSORPTION:
+            absorption_iter = optarg ? std::stoi(optarg) : 3;
+            run_scaling = true;
+            break;
         case OPT_SMOOTH_G:
             smooth_g_window_arg = optarg ? std::stoi(optarg) : 9;
             break;
@@ -748,6 +756,7 @@ int main(int argc, char **argv) {
     scaling_settings.SetPartialityModel(partiality_model);
     scaling_settings.Combine3D(combine_3d);
     scaling_settings.ScaleFulls(scale_fulls);
+    scaling_settings.AbsorptionIter(absorption_iter);
     scaling_settings.SmoothGWindow(smooth_g_window_arg.value_or(combine_3d ? 9 : 0));
     if (d_min_scale_merge)
         scaling_settings.HighResolutionLimit_A(d_min_scale_merge.value());