faster-diffbloch 0.1.0__py3-none-any.whl

faster_diffbloch/native/native_scattering.c

@@ -0,0 +1,233 @@
+#include "native_scattering.h"
+#include <math.h>
+#include <stdlib.h>
+#include <string.h>
+#include <dispatch/dispatch.h>
+
+#define PI 3.14159265358979323846
+#define TWO_PI 6.283185307179586
+#define MINUS_TWO_PI_SQ -19.739208802178716
+
+static inline double native_lobato(const double* coeff_a, const double* coeff_b, int32_t base, double g2) {
+    double f = 0.0;
+    for (int i = 0; i < 5; i++) {
+        double b = coeff_b[base + i];
+        double denom = 1.0 + b * g2;
+        f += coeff_a[base + i] * (2.0 + b * g2) / (denom * denom);
+    }
+    return f;
+}
+
+static inline double native_dwf_exponent(const double* uij, int32_t base, double h0, double h1, double h2) {
+    double quad = h0 * (uij[base] * h0 + uij[base + 1] * h1 + uij[base + 2] * h2)
+                + h1 * (uij[base + 3] * h0 + uij[base + 4] * h1 + uij[base + 5] * h2)
+                + h2 * (uij[base + 6] * h0 + uij[base + 7] * h1 + uij[base + 8] * h2);
+    return MINUS_TWO_PI_SQ * quad;
+}
+
+typedef struct {
+    double value;
+    double derivative;
+} AbsorptiveResult;
+
+static inline AbsorptiveResult native_absorptive_factor(
+    const NativeCase* c, int32_t atom, int32_t slot, double b, int clamped
+) {
+    double width = c->abs_knot_width[atom];
+    double t = (b - c->abs_knot_lo[atom]) / width;
+    double t2 = t * t;
+    double t3 = t2 * t;
+    double raw = (2.0 * t3 - 3.0 * t2 + 1.0) * c->abs_y0[slot]
+               + (t3 - 2.0 * t2 + t) * width * c->abs_d0[slot]
+               + (-2.0 * t3 + 3.0 * t2) * c->abs_y1[slot]
+               + (t3 - t2) * width * c->abs_d1[slot];
+    if (raw <= 0.0) {
+        return (AbsorptiveResult){ 0.0, 0.0 };
+    }
+    if (clamped) {
+        return (AbsorptiveResult){ c->abs_c_over_v * raw, 0.0 };
+    }
+    double slope = (6.0 * t2 - 6.0 * t) * c->abs_y0[slot]
+                 + (3.0 * t2 - 4.0 * t + 1.0) * width * c->abs_d0[slot]
+                 + (-6.0 * t2 + 6.0 * t) * c->abs_y1[slot]
+                 + (3.0 * t2 - 2.0 * t) * width * c->abs_d1[slot];
+    return (AbsorptiveResult){
+        c->abs_c_over_v * raw,
+        c->abs_c_over_v * slope / width
+    };
+}
+
+void native_parallel_structure_factors(
+    const NativeCase* c,
+    double _Complex* fgb,
+    NativePhaseCache* cache
+) {
+    int32_t n_atoms = c->n_atoms;
+    int32_t n_grid = c->n_grid;
+
+    double* b_iso = (double*)malloc(n_atoms * sizeof(double));
+    int32_t* clamped = (int32_t*)malloc(n_atoms * sizeof(int32_t));
+
+    for (int a = 0; a < n_atoms; a++) {
+        clamped[a] = 0;
+        if (c->absorption == 1) {
+            int32_t base = a * 9;
+            double b = 0.0;
+            for (int i = 0; i < 9; i++) {
+                b += c->abs_db_du[i] * c->uij[base + i];
+            }
+            if (b < 0.1) { b = 0.1; clamped[a] = 1; }
+            if (b > 4.0) { b = 4.0; clamped[a] = 1; }
+            b_iso[a] = b;
+        } else {
+            b_iso[a] = 0.0;
+        }
+    }
+
+    size_t num_chunks = 32;
+    if (n_grid < 100) num_chunks = 1;
+
+    dispatch_apply(num_chunks, dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_HIGH, 0), ^(size_t chunk) {
+        int32_t start_m = (int32_t)(chunk * n_grid / num_chunks);
+        int32_t end_m = (int32_t)((chunk + 1) * n_grid / num_chunks);
+
+        for (int32_t m = start_m; m < end_m; m++) {
+            double g2 = c->grid_g[m] * c->grid_g[m];
+            double h0 = (double)c->grid_hkl[m * 3];
+            double h1 = (double)c->grid_hkl[m * 3 + 1];
+            double h2 = (double)c->grid_hkl[m * 3 + 2];
+            double re = 0.0;
+            double im = 0.0;
+
+            for (int a = 0; a < n_atoms; a++) {
+                double form = native_lobato(c->lobato_a, c->lobato_b, a * 5, g2);
+                double dwf = exp(native_dwf_exponent(c->uij, a * 9, h0, h1, h2));
+                double amplitude = form * dwf / c->volume;
+                double phase = TWO_PI * (c->positions[a * 3] * h0
+                                       + c->positions[a * 3 + 1] * h1
+                                       + c->positions[a * 3 + 2] * h2);
+                double cp = cos(phase);
+                double sp = sin(phase);
+                int32_t slot = m * n_atoms + a;
+
+                double imag_amplitude = 0.0;
+                double dimag_db = 0.0;
+                if (c->absorption == 1) {
+                    int32_t abs_slot = a * n_grid + m;
+                    AbsorptiveResult fp = native_absorptive_factor(c, a, abs_slot, b_iso[a], clamped[a]);
+                    imag_amplitude = fp.value * dwf / c->volume;
+                    dimag_db = fp.derivative * dwf / c->volume;
+                }
+
+                cache->amplitude[slot] = amplitude;
+                cache->cosine[slot] = cp;
+                cache->sine[slot] = sp;
+                cache->imag_amplitude[slot] = imag_amplitude;
+                cache->dimag_db[slot] = dimag_db;
+
+                double occ = c->occupancies[a];
+                re += occ * (amplitude * cp - imag_amplitude * sp);
+                im += occ * (amplitude * sp + imag_amplitude * cp);
+            }
+            fgb[m] = re + I * im;
+        }
+    });
+
+    free(b_iso);
+    free(clamped);
+}
+
+void native_parallel_structure_factors_backward(
+    const NativeCase* c,
+    const NativePhaseCache* cache,
+    const double _Complex* fbar,
+    double* grad_positions,
+    double* grad_occupancies,
+    double* grad_uij
+) {
+    int32_t n_atoms = c->n_atoms;
+    int32_t n_grid = c->n_grid;
+
+    for (int i = 0; i < n_atoms * 3; i++) grad_positions[i] = 0.0;
+    for (int i = 0; i < n_atoms; i++) grad_occupancies[i] = 0.0;
+    for (int i = 0; i < n_atoms * 9; i++) grad_uij[i] = 0.0;
+
+    size_t num_chunks = 16;
+    if (n_grid < 100) num_chunks = 1;
+
+    double* local_gpos = (double*)calloc(num_chunks * n_atoms * 3, sizeof(double));
+    double* local_gocc = (double*)calloc(num_chunks * n_atoms, sizeof(double));
+    double* local_guij = (double*)calloc(num_chunks * n_atoms * 9, sizeof(double));
+
+    dispatch_apply(num_chunks, dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_HIGH, 0), ^(size_t chunk) {
+        int32_t start_m = (int32_t)(chunk * n_grid / num_chunks);
+        int32_t end_m = (int32_t)((chunk + 1) * n_grid / num_chunks);
+
+        double* t_gpos = local_gpos + chunk * n_atoms * 3;
+        double* t_gocc = local_gocc + chunk * n_atoms;
+        double* t_guij = local_guij + chunk * n_atoms * 9;
+
+        for (int32_t m = start_m; m < end_m; m++) {
+            double h0 = (double)c->grid_hkl[m * 3];
+            double h1 = (double)c->grid_hkl[m * 3 + 1];
+            double h2 = (double)c->grid_hkl[m * 3 + 2];
+            double bar_re = creal(fbar[m]);
+            double bar_im = cimag(fbar[m]);
+
+            for (int a = 0; a < n_atoms; a++) {
+                int32_t slot = m * n_atoms + a;
+                double amplitude = cache->amplitude[slot];
+                double imag_amplitude = cache->imag_amplitude[slot];
+                double cp = cache->cosine[slot];
+                double sp = cache->sine[slot];
+                double occ = c->occupancies[a];
+
+                double along = bar_re * cp + bar_im * sp;
+                double across = bar_im * cp - bar_re * sp;
+
+                double projected = amplitude * along + imag_amplitude * across;
+                t_gocc[a] += projected;
+
+                double rotated = occ * (amplitude * across - imag_amplitude * along) * TWO_PI;
+                t_gpos[a * 3] += rotated * h0;
+                t_gpos[a * 3 + 1] += rotated * h1;
+                t_gpos[a * 3 + 2] += rotated * h2;
+
+                double scaled = MINUS_TWO_PI_SQ * occ * projected;
+                int32_t base = a * 9;
+                if (c->absorption == 1) {
+                    double through_b = occ * cache->dimag_db[slot] * across;
+                    for (int i = 0; i < 9; i++) {
+                        t_guij[base + i] += through_b * c->abs_db_du[i];
+                    }
+                }
+                double xy = scaled * h0 * h1;
+                double xz = scaled * h0 * h2;
+                double yz = scaled * h1 * h2;
+                t_guij[base]     += scaled * h0 * h0;
+                t_guij[base + 1] += xy;
+                t_guij[base + 2] += xz;
+                t_guij[base + 3] += xy;
+                t_guij[base + 4] += scaled * h1 * h1;
+                t_guij[base + 5] += yz;
+                t_guij[base + 6] += xz;
+                t_guij[base + 7] += yz;
+                t_guij[base + 8] += scaled * h2 * h2;
+            }
+        }
+    });
+
+    // Reduce thread-local accumulators
+    for (size_t chunk = 0; chunk < num_chunks; chunk++) {
+        const double* t_gpos = local_gpos + chunk * n_atoms * 3;
+        const double* t_gocc = local_gocc + chunk * n_atoms;
+        const double* t_guij = local_guij + chunk * n_atoms * 9;
+        for (int i = 0; i < n_atoms * 3; i++) grad_positions[i] += t_gpos[i];
+        for (int i = 0; i < n_atoms; i++) grad_occupancies[i] += t_gocc[i];
+        for (int i = 0; i < n_atoms * 9; i++) grad_uij[i] += t_guij[i];
+    }
+
+    free(local_gpos);
+    free(local_gocc);
+    free(local_guij);
+}

