Jungfraujoch/rugnux_vs_xds.py

#!/usr/bin/env python3
"""
rugnux_vs_xds.py -- run rugnux over a directory of test crystals and tabulate its
merging statistics against the XDS reference (CORRECT.LP) next to each dataset.

For every <data-root>/<crystal>/ that has both a *_master.h5 and an XDS CORRECT.LP it:
  * reads the XDS reference (space group, resolution range, Friedel setting, merge
    stats, ISa);
  * runs rugnux de-novo with the SAME high-resolution limit and Friedel/anomalous
    setting as XDS, so the two are comparable -- the space group is NOT forced, it is
    the thing we compare;
  * parses rugnux's own mmCIF output (the stable interface -- NOT the console log,
    which is meant for humans and may be trimmed/reformatted);
  * prints one compact table: reflections, observations, space group (H-M + number),
    R_meas / CC1/2 overall, R_meas low shell, CC1/2 high shell, ISa -- XDS vs rugnux.

Self-contained: on first run it builds a private venv, pip-installs gemmi and re-execs.
All rugnux console output goes to a per-crystal log; only the table is printed.
"""

import os
import sys
import subprocess
import pathlib
import venv

_SELF = os.path.abspath(__file__)


# --------------------------------------------------------------------------- #
#  venv bootstrap (gemmi is the only third-party dependency)
# --------------------------------------------------------------------------- #
def _venv_dir():
    if os.environ.get("RUGNUX_CMP_VENV"):
        return pathlib.Path(os.environ["RUGNUX_CMP_VENV"])
    cache = os.environ.get("XDG_CACHE_HOME") or os.path.join(os.path.expanduser("~"), ".cache")
    return pathlib.Path(cache) / "rugnux_vs_xds" / "venv"


