Jungfraujoch/common/DetectorModules.cpp

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

#include "DetectorModules.h"
#include "JFJochException.h"
#include "DetectorGeometry.h"
#include "Definitions.h"
#include "PixelMask.h"

#include <algorithm>
#include <cmath>
#include <tuple>

DetectorModule DetectorModule::AsImaged(int64_t x0, int64_t y0, int64_t width, int64_t height) {
    DetectorModule m;
    m.x0 = x0;
    m.y0 = y0;
    m.width = width;
    m.height = height;
    m.ox = static_cast<double>(x0);
    m.oy = static_cast<double>(y0);
    return m;
}

int64_t DetectorModule::ReadoutX0() const {
    return (fast == Direction::Xneg || slow == Direction::Xneg) ? x0 + width - 1 : x0;
}

int64_t DetectorModule::ReadoutY0() const {
    return (fast == Direction::Yneg || slow == Direction::Yneg) ? y0 + height - 1 : y0;
}

std::pair<int64_t, int64_t> DirectionStep(DetectorModuleGeometry::Direction d) {
    switch (d) {
        case DetectorModuleGeometry::Direction::Xneg: return {-1, 0};
        case DetectorModuleGeometry::Direction::Ypos: return {0, 1};
        case DetectorModuleGeometry::Direction::Yneg: return {0, -1};
        default: return {1, 0};
    }
}

std::string DirectionName(DetectorModuleGeometry::Direction d) {
    switch (d) {
        case DetectorModuleGeometry::Direction::Xneg: return "Xn";
        case DetectorModuleGeometry::Direction::Ypos: return "Yp";
        case DetectorModuleGeometry::Direction::Yneg: return "Yn";
        default: return "Xp";
    }
}

DetectorModuleGeometry::Direction DirectionOfStep(double dx, double dy) {
    if (std::fabs(dx) > std::fabs(dy))
        return dx > 0 ? DetectorModuleGeometry::Direction::Xpos : DetectorModuleGeometry::Direction::Xneg;
    return dy > 0 ? DetectorModuleGeometry::Direction::Ypos : DetectorModuleGeometry::Direction::Yneg;
}

DetectorModuleGeometry::Direction DirectionFromName(const std::string &name) {
    if (name == "Xp") return DetectorModuleGeometry::Direction::Xpos;
    if (name == "Xn") return DetectorModuleGeometry::Direction::Xneg;
    if (name == "Yp") return DetectorModuleGeometry::Direction::Ypos;
    if (name == "Yn") return DetectorModuleGeometry::Direction::Yneg;
    throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Not a module direction: " + name);
}

ModulePlacement PlacementOf(const DetectorModule &m) {
    ModulePlacement p;
    p.x0 = m.ReadoutX0();
    p.y0 = m.ReadoutY0();
    p.fast = m.fast;
    p.slow = m.slow;
    std::tie(p.x0_plane, p.y0_plane) = m.ImageToPlane(static_cast<double>(p.x0), static_cast<double>(p.y0));
    // One readout pixel along an axis is one image pixel along its direction.
    const auto [ufx, ufy] = DirectionStep(m.fast);
    const auto [usx, usy] = DirectionStep(m.slow);
    p.fast_step = {ufx * m.fx + ufy * m.sx, ufx * m.fy + ufy * m.sy};
    p.slow_step = {usx * m.fx + usy * m.sx, usx * m.fy + usy * m.sy};
    return p;
}

bool DetectorModule::ContainsImagePixel(double x, double y) const {
    // A pixel's centre is at its index, so the module covers [x0 - 0.5, x0 + width - 0.5).
    return x >= static_cast<double>(x0) - 0.5 && x < static_cast<double>(x0 + width) - 0.5
           && y >= static_cast<double>(y0) - 0.5 && y < static_cast<double>(y0 + height) - 0.5;
}

std::pair<double, double> DetectorModule::ImageToPlane(double x, double y) const {
    const double u = x - static_cast<double>(x0), v = y - static_cast<double>(y0);
    return {ox + u * fx + v * sx, oy + u * fy + v * sy};
}

