diff --git a/Demos/RealData.py b/Demos/RealData.py index 30ae1df4b..8fe1befe6 100644 --- a/Demos/RealData.py +++ b/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() diff --git a/tests/conftest.py b/tests/conftest.py index 5132b1254..b400c0881 100644 --- a/tests/conftest.py +++ b/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"])) diff --git a/tests/test_RecToolsIRCuPy.py b/tests/test_RecToolsIRCuPy.py index 8c909e503..de90c3cb0 100644 --- a/tests/test_RecToolsIRCuPy.py +++ b/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,7 +849,7 @@ 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) @@ -822,6 +857,7 @@ def test_FISTA_OS_PWLS_regul_PDTV_cp_3D(angles, raw_data, flats, darks): 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) diff --git a/tomobar/cuda_kernels/stripe_weighted_least_squares.cu b/tomobar/cuda_kernels/stripe_weighted_least_squares.cu new file mode 100644 index 000000000..0366c05b9 --- /dev/null +++ b/tomobar/cuda_kernels/stripe_weighted_least_squares.cu @@ -0,0 +1,24 @@ +#include + +template +__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(tz) * dimY * dimX + static_cast(ty) * dimX + static_cast(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); +} \ No newline at end of file diff --git a/tomobar/data_fidelities.py b/tomobar/data_fidelities.py index 10815192a..a954d7873 100644 --- a/tomobar/data_fidelities.py +++ b/tomobar/data_fidelities.py @@ -2,6 +2,7 @@ import cupy as cp from typing import Optional +from tomobar.cuda_kernels import load_cuda_module def grad_data_term( @@ -10,8 +11,8 @@ def grad_data_term( b: cp.ndarray, use_os: bool, sub_ind: int, - indVec: Optional[cp.ndarray] = None, w: Optional[cp.ndarray] = None, + w_sum: Optional[cp.ndarray] = None, ) -> cp.ndarray: """Calculation of the gradient of the data fidelity term Args: @@ -21,20 +22,38 @@ def grad_data_term( sub_ind (int): index for the ordered-subset approach. indVec (Optional, cp.ndarray): Array of indices for the OS-model. w (Optional, cp.ndarray): weights for Penalised-Weighted LS. + w_sum (Optional, cp.ndarray): sum_weights for SWLS. Returns: cp.ndarray: gradient of the data fidelity as a 3D CuPy array. """ + kernel_name = "stripe_weighted_least_squares_float" + module = load_cuda_module("stripe_weighted_least_squares") + stripe_weighted_least_squares = module.get_function(kernel_name) + if self.data_fidelity in ["LS", "PWLS"]: # Least-Squares (LS) res = self._Ax(x, sub_ind, use_os) - b if w is not None: # Penalised-Weighted least squares - if use_os: - cp.multiply(res, w[:, indVec, :], out=res) - else: - cp.multiply(res, w, out=res) + cp.multiply(res, w, out=res) elif self.data_fidelity == "KL": # Kullback-Leibler term. Note that b in that case should be given as pre-log data (raw) res = 1 - b / cp.clip(self._Ax(x, sub_ind, use_os), 1e-8, None) + elif self.data_fidelity == "SWLS": + res = self._Ax(x, sub_ind, use_os) - b + weights_mul_res = cp.multiply(w, res) + weights_dot_res = cp.sum(weights_mul_res, axis=1) + + dz, dy, dx = res.shape + block_dims = (128, 1, 1) + grid_dims = tuple( + (res.T.shape[i] + block_dims[i] - 1) // block_dims[i] for i in range(3) + ) + stripe_weighted_least_squares( + grid_dims, + block_dims, + (res, w, weights_mul_res, weights_dot_res, w_sum, dx, dy, dz), + ) + return self._Atb(res, sub_ind, use_os) diff --git a/tomobar/methodsIR_CuPy.py b/tomobar/methodsIR_CuPy.py index 72bab6607..8f6138344 100644 --- a/tomobar/methodsIR_CuPy.py +++ b/tomobar/methodsIR_CuPy.py @@ -367,6 +367,12 @@ def __common_initialisation( self.Atools.detectors_x_pad, cupyrun=True, ) + if "projection_raw_data" in _data_upd_: + _data_upd_["projection_raw_data"] = _apply_horiz_detector_padding( + _data_upd_["projection_raw_data"], + self.Atools.detectors_x_pad, + cupyrun=True, + ) if _algorithm_upd_.get("lipschitz_const") is None: _algorithm_upd_["lipschitz_const"] = self.powermethod(_data_upd_) @@ -389,14 +395,23 @@ def __common_initialisation( use_os = self.OS_number > 1 - if _data_["data_fidelity"] in ["PWLS"]: - w = cp.asarray(_data_upd_["projection_data"]) # weights for PWLS model + if _data_["data_fidelity"] in ["PWLS", "SWLS"]: + w = cp.asarray( + _data_upd_["projection_raw_data"], dtype=np.float32, order="C" + ) # weights for PWLS model w = cp.maximum(w, 1e-6) w /= w.max() else: w = None - return (_data_upd_, _algorithm_upd_, _regularisation_upd_, x0, w, use_os) + return ( + _data_upd_, + _algorithm_upd_, + _regularisation_upd_, + x0, + w, + use_os, + ) def FISTA( self, @@ -437,11 +452,38 @@ def FISTA( L_const_inv = 1.0 / _algorithm_upd_["lipschitz_const"] - proj_data = _data_upd_["projection_data"] indVec = None t = cp.float32(1.0) X_t = cp.copy(x0) X = cp.copy(x0) + weight_subset_sum = None + + if use_os: + proj_data = [None] * self.OS_number + weights = [None] * self.OS_number + weight_sums = [None] * self.OS_number + + for sub_ind in range(self.OS_number): + # select a specific set of indeces for the subset (OS) + indVec = self.Atools.newInd_Vec[sub_ind, :] + if indVec[self.Atools.NumbProjBins - 1] == 0: + indVec = indVec[:-1] # shrink vector size + + proj_data[sub_ind] = _data_upd_["projection_data"][:, indVec, :] + weights[sub_ind] = None if w is None else w[:, indVec, :] + weight_subset = weights[sub_ind] + + if _data_upd_["data_fidelity"] == "SWLS": + weight_sums[sub_ind] = ( + cp.sum(weight_subset, axis=1) + _data_upd_["beta_SWLS"] + ) + else: + proj_data_subset = _data_upd_["projection_data"] + weight_subset = w + if _data_upd_["data_fidelity"] == "SWLS": + weight_subset_sum = ( + cp.sum(weight_subset, axis=1) + _data_upd_["beta_SWLS"] + ) # FISTA iterations for _ in range(_algorithm_upd_["iterations"]): @@ -450,14 +492,18 @@ def FISTA( X_old = X t_old = t if use_os: - # select a specific set of indeces for the subset (OS) - indVec = self.Atools.newInd_Vec[sub_ind, :] - if indVec[self.Atools.NumbProjBins - 1] == 0: - indVec = indVec[:-1] # shrink vector size - proj_data = _data_upd_["projection_data"][:, indVec, :] + proj_data_subset = proj_data[sub_ind] + weight_subset = weights[sub_ind] + weight_subset_sum = weight_sums[sub_ind] grad_data = grad_data_term( - self, X_t, proj_data, use_os, sub_ind, indVec, w + self, + X_t, + proj_data_subset, + use_os, + sub_ind, + weight_subset, + weight_subset_sum, ) X = X_t - L_const_inv * grad_data @@ -512,7 +558,34 @@ def ADMM( ) = self.__common_initialisation( _data_, _algorithm_, _regularisation_, method_run="ADMM" ) - proj_data = _data_upd_["projection_data"] + weight_subset_sum = None + + if use_os: + proj_data = [None] * self.OS_number + weights = [None] * self.OS_number + weight_sums = [None] * self.OS_number + + for sub_ind in range(self.OS_number): + # select a specific set of indeces for the subset (OS) + indVec = self.Atools.newInd_Vec[sub_ind, :] + if indVec[self.Atools.NumbProjBins - 1] == 0: + indVec = indVec[:-1] # shrink vector size + + proj_data[sub_ind] = _data_upd_["projection_data"][:, indVec, :] + weights[sub_ind] = None if w is None else w[:, indVec, :] + + weight_subset = weights[sub_ind] + if _data_upd_["data_fidelity"] == "SWLS": + weight_sums[sub_ind] = ( + cp.sum(weight_subset, axis=1) + _data_upd_["beta_SWLS"] + ) + else: + proj_data_subset = _data_upd_["projection_data"] + weight_subset = w + if _data_upd_["data_fidelity"] == "SWLS": + weight_subset_sum = ( + cp.sum(weight_subset, axis=1) + _data_upd_["beta_SWLS"] + ) indVec = None x = x0.copy() @@ -531,15 +604,19 @@ def ADMM( for iter_no in range(_algorithm_upd_["iterations"]): for sub_ind in range(self.OS_number): if use_os: - # select a specific set of indeces for the subset (OS) - indVec = self.Atools.newInd_Vec[sub_ind, :] - if indVec[self.Atools.NumbProjBins - 1] == 0: - indVec = indVec[:-1] # shrink vector size - proj_data = _data_upd_["projection_data"][:, indVec, :] + proj_data_subset = proj_data[sub_ind] + weight_subset = weights[sub_ind] + weight_subset_sum = weight_sums[sub_ind] # ---- z-update (linearized data term) ---- grad_data = grad_data_term( - self, z, proj_data, use_os, sub_ind, indVec, w + self, + z, + proj_data_subset, + use_os, + sub_ind, + weight_subset, + weight_subset_sum, ) grad_admm = _algorithm_upd_["ADMM_rho_const"] * (z - x + u) diff --git a/tomobar/supp/dicts.py b/tomobar/supp/dicts.py index d5d4fb3e7..3e9968ca4 100644 --- a/tomobar/supp/dicts.py +++ b/tomobar/supp/dicts.py @@ -23,8 +23,10 @@ def dicts_check( Keyword Args: _data_['projection_data'] (ndarray): Can be either projection data after negative log or raw data given as a 3D CuPy array. + _data_['projection_raw_data'] (ndarray): Pre-log/pre flat-dark normalised raw data as a 3D CuPy array. Required for PWLS and SWLS fidelity terms. _data_['data_axes_labels_order'] (list, None). The order of the axes labels for the input data. The default data labels are: ["detY", "angles", "detX"]. - _data_['data_fidelity'] (str). Data fidelity given as 'LS' (Least Squares), 'PWLS' (Penalised Weightes LS), 'KL' (Kullback Leilbler). Defaults to 'LS'. + _data_['data_fidelity'] (str). Data fidelity given as 'LS' (Least Squares), 'PWLS' (Penalised Weighted LS), 'SWLS' (Stripe-Weighted PWLS), 'KL' (Kullback Leilbler). Defaults to 'LS'. + _data_['beta_SWLS'] (float): When SWLS data term is selected, this parameter controls the weight with which the stripes/rings will be suppressed. Defaults to 0.1. _algorithm_['iterations'] (int): The number of iterations for the reconstruction algorithm. _algorithm_['nonnegativity'] (bool): Enable nonnegativity for the solution. Defaults to False. @@ -51,43 +53,62 @@ def dicts_check( correct_labels_order2D = ["angles", "detX"] data2dinput = False + # -------- dealing with _data_ dictionary ------------ if _data_ is None: - raise NameError("The data dictionary must be always provided") - else: - # -------- dealing with _data_ dictionary ------------ - if _data_.get("projection_data") is None: - raise NameError("'projection_data' needs to be provided") - if _data_["projection_data"].ndim == 2: - data2dinput = True + raise NameError("The data dictionary must be provided") + if _data_.get("projection_data") is None: + raise NameError("'projection_data' must be provided") + + if _data_.get("data_fidelity") is None: + _data_["data_fidelity"] = "LS" + if _data_["data_fidelity"] not in ["LS", "PWLS", "KL", "SWLS"]: + raise ValueError( + "_data_['data_fidelity'] should be provided as 'LS', 'PWLS', 'SWLS' or 'KL'." + ) + if ( + _data_["data_fidelity"] in ["PWLS", "SWLS"] + and _data_.get("projection_raw_data") is None + ): + raise ValueError( + "For 'PWLS' and 'SWLS' data terms, 'projection_raw_data' must be provided (same dimension as 'projection_data')" + ) - if "data_axes_labels_order" not in _data_: - _data_["data_axes_labels_order"] = None + self.data_fidelity = _data_["data_fidelity"] - if _data_["data_axes_labels_order"] is not None: - if data2dinput: - correct_labels_order = correct_labels_order2D + if _data_["projection_data"].ndim == 2: + data2dinput = True - _data_["projection_data"] = _data_dims_swapper( - _data_["projection_data"], + if "data_axes_labels_order" not in _data_: + _data_["data_axes_labels_order"] = None + + if _data_["data_axes_labels_order"] is not None: + if data2dinput: + correct_labels_order = correct_labels_order2D + + _data_["projection_data"] = _data_dims_swapper( + _data_["projection_data"], + _data_["data_axes_labels_order"], + correct_labels_order, + ) + if "projection_raw_data" in _data_: + _data_["projection_raw_data"] = _data_dims_swapper( + _data_["projection_raw_data"], _data_["data_axes_labels_order"], correct_labels_order, ) - # we need to reset the swap option here as the data already been modified so we don't swap it again in the method itself - _data_["data_axes_labels_order"] = None + # we need to reset the swap option here as the data already been modified so we don't swap it again in the method itself + _data_["data_axes_labels_order"] = None - if data2dinput: - _data_["projection_data"] = cp.expand_dims( - _data_["projection_data"], axis=0 + if data2dinput: + _data_["projection_data"] = cp.expand_dims(_data_["projection_data"], axis=0) + if "projection_raw_data" in _data_: + _data_["projection_raw_data"] = cp.expand_dims( + _data_["projection_raw_data"], axis=0 ) - if _data_.get("data_fidelity") is None: - _data_["data_fidelity"] = "LS" - if _data_["data_fidelity"] not in {"LS", "PWLS", "KL"}: - raise ValueError( - "_data_['data_fidelity'] should be provided as 'LS', 'PWLS', 'KL'." - ) - else: - self.data_fidelity = _data_["data_fidelity"] + if _data_["data_fidelity"] == "SWLS": + if _data_.get("beta_SWLS") is None: + _data_["beta_SWLS"] = 0.1 if self.OS_number > 1: if method_run in {"SIRT", "CGLS", "Landweber"}: @@ -98,7 +119,7 @@ def dicts_check( # ---------- dealing with _algorithm_ -------------- if _algorithm_ is None: _algorithm_ = {} - if method_run in {"SIRT", "CGLS", "power", "Landweber", "OSEM"}: + if method_run in ["SIRT", "CGLS", "power", "Landweber", "OSEM"]: _algorithm_["lipschitz_const"] = 0 # bypass Lipshitz const calculation if _algorithm_.get("iterations") is None: if method_run == "SIRT": @@ -160,7 +181,7 @@ def dicts_check( _regularisation_ = {} if bool(_regularisation_) is False: _regularisation_["method"] = None - if method_run in {"FISTA", "ADMM", "OSEM"}: + if method_run in ["FISTA", "ADMM", "OSEM"]: # regularisation parameter (main) if "regul_param" not in _regularisation_: _regularisation_["regul_param"] = 0.001