AOT-biomaps 2.9.212__py3-none-any.whl → 2.9.233__py3-none-any.whl

AOT_biomaps/AOT_Recon/AOT_Optimizers/MLEM.py

@@ -5,6 +5,10 @@ import torch
 import numpy as np
 import os
 from tqdm import trange
+import cupy as cp
+import cupyx.scipy.sparse as cpsparse
+import gc
+
 
 def MLEM(
     SMatrix,
@@ -13,9 +17,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,
+    useSparseSMatrix=True,
+    Z=350,
 ):
     """
     Unified MLEM algorithm for Acousto-Optic Tomography.
@@ -37,7 +44,7 @@ def MLEM(
         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)):
+            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:
@@ -47,20 +54,20 @@ def MLEM(
             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)
-            else:
-                return _MLEM_single_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves)
+                if useSparseSMatrix:
+                    return _MLEM_sparseCSR(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device.index, max_saves, denominator_threshold, Z, show_logs)
+                else:
+                    return _MLEM_single_GPU(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, device, max_saves, denominator_threshold,show_logs)
         else:
             if use_numba:
-                return _MLEM_CPU_numba(SMatrix, y, numIterations, isSavingEachIteration, tumor_str, max_saves)
+                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)
+                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 +89,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 +97,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 +124,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 +144,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 +171,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 +192,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 +220,114 @@ 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(
+    SMatrix,
+    y,
+    numIterations,
+    isSavingEachIteration,
+    tumor_str,
+    device_index,
+    max_saves,
+    denominator_threshold,
+    Z,
+    show_logs=True,
+):
+    """
+    MLEM implementation using CuPy with sparse CSR matrix on a single GPU.
+    Caution : SMatrix must be a cupyx.scipy.sparse.csr_matrix please sparse it before using.
+    """
+    try:
+        cp.cuda.Device(device_index).use()
+        dtype = cp.float32
+        eps = cp.finfo(dtype).eps
+
+        # --- Préparation de la matrice et des données ---
+        if not isinstance(SMatrix, cpsparse.csr_matrix):
+            SMatrix = cpsparse.csr_matrix(SMatrix, dtype=dtype)
+        else:
+            SMatrix = SMatrix.astype(dtype)
+
+        if not isinstance(y, cp.ndarray):
+            y_cupy = cp.asarray(y, dtype=dtype)
+        else:
+            y_cupy = y.astype(dtype)
+
+        TN, ZX = SMatrix.shape
+        X = ZX // Z
+
+        # Initialisation du volume reconstruit
+        theta_flat = cp.full(ZX, 0.1, dtype=dtype)
+
+        # Facteur de normalisation
+        norm_factor = cp.maximum(SMatrix.sum(axis=0).ravel(), 1e-6)
+        norm_factor_inv = 1.0 / norm_factor
+
+        # Gestion des indices de sauvegarde
+        if numIterations <= max_saves:
+            save_indices = list(range(numIterations))
+        else:
+            step = max(1, 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 = []
+
+        description = f"AOT-BioMaps -- ML-EM (sparse CSR) ---- {tumor_str} TUMOR ---- GPU {device_index}"
+
+        iterator = trange(numIterations, desc=description, ncols=100) if show_logs else range(numIterations)
+
+        # --- Boucle principale MLEM ---
+        for it in iterator:
+            # Étape 1 : Projection
+            q_flat = SMatrix.dot(theta_flat)
+            q_flat = cp.maximum(q_flat, denominator_threshold)
+
+            # Étape 2 : Rapport y / (A*L)
+            e_flat = y_cupy / q_flat
+
+            # Étape 3 : Rétroprojection (A.T * e)
+            c_flat = SMatrix.T.dot(e_flat)
+
+            # Étape 4 : Mise à jour
+            theta_flat = theta_flat * (norm_factor_inv * c_flat)
+            theta_flat = cp.maximum(theta_flat, 0)
+
+            # Sauvegarde éventuelle
+            if isSavingEachIteration and it in save_indices:
+                saved_theta.append(theta_flat.reshape(Z, X).get())  # transfert CPU
+                saved_indices.append(it)
+
+            # Libération mémoire GPU
+            del q_flat, e_flat, c_flat
+            cp.get_default_memory_pool().free_all_blocks()
+            gc.collect()
+
+            # Vérif convergence toutes les 10 itérations
+            if it % 10 == 0 and show_logs:
+                rel_change = cp.abs(theta_flat - theta_flat).max() / (theta_flat.max() + eps)
+                if rel_change < 1e-4:
+                    print(f"Convergence atteinte à l’itération {it}")
+                    break
+
+        # --- Fin : récupération du résultat ---
+        result = theta_flat.reshape(Z, X).get()  # Retour sur CPU
+        del theta_flat, norm_factor, norm_factor_inv
+        cp.get_default_memory_pool().free_all_blocks()
+        gc.collect()
+
+        if isSavingEachIteration:
+            return saved_theta, saved_indices
+        else:
+            return result, None
+
+    except Exception as e:
+        print(f"Erreur dans _MLEM_single_GPU_CuPy: {type(e).__name__}: {e}")
+        cp.get_default_memory_pool().free_all_blocks()
+        gc.collect()
+        return None, None
+

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