Jungfraujoch/tools/battery/paper_plots.py

#!/usr/bin/env python3
"""paper_plots.py -- three battery figures of the open arm: R-free and high-resolution limit
(deposition vs Rugnux) and the processing time per dataset.

    tools/battery/paper_plots.py RUN [RUN ...] --out DIR [--time-runs RUN [RUN ...]] [--uncontended RUN]

RUN is a battery run directory (or its name under the site's runs_root); later runs fill in sets
an earlier one lacks or that failed to run there (a `-redo` of CUDA casualties). --time-runs draws
the time histogram from other runs than the two quality figures - a run whose sets had the GPU to
themselves, where the runs for the quality figures had to share it. Only the open arm is read. Everything the figures are drawn from is written
beside them, so a figure can be audited and redrawn: rfree.dat, dmin.dat (value pairs with the set
and its beamline), time.dat (set, arm, seconds, images, other GPU processes), excluded.txt (every
set left out, with the reason), summary.txt (the runs, their binary, the counts and the timing
statistics), figures.gp (the gnuplot script) and the rendered figures: a single-column vector PDF
with embedded fonts and the same figure as a 600 d.p.i. PNG. gnuplot and pdftoppm (poppler) must be
on PATH; gemmi reads the facility and beamline from each entry's cached mmCIF (_diffrn_source).

Left figure: R-free of the deposited model after the battery's REFMAC check on the depositor's data
(x) and on Rugnux's data (y). By default the refined protocol (model_check.py: the same model
refined against each data set at its own resolution, R-free over the free reflections both data sets
share); --rfree first-cycle draws the first-cycle rigid-body numbers at the common resolution instead.
Right figure: the deposited high-resolution limit (x) against Rugnux's own cut
(y), log axes so the low-resolution entries show. Small molecules (no deposited model) and sets
whose verdict is `fail` (a wrong lattice: nothing to compare) are left out.

Third figure: a histogram of the processing time of every open-arm set that ran to a report, whatever its verdict, with the median marked. The time is `rugnux_wall_s`, the
WALL_TIME rugnux itself reports over the whole invocation (reading the images through to the
written merge) - not the runner's `elapsed_s`, which also counts the wait for a GPU slot under
--gpulock (`wall_s` is the same value with `elapsed_s` as the fallback). Linear 25 s bins: the
distribution is long-tailed but spans little more than a decade, and a linear axis keeps "how long
does a dataset take" readable. Two caveats the reader must know, both written to summary.txt:
each set runs once, so a time includes reading the images from disk unless they were still in the
page cache (a cold read of a 10 GB dataset from the workstation's HDD costs of the order of a
minute), and a run under --gpulock can share the GPU with other processes - `gpu_others` counts
them as seen just before the set started. Sets that ran with a shared GPU are drawn as the lighter
top segment of their bar and are never dropped; with --uncontended RUN (a run of the same sets
with the GPU to itself) summary.txt also reports how much slower they were than there. Sets that
crashed have no WALL_TIME and go to excluded.txt.
"""
import argparse
import collections
import json
import math
import os
import re
import shutil
import subprocess
import sys

import gemmi

RUNS_ROOT = "/data/battery/runs"
RFREE_BAND = 0.05     # |R-free Rugnux - R-free depositor| counted as "within"
TIME_BIN = 25         # s, histogram bin width

