Jungfraujoch/tools/battery/score.py

"""Scoring one rugnux run against its reference.

One dataset, one verdict, one cause, decided in this order: did it run -> is the lattice right ->
is the symmetry right. Data quality is not scored: it is reported for a human to judge. The order matters:
a halved axis also makes the screw along it unobservable, and counting that row twice would hide
the defect that cost the reflections.
"""
import itertools
import os
import re

import gemmi

import sgequiv

REPORT_LINE = re.compile(r"^([A-Z_0-9]+)=\s*(.*)$")


def read_report(path):
    """KEY= value lines of a rugnux _report.txt, plus the lowest-resolution shell of its two shell
    tables: LOWRES_D / LOWRES_R_MEAS from the table over its own range, REFRES_LOWRES_D /
    REFRES_LOWRES_R_MEAS from the one over the reference range (it follows REFRES_RANGE=)."""
    out = {}
    if not os.path.exists(path):
        return out
    header = None
    for line in open(path, errors="replace"):
        line = line.rstrip("\n")
        m = REPORT_LINE.match(line)
        if m:
            out[m.group(1)] = m.group(2).strip()
            continue
        f = line.lstrip("#").split()
        if f[:2] == ["d_min", "N_obs"]:
            header = f
        elif header and f and re.match(r"^\d+\.\d+$", f[0]):
            prefix = "REFRES_" if "REFRES_RANGE" in out else ""
            out[prefix + "LOWRES_D"] = f[0]
            r_meas = f[header.index("R_meas")].rstrip("%")          # "3.7%", or "-" for none
            out[prefix + "LOWRES_R_MEAS"] = "" if r_meas == "-" else str(round(float(r_meas) / 100, 4))
            header = None                  # only the first row of each table
    return out


# REFRES_* keys of the report read as numbers, stored lower-case under the same name
REFRES_KEYS = ("refres_shells_past_limit", "refres_unique_reflections", "refres_completeness",
               "refres_multiplicity", "refres_i_over_sigma", "refres_r_meas", "refres_cc_half",
               "refres_isa")


def fnum(s):
    try:
        return float(s.split()[0])
    except (AttributeError, ValueError, IndexError):
        return None


def space_group(name, sgno):
    """The space group in the SETTING the cell was given in: by name when there is one (I 1 2 1 and
    C 1 2 1 are both number 5, with different centring vectors), else the standard setting."""
    sg = gemmi.find_spacegroup_by_name(name) if name else None
    return sg or gemmi.find_spacegroup_by_number(int(sgno or 1))


def primitive_reduced(cell, sg):
    """Niggli-reduced PRIMITIVE cell of a cell given in space group sg (a gemmi.SpaceGroup), and
    its volume.

    Setting-invariant: two settings of one lattice (C2 vs I2, a/c swapped, beta vs 180-beta)
    reduce to the same cell, so they compare equal."""
    gv = gemmi.GruberVector(gemmi.UnitCell(*cell), sg)
    gv.niggli_reduce()
    red = gv.cell_parameters()
    return red, gemmi.UnitCell(*red).volume


def lattice_match(cell, sg, ref_cell, ref_sg):
    """(primitive volume ratio, largest relative deviation of the reduced edges)."""
    red, vol = primitive_reduced(cell, sg)
    ref_red, ref_vol = primitive_reduced(ref_cell, ref_sg)
    edge_dev = max(abs(a - b) / b for a, b in zip(sorted(red[:3]), sorted(ref_red[:3])))
    return vol / ref_vol, edge_dev


def cell_dev_pct(cell, ref):
    """Largest relative deviation of a, b, c in percent, minimised over the six axis orders, so
    P212121 with a and c exchanged is not scored as a 6% error."""
    return min(max(abs(p[i] - ref[i]) / ref[i] * 100 for i in range(3))
               for p in itertools.permutations(cell[:3]))


# XDS prints CC1/2 in % with one decimal; its value is taken at the bottom of that rounding, so a
# printed 100.0 does not demand a perfect CC1/2 of rugnux
XDS_CC_HALF_ROUNDING = 0.0005


