Compute a full array of Clebsch-Gordan coefficients in parallel

Author: Luthaf

Cargo.toml

@@ -19,6 +19,7 @@ lazy_static = "1"
 smallvec = "1"
 parking_lot = "0.12"
 lru = "0.7"
+rayon = "1"
 
 [dev-dependencies]
 approx = "0.5"

MANIFEST.in

@@ -1,9 +1,9 @@
 recursive-include src *
 include Cargo.*
 include pyproject.toml
+recursive-include benchmarks *
 
 recursive-exclude .github *
-recursive-exclude benchmarks *
 
 recursive-exclude python/tests *
 recursive-exclude python/wigners.egg-info *

pyproject.toml

@@ -12,6 +12,7 @@ legacy_tox_ini = """
 [testenv]
 deps =
     discover
+    numpy
 
 commands =
     discover -p "*.py" -s python/tests

python/tests/wigner.py

@@ -1,6 +1,7 @@
 import math
 import unittest
 
+import numpy as np
 import wigners
 
 
@@ -29,3 +30,18 @@ def test_clebsch_gordan(self):
             wigners.clebsch_gordan(j1=1, m1=1, j2=1, m2=0, j3=2, m3=1),
             1 / math.sqrt(2.0),
         )
+
+    def test_clebsch_gordan_array(self):
+        j1 = 12
+        j2 = 15
+        j3 = 8
+
+        expected = np.zeros((2 * j1 + 1, 2 * j2 + 1, 2 * j3 + 1), dtype=np.float64)
+        for m1 in range(-j1, j1 + 1):
+            for m2 in range(-j2, j2 + 1):
+                for m3 in range(-j3, j3 + 1):
+                    expected[j1 + m1, j2 + m2, j3 + m3] = wigners.clebsch_gordan(
+                        j1, m1, j2, m2, j3, m3
+                    )
+
+        self.assertTrue(np.all(wigners.clebsch_gordan_array(j1, j2, j3) == expected))

python/wigners/__init__.py

@@ -1,5 +1,6 @@
 import os
 import ctypes
+import numpy as np
 
 __lib = ctypes.CDLL(os.path.join(os.path.dirname(__file__), "_wigners.so"))
 
@@ -24,17 +25,56 @@
 __lib.clebsch_gordan.restype = ctypes.c_double
 
 
+__lib.clebsch_gordan_array_c.argtypes = [
+    ctypes.c_uint32,
+    ctypes.c_uint32,
+    ctypes.c_uint32,
+    ctypes.POINTER(ctypes.c_double),
+]
+__lib.clebsch_gordan_array_c.restype = None
+
+
 __lib.clear_wigner_3j_cache.argtypes = []
 __lib.clear_wigner_3j_cache.restype = None
 
 
 def wigner_3j(j1: int, j2: int, j3: int, m1: int, m2: int, m3: int) -> float:
+    """
+    Compute a single Wigner 3j coefficient, corresponding to:
+
+    .. code-block::
+
+        | j1  j2  j3 |
+        | m1  m2  m3 |
+    """
     return __lib.wigner_3j(j1, j2, j3, m1, m2, m3)
 
 
 def clebsch_gordan(j1: int, m1: int, j2: int, m2: int, j3: int, m3: int) -> float:
+    """
+    Compute a single Clebsch-Gordan coefficient, corresponding to:
+
+    .. code-block::
+
+        <j1 m1  j2 m2 | j3 m3>
+    """
     return __lib.clebsch_gordan(j1, m1, j2, m2, j3, m3)
 
 
+def clebsch_gordan_array(j1: int, j2: int, j3: int) -> np.ndarray:
+    """
+    Compute a full array of Clebsch-Gordan coefficient for the three given
+    ``j``.
+
+    The result is a 3-dimensional array with shape ``(2 * j1 + 1, 2 * j2 + 1, 2
+    * j3 + 1)``.
+    """
+    array = np.zeros((2 * j1 + 1, 2 * j2 + 1, 2 * j3 + 1), dtype=np.float64)
+    ptr = array.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
+    __lib.clebsch_gordan_array_c(j1, j2, j3, ptr, array.size)
+    return array
+
+
 def clear_wigner_3j_cache():
+    """Clear the LRU cache of Wigner 3j symbols"""
     return __lib.clear_wigner_3j_cache()

setup.cfg

@@ -27,6 +27,9 @@ zip_safe = False
 packages = find:
 package_dir =
     = python
+install_requires =
+    numpy
+
 
 [options.packages.find]
 where = python

src/wigner_3j.rs

@@ -1,6 +1,7 @@
 use parking_lot::Mutex;
 
 use lru::LruCache;
+use rayon::prelude::*;
 
 use crate::primes::{factorial, PrimeFactorization};
 use crate::rational::Rational;
@@ -117,6 +118,42 @@ pub extern fn clebsch_gordan(j1: u32, m1: i32, j2: u32, m2: i32, j3: u32, m3: i3
     }
 }
 
+
+/// Compute the full array of Clebsch-Gordan coefficients for the three given
+/// `j`.
+///
+/// Data will be written to `output`, which can be interpreted as a row-major
+/// 3-dimensional array with shape `(2 * j1 + 1, 2 * j2 + 1, 2 * j3 + 1)`.
+pub fn clebsch_gordan_array(j1: u32, j2: u32, j3: u32, output: &mut [f64]) {
+    let j1_size = 2 * j1 + 1;
+    let j2_size = 2 * j2 + 1;
+    let j3_size = 2 * j3 + 1;
+
+    let size = (j1_size * j2_size * j3_size) as usize;
+    if output.len() != size {
+        panic!(
+            "invalid output size, expected to have space for {} entries, but got {}",
+            size, output.len()
+        );
+    }
+
+    output.par_iter_mut().enumerate().for_each(|(i, o)| {
+        let i = i as u32;
+        let m1 = ((i / j3_size) / j2_size) as i32 - j1 as i32;
+        let m2 = ((i / j3_size) % j2_size) as i32 - j2 as i32;
+        let m3 = (i % j3_size) as i32 - j3 as i32;
+
+        *o = clebsch_gordan(j1, m1, j2, m2, j3, m3);
+    })
+}
+
+/// Same function as `clebsch_gordan_array`, but can be called directly from C
+#[no_mangle]
+pub unsafe extern fn clebsch_gordan_array_c(j1: u32, j2: u32, j3: u32, data: *mut f64, len: u64) {
+    let slice = std::slice::from_raw_parts_mut(data, len as usize);
+    clebsch_gordan_array(j1, j2, j3, slice);
+}
+
 /// check the triangle condition on j1, j2, j3, i.e. `|j1 - j2| <= j3 <= j1 + j2`
 fn triangle_condition(j1: u32, j2: u32, j3: u32) -> bool {
     return (j3 <= j1 + j2) && (j1 <= j2 + j3) && (j2 <= j3 + j1);