# Sized and lettered to the IUCr artwork guide (journals.iucr.org/services/help/artwork/guide.html):
# one column is 8.85 cm, lettering upright in a standard font (Helvetica, embedded by cairo) and in
# the 1.5-3 mm range of the notes for authors (10 pt: capitals ~2.5 mm),
# line weights 0.35-1.5 pt at final size (pdfcairo lw 1 = 0.5 pt), no grid, a colour that stays
# legible in greyscale. Ticks are short, outward and not mirrored.
GNUPLOT = r"""
set terminal pdfcairo size 8.8cm,8.8cm font "Helvetica,10"
set encoding utf8
set size ratio 1
set border lw 1.5
set tics out scale 0.4 nomirror
set key off
set style line 1 lc rgb "black" lw 1.5 dt 2
set style line 2 pt 7 ps 0.35 lc rgb "#0072b2"
# counts: names left-aligned, numbers right-aligned in a column beside them
set label 11 "Structures\nFacilities\nBeamlines" at graph 0.05, graph 0.93 left front
set label 12 "%(n)d\n%(nf)d\n%(nb)d" at graph 0.31, graph 0.93 right front

# --- R-free
set output "%(out)s/rfree.pdf"
set xlabel "{/:Italic R}_{free}, depositor data"
set ylabel "{/:Italic R}_{free}, Rugnux data"
set xrange [%(rlo).2f:%(rhi).2f]; set yrange [%(rlo).2f:%(rhi).2f]
set xtics 0.1; set ytics 0.1
set format x "%%.1f"; set format y "%%.1f"
set label 2 "Within ±%(band).2f: %(within)d/%(n)d" at graph 0.95, graph 0.07 right front
plot x ls 1, "%(out)s/rfree.dat" using 1:2 ls 2
unset label 2

# --- high-resolution limit
set output "%(out)s/dmin.pdf"
set logscale xy
set xlabel "{/:Italic d}_{min}, deposition (Å)"
set ylabel "{/:Italic d}_{min}, Rugnux (Å)"
set xrange [%(dlo).2f:%(dhi).2f]; set yrange [%(dlo).2f:%(dhi).2f]
set xtics (0.5, 1, 2, 4, 8); set ytics (0.5, 1, 2, 4, 8)
set format x "%%.1f"; set format y "%%.1f"
set label 12 "%(dn)d\n%(dnf)d\n%(dnb)d" at graph 0.31, graph 0.93 right front
set label 3 "Rugnux deeper: %(deeper)d" at graph 0.95, graph 0.07 right front
set label 4 "Deposition deeper: %(shallower)d" at graph 0.05, graph 0.68 left front
plot x ls 1, "%(out)s/dmin.dat" using 1:2 ls 2
unset label 3; unset label 4

# --- processing time: all sets in the lighter fill, those with the GPU to themselves over them in
# full colour, so the lighter top of a bar is the sets that shared the GPU
set output "%(out)s/time.pdf"
unset logscale
set xlabel "Processing time (s)"
set ylabel "Datasets"
set xrange [0:%(thi)d]; set yrange [0:%(tyhi)d]
set xtics 100; set ytics 10
set format x "%%.0f"; set format y "%%.0f"
set boxwidth %(bin)d absolute
set style fill solid noborder
set label 11 "Datasets" at graph 0.64, graph 0.93 left front
set label 12 "%(tn)d" at graph 0.95, graph 0.93 right front
%(tkey)s
set arrow 1 from %(tmed).1f, graph 0 to %(tmed).1f, graph 1 nohead ls 1 front
set label 5 "median %(tmed).0f s" at %(tmed).1f + 0.02 * %(thi)d, graph 0.93 left front
bin(x) = (floor(x / %(bin)d) + 0.5) * %(bin)d
plot "%(out)s/time.dat" using (bin($3)):(1) smooth frequency with boxes lc rgb "#8fc0e0" title "GPU shared", \
     "" using (bin($3)):($5 > 0 ? 0 : 1) smooth frequency with boxes lc rgb "#0072b2" title "GPU exclusive"
"""


def load_rows(runs):
    """Rows of the open arm, a later run's row replacing one that did not run."""
    rows, manifests = {}, []
    for run in runs:
        path = run if os.path.isdir(run) else os.path.join(RUNS_ROOT, run)
        manifests.append((path, json.load(open(os.path.join(path, "manifest.json")))))
        for r in json.load(open(os.path.join(path, "results.json"))):
            if r["arm"] != "open":
                continue
            if r["set"] not in rows or rows[r["set"]].get("exit_code"):
                rows[r["set"]] = r
    return rows, manifests


def percentile(values, p):
    v = sorted(values)
    return v[round(p * (len(v) - 1))]


def source(model_path):
    """(facility, beamline) from the entry's _diffrn_source, facility without its operator suffix
    ("PETRA III, EMBL c/o DESY" -> "PETRA III"); None where the entry does not record one."""
    block = gemmi.cif.read(model_path).sole_block()

    def value(tag):
        v = block.find_value(tag)
        return None if v in (None, "?", ".") else gemmi.cif.as_string(v).strip()

    site = value("_diffrn_source.pdbx_synchrotron_site")
    if not site:
        return None
    beamline = value("_diffrn_source.pdbx_synchrotron_beamline") or "?"
    facility = site.split(",")[0].strip().upper()
    return facility, facility + " " + re.sub(r"\s+", "", beamline.upper())


