Array API compatibility

Author: wahln

examples/pencilbeam_proton.py

@@ -23,10 +23,14 @@
     fluence_optimization,
     plot_slice,
     load_tg119,
+    xp_utils,
 )
 
 from pyRadPlan.optimization.objectives import SquaredDeviation, SquaredOverdosing, MeanDose
 
+
+xp_utils.PREFER_GPU = True
+xp_utils.PREFERRED_CPU_ARRAY_BACKEND = "numpy"
 logging.basicConfig(level=logging.INFO)
 
 # %%
@@ -39,6 +43,7 @@
 # Create a plan object
 pln = IonPlan(radiation_mode="protons", machine="Generic")
 pln.prop_opt = {"solver": "scipy"}
+pln.prop_dose_calc = {"dose_grid": ct.grid}
 
 # Generate Steering Geometry ("stf")
 stf = generate_stf(ct, cst, pln)

examples/raytracer_rad_depth.py

@@ -28,6 +28,11 @@
 from pyRadPlan.visualization import plot_slice
 from pyRadPlan.io import load_patient
 
+from pyRadPlan import xp_utils
+
+xp_utils.PREFERRED_GPU_ARRAY_BACKEND = "cupy"
+xp_utils.PREFER_GPU = True
+
 # Configure the Logger to show you debug information
 logging.basicConfig(level=logging.INFO)
 logging.getLogger("pyRadPlan").setLevel(logging.DEBUG)
@@ -63,7 +68,7 @@
 # with a single beam at 0 degrees
 stfgen = StfGeneratorIMPT(pln)
 stfgen.bixel_width = 5.0
-stfgen.gantry_angles = [90.0]
+stfgen.gantry_angles = [0.0]
 
 # We generate the beam geometry on the CT and CST
 stf = stfgen.generate(ct, cst)
@@ -72,6 +77,8 @@
 # For that, we use a default HU->rSP table to convert our CT to water-equivalent thickness and then
 # call a voxel-wise RayTracing algorithm (proposed by Siddon) to calculate the radiological depth.
 rt = RayTracerSiddon([ct.compute_wet(default_hlut())])
+rt.lateral_cut_off = 150.0
+rt.precision = np.float32
 rt.debug_core_performance = True
 rad_depth_cubes = rt.trace_cubes(stf[0])
 

matRad

@@ -37,6 +37,7 @@
 from .analysis._dvh import DVH, DVHCollection
 from .visualization import plot_slice
 from .io import load_patient, load_tg119
+from .core import xp_utils
 
 try:
     __version__ = version(__name__)
