Implement SWLS proximal operator for 3D

Demos/RealData.py

@@ -18,7 +18,7 @@
 from tomobar.methodsIR_CuPy import RecToolsIRCuPy
 from tomobar.methodsDIR_CuPy import RecToolsDIRCuPy
 
-# load dendritic data
+# load raw X-ray data of dendrites
 datadict = scipy.io.loadmat("../data/DendrRawData.mat")
 # extract data (print(datadict.keys()))
 dataRaw = datadict["data_raw3D"]
@@ -31,6 +31,7 @@
     dataRaw, flats[:, np.newaxis, :], darks[:, np.newaxis, :], axis=1
 )
 data_norm_cupy = cp.asarray(data_norm[:, :, 5:10], order="C")
+data_raw_cupy = cp.asarray(dataRaw[:, :, 5:10], order="C")
 
 detectorHoriz = cp.size(data_norm_cupy, 0)
 detectorVert = cp.size(data_norm_cupy, 2)
@@ -173,19 +174,20 @@
 _data_ = {
     "data_fidelity": "PWLS",
     "projection_data": data_norm_cupy,  # Normalised projection data
+    "projection_raw_data": data_raw_cupy,  # raw data for PWLS/SWLS models
     "data_axes_labels_order": data_labels3D,
 }
 tic = timeit.default_timer()
 ####################### Creating the algorithm dictionary: #######################
 _algorithm_ = {
-    "iterations": 25,
+    "iterations": 20,
     "recon_mask_radius": 2.0,
 }  # The number of iterations
 
 ##### creating regularisation dictionary: #####
 _regularisation_ = {
     "method": "PD_TV",  # Selected regularisation method
-    "regul_param": 0.000002,  # Regularisation parameter
+    "regul_param": 0.0000005,  # Regularisation parameter
     "iterations": 50,  # The number of regularisation iterations
     "half_precision": True,  # enabling half-precision calculation
 }
@@ -222,6 +224,7 @@
 _data_ = {
     "data_fidelity": "PWLS",
     "projection_data": data_norm_cupy,  # Normalised projection data
+    "projection_raw_data": data_raw_cupy,  # raw data for PWLS/SWLS models
     "data_axes_labels_order": data_labels3D,
 }
 
@@ -255,3 +258,104 @@
 plt.show()
 # fig.savefig('dendr_ADMM.png', dpi=200)
 # %%
