AOT-biomaps 2.9.176__py3-none-any.whl → 2.9.279__py3-none-any.whl

AOT_biomaps/AOT_Recon/AOT_Optimizers/MLEM.py

@@ -1,10 +1,18 @@
 from AOT_biomaps.AOT_Recon.ReconTools import _forward_projection, _backward_projection, check_gpu_memory, calculate_memory_requirement
 from AOT_biomaps.Config import config
+from AOT_biomaps.AOT_Recon.AOT_SparseSMatrix.SparseSMatrix_SELL import SparseSMatrix_SELL
+from AOT_biomaps.AOT_Recon.AOT_SparseSMatrix.SparseSMatrix_CSR import SparseSMatrix_CSR
+from AOT_biomaps.AOT_Recon.ReconEnums import SMatrixType
 import numba
 import torch
 import numpy as np
 import os
 from tqdm import trange
+import cupy as cp
+import cupyx.scipy.sparse as cpsparse
+import gc
+import pycuda.driver as drv
+
 
 def MLEM(
     SMatrix,
@@ -13,9 +21,12 @@ def MLEM(
     isSavingEachIteration=True,
     withTumor=True,
     device=None,
-    use_multi_gpu=False,
     use_numba=False,
+    denominator_threshold=1e-6,
     max_saves=5000,
+    show_logs=True,
+    smatrixType=SMatrixType.SELL,
+    Z=350,
 ):
     """
     Unified MLEM algorithm for Acousto-Optic Tomography.
@@ -33,34 +44,38 @@ def MLEM(
     Returns:
         Reconstructed image(s) and iteration indices (if isSavingEachIteration)
     """
-    try:
-        tumor_str = "WITH" if withTumor else "WITHOUT"
-        # Auto-select device and method
-        if device is None:
-            if torch.cuda.is_available() and check_gpu_memory(config.select_best_gpu(), calculate_memory_requirement(SMatrix, y)):
-                device = torch.device(f"cuda:{config.select_best_gpu()}")
-                use_gpu = True
-            else:
-                device = torch.device("cpu")
-                use_gpu = False
+    # try:
+    tumor_str = "WITH" if withTumor else "WITHOUT"
+    # Auto-select device and method
+    if device is None:
+        if torch.cuda.is_available() and check_gpu_memory(config.select_best_gpu(), calculate_memory_requirement(SMatrix, y), show_logs=show_logs):
+            device = torch.device(f"cuda:{config.select_best_gpu()}")
+            use_gpu = True
         else:
-            use_gpu = device.type == "cuda"
-        # Dispatch to the appropriate implementation
-        if use_gpu:
-            if use_multi_gpu and torch.cuda.device_count() > 1:
-                return _MLEM_multi_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves)
+            device = torch.device("cpu")
+            use_gpu = False
+    else:
+        use_gpu = device.type == "cuda"
+    # Dispatch to the appropriate implementation
+    if use_gpu:
+            if smatrixType == SMatrixType.CSR:
+                return MLEM_sparseCSR_pycuda(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves, denominator_threshold, Z, show_logs)
+            elif smatrixType == SMatrixType.SELL:
+                return MLEM_sparseSELL_pycuda(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves, denominator_threshold, show_logs)
+            elif smatrixType == SMatrixType.DENSE:
+                return _MLEM_single_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves, denominator_threshold,show_logs)
             else:
-                return _MLEM_single_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves)
+                raise ValueError("Unsupported SMatrixType for GPU MLEM.")
+    else:
+        if use_numba:
+            return _MLEM_CPU_numba(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves, denominator_threshold, show_logs)
         else:
-            if use_numba:
-                return _MLEM_CPU_numba(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves)
-            else:
-                return _MLEM_CPU_opti(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves)
-    except Exception as e:
-        print(f"Error in MLEM: {type(e).__name__}: {e}")
-        return None, None
+            return _MLEM_CPU_opti(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves, denominator_threshold, show_logs)
+    # except Exception as e:
+    #     print(f"Error in MLEM: {type(e).__name__}: {e}")
+    #     return None, None
 
-def _MLEM_single_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves=5000):
+def _MLEM_single_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves, denominator_threshold, show_logs=True):
     try:
         eps = torch.finfo(torch.float32).eps
         T, Z, X, N = SMatrix.shape
