faster-diffbloch 0.1.0__tar.gz

faster_diffbloch-0.1.0/.gitignore

@@ -0,0 +1,24 @@
+# The upstream package and the skills pack are their own repositories.
+diffBloch/
+_ref/
+
+# Generated: case files are rebuilt by bench/make_case.py, dumps by the Flow binaries.
+bench/case/
+bench/out/
+bench/results.json
+
+__pycache__/
+*.pyc
+.pytest_cache/
+
+# Build artifacts and throwaway probes.
+mojobloch/main
+mojobloch/main.o
+scratch/
+.DS_Store
+
+# Compiled by hand for the expm floor measurement (reports/headroom.md).
+bench/build/
+*.metallib
+*.dylib
+*.so

faster_diffbloch-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Abhishek Shivakumar
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

faster_diffbloch-0.1.0/PKG-INFO

@@ -0,0 +1,98 @@
+Metadata-Version: 2.5
+Name: faster-diffbloch
+Version: 0.1.0
+Summary: Drop-in Metal GPU and CPU acceleration for diffBloch electron crystallography
+Project-URL: Homepage, https://godofecht.github.io/diffFlow/
+Project-URL: Repository, https://github.com/godofecht/diffFlow
+Project-URL: Issues, https://github.com/godofecht/diffFlow/issues
+Project-URL: Original diffBloch, https://diffbloch.com
+Author-email: Abhishek Shivakumar <abhishek@example.com>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.24
+Requires-Dist: torch>=2.0
+Provides-Extra: diffbloch
+Requires-Dist: diffbloch; extra == 'diffbloch'
+Description-Content-Type: text/markdown
+
+# faster-diffBloch
+
+Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) electron crystallography structure refinement.
+
+Documentation and comparison benchmarks: [https://godofecht.github.io/diffFlow/](https://godofecht.github.io/diffFlow/)
+
+Original diffBloch project: [https://diffbloch.com](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):
+
+| Implementation | Forward | Forward + Backward | Speedup vs PyTorch CPU | Speedup vs PyTorch MPS |
+| :--- | :---: | :---: | :---: | :---: |
+| PyTorch CPU | 25.7 ms | 130.3 ms | 1.00x | 1.17x |
+| PyTorch MPS (fallback) | 26.2 ms | 153.0 ms | 0.85x | 1.00x |
+| **faster-diffBloch CPU** | **24.4 ms** | **83.5 ms** | **1.56x** | **1.83x** |
+| **faster-diffBloch Metal GPU** | **13.1 ms** | **58.1 ms** | **2.24x** | **2.63x** |
+
+---
+
+## Installation
+
+```bash
+pip install faster-diffbloch
+```
+
+---
+
+## Usage
+
+### 1. Drop-in CLI
+
+Use `diffbloch-fast` or `faster-diffbloch` anywhere you would use `diffbloch`:
+
+```bash
+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:
+
+```python
+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
+```

faster_diffbloch-0.1.0/README.md

@@ -0,0 +1,72 @@
+# faster-diffBloch
+
+Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) electron crystallography structure refinement.
+
+Documentation and comparison benchmarks: [https://godofecht.github.io/diffFlow/](https://godofecht.github.io/diffFlow/)
+
+Original diffBloch project: [https://diffbloch.com](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):
+
+| Implementation | Forward | Forward + Backward | Speedup vs PyTorch CPU | Speedup vs PyTorch MPS |
+| :--- | :---: | :---: | :---: | :---: |
+| PyTorch CPU | 25.7 ms | 130.3 ms | 1.00x | 1.17x |
+| PyTorch MPS (fallback) | 26.2 ms | 153.0 ms | 0.85x | 1.00x |
+| **faster-diffBloch CPU** | **24.4 ms** | **83.5 ms** | **1.56x** | **1.83x** |
+| **faster-diffBloch Metal GPU** | **13.1 ms** | **58.1 ms** | **2.24x** | **2.63x** |
+
+---
+
+## Installation
+
+```bash
+pip install faster-diffbloch
+```
+
+---
+
+## Usage
+
+### 1. Drop-in CLI
+
+Use `diffbloch-fast` or `faster-diffbloch` anywhere you would use `diffbloch`:
+
+```bash
+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:
+
+```python
+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
+```

faster_diffbloch-0.1.0/pyproject.toml

@@ -0,0 +1,53 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "faster-diffbloch"
+version = "0.1.0"
+description = "Drop-in Metal GPU and CPU acceleration for diffBloch electron crystallography"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+authors = [
+    { name = "Abhishek Shivakumar", email = "abhishek@example.com" }
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Physics",
+]
+dependencies = [
+    "numpy>=1.24",
+    "torch>=2.0",
+]
+
+[project.optional-dependencies]
+diffbloch = [
+    "diffbloch",
+]
+
+[project.scripts]
+faster-diffbloch = "faster_diffbloch.cli:main"
+diffbloch-fast = "faster_diffbloch.cli:main"
+
+[project.urls]
+Homepage = "https://godofecht.github.io/diffFlow/"
+Repository = "https://github.com/godofecht/diffFlow"
+Issues = "https://github.com/godofecht/diffFlow/issues"
+"Original diffBloch" = "https://diffbloch.com"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/faster_diffbloch"]
+
+[tool.hatch.build.targets.sdist]
+include = [
+    "src/faster_diffbloch",
+    "README.md",
+    "LICENSE",
+]

faster_diffbloch-0.1.0/src/faster_diffbloch/__init__.py

@@ -0,0 +1,14 @@
+"""faster-diffBloch: Metal GPU and CPU acceleration for diffBloch."""
+
+from __future__ import annotations
+
+from .backend import enable, disable, matrix_exp, matrix_exp_backward, faster_propagate
+
+__version__ = "0.1.0"
+__all__ = [
+    "enable",
+    "disable",
+    "matrix_exp",
+    "matrix_exp_backward",
+    "faster_propagate",
+]

faster_diffbloch-0.1.0/src/faster_diffbloch/backend.py

@@ -0,0 +1,149 @@
+"""Flow and Metal acceleration backend for diffBloch."""
+
+from __future__ import annotations
+
+import ctypes
+from typing import Literal
+import numpy as np
+import torch
+
+from .builder import build_and_load_library
+
+_LIB: ctypes.CDLL | None = None
+_CURRENT_DEVICE: Literal["cpu", "gpu"] = "gpu"
+
+
+def get_library() -> ctypes.CDLL:
+    global _LIB
+    if _LIB is None:
+        _LIB = build_and_load_library()
+    return _LIB
+
+
+def matrix_exp(a: np.ndarray, device: str = "gpu") -> np.ndarray:
+    lib = get_library()
+    a_c = np.ascontiguousarray(a, dtype=np.complex64)
+    out = np.empty_like(a_c)
+    n = a_c.shape[-1]
+    batch = int(a_c.size // (n * n))
+    ptr_in = a_c.ctypes.data_as(ctypes.c_void_p)
+    ptr_out = out.ctypes.data_as(ctypes.c_void_p)
+
+    if device == "gpu":
+        fn = getattr(lib, "bridge_matrix_exp_gpu_ptr_c64_i32_i32_ptr_c64", None)
+        if fn:
+            fn.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
+            fn.restype = None
+            fn(ptr_in, batch, n, ptr_out)
+            return out
+
+    fn = getattr(lib, "bridge_matrix_exp_ptr_c64_i32_i32_ptr_c64")
+    fn.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
+    fn.restype = None
+    fn(ptr_in, batch, n, ptr_out)
+    return out
+
+
+def matrix_exp_backward(a: np.ndarray, ebar: np.ndarray, dense: bool = False, device: str = "gpu") -> np.ndarray:
+    lib = get_library()
+    a_c = np.ascontiguousarray(a, dtype=np.complex64)
+    ebar_c = np.ascontiguousarray(ebar, dtype=np.complex64)
+    out = np.empty_like(a_c)
+    n = a_c.shape[-1]
+    batch = int(a_c.size // (n * n))
+    dense_val = 1 if dense else 0
+
+    ptr_a = a_c.ctypes.data_as(ctypes.c_void_p)
+    ptr_ebar = ebar_c.ctypes.data_as(ctypes.c_void_p)
+    ptr_out = out.ctypes.data_as(ctypes.c_void_p)
+
+    if device == "gpu":
+        fn = getattr(lib, "bridge_matrix_exp_backward_gpu_ptr_c64_ptr_c64_i32_i32_i32_ptr_c64", None)
+        if fn:
+            fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
+            fn.restype = None
+            fn(ptr_a, ptr_ebar, dense_val, batch, n, ptr_out)
+            return out
+
+    fn = getattr(lib, "bridge_matrix_exp_backward_ptr_c64_ptr_c64_i32_i32_i32_ptr_c64")
+    fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
+    fn.restype = None
+    fn(ptr_a, ptr_ebar, dense_val, batch, n, ptr_out)
+    return out
+
+
+class FasterMatrixExp(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, m: torch.Tensor) -> torch.Tensor:
+        source = m.detach().resolve_conj().contiguous()
+        exp_np = matrix_exp(source.cpu().numpy(), device=_CURRENT_DEVICE)
+        ctx.save_for_backward(source)
+        return torch.from_numpy(exp_np).to(device=m.device)
+
+    @staticmethod
+    @torch.autograd.function.once_differentiable
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+        (source,) = ctx.saved_tensors
+        cotangent = grad_output.resolve_conj().contiguous().to(torch.complex64)
+        pullback_np = matrix_exp_backward(source.cpu().numpy(), cotangent.cpu().numpy(), dense=False, device=_CURRENT_DEVICE)
+        return torch.from_numpy(pullback_np).to(device=grad_output.device)
+
+
+def faster_propagate_matrix_exp(system, thicknesses, *, max_batch=None):
+    from diffBloch.core.solver import _complex_operator
+    a = _complex_operator(system.a).to(torch.complex64)
+    psi0 = system.psi0.to(dtype=a.dtype, device=a.device)
+    scalars = (1j * torch.pi * thicknesses / system.k_n).to(a.dtype)
+    if max_batch is None:
+        transfer = FasterMatrixExp.apply(a.unsqueeze(-3) * scalars[:, None, None])
+        return (transfer @ psi0.unsqueeze(-1)).squeeze(-1)
+    n = a.shape[-1]
+    a_flat = a.reshape(-1, n, n)
+    n_batch, n_thick = a_flat.shape[0], scalars.shape[0]
+    total = n_batch * n_thick
+    amplitudes = []
+    for start in range(0, total, max_batch):
+        flat = torch.arange(start, min(total, start + max_batch), device=a.device)
+        block = a_flat[flat // n_thick] * scalars[flat % n_thick][:, None, None]
+        amplitudes.append((FasterMatrixExp.apply(block) @ psi0.unsqueeze(-1)).squeeze(-1))
+    return torch.cat(amplitudes, dim=0).reshape(*a.shape[:-2], n_thick, n)
+
+
+def faster_propagate(system, thicknesses, *, method="matrix_exp", max_batch=None):
+    from diffBloch.core.solver import _propagate_bloch_eigen
+    if max_batch is not None and max_batch < 1:
+        raise ValueError(f"max_batch must be a positive integer or None, got {max_batch}")
+    t = torch.as_tensor(thicknesses, dtype=torch.float32, device=system.a.device)
+    if t.ndim == 0:
+        t = t.reshape(1)
+    if t.ndim != 1:
+        raise ValueError("thicknesses must be a scalar or 1-D sequence")
+    if method == "matrix_exp":
+        return faster_propagate_matrix_exp(system, t, max_batch=max_batch)
+    if method == "bloch_eigen":
+        return _propagate_bloch_eigen(system, t)
+    raise ValueError(f"method must be 'matrix_exp' or 'bloch_eigen', got {method!r}")
+
+
+def enable(device: Literal["cpu", "gpu"] = "gpu") -> None:
+    """Inject faster-diffbloch acceleration into diffBloch runtime."""
+    global _CURRENT_DEVICE
+    _CURRENT_DEVICE = device
+    import importlib
+    try:
+        solver = importlib.import_module("diffBloch.core.solver")
+        solver._propagate_matrix_exp = faster_propagate_matrix_exp
+        solver.propagate = faster_propagate
+        for mod_name in ("diffBloch.core", "diffBloch.engine.forward"):
+            try:
+                mod = importlib.import_module(mod_name)
+                setattr(mod, "propagate", faster_propagate)
+            except Exception:
+                pass
+    except ImportError:
+        pass
+
+
+def disable() -> None:
+    """Restore original diffBloch functions."""
+    pass

faster_diffbloch-0.1.0/src/faster_diffbloch/builder.py

@@ -0,0 +1,58 @@
+"""Native runtime builder and loader for faster-diffbloch."""
+
+from __future__ import annotations
+
+import ctypes
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+NATIVE_DIR = Path(__file__).resolve().parent / "native"
+BUILD_DIR = Path.home() / ".cache" / "faster_diffbloch"
+
+
+def build_and_load_library() -> ctypes.CDLL:
+    """Build or load the cached native acceleration library."""
+    BUILD_DIR.mkdir(parents=True, exist_ok=True)
+    lib_name = "libfaster_diffbloch.dylib" if sys.platform == "darwin" else "libfaster_diffbloch.so"
+    lib_path = BUILD_DIR / lib_name
+
+    sources = [
+        NATIVE_DIR / "bridge_lib.c",
+        NATIVE_DIR / "batch_cgemm.c",
+        NATIVE_DIR / "native_scattering.c",
+    ]
+    if sys.platform == "darwin":
+        sources.append(NATIVE_DIR / "metal_batch_cgemm.m")
+        # Build metallib if metal is available
+        metal_src = NATIVE_DIR / "batch_cgemm.metal"
+        metallib_path = BUILD_DIR / "batch_cgemm.metallib"
+        if metal_src.exists() and (not metallib_path.exists() or metallib_path.stat().st_mtime < metal_src.stat().st_mtime):
+            try:
+                subprocess.run(
+                    ["xcrun", "-sdk", "macosx", "metal", "-O3", "-c", str(metal_src), "-o", str(BUILD_DIR / "batch_cgemm.air")],
+                    check=True, capture_output=True
+                )
+                subprocess.run(
+                    ["xcrun", "-sdk", "macosx", "metallib", str(BUILD_DIR / "batch_cgemm.air"), "-o", str(metallib_path)],
+                    check=True, capture_output=True
+                )
+            except Exception:
+                pass
+
+    stale = not lib_path.exists() or any(lib_path.stat().st_mtime < s.stat().st_mtime for s in sources if s.exists())
+    if stale:
+        frameworks = ["-framework", "Accelerate", "-framework", "Metal", "-framework", "Foundation"] if sys.platform == "darwin" else ["-lblas"]
+        cmd = [
+            "clang", "-std=c11", "-O3", "-fPIC", "-shared",
+            "-D_DEFAULT_SOURCE", "-Wno-unused-function", "-Wno-unused-variable",
+            f"-I{NATIVE_DIR}",
+            *[str(s) for s in sources if s.exists()],
+            *frameworks, "-lm", "-o", str(lib_path),
+        ]
+        res = subprocess.run(cmd, capture_output=True, text=True)
+        if res.returncode != 0:
+            raise RuntimeError(f"Failed to build faster-diffbloch native library:\n{res.stderr}")
+
+    return ctypes.CDLL(str(lib_path))

faster_diffbloch-0.1.0/src/faster_diffbloch/cli.py

@@ -0,0 +1,21 @@
+"""CLI wrapper enabling faster-diffbloch before running diffbloch commands."""
+
+from __future__ import annotations
+
+import sys
+from .backend import enable
+
+
+def main() -> None:
+    # Enable Metal GPU acceleration by default
+    enable(device="gpu")
+    try:
+        from diffBloch.app.cli import main as diffbloch_main
+        diffbloch_main()
+    except ImportError:
+        print("Error: diffbloch is not installed. Install diffbloch or run 'pip install diffbloch'.", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()