faster-diffbloch 0.1.6__tar.gz → 0.1.7__tar.gz

{faster_diffbloch-0.1.6 → faster_diffbloch-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: faster-diffbloch
-Version: 0.1.6
+Version: 0.1.7
 Summary: Drop-in Metal GPU (macOS) and optimized CPU (macOS/Linux) acceleration for diffBloch
 Project-URL: Homepage, https://godofecht.github.io/diffFlow/
 Project-URL: Documentation, https://godofecht.github.io/diffFlow/
@@ -35,7 +35,7 @@ Description-Content-Type: text/markdown
 
 Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) 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.
+**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.
 
 | Package | Documentation | Source repository | Original project |
 | :--- | :--- | :--- | :--- |
@@ -84,34 +84,31 @@ of three runs each.
 
 | Beams | PyTorch | CPU path | Metal path | CPU | Metal |
 | :---: | :---: | :---: | :---: | :---: | :---: |
-| 31 | 1.52 ms | 1.42 ms | 1.95 ms | 1.07x | 0.78x |
-| 61 | 2.11 ms | 2.07 ms | 2.20 ms | 1.02x | 0.96x |
-| 91 | 3.87 ms | 3.61 ms | 3.70 ms | 1.07x | 1.05x |
-| 163 | 10.46 ms | 8.60 ms | 6.27 ms | 1.22x | 1.67x |
-| 579 | 123.30 ms | 89.17 ms | 47.63 ms | 1.38x | **2.59x** |
-
-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.
+| 31 | 1.47 ms | 0.72 ms | 1.07 ms | 2.03x | 1.37x |
+| 61 | 2.00 ms | 1.33 ms | 1.33 ms | 1.50x | 1.50x |
+| 91 | 3.68 ms | 2.24 ms | 1.67 ms | 1.65x | 2.21x |
+| 163 | 10.34 ms | 5.78 ms | 3.05 ms | 1.79x | **3.39x** |
+| 579 | 117.44 ms | 81.55 ms | 39.87 ms | 1.44x | 2.95x |
+
+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 avoids the framework
-overhead and goes further. At 579 beams, forward plus backward, minimum of five
-runs:
+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:
 
 | Implementation | Forward + Backward |
 | :--- | :---: |
@@ -120,7 +117,8 @@ runs:
 | Flow port, C backend | 82.84 ms |
 | Flow port, MLIR backend | 79.74 ms |
 
-Running the port directly is how those are reached. See the
+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](https://godofecht.github.io/diffFlow/).
 
 These figures come from one workload on one machine. Other crystals, hardware

{faster_diffbloch-0.1.6 → faster_diffbloch-0.1.7}/README.md

@@ -2,7 +2,7 @@
 
 Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) 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.
+**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.
 
 | Package | Documentation | Source repository | Original project |
 | :--- | :--- | :--- | :--- |
@@ -51,34 +51,31 @@ of three runs each.
 
 | Beams | PyTorch | CPU path | Metal path | CPU | Metal |
 | :---: | :---: | :---: | :---: | :---: | :---: |
-| 31 | 1.52 ms | 1.42 ms | 1.95 ms | 1.07x | 0.78x |
-| 61 | 2.11 ms | 2.07 ms | 2.20 ms | 1.02x | 0.96x |
-| 91 | 3.87 ms | 3.61 ms | 3.70 ms | 1.07x | 1.05x |
-| 163 | 10.46 ms | 8.60 ms | 6.27 ms | 1.22x | 1.67x |
-| 579 | 123.30 ms | 89.17 ms | 47.63 ms | 1.38x | **2.59x** |
-
-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.
+| 31 | 1.47 ms | 0.72 ms | 1.07 ms | 2.03x | 1.37x |
+| 61 | 2.00 ms | 1.33 ms | 1.33 ms | 1.50x | 1.50x |
+| 91 | 3.68 ms | 2.24 ms | 1.67 ms | 1.65x | 2.21x |
+| 163 | 10.34 ms | 5.78 ms | 3.05 ms | 1.79x | **3.39x** |
+| 579 | 117.44 ms | 81.55 ms | 39.87 ms | 1.44x | 2.95x |
+
+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 avoids the framework
-overhead and goes further. At 579 beams, forward plus backward, minimum of five
-runs:
+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:
 
 | Implementation | Forward + Backward |
 | :--- | :---: |
@@ -87,7 +84,8 @@ runs:
 | Flow port, C backend | 82.84 ms |
 | Flow port, MLIR backend | 79.74 ms |
 