+print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
+print("Reconstructing with FISTA OS-SWLS-TV (PD) method %%%%%%%%%%%")
+print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
+####################### Creating the data dictionary: #######################
+RectoolsCuPy = RecToolsIRCuPy(
+    DetectorsDimH=detectorHoriz,  # Horizontal detector dimension
+    DetectorsDimH_pad=padding_value,  # Padding size of horizontal detector
+    DetectorsDimV=detectorVert,  # Vertical detector dimension (3D case)
+    CenterRotOffset=None,  # Center of Rotation scalar
+    AnglesVec=angles_rad,  # A vector of projection angles in radians
+    ObjSize=N_size,  # Reconstructed object dimensions (scalar)
+    device_projector=0,
+    OS_number=6,  # The number of ordered subsets
+)
+
+_data_ = {
+    "data_fidelity": "SWLS",
+    "beta_SWLS": 1.0,  # stripe-suppressing parameter for SWLS
+    "projection_data": data_norm_cupy,  # Normalised projection data
+    "projection_raw_data": data_raw_cupy,  # raw data for PWLS/SWLS models
+    "data_axes_labels_order": data_labels3D,
+}
+
+tic = timeit.default_timer()
+####################### Creating the algorithm dictionary: #######################
+_algorithm_ = {
+    "iterations": 20,
+    "recon_mask_radius": 2.0,
+}  # The number of iterations
+
+##### creating regularisation dictionary: #####
+_regularisation_ = {
+    "method": "PD_TV",  # Selected regularisation method
+    "regul_param": 0.0000005,  # Regularisation parameter
+    "iterations": 50,  # The number of regularisation iterations
+    "half_precision": True,  # enabling half-precision calculation
+}
+
+# RUN THE FISTA METHOD:
+RecFISTA_os_swls_tv = RectoolsCuPy.FISTA(_data_, _algorithm_, _regularisation_)
+toc = timeit.default_timer()
+Run_time = toc - tic
+print("FISTA OS-SWLS-TV (PD) reconstruction done in {} seconds".format(Run_time))
+swls_tv_np = cp.asnumpy(RecFISTA_os_swls_tv)
+
+fig = plt.figure()
+plt.imshow(swls_tv_np[3, :, :], vmin=0, vmax=0.003, cmap="gray")
+plt.title("FISTA OS-SWLS-TV (PD) reconstruction")
+plt.show()
+# fig.savefig('dendr_SWLS.png', dpi=200)
+# %%
+print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
+print("Reconstructing with ADMM OS-SWLS-TV (PD) method %%%%%%%%%%%%%")
+print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
+RectoolsCuPy = RecToolsIRCuPy(
+    DetectorsDimH=detectorHoriz,  # Horizontal detector dimension
+    DetectorsDimH_pad=padding_value,  # Padding size of horizontal detector
+    DetectorsDimV=detectorVert,  # Vertical detector dimension (3D case)
+    CenterRotOffset=None,  # Center of Rotation scalar
+    AnglesVec=angles_rad,  # A vector of projection angles in radians
+    ObjSize=N_size,  # Reconstructed object dimensions (scalar)
+    device_projector=0,
+    OS_number=12,  # The number of ordered subsets
+)
+####################### Creating the data dictionary: #######################
+_data_ = {
+    "data_fidelity": "SWLS",
+    "beta_SWLS": 1.0,  # stripe-suppressing parameter for SWLS
+    "projection_data": data_norm_cupy,  # Normalised projection data
+    "projection_raw_data": data_raw_cupy,  # raw data for PWLS/SWLS models
+    "data_axes_labels_order": data_labels3D,
+}
+
+#################### Creating the algorithm dictionary: #######################
+_algorithm_ = {
+    "initialise": None,
+    "iterations": 20,
+    "ADMM_rho_const": 0.9,
+    "ADMM_relax_par": 1.7,
+    "recon_mask_radius": 2.0,
+}  # The number of iterations
+
+##### creating regularisation dictionary: #####
+_regularisation_ = {
+    "method": "PD_TV",  # Selected regularisation method
+    "regul_param": 0.002,  # Regularisation parameter
+    "iterations": 40,  # The number of regularisation iterations
+    "half_precision": True,  # enabling half-precision calculation
+}
+
+# RUN THE ADMM-OS-TV METHOD:
+tic = timeit.default_timer()
+RecADMM_os_tv = RectoolsCuPy.ADMM(_data_, _algorithm_, _regularisation_)
+toc = timeit.default_timer()
+Run_time = toc - tic
+print("ADMM-OS-TV (PD) reconstruction done in {} seconds".format(Run_time))
+
+fig = plt.figure()
+plt.imshow(cp.asnumpy((RecADMM_os_tv[3, :, :])), vmin=0, vmax=0.003, cmap="gray")
+plt.title("ADMM OS-TV (PD) reconstruction")
+plt.show()

tests/conftest.py

@@ -94,15 +94,18 @@ def raw_data(data_raw_file):
     return np.float32(np.copy(data_raw_file["data"]))
 
 
+@pytest.fixture
+def raw_data_cupy(raw_data):
+    return cp.asarray(raw_data)
+
+
 @pytest.fixture
 def flats(data_raw_file):
     return np.float32(np.copy(data_raw_file["flats"]))
 
 
 @pytest.fixture
-def darks(
-    data_raw_file,
-):
+def darks(data_raw_file):
     return np.float32(np.copy(data_raw_file["darks"]))
 
 

tests/test_RecToolsIRCuPy.py

@@ -323,6 +323,39 @@ def test_FISTA_cp_lc_known_3D(data_cupy, angles, ensure_clean_memory):
     assert Iter_rec.shape == (128, 160, 160)
 
 
+def test_FISTA_SWLS_cp_3D(data_cupy, raw_data_cupy, angles, ensure_clean_memory):
+    detX = cp.shape(data_cupy)[2]
+    detY = cp.shape(data_cupy)[1]
+    N_size = detX
+    RecTools = RecToolsIRCuPy(
+        DetectorsDimH=detX,  # Horizontal detector dimension
+        DetectorsDimH_pad=0,  # Padding size of horizontal detector
+        DetectorsDimV=detY,  # Vertical detector dimension (3D case)
+        CenterRotOffset=0.0,  # Center of Rotation scalar or a vector
+        AnglesVec=angles,  # A vector of projection angles in radians
+        ObjSize=N_size,  # Reconstructed object dimensions (scalar)
+        device_projector=0,  # define the device
+    )
+
+    _data_ = {
+        "data_fidelity": "SWLS",
+        "projection_data": data_cupy,
+        "projection_raw_data": raw_data_cupy,
+        "beta_SWLS": 1.0,
+        "data_axes_labels_order": ["angles", "detY", "detX"],
+    }  # data dictionary
+    # calculate Lipschitz constant
+    lc = RecTools.powermethod(_data_)
+    _algorithm_ = {"iterations": 10, "lipschitz_const": lc}
+    Iter_rec = RecTools.FISTA(_data_, _algorithm_)
+
+    Iter_rec = Iter_rec.get()
+    assert_allclose(np.min(Iter_rec), -0.0061533, rtol=1e-04)
+    assert_allclose(np.max(Iter_rec), 0.008243, rtol=1e-04)
+    assert Iter_rec.dtype == np.float32
+    assert Iter_rec.shape == (128, 160, 160)
+
+
 def test_FISTA_cp_3D(data_cupy, angles, ensure_clean_memory):
     detX = cp.shape(data_cupy)[2]
     detY = cp.shape(data_cupy)[1]
