Improve benchmarking (memory & time) evaluation

.gitignore

@@ -13,4 +13,7 @@ build/
 .*.swp
 
 evaluation/cache/
-evaluation/results/
+evaluation/results/*
+!evaluation/results/performance*.csv
+!evaluation/results/performance*.png
+evaluation/results/*_quick*

evaluation/README.md

@@ -4,7 +4,7 @@ Scripts to benchmark and validate tsdate **by hand**. They are not part of the
 package or of the unit tests (pytest only collects `tests/`), and need extra
 dependencies: `pip install -r evaluation/requirements.txt`.
 
-Run from the repository root; results go to `evaluation/results/` (git-ignored),
+Run from the repository root; results go to `evaluation/results/` (git-ignored, except `performance*.csv`/`.png`, which are committed to track progress),
 and simulations are cached in `evaluation/cache/`. Simulations use a discrete-time
 Wright-Fisher model for the most recent 20 generations (so node times are integers
 there, and the coalescent approximation does not give sub-generation nodes), then the
@@ -13,18 +13,63 @@ standard coalescent.
 | Script | Purpose |
 |---|---|
 | `accuracy.py` | Sweep a parameter (`sample_size`, `length`, `mutation_rate`, `max_iterations`) for `variational_gamma`, comparing estimated to true node times. `--modes true` dates the true topology; `inferred` dates a tsinfer topology and compares to best-matching true nodes via [tscompare](https://github.com/tskit-dev/tscompare) (see issue #301). Writes one row per node (`accuracy_*.csv`). It also compares mutation times (`mutations_*.csv`): sites with a single mutation are matched by position, and tsdate's posterior mean mutation time is compared with the midpoint of the true edge the mutation sits on (the best any method can do; the exact simulated time is also saved as `true_time_exact`). `plot.py` scores scatter, RMSE and bias against the midpoint but PIT and interval coverage against the exact time, since the posterior is for the realised mutation time; `--mutation-truth midpoint` or `exact` forces one truth everywhere. Mutation comparison needs no node matching in either mode. |
+| `afs.py` | Compare the observed site allele frequency spectrum with the expected one from branch lengths (branch AFS x mutation rate) in genomic windows, for the true times, the true topology dated by tsdate, and a tsinfer topology dated by tsdate. Writes `afs_*.csv`. |
 | `performance.py` | Wall time and peak RSS, each case in a fresh process. |
+| `compare.py` | Compares two `performance.py` CSVs case by case (see below). |
 | `plot.py` | Reads `results/*.csv` and writes PNGs: error/bias/interval coverage vs parameter (`accuracy_*_summary.png`: top row node times, bottom row mutation times), estimated-vs-true scatter (`accuracy_*_scatter.png`: columns are swept values, rows are node/mutation times for true/tsinferred topology), PIT calibration histograms (`accuracy_*_calibration.png`: node row, mutation row), AFS observed-vs-expected (`afs_*.png`), time and memory scaling. |
-| `afs.py` | Compare the observed site allele frequency spectrum with the expected one from branch lengths (branch AFS x mutation rate) in genomic windows, for the true times, the true topology dated by tsdate, and a tsinfer topology dated by tsdate. Writes `afs_*.csv`. |
 | `validation.py` | Stub: numerical checks too slow for the unit tests (issue #384). |
 
 Each CSV row records the tsdate version, git SHA and timestamp.
 
 ```
 python evaluation/accuracy.py --sweep sample_size --modes true inferred --seeds 3 --max-seeds 250 --cores 8
-python evaluation/performance.py --sizes 100 500 1000
+python evaluation/performance.py
 python evaluation/plot.py
 ```
+### Comparing performance across commits
+
+`performance.py` records the git SHA in every row (suffixed `-dirty` if tracked files
+have uncommitted changes), and `--output` names the CSV. To test whether a change
+affects speed or memory, run the same sweep on each version and compare:
+
+```
+git checkout main
+python evaluation/performance.py --repeats 3 --output base.csv
+git checkout my-branch
+python evaluation/performance.py --repeats 3 --output new.csv
+python evaluation/compare.py evaluation/results/base.csv evaluation/results/new.csv
+```
+
+For each (length, sample size) it prints the median time and memory increase for both
+versions and the ratio new/base (below 1 is better). Differences smaller than the
+min-to-max spread of the repeats are shown as `~` (noise), so use at least 3 repeats and
+keep the machine otherwise idle. Use identical options (`--sizes`, `--lengths`,
+`--inferred`) for both runs; only cases present in both files are compared. Memory is
+the RSS increase during dating, not total process memory.
+
+#### Comparing against the stored results
+
+`evaluation/results/performance.csv` and `performance.png` are committed, so they act as
+a reference to check progress against. Each row records the machine (`machine`, `cpu`,
+`cpu_cores`, `ram_gb`, `os`, `python`) as well as the tsdate version and git SHA, and
+`compare.py` prints the machine of each file. To test the current code against the
+stored run, write the new results to a *different* file (otherwise you overwrite the
+stored reference) and compare:
+
+```
+python evaluation/performance.py --repeats 3 --output new.csv   # same options as stored run
+python evaluation/compare.py evaluation/results/performance.csv evaluation/results/new.csv
+```
+
+The stored run's sweep (sizes, lengths, `--inferred`) is in the CSV, so use the same
+options; only cases present in both files are compared. If the CPUs differ, `compare.py`
+prints a note: the ratios then include hardware effects (a fanless laptop may also
+throttle), so interpret them with care. To compare with an older commit instead, use
+`git show <commit>:evaluation/results/performance.csv > old.csv`. `new.csv` is
+git-ignored. To update the stored reference, re-run without `--output`, run `plot.py`,
+and commit the two `performance.*` files. The plot shows mean ± standard error over the
+repeats (use 5 or more for a reference run).
+
 Add `--quick` to `accuracy.py` / `performance.py` for a fast smoke test. Quick runs write
 to separate `*_quick.csv` files (so they never overwrite full results), and `plot.py`
 labels their figures "quick run" and names them `*_quick*.png`.

evaluation/common.py

@@ -4,6 +4,7 @@
 import hashlib
 import os
 import pathlib
+import platform
 import subprocess
 
 import msprime
@@ -16,6 +17,56 @@
 RECOMB_RATE = 1e-8
 
 
+def cpu_name():
+    try:
+        if platform.system() == "Darwin":
+            cmd = ["sysctl", "-n", "machdep.cpu.brand_string"]
+            return subprocess.check_output(cmd, text=True).strip()
+        if platform.system() == "Linux":
+            with open("/proc/cpuinfo") as f:
+                for line in f:
+                    if line.startswith("model name"):
+                        return line.split(":", 1)[1].strip()
+    except Exception:
+        pass
+    return platform.processor() or platform.machine()
+
+
+def machine_model():
+    """Hardware model, e.g. "MacBook Air (Mac14,2)"; useful as laptops may throttle."""
+    try:
+        if platform.system() == "Darwin":
+            out = subprocess.check_output(
+                ["system_profiler", "SPHardwareDataType"],
+                text=True,
+                stderr=subprocess.DEVNULL,
+            )
+            info = dict(
+                (k.strip(), v.strip())
+                for k, _, v in (line.partition(":") for line in out.splitlines())
+            )
+            return f"{info['Model Name']} ({info['Model Identifier']})"
+        if platform.system() == "Linux":
+            dmi = pathlib.Path("/sys/devices/virtual/dmi/id")
+            parts = [
+                (dmi / f).read_text().strip() for f in ("sys_vendor", "product_name")
+            ]
+            return " ".join(parts)
+    except Exception:
+        pass
+    return ""
+
+
+def total_ram_gb():
+    try:
+        if platform.system() == "Darwin":
+            cmd = ["sysctl", "-n", "hw.memsize"]
+            return round(int(subprocess.check_output(cmd, text=True)) / 2**30)
+        return round(os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / 2**30)
+    except Exception:
+        return None
+
+
 def provenance():
     """Columns identifying the code version, to allow comparing runs over commits."""
     import tsdate
@@ -24,11 +75,23 @@ def provenance():
         sha = subprocess.check_output(
             ["git", "rev-parse", "--short", "HEAD"], cwd=HERE, text=True
         ).strip()
+        dirty = subprocess.check_output(
+            ["git", "status", "--porcelain", "--untracked-files=no", str(HERE.parent)],
+            cwd=HERE,
+            text=True,
+        ).strip()
+        sha += "-dirty" if dirty else ""
     except Exception:
         sha = ""
     return {
         "tsdate_version": tsdate.__version__,
         "git_sha": sha,
+        "machine": machine_model(),
+        "cpu": cpu_name(),
+        "cpu_cores": os.cpu_count(),
+        "ram_gb": total_ram_gb(),
+        "os": f"{platform.system()} {platform.release()}",
+        "python": platform.python_version(),
         "timestamp": datetime.datetime.now().isoformat(timespec="seconds"),
     }
 

evaluation/compare.py

@@ -0,0 +1,69 @@
+"""
+Compare two performance runs (e.g. before and after a change) case by case.
+
+    python evaluation/compare.py results/base.csv results/new.csv
+
+Prints median time and memory increase for each (length, sample size) case, and the
+ratio new/base (below 1 is an improvement). Differences within the run-to-run spread
+of the repeats are marked "~" rather than as a change, so use --repeats >= 3 in
+performance.py for this to be meaningful. See evaluation/README.md for the workflow.
+"""
+
+import argparse
+
+import pandas as pd
+
+KEYS = ["length", "sample_size", "inferred"]
+METRICS = ["time_s", "extra_rss_mb"]
+
+
+def summarise(path):
+    df = pd.read_csv(path)
+    g = df.groupby(KEYS)
+    out = g[METRICS].median()
+    for m in METRICS:
+        out[m + "_spread"] = g[m].max() - g[m].min()
+    label = ",".join(sorted(df.git_sha.dropna().astype(str).unique()))
+    machine = ""
+    cols = [c for c in ("machine", "cpu") if c in df]
+    if cols:
+        machine = "; ".join(
+            sorted(set(df[cols].fillna("").astype(str).agg(" ".join, axis=1)))
+        )
+    return out, label, machine
+
+
+def flag(ratio, base, new, base_spread, new_spread):
+    if abs(new - base) <= max(base_spread, new_spread):
+        return "~"
+    return "faster/smaller" if ratio < 1 else "SLOWER/LARGER"
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    p.add_argument("base", help="performance CSV of the reference version")
+    p.add_argument("new", help="performance CSV of the version being tested")
+    args = p.parse_args()
+    base, base_label, base_machine = summarise(args.base)
+    new, new_label, new_machine = summarise(args.new)
+    print(f"base: {args.base} ({base_label}) {base_machine}")
+    print(f"new:  {args.new} ({new_label}) {new_machine}\n")
+    if base_machine != new_machine:
+        print("Note: CPUs differ or are unrecorded, so ratios include hardware effects\n")
+    both = base.join(new, how="inner", suffixes=("_b", "_n"))
+    if both.empty:
+        raise SystemExit("No cases in common between the two files")
+    names = {"time_s": "time (s)", "extra_rss_mb": "memory increase (MB)"}
+    for m in METRICS:
+        print(names[m])
+        print(f"{'length':>9} {'n':>6} {'base':>9} {'new':>9} {'new/base':>9}  verdict")
+        for (length, n, _), r in both.iterrows():
+            b, w = r[m + "_b"], r[m + "_n"]
+            ratio = w / b if b else float("nan")
+            verdict = flag(ratio, b, w, r[m + "_spread_b"], r[m + "_spread_n"])
+            print(f"{length:>9.0e} {n:>6} {b:>9.2f} {w:>9.2f} {ratio:>9.2f}  {verdict}")
+        print()
+
+
+if __name__ == "__main__":
+    main()

evaluation/performance.py

@@ -1,26 +1,26 @@
 """
-Time and peak memory of tsdate.variational_gamma. Each case runs in a fresh
-(spawned) process so that peak RSS is not polluted by earlier cases. Writes
-results/performance.csv; use plot.py to make figures.
+Time and memory of tsdate.variational_gamma. Each case runs in a fresh (spawned)
+process so that memory is not polluted by earlier cases. Memory is the increase in
+resident set size (RSS) during dating over the level just before it, sampled from
+the parent process (the process itself would be starved by jitted code holding the
+GIL). Writes results/performance.csv; use plot.py to make figures.
 
 Example:
-  python evaluation/performance.py --sizes 100 500 1000 --lengths 1e6 --repeats 3
+  python evaluation/performance.py --sizes 10 50 100 500 1000 --lengths 1e6 1e7
 """
 
 import argparse
 import multiprocessing
-import resource
-import sys
+import os
+import queue as queue_lib
 import time
 
 import common
 import pandas as pd
+import psutil
 import tqdm
 
-
-def _peak_rss_mb():
-    rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
-    return rss / 1e6 if sys.platform == "darwin" else rss / 1e3  # bytes vs KB
+POLL_S = 0.01
 
 
 def _case(queue, n, length, mutation_rate, seed, inferred):
@@ -33,45 +33,79 @@ def _case(queue, n, length, mutation_rate, seed, inferred):
     tsdate.date(
         common.simulate(5, 1e4, 1e-8, 1), mutation_rate=1e-8, method="variational_gamma"
     )
-    baseline = _peak_rss_mb()
+    queue.put(("start", None))
     start = time.perf_counter()
     tsdate.date(ts, mutation_rate=mutation_rate, method="variational_gamma")
     elapsed = time.perf_counter() - start
     queue.put(
-        {
-            "num_nodes": ts.num_nodes,
-            "num_edges": ts.num_edges,
-            "num_trees": ts.num_trees,
-            "num_mutations": ts.num_mutations,
-            "time_s": elapsed,
-            "peak_rss_mb": _peak_rss_mb(),
-            "baseline_rss_mb": baseline,
-        }
+        (
+            "done",
+            {
+                "num_nodes": ts.num_nodes,
+                "num_edges": ts.num_edges,
+                "num_trees": ts.num_trees,
+                "num_mutations": ts.num_mutations,
+                "time_s": elapsed,
+            },
+        )
     )
 
 
+def _rss_mb(ps):
+    try:
+        return ps.memory_info().rss / 1e6
+    except psutil.NoSuchProcess:
+        return None
+
+
 def run_case(**kwargs):
     ctx = multiprocessing.get_context("spawn")
     queue = ctx.Queue()
     proc = ctx.Process(target=_case, args=(queue,), kwargs=kwargs)
     proc.start()
-    result = queue.get()
+    ps = psutil.Process(proc.pid)
+    baseline = peak = None
+    while True:
+        try:
+            kind, result = queue.get(timeout=POLL_S)
+        except queue_lib.Empty:
+            rss = _rss_mb(ps)
+            if rss is not None and baseline is not None:
+                peak = max(peak, rss)
+            continue
+        if kind == "start":
+            baseline = peak = _rss_mb(ps)
+        else:
+            break
     proc.join()
-    return result
+    return dict(
+        result,
+        baseline_rss_mb=baseline,
+        peak_rss_mb=peak,
+        extra_rss_mb=peak - baseline,
+    )
 
 
 def main():
     p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
-    p.add_argument("--sizes", nargs="+", type=int, default=[100, 500, 1000])
-    p.add_argument("--lengths", nargs="+", type=float, default=[1e6])
+    p.add_argument("--sizes", nargs="+", type=int, default=[10, 50, 100, 500, 1000])
+    p.add_argument("--lengths", nargs="+", type=float, default=[1e6, 1e7])
     p.add_argument("--mutation-rate", type=float, default=1e-8)
     p.add_argument("--repeats", type=int, default=3)
     p.add_argument("--inferred", action="store_true", help="Use a tsinfer topology")
+    p.add_argument(
+        "--output",
+        help="Output CSV name in results/ (default performance.csv); give each "
+        "commit its own name so runs can be compared with compare.py",
+    )
     p.add_argument("--quick", action="store_true", help="Tiny smoke-test sizes")
     args = p.parse_args()
     if args.quick:
         args.sizes, args.lengths, args.repeats = [10, 30], [1e5], 1
 
+    # Compiled code is not part of the timing or the memory increase (it is warmed up
+    # first), so cache it on disk to avoid recompiling in every fresh process
+    os.environ.setdefault("TSDATE_ENABLE_NUMBA_CACHE", "1")
     prov = common.provenance()
     rows = []
     total = len(args.sizes) * len(args.lengths) * args.repeats
@@ -100,13 +134,13 @@ def main():
                 )
                 tqdm.tqdm.write(
                     f"n={n} L={length:g} seed={seed}: {res['time_s']:.2f}s, "
-                    f"{res['peak_rss_mb']:.0f} MB"
+                    f"+{res['extra_rss_mb']:.0f} MB"
                 )
                 progress.update()
     progress.close()
     common.RESULTS.mkdir(exist_ok=True)
     suffix = "_quick" if args.quick else ""
-    out = common.RESULTS / f"performance{suffix}.csv"
+    out = common.RESULTS / (args.output or f"performance{suffix}.csv")
     pd.DataFrame(rows).to_csv(out, index=False)
     print("wrote", out)