def counts(sets, where):
    known = [where[s] for s in sets if where[s]]
    return len(sets), len({f for f, _ in known}), len({b for _, b in known})


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("runs", nargs="+", help="battery run directories or names (later ones fill in)")
    ap.add_argument("--out", required=True, help="output directory")
    ap.add_argument("--rfree", choices=["refined", "first-cycle"], default="refined",
                    help="which REFMAC R-free pair to draw (default: refined, on the shared free set)")
    ap.add_argument("--time-runs", nargs="+", help="runs to draw the time histogram from "
                                                  "(default: the runs above)")
    ap.add_argument("--uncontended", help="a run of the same sets with the GPU to itself, to size the "
                                          "slowdown of the sets that shared it")
    a = ap.parse_args()
    for tool in ("gnuplot", "pdftoppm"):
        if not shutil.which(tool):
            sys.exit(tool + " not found on PATH")
    os.makedirs(a.out, exist_ok=True)
    out = os.path.abspath(a.out)

    rows, manifests = load_rows(a.runs)
    time_rows, time_manifests = load_rows(a.time_runs) if a.time_runs else (rows, [])
    excluded, rfree, dmin, where, times = [], [], [], {}, []
    for s, r in sorted(time_rows.items()):
        if r.get("rugnux_wall_s") is None:
            excluded.append((s, "time", "no WALL_TIME: " + str(r.get("reason"))))
        else:
            times.append((s, r["arm"], r["rugnux_wall_s"], r.get("images"), r.get("gpu_others")))
    for s, r in sorted(rows.items()):
        if not r.get("model"):
            excluded.append((s, "both", "no deposited model (small molecule or unpublished)"))
            continue
        if r.get("verdict") == "fail":
            excluded.append((s, "both", "verdict fail: " + str(r.get("reason"))))
            continue
        where[s] = source(r["model"])
        if a.rfree == "refined":
            dep, rug = r.get("refmac_refined_rfree_shared_depdata"), r.get("refmac_refined_rfree_shared")
        else:
            dep, rug = r.get("refmac_rfree_depdata"), r.get("refmac_rfree_depflags")
        if dep and rug:
            rfree.append((dep, rug, s))
        else:
            excluded.append((s, "rfree", "no REFMAC pair: " + str(r.get("refmac_refined_reason") or r.get("refmac_reason")
                                                                   or r.get("dep_reason") or "depositor data not scored")))
        if r.get("d_min_ref") and r.get("d_min"):
            dmin.append((r["d_min_ref"], r["d_min"], s))
        else:
            excluded.append((s, "dmin", "no deposited or Rugnux d_min"))

    bl = lambda s: (where[s][1] if where[s] else "not-recorded").replace(" ", "_")
    with open(os.path.join(out, "rfree.dat"), "w") as fh:
        fh.write("# rfree_depositor_data rfree_rugnux_data set beamline\n")
        fh.writelines(f"{x:.4f} {y:.4f} {s} {bl(s)}\n" for x, y, s in rfree)
    with open(os.path.join(out, "dmin.dat"), "w") as fh:
        fh.write("# dmin_deposition dmin_rugnux set beamline\n")
        fh.writelines(f"{x:.3f} {y:.3f} {s} {bl(s)}\n" for x, y, s in dmin)
    with open(os.path.join(out, "time.dat"), "w") as fh:
        fh.write("# set arm seconds images gpu_others\n")
        fh.writelines(f"{s} {arm} {t:.2f} {'nan' if im is None else round(im)} {'nan' if g is None else g}\n"
                      for s, arm, t, im, g in times)
    with open(os.path.join(out, "excluded.txt"), "w") as fh:
        fh.writelines(f"{s}\t{which}\t{why}\n" for s, which, why in excluded)

    n, nf, nb = counts([s for _, _, s in rfree], where)
    dn, dnf, dnb = counts([s for _, _, s in dmin], where)
    within = sum(abs(y - x) <= RFREE_BAND for x, y, _ in rfree)
    deeper = sum(y < x for x, y, _ in dmin)
    values = [v for x, y, _ in rfree for v in (x, y)]
    secs = [t for _, _, t, _, _ in times]
    tmed = percentile(secs, 0.5)
    shared = [s for s, _, _, _, g in times if g]
    tallest = max(collections.Counter(math.floor(t / TIME_BIN) for t in secs).values())
    # the shared/exclusive key only means something when some set did share the GPU
    tkey = ("set key at graph 0.95, graph 0.76 right Left reverse samplen 1.5 spacing 1.1" if shared
            else "set key off")
    params = dict(out=out, tkey=tkey, n=n, nf=nf, nb=nb, dn=dn, dnf=dnf, dnb=dnb, within=within, band=RFREE_BAND,
                  deeper=deeper, shallower=sum(y > x for x, y, _ in dmin),
                  rlo=0.05 * math.floor((min(values) - 0.01) / 0.05), rhi=0.05 * math.ceil((max(values) + 0.01) / 0.05),
                  dlo=0.5, dhi=8.0,
                  bin=TIME_BIN, thi=TIME_BIN * math.ceil(max(secs) / TIME_BIN), tyhi=10 * math.ceil((tallest + 1) / 10),
                  tmed=tmed, tn=len(times))
    gp = os.path.join(out, "figures.gp")
    open(gp, "w").write(GNUPLOT % params)
    subprocess.run(["gnuplot", gp], check=True)
    for fig in ("rfree", "dmin", "time"):    # the PNG is the PDF rasterised at 600 d.p.i.
        subprocess.run(["pdftoppm", "-r", "600", "-png", "-singlefile", os.path.join(out, fig + ".pdf"),
                        os.path.join(out, fig)], check=True)

    with open(os.path.join(out, "summary.txt"), "w") as fh:
        for path, m in manifests + [(p, dict(m, label=f"{m.get('label')} (time figure only)"))
                                    for p, m in time_manifests]:
            b = m.get("binary", {})
            fh.write(f"run {path}\n  label {m.get('label')}  binary git {b.get('source_head')} "
                     f"dirty {b.get('source_dirty')}  sha256 {b.get('sha256')}\n  flags {b.get('flags')}  "
                     f"runner git {m.get('runner_git')} dirty {m.get('runner_dirty')}\n")
        fh.write(f"R-free: {n} structures, {nf} facilities, {nb} beamlines; within +-{RFREE_BAND}: {within}\n"
                 f"d_min: {dn} structures, {dnf} facilities, {dnb} beamlines; Rugnux reaches further on "
                 f"{deeper}\n")
        queue_wait = max(time_rows[s]["elapsed_s"] - t for s, _, t, _, _ in times)
        fh.write(f"time: {len(times)} open-arm sets; field rugnux_wall_s "
                 f"(rugnux's WALL_TIME over the whole invocation, not the runner's elapsed_s, which adds the wait "
                 f"for a GPU slot: here at most {queue_wait:.1f} s)\n"
                 f"  median {tmed:.1f} s, mean {sum(secs) / len(secs):.1f} s, p10 {percentile(secs, 0.1):.1f}, "
                 f"p25 {percentile(secs, 0.25):.1f}, p75 {percentile(secs, 0.75):.1f}, p90 {percentile(secs, 0.9):.1f}, "
                 f"min {min(secs):.1f}, max {max(secs):.1f}\n"
                 f"  each set ran once: a time includes reading the images from disk unless they were still in the "
                 f"page cache\n"
                 f"  GPU shared with other processes (gpu_others > 0) on {len(shared)} sets, kept and drawn as the "
                 f"lighter segment; median without them "
                 f"{percentile([t for s, _, t, _, _ in times if s not in shared], 0.5):.1f} s\n")
        if a.uncontended:
            other, _ = load_rows([a.uncontended])
            ratio = lambda names: percentile([t / other[s]["rugnux_wall_s"] for s, _, t, _, _ in times
                                              if s in names and other.get(s, {}).get("rugnux_wall_s")], 0.5)
            exclusive = [s for s, _, _, _, g in times if not g]
            fh.write(f"  against {a.uncontended}: median time ratio {ratio(shared):.2f} on the shared-GPU sets, "
                     f"{ratio(exclusive):.2f} on the rest\n")
        fh.write("excluded (set, figure, reason): excluded.txt\n")
    print(open(os.path.join(out, "summary.txt")).read(), end="")


if __name__ == "__main__":
    main()