Update Cython wrapper to use NumPy memoryview interface (#31)

* Use memoryview rather than NumPy array buffer interface

* Further updates to use memoryview rather than buffer interface

* Add test for transforms on real signals

* Update additional reality tests

* Update Actions versions

* Add CMake env variable to avoid FFTW build error

Author: matt-graham

.github/workflows/cmake.yml

@@ -19,10 +19,10 @@ jobs:
         python-version: [3.8]
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
 

.github/workflows/python.yml

@@ -1,5 +1,10 @@
 name: Python Build
 
+env:
+  # Set to avoid errors when building FFTW with CMake v4+
+  # https://github.com/FFTW/fftw3/issues/381
+  CMAKE_POLICY_VERSION_MINIMUM: 3.5
+
 "on":
   push:
     branches: ["main"]
@@ -13,9 +18,9 @@ jobs:
     name: python source distribution
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python 3.8
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: 3.8
 
@@ -33,7 +38,7 @@ jobs:
       - name: run pytest
         run: pytest tests/
 
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         if: ${{ startsWith(github.ref, 'refs/tags') }}
         with:
           path: ./dist/*.tar.gz
@@ -49,7 +54,7 @@ jobs:
         python-version: [3.8]
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
 
       - uses: actions/setup-python@v4
         name: Install Python
@@ -71,7 +76,7 @@ jobs:
       - name: run pytests
         run: pytest tests
 
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         if: ${{ startsWith(github.ref, 'refs/tags') }}
         with:
           path: ./dist/*.whl
@@ -84,7 +89,7 @@ jobs:
     needs: [build_wheels, from-sdist]
     steps:
       - name: Download wheels and sdist
-        uses: actions/download-artifact@v3
+        uses: actions/download-artifact@v4
 
       - name: Move wheels and source distribution to dist/
         run: |

src/so3/bindings.pyx

@@ -1,11 +1,6 @@
 # cython: language_level=3
 
-# import both numpy and the Cython declarations for numpy
 import numpy as np
-cimport numpy as np
-
-# if you want to use the Numpy-C-API from Cython
-np.import_array()
 
 #----------------------------------------------------------------------------------------------------#
 
@@ -317,32 +312,48 @@ def is_elmn_non_zero(int el, int m, int n, so3_parameters):
 
 
 # forward and inverse for MW and MWSS for complex functions
-def inverse(np.ndarray[ double complex, ndim=1, mode="c"] flmn not None, so3_parameters not None):
-    cdef so3_parameters_t parameters=create_parameter_struct(so3_parameters)
-
+def inverse(double complex[::1] flmn not None, so3_parameters not None):
+    cdef so3_parameters_t parameters = create_parameter_struct(so3_parameters)
     if so3_parameters.reality:
-        f_length = f_size(so3_parameters)
-        f = np.zeros([f_length,], dtype=float)
-        so3_core_inverse_via_ssht_real(<double *> np.PyArray_DATA(f), <const double complex*> np.PyArray_DATA(flmn), &parameters)
+        return np.array(_inverse_real(flmn, parameters))
     else:
-        f_length = f_size(so3_parameters)
-        f = np.zeros([f_length,], dtype=complex)
-        so3_core_inverse_via_ssht(<double complex*> np.PyArray_DATA(f), <const double complex*> np.PyArray_DATA(flmn), &parameters)
+        return np.array(_inverse_complex(flmn, parameters))
 
+cdef double[::1] _inverse_real(double complex[::1] flmn, so3_parameters_t parameters):
+    cdef int f_length = so3_sampling_f_size(&parameters)
+    cdef double[::1] f = np.zeros([f_length,], dtype=float)
+    so3_core_inverse_via_ssht_real(&f[0], &flmn[0], &parameters)
     return f
 
-def forward(np.ndarray[ double complex, ndim=1, mode="c"] f not None, so3_parameters not None):
-    cdef so3_parameters_t parameters=create_parameter_struct(so3_parameters)
+cdef double complex[::1] _inverse_complex(double complex[::1] flmn, so3_parameters_t parameters):
+    cdef int f_length = so3_sampling_f_size(&parameters)
+    cdef double complex[::1] f = np.zeros([f_length,], dtype=complex)
+    so3_core_inverse_via_ssht(&f[0], &flmn[0], &parameters)
+    return f
 
-    if so3_parameters.reality:
-        flmn_length = flmn_size(so3_parameters)
-        flmn = np.zeros([flmn_length,], dtype=float)
-        so3_core_forward_via_ssht_real(<double complex*> np.PyArray_DATA(flmn), <const double *> np.PyArray_DATA(f), &parameters)
+ctypedef fused real_or_complex:
+    double
+    double complex
+
+def forward(real_or_complex[::1] f not None, so3_parameters not None):
+    cdef so3_parameters_t parameters = create_parameter_struct(so3_parameters)
+    if real_or_complex is double:
+        if not so3_parameters.reality:
+            raise ValueError("f has a real data type but reality flag not set")
+        return np.array(_forward_real(f, parameters))
     else:
-        flmn_length = flmn_size(so3_parameters)
-        flmn = np.zeros([flmn_length,], dtype=complex)
-        so3_core_forward_via_ssht(<double complex*> np.PyArray_DATA(flmn), <const double complex*> np.PyArray_DATA(f), &parameters)
+        return np.array(_forward_complex(f, parameters))
+    
+cdef double complex[::1] _forward_real(double[::1] f, so3_parameters_t parameters):
+    cdef int flmn_length = so3_sampling_flmn_size(&parameters)
+    cdef double complex[::1] flmn = np.zeros([flmn_length,], dtype=complex)
+    so3_core_forward_via_ssht_real(&flmn[0], &f[0], &parameters)
+    return flmn
 
+cdef double complex[::1] _forward_complex(double complex[::1] f, so3_parameters_t parameters):
+    cdef int flmn_length = so3_sampling_flmn_size(&parameters)
+    cdef double complex[::1] flmn = np.zeros([flmn_length,], dtype=complex)
+    so3_core_forward_via_ssht(&flmn[0], &f[0], &parameters)
     return flmn
 
 # convolution both in real and harmonic space and helper params function
@@ -368,9 +379,9 @@ def get_convolved_parameters(f_so3_parameters, g_so3_parameters):
     return SO3Parameters().from_dict(h_parameters)
 
 def convolve(
-    np.ndarray[ double complex, ndim=1, mode="c"] f not None, 
+    double complex[::1] f not None,
     f_parameters,
-    np.ndarray[ double complex, ndim=1, mode="c"] g not None, 
+    double complex[::1] g not None,
     g_parameters
     ):
 
@@ -381,22 +392,22 @@ def convolve(
     cdef so3_parameters_t h_parameters_struct=create_parameter_struct(h_parameters)
 
     h_length = f_size(h_parameters)
-    h = np.zeros([h_length,], dtype=complex)
+    cdef double complex[::1] h = np.zeros([h_length,], dtype=complex)
 
     so3_conv_convolution(
-        <double complex *> np.PyArray_DATA(h),
+        &h[0],
         &h_parameters_struct,
-        <const double complex *> np.PyArray_DATA(f),
+        &f[0],
         &f_parameters_struct,
-        <const double complex *> np.PyArray_DATA(g),
+        &g[0],
         &g_parameters_struct
     )
     return h, h_parameters
 
 def convolve_harmonic(    
-    np.ndarray[ double complex, ndim=1, mode="c"] flmn not None, 
+    double complex[::1] flmn not None,
     f_parameters,
-    np.ndarray[ double complex, ndim=1, mode="c"] glmn not None, 
+    double complex[::1] glmn not None,
     g_parameters
     ):
 
@@ -408,34 +419,29 @@ def convolve_harmonic(
     cdef so3_parameters_t h_parameters_struct=create_parameter_struct(h_parameters)
 
     hlmn_length = flmn_size(h_parameters)
-    hlmn = np.zeros([hlmn_length,], dtype=complex)
+    cdef double complex[::1] hlmn = np.zeros([hlmn_length,], dtype=complex)
 
     so3_conv_convolution(
-        <double complex *> np.PyArray_DATA(hlmn),
+        &hlmn[0],
         &h_parameters_struct,
-        <const double complex *> np.PyArray_DATA(flmn),
+        &flmn[0],
         &f_parameters_struct,
-        <const double complex *> np.PyArray_DATA(glmn),
+        &glmn[0],
         &g_parameters_struct
     )
     return hlmn, h_parameters
 
 def s2toso3_harmonic_convolution(
     h_so3_parameters,
-    np.ndarray[ double complex, ndim=1, mode="c"] flm not None,
-    np.ndarray[ double complex, ndim=1, mode="c"] glm not None):
+    double complex[::1] flm not None,
+    double complex[::1] glm not None):
 
     cdef so3_parameters_t h_parameters=create_parameter_struct(h_so3_parameters)
 
     hlmn_length = flmn_size(h_so3_parameters)
-    hlmn = np.zeros([hlmn_length,], dtype=complex)
+    cdef double complex[::1] hlmn = np.zeros([hlmn_length,], dtype=complex)
 
-    so3_conv_s2toso3_harmonic_convolution(
-    <double complex *> np.PyArray_DATA(hlmn), 
-    &h_parameters,
-    <const double complex *> np.PyArray_DATA(flm),
-    <const double complex *> np.PyArray_DATA(glm)
-    )
+    so3_conv_s2toso3_harmonic_convolution(&hlmn[0], &h_parameters, &flm[0], &glm[0])
     return hlmn
 
 def test_func():

tests/test_pyso3.py

@@ -1,18 +1,20 @@
 import numpy as np
+import pytest
 from pytest import approx
 
 import so3
 
 
-def test_pyso3():
-
-    f_p = so3.create_parameter_dict(16, 8)
+@pytest.mark.parametrize("L, N, L0", [(16, 8, 0), (16, 8, 1), (4, 1, 0), (3, 2, 0)])
+@pytest.mark.parametrize("reality", [False, True])
+def test_pyso3(L, N, L0, reality):
+    f_p = so3.create_parameter_dict(L, N, L0, reality=reality * 1)
     f_length = so3.f_size(f_p)
     rng = np.random.default_rng()
-    f_before = rng.normal(size=(f_length, 2)) @ [1, 1j]
+    multiplier = [1, 1] if reality else [1, 1j]
+    f_before = rng.normal(size=(f_length, 2)) @ multiplier
     flmn = so3.forward(f_before, f_p)
     f_before = so3.inverse(flmn, f_p)
-
     flmn = so3.forward(f_before, f_p)
 
     f = so3.inverse(flmn, f_p)