Jungfraujoch/tools/battery/inputs.py

"""Finding a dataset's input file, and reading an XDS CORRECT.LP as a reference.

Both run only when a manifest is (re)built - `battery.py discover` and `battery.py refs`. A battery
run itself never searches: it runs the input the manifest names, so two runs of one manifest always
process the same files.
"""
import glob
import os
import re

# XDS writes calibration images beside the frames; none of them is a sweep.
XDS_AUX = ("-CORRECTIONS", "BKGINIT", "BKGPIX", "BLANK", "GAIN", "DECAY", "ABS.", "MODPIX", "FRAME.")


def master_frames(path):
    """Number of frames a NXmx master describes.

    A sweep is chosen by frame count, never by file size: a master with a few embedded screening
    frames can outweigh the master of a real sweep whose frames sit in external data files."""
    import h5py
    try:
        with h5py.File(path, "r") as h:
            det = "entry/instrument/detector/detectorSpecific"
            if f"{det}/nimages" in h:
                n = int(h[f"{det}/nimages"][()])
                if f"{det}/ntrigger" in h:
                    n *= max(1, int(h[f"{det}/ntrigger"][()]))
                return n
            omega = "entry/sample/goniometer/omega"
            if omega in h and h[omega].shape:
                return int(h[omega].shape[0])
            n = 0
            for k in h.get("entry/data", {}):
                try:
                    n += int(h["entry/data"][k].shape[0])
                except (KeyError, OSError):   # an external link whose file is missing
                    pass
            return n
    except OSError:
        return 0


def detector_frame(path):
    """'marCCD' or 'SMV' if the file's header says so, else None. These formats have no reliable
    extension (.mccd, .img, .001, ...; .img is SMV at one site and marCCD at the next)."""
    try:
        with open(path, "rb") as fh:
            head = fh.read(64)
            if head[:2] in (b"II", b"MM"):
                fh.seek(1028)
                return "marCCD" if fh.read(3) == b"MMX" else None
            if head[:1] == b"{" and b"HEADER_BYTES" in head:
                return "SMV"
    except OSError:
        pass
    return None


def find_input(d):
    """(input file, format, note) for the largest sweep under directory d, or (None, None, why).

    Frames are grouped per directory by template (the last run of digits in the name), and the
    biggest group by frame count wins; a tie goes to the alphabetically first. Groups are never
    merged across directories - two crystals can have identically named frames side by side."""
    masters = [m for m in glob.glob(f"{d}/**/*master*.h5", recursive=True)
               if not os.path.basename(m).startswith("._")]
    if not masters:
        # Diamond writes a NeXus .nxs master beside _NNNNNN.h5 data files
        masters = [m for m in glob.glob(f"{d}/**/*.nxs", recursive=True)
                   if not os.path.basename(m).startswith("._") and "meta" not in os.path.basename(m)]
    if masters:
        best = min(masters, key=lambda m: (-master_frames(m), m))
        return best, "nxs" if best.endswith(".nxs") else "h5", \
            f"master, {master_frames(best)} frames, {len(masters)} master(s)"

    groups = {}
    for p in glob.glob(f"{d}/**/*", recursive=True):
        name = os.path.basename(p)
        if not os.path.isfile(p) or name.startswith("._") or any(a in name.upper() for a in XDS_AUX):
            continue
        key = (os.path.dirname(p), re.sub(r"\d+(?=\D*$)", "#", name))
        groups.setdefault(key, []).append(p)
    for key, files in sorted(groups.items(), key=lambda kv: (-len(kv[1]), kv[0])):
        if len(files) < 10:
            break
        files.sort()
        if files[0].endswith((".cbf", ".cbf.gz")):
            fmt = "cbf"
        else:
            fmt = detector_frame(files[0])
        if fmt:
            return files[0], fmt, f'{fmt} "{key[1]}" x{len(files)} of {len(groups)} group(s)'
    return None, None, "no readable sweep found"


def input_format(path):
    """Format tag of an input file named in a manifest."""
    if path.endswith(".h5"):
        return "h5"
    if path.endswith(".nxs"):
        return "nxs"
    if path.endswith((".cbf", ".cbf.gz")):
        return "cbf"
    return detector_frame(path) or "other"