@@ -49,6 +50,7 @@
 
 __all__ = [
     "__version__",
+    "xp_utils",
     "Plan",
     "IonPlan",
     "PhotonPlan",

pyRadPlan/__init__.py

@@ -10,6 +10,8 @@
 from pydantic.alias_generators import to_camel
 from copy import deepcopy
 
+import array_api_compat
+
 
 class PyRadPlanBaseModel(BaseModel):
     """
@@ -45,7 +47,7 @@ def __eq__(self, other: Any) -> bool:
         """
         try:
             return super().__eq__(other)
-        except ValueError:
+        except (ValueError, TypeError):
             if self.__dict__.keys() != other.__dict__.keys():
                 return False
             stack = [(self.__dict__, other.__dict__)]
@@ -56,12 +58,7 @@ def __eq__(self, other: Any) -> bool:
                 for key in dict_a:
                     if isinstance(dict_a[key], dict) and isinstance(dict_b[key], dict):
                         stack.append((dict_a[key], dict_b[key]))
-                    elif isinstance(dict_a[key], np.ndarray) and isinstance(
-                        dict_b[key], np.ndarray
-                    ):
-                        if not np.array_equal(dict_a[key], dict_b[key]):
-                            return False
-                    elif dict_a[key] != dict_b[key]:
+                    elif not self._eq_dict_entry(dict_a, dict_b, key):
                         return False
             return True
 
@@ -76,6 +73,55 @@ def __ne__(self, other: Any) -> bool:
         else:
             return True
 
+    def _eq_dict_entry(self, dict_a: dict, dict_b: dict, key: Any) -> bool:
+        """Compare two dictionary entries for equality."""
+
+        if (
+            isinstance(dict_a[key], np.ndarray)
+            and isinstance(dict_b[key], np.ndarray)
+            and dict_a[key].dtype == object
+            and dict_b[key].dtype == object
+        ):
+            return self._eq_object_arrays(dict_a[key], dict_b[key])
+        elif array_api_compat.is_array_api_obj(dict_a[key]) and array_api_compat.is_array_api_obj(
+            dict_b[key]
+        ):
+            try:
+                xp = array_api_compat.array_namespace(dict_a[key], dict_b[key])
+            except (ValueError, TypeError):
+                return False
+            if dict_a[key].shape != dict_b[key].shape or not xp.all(dict_a[key] == dict_b[key]):
+                return False
+        elif dict_a[key] != dict_b[key]:
+            return False
+
+        return True
+
+    def _eq_object_arrays(self, obj_array_a: np.ndarray, obj_array_b: np.ndarray) -> bool:
+        """Compare two object arrays for equality."""
+        assert obj_array_a.dtype == object and obj_array_b.dtype == object
+
+        if obj_array_a.shape != obj_array_b.shape:
+            return False
+
+        for a, b in zip(obj_array_a.flat, obj_array_b.flat):
+            if a is None and b is None:
+                continue
+            elif isinstance(a, np.ndarray) and isinstance(b, np.ndarray):
+                if a.shape != b.shape or a.dtype != b.dtype or not np.all(a == b):
+                    return False
+            elif array_api_compat.is_array_api_obj(a) and array_api_compat.is_array_api_obj(b):
+                try:
+                    xp = array_api_compat.array_namespace(a, b)
+                except (ValueError, TypeError):
+                    return False
+                if not xp.all(a == b):
+                    return False
+            elif any(a != b):
+                return False
+
+        return True
+
     def to_matrad(self, context: Union[str, dict] = "mat-file") -> Any:
         """
         Perform matRad compatible serialization.

pyRadPlan/core/datamodel.py

@@ -3,10 +3,16 @@
 from typing import Literal
 import SimpleITK as sitk
 import numpy as np
+import array_api_compat
+from array_api_compat import numpy as xnp
 from ._grids import Grid
 
+from numpy.typing import NDArray
+from ..core.xp_utils.typing import Array
+from ..core.xp_utils import to_numpy, from_numpy
 
-def sitk_mask_to_linear_indices(mask: sitk.Image, order="sitk") -> np.ndarray:
+
+def sitk_mask_to_linear_indices(mask: sitk.Image, order="sitk") -> NDArray:
     """
     Convert a SimpleITK mask to linear indices.
 
@@ -20,7 +26,7 @@ def sitk_mask_to_linear_indices(mask: sitk.Image, order="sitk") -> np.ndarray:
 
     Returns
     -------
-    np.ndarray
+    NDArray
         A 1D numpy array of linear indices where the mask is non-zero.
 
     Raises
@@ -37,16 +43,14 @@ def sitk_mask_to_linear_indices(mask: sitk.Image, order="sitk") -> np.ndarray:
     raise ValueError("Invalid ordering. Must be 'sitk' or 'numpy'.")
 
 
-def linear_indices_to_sitk_mask(
-    indices: np.ndarray, ref_image: sitk.Image, order="sitk"
-) -> sitk.Image:
+def linear_indices_to_sitk_mask(indices: Array, ref_image: sitk.Image, order="sitk") -> sitk.Image:
     """
     Convert linear indices to a SimpleITK mask.
 
     Parameters
     ----------
-    indices : np.ndarray
-        A 1D numpy array of linear indices where the mask is non-zero.
+    indices : Array
+        A 1D Array API conform array of linear indices where the mask is non-zero.
     ref_image : sitk.Image
         The reference image on which the mask is defined.
     order : str, optional
@@ -64,7 +68,9 @@ def linear_indices_to_sitk_mask(
         If the ordering is not 'sitk' or 'numpy'.
     """
 
-    arr = np.zeros_like(sitk.GetArrayViewFromImage(ref_image), dtype=np.uint8)
+    indices = to_numpy(indices)
+
+    arr: NDArray = xnp.zeros_like(sitk.GetArrayViewFromImage(ref_image), dtype=xnp.uint8)
 
     if order == "sitk":
         arr.T.flat[indices] = 1
@@ -80,18 +86,18 @@ def linear_indices_to_sitk_mask(
 
 
 def linear_indices_to_grid_coordinates(
-    indices: np.ndarray,
+    indices: Array,
     grid: Grid,
     index_type: Literal["numpy", "sitk"] = "numpy",
     dtype: np.dtype = np.float64,
-) -> np.ndarray:
+) -> Array:
     """
     Convert linear indices to gridcoordinates.
 
     Parameters
     ----------
-    indices : np.ndarray
-        A 1D numpy array of linear indices where the mask is non-zero.
+    indices : Array
+        A 1D Array API conform array of linear indices where the mask is non-zero.
     grid : Grid
         The image grid on which the indices lie.
     index_type : Literal["numpy", "sitk"], optional
@@ -102,10 +108,12 @@ def linear_indices_to_grid_coordinates(
 
     Returns
     -------
-    np.ndarray
-        A 2D numpy array of image coordinates.
+    Array
+        A 2D Array API conform array array of image coordinates.
     """
 
+    xp = array_api_compat.array_namespace(indices)
+
     # this is a manual reimplementation of np.unravel_index
     # to avoid the overhead of creating a tuple of arrays
     if index_type == "numpy":
@@ -117,6 +125,8 @@ def linear_indices_to_grid_coordinates(
     else:
         raise ValueError("Invalid index type. Must be 'numpy' or 'sitk'.")
 
+    indices = to_numpy(indices)
+
     v = np.empty((3, np.asarray(indices).size), dtype=dtype)
     tmp, v[order[0]] = np.divmod(indices, d2)
     v[order[2]], v[order[1]] = np.divmod(tmp, d1)
@@ -128,22 +138,22 @@ def linear_indices_to_grid_coordinates(
 
     physical_point = origin + np.matmul(np.matmul(grid.direction, spacing_diag), v).T
 
-    return physical_point
+    return from_numpy(xp, physical_point)
 
 
 def linear_indices_to_image_coordinates(
-    indices: np.ndarray,
+    indices: Array,
     image: sitk.Image,
     index_type: Literal["numpy", "sitk"] = "numpy",
     dtype: np.dtype = np.float64,
-) -> np.ndarray:
+) -> Array:
     """
     Convert linear indices to image coordinates.
 
     Parameters
     ----------
-    indices : np.ndarray
-        A 1D numpy array of linear indices where the mask is non-zero.
+    indices : Array
+        A 1D Array API conform array of linear indices where the mask is non-zero.
     image : sitk.Image
         The reference image on which the mask is defined.
     index_type : Literal["numpy", "sitk"], optional
@@ -154,8 +164,8 @@ def linear_indices_to_image_coordinates(
 
     Returns
     -------
-    np.ndarray
-        A 2D numpy array of image coordinates.
+    Array
+        A 2D Array API conform array of image coordinates.
     """
 
     grid = Grid.from_sitk_image(image)

pyRadPlan/core/np2sitk.py

@@ -0,0 +1,90 @@
+"""Provide information about the available compute backends."""
+
+import importlib
+
+from typing import Optional
+
+try:
+    import cupy as cp
+except ImportError:
+    cp = None
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+from .helpers import (
+    get_current_stream,
+    create_stream,
+    synchronize,
+    record_event,
+    elapsed_time,
+    to_numpy,
+    from_numpy,
+    to_namespace,
+)
+
+from .typing import Array, ArrayNamespace
+
+
+# Check if GPU is available
+def cupy_available() -> bool:
+    """Check if CuPy is available and a compatible GPU is present."""
+    return cp is not None and cp.cuda.is_available()
+
+
+def pytorch_available() -> bool:
+    """Check if PyTorch is available."""
+    return torch is not None
+
+
+def pytorch_gpu_available() -> bool:
+    """Check if PyTorch is available and a compatible GPU is present."""
+    return torch is not None and torch.cuda.is_available()
+
+
+PREFERRED_CPU_ARRAY_BACKEND: str = "numpy"
+PREFERRED_GPU_ARRAY_BACKEND: str = "cupy" if cupy_available() else None
+PREFER_GPU = True
+
+
+def choose_array_api_namespace(namespace: Optional[str] = None) -> ArrayNamespace:
+    """
+    Get the name of the preferred Array API conform computational namespace / backend.
+
+    Parameters
+    ----------
+    namespace : Optional[str], optional
+        The name of the desired backend. If None, the preferred backend is used.
+    """
+    if namespace is None:
+        if PREFER_GPU and PREFERRED_GPU_ARRAY_BACKEND is not None:
+            namespace = PREFERRED_GPU_ARRAY_BACKEND
+        else:
+            namespace = PREFERRED_CPU_ARRAY_BACKEND
+    else:
+        namespace = namespace.lower()
+
+    try:
+        return importlib.import_module(f"array_api_compat.{namespace}")
+    except ModuleNotFoundError:
+        return importlib.import_module(namespace)
+
+
+__all__ = [
+    "cp",
+    "torch",
+    "cupy_available",
+    "torch_gpu_available",
+    "get_current_stream",
+    "create_stream",
+    "record_event",
+    "synchronize",
+    "elapsed_time",
+    "to_numpy",
+    "from_numpy",
+    "to_namespace",
+    "Array",
+    "ArrayNamespace",
+]