Add CuPy interface

docs/src/installation.rst

@@ -11,7 +11,7 @@ The Python package can be installed with pip by simply running
 
     pip install sphericart
 
-This basic package makes use of NumPy. Implementations supporting PyTorch and JAX can be installed with
+This basic package makes use of NumPy and CuPy. Implementations supporting PyTorch and JAX can be installed with
 
 .. code-block:: bash
 
@@ -37,9 +37,10 @@ Before installing the JAX version of ``sphericart``, make sure you already have
 library installed according to the `official JAX installation instructions
 <https://jax.readthedocs.io/en/latest/installation.html>`_.
 
-In addition, if you want to use the CUDA functionalities of sphericart (either with torch
-or JAX), make sure you have installed the CUDA toolkit and set up the environment variables
-``CUDA_HOME``, ``LD_LIBRARY_FLAGS``, and ``PATH`` accordingly.
+In addition, if you want to use the CUDA functionalities of sphericart (with CuPy,
+torch or JAX), make sure you have installed the CUDA toolkit
+and set up the environment variables ``CUDA_HOME``, ``LD_LIBRARY_FLAGS``, and ``PATH``
+accordingly. In case you want to use CuPy, it should be installed separately.
 
 
 Julia package

docs/src/python-api.rst

@@ -1,6 +1,8 @@
 Python API
 ==========
 
+The Python calculators accept either NumPy arrays on CPU or CuPy arrays on CUDA.
+
 .. autoclass:: sphericart.SphericalHarmonics
     :members:
 

docs/src/python-examples.rst

@@ -8,3 +8,8 @@ difference that, in line with pythonic mores, the
 
 .. literalinclude:: ../../examples/python/example.py
     :language: python
+
+The same calculators also accept CuPy arrays:
+
+.. literalinclude:: ../../examples/python/cupy.py
+    :language: python

examples/python/cupy.py

@@ -0,0 +1,39 @@
+import argparse
+
+import cupy as cp
+import numpy as np
+
+import sphericart
+
+docstring = """
+An example of the CuPy interface of the `sphericart` library.
+
+Computes Cartesian spherical harmonics on the GPU for a random array of 3D
+points and compares the result against the NumPy CPU backend.
+"""
+
+
+def sphericart_cupy_example(l_max=10, n_samples=10000):
+    xyz_cpu = np.random.rand(n_samples, 3)
+    xyz_gpu = cp.asarray(xyz_cpu)
+
+    sh_calculator = sphericart.SphericalHarmonics(l_max)
+
+    sh_cpu = sh_calculator.compute(xyz_cpu)
+    sh_gpu = sh_calculator.compute(xyz_gpu)
+
+    print(
+        "CPU vs GPU relative error: %12.8e"
+        % (np.linalg.norm(sh_cpu - cp.asnumpy(sh_gpu)) / np.linalg.norm(sh_cpu))
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=docstring)
+
+    parser.add_argument("-l", type=int, default=10, help="maximum angular momentum")
+    parser.add_argument("-s", type=int, default=1000, help="number of samples")
+
+    args = parser.parse_args()
+
+    sphericart_cupy_example(args.l, args.s)

python/src/sphericart/_c_lib.py

@@ -22,6 +22,22 @@ class sphericart_solid_harmonics_calculator_f_t(ctypes.c_void_p):
     pass
 
 
+class sphericart_cuda_spherical_harmonics_calculator_t(ctypes.c_void_p):
+    pass
+
+
+class sphericart_cuda_spherical_harmonics_calculator_f_t(ctypes.c_void_p):
+    pass
+
+
+class sphericart_cuda_solid_harmonics_calculator_t(ctypes.c_void_p):
+    pass
+
+
+class sphericart_cuda_solid_harmonics_calculator_f_t(ctypes.c_void_p):
+    pass
+
+
 def setup_functions(lib):
     lib.sphericart_spherical_harmonics_new.restype = (
         sphericart_spherical_harmonics_calculator_t
@@ -221,6 +237,168 @@ def setup_functions(lib):
         sphericart_solid_harmonics_calculator_f_t,
     ]
 
