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 <optional>

#include "../bragg_prediction/BraggPrediction.h"
#include "../common/DiffractionExperiment.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 * pol + I_bkg
//
// in raw detector counts (pol = per-reflection polarization correction, I_bkg =
// local per-shoebox background read from the image).
//
// 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.
//
// On shoeboxes, tails, and gatekeeping
// ------------------------------------
// We deliberately do NOT chase the full spot. A shoebox only needs to cover
// enough of a reflection for the fit to be meaningful; clipped tails are fine,
// because the partiality term already downweights whatever falls outside the
// well-modelled core. The premise - especially for serial crystallography - is
// that the problem is rarely *missing* information; it is *failing to gate out*
// information that is not meaningful. As long as the model knows a piece is
// missing (low partiality), it is safe to leave it missing. That flips the usual
// trade-off: rather than shrinking boxes to avoid contamination, we can grow them
// and let partiality decide, per pixel and per reflection, what actually carries
// signal. (For downstream integration this pixel-level gating is the point - keep
// only meaningful pixels, instead of a fixed geometric mask.) The bandwidth term
// below is part of the same idea: it tells the model where the radial tails are
// *expected* to be, so it can weight rather than blindly include them.
//
// X-ray bandwidth (optional)
// --------------------------
// A finite bandwidth thickens the Ewald shell radially and smears spots along the
// radial direction, growing like 1/d^2 (the pink-beam/DMM signature). It enters
// as a fixed, resolution-dependent addition to the radial width R0 (see
// PixelRefineData::bandwidth and PixelRefine.cpp). It is OFF by default
// (bandwidth = 0, monochromatic); set it for DMM-type data, leave it for Si.
//
// Status: experimental prototype. The forward model (esp. the still-image
// partiality normalization) is deliberately simple and expected to evolve: it
// works in raw detector counts with a local per-shoebox background and a
// per-reflection polarization correction (no per-pixel solid-angle/Lorentz
// weighting), matching the classical integrator. 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)

    // Relative X-ray bandwidth (sigma of dlambda/lambda), e.g. ~0.004 for a 1%
    // FWHM DMM, ~1e-4 for Si(111). Adds a resolution-dependent radial broadening
    // to R[0]. 0 = monochromatic (the term switches off entirely).
    double bandwidth = 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_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 signal box (peak region that enters the fit)
    // Half-size of the local-background sampling box. Background is estimated from
    // the ring shoebox_radius < |dx|,|dy| <= bkg_outer_radius around each spot
    // (excluding pixels belonging to any predicted spot core), mirroring the local
    // shoebox background of BraggIntegrate2D. Must be > shoebox_radius.
    int    bkg_outer_radius = 6;
    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;
    double cc = NAN;        // per-image CC of scaled intensities vs reference
    int64_t cc_n = 0;       // number of reflections in the CC
};

class PixelRefine {
    const size_t xpixel, ypixel;
    const DiffractionExperiment &experiment;

    const HKLKeyGenerator hkl_key_generator;
    std::map<HKLKey, double> reference_data;

    // Fills the Ceres parameter blocks (geometry + symmetry-aware lattice
    // parametrization) from the current model state. Shared by Run and
    // PredictImage so both walk identical geometry/lattice code.
    void BuildParameterBlocks(const PixelRefineData &data,
                              double beam[2], double &dist_mm,
                              double detector_rot[2],
                              double latt_vec0[3], double latt_vec1[3], double latt_vec2[3]) const;

    // Bragg-prediction settings shared by Run/PredictImage/ChiSquaredImage so all
    // three predict an identical reflection set. Critically it sets the radial
    // Ewald-acceptance band (ewald_dist_cutoff) and the bandwidth broadening - see
    // the definition for why these must track the partiality model, not the
    // BraggPredictionSettings struct defaults.
    BraggPredictionSettings BuildPredictionSettings(const PixelRefineData &data) const;
public:
    PixelRefine(const DiffractionExperiment &experiment,
                const std::vector<MergedReflection> &reference);

    // The BraggPrediction is supplied per call (it is mutated): this keeps a
    // single PixelRefine instance usable from several threads, each passing its
    // own prediction buffer. Only `data` is written; PixelRefine state is const.
    // The image is in raw detector counts (masked/saturated pixels carry the type
    // sentinel); background is estimated locally per shoebox from the image itself.
    template<class T>
    void Run(const T *image,
             BraggPrediction &prediction,
             PixelRefineData &data);

    // Render the forward model as a full detector image (raw detector counts, so
    // it overlays directly on the original image). Uses the *same* per-pixel
    // model path (PixelResidual::Model) as the optimizer, evaluated in double
    // precision - slow but exact. For each reference reflection it adds the Bragg
    // signal over its shoebox; with include_background it also lays down the local
    // per-shoebox background read from the supplied image - the same background the
    // fit uses. Diagnostic tool, not on the hot path.
    template<class T>
    std::vector<float> PredictImage(const T *image,
                                    BraggPrediction &prediction,
                                    const PixelRefineData &data,
                                    bool include_background = true) const;

    // Render the per-pixel chi-square (cost density) that the optimizer actually
    // minimizes: for every shoebox pixel that enters the fit it stores the squared
    // weighted residual ((I_pred - I_obs)/sigma)^2 in raw counts - identical to the
    // Ceres residual_i^2 - accumulating where shoeboxes overlap. Pixels that are not
    // part of any shoebox stay 0; masked/saturated pixels (skipped by the fit) also
    // stay 0. Summing the image gives ~2*final_cost. Diagnostic tool.
    template<class T>
    std::vector<float> ChiSquaredImage(const T *image,
                                       BraggPrediction &prediction,
                                       const PixelRefineData &data) const;

    // Reference (merged) intensity used as the fixed hypothesis for a reflection,
    // or nullopt if this hkl is not in the reference. Lets callers show the fitted
    // estimate next to the reference it was scaled against.
    std::optional<double> ReferenceIntensity(const Reflection &r) const {
        const auto it = reference_data.find(hkl_key_generator(r));
        if (it == reference_data.end())
            return std::nullopt;
        return it->second;
    }
};