@@ -82,7 +97,6 @@ def _MLEM_single_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str
             .reshape(-1)
         )
         description = f"AOT-BioMaps -- ML-EM ---- {tumor_str} TUMOR ---- GPU {torch.cuda.current_device()}"
-
         # Calculate save indices
         if numIterations <= max_saves:
             save_indices = list(range(numIterations))
@@ -91,20 +105,21 @@ def _MLEM_single_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str
             save_indices = list(range(0, numIterations, step))
             if save_indices[-1] != numIterations - 1:
                 save_indices.append(numIterations - 1)
-
         saved_theta = []
         saved_indices = []
-
         with torch.no_grad():
-            for it in trange(numIterations, desc=description):
+            # Utilise range si show_logs=False, sinon trange
+            iterator = trange(numIterations, desc=description) if show_logs else range(numIterations)
+            for it in iterator:
                 q_flat = A_flat @ theta_flat
-                e_flat = y_flat / (q_flat + eps)
+                # Appliquer le seuil : si q_flat < denominator_threshold, on met e_flat à 1 (comme dans le code C++)
+                mask = q_flat >= denominator_threshold
+                e_flat = torch.where(mask, y_flat / (q_flat + eps), torch.ones_like(q_flat))
                 c_flat = A_flat.T @ e_flat
                 theta_flat = (theta_flat / (norm_factor_flat + eps)) * c_flat
                 if isSavingEachIteration and it in save_indices:
                     saved_theta.append(theta_flat.reshape(Z, X).clone())
                     saved_indices.append(it)
-
         # Free memory
         del A_flat, y_flat, norm_factor_flat
         torch.cuda.empty_cache()
@@ -117,74 +132,15 @@ def _MLEM_single_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str
         torch.cuda.empty_cache()
         return None, None
 
-def _MLEM_multi_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves=5000):
-    try:
-        num_gpus = torch.cuda.device_count()
-        device = torch.device('cuda:0')
-        T, Z, X, N = SMatrix.shape
-        A_matrix_torch = torch.tensor(SMatrix, dtype=torch.float32).to(device).permute(0, 3, 1, 2).reshape(T * N, Z * X)
-        y_torch = torch.tensor(y, dtype=torch.float32).to(device).reshape(-1)
-        A_split = torch.chunk(A_matrix_torch, num_gpus, dim=0)
-        y_split = torch.chunk(y_torch, num_gpus)
-        theta_0 = torch.ones((Z, X), dtype=torch.float32, device=device)
-        theta_list = [theta_0.clone().to(device) for _ in range(num_gpus)]
-        normalization_factor = A_matrix_torch.sum(dim=0).reshape(Z, X).to(device)
-
-        # Calculate save indices
-        if numIterations <= max_saves:
-            save_indices = list(range(numIterations))
-        else:
-            step = numIterations // max_saves
-            save_indices = list(range(0, numIterations, step))
-            if save_indices[-1] != numIterations - 1:
-                save_indices.append(numIterations - 1)
-
-        saved_theta = [theta_0.cpu().numpy()]
-        saved_indices = [0]
-        description = f"AOT-BioMaps -- ML-EM ---- {tumor_str} TUMOR ---- processing on multi-GPU ({num_gpus} GPUs) ----"
-
-        for it in trange(numIterations, desc=description):
-            theta_p_list = []
-            for i in range(num_gpus):
-                with torch.cuda.device(f'cuda:{i}'):
-                    theta_p = theta_list[i].to(f'cuda:{i}')
-                    A_i = A_split[i].to(f'cuda:{i}')
-                    y_i = y_split[i].to(f'cuda:{i}')
-                    q_flat = A_i @ theta_p.reshape(-1)
-                    e_flat = y_i / (q_flat + torch.finfo(torch.float32).tiny)
-                    c_flat = A_i.T @ e_flat
-                    theta_p_plus_1_flat = (theta_p.reshape(-1) / (normalization_factor.to(f'cuda:{i}').reshape(-1) + torch.finfo(torch.float32).tiny)) * c_flat
-                    theta_p_plus_1 = theta_p_plus_1_flat.reshape(Z, X)
-                    theta_p_list.append(theta_p_plus_1)
-            for i in range(num_gpus):
-                theta_list[i] = theta_p_list[i].to('cuda:0')
-            if isSavingEachIteration and it in save_indices:
-                saved_theta.append(torch.stack(theta_p_list).mean(dim=0).cpu().numpy())
-                saved_indices.append(it + 1)
-
-        del A_matrix_torch, y_torch, A_split, y_split, theta_0, normalization_factor
-        for i in range(num_gpus):
-            torch.cuda.empty_cache()
-        if not isSavingEachIteration:
-            return torch.stack(theta_p_list).mean(dim=0).cpu().numpy(), None
-        else:
-            return saved_theta, saved_indices
-    except Exception as e:
-        print(f"Error in multi-GPU MLEM: {type(e).__name__}: {e}")
-        del A_matrix_torch, y_torch, A_split, y_split, theta_0, normalization_factor
-        for i in range(num_gpus):
-            torch.cuda.empty_cache()
-        return None, None
-
-def _MLEM_CPU_numba(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves=5000):
+def _MLEM_CPU_numba(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves, denominator_threshold, show_logs=True):
     try:
         numba.set_num_threads(os.cpu_count())
