Jungfraujoch/image_analysis/pixel_refinement/PixelRefine.h

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

#pragma once

#include "../bragg_prediction/BraggPrediction.h"
#include "../common/DiffractionExperiment.h"
#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/AzimuthalIntegrationProfile.h"
#include "../scale_merge/HKLKey.h"

// =============================================================================
// PixelRefine — one optimization to rule geometry, integration and scaling
// =============================================================================
//
// Intent
// ------
// Classical crystallographic data processing is a one-way pipeline:
//
//     spot finding -> indexing -> geometry refinement -> integration -> scaling -> merging
//
// Each stage consumes the previous stage's output and never talks back. The
// integrator trusts the refined geometry; the scaler trusts the integrated
// intensities; nothing downstream is ever allowed to correct an upstream
// parameter. Post-refinement and profile fitting were the field's partial
// answers to this: post-refinement lets merged intensities nudge per-image
// orientation/cell/mosaicity, and profile fitting lets a learned spot shape
// improve weak-reflection intensities. But both are narrow back-channels bolted
// onto a feed-forward pipeline — there is no end-to-end gradient that flows from
// the raw detector pixels all the way back to every parameter at once.
//
// PixelRefine is an experiment in doing the whole thing as a *single* least
// squares problem. We write down, for every pixel in a reflection's shoebox, the
// expected counts as an explicit forward model
//
//     I_pred(pixel) = G * I_true * B_term * P_radial * P_tangential + I_bkg
//
// and let Ceres autodiff back-propagate the per-pixel residuals into ALL of:
//   * detector geometry  (beam centre, distance, tilt)
//   * crystal orientation + unit cell
//   * overall scale G and Debye-Waller B
//   * the reciprocal-space spot widths R = (radial, tangential)
// simultaneously. Geometry refinement, profile-fitted integration and scaling
// then stop being separate stages: they are different parameters of one model,
// coupled through the same pixels, with full backpropagation between them. Once
// the model is differentiable end-to-end, things that used to need bespoke code
// — mosaicity refinement, profile fitting, partiality — fall out "for free" as
// extra parameters of the same forward model.
//
// How I_true enters
// -----------------
// I_true is NOT refined here. It is a *fixed hypothesis* for the duration of a
// pass: the current best merged estimate of each reflection's full intensity.
// The intended outer loop is iterative, like EM / self-consistent field:
//
//     repeat over the whole dataset:
//         run PixelRefine on every image with the current I_true reference
//         re-merge the resulting intensities -> new I_true
//     until the reference stops changing
//
// So a single PixelRefine call answers "given this intensity hypothesis, what
// geometry/scale/profile best explains these pixels, and what intensities do I
// read back out?", and the dataset-level loop refines the hypothesis itself.
//
// Inner predict<->refine loop
// ---------------------------
// Within one image we also iterate (max_iterations): Bragg prediction places the
// shoeboxes, we refine, the refined geometry/cell feed the next prediction, etc.
// Initially the model geometry equals the experiment's DiffractionGeometry, but
// as refinement proceeds it diverges, so later predictions must use the *refined*
// data.geom rather than the static experiment geometry.
//
// Status: experimental prototype. The forward model (esp. the still-image
// Lorentz/partiality normalization) is deliberately simple and expected to
// evolve. See PixelRefine.cpp for the physics conventions and known caveats.
// =============================================================================

struct PixelRefineData {
    // --- model state (input as initial guess, output as refined result) ---
    DiffractionGeometry geom;
    CrystalLattice latt;
    gemmi::CrystalSystem crystal_system = gemmi::CrystalSystem::Triclinic;
    char centering = 'P';

    double B_factor = 0.0;      // Debye-Waller B (A^2)
    double scale_factor = 1.0;  // overall scale G
    double R[2] = {0.005, 0.005}; // R[0] = radial (partiality) width, R[1] = tangential (profile) width (A^-1)

    // Goniometer: for a still image keep angle_deg = 0. For a wedge, pass the
    // angle (deg) of the slice centre relative to the reference (phi=0) frame.
    double angle_deg = 0.0;

    // --- what to refine ---
    bool refine_orientation     = true;  // crystal orientation (p0)
    bool refine_unit_cell       = false; // cell lengths + angles
    bool refine_beam_center     = false;
    bool refine_distance        = false;
    bool refine_detector_angles = false;
    bool refine_rotation_axis   = false;
    bool refine_scale           = true;
    bool refine_B               = false;
    bool refine_R               = true;

    double max_time_s = 5.0;
    int    shoebox_radius = 3; // half-size of the per-reflection pixel box
    int    max_iterations = 3; // inner predict<->refine cycles (re-predict with refined geom/latt)

    // --- output ---
    std::vector<Reflection> reflections; // profile-fitted integration result
    bool   solved = false;
    double final_cost = NAN;
    size_t residual_count = 0;
};

class PixelRefine {
    BraggPrediction &prediction;
    const AzimuthalIntegrationMapping &mapping;
    const size_t xpixel, ypixel;
    const DiffractionExperiment &experiment;

    const HKLKeyGenerator hkl_key_generator;
    std::map<HKLKey, double> reference_data;
public:
    PixelRefine(const DiffractionExperiment &experiment,
                const AzimuthalIntegrationMapping &mapping,
                const std::vector<MergedReflection> &reference,
                BraggPrediction &prediction);

    template<class T>
    void Run(const T *image,
             const AzimuthalIntegrationProfile &profile,
             PixelRefineData &data);
};