@@ -776,6 +809,7 @@ def test_FISTA_OS_regul_ROFTV_cp_3D(data_cupy, angles, ensure_clean_memory):
 def test_FISTA_OS_PWLS_regul_PDTV_cp_3D(angles, raw_data, flats, darks):
     normalised = normaliser(raw_data, flats, darks)
     normalised_cp = cp.asarray(normalised)
+    raw_data_cp = cp.asarray(raw_data)
 
     detX = cp.shape(normalised_cp)[2]
     detY = cp.shape(normalised_cp)[1]
@@ -794,6 +828,7 @@ def test_FISTA_OS_PWLS_regul_PDTV_cp_3D(angles, raw_data, flats, darks):
     _data_ = {
         "data_fidelity": "PWLS",
         "projection_data": normalised_cp,
+        "projection_raw_data": raw_data_cp,
         "data_axes_labels_order": ["angles", "detY", "detX"],
     }
     # calculate Lipschitz constant
@@ -814,14 +849,15 @@ def test_FISTA_OS_PWLS_regul_PDTV_cp_3D(angles, raw_data, flats, darks):
     Iter_rec = Iter_rec.get()
 
     assert 4000 <= lc <= 5000
-    assert_allclose(np.max(Iter_rec), 0.03565, rtol=1e-03)
+    assert_allclose(np.max(Iter_rec), 0.027205, rtol=1e-03)
     assert Iter_rec.dtype == np.float32
     assert Iter_rec.shape == (128, 160, 160)
 
 
 def test_FISTA_OS_PWLS_regul_ROFTV_cp_3D(angles, raw_data, flats, darks):
     normalised = normaliser(raw_data, flats, darks)
     normalised_cp = cp.asarray(normalised)
+    raw_data_cp = cp.asarray(raw_data)
 
     detX = cp.shape(normalised_cp)[2]
     detY = cp.shape(normalised_cp)[1]
@@ -839,6 +875,7 @@ def test_FISTA_OS_PWLS_regul_ROFTV_cp_3D(angles, raw_data, flats, darks):
     # data dictionary
     _data_ = {
         "data_fidelity": "PWLS",
+        "projection_raw_data": raw_data_cp,
         "projection_data": normalised_cp,
         "data_axes_labels_order": ["angles", "detY", "detX"],
     }
@@ -860,7 +897,7 @@ def test_FISTA_OS_PWLS_regul_ROFTV_cp_3D(angles, raw_data, flats, darks):
     Iter_rec = Iter_rec.get()
 
     assert 5000 <= lc <= 6000
-    assert_allclose(np.max(Iter_rec), 0.032535, rtol=1e-03)
+    assert_allclose(np.max(Iter_rec), 0.027137, rtol=1e-03)
     assert Iter_rec.dtype == np.float32
     assert Iter_rec.shape == (128, 160, 160)
 
@@ -985,9 +1022,13 @@ def test_ADMM_OS_cp_3D(data_cupy, angles, test_case, ensure_clean_memory):
     assert Iter_rec.shape == (128, 160, 160)
 
 
-def test_ADMM_OS_PWLS_warmstart_pad_cp_3D(data_cupy, angles, ensure_clean_memory):
+def test_ADMM_OS_PWLS_warmstart_pad_cp_3D(
+    data_cupy, raw_data, angles, ensure_clean_memory
+):
     detX = cp.shape(data_cupy)[2]
     detY = cp.shape(data_cupy)[1]