-        q_p = np.zeros((SMatrix.shape[0], SMatrix.shape[3]))
-        c_p = np.zeros((SMatrix.shape[1], SMatrix.shape[2]))
-        theta_p_0 = np.ones((SMatrix.shape[1], SMatrix.shape[2]))
+        q_p = np.zeros((SMatrix.shape[0], SMatrix.shape[3]), dtype=np.float32)
+        c_p = np.zeros((SMatrix.shape[1], SMatrix.shape[2]), dtype=np.float32)
+        theta_p_0 = np.ones((SMatrix.shape[1], SMatrix.shape[2]), dtype=np.float32)
         matrix_theta = [theta_p_0]
         saved_indices = [0]
-        normalization_factor = np.sum(SMatrix, axis=(0, 3))
+        normalization_factor = np.sum(SMatrix, axis=(0, 3)).astype(np.float32)
 
         # Calculate save indices
         if numIterations <= max_saves:
@@ -196,14 +152,20 @@ def _MLEM_CPU_numba(SMatrix, y, numIterations, isSavingEachIteration, tumor_str,
                 save_indices.append(numIterations - 1)
 
         description = f"AOT-BioMaps -- ML-EM ---- {tumor_str} TUMOR ---- processing on multithread CPU ({numba.config.NUMBA_DEFAULT_NUM_THREADS} threads) ----"
+        iterator = trange(numIterations, desc=description) if show_logs else range(numIterations)
 
-        for it in trange(numIterations, desc=description):
+        for it in iterator:
             theta_p = matrix_theta[-1]
             _forward_projection(SMatrix, theta_p, q_p)
-            e_p = y / (q_p + 1e-8)
+
+            # Appliquer le seuil : si q_p < denominator_threshold, on met e_p à 1
+            mask = q_p >= denominator_threshold
+            e_p = np.where(mask, y / (q_p + 1e-8), 1.0)
+
             _backward_projection(SMatrix, e_p, c_p)
             theta_p_plus_1 = theta_p / (normalization_factor + 1e-8) * c_p
-            if isSavingEachIteration and it in save_indices:
+
+            if isSavingEachIteration and (it + 1) in save_indices:
                 matrix_theta.append(theta_p_plus_1)
                 saved_indices.append(it + 1)
             else:
@@ -217,7 +179,7 @@ def _MLEM_CPU_numba(SMatrix, y, numIterations, isSavingEachIteration, tumor_str,
         print(f"Error in Numba CPU MLEM: {type(e).__name__}: {e}")
         return None, None
 
-def _MLEM_CPU_opti(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves=5000):
+def _MLEM_CPU_opti(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves, denominator_threshold, show_logs=True):
     try:
         T, Z, X, N = SMatrix.shape
         A_flat = SMatrix.astype(np.float32).transpose(0, 3, 1, 2).reshape(T * N, Z * X)
@@ -238,16 +200,22 @@ def _MLEM_CPU_opti(SMatrix, y, numIterations, isSavingEachIteration, tumor_str,
                 save_indices.append(numIterations - 1)
 
         description = f"AOT-BioMaps -- ML-EM ---- {tumor_str} TUMOR ---- processing on single CPU (optimized) ----"
+        iterator = trange(numIterations, desc=description) if show_logs else range(numIterations)
 
-        for it in trange(numIterations, desc=description):
+        for it in iterator:
             theta_p = matrix_theta[-1]
             theta_p_flat = theta_p.reshape(-1)
             q_flat = A_flat @ theta_p_flat
-            e_flat = y_flat / (q_flat + np.finfo(np.float32).tiny)
+
+            # Appliquer le seuil : si q_flat < denominator_threshold, on met e_flat à 1
+            mask = q_flat >= denominator_threshold
+            e_flat = np.where(mask, y_flat / (q_flat + np.finfo(np.float32).tiny), 1.0)
+
             c_flat = A_flat.T @ e_flat
             theta_p_plus_1_flat = theta_p_flat / (normalization_factor_flat + np.finfo(np.float32).tiny) * c_flat
             theta_p_plus_1 = theta_p_plus_1_flat.reshape(Z, X)
-            if isSavingEachIteration and it in save_indices:
+
+            if isSavingEachIteration and (it + 1) in save_indices:
                 matrix_theta.append(theta_p_plus_1)
                 saved_indices.append(it + 1)
             else:
@@ -260,3 +228,239 @@ def _MLEM_CPU_opti(SMatrix, y, numIterations, isSavingEachIteration, tumor_str,
     except Exception as e:
         print(f"Error in optimized CPU MLEM: {type(e).__name__}: {e}")
         return None, None
+    
+def MLEM_sparseCSR_pycuda(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves, denominator_threshold, show_logs=True):
+    """
+    SMatrix: instance of SparseMatrixGPU (already allocated)
+    y: measured data (1D np.float32 of length TN)
+
+    Assumptions:
+    - SMatrix.values_gpu and SMatrix.col_ind_gpu and SMatrix.row_ptr_gpu are device pointers
+    - SMatrix.norm_factor_inv_gpu exists
+    - SMatrix.ctx is the PyCUDA context for the target GPU.
+    """
+    
+    # We use a final_result placeholder to ensure it's defined outside the try block
+    final_result = None
+    
+    try:
+        if not isinstance(SMatrix, SparseSMatrix_CSR):
+            raise TypeError("SMatrix must be a SparseSMatrix_CSR object")
+
+        # --- CONTEXT FIX: Push the context associated with SMatrix ---
+        # This ensures all subsequent PyCUDA operations use the correct GPU/context.
+        if SMatrix.ctx:
+            SMatrix.ctx.push()
+        # -----------------------------------------------------------
+
+        dtype = np.float32
+        TN = SMatrix.N * SMatrix.T
+        ZX = SMatrix.Z * SMatrix.X
+        # Ensure Z and X are correctly defined for reshaping
+        Z = SMatrix.Z 
+        X = SMatrix.X
+
+        if show_logs:
+            # We assume SMatrix was initialized using the correct device index.
+            print(f"Executing on GPU device index: {SMatrix.device.primary_context.device.name()}")
+            print(f"Dim X: {X}, Dim Z: {Z}, TN: {TN}, ZX: {ZX}")
+
+        # streams
+        stream = drv.Stream()
+
+        # allocate device buffers
+        y = y.T.flatten().astype(np.float32)
+        y_gpu = drv.mem_alloc(y.nbytes)
+        drv.memcpy_htod_async(y_gpu, y.astype(dtype), stream)
+
+        theta_flat_gpu = drv.mem_alloc(ZX * np.dtype(dtype).itemsize)
+        initial_theta = np.full(ZX, 0.1, dtype=dtype)
+        drv.memcpy_htod_async(theta_flat_gpu, initial_theta, stream)
+
+        norm_factor_inv_gpu = SMatrix.norm_factor_inv_gpu
+
+        q_flat_gpu = drv.mem_alloc(TN * np.dtype(dtype).itemsize)
+        e_flat_gpu = drv.mem_alloc(TN * np.dtype(dtype).itemsize)
+        c_flat_gpu = drv.mem_alloc(ZX * np.dtype(dtype).itemsize)
+
+        # Assuming the cubin file is found globally or managed by the caller
+        mod = drv.module_from_file('AOT_biomaps_kernels.cubin')
+        projection_kernel = mod.get_function('projection_kernel__CSR')
+        backprojection_kernel = mod.get_function('backprojection_kernel__CSR')
+        ratio_kernel = mod.get_function('ratio_kernel')
+        update_kernel = mod.get_function('update_theta_kernel')
+        block_size = 256
+
+        saved_theta, saved_indices = [], []
+        if numIterations <= max_saves:
+            save_indices = list(range(numIterations))
+        else:
+            save_indices = list(range(0, numIterations, max(1, numIterations // max_saves)))
+            if save_indices[-1] != numIterations - 1:
+                save_indices.append(numIterations - 1)
+
+        description = f"AOT-BioMaps -- ML-EM (CSR-sparse SMatrix) ---- {tumor_str} TUMOR ---- GPU {torch.cuda.current_device()}"
+        iterator = trange(numIterations, desc=description) if show_logs else range(numIterations)
+        for it in iterator:
+            # projection: q = A * theta
+            projection_kernel(q_flat_gpu, SMatrix.values_gpu, SMatrix.row_ptr_gpu, SMatrix.col_ind_gpu,
+                              theta_flat_gpu, np.int32(TN),
+                              block=(block_size, 1, 1), grid=((TN + block_size - 1) // block_size, 1, 1),
+                              stream=stream)
+
+            # ratio: e = y / max(q, threshold)
+            ratio_kernel(e_flat_gpu, y_gpu, q_flat_gpu, np.float32(denominator_threshold), np.int32(TN),
+                              block=(block_size, 1, 1), grid=((TN + block_size - 1) // block_size, 1, 1), stream=stream)
+
+            # backprojection: c = A^T * e
+            drv.memset_d32_async(c_flat_gpu, 0, ZX, stream)
+            backprojection_kernel(c_flat_gpu, SMatrix.values_gpu, SMatrix.row_ptr_gpu, SMatrix.col_ind_gpu,
+                                  e_flat_gpu, np.int32(TN),
+                                  block=(block_size, 1, 1), grid=((TN + block_size - 1) // block_size, 1, 1), stream=stream)
+
+            # update: theta *= norm_factor_inv * c
+            update_kernel(theta_flat_gpu, c_flat_gpu, norm_factor_inv_gpu, np.int32(ZX),
+                          block=(block_size, 1, 1), grid=((ZX + block_size - 1) // block_size, 1, 1), stream=stream)
+
+            if show_logs and (it % 10 == 0 or it == numIterations - 1):
+                drv.Context.synchronize()
+
+            if isSavingEachIteration and it in save_indices:
+                theta_host = np.empty(ZX, dtype=dtype)
+                drv.memcpy_dtoh(theta_host, theta_flat_gpu)
+                saved_theta.append(theta_host.reshape(Z, X))
+                saved_indices.append(it)
+
+        drv.Context.synchronize()
+
+        final_result = np.empty(ZX, dtype=dtype)
+        drv.memcpy_dtoh(final_result, theta_flat_gpu)
+        final_result = final_result.reshape(Z, X)
+
+        # free local allocations
+        y_gpu.free(); q_flat_gpu.free(); e_flat_gpu.free(); c_flat_gpu.free(); theta_flat_gpu.free()
+
+        return (saved_theta, saved_indices) if isSavingEachIteration else (final_result, None)
+
+    except Exception as e:
+        print(f"Error in MLEM_sparseCSR_pycuda: {type(e).__name__}: {e}")
+        gc.collect()
+        return None, None
+    
+    finally:
+        # --- CONTEXT FIX: Pop the context ---
+        if SMatrix and hasattr(SMatrix, 'ctx') and SMatrix.ctx:
+            SMatrix.ctx.pop()
+        # ------------------------------------
+
+def MLEM_sparseSELL_pycuda(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves, denominator_threshold, show_logs=True):
+    """
+    MLEM using SELL-C-σ kernels already present on device.
+    y must be float32 length TN.
+    """
+    final_result = None
+
+    try:
+        # check if SMatrix is SparseSMatrix_SELL object
+        if not isinstance(SMatrix, SparseSMatrix_SELL):
+            raise TypeError("SMatrix must be a SparseSMatrix_SELL object")
+        if SMatrix.sell_values_gpu is None:
+            raise RuntimeError("SELL not built. Call allocate_sell_c_sigma_direct() first.")
+            
+        # --- CONTEXT FIX: Push the context associated with SMatrix ---
+        # This ensures all subsequent PyCUDA operations use the correct GPU/context.
+        if SMatrix.ctx:
+            SMatrix.ctx.push()
+        # -----------------------------------------------------------
+
+        TN = int(SMatrix.N * SMatrix.T)
+        ZX = int(SMatrix.Z * SMatrix.X)
+        dtype = np.float32
+        block_size = 256
+
+        mod = drv.module_from_file('AOT_biomaps_kernels.cubin')
+        proj = mod.get_function("projection_kernel__SELL")
+        backproj = mod.get_function("backprojection_kernel__SELL")
+        ratio = mod.get_function("ratio_kernel")
+        update = mod.get_function("update_theta_kernel")
+
+        stream = drv.Stream()
+
+        # device buffers
+        y = y.T.flatten().astype(np.float32)
+        y_gpu = drv.mem_alloc(y.nbytes)
+        drv.memcpy_htod_async(y_gpu, y.astype(dtype), stream)
+
+        theta_gpu = drv.mem_alloc(ZX * np.dtype(dtype).itemsize)
+        drv.memcpy_htod_async(theta_gpu, np.full(ZX, 0.1, dtype=dtype), stream)
+
+        q_gpu = drv.mem_alloc(TN * np.dtype(dtype).itemsize)
+        e_gpu = drv.mem_alloc(TN * np.dtype(dtype).itemsize)
+        c_gpu = drv.mem_alloc(ZX * np.dtype(dtype).itemsize)
+
+        slice_ptr_gpu = SMatrix.slice_ptr_gpu
+        slice_len_gpu = SMatrix.slice_len_gpu
+        slice_height = np.int32(SMatrix.slice_height)
+
+        grid_rows = ((TN + block_size - 1) // block_size, 1, 1)
+        grid_cols = ((ZX + block_size - 1) // block_size, 1, 1)
+
+        saved_theta, saved_indices = [], []
+        if numIterations <= max_saves:
+            save_indices = list(range(numIterations))
+        else:
+            save_indices = list(range(0, numIterations, max(1, numIterations // max_saves)))
+            if save_indices[-1] != numIterations - 1:
+                save_indices.append(numIterations - 1)
+
+        description = f"AOT-BioMaps -- ML-EM (SELL-c-σ-sparse SMatrix) ---- {tumor_str} TUMOR ---- GPU {torch.cuda.current_device()}"
+        iterator = trange(numIterations, desc=description) if show_logs else range(numIterations)
+        for it in iterator:
+            # projection
+            proj(q_gpu, SMatrix.sell_values_gpu, SMatrix.sell_colinds_gpu, slice_ptr_gpu, slice_len_gpu,
+                    theta_gpu, np.int32(TN), slice_height,
+                    block=(block_size,1,1), grid=grid_rows, stream=stream)
+
+            # ratio
+            ratio(e_gpu, y_gpu, q_gpu, np.float32(denominator_threshold), np.int32(TN),
+                    block=(block_size,1,1), grid=grid_rows, stream=stream)
+
+            # zero c
+            drv.memset_d32_async(c_gpu, 0, ZX, stream)
+
+            # backprojection accumulate
+            backproj(SMatrix.sell_values_gpu, SMatrix.sell_colinds_gpu, slice_ptr_gpu, slice_len_gpu,
+                        e_gpu, c_gpu, np.int32(TN), slice_height,
+                        block=(block_size,1,1), grid=grid_rows, stream=stream)
+
+            # update
+            update(theta_gpu, c_gpu, SMatrix.norm_factor_inv_gpu, np.int32(ZX),
+                    block=(block_size,1,1), grid=grid_cols, stream=stream)
+
+            stream.synchronize()
+            if isSavingEachIteration and it in save_indices:
+                out = np.empty(ZX, dtype=np.float32)
+                drv.memcpy_dtoh(out, theta_gpu)
+                saved_theta.append(out.reshape((SMatrix.Z, SMatrix.X)))
+                saved_indices.append(it)
+
+        # final copy
+        res = np.empty(ZX, dtype=np.float32)
+        drv.memcpy_dtoh(res, theta_gpu)
+
+        # free temporaries
+        y_gpu.free(); q_gpu.free(); e_gpu.free(); c_gpu.free(); theta_gpu.free()
+        
+        final_result = res.reshape((SMatrix.Z, SMatrix.X))
+        return (saved_theta, saved_indices) if isSavingEachIteration else (final_result, None)
+        
+    except Exception as e:
+        print(f"Error in MLEM_sparseSELL_pycuda: {type(e).__name__}: {e}")
+        gc.collect()
+        return None, None
+    
+    finally:
+        # --- CONTEXT FIX: Pop the context ---
+        if SMatrix and hasattr(SMatrix, 'ctx') and SMatrix.ctx:
+            SMatrix.ctx.pop()
+        # ------------------------------------

AOT_biomaps/AOT_Recon/AOT_Optimizers/PDHG.py

@@ -27,7 +27,7 @@ def CP_TV(
     Works on both CPU and GPU.
     Args:
         SMatrix: System matrix (shape: T, Z, X, N)
-        y: Measurement data (shape: T, X, N)
+        y: Measurement data (shape: T, N)
         alpha: Regularization parameter for TV
         theta: Relaxation parameter (1.0 for standard Chambolle-Pock)
         numIterations: Number of iterations

AOT_biomaps/AOT_Recon/AOT_PotentialFunctions/RelativeDifferences.py

@@ -9,26 +9,24 @@ def _Omega_RELATIVE_DIFFERENCE_CPU(theta_flat, index, values, gamma):
     theta_k = theta_flat[k_idx]
     diff = theta_k - theta_j
     abs_diff = np.abs(diff)
-
     denom = theta_k + theta_j + gamma * abs_diff + 1e-8
     num = diff ** 2
-
+    psi_pair = num / denom
+    psi_pair = values * psi_pair
     # First derivative ∂U/∂θ_j
     dpsi = (2 * diff * denom - num * (1 + gamma * np.sign(diff))) / (denom ** 2)
     grad_pair = values * (-dpsi)  # Note the negative sign: U contains ψ(θ_k, θ_j), seeking ∂/∂θ_j
-
     # Second derivative ∂²U/∂θ_j² (numerically stable, approximate treatment)
     d2psi = (2 * denom ** 2 - 4 * diff * denom * (1 + gamma * np.sign(diff))
                 + 2 * num * (1 + gamma * np.sign(diff)) ** 2) / (denom ** 3 + 1e-8)
     hess_pair = values * d2psi
-
     grad_U = np.zeros_like(theta_flat)
     hess_U = np.zeros_like(theta_flat)
-
     np.add.at(grad_U, j_idx, grad_pair)
     np.add.at(hess_U, j_idx, hess_pair)
-
-    return grad_U, hess_U
+    # Compute U_value
+    U_value = 0.5 * np.sum(psi_pair)
+    return grad_U, hess_U, U_value
 
 def _Omega_RELATIVE_DIFFERENCE_GPU(theta_flat, index, values, device, gamma):
     j_idx, k_idx = index
@@ -38,26 +36,24 @@ def _Omega_RELATIVE_DIFFERENCE_GPU(theta_flat, index, values, device, gamma):
     abs_diff = torch.abs(diff)
     denom = theta_k + theta_j + gamma * abs_diff + 1e-8
     num = diff ** 2
-
+    psi_pair = num / denom
+    psi_pair = values * psi_pair
     # Compute gradient contributions
     dpsi = (2 * diff * denom - num * (1 + gamma * torch.sign(diff))) / (denom ** 2)
     grad_pair = values * (-dpsi)
-
     # Compute Hessian contributions
     d2psi = (2 * denom ** 2 - 4 * diff * denom * (1 + gamma * torch.sign(diff))
              + 2 * num * (1 + gamma * torch.sign(diff)) ** 2) / (denom ** 3 + 1e-8)
     hess_pair = values * d2psi
-
     # Initialize gradient and Hessian on the correct device
     grad_U = torch.zeros_like(theta_flat, device=device)
     hess_U = torch.zeros_like(theta_flat, device=device)
-
     # Accumulate gradient contributions
     grad_U.index_add_(0, j_idx, grad_pair)
     grad_U.index_add_(0, k_idx, -grad_pair)
-
     # Accumulate Hessian contributions
     hess_U.index_add_(0, j_idx, hess_pair)
     hess_U.index_add_(0, k_idx, hess_pair)
-
-    return grad_U, hess_U
+    # Compute U_value
+    U_value = 0.5 * psi_pair.sum()
+    return grad_U, hess_U, U_value