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

AOT_biomaps/AOT_Recon/ReconEnums.py

@@ -60,7 +60,17 @@ class OptimizerType(Enum):
     This optimizer is compatible with both histogram and list-mode data.
     This optimizer is compatible with both emission and transmission data.
     """
-    LS_TV = 'LS_TV'
+    CP_TV = 'CP_TV'
+    """
+    This optimizer implements the Chambolle-Pock algorithm for total variation regularization.
+    It is suitable for problems where the objective function includes a total variation term.
+    It is particularly effective for preserving edges while reducing noise in the reconstructed image.
+    """ 
+    CP_KL = 'CP_KL'
+    """
+    This optimizer implements the Kullback-Leibler divergence for regularization.
+    It is suitable for problems where the objective function includes a Kullback-Leibler divergence term.
+    """
     LS = 'LS'
     """
     This optimizer implements the standard Landweber algorithm for least-squares optimization.
@@ -344,8 +354,33 @@ class NoiseType(Enum):
     - None: No noise is applied.
     """
     POISSON = 'poisson'
-    """Poisson noise, typically used for emission data."""
+    """Poisson noise."""
     GAUSSIAN = 'gaussian'
-    """Gaussian noise, typically used for transmission data."""
+    """Gaussian noise."""
     None_ = 'none'
     """No noise is applied."""
+
+class SMatrixType(Enum):
+    """
+    Enum for different sparsing methods used in reconstructions.
+    
+    Selection of sparsing methods:
+    - Thresholding: Sparsing based on a threshold value.
+    - TopK: Sparsing by retaining the top K values.
+    - None: No sparsing is applied.
+    """
+    DENSE = 'DENSE'
+    """No sparsing is applied."""
+    CSR = 'CSR'
+    """Sparsing based on a threshold value."""
+    COO = 'COO'
+    """Sparsing by retaining the top K values."""
+    SELL = 'SELL'
+    """Sparsing using sell C sigma method.
+    Optimized variant of ELLPACK, dividing the matrix into fixed-size "chunks" of `C` rows.
+    Non-zero elements are sorted by column within each chunk to improve memory coalescing on GPUs.
+    Rows are padded with zeros to align their length to the longest row in the chunk.
+    ** Ref : Kreutzer, M., Hager, G., Wellein, G., Fehske, H., & Bishop, A. R. (2014).
+          "A Unified Sparse Matrix Data Format for Efficient General Sparse Matrix-Vector Multiply on Modern Processors".
+          ACM Transactions on Mathematical Software, 41(2), 1–24. DOI: 10.1145/2592376.
+    """

AOT_biomaps/AOT_Recon/ReconTools.py

@@ -1,9 +1,11 @@
 import os
+from AOT_biomaps.AOT_Recon.AOT_SparseSMatrix import SparseSMatrix_CSR, SparseSMatrix_SELL
 import torch
 import numpy as np
+import pycuda.driver as drv
 from numba import njit, prange
 from torch_sparse import coalesce
-import torch.nn.functional as F
+from scipy.signal.windows import hann
 
 def load_recon(hdr_path):
     """
@@ -78,7 +80,7 @@ def load_recon(hdr_path):
     rescale_offset = float(header.get('data rescale offset', 0))
     image = image * rescale_slope + rescale_offset
     