+    raw_data_cp = cp.asarray(raw_data)
+
     N_size = 128
     pad_value = 17
     RecTools = RecToolsIRCuPy(
@@ -1007,6 +1048,7 @@ def test_ADMM_OS_PWLS_warmstart_pad_cp_3D(data_cupy, angles, ensure_clean_memory
     _data_ = {
         "data_fidelity": "PWLS",
         "projection_data": data_cupy,
+        "projection_raw_data": raw_data_cp,
         "data_axes_labels_order": ["angles", "detY", "detX"],
     }
     _algorithm_ = {
@@ -1025,6 +1067,61 @@ def test_ADMM_OS_PWLS_warmstart_pad_cp_3D(data_cupy, angles, ensure_clean_memory
     Iter_rec = RecTools.ADMM(_data_, _algorithm_, _regularisation_)
 
     Iter_rec = Iter_rec.get()
-    assert_allclose(np.max(Iter_rec), 0.031305, rtol=1e-03)
+    assert_allclose(np.max(Iter_rec), 0.02733, rtol=1e-03)
     assert Iter_rec.dtype == np.float32
     assert Iter_rec.shape == (128, 128, 128)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    zip(
+        ADMM_TEST_CASES,
+        [
+            -0.000639,
+            -0.000632,
+            -0.000487,
+        ],
+        [
+            0.000919,
+            0.000915,
+            0.000841,
+        ],
+    ),
+    ids=ADMM_TEST_IDS,
+)
+def test_ADMM_SWLS_cp_3D(
+    data_cupy, raw_data_cupy, angles, test_case, ensure_clean_memory
+):
+    regularization, expected_min, expected_max = test_case
+    detX = cp.shape(data_cupy)[2]
+    detY = cp.shape(data_cupy)[1]
+    N_size = detX
+    RecTools = RecToolsIRCuPy(
+        DetectorsDimH=detX,  # Horizontal detector dimension
+        DetectorsDimH_pad=0,  # Padded size of horizontal detector with edge values
+        DetectorsDimV=detY,  # Vertical detector dimension (3D case)
+        CenterRotOffset=0.0,  # Center of Rotation scalar or a vector
+        AnglesVec=angles,  # A vector of projection angles in radians
+        ObjSize=N_size,  # Reconstructed object dimensions (scalar)
+        device_projector=0,  # define the device
+    )
+
+    _data_ = {
+        "data_fidelity": "SWLS",
+        "projection_data": data_cupy,
+        "projection_raw_data": raw_data_cupy,
+        "beta_SWLS": 1.0,
+        "data_axes_labels_order": ["angles", "detY", "detX"],
+    }
+    _algorithm_ = {
+        "iterations": 2,
+        "ADMM_rho_const": 1.0,
+        "ADMM_relax_par": 1.6,
+    }
+
+    Iter_rec = RecTools.ADMM(_data_, _algorithm_, regularization)
+
+    assert_allclose(cp.min(Iter_rec).get(), expected_min, rtol=0, atol=eps)
+    assert_allclose(cp.max(Iter_rec).get(), expected_max, rtol=0, atol=eps)
+    assert Iter_rec.dtype == np.float32
+    assert Iter_rec.shape == (128, 160, 160)

tomobar/cuda_kernels/stripe_weighted_least_squares.cu

@@ -0,0 +1,24 @@
+#include <cuda_fp16.h>
+
+template <typename T>
+__device__ __forceinline__ void stripe_weighted_least_squares(T *res, T *weights, T *weights_mul_res, T *weights_dot_res, T *weight_sum, int dimX, int dimY, int dimZ)
+{
+    const long tx = blockDim.x * blockIdx.x + threadIdx.x;
+    const long ty = blockDim.y * blockIdx.y + threadIdx.y;
+    const long tz = blockDim.z * blockIdx.z + threadIdx.z;
+
+    if (tx >= dimX || ty >= dimY || tz >= dimZ)
+    {
+        return;
+    }
+
+    const long long index = static_cast<long long>(tz) * dimY * dimX + static_cast<long long>(ty) * dimX + static_cast<long long>(tx);
+    const long long collapsed_projection_index = tz * dimX + tx;
+
+    res[index] = weights_mul_res[index] - 1.0 / weight_sum[collapsed_projection_index] * weights_dot_res[collapsed_projection_index] * weights[index];
+}
+
+extern "C" __global__ void stripe_weighted_least_squares_float(float *res, float *weights, float *weights_mul_res, float *weights_dot_res, float *weight_sum, int dimX, int dimY, int dimZ)
+{
+    stripe_weighted_least_squares(res, weights, weights_mul_res, weights_dot_res, weight_sum, dimX, dimY, dimZ);
+}