def parse_correct_lp(path):
    """Reference values from an XDS CORRECT.LP: the space group and cell CORRECT refined, the
    resolution range it merged (dmin_xds, dmax), the reference d_min read off its statistics
    (dmin, dmin_rule - see reference_dmin), its overall merging statistics, and the R_meas of its
    lowest-resolution shell (r_meas_low, over dmax to dmin_low) and the CC1/2 of its highest
    (cc_half_high).

    CORRECT.LP states the group and cell twice - first as the INPUT to INTEGRATE (routinely P1),
    then as CORRECT's own answer. The last SPACE_GROUP_NUMBER and the refined UNIT CELL PARAMETERS
    are the ones the statistics were computed in."""
    txt = open(path, errors="replace").read()
    ref = {}
    sgs = re.findall(r"SPACE_GROUP_NUMBER=\s*(\d+)", txt)
    ref["sgno"] = int(sgs[-1]) if sgs else None
    cells = re.findall(r"UNIT CELL PARAMETERS\s+([\d.\s]+)", txt)
    if cells:
        ref["cell"] = [float(x) for x in cells[-1].split()[:6]]
    else:
        m = re.findall(r"UNIT_CELL_CONSTANTS=\s*([\d.\s]+)", txt)
        ref["cell"] = [float(x) for x in m[-1].split()[:6]] if m else None
    m = re.search(r"FRIEDEL'S_LAW=\s*(TRUE|FALSE)", txt)
    ref["anomalous"] = bool(m and m.group(1) == "FALSE")
    m = re.search(r"^\s*a\s+b\s+ISa\s*\n\s*[\d.Ee+-]+\s+[\d.Ee+-]+\s+([\d.]+)", txt, re.M)
    ref["isa"] = float(m.group(1)) if m else None

    # the last "... AS FUNCTION OF RESOLUTION" table; shells from low to high resolution, then total
    seg = txt[txt.rfind("AS FUNCTION OF RESOLUTION"):]
    num = lambda x: float(x.rstrip("*%"))
    shells, total = [], None
    for line in seg.splitlines():
        f = line.split()
        if f and f[0] == "total" and len(f) >= 11:
            total = f
            break
        if f and re.match(r"^\d+\.\d+$", f[0]) and len(f) >= 11:
            shells.append(f)
    if total:
        ref["completeness"] = num(total[4])
        ref["r_meas"] = round(num(total[9]) / 100, 4)
        ref["cc_half"] = round(num(total[10]) / 100, 4)
        ref["multiplicity"] = round(int(total[1]) / int(total[2]), 2) if int(total[2]) else None
    if shells:
        # the lowest-resolution shell, dmax to dmin_low: strong reflections only, so its R_meas
        # is the merge's accuracy rather than where the signal fades
        ref["dmin_low"] = float(shells[0][0])
        ref["r_meas_low"] = round(num(shells[0][9]) / 100, 4)
        ref["cc_half_high"] = round(num(shells[-1][10]) / 100, 4)    # the highest-resolution shell

    # The range XDS merged: the explicit INCLUDE_RESOLUTION_RANGE, or - when XDS.INP left the high
    # limit at 0.0, "the whole detector" - the finest shell actually tabulated. The pooled statistics
    # above are over this range.
    m = re.search(r"INCLUDE_RESOLUTION_RANGE=\s*([\d.]+)\s+([\d.]+)", txt)
    include_low, include_high = (float(m.group(1)), float(m.group(2))) if m else (0.0, 0.0)
    table_high = float(shells[-1][0]) if shells else 0.0
    dmin_xds = (include_high if include_high > 0 else table_high) or None
    ref["dmin"], ref["dmin_rule"] = reference_dmin(
        [(float(f[0]), num(f[10]) / 100, f[10].endswith("*")) for f in shells], include_low, dmin_xds)
    ref["dmin_xds"] = dmin_xds
    ref["dmax"] = include_low or None
    return ref


# The CC1/2 target rugnux's own resolution cutoff uses (--resolution-cc-target, default 0.30)
CC_TARGET = 0.30


def reference_dmin(shells, dmax, dmin_xds):
    """(reference d_min, rule) from XDS's shell table [(shell high-resolution limit, CC1/2,
    CC1/2 significant)], low to high resolution.

    XDS's merged range is a resolution only when someone chose it. Left at "the whole detector"
    (INCLUDE_RESOLUTION_RANGE 0.0, or a limit at the detector edge) it is the detector's reach, not
    the crystal's: scored against, it is a meaningless number, and forced on rugnux it makes rugnux
    merge shells of pure noise. XDS itself says which case it is - it marks each shell's CC1/2 with
    '*' when it is significant at the 0.1% level:

      xds_range      the finest shell's CC1/2 is significant: XDS's limit stands.
      cc_half_0.30   it is not - XDS merged past its own signal. The reference is then where CC1/2
                     falls through the target rugnux's own cutoff uses, interpolated in 1/d^2
                     between the centres of the last shell above it and the first below it."""
    if not shells or shells[-1][2]:
        return dmin_xds, "xds_range"
    below = next((i for i, (_, cc, _) in enumerate(shells) if cc < CC_TARGET), None)
    if not below:                        # None, or already the first shell: no fall-off to read
        return dmin_xds, "xds_range"
    s = [1 / d ** 2 for d, _, _ in shells]
    edges = [1 / dmax ** 2 if dmax else 0.0] + s
    centre = [(edges[i] + edges[i + 1]) / 2 for i in range(len(s))]
    cc_hi, cc_lo = shells[below - 1][1], shells[below][1]
    frac = (cc_hi - CC_TARGET) / (cc_hi - cc_lo)
    s_cross = centre[below - 1] + frac * (centre[below] - centre[below - 1])
    return round(1 / s_cross ** 0.5, 3), f"cc_half_{CC_TARGET:.2f}"