diff --git a/environment-dev-arm.yml b/environment-dev-arm.yml index d44ba17e..4b5c80d1 100644 --- a/environment-dev-arm.yml +++ b/environment-dev-arm.yml @@ -13,7 +13,7 @@ dependencies: - sympy - pymc>=5 - pytensor - - astra-toolbox + - astra-toolbox>=2.2.0 - matplotlib - ipython - pytest diff --git a/environment-dev-gpu.yml b/environment-dev-gpu.yml index 8436d228..cb19a0b5 100644 --- a/environment-dev-gpu.yml +++ b/environment-dev-gpu.yml @@ -10,6 +10,7 @@ dependencies: - scipy>=1.13.0 - cupy - sympy + - astra-toolbox>=2.3.0 - matplotlib - ipython - pytest diff --git a/environment-dev.yml b/environment-dev.yml index fdfb6130..f719e8d9 100644 --- a/environment-dev.yml +++ b/environment-dev.yml @@ -13,7 +13,7 @@ dependencies: - sympy - pymc>=5 - pytensor - - astra-toolbox + - astra-toolbox>=2.3.0 - matplotlib - ipython - pytest diff --git a/pylops/medical/ct.py b/pylops/medical/ct.py index d6952072..db31a8d4 100644 --- a/pylops/medical/ct.py +++ b/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) diff --git a/pytests/test_ct.py b/pytests/test_ct.py index 4ebf459c..052bb8e0 100755 --- a/pytests/test_ct.py +++ b/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) diff --git a/requirements-dev-gpu.txt b/requirements-dev-gpu.txt index 22ed24a9..6d011e53 100644 --- a/requirements-dev-gpu.txt +++ b/requirements-dev-gpu.txt @@ -4,6 +4,7 @@ cupy-cuda12x torch numba sympy +astra-toolbox>=2.3.0 matplotlib ipython pytest diff --git a/requirements-dev.txt b/requirements-dev.txt index 091c06a2..f10a4002 100644 --- a/requirements-dev.txt +++ b/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 diff --git a/requirements-doc.txt b/requirements-doc.txt index a445cec1..a1ec56d3 100644 --- a/requirements-doc.txt +++ b/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