faster-diffbloch

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

Up to 3.4x faster than PyTorch on the forward-plus-backward pass, 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.

Two kernels carry this. The matrix exponential is O(N^3) and dominates at large beam counts. The structure factors dominate at small ones: at 31 beams they are about 80% of the forward pass while the exponential is 15%. Both now run natively, which is why the speedup holds across the range rather than appearing only at scale.

Accuracy is not traded for any of it. The structure factors run in float64 and agree with diffBloch to 4e-15, their gradients to 1e-14, and the systematic absences land on the same exact zeros. End to end the loss is bit-identical to diffBloch's own.

bench/pkgbench.py reproduces the table.

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 runs on the CPU and avoids the framework overhead. At 579 beams, forward plus backward, minimum of five runs:

The port is the fastest CPU route. This package on Metal is faster still, at 39.87 ms for the same case, because the port has no GPU path. 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