faster-diffBloch

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

Documentation and comparison benchmarks: https://godofecht.github.io/diffFlow/

Original diffBloch project: https://diffbloch.com

Why faster-diffBloch?

1. Native Metal GPU Execution: PyTorch MPS lacks a native GPU kernel for aten::linalg_matrix_exp, which causes PyTorch to fall back to CPU execution with host-device memory transfers. faster-diffBloch executes matrix exponentials directly on Apple Silicon Metal with zero-copy unified memory.

2. Blocked-Pair Adjoint Formulation: Standard matrix exponential autograd embeds the operator into a $2N \times 2N$ block matrix, costing $8 \times N^3$ FLOPs. faster-diffBloch evaluates the pullback in the block-triangular pair algebra $(Y_a Y_b, Y_a L_b + L_a Y_b)$, reducing the work to $3 \times N^3$ FLOPs (2.67x fewer products).

3. Bit-for-Bit Validation: Passes all 738 unit tests in diffBloch and reproduces the experimental 99-rotation quartz dataset ($R_{\text{obs}} = 0.0486$).

Performance

Forward and backward timing comparison on Apple Silicon (M4 Max) at $N=579$ beams (CsPbBr3 scale):

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
faster_diffbloch.enable(device="gpu")

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

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