def _bootstrap():
    vdir = _venv_dir()
    py = vdir / "bin" / "python"
    if not py.exists():
        venv.create(vdir, with_pip=True)
    have_gemmi = subprocess.run([str(py), "-c", "import gemmi"],
                                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0
    if not have_gemmi:
        subprocess.run([str(py), "-m", "pip", "install", "-q", "--disable-pip-version-check", "gemmi"],
                       check=True)
    os.environ["RUGNUX_CMP_BOOTSTRAPPED"] = "1"
    os.execv(str(py), [str(py), _SELF, *sys.argv[1:]])


if __name__ == "__main__" and not os.environ.get("RUGNUX_CMP_BOOTSTRAPPED"):
    _bootstrap()

# ---- past this point we are running inside the venv -----------------------
import argparse
import re
import shutil
import time
import datetime
import gemmi


# --------------------------------------------------------------------------- #
#  XDS reference (CORRECT.LP)
# --------------------------------------------------------------------------- #
def parse_xds(correct_lp):
    """Pull the final merge statistics out of an XDS CORRECT.LP."""
    txt = pathlib.Path(correct_lp).read_text(errors="replace")
    r = {}

    # The first SPACE_GROUP_NUMBER is only the INTEGRATE-step group (often P1); the LAST one is
    # the symmetry CORRECT actually merged/scaled in and wrote to XDS_ASCII.HKL -- that is the
    # XDS space group we compare (and the one the stats below are computed in).
    sgs = re.findall(r"SPACE_GROUP_NUMBER=\s*(\d+)", txt)
    r["sg"] = int(sgs[-1]) if sgs else None

    m = re.search(r"FRIEDEL'S_LAW=\s*(TRUE|FALSE)", txt)
    r["anomalous"] = (m.group(1) == "FALSE") if m else False

    m = re.search(r"INCLUDE_RESOLUTION_RANGE=\s*([\d.]+)\s+([\d.]+)", txt)
    include_low = float(m.group(1)) if m else 0.0
    include_high = float(m.group(2)) if m else 0.0

    m = re.search(r"^\s*a\s+b\s+ISa\s*\n\s*[\d.Ee+-]+\s+[\d.Ee+-]+\s+([\d.]+)", txt, re.M)
    r["isa"] = float(m.group(1)) if m else None

    # Final "... AS FUNCTION OF RESOLUTION" table (last occurrence); rows run from the
    # lowest-resolution shell down to the highest, then a "total" line.
    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 not f:
            continue
        if f[0] == "total" and len(f) >= 11:
            total = f
            break
        if re.match(r"^\d+\.\d+$", f[0]) and len(f) >= 11:
            shells.append(f)
    if total:
        r["obs"], r["uniq"] = int(total[1]), int(total[2])
        r["rmeas"], r["cc"] = num(total[9]), num(total[10])
    if shells:
        r["rmeas_lo"] = num(shells[0][9])    # first row  = lowest resolution
        r["cc_hi"] = num(shells[-1][10])     # last row   = highest resolution
    # Effective XDS high-resolution limit to hand to rugnux: prefer the explicit
    # INCLUDE_RESOLUTION_RANGE, but fall back to the finest shell actually tabulated
    # (XDS.INP often leaves the high limit at 0.0 = "use the full detector range").
    table_high = float(shells[-1][0]) if shells else 0.0
    r["dmin"] = (include_high if include_high > 0 else table_high) or None
    # Same for the LOW-resolution limit. XDS configurations normally cut at 50 A while rugnux
    # cuts at its own default, so without matching this the "lowest shell" columns below are not
    # the same shell in the two programs and R_meas_lo is not comparable.
    r["dmax"] = include_low or None

    # Mosaicity for comparison only. XDS refines this per run and the last value is the one it
    # settled on; it is a useful sanity check on rugnux's own estimate but NOT ground truth --
    # XDS produces outliers here too, so read a big disagreement as "look at this crystal",
    # not as "rugnux is wrong".
    mos = re.findall(r"REFLECTING_RANGE_E\.S\.D\.=\s*([0-9.]+)", txt)
    r["mosaicity"] = float(mos[-1]) if mos else None
    return r


def xds_integrate_mosaicity(integrate_lp):
    """Mean per-batch SIGMAR from INTEGRATE.LP - XDS's INTEGRATION-STAGE estimate.

    This, not CORRECT.LP's REFLECTING_RANGE_E.S.D., is the apples-to-apples comparison for
    rugnux's per-image mosaicity: both are fitted from the images during integration. The
    CORRECT.LP value is XDS's POST-REFINED mosaicity, and the two can differ by more than 2x.
    """
    if not integrate_lp.exists():
        return None
    vals = []
    for line in integrate_lp.read_text(errors="replace").splitlines():
        if line.strip().startswith("SIGMAR (degree)"):
            for f in line.split()[2:]:
                try:
                    vals.append(float(f))
                except ValueError:
                    pass
    return sum(vals) / len(vals) if vals else None


# --------------------------------------------------------------------------- #
#  rugnux result (mmCIF)
# --------------------------------------------------------------------------- #
def rugnux_mosaicity(workdir, name):
    """Median per-image mosaicity (deg) from rugnux's <prefix>_image.dat, or None."""
    import statistics
    hits = sorted(workdir.glob(f"{name}*_image.dat"))
    if not hits:
        return None
    vals = []
    for line in hits[-1].read_text(errors="replace").splitlines():
        if line.startswith("#"):
            continue
        f = line.split()
        if len(f) >= 3:
            try:
                v = float(f[2])
            except ValueError:
                continue
            if v == v and v > 0:          # skip nan / unfitted frames
                vals.append(v)
    return statistics.median(vals) if vals else None


def parse_rugnux_cif(cif_path):
    """Pull the same statistics out of rugnux's mmCIF (fractions -> percent)."""
    block = gemmi.cif.read(str(cif_path)).sole_block()

    def val(tag):
        v = block.find_value(tag)
        return gemmi.cif.as_string(v) if v is not None else None

    def fnum(tag, scale=1.0):
        v = val(tag)
        try:
            return float(v) * scale
        except (TypeError, ValueError):
            return None

    def inum(tag):
        v = val(tag)
        try:
            return int(float(v))
        except (TypeError, ValueError):
            return None

    r = {
        "sg": inum("_symmetry.Int_Tables_number"),
        "uniq": inum("_reflns.number_obs"),
        "obs": inum("_reflns.pdbx_number_measured_all"),
        "rmeas": fnum("_reflns.pdbx_Rrim_I_all", 100.0),
        "cc": fnum("_reflns.pdbx_CC_half", 100.0),
        "isa": fnum("_reflns.jfjoch_diffrn_ISa"),
    }

    # One unreadable column must not discard the whole shell. An outer shell whose R_meas is
    # undefined - which happens exactly when the mean intensity there has gone non-positive, i.e. on
    # the WORST arm - used to take its CC1/2 down with it, so the comparison silently fell back to
    # the next shell in and reported that arm's high-resolution CC1/2 as better than it is. Each
    # column is now read on its own, and each statistic comes from the outermost (or innermost)
    # shell that actually has it.
    def maybe(v):
        try:
            return float(v)
        except (TypeError, ValueError):
            return None

    shells = []
    for row in block.find("_reflns_shell.", ["d_res_high", "pdbx_Rrim_I_all", "pdbx_CC_half"]):
        d = maybe(row[0])
        if d is not None:
            shells.append((d, maybe(row[1]), maybe(row[2])))
    if shells:
        shells.sort(key=lambda s: s[0])
        cc = next((s[2] for s in shells if s[2] is not None), None)          # highest resolution
        rmeas = next((s[1] for s in reversed(shells) if s[1] is not None), None)   # lowest
        if cc is not None:
            r["cc_hi"] = cc * 100.0
        if rmeas is not None:
            r["rmeas_lo"] = rmeas * 100.0
    return r


# --------------------------------------------------------------------------- #
#  running rugnux
# --------------------------------------------------------------------------- #
def find_rugnux(explicit):
    for cand in (explicit, os.environ.get("RUGNUX")):
        if cand and (shutil.which(cand) or os.path.exists(cand)):
            return shutil.which(cand) or cand
    if shutil.which("rugnux"):
        return shutil.which("rugnux")
    repo = pathlib.Path(_SELF).parent
    # The rugnux CLI target builds to build*/rugnux/rugnux. The old build*/tools/rugnux paths are
    # stale leftovers from a previous layout that cmake no longer updates, so pick the most recently
    # built binary among the candidates rather than a fixed order (avoids silently running old code).
    cands = [repo / rel for rel in ("build_gpu/rugnux/rugnux", "build/rugnux/rugnux",
                                    "cmake-build-release-nogpu/rugnux/rugnux")]
    existing = [c for c in cands if c.exists()]
    if existing:
        return str(max(existing, key=lambda p: p.stat().st_mtime))
    return None


def run_rugnux(master, workdir, name, xds, rugnux_bin, threads, timeout, reuse, extra_args="", force_sg=False):
    """Run rugnux; return (cif_path, error). Resolution + anomalous matched to XDS."""
    workdir.mkdir(parents=True, exist_ok=True)
    cif = workdir / f"{name}.cif"
    if reuse and cif.exists():
        return cif, None, None   # no fresh timing when we skip the run

    # rugnux always writes both MTZ and CIF now (the old --scaling-output flag was removed).
    cmd = [rugnux_bin, "-o", name]
    if xds.get("anomalous"):
        cmd.append("-A")
    if xds.get("dmin"):
        cmd += ["--scaling-high-resolution", f"{xds['dmin']:.3f}"]
    if xds.get("dmax"):
        cmd += ["--scaling-low-resolution", f"{xds['dmax']:.3f}"]
    if threads:
        cmd += ["-N", str(threads)]
    if force_sg and xds.get("sg"):
        cmd += ["-S", str(xds["sg"])]
    if extra_args:
        import shlex
        cmd += shlex.split(extra_args)
    cmd.append(str(master))

    log = workdir / "rugnux.log"
    t0 = time.monotonic()
    with open(log, "w") as lf:
        lf.write("$ " + " ".join(cmd) + "\n\n")
        lf.flush()
        try:
            p = subprocess.run(cmd, cwd=str(workdir), stdout=lf,
                               stderr=subprocess.STDOUT, timeout=timeout)
        except subprocess.TimeoutExpired:
            return None, f"timeout after {timeout}s", float(timeout)
    dt = time.monotonic() - t0
    if p.returncode != 0:
        return None, f"exit {p.returncode} (see {log})", dt
    if cif.exists():
        return cif, None, dt
    hits = sorted(workdir.glob("*.cif"))
    return (hits[-1], None, dt) if hits else (None, f"no .cif produced (see {log})", dt)


# --------------------------------------------------------------------------- #
#  rendering
# --------------------------------------------------------------------------- #
def sg_name(n):
    sg = gemmi.find_spacegroup_by_number(n) if n else None
    if not sg:
        return "?"
    try:
        return sg.short_name()
    except Exception:
        return sg.hm.replace(" ", "")


def sym_key(n):
    """Merge-relevant symmetry key = point group + lattice centring. Treats screw-axis /
    enantiomorph variants as equal (XDS reports e.g. P422=89 where rugnux finds P4_3 2_1 2=96;
    I23=197 vs I2_1 3=199) while keeping genuinely different lattices (P2 vs C2) distinct."""
    sg = gemmi.find_spacegroup_by_number(n) if n else None
    return f"{sg.point_group_hm()}{sg.hm[0]}" if sg else None


W = dict(src=6, refl=8, obs=9, sg=14, rmeas=7, cc=6, rlo=8, chi=7, isa=6, mos=15, time=8)


def fmt_dur(s):
    """Compact wall-clock: 42s / 2m05s / 1h03m."""
    if s is None:
        return None
    s = int(round(s))
    if s < 60:
        return f"{s}s"
    if s < 3600:
        return f"{s // 60}m{s % 60:02d}s"
    return f"{s // 3600}h{(s % 3600) // 60:02d}m"


def _int(v, w):
    return "-".rjust(w) if v is None else f"{v:>{w}d}"


def _pct(v, w):
    return "-".rjust(w) if v is None else f"{v:.1f}%".rjust(w)


def _isa(v, w):
    return "-".rjust(w) if v is None else f"{v:.2f}".rjust(w)


def _sg(n, w):
    return (f"{sg_name(n)} ({n})" if n else "-").ljust(w)


def _mos(v, w, v2=None):
    a = "-" if v is None else f"{v:.4f}"
    if v2 is not None:
        a += f"|{v2:.4f}"
    return a.rjust(w)


def _dur(v, w):
    return (fmt_dur(v) or "-").rjust(w)


def row_line(src, d, elapsed=None):
    d = d or {}
    return (f"  {src:<{W['src']}} "
            f"{_int(d.get('uniq'), W['refl'])} "
            f"{_int(d.get('obs'), W['obs'])}  "
            f"{_sg(d.get('sg'), W['sg'])} "
            f"{_pct(d.get('rmeas'), W['rmeas'])} "
            f"{_pct(d.get('cc'), W['cc'])} "
            f"{_pct(d.get('rmeas_lo'), W['rlo'])} "
            f"{_pct(d.get('cc_hi'), W['chi'])} "
            f"{_isa(d.get('isa'), W['isa'])} "
            f"{_mos(d.get('mosaicity'), W['mos'], d.get('mosaicity_mle'))} "
            f"{_dur(elapsed, W['time'])}")


def header_line():
    return (f"  {'Source':<{W['src']}} {'Refl':>{W['refl']}} {'Obs':>{W['obs']}}  "
            f"{'Space group':<{W['sg']}} {'Rmeas':>{W['rmeas']}} {'CC1/2':>{W['cc']}} "
            f"{'Rmeas_lo':>{W['rlo']}} {'CC1/2_hi':>{W['chi']}} {'ISa':>{W['isa']}} "
            f"{'Mosaic':>{W['mos']}} {'Time':>{W['time']}}")


# --------------------------------------------------------------------------- #
#  main
# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser(description="Compare rugnux merge stats against XDS CORRECT.LP.")
    ap.add_argument("data_root", nargs="?", default="/data/rotation_test",
                    help="directory of <crystal>/ subdirs (default: /data/rotation_test)")
    ap.add_argument("--rugnux", help="path to the rugnux binary (else $RUGNUX / PATH / build*/tools)")
    ap.add_argument("--workdir", default="rugnux_cmp", help="where rugnux outputs go (default: ./rugnux_cmp)")
    ap.add_argument("--only", help="comma-separated substrings; only matching crystals are run")
    ap.add_argument("--threads", type=int, help="rugnux -N (default: all hardware threads)")
    ap.add_argument("--timeout", type=int, default=3600, help="per-crystal timeout, seconds (default: 3600)")
    ap.add_argument("--reuse", action="store_true", help="skip rugnux where its .cif already exists")
    ap.add_argument("--xds-only", action="store_true", help="only parse+print the XDS side (no rugnux run)")
    ap.add_argument("--progress", action="store_true", help="print per-crystal progress to stderr")
    ap.add_argument("--force-xds-sg", action="store_true",
                    help="pin each crystal to the XDS reference space group with -S, instead of "
                         "letting rugnux determine it de novo (exercises the fixed-group path)")
    ap.add_argument("--extra-args", default="", help="extra args appended to every rugnux command (e.g. '--rotation-refine-geometry')")
    args = ap.parse_args()

    root = pathlib.Path(args.data_root)
    if not root.is_dir():
        sys.exit(f"data root not found: {root}")

    rugnux_bin = None if args.xds_only else find_rugnux(args.rugnux)
    if not args.xds_only and not rugnux_bin:
        sys.exit("rugnux binary not found -- pass --rugnux PATH or set $RUGNUX")

    filters = [s for s in (args.only or "").split(",") if s]
    workdir = pathlib.Path(args.workdir)

    crystals, skipped = [], []
    for name in sorted(os.listdir(root)):
        d = root / name
        if not d.is_dir():
            continue
        if filters and not any(f in name for f in filters):
            continue
        correct = d / "CORRECT.LP"
        master = sorted(d.glob("*_master.h5"))
        if not correct.exists() or not master:
            skipped.append(name)
            continue
        crystals.append((name, master[0], correct))

    results = []
    for i, (name, master, correct) in enumerate(crystals, 1):
        xds = parse_xds(correct)
        xds["mosaicity_mle"] = xds_integrate_mosaicity(correct.parent / "INTEGRATE.LP")
        rug, err, elapsed = ({}, None, None)
        if not args.xds_only:
            if args.progress:
                print(f"[{i}/{len(crystals)}] rugnux {name} "
                      f"(dmin={xds.get('dmin')}, dmax={xds.get('dmax')}, "
                      f"anom={xds.get('anomalous')}) ...",
                      file=sys.stderr, flush=True)
            cif, err, elapsed = run_rugnux(master, workdir / name, name, xds,
                                           rugnux_bin, args.threads, args.timeout, args.reuse,
                                           args.extra_args, args.force_xds_sg)
            if cif:
                try:
                    rug = parse_rugnux_cif(cif)
                    rug["mosaicity"] = rugnux_mosaicity(workdir / name, name)
                except Exception as e:  # keep the table alive if one cif is unreadable
                    err = f"cif parse failed: {e}"
        results.append((name, xds, rug, err, elapsed))

    # -------- print the table --------
    hdr = header_line()
    rule = "  " + "-" * (len(hdr) - 2)
    today = datetime.date.today().isoformat()
    print()
    print(f"  rugnux vs XDS   ·   {root}   ·   {today}")
    if not args.xds_only:
        print(f"  rugnux: {rugnux_bin}")
    print(f"  rugnux run de-novo, resolution range (both limits) + Friedel matched to XDS.")
    print(f"  XDS SG = symmetry")
    print(f"  CORRECT merged in; match is point-group level (screw/enantiomorph ignored).  [ - = missing ]")
    print(f"  Mosaic (deg): XDS cell = post-refined (CORRECT.LP) | integration-stage MLE (INTEGRATE.LP);")
    print(f"  rugnux cell = median per-image mosaicity.  Compare rugnux against the INTEGRATION-STAGE")
    print(f"  number - the post-refined one measures something else and can differ by >2x.")
    print(f"  Reference only: XDS produces outliers here too, badly so on weak / low-resolution data.")
    print()
    print(hdr)
    print(rule)

    n_match = 0
    for name, xds, rug, err, elapsed in results:
        if args.xds_only:
            mark = ""
        elif err:
            mark = "  FAIL"
        else:
            xk, rk = sym_key(xds.get("sg")), sym_key(rug.get("sg"))
            same = xk is not None and xk == rk
            n_match += same
            mark = "  SG OK  ✅" if same else "  SG DIFF ❌"
        print(f"● {name}{mark}")
        print(row_line("XDS", xds))
        if args.xds_only:
            continue
        if err:
            print(f"  {'rugnux':<{W['src']}} (failed: {err})  [{fmt_dur(elapsed) or '-'}]")
        else:
            print(row_line("rugnux", rug, elapsed))
    print(rule)

    if not args.xds_only:
        n_fail = sum(1 for *_, e, _ in results if e)
        total = sum(e for *_, e in results if e)
        print(f"  {len(results)} crystals   ·   space group OK {n_match}/{len(results)}"
              f"   ·   rugnux failed {n_fail}   ·   total rugnux time {fmt_dur(total) or '-'}")
    if skipped:
        print(f"  skipped (no CORRECT.LP or no *_master.h5): {', '.join(skipped)}")
    print()


if __name__ == "__main__":
    main()