Jungfraujoch/reader/MiniCBF.cpp

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

#include "MiniCBF.h"

#include <zlib.h>

#include <algorithm>
#include <cctype>
#include <cmath>
#include <cstring>
#include <fstream>
#include <limits>
#include <map>
#include <mutex>
#include <regex>
#include <stdexcept>

#include "../common/JFJochException.h"
#include "../compression/CBFByteOffset.h"

namespace minicbf {

namespace {

// One capture group, first match, or nothing. The headers are a few kB, so a regex per field is
// cheap and keeps each rule next to the thing it reads - but COMPILING one is not: measured at ~60 us
// here, and a header parse runs 26 of them. Parsing happens once per image read across a whole sweep,
// so the same two dozen patterns were being recompiled thousands of times. They are compiled once and
// kept; std::map never invalidates a reference, so the pointer outlives the lock.
std::optional<std::string> Match(const std::string &text, const char *pattern) {
    static std::mutex cache_mutex;
    static std::map<std::string, std::regex> cache;
    const std::regex *re;
    {
        const std::lock_guard<std::mutex> lock(cache_mutex);
        auto it = cache.find(pattern);
        if (it == cache.end())
            it = cache.emplace(pattern, std::regex(pattern)).first;
        re = &it->second;
    }
    std::smatch m;
    if (!std::regex_search(text, m, *re) || m.size() < 2)
        return {};
    return m[1].str();
}

// The captures are character classes, not number grammars: "[\d.eE+-]+" matches "." and "+-", and
// "(\d+)" matches a digit string too long for int64. std::stod and std::stoll answer both with a raw
// std:: exception, which would leave the format probe below - CanRead catches JFJochException only -
// and reach the caller as an unhandled throw from merely LOOKING at a file. A header field that does
// not parse is a malformed header, so say that.
double Num(const std::string &text, const char *pattern, double fallback = 0.0) {
    const auto s = Match(text, pattern);
    if (!s.has_value())
        return fallback;
    try {
        return std::stod(*s);
    } catch (const std::exception &) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Malformed number in CBF header: '" + *s + "'");
    }
}

int64_t Int(const std::string &text, const char *pattern, int64_t fallback = 0) {
    const auto s = Match(text, pattern);
    if (!s.has_value())
        return fallback;
    try {
        return std::stoll(*s);
    } catch (const std::exception &) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Malformed integer in CBF header: '" + *s + "'");
    }
}

// A goniometer angle, or nothing where the head has no such axis. Writers spell that -9999, and
// taking a sentinel for an angle would put the head somewhere it never was.
std::optional<double> Angle(const std::string &text, const char *pattern) {
    const auto s = Match(text, pattern);
    if (!s.has_value())
        return {};
    double v;
    try {
        v = std::stod(*s);
    } catch (const std::exception &) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Malformed angle in CBF header: '" + *s + "'");
    }
    if (v < -9998.0)
        return {};
    return v;
}

std::string Field(const std::map<std::string, std::string> &row, const std::string &name) {
    const auto it = row.find(name);
    return (it == row.end()) ? std::string() : it->second;
}

// The imgCIF axis table: which axis turns or translates in which laboratory direction, and which two
// axes the image's columns and rows run along. Absent from most headers, which say nothing about any
// of this and are left exactly as they were read before.
//
// The element vectors are stated in the frame of the axis they depend on. Between them and the
// detector's own rotation every header seen has translations only, so they describe the image in the
// unswung detector frame - which is where the image orientation belongs, with the arm applied on top.
// Following Hammersley, Bernstein & Westbrook (2006) Int. Tables Cryst. G, 444-458
void ParseAxisTable(const std::string &text, Header &h) {
    const auto axes = imgcif::ParseLoop(text, "_axis.id");
    if (axes.empty())
        return;

    std::map<std::string, std::array<double, 3>> vector_of;
    for (const auto &row: axes) {
        const std::string id = Field(row, "_axis.id");
        const std::string type = Field(row, "_axis.type");
        const std::string equipment = Field(row, "_axis.equipment");
        const std::string depends_on = Field(row, "_axis.depends_on");
        std::array<double, 3> v{};
        try {
            for (int i = 0; i < 3; i++)
                v[i] = std::stod(Field(row, "_axis.vector[" + std::to_string(i + 1) + "]"));
        } catch (const std::exception &) {
            continue;   // "." for a vector: the table states no direction for this axis
        }
        vector_of[id] = v;

        // The base spindle and the detector arm are the rotations that hang off nothing: everything
        // further in is carried by them. Naming neither, so a beamline is free to call them anything.
        if ((type == "rotation") && (depends_on == ".")) {
            if (equipment == "goniometer")
                h.spindle_axis = v;
            else if (equipment == "detector")
                h.detector_axis = v;
        } else if ((type == "rotation") && (equipment == "goniometer")) {
            h.inner_spindle_axis = v;   // carried by the base axis: a fixed-chi or kappa inclination
        }
    }

    // Which axis the fast index runs along, and which the slow, through the two tables that say so.
    std::map<std::string, std::string> axis_of_set;
    for (const auto &row: imgcif::ParseLoop(text, "_array_structure_list_axis.axis_set_id"))
        axis_of_set[Field(row, "_array_structure_list_axis.axis_set_id")]
                = Field(row, "_array_structure_list_axis.axis_id");

    for (const auto &row: imgcif::ParseLoop(text, "_array_structure_list.array_id")) {
        const std::string set = Field(row, "_array_structure_list.axis_set_id");
        const auto id = axis_of_set.contains(set) ? axis_of_set[set] : set;
        const auto it = vector_of.find(id);
        if (it == vector_of.end())
            continue;
        std::array<double, 3> v = it->second;
        if (Field(row, "_array_structure_list.direction") == "decreasing")
            for (double &c: v)
                c = -c;
        const std::string index = Field(row, "_array_structure_list.index");
        if (index == "1")
            h.fast_direction = v;
        else if (index == "2")
            h.slow_direction = v;
    }
}

}  // namespace