+    lib.sphericart_cuda_spherical_harmonics_new.restype = (
+        sphericart_cuda_spherical_harmonics_calculator_t
+    )
+    lib.sphericart_cuda_spherical_harmonics_new.argtypes = [ctypes.c_size_t]
+
+    lib.sphericart_cuda_spherical_harmonics_new_f.restype = (
+        sphericart_cuda_spherical_harmonics_calculator_f_t
+    )
+    lib.sphericart_cuda_spherical_harmonics_new_f.argtypes = [ctypes.c_size_t]
+
+    lib.sphericart_cuda_spherical_harmonics_delete.restype = None
+    lib.sphericart_cuda_spherical_harmonics_delete.argtypes = [
+        sphericart_cuda_spherical_harmonics_calculator_t
+    ]
+
+    lib.sphericart_cuda_spherical_harmonics_delete_f.restype = None
+    lib.sphericart_cuda_spherical_harmonics_delete_f.argtypes = [
+        sphericart_cuda_spherical_harmonics_calculator_f_t
+    ]
+
+    lib.sphericart_cuda_spherical_harmonics_compute_array.restype = None
+    lib.sphericart_cuda_spherical_harmonics_compute_array.argtypes = [
+        sphericart_cuda_spherical_harmonics_calculator_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_spherical_harmonics_compute_array_f.restype = None
+    lib.sphericart_cuda_spherical_harmonics_compute_array_f.argtypes = [
+        sphericart_cuda_spherical_harmonics_calculator_f_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_spherical_harmonics_compute_array_with_gradients.restype = None
+    lib.sphericart_cuda_spherical_harmonics_compute_array_with_gradients.argtypes = [
+        sphericart_cuda_spherical_harmonics_calculator_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_spherical_harmonics_compute_array_with_gradients_f.restype = (
+        None
+    )
+    lib.sphericart_cuda_spherical_harmonics_compute_array_with_gradients_f.argtypes = [
+        sphericart_cuda_spherical_harmonics_calculator_f_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_spherical_harmonics_compute_array_with_hessians.restype = None
+    lib.sphericart_cuda_spherical_harmonics_compute_array_with_hessians.argtypes = [
+        sphericart_cuda_spherical_harmonics_calculator_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_spherical_harmonics_compute_array_with_hessians_f.restype = None
+    lib.sphericart_cuda_spherical_harmonics_compute_array_with_hessians_f.argtypes = [
+        sphericart_cuda_spherical_harmonics_calculator_f_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_solid_harmonics_new.restype = (
+        sphericart_cuda_solid_harmonics_calculator_t
+    )
+    lib.sphericart_cuda_solid_harmonics_new.argtypes = [ctypes.c_size_t]
+
+    lib.sphericart_cuda_solid_harmonics_new_f.restype = (
+        sphericart_cuda_solid_harmonics_calculator_f_t
+    )
+    lib.sphericart_cuda_solid_harmonics_new_f.argtypes = [ctypes.c_size_t]
+
+    lib.sphericart_cuda_solid_harmonics_delete.restype = None
+    lib.sphericart_cuda_solid_harmonics_delete.argtypes = [
+        sphericart_cuda_solid_harmonics_calculator_t
+    ]
+
+    lib.sphericart_cuda_solid_harmonics_delete_f.restype = None
+    lib.sphericart_cuda_solid_harmonics_delete_f.argtypes = [
+        sphericart_cuda_solid_harmonics_calculator_f_t
+    ]
+
+    lib.sphericart_cuda_solid_harmonics_compute_array.restype = None
+    lib.sphericart_cuda_solid_harmonics_compute_array.argtypes = [
+        sphericart_cuda_solid_harmonics_calculator_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_solid_harmonics_compute_array_f.restype = None
+    lib.sphericart_cuda_solid_harmonics_compute_array_f.argtypes = [
+        sphericart_cuda_solid_harmonics_calculator_f_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_solid_harmonics_compute_array_with_gradients.restype = None
+    lib.sphericart_cuda_solid_harmonics_compute_array_with_gradients.argtypes = [
+        sphericart_cuda_solid_harmonics_calculator_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_solid_harmonics_compute_array_with_gradients_f.restype = None
+    lib.sphericart_cuda_solid_harmonics_compute_array_with_gradients_f.argtypes = [
+        sphericart_cuda_solid_harmonics_calculator_f_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_solid_harmonics_compute_array_with_hessians.restype = None
+    lib.sphericart_cuda_solid_harmonics_compute_array_with_hessians.argtypes = [
+        sphericart_cuda_solid_harmonics_calculator_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.POINTER(ctypes.c_double),
+        ctypes.c_void_p,
+    ]
+
+    lib.sphericart_cuda_solid_harmonics_compute_array_with_hessians_f.restype = None
+    lib.sphericart_cuda_solid_harmonics_compute_array_with_hessians_f.argtypes = [
+        sphericart_cuda_solid_harmonics_calculator_f_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_size_t,
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.POINTER(ctypes.c_float),
+        ctypes.c_void_p,
+    ]
+
 
 class LibraryFinder(object):
     def __init__(self):

python/src/sphericart/_dispatch.py

@@ -0,0 +1,60 @@
+import ctypes
+
+import numpy as np
+
+
+try:
+    import cupy as cp
+    from cupy import ndarray as cupy_ndarray
+except ImportError:
+    cp = None
+
+    class cupy_ndarray:  # type: ignore[no-redef]
+        pass
+
+
+def is_array(array):
+    return isinstance(array, np.ndarray) or isinstance(array, cupy_ndarray)
+
+
+def is_cupy_array(array):
+    return isinstance(array, cupy_ndarray)
+
+
+def make_contiguous(array):
+    if isinstance(array, np.ndarray):
+        return np.ascontiguousarray(array)
+    if isinstance(array, cupy_ndarray):
+        return cp.ascontiguousarray(array)
+    raise TypeError(f"only numpy and cupy arrays are supported, found {type(array)}")
+
+
+def empty_like(shape, array):
+    if isinstance(array, np.ndarray):
+        return np.empty(shape, dtype=array.dtype)
+    if isinstance(array, cupy_ndarray):
+        return cp.empty(shape, dtype=array.dtype)
+    raise TypeError(f"only numpy and cupy arrays are supported, found {type(array)}")
+
+
+def get_pointer(array):
+    if array.dtype == np.float32:
+        ptr_type = ctypes.POINTER(ctypes.c_float)
+    elif array.dtype == np.float64:
+        ptr_type = ctypes.POINTER(ctypes.c_double)
+    else:
+        raise TypeError(
+            f"only float32 and float64 arrays are supported, found {array.dtype}"
+        )
+
+    if isinstance(array, np.ndarray):
+        return array.ctypes.data_as(ptr_type)
+    if isinstance(array, cupy_ndarray):
+        return ctypes.cast(array.data.ptr, ptr_type)
+    raise TypeError(f"only numpy and cupy arrays are supported, found {type(array)}")
+
+
+def get_cuda_stream(array):
+    if not isinstance(array, cupy_ndarray):
+        return None
+    return ctypes.c_void_p(cp.cuda.get_current_stream().ptr)