faster-diffBloch

faster-diffbloch provides drop-in Apple Silicon Metal GPU and optimized CPU acceleration for diffBloch differentiable electron crystallography structure refinement.

Operating System & Platform Support

Runtime platform guards detect the host operating system. If device="gpu" is selected on Linux or Windows, faster-diffbloch logs an informative warning and automatically executes on the optimized CPU backend.

Installation

pip install faster-diffbloch

Quickstart

Use the drop-in CLI command diffbloch-fast anywhere you run diffbloch:

# Run inference on all experimental rotations with Metal GPU
diffbloch-fast infer examples/Colmey_et_al_2026/data/quartz-no-abs

# Run gradient refinement
diffbloch-fast refine examples/Colmey_et_al_2026/data/quartz-no-abs

Or enable acceleration in Python scripts with a single function call:

import faster_diffbloch

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

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

# Import standard diffBloch modules
from diffBloch.core.solver import propagate
# Calls now execute via zero-copy accelerated kernels

1. Metal GPU vs PyTorch MPS Fallback

PyTorch lacks a native GPU kernel for aten::linalg_matrix_exp on Apple Silicon MPS. When running on MPS, PyTorch triggers an internal fallback:

1. Copies matrix tensors from GPU memory to CPU host memory.
2. Executes single-threaded CPU ATen matrix exponential routines.
3. Copies result tensors back to GPU memory.

faster-diffBloch implements native Metal compute shaders on Apple Silicon unified memory (MTLResourceStorageModeShared). The entire degree-18 Taylor scaling-and-squaring sequence and block-triangular adjoint passes execute in a single GPU command buffer without host copies.

2. Blocked-Pair Adjoint vs 2N×2N Dense Embedding

The original diffBloch documentation (Devices and Scaling) reports that matrix_exp backward passes cost ~6x more than forward passes on large structures.

PyTorch computes matrix exponential adjoints by embedding the $N \times N$ system into a $2N \times 2N$ block triangular matrix:

exp([[M^H, Ebar], [0, M^H]])

Because matrix multiplication complexity scales as $\mathcal{O}((2N)^3) = 8 \mathcal{O}(N^3)$, every squaring step costs 8 matrix multiplications.

faster-diffBloch computes the pullback directly in the closed pair algebra:

(Y1, L1) * (Y2, L2) = (Y1 * Y2, Y1 * L2 + L1 * Y2)

This requires only 3 matrix multiplications per squaring step, reducing backward computational cost by 2.67x and lowering the backward/forward ratio from 6.14x–7.98x down to 2.72x.

3. Performance Benchmarks

Measured on Apple Silicon M4 Max (minimum over 5 trials):

4. Experimental Dataset Accuracy

Evaluated across all 99 experimental crystal diffraction patterns on alpha-quartz (Colmey_et_al_2026/data/quartz-no-abs):

Passes all 738 unit tests in the diffBloch test suite.