faster-diffbloch 0.1.1__tar.gz → 0.1.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,14 +1,18 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: faster-diffbloch
3
- Version: 0.1.1
3
+ Version: 0.1.2
4
4
  Summary: Drop-in Metal GPU (macOS) and optimized CPU (macOS/Linux) acceleration for diffBloch
5
5
  Project-URL: Homepage, https://godofecht.github.io/diffFlow/
6
+ Project-URL: Documentation, https://godofecht.github.io/diffFlow/
6
7
  Project-URL: Repository, https://github.com/godofecht/diffFlow
7
8
  Project-URL: Issues, https://github.com/godofecht/diffFlow/issues
9
+ Project-URL: PyPI, https://pypi.org/project/faster-diffbloch/
10
+ Project-URL: piwheels, https://www.piwheels.org/project/faster-diffbloch/
8
11
  Project-URL: Original diffBloch, https://diffbloch.com
9
12
  Author-email: Abhishek Shivakumar <abhishek@example.com>
10
13
  License-Expression: MIT
11
14
  License-File: LICENSE
15
+ Keywords: Metal GPU,PyTorch acceleration,diffBloch,electron crystallography,electron diffraction,structure refinement
12
16
  Classifier: Development Status :: 4 - Beta
13
17
  Classifier: Intended Audience :: Science/Research
14
18
  Classifier: License :: OSI Approved :: MIT License
@@ -27,13 +31,15 @@ Provides-Extra: diffbloch
27
31
  Requires-Dist: diffbloch; extra == 'diffbloch'
28
32
  Description-Content-Type: text/markdown
29
33
 
30
- # faster-diffBloch
34
+ # faster-diffbloch
31
35
 
32
36
  Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) electron crystallography structure refinement.
33
37
 
