Jungfraujoch/image_analysis/geom_refinement/GeometryRefiner.h

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

#pragma once

#include <array>
#include <cstdint>
#include <vector>

#include "../../common/CrystalLattice.h"
#include "../../common/DiffractionGeometry.h"
#include "../../common/ModuleShift.h"
#include "../../common/UnitCell.h"
#include "gemmi/symmetry.hpp"

// Offline global (multi-frame) geometry refinement for serial stills. Pass 1 indexes many frames
// independently; this joint bundle adjustment then determines the ONE shared detector geometry (beam
// centre, distance, cell) that per-image refinement cannot reach from a single sparse still, so
// re-indexing with it converts more frames. Mirrors the validated Python prototype: one Ceres problem
// with shared beam/distance/cell blocks and a per-frame orientation block, robust loss, and a prior
// on the input cell (GeometryRefinerSettings::cell_length_sigma) that settles the distance<->cell-scale
// degeneracy where the spots leave it open. The detector tilt (rot1/rot2) is refined only on request (refine_tilt), and then with the beam
// written as the direct beam, which the tilt does not move.

struct GeomRefineSpot {
    float x = 0, y = 0;         // observed (raw) spot position, pixels
    int32_t h = 0, k = 0, l = 0; // integer Miller index assigned in pass 1
    int16_t module = -1;        // detector module the spot lies on (GeometryRefinerSettings::modules); -1: none
};


struct GeomRefineFrame {
    CrystalLattice lattice;             // per-frame indexed lattice (real space), for the orientation seed
    std::vector<GeomRefineSpot> spots;  // indexed spots on this frame
    std::vector<GeomRefineSpot> unindexed;  // the frame's other spots (no index); see ScanModuleShifts
    int32_t ordinal = 0;                // image this came from; ties are broken on it, so the
                                        // selected bundle does not depend on collection order
};

struct GeometryRefinerSettings {
    gemmi::CrystalSystem crystal_system = gemmi::CrystalSystem::Triclinic;
    // The input cell is a prior on the refined one, not a constraint: its lengths and the angles the crystal
    // system leaves free (beta of a monoclinic cell, all three of a triclinic one) are refined, each held to
    // the input value by a Gaussian of this width - relative for the lengths, in degrees for the angles -
    // weighed against the spots' scatter of about a pixel. A sigma of 1% / 1 degree covers what a cell from
    // another crystal, temperature or setting is usually off by. On a bundle of strong frames the spots
    // outweigh it by orders of magnitude wherever they determine the cell, so it decides only what the spots
    // leave undetermined: the cell scale near the beam, where it trades with the distance.
    double cell_length_sigma = 0.01;
    double cell_angle_sigma_deg = 1.0;
    // With the distance held, the cell is no longer degenerate with it and is refined from the spots.
    bool refine_distance = true;
    int rounds = 3;                 // hkl-reassignment / tolerance-tightening rounds
    double beam_range_px = 8.0;     // beam bound around the nominal value
    double distance_range_mm = 8.0; // distance bound around the nominal value
    int min_spots_total = 100;      // refuse to refine below this many indexed spots
    int num_threads = 1;
    // false: beam, distance and cell are held where nominal_geom and input_cell put them and only the
    // crystal orientations follow - the fit quality of a geometry decided elsewhere, on frames it was
    // not fitted to.
    bool refine_beam_distance = true;
    // The detector's modules, with where each is placed to start from; empty: the image is one rigid
    // detector. With refine_modules, every module that has at least min_module_spots spots on its frames'
    // lattices where it starts is moved as a rigid block (shift and turn in the detector plane); the others
    // go back to where they were given (no move), whatever their start.
    std::vector<ModuleShift> modules;
    bool refine_modules = false;
    int min_module_spots = 100;
    // Modules held where they were given whatever their spots (as modules; empty: none held).
    std::vector<bool> modules_held;
    // Refine the detector tilt (PONI rot1, rot2), within 2 degrees of the nominal one, together with the beam,
    // the distance and (with refine_modules) the modules. The beam is then fitted as the direct beam, and its
    // box (beam_range_px) is about where the nominal geometry puts the direct beam: in PONI terms a tilt moves
    // the whole pattern as a moved beam does, by D tan(tilt) / pixel, more than the box would hold.
    bool refine_tilt = false;
};

struct GeometryRefinerResult {
    bool ok = false;
    double beam_x_px = 0, beam_y_px = 0;   // refined beam centre, the PONI (DiffractionGeometry::BeamX_pxl)
    double distance_mm = 0;
    UnitCell cell{};                       // refined cell
    int frames_used = 0;
    int spots_used = 0;
    double median_residual_px = 0;         // final per-spot detector residual (fit quality)
    // The same over EVERY spot of the frames, each spot's residual to its nearest predicted reflection
    // capped at 2 px: unlike the median of the spots that fit, it does not improve by fitting fewer of them,
    // so two geometries can be compared on the same spots.
    double capped_rms_px = 0;
    std::vector<ModuleShift> modules;      // the modules' placement at the end (settings.modules if not refined)
    double rot1_rad = 0, rot2_rad = 0;     // detector tilt at the end (the nominal one where not refined)
    // Per module (as settings.modules): rms of every spot's residual capped at 2 px, and the mean
    // resolution of its spots - what a module's fit rests on.
    std::vector<double> module_capped_rms_px;
    std::vector<double> module_mean_d_A;
    // Per frame, per spot (as frames[i].spots): where the refined geometry predicts the spot's reflection on
    // the detector, as the spot would be found (before its module's move); NaN where it predicts none.
    std::vector<std::vector<std::array<float, 2>>> predicted_px;
    // Per frame (as frames): its lattice at the end, the frame's orientation with the refined cell.
    std::vector<CrystalLattice> lattices;
};

// nominal_geom supplies the starting beam/distance/tilt/wavelength/pixel size; input_cell is where the cell
// starts and the centre of its prior. Returns ok=false (and leaves the caller's geometry untouched) on any failure.
GeometryRefinerResult RefineGlobalGeometry(const DiffractionGeometry &nominal_geom,
                                           const UnitCell &input_cell,
                                           const std::vector<GeomRefineFrame> &frames,
                                           const GeometryRefinerSettings &settings);

// A coarse search for where each module is, ahead of RefineGlobalGeometry: a module misplaced by more than a
// fraction of the spot spacing has its spots left unindexed (or indexed wrongly), so the refinement has
// nothing to pull it with. Each module is shifted over a grid of +-range_px, and at each shift every spot on
// it (indexed or not) is scored against its frame's lattice as indexed: the shift that puts the most spots on
// the lattices wins. The frames' lattices come from the other modules, so this works while most of the
// detector is roughly right. Turns are not searched.
struct ModuleScan {
    double dx = 0, dy = 0;     // the best shift found
    int on_lattice_at_zero = 0;
    int on_lattice_at_best = 0;
};

std::vector<ModuleScan> ScanModuleShifts(const DiffractionGeometry &geom,
                                         const std::vector<GeomRefineFrame> &frames,
                                         const std::vector<ModuleShift> &modules,
                                         double range_px, size_t nthreads = 1);