-Running the port directly is how those are reached. See the
+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](https://godofecht.github.io/diffFlow/).
 
 These figures come from one workload on one machine. Other crystals, hardware

{faster_diffbloch-0.1.6 → faster_diffbloch-0.1.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "faster-diffbloch"
-version = "0.1.6"
+version = "0.1.7"
 description = "Drop-in Metal GPU (macOS) and optimized CPU (macOS/Linux) acceleration for diffBloch"
 readme = "README.md"
 requires-python = ">=3.10"

{faster_diffbloch-0.1.6 → faster_diffbloch-0.1.7}/src/faster_diffbloch/__init__.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from .backend import enable, disable, matrix_exp, matrix_exp_backward, faster_propagate
 
-__version__ = "0.1.6"
+__version__ = "0.1.7"
 __all__ = [
     "enable",
     "disable",

{faster_diffbloch-0.1.6 → faster_diffbloch-0.1.7}/src/faster_diffbloch/backend.py

@@ -200,6 +200,30 @@ def enable(device: Literal["cpu", "gpu"] = "gpu") -> None:
     except ImportError:
         pass
 
+    # Structure factors are the larger cost below roughly 163 beams, where the
+    # matrix exponential is small. Both call sites import the name at module
+    # scope, so each module's own attribute has to be replaced.
+    try:
+        from .scattering import faster_structure_factors, remember_original
+
+        scattering = importlib.import_module("diffBloch.core.scattering")
+        remember_original(scattering.structure_factors)
+        for mod_name in (
+            "diffBloch.core.scattering",
+            "diffBloch.core",
+            "diffBloch.engine.forward",
+            "diffBloch.preprocess.experiment",
+        ):
+            try:
+                mod = importlib.import_module(mod_name)
+                if hasattr(mod, "structure_factors"):
+                    _remember(mod, "structure_factors")
+                    setattr(mod, "structure_factors", faster_structure_factors)
+            except Exception:
+                pass
+    except ImportError:
+        pass
+
 
 def _remember(module, name: str) -> None:
     """Record a module attribute the first time it is replaced.
@@ -231,3 +255,9 @@ def disable() -> None:
         else:
             setattr(module, attr, original)
     _ORIGINALS.clear()
+    try:
+        from .scattering import forget_original
+
+        forget_original()
+    except ImportError:
+        pass

faster_diffbloch-0.1.7/src/faster_diffbloch/scattering.py

@@ -0,0 +1,271 @@
+"""Structure factors through the native path.
+
+The native library already carries `native_parallel_structure_factors` and its
+backward, but nothing called them, so `enable()` only ever replaced the matrix
+exponential. That is the wrong half of the work at small beam counts: at 31
+beams the exponential is 15% of the forward pass and the structure factors are
+80%, which is why the accelerated package barely beat PyTorch there.
+
+This wires the native pair in behind a `torch.autograd.Function`.
+
+Scope. The native kernel covers the elastic case in float64 on the CPU. Anything
+else (absorption, a different dtype, a non-CPU tensor, a cutoff mode other than
+`hard`) falls back to diffBloch's own implementation, so behaviour is unchanged
+where the fast path does not apply.
+
+Two details of diffBloch's definition have to be reproduced exactly. The cutoff
+window multiplies every atom's contribution at a given `g`, so it factors out of
+the sum over atoms and can be applied to the result. The `zero_threshold` snap
+happens before the division by the cell volume, and the native kernel divides
+inside, so the comparison is made against the rescaled value.
+"""
+
+from __future__ import annotations
+
+import ctypes
+from typing import Optional
+
+import numpy as np
+import torch
+
+from .backend import get_library
+
+
+class NativeCase(ctypes.Structure):
+    _fields_ = [
+        ("n_atoms", ctypes.c_int32),
+        ("n_grid", ctypes.c_int32),
+        ("absorption", ctypes.c_int32),
+        ("volume", ctypes.c_double),
+        ("abs_c_over_v", ctypes.c_double),
+        ("positions", ctypes.POINTER(ctypes.c_double)),
+        ("occupancies", ctypes.POINTER(ctypes.c_double)),
+        ("uij", ctypes.POINTER(ctypes.c_double)),
+        ("lobato_a", ctypes.POINTER(ctypes.c_double)),
+        ("lobato_b", ctypes.POINTER(ctypes.c_double)),
+        ("grid_hkl", ctypes.POINTER(ctypes.c_int32)),
+        ("grid_g", ctypes.POINTER(ctypes.c_double)),
+        ("abs_db_du", ctypes.POINTER(ctypes.c_double)),
+        ("abs_knot_lo", ctypes.POINTER(ctypes.c_double)),
+        ("abs_knot_width", ctypes.POINTER(ctypes.c_double)),
+        ("abs_y0", ctypes.POINTER(ctypes.c_double)),
+        ("abs_y1", ctypes.POINTER(ctypes.c_double)),
+        ("abs_d0", ctypes.POINTER(ctypes.c_double)),
+        ("abs_d1", ctypes.POINTER(ctypes.c_double)),
+    ]
+
+
+class PhaseCache(ctypes.Structure):
+    _fields_ = [
+        (name, ctypes.POINTER(ctypes.c_double))
+        for name in ("amplitude", "cosine", "sine", "imag_amplitude", "dimag_db")
+    ]
+
+
+_VOID = ctypes.POINTER(ctypes.c_double * 0)
+
+
+def _f64(a) -> np.ndarray:
+    return np.ascontiguousarray(a, dtype=np.float64).reshape(-1)
+
+
+def _ptr(a: np.ndarray):
+    return a.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
+
+
+_LOBATO: dict[int, tuple[np.ndarray, np.ndarray]] = {}
+
+
+def _lobato_for(numbers: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """The five Lobato a/b coefficients per atom, laid out as the kernel reads them."""
+    if not _LOBATO:
+        from diffBloch.core.scattering import _lobato_table
+
+        for z, (a, b) in _lobato_table().items():
+            _LOBATO[int(z)] = (np.asarray(a, dtype=np.float64),
+                               np.asarray(b, dtype=np.float64))
+    a_out = np.empty((numbers.shape[0], 5), dtype=np.float64)
+    b_out = np.empty((numbers.shape[0], 5), dtype=np.float64)
+    for i, z in enumerate(numbers.tolist()):
+        a, b = _LOBATO[int(z)]
+        a_out[i] = a
+        b_out[i] = b
+    return a_out, b_out
+
+
+class _Native:
+    """Holds every buffer the kernel points at, so nothing is collected mid-call."""
+
+    def __init__(self, positions, numbers, occupancies, uij_star, hkl, g, volume):
+        self.n_atoms = int(positions.shape[0])
+        self.n_grid = int(hkl.shape[0])
+        self.pos = _f64(positions)
+        self.occ = _f64(occupancies)
+        self.uij = _f64(uij_star)
+        la, lb = _lobato_for(np.asarray(numbers))
+        self.la, self.lb = _f64(la), _f64(lb)
+        self.hkl = np.ascontiguousarray(hkl, dtype=np.int32).reshape(-1)
+        self.g = _f64(g)
+
+        c = NativeCase()
+        c.n_atoms, c.n_grid, c.absorption = self.n_atoms, self.n_grid, 0
+        c.volume, c.abs_c_over_v = float(volume), 0.0
+        c.positions, c.occupancies, c.uij = _ptr(self.pos), _ptr(self.occ), _ptr(self.uij)
+        c.lobato_a, c.lobato_b = _ptr(self.la), _ptr(self.lb)
+        c.grid_hkl = self.hkl.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
+        c.grid_g = _ptr(self.g)
+        self.case = c
+
+        n = self.n_grid * self.n_atoms
+        self.cache_arrays = {}
+        cache = PhaseCache()
+        for name in ("amplitude", "cosine", "sine", "imag_amplitude", "dimag_db"):
+            buf = np.zeros(n, dtype=np.float64)
+            self.cache_arrays[name] = buf
+            setattr(cache, name, _ptr(buf))
+        self.cache = cache
+
+    def forward(self) -> np.ndarray:
+        lib = get_library()
+        fgb = np.zeros(self.n_grid, dtype=np.complex128)
+        lib.native_parallel_structure_factors(
+            ctypes.byref(self.case),
+            fgb.ctypes.data_as(_VOID),
+            ctypes.byref(self.cache),
+        )
+        return fgb
+
+    def backward(self, fbar: np.ndarray):
+        lib = get_library()
+        gp = np.zeros(self.n_atoms * 3, dtype=np.float64)
+        go = np.zeros(self.n_atoms, dtype=np.float64)
+        gu = np.zeros(self.n_atoms * 9, dtype=np.float64)
+        fbar_c = np.ascontiguousarray(fbar, dtype=np.complex128)
+        lib.native_parallel_structure_factors_backward(
+            ctypes.byref(self.case),
+            ctypes.byref(self.cache),
+            fbar_c.ctypes.data_as(_VOID),
+            _ptr(gp), _ptr(go), _ptr(gu),
+        )
+        return gp.reshape(self.n_atoms, 3), go, gu.reshape(self.n_atoms, 3, 3)
+
+
+class FasterStructureFactors(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, positions, occupancies, uij_star, numbers, hkl, g,
+                volume, window, zero_threshold):
+        native = _Native(
+            positions.detach().cpu().numpy(), numbers, occupancies.detach().cpu().numpy(),
+            uij_star.detach().cpu().numpy(), hkl, g, volume,
+        )
+        raw = native.forward()
+        # diffBloch snaps before dividing by the volume; the kernel divides inside.
+        unmasked = raw * volume * window
+        keep_re = np.abs(unmasked.real) >= zero_threshold
+        keep_im = np.abs(unmasked.imag) >= zero_threshold
+        out = (np.where(keep_re, unmasked.real, 0.0)
+               + 1j * np.where(keep_im, unmasked.imag, 0.0)) / volume
+
+        ctx.native = native
+        ctx.window = window
+        ctx.keep = (keep_re, keep_im)
+        ctx.shapes = (positions.shape, occupancies.shape, uij_star.shape)
+        ctx.meta = (positions.dtype, positions.device)
+        return torch.from_numpy(out).to(device=positions.device)
+
+    @staticmethod
+    @torch.autograd.function.once_differentiable
+    def backward(ctx, grad_output):
+        keep_re, keep_im = ctx.keep
+        gout = grad_output.detach().cpu().numpy()
+        # Undo the snap and the per-g window, so the cotangent matches what the
+        # kernel differentiates: the plain sum over atoms, divided by the volume.
+        fbar = (np.where(keep_re, gout.real, 0.0)
+                + 1j * np.where(keep_im, gout.imag, 0.0)) * ctx.window
+        gp, go, gu = ctx.native.backward(fbar)
+        dtype, device = ctx.meta
+        p_shape, o_shape, u_shape = ctx.shapes
+        return (
+            torch.from_numpy(gp).to(dtype=dtype, device=device).reshape(p_shape),
+            torch.from_numpy(go).to(dtype=dtype, device=device).reshape(o_shape),
+            torch.from_numpy(gu).to(dtype=dtype, device=device).reshape(u_shape),
+            None, None, None, None, None, None,
+        )
+
+
+def _shapes_ok(positions, numbers, occupancies, uij_star, cell_volume) -> bool:
+    """Whether the inputs are the shapes diffBloch accepts.
+
+    A malformed call is handed to diffBloch rather than rejected here, so the
+    error it raises is its own, with its own message. Validating separately
+    would mean keeping two copies of the same rules in step.
+    """
+    if positions.ndim != 2 or positions.shape[1] != 3:
+        return False
+    n_atoms = positions.shape[0]
+    if tuple(numbers.shape) != (n_atoms,) or tuple(occupancies.shape) != (n_atoms,):
+        return False
+    if tuple(uij_star.shape) != (n_atoms, 3, 3):
+        return False
+    return cell_volume > 0.0
+
+
+def _usable(positions, numbers, occupancies, uij_star, cell_volume,
+            absorption, cutoff) -> bool:
+    if absorption is not None and getattr(absorption, "enabled", False):
+        return False
+    if cutoff != "hard":
+        return False
+    for t in (positions, occupancies, uij_star):
+        if t.dtype != torch.float64 or t.device.type != "cpu":
+            return False
+    if not _shapes_ok(positions, numbers, occupancies, uij_star, cell_volume):
+        return False
+    return get_library() is not None
+
+
+def faster_structure_factors(
+    positions, numbers, occupancies, uij_star, hkl, reciprocal_basis, cell_volume,
+    *, g_max, cutoff="hard", zero_threshold=1e-12, absorption=None, energy=None,
+):
+    """diffBloch's `structure_factors`, over the native kernel where it applies."""
+    from diffBloch.core.scattering import (
+        _g_vector_lengths,
+        structure_factors as diffbloch_structure_factors,
+    )
+
+    original = _ORIGINAL[0] or diffbloch_structure_factors
+    if absorption is None:
+        from diffBloch.specs import NO_ABSORPTION
+
+        absorption = NO_ABSORPTION
+
+    if not _usable(positions, numbers, occupancies, uij_star, cell_volume,
+                   absorption, cutoff):
+        return original(
+            positions, numbers, occupancies, uij_star, hkl, reciprocal_basis,
+            cell_volume, g_max=g_max, cutoff=cutoff, zero_threshold=zero_threshold,
+            absorption=absorption, energy=energy,
+        )
+
+    g = _g_vector_lengths(hkl, reciprocal_basis)
+    window = (g <= g_max).to(torch.float64).detach().cpu().numpy()
+    return FasterStructureFactors.apply(
+        positions, occupancies, uij_star,
+        np.asarray(numbers.detach().cpu().numpy()),
+        hkl.detach().cpu().numpy(),
+        g.detach().cpu().numpy(),
+        float(cell_volume), window, float(zero_threshold),
+    )
+
+
+_ORIGINAL: list[Optional[object]] = [None]
+
+
+def remember_original(fn) -> None:
+    if _ORIGINAL[0] is None:
+        _ORIGINAL[0] = fn
+
+
+def forget_original() -> None:
+    _ORIGINAL[0] = None