bool ScansPhi(const Header &h) {
    if (h.phi_increment_deg != 0.0 || h.omega_increment_deg != 0.0 || h.chi_increment_deg != 0.0)
        return h.phi_increment_deg != 0.0;
    std::string name = h.axis_name;
    std::transform(name.begin(), name.end(), name.begin(),
                   [](unsigned char c) { return std::tolower(c); });
    return name.starts_with("phi");
}

std::optional<size_t> FindBinarySection(const uint8_t *data, size_t size) {
    if (size < sizeof(BINARY_SEPARATOR))
        return {};
    const auto *end = data + size;
    const auto *hit = std::search(data, end, std::begin(BINARY_SEPARATOR), std::end(BINARY_SEPARATOR));
    if (hit == end)
        return {};
    return static_cast<size_t>(hit - data) + sizeof(BINARY_SEPARATOR);
}

Header ParseHeader(const char *data, size_t size) {
    const std::string t(data, size);
    Header h;

    h.detector       = Match(t, R"(#\s*Detector:\s*([^\r\n]+))").value_or("PILATUS");
    h.pixel_x_m      = Num(t, R"(#\s*Pixel_size\s+([\d.eE+-]+)\s*m)");
    h.pixel_y_m      = Num(t, R"(#\s*Pixel_size\s+[\d.eE+-]+\s*m\s*x\s*([\d.eE+-]+)\s*m)");
    // The unit is optional: one ALBA set writes "thickness 0.001000" with no " m" after it.
    h.thickness_m    = Num(t, R"(sensor,\s*thickness\s+([\d.eE+-]+))");
    h.distance_m     = Num(t, R"(#\s*Detector_distance\s+([\d.eE+-]+))");
    h.beam_x_px      = Num(t, R"(#\s*Beam_xy\s*\(\s*([\d.eE+-]+))");
    h.beam_y_px      = Num(t, R"(#\s*Beam_xy\s*\([^,]+,\s*([\d.eE+-]+))");
    h.wavelength_A   = Num(t, R"(#\s*Wavelength\s+([\d.eE+-]+))");
    h.start_angle_deg     = Num(t, R"(#\s*Start_angle\s+([\d.eE+-]+))");
    h.angle_increment_deg = Num(t, R"(#\s*Angle_increment\s+([\d.eE+-]+))");
    h.two_theta_deg  = Num(t, R"(#\s*Detector_2theta\s+([\d.eE+-]+))");
    // The whitespace after each name is what keeps "Chi_increment" out of "Chi".
    h.chi_deg        = Angle(t, R"(#\s*Chi\s+([\d.eE+-]+))");
    h.omega_deg      = Angle(t, R"(#\s*Omega\s+([\d.eE+-]+))");
    h.kappa_deg      = Angle(t, R"(#\s*Kappa\s+([\d.eE+-]+))");
    h.chi_increment_deg   = Angle(t, R"(#\s*Chi_increment\s+([\d.eE+-]+))").value_or(0.0);
    h.phi_increment_deg   = Angle(t, R"(#\s*Phi_increment\s+([\d.eE+-]+))").value_or(0.0);
    h.omega_increment_deg = Angle(t, R"(#\s*Omega_increment\s+([\d.eE+-]+))").value_or(0.0);
    h.exposure_s     = Num(t, R"(#\s*Exposure_time\s+([\d.eE+-]+))");
    h.period_s       = Num(t, R"(#\s*Exposure_period\s+([\d.eE+-]+))");
    h.count_cutoff   = Int(t, R"(#\s*Count_cutoff\s+(\d+))");
    if (Match(t, R"(#\s*Polarization\s+([\d.eE+-]+))").has_value())
        h.polarization_fraction = Num(t, R"(#\s*Polarization\s+([\d.eE+-]+))");
    h.axis_name      = Match(t, R"(#\s*Oscillation_axis\s+(\S+))").value_or("omega");
    // "+SLOW" / "+FAST" on that same line: which of the image's two directions the spindle runs
    // along. Some writers state that instead of an axis name, and it is the only thing a header with
    // no axis table says about the spindle's direction at all.
    h.spindle_along_slow = Match(t, R"(#\s*Oscillation_axis[^\r\n]*\+(SLOW|slow))").has_value();
    ParseAxisTable(t, h);

    // "# Silicon sensor, ..." / "# CdTe sensor, ...". The rest of the code compares the material
    // against "CdTe" (BraggIntegrationEngine), so an unnormalised "Silicon" would silently give a
    // CdTe sensor silicon's attenuation length.
    if (const auto m = Match(t, R"(#\s*(\w+)\s+sensor,)")) {
        std::string s = *m;
        std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return std::tolower(c); });
        h.material = (s == "cdte") ? "CdTe" : "Si";
    }

    h.nx    = Int(t, R"(X-Binary-Size-Fastest-Dimension:\s*(\d+))");
    h.ny    = Int(t, R"(X-Binary-Size-Second-Dimension:\s*(\d+))");
    h.nelem = Int(t, R"(X-Binary-Number-of-Elements:\s*(\d+))");

    const auto conv = Match(t, R"RE(conversions\s*=\s*"([^"]+)")RE").value_or("");
    h.byte_offset = conv.find("x-CBF_BYTE_OFFSET") != std::string::npos;

    if (!Match(t, R"(#\s*(Pixel_size))").has_value()) {   // Match returns the first group
        h.imgcif = imgcif::Parse(t);
        if (h.imgcif.has_value()) {
            h.detector = h.imgcif->detector;
            h.pixel_x_m = std::abs(h.imgcif->fast.increment) * 1e-3;
            h.pixel_y_m = std::abs(h.imgcif->slow.increment) * 1e-3;
            h.wavelength_A = h.imgcif->wavelength_A;
            h.count_cutoff = static_cast<int64_t>(h.imgcif->overload);
            h.exposure_s = h.imgcif->exposure_s;
            h.thickness_m = 0;
        }
    }

    if (h.period_s <= 0.0)
        h.period_s = h.exposure_s;

    return h;
}

namespace {

// Whether the file is gzip-wrapped, from its first two bytes rather than its name - a sweep is
// named consistently but the magic is what decides how to read it.
bool IsGzip(const std::string &path) {
    std::ifstream f(path, std::ios::binary);
    unsigned char m[2] = {0, 0};
    if (f)
        f.read(reinterpret_cast<char *>(m), sizeof(m));
    return m[0] == 0x1f && m[1] == 0x8b;
}

// The gzip path. EMBL Hamburg's beamlines write .cbf.gz by default, so a sweep from there is
// otherwise twice the disk and a decompression pass before anything can read a frame. The
// decompressed size is not in the stream, so the buffer grows in chunks instead of being sized
// up front the way the plain path can. It grows only past what it already holds: a buffer reused
// frame to frame is the previous frame's size, and every resize above that zero-fills the bytes
// it adds - tens of megabytes a frame of a 16 Mpx sweep, all of them overwritten by the read.
template <class Buffer>
void SlurpGz(const std::string &path, size_t max_bytes, Buffer &buf) {
    gzFile g = gzopen(path.c_str(), "rb");
    if (g == nullptr)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot open gzipped CBF file " + path);
    constexpr size_t CHUNK = 1u << 20;
    size_t at = 0;
    while (at < max_bytes) {
        const size_t want = std::min(CHUNK, max_bytes - at);
        if (buf.size() < at + want)
            buf.resize(at + want);
        const int got = gzread(g, buf.data() + at, static_cast<unsigned>(want));
        if (got < 0) {
            int err = 0;
            const std::string msg = gzerror(g, &err);
            gzclose(g);
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Error decompressing " + path + ": " + msg);
        }
        at += static_cast<size_t>(got);
        if (static_cast<size_t>(got) < want)
            break;                      // end of the stream
    }
    buf.resize(at);
    gzclose(g);
}

template <class Buffer>
void Slurp(const std::string &path, size_t max_bytes, Buffer &buf) {
    // The plain path is kept exactly as it was rather than reading everything through zlib: a
    // gzFile over an uncompressed file works, but it copies every byte through zlib's own buffer,
    // and the overwhelming majority of sweeps are not compressed.
    if (IsGzip(path)) {
        SlurpGz(path, max_bytes, buf);
        return;
    }
    std::ifstream f(path, std::ios::binary);
    if (!f)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot open CBF file " + path);
    f.seekg(0, std::ios::end);
    const auto file_size = static_cast<size_t>(f.tellg());
    f.seekg(0, std::ios::beg);
    buf.resize(std::min(file_size, max_bytes));
    f.read(reinterpret_cast<char *>(buf.data()), static_cast<std::streamsize>(buf.size()));
    buf.resize(static_cast<size_t>(f.gcount()));
}

// Enough to reach the separator on any header seen in the wild (the longest measured is ~6.3 kB).
// The first read is kept small because a sweep's startup reads every frame's header: on a spinning
// disk 256 kB from each of 2400 files took twice as long as 16 kB (12 s against 6 s). A header
// that does not end within it is read again with the large probe, so the result never depends on it.
constexpr size_t HEADER_FIRST_PROBE_BYTES = 16 * 1024;
constexpr size_t HEADER_PROBE_BYTES = 256 * 1024;

}  // namespace

Header ReadHeader(const std::string &path) {
    std::vector<uint8_t> buf;
    Slurp(path, HEADER_FIRST_PROBE_BYTES, buf);
    auto start = FindBinarySection(buf.data(), buf.size());
    if (!start.has_value()) {
        Slurp(path, HEADER_PROBE_BYTES, buf);
        start = FindBinarySection(buf.data(), buf.size());
    }
    if (!start.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No CBF binary section in " + path);
    return ParseHeader(reinterpret_cast<const char *>(buf.data()),
                       *start - sizeof(BINARY_SEPARATOR));
}

namespace {

// One read of the file, header parsed, dimensions checked. The caller supplies where the pixels go.
template <class Buffer>
Header ReadCommon(const std::string &path, Buffer &buf, size_t &binary_start) {
    Slurp(path, std::numeric_limits<size_t>::max(), buf);
    const auto start = FindBinarySection(buf.data(), buf.size());
    if (!start.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No CBF binary section in " + path);

    const Header h = ParseHeader(reinterpret_cast<const char *>(buf.data()),
                                 *start - sizeof(BINARY_SEPARATOR));
    if (!h.byte_offset)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Unsupported CBF compression in " + path + " (only x-CBF_BYTE_OFFSET)");
    if (h.nelem <= 0 || h.nx <= 0 || h.ny <= 0 || h.nelem != h.nx * h.ny)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Inconsistent image dimensions in " + path);
    // Without a pixel size there is no geometry at all - every resolution, every scattering vector
    // and the beam centre in millimetres all scale by it - and the default of 0 collapses all of
    // them silently. A byte-offset CBF carrying no "# Pixel_size" line is not a detector image from
    // this family at all; XDS writes correction files in exactly that shape. Refuse it here rather
    // than let it through with a geometry of zero.
    if (!(h.pixel_x_m > 0.0) || !(h.pixel_y_m > 0.0))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + " has no pixel size in its header (no '# Pixel_size' line); "
                              "it carries a CBF binary section but is not a detector image");

    binary_start = *start;
    return h;
}

}  // namespace

Header Read(const std::string &path, std::vector<int32_t> &out) {
    std::vector<uint8_t> buf;
    size_t start = 0;
    const Header h = ReadCommon(path, buf, start);
    out.resize(static_cast<size_t>(h.nelem));
    DecodeCBFByteOffset(buf.data() + start, buf.size() - start, out.data(),
                     static_cast<size_t>(h.nelem));
    return h;
}

Header ReadCompressed(const std::string &path, RawByteBuffer &buf, size_t &binary_start) {
    return ReadCommon(path, buf, binary_start);
}

Header ReadInto(const std::string &path, int32_t *out, size_t capacity, std::vector<uint8_t> &scratch) {
    size_t start = 0;
    const Header h = ReadCommon(path, scratch, start);
    if (static_cast<size_t>(h.nelem) > capacity)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "CBF image does not fit the supplied buffer: " + path);
    DecodeCBFByteOffset(scratch.data() + start, scratch.size() - start, out, static_cast<size_t>(h.nelem));
    return h;
}

}  // namespace minicbf