modification to the supplementary function

Author: dkazanc

httomolibgpu/misc/supp_func.py

@@ -39,7 +39,7 @@ def _naninfs_check(
     Parameters
     ----------
     data : cp.ndarray
-        Input CuPy array either float32 or uint16 data type.
+        Input CuPy or Numpy array either float32 or uint16 data type.
     correction : bool
         If correction is enabled then Inf's and NaN's will be replaced by zeros.
     verbosity : bool
@@ -50,7 +50,7 @@ def _naninfs_check(
     Returns
     -------
     ndarray
-        Corrected (or not) CuPy array.
+        Uncorrected or corrected (nans and infs converted to zeros) input array.
     """
     if cupy_run:
         xp = cp.get_array_module(data)
@@ -82,7 +82,7 @@ def _zeros_check(
     Parameters
     ----------
     data : cp.ndarray
-        Input CuPy array either float32 or uint16 data type.
+        Input CuPy or Numpy array either float32 or uint16 data type.
     verbosity : bool
         If enabled, then the printing of the warning happens when data contains infs or nans
     percentage_threshold: float:
@@ -111,3 +111,41 @@ def _zeros_check(
             )
 
     return warning_zeros
+
+
+def data_checker(
+    data: cp.ndarray,
+    verbosity: bool = True,
+    method_name: Optional[str] = None,
+) -> bool:
+    """
+    Function that performs the variety of checks on input data, in some cases also correct the data and prints warnings.
+    Currently it checks for: the presence of infs and nans in data; the number of zero elements. 
+
+    Parameters
+    ----------
+    data : xp.ndarray
+        Input CuPy or Numpy array either float32 or uint16 data type.
+    verbosity : bool
+        If enabled, then the printing of the warning happens when data contains infs or nans
+    method_name : str, optional.
+        Method's name for which input data is tested.
+
+    Returns
+    -------
+    cp.ndarray
+        Returns corrected or not data array
+    """
+
+    data = _naninfs_check(
+        data, correction=True, verbosity=verbosity, method_name=method_name
+    )
+
+    _zeros_check(
+        data,
+        verbosity=verbosity,
+        percentage_threshold=50,
+        method_name=method_name,
+    )
+
+    return data

tests/test_misc/test_supp_func.py

@@ -4,6 +4,7 @@
 from httomolibgpu.misc.supp_func import (
     _naninfs_check,
     _zeros_check,
+    data_checker,
 )
 from numpy.testing import assert_allclose, assert_equal
 
@@ -140,3 +141,50 @@ def test_zeros_check3_numpy():
     warning_zeros = _zeros_check(_data_input.copy())
 
     assert warning_zeros == False
+
+
+def test_data_checker_numpy():
+    _data_input = np.ones(shape=(10, 10, 10), dtype=np.float32)
+    _data_input[1, 1, 1] = -np.inf
+    _data_input[1, 1, 2] = np.inf
+    _data_input[1, 1, 3] = np.nan
+
+    _data_output = data_checker(_data_input.copy())
+
+    assert_equal(
+        _data_output[1, 1, 1],
+        0,
+    )
+    assert_equal(
+        _data_output[1, 1, 2],
+        0,
+    )
+    assert_equal(
+        _data_output[1, 1, 3],
+        0,
+    )
+    assert _data_output.dtype == _data_input.dtype
+    assert _data_output.shape == (10, 10, 10)
+
+def test_data_checker():
+    _data_input = cp.ones(shape=(10, 10, 10), dtype=cp.float32)*100
+    _data_input[1, 1, 1] = -cp.inf
+    _data_input[1, 1, 2] = cp.inf
+    _data_input[1, 1, 3] = cp.nan
+
+    _data_output = data_checker(_data_input.copy())
+
+    assert_equal(
+        _data_output[1, 1, 1].get(),
+        0.0,
+    )
+    assert_equal(
+        _data_output[1, 1, 2].get(),
+        0.0,
+    )
+    assert_equal(
+        _data_output[1, 1, 3].get(),
+        0.0,
+    )
+    assert _data_output.dtype == _data_input.dtype
+    assert _data_output.shape == (10, 10, 10)