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 -> (XDS arms only) is the merge as good as XDS's by CC1/2. 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
CC_HALF_NOISE_LIMIT = 2.0


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
    (CC_HALF_NOISE_LIMIT) is a merge as noisy as XDS's would be with half of its observations.
    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)


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"]:
            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']}")

    # Merge quality is judged only against XDS, by CC1/2 over XDS's own range (cc_half_noise_ratio).
    # The open arm has no merge criterion.
    if r["cc_half_noise_ratio"] is not None and r["cc_half_noise_ratio"] > CC_HALF_NOISE_LIMIT:
        return verdict("fail", "merge", f"CC1/2 {r['refres_cc_half']:.4f} vs XDS {r['cc_half_ref']:.3f} "
                                        f"over XDS's range: {r['cc_half_noise_ratio']:.1f}x XDS's half-set noise")
    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)