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

{faster_diffbloch-0.1.2 → faster_diffbloch-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: faster-diffbloch
-Version: 0.1.2
+Version: 0.1.4
 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/
@@ -9,7 +9,7 @@ Project-URL: Issues, https://github.com/godofecht/diffFlow/issues
 Project-URL: PyPI, https://pypi.org/project/faster-diffbloch/
 Project-URL: piwheels, https://www.piwheels.org/project/faster-diffbloch/
 Project-URL: Original diffBloch, https://diffbloch.com
-Author-email: Abhishek Shivakumar <abhishek@example.com>
+Author-email: Abhishek Shivakumar <abhishek.shivakumar@gmail.com>
 License-Expression: MIT
 License-File: LICENSE
 Keywords: Metal GPU,PyTorch acceleration,diffBloch,electron crystallography,electron diffraction,structure refinement
@@ -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.
 
-**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.
+A Metal GPU path for Apple Silicon and an optimized CPU path for macOS and Linux, measured against diffBloch running on PyTorch.
 
 | Package | Documentation | Source repository | Original project |
 | :--- | :--- | :--- | :--- |
@@ -63,38 +63,58 @@ the expensive propagation and gradient work onto an accelerated native backend.
 You get the same scientific calculation, with a faster path on supported hardware
 and a safe PyTorch fallback when the native backend is unavailable.
 
-The package is validated against the diffBloch test suite and reference results.
-It passes all 738 diffBloch tests, 55 additional conformance tests, and reproduces
-the experimental quartz refinement result.
+The package is validated against diffBloch's own test suite. With acceleration
+enabled it passes all 738 diffBloch tests and 55 conformance tests that compare
+the accelerated results against diffBloch's, field by field, including gradients.
+The diffBloch run makes 4855 calls into the native matrix exponential, so the
+accelerated path is exercised rather than skipped.
 
 ---
 
 ## Performance
 
-End-to-end forward-plus-backward timing on an Apple M4 Max, measured across the
-same diffBloch workload at different beam counts:
+Forward-plus-backward timing on an Apple M4 Max, single-threaded PyTorch,
+running the same diffBloch workload through `bench/compare_faster.py`. This is
+what `pip install faster-diffbloch` gives you: the accelerated matrix
+exponential bridged into diffBloch, with PyTorch handling everything else.
 
-| Beams | faster-diffBloch | PyTorch | Speedup |
+| Beams | faster-diffbloch | PyTorch | Speedup |
 | :---: | :---: | :---: | :---: |
-| 31 | 0.90 ms | 1.62 ms | **1.80×** |
-| 61 | 1.37 ms | 2.24 ms | **1.63×** |
-| 91 | 2.84 ms | 3.99 ms | **1.41×** |
-| 163 | 7.67 ms | 10.98 ms | **1.43×** |
-| 579 | 89.81 ms | 135.77 ms | **1.51×** |
-
-That is **1.41×–1.80× faster than PyTorch CPU** across the measured cases,
-including the largest 579-beam case. The same benchmark is **1.21× faster than
-the Mojo reference at 31 beams**, remains ahead through 163 beams, and is
-essentially level at 579 beams.
-
-On Apple Silicon, the Metal GPU path is also available. In the package's M4 Max
-comparison run at 579 beams, it measured **2.24× faster than PyTorch CPU for the
-forward pass** and **2.63× faster than PyTorch MPS for forward plus backward**.
-
-These figures are workload benchmarks, not a promise that every crystal,
-hardware configuration, or beam count will see the same result. The CPU table
-uses single-threaded PyTorch and the same Apple Accelerate environment for a
-like-for-like comparison.
+| 31 | 1.48 ms | 1.54 ms | 1.05x |
+| 61 | 1.93 ms | 2.07 ms | 1.07x |
+| 91 | 3.48 ms | 3.69 ms | 1.06x |
+| 163 | 8.76 ms | 10.14 ms | 1.16x |
+
+Best of three alternating runs per mode. The gain grows with beam count,
+because the matrix exponential takes a larger share of the work as the system
+grows and the fixed ctypes and PyTorch overhead matters less.
+
+### 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 entirely and is considerably faster. At 579 beams, forward plus
+backward, minimum of five runs:
+
+| Implementation | Forward + Backward |
+| :--- | :---: |
+| PyTorch | 117.21 ms |
+| Mojo port | 85.52 ms |
+| Flow port, C backend | 82.84 ms |
+| Flow port, MLIR backend | 79.74 ms |
+
+Those numbers are not what this package delivers. They are reachable by
+running the port directly, and they are why the package exists. See the
+[diffFlow documentation](https://godofecht.github.io/diffFlow/).
+
+These figures come from one workload on one machine. Other crystals, hardware
+configurations and beam counts will differ.
+
+### Metal GPU
+
+The Metal path is available on Apple Silicon through `enable(device="gpu")`.
+The table below is from an earlier run and has not been reproduced against the
+current package, so treat it as indicative.
 
 ### Package benchmark snapshot
 

{faster_diffbloch-0.1.2 → faster_diffbloch-0.1.4}/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.
 
-**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.
+A Metal GPU path for Apple Silicon and an optimized CPU path for macOS and Linux, measured against diffBloch running on PyTorch.
 
 | Package | Documentation | Source repository | Original project |
 | :--- | :--- | :--- | :--- |
@@ -30,38 +30,58 @@ the expensive propagation and gradient work onto an accelerated native backend.
 You get the same scientific calculation, with a faster path on supported hardware
 and a safe PyTorch fallback when the native backend is unavailable.
 
-The package is validated against the diffBloch test suite and reference results.
-It passes all 738 diffBloch tests, 55 additional conformance tests, and reproduces
-the experimental quartz refinement result.
+The package is validated against diffBloch's own test suite. With acceleration
+enabled it passes all 738 diffBloch tests and 55 conformance tests that compare
+the accelerated results against diffBloch's, field by field, including gradients.
+The diffBloch run makes 4855 calls into the native matrix exponential, so the
+accelerated path is exercised rather than skipped.
 
 ---
 
 ## Performance
 
-End-to-end forward-plus-backward timing on an Apple M4 Max, measured across the
-same diffBloch workload at different beam counts:
+Forward-plus-backward timing on an Apple M4 Max, single-threaded PyTorch,
+running the same diffBloch workload through `bench/compare_faster.py`. This is
+what `pip install faster-diffbloch` gives you: the accelerated matrix
+exponential bridged into diffBloch, with PyTorch handling everything else.
 
-| Beams | faster-diffBloch | PyTorch | Speedup |
+| Beams | faster-diffbloch | PyTorch | Speedup |
 | :---: | :---: | :---: | :---: |
-| 31 | 0.90 ms | 1.62 ms | **1.80×** |
-| 61 | 1.37 ms | 2.24 ms | **1.63×** |
-| 91 | 2.84 ms | 3.99 ms | **1.41×** |
-| 163 | 7.67 ms | 10.98 ms | **1.43×** |
-| 579 | 89.81 ms | 135.77 ms | **1.51×** |
-
-That is **1.41×–1.80× faster than PyTorch CPU** across the measured cases,
-including the largest 579-beam case. The same benchmark is **1.21× faster than
-the Mojo reference at 31 beams**, remains ahead through 163 beams, and is
-essentially level at 579 beams.
-
-On Apple Silicon, the Metal GPU path is also available. In the package's M4 Max
-comparison run at 579 beams, it measured **2.24× faster than PyTorch CPU for the
-forward pass** and **2.63× faster than PyTorch MPS for forward plus backward**.
-
-These figures are workload benchmarks, not a promise that every crystal,
-hardware configuration, or beam count will see the same result. The CPU table
-uses single-threaded PyTorch and the same Apple Accelerate environment for a
-like-for-like comparison.
+| 31 | 1.48 ms | 1.54 ms | 1.05x |
+| 61 | 1.93 ms | 2.07 ms | 1.07x |
+| 91 | 3.48 ms | 3.69 ms | 1.06x |
+| 163 | 8.76 ms | 10.14 ms | 1.16x |
+
+Best of three alternating runs per mode. The gain grows with beam count,
+because the matrix exponential takes a larger share of the work as the system
+grows and the fixed ctypes and PyTorch overhead matters less.
+
+### 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 entirely and is considerably faster. At 579 beams, forward plus
+backward, minimum of five runs:
+
+| Implementation | Forward + Backward |
+| :--- | :---: |
+| PyTorch | 117.21 ms |
+| Mojo port | 85.52 ms |
+| Flow port, C backend | 82.84 ms |
+| Flow port, MLIR backend | 79.74 ms |
+
+Those numbers are not what this package delivers. They are reachable by
+running the port directly, and they are why the package exists. See the
+[diffFlow documentation](https://godofecht.github.io/diffFlow/).
+
+These figures come from one workload on one machine. Other crystals, hardware
+configurations and beam counts will differ.
+
+### Metal GPU
+
+The Metal path is available on Apple Silicon through `enable(device="gpu")`.
+The table below is from an earlier run and has not been reproduced against the
+current package, so treat it as indicative.
 
 ### Package benchmark snapshot
 

{faster_diffbloch-0.1.2 → faster_diffbloch-0.1.4}/pyproject.toml

@@ -4,14 +4,14 @@ build-backend = "hatchling.build"
 
 [project]
 name = "faster-diffbloch"
-version = "0.1.2"
+version = "0.1.4"
 description = "Drop-in Metal GPU (macOS) and optimized CPU (macOS/Linux) acceleration for diffBloch"
 readme = "README.md"
 requires-python = ">=3.10"
 keywords = ["diffBloch", "electron diffraction", "electron crystallography", "structure refinement", "PyTorch acceleration", "Metal GPU"]
 license = "MIT"
 authors = [
-    { name = "Abhishek Shivakumar", email = "abhishek@example.com" }
+    { name = "Abhishek Shivakumar", email = "abhishek.shivakumar@gmail.com" }
 ]
 classifiers = [
     "Development Status :: 4 - Beta",

{faster_diffbloch-0.1.2 → faster_diffbloch-0.1.4}/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.2"
+__version__ = "0.1.4"
 __all__ = [
     "enable",
     "disable",

{faster_diffbloch-0.1.2 → faster_diffbloch-0.1.4}/src/faster_diffbloch/backend.py

@@ -20,6 +20,11 @@ from .builder import (
 _LIB: ctypes.CDLL | None = None
 _CURRENT_DEVICE: Literal["cpu", "gpu"] = "gpu"
 
+# What enable() replaced, keyed by (module name, attribute), so disable() can
+# put it back. A sentinel marks an attribute that did not exist before.
+_MISSING = object()
+_ORIGINALS: dict[tuple[str, str], object] = {}
+
 
 def get_library() -> ctypes.CDLL | None:
     global _LIB
@@ -181,11 +186,14 @@ def enable(device: Literal["cpu", "gpu"] = "gpu") -> None:
     import importlib
     try:
         solver = importlib.import_module("diffBloch.core.solver")
+        _remember(solver, "_propagate_matrix_exp")
+        _remember(solver, "propagate")
         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)
+                _remember(mod, "propagate")
                 setattr(mod, "propagate", faster_propagate)
             except Exception:
                 pass
@@ -193,6 +201,33 @@ def enable(device: Literal["cpu", "gpu"] = "gpu") -> None:
         pass
 
 
+def _remember(module, name: str) -> None:
+    """Record a module attribute the first time it is replaced.
+
+    Only the first value is kept, so calling enable() twice does not record
+    the accelerated function as the thing to restore.
+    """
+    key = (module.__name__, name)
+    if key not in _ORIGINALS:
+        _ORIGINALS[key] = getattr(module, name, _MISSING)
+
+
 def disable() -> None:
-    """Restore original diffBloch functions."""
-    pass
+    """Put diffBloch's own functions back.
+
+    Undoes what enable() replaced, so a later call runs stock diffBloch.
+    Calling this without a preceding enable() does nothing.
+    """
+    import importlib
+
+    for (mod_name, attr), original in list(_ORIGINALS.items()):
+        try:
+            module = importlib.import_module(mod_name)
+        except ImportError:
+            continue
+        if original is _MISSING:
+            if hasattr(module, attr):
+                delattr(module, attr)
+        else:
+            setattr(module, attr, original)
+    _ORIGINALS.clear()