-    return image.T
+    return image
 
 def mse(y_true, y_pred):
     """
@@ -150,20 +152,82 @@ def ssim(img1, img2, win_size=7, k1=0.01, k2=0.03, L=1.0):
     return np.mean(ssim_map)
 
 def calculate_memory_requirement(SMatrix, y):
-    """Calculate the memory requirement for the given matrices in GB."""
-    num_elements_SMatrix = SMatrix.size
-    num_elements_y = y.size
-    num_elements_theta = SMatrix.shape[1] * SMatrix.shape[2]  # Assuming theta has shape (Z, X)
+    """
+    Calcule la mémoire requise (en Go) pour :
+    - SMatrix : Matrice (np.ndarray, CuPy CSR, SparseSMatrix_CSR ou SparseSMatrix_SELL)
+    - y : vecteur (NumPy ou CuPy, float32)
+
+    Args:
+        SMatrix: Matrix object (np.ndarray, cpsparse.csr_matrix, SparseSMatrix_CSR, or SparseSMatrix_SELL)
+        y: Vector (float32)
+    """
+    total_bytes = 0
+
+    # --- 1. Memory for SMatrix ---
+    
+    # 1.1. Custom Sparse Matrix (SELL/CSR)
+    if isinstance(SMatrix, (SparseSMatrix_SELL, SparseSMatrix_CSR)):
+        # We rely on the getMatrixSize method, which we fixed to track all host/GPU bytes.
+        # This is the most reliable way to estimate memory for custom GPU-backed structures.
+        try:
+            matrix_size_gb = SMatrix.getMatrixSize()
+            if isinstance(matrix_size_gb, dict) and 'error' in matrix_size_gb:
+                raise ValueError(f"SMatrix allocation error: {matrix_size_gb['error']}")
+            
+            # Convert GB back to bytes (1 GB = 1024^3 bytes)
+            size_SMatrix = matrix_size_gb * (1024 ** 3)
+            total_bytes += size_SMatrix
+            print(f"SMatrix (Custom Sparse) size: {matrix_size_gb:.3f} GB")
+
+        except AttributeError:
+            raise AttributeError("Custom Sparse Matrix must implement the getMatrixSize() method.")
+    
+    # 1.2. NumPy Dense Array (Standard)
+    elif isinstance(SMatrix, np.ndarray):
+        # Dense NumPy array (float32)
+        size_SMatrix = SMatrix.nbytes
+        total_bytes += size_SMatrix
+        print(f"SMatrix (NumPy Dense) size: {size_SMatrix / (1024 ** 3):.3f} GB")
+
+    # 1.3. CuPy CSR Matrix (Standard Sparse CuPy)
+    # Note: Requires CuPy to be imported, which is usually done outside this function.
+    # Assuming 'cpsparse.csr_matrix' is available in the environment if this path is taken.
+    elif 'cupy.sparse' in str(type(SMatrix)): # Using string check for type safety outside CuPy context
+        # CuPy CSR matrix structure: data (float32), indices (int32), indptr (int32)
+        nnz = SMatrix.nnz
+        num_rows = SMatrix.shape[0]
+        size_data = nnz * 4        # float32 = 4 bytes
+        size_indices = nnz * 4     # int32 = 4 bytes
+        size_indptr = (num_rows + 1) * 4 # int32 = 4 bytes
+        size_SMatrix = size_data + size_indices + size_indptr
+        total_bytes += size_SMatrix
+        print(f"SMatrix (CuPy CSR) size: {size_SMatrix / (1024 ** 3):.3f} GB")
+
+    else:
+        raise ValueError("SMatrix must be a np.ndarray, cpsparse.csr_matrix, or a custom SparseSMatrix object (CSR/SELL).")
+
+    # --- 2. Memory for Vector y ---
+    
+    # Check if y is a CuPy array or NumPy array (assuming float32 based on docstring)
+    if hasattr(y, 'nbytes'):
+        size_y = y.nbytes
+        total_bytes += size_y
+        print(f"Vector y size: {size_y / (1024 ** 3):.3f} GB")
+    else:
+        # Fallback if object doesn't expose nbytes (e.g., custom buffer), but usually array objects do.
+        raise ValueError("Vector y must be an array type exposing the .nbytes attribute.")
+
 
-    # Calculate total memory requirement in GB
-    total_memory = (num_elements_SMatrix + num_elements_y + num_elements_theta) * 32 / 8 / 1024**3
-    return total_memory
+    # --- 3. Final Result ---
+    return total_bytes / (1024 ** 3)
 
-def check_gpu_memory(device_index, required_memory):
+
+def check_gpu_memory(device_index, required_memory, show_logs=True):
     """Check if enough memory is available on the specified GPU."""
-    free_memory, total_memory = torch.cuda.mem_get_info(f"cuda:{device_index}")
+    free_memory, _ = torch.cuda.mem_get_info(f"cuda:{device_index}")
     free_memory_gb = free_memory / 1024**3
-    print(f"Free memory on GPU {device_index}: {free_memory_gb:.2f} GB, Required memory: {required_memory:.2f} GB")
+    if show_logs:
+        print(f"Free memory on GPU {device_index}: {free_memory_gb:.2f} GB, Required memory: {required_memory:.2f} GB")
     return free_memory_gb >= required_memory
 
 @njit(parallel=True)
@@ -210,86 +274,53 @@ def _build_adjacency_sparse(Z, X, device, corner=(0.5 - np.sqrt(2) / 4) / np.sqr
     return index, values
 
 
-def power_method(P, PT, data, Z, X, n_it=10, isGPU=False):
-    x = PT(data)
-    x = x.reshape(Z, X)
+def power_method(P, PT, data, Z, X, n_it=10):
+    x = torch.randn(Z * X, device=data.device)
+    x = x / torch.norm(x)
     for _ in range(n_it):
-        grad = gradient_gpu(x) if isGPU else gradient_cpu(x)
-        div = div_gpu(grad) if isGPU else div_cpu(grad)
-        x = PT(P(x.ravel())) - div.ravel()
-        s = torch.sqrt(torch.sum(x**2))
-        x /= s
-        x = x.reshape(Z, X)
-    return torch.sqrt(s)
+        Ax = P(x)
+        ATax = PT(Ax)
+        x = ATax / torch.norm(ATax)
+    ATax = PT(P(x))
+    return torch.sqrt(torch.dot(x, ATax))
 
 def proj_l2(p, alpha):
-    norm = torch.sqrt(torch.sum(p**2, dim=0, keepdim=True))
-    return p * alpha / torch.max(norm, torch.tensor(alpha, device=p.device))
+    if alpha <= 0:
+        return torch.zeros_like(p)
+    norm = torch.sqrt(torch.sum(p**2, dim=0, keepdim=True) + 1e-12)
+    return p * torch.min(norm, torch.tensor(alpha, device=p.device)) / (norm + 1e-12)
 
-def norm2sq(x):
-    return torch.sum(x**2)
-
-def norm1(x):
-    return torch.sum(torch.abs(x))
-
-def gradient_cpu(x):
+def gradient(x):
     grad_x = torch.zeros_like(x)
     grad_y = torch.zeros_like(x)
-
-    grad_x[:-1, :] = x[1:, :] - x[:-1, :]
-    grad_y[:, :-1] = x[:, 1:] - x[:, :-1]
-
+    grad_x[:, :-1] = x[:, 1:] - x[:, :-1]  # Gradient horizontal
+    grad_y[:-1, :] = x[1:, :] - x[:-1, :]   # Gradient vertical
     return torch.stack((grad_x, grad_y), dim=0)
 
-def div_cpu(x):
+def div(x):
     if x.dim() == 3:
-        x = x.unsqueeze(0)  # Devient [1, 2, H, W]
+        x = x.unsqueeze(0)  # Ajoute une dimension batch si nécessaire
 
-    gx = x[:, 0:1, :, :]  # gradient horizontal
-    gy = x[:, 1:2, :, :]  # gradient vertical
+    gx = x[:, 0, :, :]  # Gradient horizontal (shape: [1, H, W] ou [H, W])
+    gy = x[:, 1, :, :]  # Gradient vertical   (shape: [1, H, W] ou [H, W])
 
-    # Définition des noyaux de divergence
-    kernel_x = torch.tensor([[[[1.0], [-1.0]]]], dtype=torch.float32)
-    kernel_y = torch.tensor([[[[1.0, -1.0]]]], dtype=torch.float32)
+    # Divergence du gradient horizontal (gx)
+    div_x = torch.zeros_like(gx)
+    div_x[:, :, 1:] += gx[:, :, :-1]  # Contribution positive (gauche)
+    div_x[:, :, :-1] -= gx[:, :, :-1] # Contribution négative (droite)
 
-    # Appliquer la convolution
-    div_x = F.conv2d(gx, kernel_x, padding=(1, 0))
-    div_y = F.conv2d(gy, kernel_y, padding=(0, 1))
+    # Divergence du gradient vertical (gy)
+    div_y = torch.zeros_like(gy)
+    div_y[:, 1:, :] += gy[:, :-1, :]  # Contribution positive (haut)
+    div_y[:, :-1, :] -= gy[:, :-1, :] # Contribution négative (bas)
 
-    # Rogner pour avoir la même taille (H, W)
-    H, W = x.shape[2:]
-    div_x = div_x[:, :, :H, :]
-    div_y = div_y[:, :, :, :W]
-
-    return -(div_x + div_y).squeeze()
-
-def gradient_gpu(x):
-    grad_x = torch.zeros_like(x)
-    grad_y = torch.zeros_like(x)
-    grad_x[:-1, :] = x[1:, :] - x[:-1, :]
-    grad_y[:, :-1] = x[:, 1:] - x[:, :-1]
-    return torch.stack((grad_x, grad_y), dim=0)
+    return -(div_x + div_y)
 
-def div_gpu(x):
-    if x.dim() == 3:
-        x = x.unsqueeze(0)  # Devient [1, 2, H, W]
-    gx = x[:, 0:1, :, :]  # gradient horizontal
-    gy = x[:, 1:2, :, :]  # gradient vertical
-
-    # Définition des noyaux de divergence
-    kernel_x = torch.tensor([[[[1.0], [-1.0]]]], dtype=torch.float32, device=x.device)
-    kernel_y = torch.tensor([[[[1.0, -1.0]]]], dtype=torch.float32, device=x.device)
-
-    # Appliquer la convolution
-    div_x = F.conv2d(gx, kernel_x, padding=(1, 0))
-    div_y = F.conv2d(gy, kernel_y, padding=(0, 1))
-
-    # Rogner pour avoir la même taille (H, W)
-    H, W = x.shape[2:]
-    div_x = div_x[:, :, :H, :]
-    div_y = div_y[:, :, :, :W]
+def norm2sq(x):
+    return torch.sum(x**2)
 
-    return -(div_x + div_y).squeeze()
+def norm1(x):
+    return torch.sum(torch.abs(x))
 
 def KL_divergence(Ax, y):
     return torch.sum(Ax - y * torch.log(Ax + 1e-10))
@@ -303,3 +334,79 @@ def prox_F_star(y, sigma, a):
 def prox_G(x, tau, K):
     return torch.clamp(x - tau * K, min=0)
 
+def filter_radon(f, N, filter_type, Fc):
+    """
+    Implémente les filtres pour la rétroprojection filtrée (iRadon).
+    Inspirée de la fonction MATLAB FilterRadon de Mamouna Bocoum.
+
+    Paramètres :
+    ------------
+    f : np.ndarray
+        Vecteur des fréquences (ex: f_t ou f_z).
+    N : int
+        Taille du filtre (longueur de f).
+    filter_type : str
+        Type de filtre : 'ram-lak', 'shepp-logan', 'cosine', 'hamming', 'hann'.
+    Fc : float
+        Fréquence de coupure.
+
+    Retourne :
+    -----------
+    FILTER : np.ndarray
+        Filtre appliqué aux fréquences.
+    """
+    FILTER = np.abs(f)
+
+    if filter_type == 'ram-lak':
+        pass  # FILTER = |f| (déjà calculé)
+    elif filter_type == 'shepp-logan':
+        # Évite la division par zéro
+        with np.errstate(divide='ignore', invalid='ignore'):
+            FILTER = FILTER * (np.sinc(2 * f / (2 * Fc)))  # sin(2πf/(2Fc))/(2πf/(4Fc)) = sinc(2f/(2Fc))
+        FILTER[np.isnan(FILTER)] = 1.0  # Pour f=0
+    elif filter_type == 'cosine':
+        FILTER = FILTER * np.cos(2 * np.pi * f / (4 * Fc))
+    elif filter_type == 'hamming':
+        FILTER = FILTER * (0.54 + 0.46 * np.cos(2 * np.pi * f / Fc))
+    elif filter_type == 'hann':
+        FILTER = FILTER * (1 + np.cos(2 * np.pi * f / (4 * Fc))) / 2
+    else:
+        raise ValueError(f"Type de filtre inconnu : {filter_type}")
+
+    # Coupure des fréquences au-delà de Fc
+    FILTER[np.abs(f) > Fc] = 0
+    # Atténuation exponentielle (optionnelle, comme dans le code MATLAB)
+    FILTER = FILTER * np.exp(-2 * (np.abs(f) / Fc)**10)
+
+    return FILTER
+
+
+def get_apodization_vector_gpu(matrix_sparse_obj):
+    """
+    Génère un vecteur de fenêtrage 2D (Hanning) pour l'apodisation 
+    de la matrice système A et le transfère sur le GPU.
+    Ce vecteur doit être multiplié par les colonnes de A (pixels Z*X).
+    """
+    Z = matrix_sparse_obj.Z
+    X = matrix_sparse_obj.X
+    
+    # 1. Génération des fenêtres 1D sur l'axe X et Z
+    # Forte apodisation latérale (X) pour cibler l'artefact de bordure.
+    fenetre_x = hann(X).astype(np.float32)
+    
+    # Fenêtre uniforme en profondeur (Z), car l'artefact est surtout latéral.
+    fenetre_z = np.ones(Z, dtype=np.float32)
+    
+    # 2. Création de la matrice de fenêtre 2D (Z, X)
+    fenetre_2d = np.outer(fenetre_z, fenetre_x)
+    
+    # 3. Vectorisation (Z*X)
+    fenetre_vectorisee = fenetre_2d.flatten()
+    
+    # 4. Transfert sur GPU (mémoire contiguë)
+    fenetre_gpu = drv.mem_alloc(fenetre_vectorisee.nbytes)
+    drv.memcpy_htod(fenetre_gpu, fenetre_vectorisee)
+    
+    print(f"✅ Vecteur de fenêtrage (Z*X={Z*X}) généré et transféré sur GPU.")
+    
+    return fenetre_gpu

AOT_biomaps/AOT_Recon/__init__.py

@@ -9,3 +9,4 @@ from .ReconTools import *
 
 
 
+