34
- Documentation and comparison benchmarks: [https://godofecht.github.io/diffFlow/](https://godofecht.github.io/diffFlow/)
38
+ **Up to 1.80× faster than PyTorch CPU** in measured end-to-end forward-plus-backward benchmarks, with a Metal GPU path for Apple Silicon and an optimized CPU path for macOS and Linux.
35
39
 
36
- Original diffBloch project: [https://diffbloch.com](https://diffbloch.com)
40
+ | Package | Documentation | Source repository | Original project |
41
+ | :--- | :--- | :--- | :--- |
42
+ | [PyPI](https://pypi.org/project/faster-diffbloch/) · [piwheels](https://www.piwheels.org/project/faster-diffbloch/) | [diffFlow documentation and benchmarks](https://godofecht.github.io/diffFlow/) | [github.com/godofecht/diffFlow](https://github.com/godofecht/diffFlow) | [diffbloch.com](https://diffbloch.com) |
37
43
 
38
44
  ---
39
45
 
@@ -52,20 +58,45 @@ Runtime platform guards automatically detect your operating system and hardware
52
58
 
53
59
  ## Why faster-diffBloch?
54
60
 
55
- 1. **Native Metal GPU Execution on macOS:**
56
- 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.
61
+ `faster-diffbloch` keeps the diffBloch workflow and public API intact while moving
62
+ the expensive propagation and gradient work onto an accelerated native backend.
63
+ You get the same scientific calculation, with a faster path on supported hardware
64
+ and a safe PyTorch fallback when the native backend is unavailable.
57
65
 
58
- 2. **Blocked-Pair Adjoint Formulation:**
59
- 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).
60
-
61
- 3. **Bit-for-Bit Validation:**
62
- Passes all 738 unit tests in diffBloch and reproduces the experimental 99-rotation quartz dataset ($R_{\text{obs}} = 0.0486$).
66
+ The package is validated against the diffBloch test suite and reference results.
67
+ It passes all 738 diffBloch tests, 55 additional conformance tests, and reproduces
68
+ the experimental quartz refinement result.
63
69
 
64
70
  ---
65
71
 
66
72
  ## Performance
67
73
 
68
- Forward and backward timing comparison on Apple Silicon (M4 Max) at $N=579$ beams (CsPbBr3 scale):
74
+ End-to-end forward-plus-backward timing on an Apple M4 Max, measured across the
75
+ same diffBloch workload at different beam counts:
76
+
77
+ | Beams | faster-diffBloch | PyTorch | Speedup |
78
+ | :---: | :---: | :---: | :---: |
79
+ | 31 | 0.90 ms | 1.62 ms | **1.80×** |
80
+ | 61 | 1.37 ms | 2.24 ms | **1.63×** |
81
+ | 91 | 2.84 ms | 3.99 ms | **1.41×** |
82
+ | 163 | 7.67 ms | 10.98 ms | **1.43×** |
83
+ | 579 | 89.81 ms | 135.77 ms | **1.51×** |
84
+
85
+ That is **1.41×–1.80× faster than PyTorch CPU** across the measured cases,
86
+ including the largest 579-beam case. The same benchmark is **1.21× faster than
87
+ the Mojo reference at 31 beams**, remains ahead through 163 beams, and is
88
+ essentially level at 579 beams.
89
+
90
+ On Apple Silicon, the Metal GPU path is also available. In the package's M4 Max
91
+ comparison run at 579 beams, it measured **2.24× faster than PyTorch CPU for the
92
+ forward pass** and **2.63× faster than PyTorch MPS for forward plus backward**.
93
+
94
+ These figures are workload benchmarks, not a promise that every crystal,
95
+ hardware configuration, or beam count will see the same result. The CPU table
96
+ uses single-threaded PyTorch and the same Apple Accelerate environment for a
97
+ like-for-like comparison.
98
+
99
+ ### Package benchmark snapshot
69
100
 
70
101
  | Implementation | Forward | Forward + Backward | Speedup vs PyTorch CPU | Speedup vs PyTorch MPS |
71
102
  | :--- | :---: | :---: | :---: | :---: |
@@ -0,0 +1,112 @@
1
+ # faster-diffbloch
2
+
3
+ Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) electron crystallography structure refinement.
4
+
5
+ **Up to 1.80× faster than PyTorch CPU** in measured end-to-end forward-plus-backward benchmarks, with a Metal GPU path for Apple Silicon and an optimized CPU path for macOS and Linux.
6
+
7
+ | Package | Documentation | Source repository | Original project |
8
+ | :--- | :--- | :--- | :--- |
9
+ | [PyPI](https://pypi.org/project/faster-diffbloch/) · [piwheels](https://www.piwheels.org/project/faster-diffbloch/) | [diffFlow documentation and benchmarks](https://godofecht.github.io/diffFlow/) | [github.com/godofecht/diffFlow](https://github.com/godofecht/diffFlow) | [diffbloch.com](https://diffbloch.com) |
10
+
11
+ ---
12
+
13
+ ## Operating System and Platform Support
14
+
15
+ | Operating System / Hardware | CPU Acceleration (`device="cpu"`) | GPU Acceleration (`device="gpu"`) | Backend Runtime |
16
+ | :--- | :---: | :---: | :--- |
17
+ | **macOS Apple Silicon (M1/M2/M3/M4/Max/Ultra)** | **Supported** | **Supported** | Native Metal Compute Shaders + Apple Accelerate BLAS |
18
+ | **macOS Intel (x86_64)** | **Supported** | Fallback to CPU | Apple Accelerate BLAS |
19
+ | **Linux (x86_64 / aarch64)** | **Supported** | Fallback to CPU | OpenBLAS / C11 BLAS |
20
+ | **Windows** | PyTorch Reference | PyTorch Reference | Pure PyTorch Reference Fallback |
21
+
22
+ 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.
23
+
24
+ ---
25
+
26
+ ## Why faster-diffBloch?
27
+
28
+ `faster-diffbloch` keeps the diffBloch workflow and public API intact while moving
29
+ the expensive propagation and gradient work onto an accelerated native backend.
30
+ You get the same scientific calculation, with a faster path on supported hardware
31
+ and a safe PyTorch fallback when the native backend is unavailable.
32
+
33
+ The package is validated against the diffBloch test suite and reference results.
34
+ It passes all 738 diffBloch tests, 55 additional conformance tests, and reproduces
35
+ the experimental quartz refinement result.
36
+
37
+ ---
38
+
39
+ ## Performance
40
+
41
+ End-to-end forward-plus-backward timing on an Apple M4 Max, measured across the
42
+ same diffBloch workload at different beam counts:
43
+
44
+ | Beams | faster-diffBloch | PyTorch | Speedup |
45
+ | :---: | :---: | :---: | :---: |
46
+ | 31 | 0.90 ms | 1.62 ms | **1.80×** |
47
+ | 61 | 1.37 ms | 2.24 ms | **1.63×** |
48
+ | 91 | 2.84 ms | 3.99 ms | **1.41×** |
49
+ | 163 | 7.67 ms | 10.98 ms | **1.43×** |
50
+ | 579 | 89.81 ms | 135.77 ms | **1.51×** |
51
+
52
+ That is **1.41×–1.80× faster than PyTorch CPU** across the measured cases,
53
+ including the largest 579-beam case. The same benchmark is **1.21× faster than
54
+ the Mojo reference at 31 beams**, remains ahead through 163 beams, and is
55
+ essentially level at 579 beams.
56
+
57
+ On Apple Silicon, the Metal GPU path is also available. In the package's M4 Max
58
+ comparison run at 579 beams, it measured **2.24× faster than PyTorch CPU for the
59
+ forward pass** and **2.63× faster than PyTorch MPS for forward plus backward**.
60
+
61
+ These figures are workload benchmarks, not a promise that every crystal,
62
+ hardware configuration, or beam count will see the same result. The CPU table
63
+ uses single-threaded PyTorch and the same Apple Accelerate environment for a
64
+ like-for-like comparison.
65
+
66
+ ### Package benchmark snapshot
67
+
68
+ | Implementation | Forward | Forward + Backward | Speedup vs PyTorch CPU | Speedup vs PyTorch MPS |
69
+ | :--- | :---: | :---: | :---: | :---: |
70
+ | PyTorch CPU | 25.7 ms | 130.3 ms | 1.00x | 1.17x |
71
+ | PyTorch MPS (fallback) | 26.2 ms | 153.0 ms | 0.85x | 1.00x |
72
+ | **faster-diffBloch CPU** | **24.4 ms** | **83.5 ms** | **1.56x** | **1.83x** |
73
+ | **faster-diffBloch Metal GPU** | **13.1 ms** | **58.1 ms** | **2.24x** | **2.63x** |
74
+
75
+ ---
76
+
77
+ ## Installation
78
+
79
+ ```bash
80
+ pip install faster-diffbloch
81
+ ```
82
+
83
+ ---
84
+
85
+ ## Usage
86
+
87
+ ### 1. Drop-in CLI
88
+
89
+ Use `diffbloch-fast` or `faster-diffbloch` anywhere you would use `diffbloch`:
90
+
91
+ ```bash
92
+ diffbloch-fast infer examples/Colmey_et_al_2026/data/quartz-no-abs
93
+ diffbloch-fast refine examples/Colmey_et_al_2026/data/quartz-no-abs
94
+ ```
95
+
96
+ ### 2. Python API Injection
97
+
98
+ Enable acceleration inside any existing diffBloch script:
99
+
100
+ ```python
101
+ import faster_diffbloch
102
+
103
+ # Enable Metal GPU acceleration (macOS Apple Silicon)
104
+ faster_diffbloch.enable(device="gpu")
105
+
106
+ # Or CPU acceleration (macOS and Linux)
107
+ faster_diffbloch.enable(device="cpu")
108
+
109
+ # Run standard diffBloch code
110
+ import diffBloch
111
+ # All propagate and matrix_exp calls now route through faster-diffBloch
112
+ ```
@@ -4,10 +4,11 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "faster-diffbloch"
7
- version = "0.1.1"
7
+ version = "0.1.2"
8
8
  description = "Drop-in Metal GPU (macOS) and optimized CPU (macOS/Linux) acceleration for diffBloch"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10"
11
+ keywords = ["diffBloch", "electron diffraction", "electron crystallography", "structure refinement", "PyTorch acceleration", "Metal GPU"]
11
12
  license = "MIT"
12
13
  authors = [
13
14
  { name = "Abhishek Shivakumar", email = "abhishek@example.com" }
@@ -41,8 +42,11 @@ diffbloch-fast = "faster_diffbloch.cli:main"
41
42
 
42
43
  [project.urls]
43
44
  Homepage = "https://godofecht.github.io/diffFlow/"
45
+ Documentation = "https://godofecht.github.io/diffFlow/"
44
46
  Repository = "https://github.com/godofecht/diffFlow"
45
47
  Issues = "https://github.com/godofecht/diffFlow/issues"
48
+ PyPI = "https://pypi.org/project/faster-diffbloch/"
49
+ piwheels = "https://www.piwheels.org/project/faster-diffbloch/"
46
50
  "Original diffBloch" = "https://diffbloch.com"
47
51
 
48
52
  [tool.hatch.build.targets.wheel]
@@ -4,7 +4,7 @@ from __future__ import annotations
4
4
 
5
5
  from .backend import enable, disable, matrix_exp, matrix_exp_backward, faster_propagate
6
6
 
7
- __version__ = "0.1.0"
7
+ __version__ = "0.1.2"
8
8
  __all__ = [
9
9
  "enable",
10
10
  "disable",
@@ -46,15 +46,15 @@ def matrix_exp(a: np.ndarray, device: str = "gpu") -> np.ndarray:
46
46
  fn = getattr(lib, "bridge_matrix_exp_gpu_ptr_c64_i32_i32_ptr_c64", None)
47
47
  if fn:
48
48
  fn.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
49
- fn.restype = None
50
- fn(ptr_in, batch, n, ptr_out)
51
- return out
49
+ fn.restype = ctypes.c_int32
50
+ if fn(ptr_in, n, batch, ptr_out) == 0:
51
+ return out
52
52
 
53
53
  fn = getattr(lib, "bridge_matrix_exp_ptr_c64_i32_i32_ptr_c64", None)
54
54
  if fn:
55
55
  fn.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
56
56
  fn.restype = None
57
- fn(ptr_in, batch, n, ptr_out)
57
+ fn(ptr_in, n, batch, ptr_out)
58
58
  return out
59
59
 
60
60
  return torch.matrix_exp(torch.from_numpy(a)).numpy()
@@ -63,10 +63,13 @@ def matrix_exp(a: np.ndarray, device: str = "gpu") -> np.ndarray:
63
63
  def matrix_exp_backward(a: np.ndarray, ebar: np.ndarray, dense: bool = False, device: str = "gpu") -> np.ndarray:
64
64
  lib = get_library()
65
65
  if lib is None:
66
- tA = torch.from_numpy(a).requires_grad_(True)
67
- tE = torch.matrix_exp(tA)
68
- tE.backward(torch.from_numpy(ebar))
69
- return tA.grad.numpy()
66
+ # Custom autograd.Function backward methods run with grad mode disabled.
67
+ # Re-enable it for the correctness fallback.
68
+ with torch.enable_grad():
69
+ tA = torch.from_numpy(a).requires_grad_(True)
70
+ tE = torch.matrix_exp(tA)
71
+ tE.backward(torch.from_numpy(ebar))
72
+ return tA.grad.numpy()
70
73
 
71
74
  a_c = np.ascontiguousarray(a, dtype=np.complex64)
72
75
  ebar_c = np.ascontiguousarray(ebar, dtype=np.complex64)
@@ -83,15 +86,15 @@ def matrix_exp_backward(a: np.ndarray, ebar: np.ndarray, dense: bool = False, de
83
86
  fn = getattr(lib, "bridge_matrix_exp_backward_gpu_ptr_c64_ptr_c64_i32_i32_i32_ptr_c64", None)
84
87
  if fn:
85
88
  fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
86
- fn.restype = None
87
- fn(ptr_a, ptr_ebar, dense_val, batch, n, ptr_out)
88
- return out
89
+ fn.restype = ctypes.c_int32
90
+ if fn(ptr_a, ptr_ebar, n, batch, dense_val, ptr_out) == 0:
91
+ return out
89
92
 
90
93
  fn = getattr(lib, "bridge_matrix_exp_backward_ptr_c64_ptr_c64_i32_i32_i32_ptr_c64", None)
91
94
  if fn:
92
95
  fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
93
96
  fn.restype = None
94
- fn(ptr_a, ptr_ebar, dense_val, batch, n, ptr_out)
97
+ fn(ptr_a, ptr_ebar, n, batch, dense_val, ptr_out)
95
98
  return out
96
99
 
97
100
  tA = torch.from_numpy(a).requires_grad_(True)
@@ -1,85 +0,0 @@
1
- # faster-diffBloch
2
-
3
- Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) electron crystallography structure refinement.
4
-
5
- Documentation and comparison benchmarks: [https://godofecht.github.io/diffFlow/](https://godofecht.github.io/diffFlow/)
6
-
7
- Original diffBloch project: [https://diffbloch.com](https://diffbloch.com)
8
-
9
- ---
10
-
11
- ## Operating System and Platform Support
12
-
13
- | Operating System / Hardware | CPU Acceleration (`device="cpu"`) | GPU Acceleration (`device="gpu"`) | Backend Runtime |
14
- | :--- | :---: | :---: | :--- |
15
- | **macOS Apple Silicon (M1/M2/M3/M4/Max/Ultra)** | **Supported** | **Supported** | Native Metal Compute Shaders + Apple Accelerate BLAS |
16
- | **macOS Intel (x86_64)** | **Supported** | Fallback to CPU | Apple Accelerate BLAS |
17
- | **Linux (x86_64 / aarch64)** | **Supported** | Fallback to CPU | OpenBLAS / C11 BLAS |
18
- | **Windows** | PyTorch Reference | PyTorch Reference | Pure PyTorch Reference Fallback |
19
-
20
- 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.
21
-
22
- ---
23
-
24
- ## Why faster-diffBloch?
25
-
26
- 1. **Native Metal GPU Execution on macOS:**
27
- 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.
28
-
29
- 2. **Blocked-Pair Adjoint Formulation:**
30
- 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).
31
-
32
- 3. **Bit-for-Bit Validation:**
33
- Passes all 738 unit tests in diffBloch and reproduces the experimental 99-rotation quartz dataset ($R_{\text{obs}} = 0.0486$).
34
-
35
- ---
36
-
37
- ## Performance
38
-
39
- Forward and backward timing comparison on Apple Silicon (M4 Max) at $N=579$ beams (CsPbBr3 scale):
40
-
41
- | Implementation | Forward | Forward + Backward | Speedup vs PyTorch CPU | Speedup vs PyTorch MPS |
42
- | :--- | :---: | :---: | :---: | :---: |
43
- | PyTorch CPU | 25.7 ms | 130.3 ms | 1.00x | 1.17x |
44
- | PyTorch MPS (fallback) | 26.2 ms | 153.0 ms | 0.85x | 1.00x |
45
- | **faster-diffBloch CPU** | **24.4 ms** | **83.5 ms** | **1.56x** | **1.83x** |
46
- | **faster-diffBloch Metal GPU** | **13.1 ms** | **58.1 ms** | **2.24x** | **2.63x** |
47
-
48
- ---
49
-
50
- ## Installation
51
-
52
- ```bash
53
- pip install faster-diffbloch
54
- ```
55
-
56
- ---
57
-
58
- ## Usage
59
-
60
- ### 1. Drop-in CLI
61
-
62
- Use `diffbloch-fast` or `faster-diffbloch` anywhere you would use `diffbloch`:
63
-
64
- ```bash
65
- diffbloch-fast infer examples/Colmey_et_al_2026/data/quartz-no-abs
66
- diffbloch-fast refine examples/Colmey_et_al_2026/data/quartz-no-abs
67
- ```
68
-
69
- ### 2. Python API Injection
70
-
71
- Enable acceleration inside any existing diffBloch script:
72
-
73
- ```python
74
- import faster_diffbloch
75
-
76
- # Enable Metal GPU acceleration (macOS Apple Silicon)
77
- faster_diffbloch.enable(device="gpu")
78
-
79
- # Or CPU acceleration (macOS and Linux)
80
- faster_diffbloch.enable(device="cpu")
81
-
82
- # Run standard diffBloch code
83
- import diffBloch
84
- # All propagate and matrix_exp calls now route through faster-diffBloch
85
- ```