AOT-biomaps 2.9.186__py3-none-any.whl → 2.9.294__py3-none-any.whl

AOT_biomaps/AOT_Recon/AOT_SparseSMatrix/SparseSMatrix_CSR.py

@@ -0,0 +1,274 @@
+import pycuda.driver as drv
+import numpy as np
+from pycuda.compiler import SourceModule
+from tqdm import trange
+import gc
+import os
+
+class SparseSMatrix_CSR:
+    """Construction d'une matrice CSR à partir d'un objet `manip`.
+    Usage:
+        S = SparseMatrixGPU(manip)
+        S.allocate()
+    Après allocate(), on a: row_ptr (host np.int64 array), row_ptr_gpu (device ptr),
+    h_col_ind, h_values, col_ind_gpu, values_gpu, norm_factor_inv.
+    """
+
+    def __init__(self, manip, block_rows=64, relative_threshold=0.3, device=0):
+        drv.init()
+        self.device = drv.Device(device)
+        self.ctx = self.device.make_context()
+        self.manip = manip
+        self.N = len(manip.AcousticFields)
+        self.T = manip.AcousticFields[0].field.shape[0]
+        self.Z = manip.AcousticFields[0].field.shape[1]
+        self.X = manip.AcousticFields[0].field.shape[2]
+        self.block_rows = block_rows
+        self.relative_threshold = relative_threshold
+        
+        # --- FIX: Résolution du chemin du .cubin (dans AOT_Recon/) ---
+        # Le fichier SparseSMatrix_CSR.py est dans AOT_Recon/AOT_SparseSMatrix/
+        # On remonte d'un répertoire pour atteindre AOT_Recon/
+        cubin_parent_dir = os.path.dirname(os.path.dirname(__file__))
+        self.module_path = os.path.join(cubin_parent_dir, "AOT_biomaps_kernels.cubin")
+        # --- FIN FIX ---
+
+        self.h_dense = None
+        self.row_ptr = None
+        self.row_ptr_gpu = None
+        self.h_col_ind = None
+        self.h_values = None
+        self.total_nnz = 0
+        self.norm_factor_inv = None
+        self.sparse_mod = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.free()
+
+    def load_precompiled_module(self):
+        """
+        Charge le module CUDA pré-compilé (.cubin) en utilisant le chemin résolu.
+        Supprime la logique de compilation JIT.
+        """
+        so_path = self.module_path # Utilise le chemin résolu dans __init__
+
+        if not os.path.exists(so_path):
+            raise FileNotFoundError(
+                f"Le module CUDA {os.path.basename(so_path)} est introuvable au chemin: {so_path}. "
+                "Assurez-vous qu'il est compilé et bien placé."
+            )
+            
+        try:
+            self.sparse_mod = drv.module_from_file(so_path)
+            print(f"✅ Module CUDA chargé depuis {so_path}")
+        except Exception as e:
+             raise RuntimeError(f"Le fichier {os.path.basename(so_path)} a été trouvé, mais PyCUDA n'a pas pu le charger. Vérifiez la compatibilité.") from e
+
+    def estimate_nnz_cpu(self):
+        """Estimation rapide (non-exacte) — utile si tu veux une estimation faible.
+        Recommandé : utiliser la passe GPU exacte (count_nnz_per_row_kernel) à la place.
+        """
+        total = 0
+        for n in range(self.N):
+            field = self.manip.AcousticFields[n].field
+            for t in range(self.T):
+                row = field[t].flatten()
+                row_max = np.max(np.abs(row))
+                thr = row_max * self.relative_threshold
+                total += np.count_nonzero(np.abs(row) > thr)
+        return int(total)
+
+    def allocate(self, kernel_module_path=None):
+        try:
+            # --- 1. Construction bloc par bloc (sans garder tout le dense si possible) ---
+            num_rows = self.N * self.T
+            num_cols = self.Z * self.X
+            bytes_float = np.dtype(np.float32).itemsize
+
+            # Charge module
+            # FIX: Toujours charger depuis self.module_path (résolu)
+            self.load_precompiled_module() 
+
+            count_nnz_kernel = self.sparse_mod.get_function('count_nnz_rows_kernel')
+            fill_csr_kernel = self.sparse_mod.get_function('fill_kernel__CSR')
+
+            # allocate host row_ptr
+            self.row_ptr = np.zeros(num_rows + 1, dtype=np.int64)
+
+            # GPU temp buffers
+            dense_block_host = np.empty((self.block_rows, num_cols), dtype=np.float32)
+            dense_block_gpu = drv.mem_alloc(self.block_rows * num_cols * bytes_float)
+            row_nnz_gpu = drv.mem_alloc(self.block_rows * np.dtype(np.int32).itemsize)
+
+            block_size = 128
+
+            # --- Count NNZ per row using GPU kernel to be consistent with filling logic ---
+            for b in trange(0, num_rows, self.block_rows, desc='Comptage NNZ'):
+                current_rows = min(self.block_rows, num_rows - b)
+                # Fill dense_block_host from manip
+                for r in range(current_rows):
+                    global_row = b + r
+                    n_idx = global_row // self.T
+                    t_idx = global_row % self.T
+                    dense_block_host[r, :] = self.manip.AcousticFields[n_idx].field[t_idx].flatten()
+                drv.memcpy_htod(dense_block_gpu, dense_block_host)
+
+                grid = ((current_rows + block_size - 1) // block_size, 1, 1)
+                # Note: Assuming 'count_nnz_per_row_kernel' is the correct name (verified by user in prior steps)
+                count_nnz_kernel(dense_block_gpu, row_nnz_gpu,
+                                 np.int32(current_rows), np.int32(num_cols),
+                                 np.float32(self.relative_threshold),
+                                 block=(block_size, 1, 1), grid=grid)
+
+                row_nnz_host = np.empty(current_rows, dtype=np.int32)
+                drv.memcpy_dtoh(row_nnz_host, row_nnz_gpu)
+                self.row_ptr[b + 1:b + current_rows + 1] = self.row_ptr[b] + np.cumsum(row_nnz_host, dtype=np.int64)
+
+            # total nnz
+            self.total_nnz = int(self.row_ptr[-1])
+            print(f"NNZ total : {self.total_nnz}")
+
+            # allocate final arrays
+            self.h_col_ind = np.zeros(self.total_nnz, dtype=np.uint32)
+            self.h_values = np.zeros(self.total_nnz, dtype=np.float32)
+
+            # copy row_ptr to device once
+            self.row_ptr_gpu = drv.mem_alloc(self.row_ptr.nbytes)
+            drv.memcpy_htod(self.row_ptr_gpu, self.row_ptr)
+
+            # allocate device arrays for final csr
+            self.col_ind_gpu = drv.mem_alloc(self.h_col_ind.nbytes)
+            self.values_gpu = drv.mem_alloc(self.h_values.nbytes)
+
+            # --- Fill CSR per-block ---
+            for b in trange(0, num_rows, self.block_rows, desc='Remplissage CSR'):
+                current_rows = min(self.block_rows, num_rows - b)
+                for r in range(current_rows):
+                    global_row = b + r
+                    n_idx = global_row // self.T
+                    t_idx = global_row % self.T
+                    dense_block_host[r, :] = self.manip.AcousticFields[n_idx].field[t_idx].flatten()
+                drv.memcpy_htod(dense_block_gpu, dense_block_host)
+
+                grid = ((current_rows + block_size - 1) // block_size, 1, 1)
+                fill_csr_kernel(dense_block_gpu,
+                                self.row_ptr_gpu,
+                                self.col_ind_gpu,
+                                self.values_gpu,
+                                np.int32(b),
+                                np.int32(current_rows),
+                                np.int32(num_cols),
+                                np.float32(self.relative_threshold),
+                                np.int64(self.total_nnz),
+                                block=(block_size, 1, 1), grid=grid)
+                drv.Context.synchronize()
+
+            # copy back
+            drv.memcpy_dtoh(self.h_col_ind, self.col_ind_gpu)
+            drv.memcpy_dtoh(self.h_values, self.values_gpu)
+            print('CSR généré ✔')
+
+            # compute normalization factor from CSR (sum per column)
+            self.compute_norm_factor_from_csr()
+
+            # free temporaries
+            dense_block_gpu.free(); row_nnz_gpu.free()
+
+        except Exception as e:
+            print(f"❌ Erreur détaillée : {e}")
+            self.free()
+            raise
+
+    def compute_norm_factor_from_csr(self):
+        ZX = self.Z * self.X
+
+        # 1) Allouer un vecteur de somme colonne sur le GPU
+        col_sum_gpu = drv.mem_alloc(ZX * np.dtype(np.float32).itemsize)
+        drv.memset_d32(col_sum_gpu, 0, ZX)
+
+        # 2) Récupérer le kernel
+        # FIX: Utiliser le nom générique 'accumulate_columns_atomic' comme dans SELL (si le binaire est partagé)
+        # Si le développeur utilise la convention __CSR, on la garde.
+        # Basé sur notre historique SELL, le nom est probablement générique 'accumulate_columns_atomic'.
+        # Je vais supposer que le nom est générique pour éviter une LogicError ici aussi.
+        acc_kernel = self.sparse_mod.get_function("accumulate_columns_atomic") 
+
+        # 3) Lancer le kernel
+        threads = 256
+        blocks = (self.total_nnz + threads - 1) // threads
+
+        acc_kernel(
+            self.values_gpu,
+            self.col_ind_gpu,
+            np.int64(self.total_nnz),
+            col_sum_gpu,
+            block=(threads,1,1),
+            grid=(blocks,1,1)
+        )
+        drv.Context.synchronize()
+
+        # 4) Récupérer le résultat
+        norm = np.empty(ZX, dtype=np.float32)
+        drv.memcpy_dtoh(norm, col_sum_gpu)
+        col_sum_gpu.free()
+
+        norm = np.maximum(norm.astype(np.float64), 1e-6)
+        self.norm_factor_inv = (1.0 / norm).astype(np.float32)
+
+        self.norm_factor_inv_gpu = drv.mem_alloc(self.norm_factor_inv.nbytes)
+        drv.memcpy_htod(self.norm_factor_inv_gpu, self.norm_factor_inv)
+
+    def getMatrixSize(self):
+        """
+        Retourne la taille totale de la matrice CSR en Go (en sommant la mémoire GPU).
+        Utilise les attributs de taille stockés pour contourner l'AttributeError de DeviceAllocation.
+        """
+        # Note: L'utilisateur doit s'assurer que self.row_ptr existe avant cet appel.
+        if self.row_ptr is None:
+            return {"error": "La matrice sparse n'est pas encore allouée."}
+            
+        total_bytes = 0
+
+        # Somme des tailles stockées (Taille calculée et attribuée dans allocate et compute_norm_factor_from_csr)
+        total_bytes += getattr(self, 'row_ptr_gpu_size', 0)
+        total_bytes += getattr(self, 'col_ind_gpu_size', 0)
+        total_bytes += getattr(self, 'values_gpu_size', 0)
+        total_bytes += getattr(self, 'norm_factor_inv_gpu_size', 0)
+        
+        return total_bytes / (1024**3)
+
+    def free(self):
+        try:
+            if hasattr(self, 'col_ind_gpu') and self.col_ind_gpu:
+                self.col_ind_gpu.free()
+            if hasattr(self, 'values_gpu') and self.values_gpu:
+                self.values_gpu.free()
+            if hasattr(self, 'row_ptr_gpu') and self.row_ptr_gpu:
+                self.row_ptr_gpu.free()
+            if hasattr(self, 'norm_factor_inv_gpu') and self.norm_factor_inv_gpu:
+                self.norm_factor_inv_gpu.free()
+            if hasattr(self, 'ctx') and self.ctx:
+                try:
+                    self.ctx.pop()
+                except Exception:
+                    pass
+                self.ctx = None
+            print('✅ Mémoire GPU libérée.')
+        except Exception as e:
+            print(f"❌ Erreur lors de la libération de la mémoire GPU : {e}")
+        
+    def compute_density(self):
+        """
+        Retourne la densité réelle de la CSR = NNZ / (num_rows * num_cols)
+        Nécessite que self.h_values et self.row_ptr existent (host).
+        """
+        if self.row_ptr is None or self.h_values is None:
+            raise RuntimeError("row_ptr et h_values requis pour calculer la densité")
+        num_rows = int(self.N * self.T)
+        num_cols = int(self.Z * self.X)
+        total_nnz = int(self.row_ptr[-1])
+        density = total_nnz / (num_rows * num_cols)
+        return density

AOT_biomaps/AOT_Recon/AOT_SparseSMatrix/SparseSMatrix_SELL.py

@@ -0,0 +1,328 @@
+import pycuda.driver as drv
+import numpy as np
+from tqdm import trange
+import os
+import gc
+
+class SparseSMatrix_SELL:
+    def __init__(self, manip, block_rows=64, relative_threshold=0.3, device=0,
+                 module_path="AOT_biomaps_kernels.cubin", slice_height=32):
+        drv.init()
+        self.device = drv.Device(device)
+        self.ctx = self.device.make_context()
+        self.manip = manip
+        self.N = len(manip.AcousticFields)
+        self.T = manip.AcousticFields[0].field.shape[0]
+        self.Z = manip.AcousticFields[0].field.shape[1]
+        self.X = manip.AcousticFields[0].field.shape[2]
+        self.block_rows = block_rows
+        self.relative_threshold = relative_threshold
+        
+        # --- PATH RESOLUTION FIX ---
+        # The cubin file is located in the parent directory (AOT_Recon/)
+        # We use os.path.dirname(os.path.dirname(__file__)) to go up one directory level.
+        cubin_parent_dir = os.path.dirname(os.path.dirname(__file__))
+        self.module_path = os.path.join(cubin_parent_dir, module_path)
+        # --- END FIX ---
+        
+        self.slice_height = slice_height
+
+        # SELL arrays (device) & Size Tracking (CRITICAL FIX: Initialized attributes)
+        self.sell_values_gpu = None
+        self.sell_colinds_gpu = None
+        self.slice_ptr = None
+        self.slice_len = None
+        self.slice_ptr_gpu = None
+        self.slice_len_gpu = None
+        
+        # Attributes to store allocated size in bytes (bypassing the problematic .size attribute)
+        self.sell_values_gpu_size = 0  
+        self.sell_colinds_gpu_size = 0 
+        self.slice_ptr_gpu_size = 0    
+        self.slice_len_gpu_size = 0    
+
+        self.total_storage = 0
+
+        self.norm_factor_inv = None
+        self.norm_factor_inv_gpu = None
+        self.norm_factor_inv_gpu_size = 0
+
+        self.sparse_mod = None
+        self.load_module()
+
+    def load_module(self):
+        """Loads the pre-compiled CUDA module (.cubin file)."""
+        
+        # Check if the file exists at the calculated absolute path
+        if not os.path.exists(self.module_path):
+            # The path is now correctly calculated to the parent directory.
+            raise FileNotFoundError(f"CUDA module {os.path.basename(self.module_path)} not found at path: {self.module_path}")
+            
+        # Try to load the module
+        try:
+            self.sparse_mod = drv.module_from_file(self.module_path)
+            print(f"Loaded CUDA module {os.path.basename(self.module_path)}")
+        except Exception as e:
+            print(f"❌ Error loading CUDA module {os.path.basename(self.module_path)}: {e}")
+            raise RuntimeError(f"File {os.path.basename(self.module_path)} was found, but PyCUDA could not load it. Check compatibility.") from e
+
+    def free(self):
+        try:
+            # Free device allocations
+            attrs = ["sell_values_gpu","sell_colinds_gpu","slice_ptr_gpu","slice_len_gpu","norm_factor_inv_gpu"]
+            for a in attrs:
+                if hasattr(self, a) and getattr(self, a) is not None:
+                    getattr(self, a).free()
+                    setattr(self, a, None)
+            
+            # Reset stored sizes
+            self.sell_values_gpu_size = 0
+            self.sell_colinds_gpu_size = 0
+            self.slice_ptr_gpu_size = 0
+            self.slice_len_gpu_size = 0
+            self.norm_factor_inv_gpu_size = 0
+            
+            if hasattr(self, 'ctx') and self.ctx:
+                try: self.ctx.pop()
+                except Exception: pass
+                self.ctx = None
+        except Exception as e:
+            print("Error freeing GPU memory:", e)
+
+    def allocate(self):
+        """
+        Build SELL-C-σ directly from manip AcousticFields in streaming blocks.
+        NOTE: This is the logic of allocate_sell_c_sigma_direct from the working class.
+        """
+        if self.sparse_mod is None:
+            raise RuntimeError("CUDA module not loaded. Check compilation.")
+            
+        # NOTE: Les noms de kernel (count_nnz_rows_kernel, fill_kernel__SELL) sont utilisés
+        # car ils sont présents dans la classe fonctionnelle.
+        count_kernel = self.sparse_mod.get_function("count_nnz_rows_kernel")
+        fill_kernel  = self.sparse_mod.get_function("fill_kernel__SELL")
+
+        num_rows = int(self.N * self.T)
+        num_cols = int(self.Z * self.X)
+        C = int(self.slice_height)
+
+        # host temporary block
+        br = int(self.block_rows)
+        bytes_per_elem = np.dtype(np.float32).itemsize
+        dense_host = np.empty((br, num_cols), dtype=np.float32)
+        
+        # Allocation 1: Dense block GPU memory
+        dense_gpu_size = dense_host.nbytes
+        dense_gpu  = drv.mem_alloc(dense_gpu_size)
+
+        # 1) count nnz per row (on host via small blocks with GPU kernel)
+        row_nnz = np.zeros(num_rows, dtype=np.int32)
+        row_nnz_gpu_block_size = br * np.dtype(np.int32).itemsize
+        row_nnz_gpu_block = drv.mem_alloc(row_nnz_gpu_block_size)
+
+        block = 256
+        for b in trange(0, num_rows, br, desc="Count NNZ per row"):
+            R = min(br, num_rows - b)
+            # fill dense_host
+            for i in range(R):
+                rg = b + i
+                n_idx = rg // self.T
+                t_idx = rg % self.T
+                dense_host[i, :] = self.manip.AcousticFields[n_idx].field[t_idx].flatten()
+            # copy only R rows
+            drv.memcpy_htod(dense_gpu, dense_host)
+            grid = ((R + block - 1) // block, 1, 1)
+            count_kernel(dense_gpu, row_nnz_gpu_block, np.int32(R), np.int32(num_cols), np.float32(self.relative_threshold),
+                         block=(block,1,1), grid=grid)
+            tmp = np.empty(R, dtype=np.int32)
+            drv.memcpy_dtoh(tmp, row_nnz_gpu_block)
+            row_nnz[b:b+R] = tmp
+
+        row_nnz_gpu_block.free()
+        dense_gpu.free()
+
+        # 2) compute per-slice maxlen and slice_ptr
+        num_slices = (num_rows + C - 1) // C
+        slice_len = np.zeros(num_slices, dtype=np.int32)
+        for s in range(num_slices):
+            r0 = s * C
+            r1 = min(num_rows, r0 + C)
+            slice_len[s] = int(np.max(row_nnz[r0:r1])) if (r1>r0) else 0
+        # slice_ptr (int64)
+        slice_ptr = np.zeros(num_slices + 1, dtype=np.int64)
+        for s in range(num_slices):
+            slice_ptr[s+1] = slice_ptr[s] + (slice_len[s] * C)
+        total_storage = int(slice_ptr[-1])
+        self.total_storage = total_storage
+        print(f"SELL: num_rows={num_rows}, num_slices={num_slices}, total_storage(padded)={total_storage}")
+
+        # allocate device SELL arrays (values float32, colinds uint32)
+        self.sell_values_gpu_size = total_storage * np.dtype(np.float32).itemsize
+        self.sell_colinds_gpu_size = total_storage * np.dtype(np.uint32).itemsize
+        
+        self.sell_values_gpu = drv.mem_alloc(self.sell_values_gpu_size)
+        self.sell_colinds_gpu = drv.mem_alloc(self.sell_colinds_gpu_size)
+        
+        # allocate slice metadata on device
+        self.slice_ptr = slice_ptr
+        self.slice_len = slice_len
+        
+        self.slice_ptr_gpu_size = self.slice_ptr.nbytes
+        self.slice_len_gpu_size = self.slice_len.nbytes
+        
+        self.slice_ptr_gpu = drv.mem_alloc(self.slice_ptr_gpu_size)
+        self.slice_len_gpu = drv.mem_alloc(self.slice_len_gpu_size)
+        
+        drv.memcpy_htod(self.slice_ptr_gpu, self.slice_ptr)
+        drv.memcpy_htod(self.slice_len_gpu, self.slice_len)
+
+        # 3) fill SELL arrays by streaming blocks again (use GPU fill kernel)
+        # reuse dense_host and allocate new dense_gpu
+        dense_host = np.empty((br, num_cols), dtype=np.float32)
+        
+        dense_gpu_2_size = dense_host.nbytes
+        dense_gpu  = drv.mem_alloc(dense_gpu_2_size)
+        
+        # we also need row_nnz on device per-block; supply global row_nnz on host but the kernel recomputes threshold
+        row_nnz_host_gpu_size = br * np.dtype(np.int32).itemsize
+        row_nnz_host_gpu = drv.mem_alloc(row_nnz_host_gpu_size)
+
+        for b in trange(0, num_rows, br, desc="Fill SELL"):
+            R = min(br, num_rows - b)
+            for i in range(R):
+                rg = b + i
+                n_idx = rg // self.T
+                t_idx = rg % self.T
+                dense_host[i, :] = self.manip.AcousticFields[n_idx].field[t_idx].flatten()
+            drv.memcpy_htod(dense_gpu, dense_host)
+            # We pass a dummy row_nnz pointer (not used in this kernel; left for API)
+            # Kernel expects rows_in_block, rows_global_offset to know where to write.
+            grid = ((R + block - 1) // block, 1, 1)
+            fill_kernel(dense_gpu,
+                        np.intp(0), # placeholder for row_nnz pointer (not used)
+                        self.slice_ptr_gpu,
+                        self.slice_len_gpu,
+                        self.sell_colinds_gpu,
+                        self.sell_values_gpu,
+                        np.int32(R),
+                        np.int32(num_cols),
+                        np.int32(b),          # rows_global_offset
+                        np.int32(C),
+                        np.float32(self.relative_threshold),
+                        block=(block,1,1), grid=grid)
+        dense_gpu.free()
+        row_nnz_host_gpu.free()
+
+        # 4) compute norm_factor_inv via GPU accumulate (col sums)
+        self.compute_norm_factor()
+    
+    def apply_apodization_gpu(self, window_vector_gpu):
+        """
+        Applique le fenêtrage directement sur self.sell_values_gpu 
+        en utilisant les indices de colonnes (pixels) pour référencer 
+        la fenêtre. Opération : A_values[i] *= W_vec[A_colinds[i]].
+        """
+        if self.sparse_mod is None:
+            raise RuntimeError("Le module CUDA n'a pas été chargé.")
+            
+        try:
+            apodize_kernel = self.sparse_mod.get_function("apply_apodisation_kernel__SELL")
+        except drv.LogicError as e:
+            raise RuntimeError(
+                f"Le kernel CUDA 'multiply_sell_by_window_kernel' est manquant dans le .cubin. "
+                f"Veuillez le compiler et l'ajouter. Erreur : {e}"
+            )
+
+        # Le total_storage inclut les éléments non-nuls et le padding SELL.
+        threads = 256
+        blocks = (self.total_storage + threads - 1) // threads
+        
+        # Lancement du kernel. Il travaille sur total_storage éléments.
+        apodize_kernel(
+            self.sell_values_gpu, 
+            self.sell_colinds_gpu, 
+            window_vector_gpu, 
+            np.int64(self.total_storage), 
+            block=(threads, 1, 1), 
+            grid=(blocks, 1, 1)
+        )
+        drv.Context.synchronize()
+        print("✅ Multiplication par le fenêtrage effectuée in-place sur GPU (SELL-C-σ).")
+    # --- Ajout de la fonction de normalisation (qui fonctionne) ---
+    def compute_norm_factor(self):
+        """
+        Accumulate column sums on GPU using accumulate_columns_atomic, then compute inverse.
+        """
+        if self.total_storage == 0:
+            raise RuntimeError("sell not built")
+        ZX = int(self.Z * self.X)
+
+        # allocate col sum on device
+        col_sum_gpu_size = ZX * np.dtype(np.float32).itemsize
+        col_sum_gpu = drv.mem_alloc(col_sum_gpu_size)
+        drv.memset_d32(col_sum_gpu, 0, ZX)
+
+        # FIX: Kernel name is "accumulate_columns_atomic"
+        acc_kernel = self.sparse_mod.get_function("accumulate_columns_atomic")
+        
+        threads = 256
+        blocks = (self.total_storage + threads - 1) // threads
+        acc_kernel(self.sell_values_gpu, self.sell_colinds_gpu, np.int64(self.total_storage), col_sum_gpu,
+                   block=(threads,1,1), grid=(blocks,1,1))
+        drv.Context.synchronize()
+
+        # copy back
+        norm_host = np.empty(ZX, dtype=np.float32)
+        drv.memcpy_dtoh(norm_host, col_sum_gpu)
+        col_sum_gpu.free()
+
+        norm = np.maximum(norm_host.astype(np.float64), 1e-6)
+        self.norm_factor_inv = (1.0 / norm).astype(np.float32)
+        if self.norm_factor_inv_gpu is not None:
+            self.norm_factor_inv_gpu.free()
+            
+        self.norm_factor_inv_gpu_size = self.norm_factor_inv.nbytes
+        self.norm_factor_inv_gpu = drv.mem_alloc(self.norm_factor_inv_gpu_size)
+        drv.memcpy_htod(self.norm_factor_inv_gpu, self.norm_factor_inv)
+
+    def compute_density(self):
+        """
+        Returns only the density of the SELL-C-σ matrix.
+        """
+        if not hasattr(self, 'slice_ptr') or self.slice_ptr is None:
+            raise RuntimeError("The SELL-C-σ matrix is not allocated.")
+
+        num_rows = self.N * self.T
+        num_cols = self.Z * self.X
+        total_elements = num_rows * num_cols
+
+        # Conservative estimate of non-zeros (excluding padding)
+        nnz_ell_estimated = int(0.9 * self.total_storage)
+
+        return nnz_ell_estimated / total_elements  # Returns only the density
+
+    def getMatrixSize(self):
+        """
+        Returns the total size of the SELL-C-σ matrix in Gigabytes (GB).
+        """
+        if self.sell_values_gpu is None:
+            return {"error": "The SELL-C-σ matrix is not yet allocated."}
+
+        total_bytes = 0
+
+        # Host-side arrays (using .nbytes which works for NumPy arrays)
+        if hasattr(self, 'slice_ptr') and self.slice_ptr is not None:
+            total_bytes += self.slice_ptr.nbytes
+        if hasattr(self, 'slice_len') and self.slice_len is not None:
+            total_bytes += self.slice_len.nbytes
+        if hasattr(self, 'norm_factor_inv') and self.norm_factor_inv is not None:
+            total_bytes += self.norm_factor_inv.nbytes
+
+        # GPU-side arrays (using the stored size attributes instead of the problematic .size)
+        total_bytes += self.sell_values_gpu_size
+        total_bytes += self.sell_colinds_gpu_size
+        total_bytes += self.slice_ptr_gpu_size
+        total_bytes += self.slice_len_gpu_size
+        total_bytes += self.norm_factor_inv_gpu_size
+        
+        return total_bytes / (1024 ** 3)  # Returns only the size in GB

AOT_biomaps/AOT_Recon/AOT_SparseSMatrix/__init__.py

@@ -0,0 +1,2 @@
+from .SparseSMatrix_CSR import *
+from .SparseSMatrix_SELL import *