def cc_half_noise_ratio(cc, cc_xds):
    """rugnux's half-set noise over XDS's, both read off CC1/2 over XDS's range; None without both.

    CC1/2 = S / (S + E), with S the variance of the signal and E that of the half-set error, so
    E / S = 1 / CC1/2 - 1 at any CC1/2. E goes as 1 / (observations): a ratio of 2 is a merge as
    noisy as XDS's would be with half of its observations. Reported, not scored.
    A CC1/2 at or below 0 is no signal at all, counted as 0.001."""
    if cc is None or cc_xds is None:
        return None
    return round((1 / max(cc, 0.001) - 1) / (1 / (cc_xds - XDS_CC_HALF_ROUNDING) - 1), 3)


_AXIS_ROW = {"a": (1, 0, 0), "b": (0, 1, 0), "c": (0, 0, 1)}


def screw_axes_differing(sg, ref_sg):
    """The axes on which two space groups disagree about a screw: the principal row one of them
    extinguishes and the other does not.

    Only ever asked of two groups of the SAME point group on the same lattice, so their centrings
    are the same and a row a centring already kills is killed in both - what is left is the screw."""
    out = set()
    for ax, hkl in _AXIS_ROW.items():
        a, b = (s.operations().is_systematically_absent(hkl) for s in (sg, ref_sg))
        if a != b:
            out.add(ax)
    return out


def point_group(sgno):
    return gemmi.find_spacegroup_by_number(int(sgno)).point_group_hm()


def sg_name(sgno):
    return gemmi.find_spacegroup_by_number(int(sgno)).hm if sgno else None


