Merge pull request #209 from DiamondLightSource/fbp2d_astra

adding fbp2d_astra algorithm and tests

Author: dkazanc

httomolibgpu/__init__.py

@@ -11,5 +11,5 @@
     remove_all_stripe,
 )
 
-from httomolibgpu.recon.algorithm import FBP, LPRec, SIRT, CGLS
+from httomolibgpu.recon.algorithm import FBP2d_astra, FBP, LPRec, SIRT, CGLS
 from httomolibgpu.recon.rotation import find_center_vo, find_center_360, find_center_pc

httomolibgpu/recon/algorithm.py

@@ -28,6 +28,8 @@
 
 from unittest.mock import Mock
 
+from tomobar.methodsDIR import RecToolsDIR
+
 if cupy_run:
     from tomobar.methodsDIR_CuPy import RecToolsDIRCuPy
     from tomobar.methodsIR_CuPy import RecToolsIRCuPy
@@ -40,15 +42,93 @@
 
 
 __all__ = [
+    "FBP2d_astra",
     "FBP",
-    "LPRec",
     "SIRT",
     "CGLS",
+    "LPRec",
 ]
 
 input_data_axis_labels = ["angles", "detY", "detX"]  # set the labels of the input data
 
 
+## %%%%%%%%%%%%%%%%%%%%%%% FBP2d_astra reconstruction %%%%%%%%%%%%%%%%%%%%%%%%%%%%  ##
+def FBP2d_astra(
+    data: np.ndarray,
+    angles: np.ndarray,
+    center: Optional[float] = None,
+    filter_type: str = "ram-lak",
+    filter_parameter: Optional[float] = None,
+    filter_d: Optional[float] = None,
+    recon_size: Optional[int] = None,
+    recon_mask_radius: float = 0.95,
+    neglog: bool = False,
+    gpu_id: int = 0,
+) -> np.ndarray:
+    """
+    Perform Filtered Backprojection (FBP) reconstruction slice-by-slice (2d) using ASTRA toolbox :cite:`van2016fast` and
+    ToMoBAR :cite:`kazantsev2020tomographic` wrappers.
+    This is a 2D recon using ASTRA's API for the FBP method, see for more parameters ASTRA's documentation here:
+    https://astra-toolbox.com/docs/algs/FBP_CUDA.html.
+
+    Parameters`
+    ----------
+    data : np.ndarray
+        Projection data as a 3d numpy array.
+    angles : np.ndarray
+        An array of angles given in radians.
+    center : float, optional
+        The center of rotation (CoR).
+    filter_type: str
+        Type of projection filter, see ASTRA's API for all available options for filters.
+    filter_parameter: float, optional
+        Parameter value for the 'tukey', 'gaussian', 'blackman' and 'kaiser' filter types.
+    filter_d: float, optional
+        D parameter value for 'shepp-logan', 'cosine', 'hamming' and 'hann' filter types.
+    recon_size : int, optional
+        The [recon_size, recon_size] shape of the reconstructed slice in pixels.
+        By default (None), the reconstructed size will be the dimension of the horizontal detector.
+    recon_mask_radius: float
+        The radius of the circular mask that applies to the reconstructed slice in order to crop
+        out some undesirable artifacts. The values outside the given diameter will be set to zero.
+        It is recommended to keep the value in the range [0.7-1.0].
+    neglog: bool
+        Take negative logarithm on input data to convert to attenuation coefficient or a density of the scanned object. Defaults to False,
+        assuming that the negative log is taken either in normalisation procedure on with Paganin filter application.
+    gpu_id : int
+        A GPU device index to perform operation on.
+
+    Returns
+    -------
+    np.ndarray
+        The FBP reconstructed volume as a numpy array.
+    """
+    data_shape = np.shape(data)
+    if recon_size is None:
+        recon_size = data_shape[2]
+
+    RecTools = _instantiate_direct_recon2d_class(
+        data, angles, center, recon_size, gpu_id
+    )
+
+    detY_size = data_shape[1]
+    reconstruction = np.empty(
+        (recon_size, detY_size, recon_size), dtype=np.float32(), order="C"
+    )
+    _take_neg_log_np(data) if neglog else data
+
+    # loop over detY slices
+    for slice_index in range(0, detY_size):
+        reconstruction[:, slice_index, :] = RecTools.FBP(
+            data[:, slice_index, :],
+            filter_type=filter_type,
+            filter_parameter=filter_parameter,
+            filter_d=filter_d,
+            recon_mask_radius=recon_mask_radius,
+        )
+    return reconstruction
+
+
 ## %%%%%%%%%%%%%%%%%%%%%%% FBP reconstruction %%%%%%%%%%%%%%%%%%%%%%%%%%%%  ##
 def FBP(
     data: cp.ndarray,
@@ -68,7 +148,7 @@ def FBP(
     Parameters
     ----------
     data : cp.ndarray
-        Projection data as a CuPy array.
+        Projection data as a 3d CuPy array.
     angles : np.ndarray
         An array of angles given in radians.
     center : float, optional
@@ -124,7 +204,7 @@ def LPRec(
     Parameters
     ----------
     data : cp.ndarray
-        Projection data as a CuPy array.
+        Projection data as a 3d CuPy array.
     angles : np.ndarray
         An array of angles given in radians.
     center : float, optional
@@ -315,6 +395,43 @@ def _instantiate_direct_recon_class(
     return RecToolsCP
 
 
+## %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%  ##
+def _instantiate_direct_recon2d_class(
+    data: np.ndarray,
+    angles: np.ndarray,
+    center: Optional[float] = None,
+    recon_size: Optional[int] = None,
+    gpu_id: int = 0,
+) -> Type:
+    """instantiate ToMoBAR's direct recon class for 2d reconstruction
+
+    Args:
+        data (cp.ndarray): data array
+        angles (np.ndarray): angles
+        center (Optional[float], optional): center of recon. Defaults to None.
+        recon_size (Optional[int], optional): recon_size. Defaults to None.
+        gpu_id (int, optional): gpu ID. Defaults to 0.
+
+    Returns:
+        Type[RecToolsDIR]: an instance of the direct recon class
+    """
+    if center is None:
+        center = data.shape[2] // 2  # making a crude guess
+    if recon_size is None:
+        recon_size = data.shape[2]
+    RecTools = RecToolsDIR(
+        DetectorsDimH=data.shape[2],  # Horizontal detector dimension
+        DetectorsDimV=None,  # 2d case
+        CenterRotOffset=data.shape[2] / 2
+        - center
+        - 0.5,  # Center of Rotation scalar or a vector
+        AnglesVec=-angles,  # A vector of projection angles in radians
+        ObjSize=recon_size,  # Reconstructed object dimensions (scalar)
+        device_projector=gpu_id,
+    )
+    return RecTools
+
+
 def _instantiate_iterative_recon_class(
     data: cp.ndarray,
     angles: np.ndarray,
@@ -361,3 +478,12 @@ def _take_neg_log(data: cp.ndarray) -> cp.ndarray:
     data[cp.isnan(data)] = 6.0
     data[cp.isinf(data)] = 0
     return data
+
+
+def _take_neg_log_np(data: np.ndarray) -> np.ndarray:
+    """Taking negative log"""
+    data[data <= 0] = 1
+    data = -np.log(data)
+    data[np.isnan(data)] = 6.0
+    data[np.isinf(data)] = 0
+    return data

tests/test_recon/test_algorithm.py

@@ -2,6 +2,7 @@
 import numpy as np
 from httomolibgpu.prep.normalize import normalize as normalize_cupy
 from httomolibgpu.recon.algorithm import (
+    FBP2d_astra,
     FBP,
     LPRec,
     SIRT,
@@ -13,6 +14,27 @@
 from cupy.cuda import nvtx
 
 
+def test_reconstruct_FBP_2d_astra(data, flats, darks, ensure_clean_memory):
+    normalised_data = normalize_cupy(data, flats, darks, cutoff=10, minus_log=True)
+    recon_size = 150
+
+    recon_data = FBP2d_astra(
+        cp.asnumpy(normalised_data),
+        np.linspace(0.0 * np.pi / 180.0, 180.0 * np.pi / 180.0, data.shape[0]),
+        79.5,
+        filter_type="shepp-logan",
+        filter_parameter=None,
+        filter_d=2.0,
+        recon_size=recon_size,
+        recon_mask_radius=0.9,
+    )
+    assert recon_data.flags.c_contiguous
+    assert_allclose(np.mean(recon_data), 0.0020, atol=1e-04)
+    assert_allclose(np.mean(recon_data, axis=(0, 2)).sum(), 0.265129, rtol=1e-05)
+    assert recon_data.dtype == np.float32
+    assert recon_data.shape == (recon_size, 128, recon_size)
+
+
 def test_reconstruct_FBP_1(data, flats, darks, ensure_clean_memory):
     recon_data = FBP(
         normalize_cupy(data, flats, darks, cutoff=10, minus_log=True),
@@ -97,9 +119,7 @@ def test_reconstruct_LPREC_1(data, flats, darks, ensure_clean_memory):
     )
     assert recon_data.flags.c_contiguous
     recon_data = recon_data.get()
-    assert_allclose(np.mean(recon_data), 0.0070263873, rtol=1e-07, atol=1e-6)
-    assert_allclose(np.mean(recon_data, axis=(0, 2)).sum(), 0.89937746, rtol=1e-05)
-    assert_allclose(np.max(recon_data), 0.098308131, rtol=1e-07, atol=1e-6)
+    assert_allclose(np.mean(recon_data), 0.0070, atol=1e-4)
     assert recon_data.dtype == np.float32
     assert recon_data.shape == (130, 128, 130)
 

zenodo-tests/test_recon/test_algorithm.py

@@ -7,6 +7,7 @@
 
 from httomolibgpu.prep.normalize import normalize
 from httomolibgpu.recon.algorithm import (
+    FBP2d_astra,
     FBP,
     LPRec,
 )
@@ -18,6 +19,33 @@
 from conftest import force_clean_gpu_memory
 
 
+def test_reconstruct_FBP2d_astra_i12_dataset1(i12_dataset1):
+    force_clean_gpu_memory()
+    projdata = i12_dataset1[0]
+    angles = i12_dataset1[1]
+    flats = i12_dataset1[2]
+    darks = i12_dataset1[3]
+    del i12_dataset1
+
+    data_normalised = normalize(projdata, flats, darks, minus_log=True)
+    del flats, darks, projdata
+    force_clean_gpu_memory()
+
+    recon_data = FBP2d_astra(
+        cp.asnumpy(data_normalised),
+        np.deg2rad(angles),
+        center=1253.75,
+        filter_type="shepp-logan",
+        filter_parameter=None,
+        filter_d=2.0,
+        recon_mask_radius=0.9,
+    )
+    assert recon_data.flags.c_contiguous
+    assert_allclose(np.sum(recon_data), 84673.68, rtol=1e-07, atol=1e-6)
+    assert recon_data.dtype == np.float32
+    assert recon_data.shape == (2560, 50, 2560)
+
+
 def test_reconstruct_FBP_i12_dataset1(i12_dataset1):
     force_clean_gpu_memory()
     projdata = i12_dataset1[0]