std::pair<double, double> DetectorModule::PlaneToImage(double px, double py) const {
    // Inverse of the 2x2 [f s]
    const double det = fx * sy - sx * fy;
    const double dx = px - ox, dy = py - oy;
    const double u = (sy * dx - sx * dy) / det;
    const double v = (-fy * dx + fx * dy) / det;
    return {static_cast<double>(x0) + u, static_cast<double>(y0) + v};
}

std::pair<double, double> DetectorModule::PlaneCentre() const {
    return ImageToPlane(static_cast<double>(x0) + 0.5 * static_cast<double>(width - 1),
                        static_cast<double>(y0) + 0.5 * static_cast<double>(height - 1));
}

void DetectorModule::Apply(const ModuleShift &shift) {
    const double c = std::cos(shift.rot_rad), s = std::sin(shift.rot_rad);
    const double ux = ox - shift.x0, uy = oy - shift.y0;
    ox = shift.x0 + c * ux - s * uy + shift.dx;
    oy = shift.y0 + s * ux + c * uy + shift.dy;
    const double nfx = c * fx - s * fy, nfy = s * fx + c * fy;
    const double nsx = c * sx - s * sy, nsy = s * sx + c * sy;
    fx = nfx; fy = nfy;
    sx = nsx; sy = nsy;
}

DetectorModules::DetectorModules(std::vector<DetectorModule> in) : modules(std::move(in)) {}

bool DetectorModules::empty() const {
    return modules.empty();
}

size_t DetectorModules::size() const {
    return modules.size();
}

const DetectorModule &DetectorModules::operator[](size_t i) const {
    return modules.at(i);
}

const std::vector<DetectorModule> &DetectorModules::Get() const {
    return modules;
}

bool DetectorModules::Moved() const {
    // Not an exact comparison: a module read back from a file went through a length and a unit vector
    // in single precision, which puts a module as imaged a few 1e-4 px off its box when it is far from
    // pixel (0, 0) on both axes.
    constexpr double tolerance = 1e-3;
    for (const auto &m: modules)
        if (std::fabs(m.ox - static_cast<double>(m.x0)) > tolerance || std::fabs(m.oy - static_cast<double>(m.y0)) > tolerance
            || std::fabs(m.fx - 1.0) > tolerance / CONVERTED_MODULE_COLS || std::fabs(m.fy) > tolerance / CONVERTED_MODULE_COLS
            || std::fabs(m.sx) > tolerance / CONVERTED_MODULE_COLS || std::fabs(m.sy - 1.0) > tolerance / CONVERTED_MODULE_COLS)
            return true;
    return false;
}

int64_t DetectorModules::ModuleOfImagePixel(double x, double y) const {
    for (size_t i = 0; i < modules.size(); i++)
        if (modules[i].ContainsImagePixel(x, y))
            return static_cast<int64_t>(i);
    return -1;
}

int64_t DetectorModules::ModuleOfPlanePoint(double px, double py) const {
    for (size_t i = 0; i < modules.size(); i++) {
        const auto [x, y] = modules[i].PlaneToImage(px, py);
        if (modules[i].ContainsImagePixel(x, y))
            return static_cast<int64_t>(i);
    }
    return -1;
}

std::pair<double, double> DetectorModules::ImageToPlane(double x, double y) const {
    const int64_t m = ModuleOfImagePixel(x, y);
    if (m < 0)
        return {x, y};
    return modules[m].ImageToPlane(x, y);
}

std::pair<double, double> DetectorModules::PlaneToImage(double px, double py) const {
    const int64_t m = ModuleOfPlanePoint(px, py);
    if (m < 0)
        return {px, py};
    return modules[m].PlaneToImage(px, py);
}

std::pair<double, double> DetectorModules::ImageToPlane(int64_t panel, double x, double y) const {
    if (panel < 0)
        return {x, y};
    return modules.at(panel).ImageToPlane(x, y);
}

