Jungfraujoch/image_analysis/geom_refinement/AssignSpotsToRings.cpp

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

#include "../../common/JFJochMath.h"
#include "AssignSpotsToRings.h"

#include <cstdint>
#include <limits>
#include <vector>
#include <cmath>
#include <tuple>
#include <algorithm>

#include "../../common/CrystalLattice.h"

FindCircleCenterResult FindCircleCenter(const std::vector<SpotToSave> &v, int64_t max_spots) {
    if (max_spots <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid spot limit");
    if (v.size() < 3)
        return {.total_votes = 0, .votes_for_beam_center = 0, .x = 0.0f, .y = 0.0f};

    // Limit to only first 250 spots, given the algorithm is N^3
    auto task_size = std::min<size_t>(v.size(), max_spots);

    // The vote grid, one bin per pixel over the spots' bounding box (see the header).
    int64_t x0 = std::numeric_limits<int64_t>::max(), x1 = std::numeric_limits<int64_t>::min();
    int64_t y0 = x0, y1 = x1;
    for (size_t i = 0; i < task_size; i++) {
        x0 = std::min<int64_t>(x0, std::floor(v[i].x));
        x1 = std::max<int64_t>(x1, std::ceil(v[i].x));
        y0 = std::min<int64_t>(y0, std::floor(v[i].y));
        y1 = std::max<int64_t>(y1, std::ceil(v[i].y));
    }
    const int64_t width = x1 - x0 + 1;
    const int64_t height = y1 - y0 + 1;
    if ((width <= 0) || (height <= 0))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid image size");

    // uint32_t, not int64_t: the most votes any bin can take is C(max_spots,3), which for the 500-spot
    // limit is 2.1e7 - three orders below what a uint32 holds - and it halves the allocation.
    std::vector<uint32_t> vote(width * height, 0);

    for (int i = 0; i < task_size; i++) {
        for (int j = i+1; j < task_size; j++) {
            for (int k = j+1; k < task_size; k++) {
                // Calculation via determinants
                float a = v[i].x * (v[j].y - v[k].y) - v[i].y * (v[j].x - v[k].x) + v[j].x * v[k].y - v[k].x * v[j].y;
                if (std::abs(a) < 1e-10)
                    continue; // Points are collinear

                float x1_sq = v[i].x*v[i].x + v[i].y*v[i].y;
                float x2_sq = v[j].x*v[j].x + v[j].y*v[j].y;
                float x3_sq = v[k].x*v[k].x + v[k].y*v[k].y;

                float bx = x1_sq * (v[j].y - v[k].y) + x2_sq * (v[k].y - v[i].y) + x3_sq * (v[i].y - v[j].y);
                float by = x1_sq * (v[k].x - v[j].x) + x2_sq * (v[i].x - v[k].x) + x3_sq * (v[j].x - v[i].x);

                const int64_t cx = std::lround(bx / (2.0f * a)) - x0;
                const int64_t cy = std::lround(by / (2.0f * a)) - y0;

                if ((cx >= 0) && (cx < width) && (cy >= 0) && (cy < height))
                    vote[cx + cy * width]++;
            }
        }
    }

    int64_t total_votes = 0;
    int64_t max_votes = 0;
    int64_t cx = 0;
    int64_t cy = 0;
    for (int64_t y = 0; y < height; y++) {
        for (int64_t x = 0; x < width; x++) {
            total_votes += vote[x + y * width];
            if (vote[x + y * width] > max_votes) {
                max_votes = vote[x + y * width];
                cx = x;
                cy = y;
            }
        }
    }

    return {
            .total_votes = total_votes,
            .votes_for_beam_center = max_votes,
            .x = static_cast<float>(cx + x0),
            .y = static_cast<float>(cy + y0)
    };
}

// Very simple 1D DBSCAN on radii (works for rings)
std::vector<std::vector<int>> ClusterSpotsIntoRings(const std::vector<float>& r, float eps, int minPts) {
    size_t n = r.size();

    std::vector<int> labels(n, -1); // -1 = unvisited, -2 = noise
    int cluster_id = 0;

    for (int i=0; i<n; i++) {
        if (labels[i] != -1) continue; // already visited

        // find neighbors within eps in radius
        std::vector<int> neighbors;
        for (int j=0; j<n; j++) {
            if (std::fabs(r[i] - r[j]) <= eps) neighbors.push_back(j);
        }

        if ((int)neighbors.size() < minPts) {
            labels[i] = -2; // noise
            continue;
        }

        // start new cluster
        labels[i] = cluster_id;
        std::vector<int> seeds = neighbors;
        for (size_t k=0; k<seeds.size(); k++) {
            int j = seeds[k];
            if (labels[j] == -2) labels[j] = cluster_id;
            if (labels[j] != -1) continue;
            labels[j] = cluster_id;

            // expand cluster
            std::vector<int> nbrs2;
            for (int m=0; m<n; m++) {
                if (std::fabs(r[j] - r[m]) <= eps) nbrs2.push_back(m);
            }
            if ((int)nbrs2.size() >= minPts) {
                seeds.insert(seeds.end(), nbrs2.begin(), nbrs2.end());
            }
        }
        cluster_id++;
    }

    // Collect results
    std::vector<std::vector<int>> clusters(cluster_id);
    for (int i=0; i<n; i++) {
        if (labels[i] >= 0)
            clusters[labels[i]].push_back(i);
    }
    return clusters;
}

float median(std::vector<float> v) {
    if (v.empty()) return std::numeric_limits<float>::quiet_NaN();
    size_t n = v.size();
    std::nth_element(v.begin(), v.begin()+n/2, v.end());
    float m = v[n/2];
    if (n % 2 == 0) {
        auto it = std::max_element(v.begin(), v.begin()+n/2);
        m = 0.5f*(m + *it);
    }
    return m;
}

std::vector<RingClusters> AnalyzeClusters(const std::vector<float>& r, const std::vector<std::vector<int>> &clusters) {
    std::vector<RingClusters> ret;

    for (const auto & cluster : clusters) {
        std::vector<float> cluster_r;
        for (const auto &idx : cluster)
            cluster_r.push_back(r[idx]);

        if (cluster_r.size() < 2) continue;
        float m = median(cluster_r);
        ret.push_back({cluster, m, -1});
    }

    // sort by observed radius
    if (!ret.empty())
        std::sort(ret.begin(), ret.end(), [](const RingClusters& a, const RingClusters& b)
        { return a.R_obs < b.R_obs; });
    return ret;
}

namespace {
    bool reflection_present(ReflectionCondition condition, int h, int k, int l) {
        const bool all_odd = (h % 2 != 0) && (k % 2 != 0) && (l % 2 != 0);
        const bool all_even = (h % 2 == 0) && (k % 2 == 0) && (l % 2 == 0);
        switch (condition) {
            case ReflectionCondition::FaceCentred:
                return all_odd || all_even;
            case ReflectionCondition::Diamond:
                return all_odd || (all_even && ((h + k + l) % 4 == 0));
            case ReflectionCondition::All:
            default:
                return true;
        }
    }
}

namespace {
// The ring enumeration itself. `present` decides which hkl the lattice actually diffracts into - a
// fixed centring condition for the built-in standards, the space group's own absences for a cell the
// user supplied.
template <typename Present>
std::vector<float> XtalRings(const UnitCell &cell, int hkl_max, Present present) {
    CrystalLattice latt(cell);

    Coord Astar = latt.Astar();
    Coord Bstar = latt.Bstar();
    Coord Cstar = latt.Cstar();

    std::vector<float> u;
    // Both signs of h and k: only for a diagonal metric does |h a* + k b*| equal |h a* - k b*|, so on a
    // triclinic cell (silver behenate) the positive octant alone misses more rings than it finds. l stays
    // non-negative because hkl and -h-k-l are the same ring.
    for (int h = -hkl_max; h <= hkl_max; h++) {
        for (int k = -hkl_max; k <= hkl_max; k++) {
            for (int l = 0; l <= hkl_max; l++) {
                if (h == 0 && k == 0 && l == 0) continue;
                if (!present(h, k, l)) continue;
                auto p = Astar * h + Bstar * k + Cstar * l;
                float Q = 2.0f * PI * p.Length();
                u.push_back(Q);
            }
        }
    }
    std::sort(u.begin(), u.end());

    // Deduplicate (since e.g. (100), (010), (001) all give sqrt(1))
    u.erase(std::unique(u.begin(), u.end(),
                        [](float a, float b){ return std::fabs(a-b) < 1e-6; }),
            u.end());

    return u;
}
}   // namespace

std::vector<float> CalculateXtalRings(const UnitCell &cell, ReflectionCondition condition, int hkl_max) {
    return XtalRings(cell, hkl_max, [condition](int h, int k, int l) {
        return reflection_present(condition, h, k, l);
    });
}

std::vector<float> CalculateXtalRings(const UnitCell &cell, const gemmi::SpaceGroup &sg, int hkl_max) {
    const gemmi::GroupOps ops = sg.operations();
    return XtalRings(cell, hkl_max, [&ops](int h, int k, int l) {
        return !ops.is_systematically_absent(gemmi::Op::Miller{h, k, l});
    });
}


std::vector<float> CalculateCubicXtalRings(float a, int hkl_max) {
    return CalculateXtalRings(UnitCell(a,a,a,90,90,90), ReflectionCondition::All, hkl_max);
}

float GuessDetectorDistance(const DiffractionGeometry& geom, float ring_radius_pxl, float d_A) {
    float sin_theta = geom.GetWavelength_A() / (2 * d_A);
    if (sin_theta < 0 || sin_theta > 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Geometry makes no sense");
    float theta = asinf(sin_theta);
    float radius_mm = ring_radius_pxl * geom.GetPixelSize_mm();
    float det_dist_mm = radius_mm / tanf(2.0f * theta);
    return det_dist_mm;
}

std::vector<RingClusters> GuessInitialGeometry(DiffractionGeometry &geom,  const std::vector<SpotToSave> &v, float largest_ring_d_A) {
    // Reset rotations. The model assumes these are very small in any case!
    geom.PoniRot1_rad(0.0).PoniRot2_rad(0.0).PoniRot3_rad(0.0);

    auto center = FindCircleCenter(v);
    if (center.votes_for_beam_center < 20)
        throw JFJochException(JFJochExceptionCategory::CalibrationError, "Beam center not found");

    geom.BeamX_pxl(center.x).BeamY_pxl(center.y);

    std::vector<float> radii(v.size());
    for (int i = 0; i < v.size(); i++)
        radii[i] = std::hypot(v[i].x - center.x, v[i].y - center.y);

    auto clusters = ClusterSpotsIntoRings(radii);
    if (clusters.empty())
        throw JFJochException(JFJochExceptionCategory::CalibrationError, "Couldn't find spot clusters");

    auto cluster_annot = AnalyzeClusters(radii, clusters);

    float det_distance = GuessDetectorDistance(geom, cluster_annot[0].R_obs, largest_ring_d_A);
    geom.DetectorDistance_mm(det_distance);
    return cluster_annot;
}

void GuessGeometry(DiffractionGeometry &geom,  const std::vector<SpotToSave> &v, const std::vector<float> &ring_q,
                   bool refine_tilt) {
    if (ring_q.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No calibrant rings given");
    auto cluster_annot = GuessInitialGeometry(geom, v, 2 * PI / ring_q[0]);

    std::vector<RingOptimizerInput> optimizer_input;

    int ring_idx = 0;
    int cluster_idx = 0;

    // Walk the observed clusters and the calibrant's rings together, in ascending q. Which of them
    // paired up used to be printed to stdout from here; it is diagnostics a caller wants back as data,
    // not something a library should write to a terminal.
    while (cluster_idx < cluster_annot.size() && ring_idx < ring_q.size()) {

        const float obs_q = 2 * PI / geom.PxlToRes(cluster_annot[cluster_idx].R_obs);
        if (std::fabs(ring_q[ring_idx] - obs_q)
                < RingMatchWindow(ring_q, ring_idx, RING_MATCH_Q_RECIPA)) {
            for (const auto &spot: cluster_annot[cluster_idx].spots)
                optimizer_input.push_back({v[spot].x, v[spot].y, ring_q[ring_idx]});
            ring_idx++;
            cluster_idx++;
        } else if (ring_q[ring_idx] < obs_q) {
            ring_idx++;
        } else {
            cluster_idx++;
        }
    }

    RingOptimizer optimizer(geom, refine_tilt);
    geom = optimizer.Run(optimizer_input);
}

float RingMatchWindow(const std::vector<float> &ring_q, size_t i, float max_window) {
    float window = max_window;
    if (i > 0)
        window = std::min(window, 0.5f * (ring_q[i] - ring_q[i - 1]));
    if (i + 1 < ring_q.size())
        window = std::min(window, 0.5f * (ring_q[i + 1] - ring_q[i]));
    return window;
}

std::vector<RingOptimizerInput> AssignSpotsToRings(const DiffractionGeometry &geom,
                                                   const std::vector<SpotToSave> &v,
                                                   const std::vector<float> &ring_q) {
    std::vector<RingOptimizerInput> optimizer_input;

    for (const auto& s: v) {
        const float q_obs = 2 * PI / geom.PxlToRes(s.x, s.y);
        // The NEAREST ring, then the window - not the first ring within a fixed window. Where two rings
        // are closer together than that window, taking the first match assigns both of them to the
        // lower-q one and biases the distance it fits.
        size_t nearest = 0;
        for (size_t i = 1; i < ring_q.size(); i++) {
            if (std::fabs(ring_q[i] - q_obs) < std::fabs(ring_q[nearest] - q_obs))
                nearest = i;
        }
        if (!ring_q.empty()
            && std::fabs(ring_q[nearest] - q_obs) < RingMatchWindow(ring_q, nearest, RING_MATCH_Q_RECIPA))
            optimizer_input.push_back({s.x, s.y, ring_q[nearest]});
    }
    return optimizer_input;
}

void OptimizeGeometry(DiffractionGeometry &geom, const std::vector<SpotToSave> &v, const std::vector<float> &ring_q,
                      bool refine_tilt, RingFitUncertainty *unc) {
    RingOptimizer optimizer(geom, refine_tilt);
    geom = optimizer.Run(AssignSpotsToRings(geom, v, ring_q), unc);
}