Update ASTRA integration

environment-dev-arm.yml

@@ -13,7 +13,7 @@ dependencies:
   - sympy
   - pymc>=5
   - pytensor
-  - astra-toolbox
+  - astra-toolbox>=2.2.0
   - matplotlib
   - ipython
   - pytest

environment-dev-gpu.yml

@@ -10,6 +10,7 @@ dependencies:
   - scipy>=1.13.0
   - cupy
   - sympy
+  - astra-toolbox>=2.3.0
   - matplotlib
   - ipython
   - pytest

environment-dev.yml

@@ -13,7 +13,7 @@ dependencies:
   - sympy
   - pymc>=5
   - pytensor
-  - astra-toolbox
+  - astra-toolbox>=2.3.0
   - matplotlib
   - ipython
   - pytest

pylops/medical/ct.py

@@ -9,6 +9,7 @@
 
 from pylops import LinearOperator
 from pylops.utils import deps
+from pylops.utils.backend import get_array_module, get_module_name, to_numpy
 from pylops.utils.decorators import reshaped
 from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray
 
@@ -19,6 +20,7 @@
 
 if astra_message is None:
     import astra
+    import astra.experimental
 
 
 class CT2D(LinearOperator):
@@ -50,9 +52,15 @@ class CT2D(LinearOperator):
         Distance between origin and detector along the source-origin line
         (only for "proj_geom_type=fanflat")
     projector_type : :obj:`int`, optional
-        Type of projection geometry (``strip``, or ``line``, or ``linear``)
+        Type of projection kernel: ``strip`` (default), ``line``, or ``linear`` for
+        ``engine=cpu``, and ``cuda`` (i.e., hardware-accelerated ``linear``) for
+        ``engine=cuda``. For ``fanflat`` geometry, ``linear`` kernel is not supported.
+    engine : :obj:`str`, optional
+        Engine used for computation (``cpu`` or ``cuda``).
     dtype : :obj:`str`, optional
-        Type of elements in input array.
+        Type of elements in input array. Note that internally all operations will be
+        performed in float32 dtype because of ASTRA compatibility, and the output will
+        be converted to the requested dtype afterwards.
     name : :obj:`str`, optional
         Name of operator (to be used by :func:`pylops.utils.describe.describe`)
 
@@ -86,54 +94,145 @@ class CT2D(LinearOperator):
     def __init__(
         self,
         dims: InputDimsLike,
-        det_width: int,
-        det_count: float,
+        det_width: float,
+        det_count: int,
         thetas: NDArray,
-        proj_geom_type: Optional[str] = "parallel",
-        source_origin_dist: float = None,
-        origin_detector_dist: float = None,
-        projector_type: Optional[str] = "strip",
-        dtype: DTypeLike = "float64",
+        proj_geom_type: str = "parallel",
+        source_origin_dist: Optional[float] = None,
+        origin_detector_dist: Optional[float] = None,
+        projector_type: Optional[str] = None,
+        engine: str = "cpu",
+        dtype: DTypeLike = "float32",
         name: str = "C",
     ) -> None:
         if astra_message is not None:
             raise NotImplementedError(astra_message)
 
-        # create volume and projection geometries
-        self.vol_geom = astra.create_vol_geom(dims)
-        if proj_geom_type == "parallel":
+        self.dims = dims
+        self.det_width = det_width
+        self.det_count = det_count
+        self.thetas = to_numpy(thetas)  # ASTRA can only consume angles as a NumPy array
+        self.proj_geom_type = proj_geom_type
+        self.source_origin_dist = source_origin_dist
+        self.origin_detector_dist = origin_detector_dist
+
+        # make "strip" projector type default for cpu engine and only allow cuda otherwise
+        if engine == "cpu":
+            if projector_type is None:
+                projector_type = "strip"
+            # fanflat geometry projectors need an appropriate suffix (unless it's "cuda")
+            if projector_type == "cuda":
+                logging.warning("'cuda' projector type specified with 'cpu' engine.")
+            elif proj_geom_type == "fanflat":
+                projector_type += "_fanflat"
+        elif engine == "cuda":
+            if projector_type in [None, "cuda"]:
+                projector_type = "cuda"
+            else:
+                raise ValueError(
+                    "Only 'cuda' projector type is supported for 'cuda' engine."
+                )
+        else:
+            raise NotImplementedError("Engine must be 'cpu' or 'cuda'")
+        self.projector_type = projector_type
+
+        # create create volume and projection geometries as well as projector
+        self._init_geometries()
+        if engine == "cuda":
+            # efficient GPU data exchange only implemented for 3D data in ASTRA, so we
+            # emulate 2D geometry as 3D case with 1 slice
+            self._init_1_slice_3d_geometries()
+
+        super().__init__(
+            dtype=np.dtype(dtype), dims=dims, dimsd=(len(thetas), det_count), name=name
+        )
+
+    def _init_geometries(self):
+        self.vol_geom = astra.create_vol_geom(self.dims)
+        if self.proj_geom_type == "parallel":
             self.proj_geom = astra.create_proj_geom(
-                proj_geom_type, det_width, det_count, thetas
+                "parallel", self.det_width, self.det_count, self.thetas
             )