std::pair<double, double> DetectorModules::PlaneToImage(int64_t panel, double px, double py) const {
    if (panel < 0)
        return {px, py};
    return modules.at(panel).PlaneToImage(px, py);
}

std::optional<PanelPoint> DetectorModules::PlaneToPanel(double px, double py) const {
    if (modules.empty())
        return PanelPoint{-1, px, py};
    for (size_t i = 0; i < modules.size(); i++) {
        const auto [x, y] = modules[i].PlaneToImage(px, py);
        if (modules[i].ContainsImagePixel(x, y))
            return PanelPoint{static_cast<int64_t>(i), x, y};
    }
    return std::nullopt;
}

PlaneBox DetectorModules::PlaneBounds(int64_t width, int64_t height) const {
    if (modules.empty())
        return {0.0, 0.0, static_cast<double>(width), static_cast<double>(height)};
    PlaneBox box{INFINITY, INFINITY, -INFINITY, -INFINITY};
    for (const auto &m: modules) {
        const double xa = static_cast<double>(m.x0) - 0.5, xb = static_cast<double>(m.x0 + m.width) - 0.5;
        const double ya = static_cast<double>(m.y0) - 0.5, yb = static_cast<double>(m.y0 + m.height) - 0.5;
        for (const auto &[x, y]: {std::pair{xa, ya}, std::pair{xb, ya}, std::pair{xa, yb}, std::pair{xb, yb}}) {
            const auto [px, py] = m.ImageToPlane(x, y);
            box.x0 = std::min(box.x0, px);
            box.y0 = std::min(box.y0, py);
            box.x1 = std::max(box.x1, px);
            box.y1 = std::max(box.y1, py);
        }
    }
    return box;
}

void DetectorModules::Apply(const std::vector<ModuleShift> &shifts) {
    for (auto &m: modules) {
        const auto [cx, cy] = m.PlaneCentre();
        const double reach = 0.5 * static_cast<double>(std::min(m.width, m.height));
        const ModuleShift *best = nullptr;
        for (const auto &s: shifts)
            if (std::hypot(s.x0 - cx, s.y0 - cy) < reach
                && (!best || std::hypot(s.x0 - cx, s.y0 - cy) < std::hypot(best->x0 - cx, best->y0 - cy)))
                best = &s;
        if (best)
            m.Apply(*best);
    }
}

std::vector<ModuleShift> DetectorModules::Centres() const {
    std::vector<ModuleShift> out;
    for (const auto &m: modules) {
        ModuleShift s;
        std::tie(s.x0, s.y0) = m.PlaneCentre();
        out.push_back(s);
    }
    return out;
}

DetectorModules ModulesAsImaged(const DetectorGeometry &geom) {
    if (geom.GetModulesNum() < 2)
        return {};
    std::vector<DetectorModule> out;
    for (int64_t m = 0; m < geom.GetModulesNum(); m++) {
        const Coord f = geom.GetFastDirection(m), s = geom.GetSlowDirection(m);
        const double cols = CONVERTED_MODULE_COLS - 1, lines = CONVERTED_MODULE_LINES - 1;
        const double xa = static_cast<double>(geom.GetX0(m)), ya = static_cast<double>(geom.GetY0(m));
        const double xb = xa + cols * f.x + lines * s.x, yb = ya + cols * f.y + lines * s.y;
        const auto x0 = static_cast<int64_t>(std::lround(std::min(xa, xb)));
        const auto y0 = static_cast<int64_t>(std::lround(std::min(ya, yb)));
        DetectorModule mod = DetectorModule::AsImaged(x0, y0, std::lround(std::abs(xb - xa)) + 1,
                                                      std::lround(std::abs(yb - ya)) + 1);
        mod.fast = DirectionOfStep(f.x, f.y);
        mod.slow = DirectionOfStep(s.x, s.y);
        out.push_back(mod);
    }
    return DetectorModules(out);
}

