faster-diffbloch

Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for diffBloch electron crystallography structure refinement.

2.59x faster than PyTorch on the forward-plus-backward pass at 579 beams, through the Metal path on Apple Silicon. An optimized CPU path covers macOS and Linux.

Operating System and Platform Support

Runtime platform guards automatically detect your operating system and hardware configuration. When device="gpu" is requested on Linux, faster-diffbloch automatically selects the optimized CPU backend with an informative warning.

Why faster-diffBloch?

faster-diffbloch keeps the diffBloch workflow and public API intact while moving the expensive propagation and gradient work onto an accelerated native backend. You get the same scientific calculation, with a faster path on supported hardware and a safe PyTorch fallback when the native backend is unavailable.

The package is validated against diffBloch's own test suite. With acceleration enabled it passes all 738 diffBloch tests and 55 conformance tests that compare the accelerated results against diffBloch's, field by field, including gradients. The diffBloch run makes 4855 calls into the native matrix exponential, so the accelerated path is exercised rather than skipped.

It also reproduces the published quartz result. Running diffbloch-fast infer over the 99-rotation Colmey et al. 2026 quartz dataset through the Metal path scores every rotation and gives a mean R_obs of 0.0485, against the 0.0486 the example documents.

Performance

Forward-plus-backward timing on an Apple M4 Max, single-threaded PyTorch, over the same diffBloch quartz workload at five beam counts. Both device paths, best of three runs each.

The speedup grows with beam count, and this is the shape of the result rather than noise. The package replaces the matrix exponential, which costs O(N^3), while structure factors and the loss stay in PyTorch. As the system grows the exponential takes a larger share of the runtime, so there is more for the accelerated path to reach.

That has two practical consequences. At 579 beams, CsPbBr3 scale, Metal runs the forward-plus-backward pass in 47.63 ms against PyTorch's 123.30 ms. Below about 91 beams Metal is slower than PyTorch, because the dispatch cost outweighs a small exponential, so enable(device="cpu") is the better choice for small systems.

Forward pass alone at 579 beams: PyTorch 25.90 ms, Metal 15.58 ms, a factor of 1.66. PyTorch MPS has no native kernel for aten::linalg_matrix_exp and falls back to the CPU with host transfers, which is the gap this closes.

The standalone Flow port

The diffFlow repository also holds a standalone port of the whole calculation, compiled from Flow rather than bridged into PyTorch. It avoids the framework overhead and goes further. At 579 beams, forward plus backward, minimum of five runs:

Running the port directly is how those are reached. See the diffFlow documentation.

These figures come from one workload on one machine. Other crystals, hardware configurations and beam counts will differ.

Installation

pip install faster-diffbloch

Usage

1. Drop-in CLI

Use diffbloch-fast or faster-diffbloch anywhere you would use diffbloch:

diffbloch-fast infer examples/Colmey_et_al_2026/data/quartz-no-abs
diffbloch-fast refine examples/Colmey_et_al_2026/data/quartz-no-abs

2. Python API Injection

Enable acceleration inside any existing diffBloch script:

import faster_diffbloch

# Enable Metal GPU acceleration (macOS Apple Silicon)
faster_diffbloch.enable(device="gpu")

# Or CPU acceleration (macOS and Linux)
faster_diffbloch.enable(device="cpu")

# Run standard diffBloch code
import diffBloch
# All propagate and matrix_exp calls now route through faster-diffBloch