-        else:
+        elif self.proj_geom_type == "fanflat":
             self.proj_geom = astra.create_proj_geom(
-                proj_geom_type,
-                det_width,
-                det_count,
-                thetas,
-                source_origin_dist,
-                origin_detector_dist,
+                "fanflat",
+                self.det_width,
+                self.det_count,
+                self.thetas,
+                self.source_origin_dist,
+                self.origin_detector_dist,
             )
-
-        # create projector
-        self.proj_id = astra.create_projector(
-            projector_type, self.proj_geom, self.vol_geom
+        self.projector_id = astra.create_projector(
+            self.projector_type, self.proj_geom, self.vol_geom
         )
-        super().__init__(
-            dtype=np.dtype(dtype), dims=dims, dimsd=(len(thetas), det_count), name=name
+
+    def _init_1_slice_3d_geometries(self):
+        """Emulate 2D geometry as 3D to be able to use zero-copy GPU data exchange in ASTRA."""
+        self._3d_vol_geom = astra.create_vol_geom(*self.dims, 1)
+        if self.proj_geom_type == "parallel":
+            self._3d_proj_geom = astra.create_proj_geom(
+                "parallel3d", 1.0, self.det_width, 1, self.det_count, self.thetas
+            )
+        elif self.proj_geom_type == "fanflat":
+            self._3d_proj_geom = astra.create_proj_geom(
+                "cone",
+                1.0,
+                self.det_width,
+                1,
+                self.det_count,
+                self.thetas,
+                self.source_origin_dist,
+                self.origin_detector_dist,
+            )
+        self._3d_projector_id = astra.create_projector(
+            "cuda3d", self._3d_proj_geom, self._3d_vol_geom
         )
 
+    def _astra_project(self, x, mode):
+        """Perform forward/backward projection using ASTRA."""
+        ncp = get_array_module(x)
+        backend = get_module_name(ncp)
+        x_dtype = x.dtype
+        if x_dtype != ncp.float32:
+            logging.warning(
+                "CT2D operator received input that is not of float32 dtype. It will "
+                "be cast into float32 internally for ASTRA compatibility."
+            )
+            x = x.astype(ncp.float32)
+        if backend == "numpy":
+            if mode == "forward":
+                y_id, y = astra.create_sino(x, self.projector_id)
+            elif mode == "backward":
+                y_id, y = astra.create_backprojection(x, self.projector_id)
+            astra.data2d.delete(y_id)
+        else:
+            # Use zero-copy GPU data exchange, which is only implemented for contiguous
+            # 3D arrays in ASTRA
+            if not x.flags["C_CONTIGUOUS"]:
+                logging.warning(
+                    "CT2D operator received input that is not of contiguous. It will be "
+                    "cast into a contiguous array internally for ASTRA compatibility."
+                )
+                x = ncp.ascontiguousarray(x)
+            x = ncp.expand_dims(x, axis=0)
+            if mode == "forward":
+                y = ncp.empty_like(x, shape=astra.geom_size(self._3d_proj_geom))
+                astra.experimental.direct_FP3D(self._3d_projector_id, x, y)
+            elif mode == "backward":
+                y = ncp.empty_like(x, shape=astra.geom_size(self._3d_vol_geom))
+                astra.experimental.direct_BP3D(self._3d_projector_id, y, x)
+        return y.astype(x_dtype)
+
     @reshaped
     def _matvec(self, x):
-        y_id, y = astra.create_sino(x, self.proj_id)
-        astra.data2d.delete(y_id)
-        return y
+        return self._astra_project(x, mode="forward")
 
     @reshaped
     def _rmatvec(self, x):
-        y_id, y = astra.create_backprojection(x, self.proj_id)
-        astra.data2d.delete(y_id)
-        return y
+        return self._astra_project(x, mode="backward")
 
     def __del__(self):
-        astra.projector.delete(self.proj_id)
+        if hasattr(self, "projector_id"):
+            astra.projector.delete(self.projector_id)
+        if hasattr(self, "_3d_projector_id"):
+            astra.projector.delete(self._3d_projector_id)

pytests/test_ct.py

@@ -1,59 +1,73 @@
 import os
 import platform
 
-import numpy as np
+if int(os.environ.get("TEST_CUPY_PYLOPS", 0)):
+    import cupy as np
+
+    backend = "cupy"
+else:
+    import numpy as np
+
+    backend = "numpy"
 import pytest
 
 from pylops.medical import CT2D
 from pylops.utils import dottest
 
-par1 = {
-    "ny": 51,
-    "nx": 30,
-    "ntheta": 20,
-    "proj_geom_type": "parallel",
-    "projector_type": "strip",
-    "dtype": "float64",
-}  # parallel, strip
-
-par2 = {
-    "ny": 51,
-    "nx": 30,
-    "ntheta": 20,
-    "proj_geom_type": "parallel",
-    "projector_type": "line",
-    "dtype": "float64",
-}  # parallel, line
-
-par3 = {
-    "ny": 51,
-    "nx": 30,
-    "ntheta": 20,
-    "proj_geom_type": "fanflat",
-    "projector_type": "strip",
-    "dtype": "float64",
-}  # fanflat, strip
-
-
-@pytest.mark.skipif(
-    int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled"
-)
-@pytest.mark.skipif(platform.system() == "Darwin", reason="Not OSX enabled")
-@pytest.mark.parametrize("par", [(par1), (par2)])
-def test_CT2D(par):
-    """Dot-test for CT2D operator"""
-    theta = np.linspace(0.0, np.pi, par["ntheta"], endpoint=False)
 
-    Cop = CT2D(
-        (par["ny"], par["nx"]),
+@pytest.fixture
+def operator(request):
+    geometry_type = request.param
+    nx = 51
+    ny = 30
+    ntheta = 20
+    theta = np.linspace(0.0, np.pi, ntheta, endpoint=False)
+    source_origin_dist = 100 if geometry_type == "fanflat" else None
+    origin_detector_dist = 0 if geometry_type == "fanflat" else None
+    return CT2D(
+        (ny, nx),
         1.0,
-        par["ny"],
+        ny,
         theta,
-        proj_geom_type=par["proj_geom_type"],
-        projector_type=par["projector_type"],
+        geometry_type,
+        source_origin_dist,
+        origin_detector_dist,
+        engine="cpu" if backend == "numpy" else "cuda",
     )
-    assert dottest(
-        Cop,
-        par["ny"] * par["ntheta"],
-        par["ny"] * par["nx"],
+
+
+@pytest.fixture
+def x(operator):
+    return np.ones(operator.dims, dtype=np.float32)
+
+
+@pytest.fixture
+def y(operator):
+    return np.ones((len(operator.thetas), operator.det_count), dtype=np.float32)
+
+
+@pytest.mark.skipif(platform.system() == "Darwin", reason="Not OSX enabled")
+@pytest.mark.parametrize("operator", ["parallel", "fanflat"], indirect=True)
+class TestCT2D:
+    def test_basic(self, operator, x, y):
+        assert not np.allclose(operator @ x, 0.0)
+        assert not np.allclose(operator.T @ y, 0.0)
+
+    @pytest.mark.skipif(
+        backend == "cupy",
+        reason="CUDA tests are failing because of severely mismatched adjoint.",
     )
+    def test_adjointess(self, operator):
+        assert dottest(operator, rtol=1e-5)
+
+    def test_non_astra_native_dtype(self, operator, x, y):
+        x = x.astype(np.float64)
+        y = y.astype(np.float64)
+        assert not np.allclose(operator @ x, 0.0)
+        assert not np.allclose(operator.T @ y, 0.0)
+
+    def test_non_contiguous_input(self, operator, x, y):
+        x = x[:, ::-1]  # non-contiguous view
+        y = y[:, ::-1]
+        assert not np.allclose(operator @ x, 0.0)
+        assert not np.allclose(operator.T @ y, 0.0)

requirements-dev-gpu.txt

@@ -4,6 +4,7 @@ cupy-cuda12x
 torch
 numba
 sympy
+astra-toolbox>=2.3.0
 matplotlib
 ipython
 pytest

requirements-dev.txt

@@ -8,7 +8,7 @@ scikit-fmm
 sympy
 devito
 # dtcwt (until numpy>=2.0.0 is supported)
-astra-toolbox
+astra-toolbox>=2.2.0
 matplotlib
 ipython
 pytest

requirements-doc.txt

@@ -9,7 +9,7 @@ scikit-fmm
 sympy
 devito
 # dtcwt (until numpy>=2.0.0 is supported)
-astra-toolbox
+astra-toolbox>=2.2.0
 matplotlib
 ipython
 pytest