Fixed beta estimation instability on gpu

README.md

@@ -12,6 +12,21 @@
 
 TorchRef is a crystallographic refinement package built entirely on PyTorch. By leveraging PyTorch's automatic differentiation and GPU acceleration, TorchRef enables seamless integration with machine learning workflows and provides a flexible, extensible framework for crystallographic structure refinement.
 
+> **Scope 
+TorchRef is a mainly a library/framework to build and experiment with. It is not intended to replace mainline refinement programs for standard problems. 
+
+# Benchmark
+
+![TorchRef AlphaFold-start refinement benchmark](paper/figure2_alphafold_start/figures/figure_af_benchmark.png)
+
+*Refinement of Phaser-placed AlphaFold models against experimental data, benchmarked on a conserved set of ~720 PDB structures (1.4–3.0 Å) with all engines starting from the same placed models and scored by a single common validator.*
+
+- **(A) R-factors (PHENIX-validated).** Starting from the AlphaFold prediction (green), TorchRef (red) drives R-work/R-free down to essentially the same cluster as REFMAC (purple) and PHENIX (blue). Median R-free is 0.3167 (TorchRef) vs 0.3166 (PHENIX) and 0.3161 (REFMAC5)
+- **(B) Geometry (RMSZ vs REFMAC restraints).** Bond, angle, chiral and main-chain B-factor RMS Z-scores. TorchRef produces valid, physically reasonable geometry; its restraints run slightly looser than PHENIX/REFMAC (bond RMSZ ≈ 1.3)
+- **(C) Wall-clock runtime.** Median runtime per structure (4 CPU cores). TorchRef (1.65 min) sits between REFMAC (0.53 min) and PHENIX (4.63 min) — ~2.8× faster than PHENIX, ~3× slower than REFMAC.
+- **(D) Convergence speed (normalized).** Fraction of the total R-free improvement achieved per macrocycle. The different programs show differing convergence behavior. 
+
+
 # Key Features
 
 - **Native PyTorch Integration**: Built on PyTorch's `nn.Module` architecture, TorchRef integrates naturally with the PyTorch ecosystem, including machine learning models, optimizers, and GPU acceleration.
@@ -55,11 +70,11 @@ pip install -e ".[dev]"
 ## Dependencies
 
 - Python ≥ 3.10
-- PyTorch ≥ 2.40
+- PyTorch ≥ 2.4
 - NumPy ≥ 2.0
 - Gemmi ≥ 0.5
-- reciprocalspaceship ≥ 0.9
-- SciPy ≥ 1.7
+- reciprocalspaceship ≥ 0.9.18
+- SciPy ≥ 1.10
 
 ## Testing
 

docs/changelog.rst

@@ -1,6 +1,12 @@
 Changelog
 =========
 
+Version 0.6.1
+-------------
+- Fixed bug in beta estimation that caused instability in GPU refinement
+- Fixed reporting bug in collection scaler were rfactor reporting would ignore masks
+
+
 Version 0.6.0
 -------------
 - Switched to per-atom cutoff radii for the electron-density sampling and added a global sigma cutoff

paper/extended_figures/exF2/collect_exF2_data.py

@@ -27,9 +27,9 @@
 BASE = Path(__file__).resolve().parent
 AF_ROOT = BASE.parent.parent / "figure2_alphafold_start"
 CROSSSCORE = AF_ROOT / "runs" / "metrics" / "fig_crossscore.csv"
-TORCHREF_ARM = AF_ROOT / "runs" / "torchref_g0p2_a0p02"
+TORCHREF_ARM = AF_ROOT / "runs" / "torchref"
 SCORER = "phenix"                       # common independent scorer (main-figure scorer)
-TORCHREF_ENGINE = "torchref_g0p2_a0p02"
+TORCHREF_ENGINE = "torchref"
 PHENIX_ENGINE = "phenix"
 REFMAC_ENGINE = "refmac"
 OUT_CSV = BASE / "data" / "exF2_rfactor_by_resolution.csv"

paper/extended_figures/exF3/plot_exF3.py

@@ -15,7 +15,7 @@
 scorers' systematic differences (PHENIX lowest, then REFMAC, then TorchRef — the
 most conservative); these cancel within a scorer. PHENIX is the main-figure scorer.
 
-TorchRef here is the locked-default arm (torchref_g0p2_a0p02; xray 1 / geometry 0.2
+TorchRef here is the canonical arm (torchref; xray 1 / geometry 0.2
 / adp 0.02). Reads figure2_alphafold_start/runs/metrics/fig_crossscore.csv
 (analysis/aggregate_crossscore.py).
 
@@ -42,7 +42,7 @@
 
 # refined-model engines (exclude the AlphaFold prediction to keep axes focused).
 # TorchRef = the locked-default arm (xray 1 / geometry 0.2 / adp 0.02).
-TORCHREF_ENGINE = "torchref_g0p2_a0p02"
+TORCHREF_ENGINE = "torchref"
 ENGINES = [
     ("refmac", "Refmac", "#762a83"),
     ("phenix", "PHENIX", "#2166ac"),

paper/figure2_alphafold_start/analysis/aggregate_weight_grid.py

@@ -9,6 +9,7 @@
 Usage:
     ./.dev/bin/python analysis/aggregate_weight_grid.py
 """
+import argparse
 import csv
 import sys
 from pathlib import Path
@@ -17,23 +18,34 @@
 import run_af_pipeline as P  # noqa: E402
 from aggregate_figure_metrics import RE_RWORK, RE_RFREE, parse_geometry  # noqa: E402
 
-MANIFEST = P.RUNS / "metrics" / "wgrid_manifest.csv"
-OUT = P.RUNS / "metrics" / "weight_grid.csv"
-
 
 def ratio(num, den):
     return (num / den) if (num is not None and den) else None
 
 
 def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--tag", default="",
+                    help="Ablation-grid namespace (matches submit_weight_grid "
+                         "--tag), e.g. 'nosimu'. Reads wgrid_<tag>_manifest.csv "
+                         "and writes weight_grid_<tag>.csv.")
+    ap.add_argument("--out-root", default=None,
+                    help="Base directory the grid was written to (matches "
+                         "submit_weight_grid --out-root). Default: figure2 runs/.")
+    args = ap.parse_args()
+    out_root = Path(args.out_root).resolve() if args.out_root else P.RUNS
+    suffix = f"_{args.tag}" if args.tag else ""
+    MANIFEST = out_root / "metrics" / f"wgrid{suffix}_manifest.csv"
+    OUT = out_root / "metrics" / f"weight_grid{suffix}.csv"
+
     if not MANIFEST.exists():
         sys.exit(f"missing {MANIFEST}; run submit_weight_grid.py first")
     manifest = list(csv.DictReader(open(MANIFEST)))
 
     rows = []
     for m in manifest:
         arm = m["arm"]
-        arm_dir = P.RUNS / arm
+        arm_dir = out_root / arm
         if not arm_dir.is_dir():
             continue
         for code_dir in sorted(p for p in arm_dir.iterdir() if p.is_dir()):

paper/figure2_alphafold_start/analysis/build_torchref_score_worklist.py

@@ -28,7 +28,7 @@
 MODELS = {
     "refmac": ("refmac", "refmac/{code}/refined.pdb"),
     "phenix": ("phenix_norb", "phenix_norb/{code}/{code}_refined_001.pdb"),
-    "torchref": ("torchref_g0p2_a0p02", "torchref_g0p2_a0p02/{code}/refined.pdb"),
+    "torchref": ("torchref", "torchref/{code}/refined.pdb"),
     "prediction": ("af_initial", "../placed/{code}_af.pdb"),
 }