faster_diffbloch/native/native_scattering.h

@@ -0,0 +1,60 @@
+#ifndef NATIVE_SCATTERING_H
+#define NATIVE_SCATTERING_H
+
+#include <stdint.h>
+#include <complex.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct {
+    int32_t n_atoms;
+    int32_t n_grid;
+    int32_t absorption;
+    double volume;
+    double abs_c_over_v;
+    const double* positions;     /* n_atoms * 3 */
+    const double* occupancies;   /* n_atoms */
+    const double* uij;           /* n_atoms * 9 */
+    const double* lobato_a;      /* n_atoms * 5 */
+    const double* lobato_b;      /* n_atoms * 5 */
+    const int32_t* grid_hkl;     /* n_grid * 3 */
+    const double* grid_g;        /* n_grid */
+    const double* abs_db_du;     /* 9 */
+    const double* abs_knot_lo;   /* n_atoms */
+    const double* abs_knot_width;/* n_atoms */
+    const double* abs_y0;        /* n_atoms * n_grid */
+    const double* abs_y1;
+    const double* abs_d0;
+    const double* abs_d1;
+} NativeCase;
+
+typedef struct {
+    double* amplitude;           /* n_grid * n_atoms */
+    double* cosine;
+    double* sine;
+    double* imag_amplitude;
+    double* dimag_db;
+} NativePhaseCache;
+
+void native_parallel_structure_factors(
+    const NativeCase* c,
+    double _Complex* fgb,
+    NativePhaseCache* cache
+);
+
+void native_parallel_structure_factors_backward(
+    const NativeCase* c,
+    const NativePhaseCache* cache,
+    const double _Complex* fbar,
+    double* grad_positions,
+    double* grad_occupancies,
+    double* grad_uij
+);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* NATIVE_SCATTERING_H */