def judge(entry, rep, run_note):
    """Score one set. entry is its manifest row (with 'ref'), rep its parsed report, run_note
    what the runner saw ('', 'timeout', 'exit N: message', 'no input')."""
    ref = dict(entry.get("ref") or {})
    if entry.get("ref_override"):          # a reference measured by hand, replacing the file's
        ref.update(entry["ref_override"])
    r = {
        "sgno": None, "sg": None, "pg": None,
        "sgno_ref": ref.get("sgno"), "sg_ref": ref.get("sg") or sg_name(ref.get("sgno")),
        "pg_ref": point_group(ref["sgno"]) if ref.get("sgno") else None,
        "cell": None, "cell_ref": ref.get("cell"), "cell_dev_pct": None, "volume_ratio": None,
        "d_min": None, "d_min_ref": ref.get("dmin"), "d_min_ref_rule": ref.get("dmin_rule"),
        "d_min_xds": ref.get("dmin_xds"), "res_gain_pct": None,
        "r_meas": fnum(rep.get("R_MEAS")), "cc_half": fnum(rep.get("CC_HALF")),
        "isa": fnum(rep.get("ISA")), "completeness": fnum(rep.get("COMPLETENESS")),
        "multiplicity": fnum(rep.get("MULTIPLICITY")), "i_over_sigma": fnum(rep.get("I_OVER_SIGMA")),
        "indexing_rate": fnum(rep.get("INDEXING_RATE")), "images": fnum(rep.get("IMAGES_PROCESSED")),
        "rugnux_wall_s": fnum(rep.get("WALL_TIME")), "rugnux_verdict": rep.get("VERDICT_TEXT"),
        "r_meas_ref": ref.get("r_meas"), "cc_half_ref": ref.get("cc_half"), "isa_ref": ref.get("isa"),
        "completeness_ref": ref.get("completeness"), "multiplicity_ref": ref.get("multiplicity"),
        "sg_relation": None,
        # rugnux's own fit of the deposited model (--model, open arm): a rigid-body placement scored
        # on rugnux's own free set, so a trend number, not the depositor's R-free
        "rfree": fnum(rep.get("R_FREE")), "rwork": fnum(rep.get("R_WORK")),
        "model_fit": rep.get("MODEL_FIT"), "cc_model": fnum(rep.get("CC_MODEL_OVERALL")),
        "sg_label": None,
    }
    # the same merge binned over the reference's range (--report-resolution, XDS arms)
    r.update({k: fnum(rep.get(k.upper())) for k in REFRES_KEYS})
    r["refres_range"] = rep.get("REFRES_RANGE")
    r["isa_ratio"] = round(r["refres_isa"] / r["isa_ref"], 4) if r["refres_isa"] and r["isa_ref"] else None
    r["r_meas_ratio"] = (round(r["refres_r_meas"] / r["r_meas_ref"], 4)
                         if r["refres_r_meas"] and r["r_meas_ref"] else None)
    # R_meas of the lowest-resolution shell: rugnux's own table, the reference-range table and XDS's
    # (reported, not scored)
    r.update(lowres_d=fnum(rep.get("LOWRES_D")), lowres_r_meas=fnum(rep.get("LOWRES_R_MEAS")),
             refres_lowres_d=fnum(rep.get("REFRES_LOWRES_D")),
             refres_lowres_r_meas=fnum(rep.get("REFRES_LOWRES_R_MEAS")),
             lowres_d_ref=ref.get("dmin_low"), lowres_r_meas_ref=ref.get("r_meas_low"))
    r["lowres_r_meas_ratio"] = (round(r["refres_lowres_r_meas"] / r["lowres_r_meas_ref"], 4)
                                if r["refres_lowres_r_meas"] and r["lowres_r_meas_ref"] else None)
    r["cc_half_noise_ratio"] = cc_half_noise_ratio(r["refres_cc_half"], r["cc_half_ref"])
    if rep.get("SPACE_GROUP_NUMBER"):
        r["sgno"] = int(fnum(rep["SPACE_GROUP_NUMBER"]))
        r["sg"] = rep.get("SPACE_GROUP_NAME") or sg_name(r["sgno"])
        # With --model the reported group can be the model's enantiomorph, a label the data did not
        # decide. The data's own answer is then the search's Sohncke group (same class up to hand).
        if rep.get("SPACE_GROUP_ENANTIOMORPH") == "ASSUMED_FROM_MODEL" and rep.get("SOHNCKE_SPACE_GROUP"):
            own = gemmi.find_spacegroup_by_name(rep["SOHNCKE_SPACE_GROUP"])
            if own and own.number != r["sgno"] and sgequiv.indistinguishable(r["sgno"], own.number):
                r["sg_label"] = r["sg"]
                r["sgno"], r["sg"] = own.number, own.hm
        r["pg"] = point_group(r["sgno"])
    if rep.get("UNIT_CELL_CONSTANTS"):
        r["cell"] = [float(x) for x in rep["UNIT_CELL_CONSTANTS"].split()[:6]]
    rng = (rep.get("INCLUDE_RESOLUTION_RANGE") or "").split()
    if len(rng) == 2:
        r["d_min"] = float(rng[1])
    if r["d_min"] and r["d_min_ref"]:
        r["res_gain_pct"] = round((r["d_min_ref"] - r["d_min"]) / r["d_min_ref"] * 100, 1) + 0.0  # no -0.0
    if r["cell"] and r["cell_ref"]:
        r["cell_dev_pct"] = round(cell_dev_pct(r["cell"], r["cell_ref"]), 3)

    def verdict(v, cause, reason):
        r.update(verdict=v, cause=cause, reason=reason)
        return r

    have_lattice = r["sgno"] is not None and r["cell"] is not None
    if entry.get("expect") == "no_lattice":
        # a no-crystal control: the right answer is to refuse
        if have_lattice:
            return verdict("fail", "false_lattice",
                           f"reported a lattice ({r['sg']}) on a set with no crystal")
        return verdict("pass", None, "no lattice reported, as expected")
    if not have_lattice:
        if run_note == "no input":
            return verdict("not_run", "no_input", "input file not found")
        if run_note == "timeout":
            return verdict("fail", "timeout", "timed out")
        if re.search(r"Error reading input|Cannot open", run_note):
            return verdict("fail", "reader", run_note)
        if re.search(r"index|lattice", run_note, re.I):
            return verdict("fail", "indexing", run_note)
        return verdict("fail", "crash", run_note or "no lattice in the report")
    if not ref.get("sgno") or not ref.get("cell"):
        return verdict("unscored", "no_reference", "no reference to score against")
    if entry.get("unscored"):              # a reference known to be wrong: run, but do not score
        return verdict("unscored", "reference_problem", entry["unscored"])

    vr, edge = lattice_match(r["cell"], space_group(rep.get("SPACE_GROUP_NAME"), r["sgno"]),
                             ref["cell"], space_group(ref.get("sg"), ref["sgno"]))
    r["volume_ratio"] = round(vr, 4)
    if not (0.95 < vr < 1.05 and edge < 0.02):
        if 0.45 <= vr <= 0.55:
            return verdict("fail", "lattice_halved", f"primitive volume ratio {vr:.2f}")
        if 1.9 <= vr <= 2.1:
            return verdict("fail", "lattice_doubled", f"primitive volume ratio {vr:.2f}")
        return verdict("fail", "lattice_other",
                       f"primitive volume ratio {vr:.2f}, reduced edges off by {edge:.1%}")

    # XDS never tests a screw axis, so against an XDS reference only the point group can be
    # judged - unless the group was measured by hand (sg_measured).
    by_sg = entry["arm"] == "open" or ref.get("sg_measured")
    if by_sg:
        rel = sgequiv.describe(ref["sgno"], r["sgno"])
        r["sg_relation"] = rel
        ok = rel in ("match", "hand only (needs anomalous)", "UNDECIDABLE from intensities")
    else:
        ok = r["pg"] == r["pg_ref"]
    if not ok:
        order = len(gemmi.find_spacegroup_by_number(r["sgno"]).operations().sym_ops)
        order_ref = len(gemmi.find_spacegroup_by_number(ref["sgno"]).operations().sym_ops)
        if r["pg"] == r["pg_ref"]:
            # A screw the data never tested is not a wrong answer, it is an unanswerable question:
            # the axial row it lives on was not recorded at all (the spindle's blind cone), so the
            # run offers every member of the set as an alternative and writes the one claiming no
            # screw, because a reflection file must carry one group. Scoring that as a failure
            # measures the sweep, not the program - the same reason the hand of an enantiomorphic
            # pair is not scored. Three conditions, all necessary: rugnux SAYS the axis is open,
            # the reference is among the groups it offered, and the ONLY axis the two differ on is
            # one it named. A screw called wrongly where the row WAS measured meets none of them
            # and stays a failure.
            undetermined = set((rep.get("SPACE_GROUP_SCREW_UNDETERMINED") or "NONE").split())
            offered = [gemmi.find_spacegroup_by_name(x.strip())
                       for x in (rep.get("SPACE_GROUP_ALTERNATIVES") or "").split("|")]
            differ = screw_axes_differing(space_group(rep.get("SPACE_GROUP_NAME"), r["sgno"]),
                                          space_group(ref.get("sg"), ref["sgno"]))
            if (undetermined != {"NONE"} and differ and differ <= undetermined
                    and any(o and o.number == ref["sgno"] for o in offered)):
                return verdict("unscored", "screw_undetermined",
                               f"{r['sg']} vs reference {r['sg_ref']}: the screw along "
                               f"{', '.join(sorted(differ))} is undeterminable from these data "
                               f"(offered {rep.get('SPACE_GROUP_ALTERNATIVES')})")
            return verdict("fail", "sym_screw", f"{r['sg']} vs reference {r['sg_ref']}")
        cause = ("sym_under" if order < order_ref else
                 "sym_over" if order > order_ref else "sym_other")
        return verdict("fail", cause, f"{r['sg']} vs reference {r['sg_ref']}")

    # Data quality (CC1/2, R_meas, resolution, ISa) is reported, never scored: that is for a human.
    note = ""
    if r["sgno"] != ref["sgno"]:
        note = f"{r['sg']} vs reference {r['sg_ref']} ({r.get('sg_relation') or 'screws not judged against XDS'})"
    return verdict("pass", None, note)