A PyTorch-based scaling and merging library for serial crystallography (SX)

In serial crystallography every reflection is seen once per crystal, with an unknown partiality, incident intensity and crystal volume. Established programs scale each crystal with a few parameters and average. TorchSX instead trains one neural network per dataset that predicts, for every observation, a scale factor and a scatter distribution from that observation's crystal- and image-level context. Merging is a posterior integral rather than an average: each merged amplitude and its uncertainty are moments of the posterior of the reflection's intensity. The network is trained on the marginal likelihood of the data.

TorchSX starts from integrated CrystFEL .stream files and writes merged MTZs. We are working on integration with dials and deep integration of Probixi.

TorchSX is under active development, if you find something behaving unexpectedly please create an issue on GitHub.

Refinement and difference maps are done with TorchRef, or any other program that reads amplitudes.

Install

pip install torchsx

Python ≥ 3.10. This brings TorchRef 0.7, torch, triton, gemmi and torchsx-native, the compiled .stream reader for your platform. TorchSX is intended to be run on a GPU, while performance and memory requirements should be quite low, it relies on triton kernels for core algorithms. We have not investigated CPU runtime or even functionality.

From a checkout, build the native package first (needs Rust, see native/README.md): pip install ./native, then pip install -e ..

Usage

torchsx.train indexed.stream --spacegroup 'P 63' --dmin 2.0 -o run/

This trains the network (400 updates by default) and merges the data with it. run/ then contains:

run/
├── merged.mtz              merged amplitudes from all crystals
├── merged_half1.mtz        the same from one half of the crystals
├── merged_half2.mtz        ... and from the other half
├── merged.metrics.json     CC1/2 and R_split between the two halves
├── best.pt                 the network at the lowest validation loss; it made the merges
├── last.pt                 the network after the last update
├── prepared.pt             the parsed observations, reused by later runs on the same input
├── run.json                the resolved settings
├── training_summary.json   updates, best step and run time
└── metrics.jsonl           training log: one line per update and per validation

Each MTZ holds:

There are no intensity columns. F-obs is already the posterior estimate of the amplitude, so it must not go through French–Wilson again; refinement programs use it as it is. With --no-merge-friedel the mates stay apart and the files carry F-obs(+), SIGF-obs(+), F-obs(-) and SIGF-obs(-). --free-from earlier.mtz takes the free set from an existing file.

The two halves split the crystals, not the training: both are merged with the same network. For an independent comparison, train two runs on disjoint sets of patterns.

Light and dark

torchsx.train --split dark=dark.stream --split light=light.stream \
  --spacegroup 'P 63' --dmin 2.0 -o run/

One network is trained on both states, and each state is merged on its own: merged_dark.mtz and merged_light.mtz, each with its two halves. A single stream with a CrystFEL custom-split file works too (--split-file).

Compared with Partialator

On the bacteriorhodopsin (bR) data of Weinert et al.¹:

R_free of models refined from one starting model against dark merges of increasing size. Each size was merged as two independent halves: points show their mean, bars their range. R_free is computed on the same reflections for both programs at each size; refinement used TorchRef with one protocol for all merges.

Light-minus-dark difference maps from two independent halves of the patterns (7,800 per state and half), each merged on its own and phased with the dark model. Shown: the end of helix E, the EF loop and helix F of the dark model (grey) and of the open state 6RPH¹ (orange). The number under each program is the correlation of its two maps over helix F.

Other commands

• torchsx.merge run/best.pt more.stream -o merged.mtz merges new or additional data with a trained network.
• torchsx.merge-stats A/merged.mtz B/merged.mtz -o stats/ compares merges: completeness, CC1/2, R_split and, for light/dark pairs, the information in ΔF (docs/merge_stats.md).
• torchsx.observation-stats indexed.stream --spacegroup 'P 63' -o figures/ plots the raw observations: intensity distributions, skewness, examples (docs/stats.md).

All options, gain and measurement settings, caching and reproducibility are in the command-line reference.

In Python, torchsx.SFXStream.read('indexed.stream', spacegroup='P 63') reads a stream with its geometry and measurement settings.

References

1. Weinert, T. et al. Proton uptake mechanism in bacteriorhodopsin captured by serial synchrotron crystallography. Science 365, 61–65 (2019). doi:10.1126/science.aaw8634
2. Seidel, H.-P., Standfuss, J. & Weinert, T. TorchRef: an open-source PyTorch framework for crystallographic refinement. J. Appl. Cryst., accepted (2026).

License

MIT. See LICENSE.