DetectorModules ModulesPlaced(const DetectorModules &target, const std::vector<ModulePlacement> &placement) {
    const auto name = [](int64_t x0, int64_t y0, DetectorModuleGeometry::Direction fast,
                         DetectorModuleGeometry::Direction slow) {
        return "x0 = " + std::to_string(x0) + ", y0 = " + std::to_string(y0) + ", " + DirectionName(fast) + "/"
               + DirectionName(slow);
    };
    std::vector<DetectorModule> out = target.Get();
    std::vector<bool> named(out.size(), false);
    for (const auto &p: placement) {
        const auto it = std::find_if(out.begin(), out.end(), [&](const DetectorModule &m) {
            return m.ReadoutX0() == p.x0 && m.ReadoutY0() == p.y0 && m.fast == p.fast && m.slow == p.slow;
        });
        if (it == out.end() || named[it - out.begin()])
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Module placement: the module " + name(p.x0, p.y0, p.fast, p.slow)
                                  + (it == out.end() ? " is not a module of this detector" : " is placed twice"));
        named[it - out.begin()] = true;
        // Image +x and +y in readout steps: image (dx, dy) from pixel 0 is readout
        // (i, j) = (dx ufx + dy ufy, dx usx + dy usy).
        const auto [ufx, ufy] = DirectionStep(p.fast);
        const auto [usx, usy] = DirectionStep(p.slow);
        DetectorModule &m = *it;
        m.fx = ufx * p.fast_step[0] + usx * p.slow_step[0];
        m.fy = ufx * p.fast_step[1] + usx * p.slow_step[1];
        m.sx = ufy * p.fast_step[0] + usy * p.slow_step[0];
        m.sy = ufy * p.fast_step[1] + usy * p.slow_step[1];
        const auto dx = static_cast<double>(m.x0 - p.x0), dy = static_cast<double>(m.y0 - p.y0);
        m.ox = p.x0_plane + dx * m.fx + dy * m.sx;
        m.oy = p.y0_plane + dx * m.fy + dy * m.sy;
    }
    for (size_t m = 0; m < out.size(); m++)
        if (!named[m])
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Module placement: the detector's module "
                                  + name(out[m].ReadoutX0(), out[m].ReadoutY0(), out[m].fast, out[m].slow)
                                  + " is not placed");
    return DetectorModules(out);
}

DetectorModules ModulesFromGapMask(const std::vector<uint32_t> &mask, int64_t width, int64_t height) {
    if (width <= 0 || height <= 0 || mask.size() != static_cast<size_t>(width * height))
        return {};
    const auto clear = [&](int64_t x, int64_t y) {
        return x >= 0 && y >= 0 && x < width && y < height
               && (mask[y * width + x] & (1u << PixelMask::ModuleGapPixelBit)) == 0;
    };
    std::vector<uint8_t> taken(mask.size(), 0);
    std::vector<DetectorModule> out;
    for (int64_t y = 0; y < height; y++)
        for (int64_t x = 0; x < width; x++) {
            if (!clear(x, y) || taken[y * width + x])
                continue;
            // The first clear pixel of a region not yet taken is a module's corner: the module runs to the
            // gap along the row and down the column, and has to be whole and bounded by gap all round.
            int64_t w = 0, h = 0;
            while (clear(x + w, y)) w++;
            while (clear(x, y + h)) h++;
            if (!((w == CONVERTED_MODULE_COLS && h == CONVERTED_MODULE_LINES)
                  || (w == CONVERTED_MODULE_LINES && h == CONVERTED_MODULE_COLS)))
                return {};
            for (int64_t yy = y - 1; yy <= y + h; yy++)
                for (int64_t xx = x - 1; xx <= x + w; xx++) {
                    const bool inside = xx >= x && xx < x + w && yy >= y && yy < y + h;
                    if (clear(xx, yy) != inside)
                        return {};
                    if (inside)
                        taken[yy * width + xx] = 1;
                }
            out.push_back(DetectorModule::AsImaged(x, y, w, h));
        }
    if (out.size() < 2)
        return {};
    return DetectorModules(out);
}