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

AOT_biomaps/AOT_Recon/AOT_Optimizers/MLEM.py

@@ -1,26 +1,86 @@
-from AOT_biomaps.AOT_Recon.ReconTools import _forward_projection, _backward_projection
+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
-
-
-def _MLEM_GPU_basic(SMatrix, y, numIterations, isSavingEachIteration, withTumor):
+import cupy as cp
+import cupyx.scipy.sparse as cpsparse
+import gc
+import pycuda.driver as drv
+
+
+def MLEM(
+    SMatrix,
+    y,
+    numIterations=100,
+    isSavingEachIteration=True,
+    withTumor=True,
+    device=None,
+    use_numba=False,
+    denominator_threshold=1e-6,
+    max_saves=5000,
+    show_logs=True,
+    smatrixType=SMatrixType.SELL,
+    Z=350,
+):
     """
-    Optimized MLEM using GPU — saves at most 1000 images, equally spaced.
-    Returns both the images and their iteration indices.
+    Unified MLEM algorithm for Acousto-Optic Tomography.
+    Works on CPU (basic, multithread, optimized) and GPU (single or multi-GPU).
+    Args:
+        SMatrix: System matrix (shape: T, Z, X, N)
+        y: Measurement data (shape: T, N)
+        numIterations: Number of iterations
+        isSavingEachIteration: If True, saves intermediate results
+        withTumor: Boolean for description only
+        device: Torch device (auto-selected if None)
+        use_multi_gpu: If True and GPU available, uses all GPUs
+        use_numba: If True and on CPU, uses multithreaded Numba
+        max_saves: Maximum number of intermediate saves (default: 5000)
+    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), show_logs=show_logs):
+            device = torch.device(f"cuda:{config.select_best_gpu()}")
+            use_gpu = True
+        else:
+            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:
+                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:
+            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, denominator_threshold, show_logs=True):
     try:
-        device = torch.device(f"cuda:{config.select_best_gpu()}")
         eps = torch.finfo(torch.float32).eps
         T, Z, X, N = SMatrix.shape
         ZX = Z * X
         TN = T * N
-
-        # Conversion to flattened matrix
         A_flat = (
             torch.from_numpy(SMatrix)
             .to(device=device, dtype=torch.float32)
@@ -36,114 +96,78 @@ def _MLEM_GPU_basic(SMatrix, y, numIterations, isSavingEachIteration, withTumor)
             .sum(dim=(0, 3))
             .reshape(-1)
         )
-
-        description = f"AOT-BioMaps -- ML-EM ---- {'WITH' if withTumor else 'WITHOUT'} TUMOR ---- GPU {torch.cuda.current_device()}"
+        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))
+        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 = []
         saved_indices = []
-
-        if isSavingEachIteration:
-            # Calculate step to save at most 1000 images
-            step = max(1, numIterations // 1000)
-            save_count = 0
-
         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 % step == 0 or it == numIterations - 1):
+                if isSavingEachIteration and it in save_indices:
                     saved_theta.append(theta_flat.reshape(Z, X).clone())
                     saved_indices.append(it)
-                    save_count += 1
-                    if save_count >= 1000:  # Hard limit
-                        break
-
         # Free memory
         del A_flat, y_flat, norm_factor_flat
         torch.cuda.empty_cache()
-
         if isSavingEachIteration:
             return [t.cpu().numpy() for t in saved_theta], saved_indices
         else:
             return theta_flat.reshape(Z, X).cpu().numpy(), None
-
     except Exception as e:
-        print("Error in MLEM:", type(e).__name__, ":", e)
+        print(f"Error in single-GPU MLEM: {type(e).__name__}: {e}")
         torch.cuda.empty_cache()
-        return None, []
-
-def _MLEM_CPU_basic(SMatrix, y, numIterations, isSavingEachIteration, withTumor):
-    try:
-        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]))
-        matrix_theta = [theta_p_0]
-        saved_indices = [0]
-        normalization_factor = np.sum(SMatrix, axis=(0, 3))
-        description = f"AOT-BioMaps -- ML-EM ---- {'WITH' if withTumor else 'WITHOUT'} TUMOR ---- processing on single CPU (basic) ----"
-
-        if isSavingEachIteration:
-            step = max(1, (numIterations - 1) // 999)
-            save_count = 1
-
-        for it in trange(numIterations, desc=description):
-            theta_p = matrix_theta[-1]
-            for _t in range(SMatrix.shape[0]):
-                for _n in range(SMatrix.shape[3]):
-                    q_p[_t, _n] = np.sum(SMatrix[_t, :, :, _n] * theta_p)
-            e_p = y / (q_p + 1e-8)
-            for _z in range(SMatrix.shape[1]):
-                for _x in range(SMatrix.shape[2]):
-                    c_p[_z, _x] = np.sum(SMatrix[:, _z, _x, :] * e_p)
-            theta_p_plus_1 = theta_p / (normalization_factor + 1e-8) * c_p
-
-            if isSavingEachIteration and (it % step == 0 or it == numIterations - 1):
-                matrix_theta.append(theta_p_plus_1)
-                saved_indices.append(it + 1)
-                save_count += 1
-                if save_count >= 1000:
-                    break
-            else:
-                matrix_theta[-1] = theta_p_plus_1
-
-        if not isSavingEachIteration:
-            return matrix_theta[-1], None
-        else:
-            return matrix_theta, saved_indices
-    except Exception as e:
-        print("Error in basic CPU MLEM:", type(e).__name__, ":", e)
         return None, None
 
-def _MLEM_CPU_multi(SMatrix, y, numIterations, isSavingEachIteration, withTumor):
+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))
-        description = f"AOT-BioMaps -- ML-EM ---- {'WITH' if withTumor else 'WITHOUT'} TUMOR ---- processing on multithread CPU ({numba.config.NUMBA_DEFAULT_NUM_THREADS} threads)----"
+        normalization_factor = np.sum(SMatrix, axis=(0, 3)).astype(np.float32)
 
-        if isSavingEachIteration:
-            step = max(1, (numIterations - 1) // 999)
-            save_count = 1
+        # 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)
 
-        for it in trange(numIterations, desc=description):
+        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 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 % step == 0 or it == numIterations - 1):
+            if isSavingEachIteration and (it + 1) in save_indices:
                 matrix_theta.append(theta_p_plus_1)
                 saved_indices.append(it + 1)
-                save_count += 1
-                if save_count >= 1000:
-                    break
             else:
                 matrix_theta[-1] = theta_p_plus_1
 
@@ -152,10 +176,10 @@ def _MLEM_CPU_multi(SMatrix, y, numIterations, isSavingEachIteration, withTumor)
         else:
             return matrix_theta, saved_indices
     except Exception as e:
-        print("Error in multi-threaded CPU MLEM:", type(e).__name__, ":", e)
+        print(f"Error in Numba CPU MLEM: {type(e).__name__}: {e}")
         return None, None
-   
-def _MLEM_CPU_opti(SMatrix, y, numIterations, isSavingEachIteration, withTumor):
+
+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)
@@ -165,27 +189,35 @@ def _MLEM_CPU_opti(SMatrix, y, numIterations, isSavingEachIteration, withTumor):
         saved_indices = [0]
         normalization_factor = np.sum(SMatrix, axis=(0, 3)).astype(np.float32)
         normalization_factor_flat = normalization_factor.reshape(-1)
-        description = f"AOT-BioMaps -- ML-EM ---- {'WITH' if withTumor else 'WITHOUT'} TUMOR ---- processing on single CPU (optimized) ----"
 
-        if isSavingEachIteration:
-            step = max(1, (numIterations - 1) // 999)
-            save_count = 1
+        # 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)
+
+        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 % step == 0 or it == numIterations - 1):
+            if isSavingEachIteration and (it + 1) in save_indices:
                 matrix_theta.append(theta_p_plus_1)
                 saved_indices.append(it + 1)
-                save_count += 1
-                if save_count >= 1000:
-                    break
             else:
                 matrix_theta[-1] = theta_p_plus_1
 
@@ -194,67 +226,241 @@ def _MLEM_CPU_opti(SMatrix, y, numIterations, isSavingEachIteration, withTumor):
         else:
             return matrix_theta, saved_indices
     except Exception as e:
-        print("Error in optimized CPU MLEM:", type(e).__name__, ":", 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)
 
-def _MLEM_GPU_multi(SMatrix, y, numIterations, isSavingEachIteration, withTumor):
+    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:
-        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)
-        saved_theta = [theta_0.cpu().numpy()]
-        saved_indices = [0]
-        description = f"AOT-BioMaps -- ML-EM ---- {'WITH' if withTumor else 'WITHOUT'} TUMOR ---- processing on multi-GPU ----"
+        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)
 
-        if isSavingEachIteration:
-            step = max(1, (numIterations - 1) // 999)
-            save_count = 1
-
-        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 % step == 0 or it == numIterations - 1):
-                saved_theta.append(torch.stack(theta_p_list).mean(dim=0).cpu().numpy())
-                saved_indices.append(it + 1)
-                save_count += 1
-                if save_count >= 1000:
-                    break
+        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)
 
-        del A_matrix_torch, y_torch, A_split, y_split, theta_0, normalization_factor
-        torch.cuda.empty_cache()
-        for i in range(num_gpus):
-            torch.cuda.empty_cache()
+            # 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)
 
-        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("Error in multi-GPU MLEM:", type(e).__name__, ":", e)
-        del A_matrix_torch, y_torch, A_split, y_split, theta_0, normalization_factor
-        torch.cuda.empty_cache()
-        for i in range(num_gpus):
-            torch.cuda.empty_cache()
+        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()
+        # ------------------------------------