A PyTorch-based crystallographic Refinement library

TorchRef is a crystallographic refinement package built entirely on PyTorch. Autograd and GPU acceleration make it composable with machine-learning workflows and cheap to extend with new targets.

Scope. TorchRef is mainly a library/framework to build and experiment with. It is not intended to replace mainline refinement programs for standard problems.

Benchmark

Refinement of Phaser-placed AlphaFold models against experimental data, on a conserved set of 723 PDB structures (1.40–3.00 Å). All engines start from the same placed models and are scored by one common validator (PHENIX).

(A) R-work and R-free per engine. Median R-free 0.3189 (TorchRef), 0.3164 (PHENIX), 0.3136 (REFMAC5); paired median difference, TorchRef minus reference, +0.0008 against PHENIX and +0.0047 against REFMAC. (B) Geometry RMS Z against REFMAC restraints, ideal 1.0: bond 0.58 (REFMAC 0.58, PHENIX 0.83), main-chain B 1.55 (REFMAC 0.90, PHENIX 1.03). (C) Wall-clock runtime per structure on 4 CPU cores: median 0.77 min (REFMAC 0.25, PHENIX 3.09). (D) Fraction of the total R-free improvement reached per macrocycle.

Key Features

• Native PyTorch Integration: Built on PyTorch's nn.Module architecture, so TorchRef composes with PyTorch models, optimizers, and devices.

• Automatic Differentiation: No hand-written gradients. Define a new refinement target's forward pass and PyTorch supplies the derivatives.

• Modular Architecture: Custom targets, restraints, and optimizers plug in without modifying core code.

• GPU Acceleration: CUDA for structure factors, scaling, and optimization. Apple Silicon works through PyTorch's MPS backend — unsupported ops fall back to CPU automatically via PYTORCH_ENABLE_MPS_FALLBACK=1, which TorchRef sets on import.

• FFT-based Structure Factors: F_calc via FFT, so large unit cells stay tractable.

Getting Started

Notebook	Description
	Quickstart — MTZ + PDB to refined structure, refined MTZ and CCP4 map; selection- and parameter-type-based refinement
	Structure factors — one-liner, FFT class, and manual voxel pipeline; standalone scaling; autograd
	Targets and weighting — standard targets, target-offset weighting, X-ray mode comparison, custom targets, driving an optimizer from a LossState

Installation

pip install torchref

For development:

git clone --filter=blob:none --sparse https://github.com/HatPdotS/TorchRef.git
cd TorchRef
git sparse-checkout set torchref tests
pip install -e ".[dev]"

This fetches ~40 MB instead of ~436 MB; most of the repository is paper/ history. Files outside the checkout are fetched on demand, so add paths later with git sparse-checkout add paper, or git sparse-checkout disable for all of it. Requires Git ≥ 2.27.

Dependencies

Python ≥ 3.10, PyTorch ≥ 2.4, NumPy ≥ 2.0, Pandas ≥ 2.0, SciPy ≥ 1.10, Gemmi ≥ 0.5, reciprocalspaceship ≥ 0.9.18, Numba ≥ 0.59, Matplotlib ≥ 3.7. pyproject.toml carries the authoritative pinned ranges; upper bounds are set one minor version above the tested maximum, so a newer dependency will refuse to install rather than fail at runtime.

Testing

pytest tests/                      # all tests
pytest tests/ --cov=torchref       # with coverage
pytest tests/unit/                 # fast unit tests only

Slow tests need --run-slow. Accelerator tests are not opt-in: they run wherever CUDA or MPS is available and are skipped when it is not.

Contributing

Contributions are welcome. Please use NumPy docstring style, add tests for new functionality, and make sure the suite passes before submitting.

Questions, feature ideas and usage discussion are welcome on the TorchRef category of the rs-station forum.

License

MIT. See LICENSE.