moving all gpu related imports under try (#152)

* moving all gpu related imports under try * adding pypi ci * fixing classifiers * correcting the publisher name * further import changes * remove nvtx to check docs * removing nvtx instances to allow docs to build, removing function doubleundescore duplications
DiamondLightSource · Sep 2, 2024 · ce1c15f · ce1c15f
1 parent 3a4524d
commit ce1c15f
Show file tree

Hide file tree

Showing 13 changed files with 117 additions and 449 deletions.
diff --git a/.github/workflows/httomolibgpu_pypi_publish.yml b/.github/workflows/httomolibgpu_pypi_publish.yml
@@ -0,0 +1,31 @@
+name: Upload Python Package to PyPI when Release is Created
+
+on:
+  release:
+    types: [created]
+
+jobs:
+  pypi-publish:
+    name: Publish release to PyPI
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/httomolibgpu
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install setuptools wheel
+          pip install --upgrade build
+      - name: Build package
+        run: |
+          python -m build
+      - name: Publish package distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
diff --git a/conda/environment.yml b/conda/environment.yml
@@ -4,7 +4,7 @@ channels:
   - httomo
 dependencies:
   - conda-forge::cupy=12.3.0
-  - conda-forge::numpy<=1.25
+  - conda-forge::numpy
   - conda-forge::nvtx
   - conda-forge::scipy
   - conda-forge::python

diff --git a/httomolibgpu/cupywrapper.py b/httomolibgpu/cupywrapper.py
@@ -1,8 +1,6 @@
-cupy_run = False
 try:
     import cupy as cp
     import nvtx
-    import cupyx
 
     try:
         cp.cuda.Device(0).compute_capability
@@ -17,5 +15,6 @@
     from unittest.mock import Mock
     import numpy as cp
 
+    cupy_run = False
+
     nvtx = Mock()
-    cupyx = Mock()
diff --git a/httomolibgpu/misc/corr.py b/httomolibgpu/misc/corr.py
@@ -28,10 +28,16 @@
 from httomolibgpu import cupywrapper
 
 cp = cupywrapper.cp
-nvtx = cupywrapper.nvtx
+cupy_run = cupywrapper.cupy_run
 
 from numpy import float32
-from httomolibgpu.cuda_kernels import load_cuda_module
+from unittest.mock import Mock
+
+if cupy_run:
+    from httomolibgpu.cuda_kernels import load_cuda_module
+else:
+    load_cuda_module = Mock()
+
 
 __all__ = [
     "median_filter",
@@ -67,19 +73,6 @@ def median_filter(
     ValueError
         If the input array is not three dimensional.
     """
-    if cupywrapper.cupy_run:
-        return __median_filter(data, kernel_size, dif)
-    else:
-        print("median_filter won't be executed because CuPy is not installed")
-        return data
-
-
-@nvtx.annotate()
-def __median_filter(
-    data: cp.ndarray,
-    kernel_size: int = 3,
-    dif: float = 0.0,
-) -> cp.ndarray:
     input_type = data.dtype
 
     if input_type not in ["float32", "uint16"]:
@@ -148,8 +141,4 @@ def remove_outlier(
     if dif <= 0.0:
         raise ValueError("Threshold value (dif) must be positive and nonzero.")
 
-    if cupywrapper.cupy_run:
-        return __median_filter(data, kernel_size, dif)
-    else:
-        print("remove_outlier won't be executed because CuPy is not installed")
-        return data
+    return median_filter(data, kernel_size, dif)
diff --git a/httomolibgpu/misc/morph.py b/httomolibgpu/misc/morph.py
@@ -24,10 +24,14 @@
 from httomolibgpu import cupywrapper
 
 cp = cupywrapper.cp
-nvtx = cupywrapper.nvtx
-cupyx = cupywrapper.cupyx
+cupy_run = cupywrapper.cupy_run
 
-from cupyx.scipy.interpolate import interpn
+from unittest.mock import Mock
+
+if cupy_run:
+    from cupyx.scipy.interpolate import interpn
+else:
+    interpn = Mock()
 
 from typing import Literal
 
@@ -59,17 +63,6 @@ def sino_360_to_180(
     cp.ndarray
         Output 3D data.
     """
-    if cupywrapper.cupy_run:
-        return __sino_360_to_180(data, overlap, rotation)
-    else:
-        print("sino_360_to_180 won't be executed because CuPy is not installed")
-        return data
-
-
-@nvtx.annotate()
-def __sino_360_to_180(
-    data: cp.ndarray, overlap: int = 0, rotation: Literal["left", "right"] = "left"
-) -> cp.ndarray:
     if data.ndim != 3:
         raise ValueError("only 3D data is supported")
 
@@ -80,8 +73,8 @@ def __sino_360_to_180(
         raise ValueError("overlap must be less than data.shape[2]")
     if overlap < 0:
         raise ValueError("only positive overlaps are allowed.")
-    
-    if rotation not in ['left', 'right']:
+
+    if rotation not in ["left", "right"]:
         raise ValueError('rotation parameter must be either "left" or "right"')
 
     n = dx // 2
@@ -128,18 +121,6 @@ def data_resampler(
     Returns:
         cp.ndarray: Up/Down-scaled 3D cupy array
     """
-    if cupywrapper.cupy_run:
-        return __data_resampler(data, newshape, axis, interpolation)
-    else:
-        print("data_resampler won't be executed because CuPy is not installed")
-        return data
-
-
-@nvtx.annotate()
-def __data_resampler(
-    data: cp.ndarray, newshape: list, axis: int = 1, interpolation: str = "linear"
-) -> cp.ndarray:    
-
     if data.ndim != 3:
         raise ValueError("only 3D data is supported")
 

diff --git a/httomolibgpu/misc/rescale.py b/httomolibgpu/misc/rescale.py
@@ -24,7 +24,6 @@
 
 cp = cupywrapper.cp
 
-nvtx = cupywrapper.nvtx
 from typing import Literal, Optional, Tuple, Union
 
 __all__ = [
@@ -70,53 +69,6 @@ def rescale_to_int(
         The original data, clipped to the range specified with the perc_range_min and
         perc_range_max, and scaled to the full range of the output integer type
     """
-    if cupywrapper.cupy_run:
-        return __rescale_to_int(data, perc_range_min, perc_range_max, bits, glob_stats)
-    else:
-        print("rescale_to_int won't be executed because CuPy is not installed")
-        return data
-
-
-@nvtx.annotate()
-def __rescale_to_int(
-    data: cp.ndarray,
-    perc_range_min: float = 0.0,
-    perc_range_max: float = 100.0,
-    bits: Literal[8, 16, 32] = 8,
-    glob_stats: Optional[Tuple[float, float, float, int]] = None,
-):
-    """
-    Rescales the data and converts it fit into the range of an unsigned integer type
-    with the given number of bits.
-
-    Parameters
-    ----------
-    data : cp.ndarray
-        Required input data array, on GPU
-    perc_range_min: float, optional
-        The lower cutoff point in the input data, in percent of the data range (defaults to 0).
-        The lower bound is computed as min + perc_range_min/100*(max-min)
-    perc_range_max: float, optional
-        The upper cutoff point in the input data, in percent of the data range (defaults to 100).
-        The upper bound is computed as min + perc_range_max/100*(max-min)
-    bits: Literal[8, 16, 32], optional
-        The number of bits in the output integer range (defaults to 8).
-        Allowed values are:
-        - 8 -> uint8
-        - 16 -> uint16
-        - 32 -> uint32
-    glob_stats: tuple, optional
-        Global statistics of the full dataset (beyond the data passed into this call).
-        It's a tuple with (min, max, sum, num_items). If not given, the min/max is
-        computed from the given data.
-
-    Returns
-    -------
-    cp.ndarray
-        The original data, clipped to the range specified with the perc_range_min and
-        perc_range_max, and scaled to the full range of the output integer type
-    """
-
     if bits == 8:
         output_dtype: Union[type[np.uint8], type[np.uint16], type[np.uint32]] = np.uint8
     elif bits == 16:

diff --git a/httomolibgpu/prep/alignment.py b/httomolibgpu/prep/alignment.py
@@ -24,10 +24,14 @@
 from httomolibgpu import cupywrapper
 
 cp = cupywrapper.cp
-nvtx = cupywrapper.nvtx
-cupyx = cupywrapper.cupyx
+cupy_run = cupywrapper.cupy_run
 
-from cupyx.scipy.ndimage import map_coordinates
+from unittest.mock import Mock
+
+if cupy_run:
+    from cupyx.scipy.ndimage import map_coordinates
+else:
+    map_coordinates = Mock()
 
 from typing import Dict, List
 
@@ -77,55 +81,6 @@ def distortion_correction_proj_discorpy(
     cp.ndarray
         3D array. Distortion-corrected array.
     """
-    if cupywrapper.cupy_run:
-        return __distortion_correction_proj_discorpy(
-            data, metadata_path, preview, order, mode
-        )
-    else:
-        print(
-            "distortion_correction_proj_discorpy won't be executed because CuPy is not installed"
-        )
-        return data
-
-
-@nvtx.annotate()
-def __distortion_correction_proj_discorpy(
-    data: cp.ndarray,
-    metadata_path: str,
-    preview: Dict[str, List[int]],
-    order: int = 1,
-    mode: str = "reflect",
-):
-    """Unwarp a stack of images using a backward model.
-
-    Parameters
-    ----------
-    data : cp.ndarray
-        3D array.
-
-    metadata_path : str
-        The path to the file containing the distortion coefficients for the
-        data.
-
-    preview : Dict[str, List[int]]
-        A dict containing three key-value pairs:
-        - a list containing the `start` value of each dimension
-        - a list containing the `stop` value of each dimension
-        - a list containing the `step` value of each dimension
-
-    order : int, optional.
-        The order of the spline interpolation.
-
-    mode : {'reflect', 'grid-mirror', 'constant', 'grid-constant', 'nearest',
-           'mirror', 'grid-wrap', 'wrap'}, optional
-        To determine how to handle image boundaries.
-
-    Returns
-    -------
-    cp.ndarray
-        3D array. Distortion-corrected image(s).
-    """   
-
     # Check if it's a stack of 2D images, or only a single 2D image
     if len(data.shape) == 2:
         data = cp.expand_dims(data, axis=0)

diff --git a/httomolibgpu/prep/normalize.py b/httomolibgpu/prep/normalize.py
@@ -24,9 +24,14 @@
 from httomolibgpu import cupywrapper
 
 cp = cupywrapper.cp
-nvtx = cupywrapper.nvtx
+cupy_run = cupywrapper.cupy_run
 
-from cupy import mean
+from unittest.mock import Mock
+
+if cupy_run:
+    from cupy import mean
+else:
+    mean = Mock()
 
 from numpy import float32
 from typing import Tuple
@@ -69,26 +74,6 @@ def normalize(
     cp.ndarray
         Normalised 3D tomographic data as a CuPy array.
     """
-    if cupywrapper.cupy_run:
-        return __normalize(
-            data, flats, darks, cutoff, minus_log, nonnegativity, remove_nans
-        )
-    else:
-        print("normalize won't be executed because CuPy is not installed")
-        return data
-
-
-@nvtx.annotate()
-def __normalize(
-    data: cp.ndarray,
-    flats: cp.ndarray,
-    darks: cp.ndarray,
-    cutoff: float = 10.0,
-    minus_log: bool = True,
-    nonnegativity: bool = False,
-    remove_nans: bool = True,
-) -> cp.ndarray:    
-
     _check_valid_input(data, flats, darks)
 
     dark0 = cp.empty(darks.shape[1:], dtype=float32)