faster_diffbloch-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,16 @@
+faster_diffbloch/__init__.py,sha256=XaxuW-wZ-zyzryZ6vdOZbabHWfBPdDfn41oTmhLmvsg,330
+faster_diffbloch/backend.py,sha256=iYGZ-Y7KEL6d9S0JSgcRkUHK2akFFW5s2PTbbKYmjyk,5866
+faster_diffbloch/builder.py,sha256=D2CWVSUHWVcj0BEV5KOavuKeXA22-RTKgImSo-nAIc0,2447
+faster_diffbloch/cli.py,sha256=kwX6W-vfRhCRS_VvLRzor-3XPkFXfgrJUYwsu84hGl0,547
+faster_diffbloch/native/batch_cgemm.c,sha256=VX_hr9oauetvtna-MbJJ4y3yIs5ct-bzLpq4Rn9nrP0,5935
+faster_diffbloch/native/batch_cgemm.h,sha256=9BOPBfAyA87sGBs3mS0oe8nfrjz4oHim2hTPnoq3P7E,3135
+faster_diffbloch/native/batch_cgemm.metal,sha256=A0UPhvY3ve-ITGY9CQBVxDOD59TzaIZBG6ABgjNr03I,11937
+faster_diffbloch/native/bridge_lib.c,sha256=v-tOdgzL78VuGWdruyKcTEbRqxVfDivK_gaUG5PjSpY,31072
+faster_diffbloch/native/metal_batch_cgemm.m,sha256=nQqewJ3OuYePLcglKf63qsD1Jqzsaa7rtZv6hvYEF28,34129
+faster_diffbloch/native/native_scattering.c,sha256=M6nCLothcCNG0BMLTCFAb3w9d8RZWQzgUpBIVsouBZg,9032
+faster_diffbloch/native/native_scattering.h,sha256=ynMJpbptbZZwRcJvo-azzDje9Cs5x1m9xsBbUkmWm5c,1488
+faster_diffbloch-0.1.0.dist-info/METADATA,sha256=1lR3t7kNdwiHMShutr7h-eKLqMgxc12QgJv-RkK-XM4,3549
+faster_diffbloch-0.1.0.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
+faster_diffbloch-0.1.0.dist-info/entry_points.txt,sha256=pPjLbGQUv0VMhCee6wfeBgMPKVgdc32ycSjmjaDx82s,106
+faster_diffbloch-0.1.0.dist-info/licenses/LICENSE,sha256=Fhvu-xKhUMhKf6kCYY_xMnvH7vMPU6y6kggs07icYMs,1076
+faster_diffbloch-0.1.0.dist-info/RECORD,,

faster_diffbloch-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.0
+Root-Is-Purelib: true
+Tag: py3-none-any

faster_diffbloch-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+diffbloch-fast = faster_diffbloch.cli:main
+faster-diffbloch = faster_diffbloch.cli:main

faster_diffbloch-0.1.0.dist-info/licenses/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.