transition_matrix_gpu.py

"""
===================================================================================
CALCUL DE LA MATRICE DE TRANSITION POUR UN MÉLANGEUR DEM (Discrete Element Method)
===================================================================================

Objectif: Construire une matrice de Markov P_ij où:
  - P[i,j] = probabilité qu'une particule passe de l'état i à l'état j entre deux timesteps
  - Les états sont définis par la position spatiale discrétisée (5x5x5 grille)

Formule: P_ij = (1/N_LT) * sum_n [ sum_p (phi_p(j,t_n) * phi_p(i,t_{n-1})) / phi(i,t_n) ]

Structure des données:
  - 6000 fichiers CSV contenant chacun ~1030 particules
  - Colonnes: Particle_ID, coordinates:0/1/2 (x,y,z), Particle_Phase_ID, etc.
  - Chaque particule est trackée par son Particle_ID unique

Auteur: OpenCode
Date: 2026
"""

# ===================================================================================
# IMPORTATIONS DES BIBLIOTHÈQUES
# ===================================================================================

import polars as pl  # Lecture efficace des fichiers CSV/Parquet
from huggingface_hub import HfFileSystem  # Accès aux fichiers distants sur HuggingFace
import torch  # Calculs GPU (CUDA)
from tqdm import tqdm  # Barre de progression
import numpy as np  # Manipulation de tableaux

# ===================================================================================
# CONNEXION AU DÉPÔT HUGGINGFACE
# ===================================================================================

fs = HfFileSystem()  # Système de fichiers pour HuggingFace

# Chemin vers les fichiers CSV de simulation DEM
folder_path = "hf://buckets/ktongue/DEM_MCM/Output Paraview"
csv_files = sorted(
    fs.glob(f"{folder_path}/*.csv")
)  # Liste triée de tous les fichiers CSV

print(f"📁 {len(csv_files)} fichiers CSV trouvés")

# ===================================================================================
# ÉTAPE 1: CALCUL DES LIMITES SPATIALES DU DOMAINE
# ===================================================================================
# On détermine les frontières min/max en x, y, z du mélangeur
# pour pouvoir discrétiser l'espace en cellules régulières

print("🔍 Calcul des limites spatiales...")

# Échantillonnage: on ne lit qu'1 fichier sur 100 pour rapidité
# (les limites spatiales ne changent pas significativement)
sample_files = csv_files[::100]
x_vals, y_vals, z_vals = [], [], []

# Parcours des fichiers échantillonnés
for f in tqdm(sample_files, desc="Scanning limits"):
    with fs.open(f, "rb") as file:
        df = pl.read_csv(file)
    # Extraction des coordonnées de toutes les particules
    x_vals.extend(df["coordinates:0"].to_list())  # Coordonnée X
    y_vals.extend(df["coordinates:1"].to_list())  # Coordonnée Y
    z_vals.extend(df["coordinates:2"].to_list())  # Coordonnée Z

# Calcul des valeurs min/max
xmin, xmax = min(x_vals), max(x_vals)
ymin, ymax = min(y_vals), max(y_vals)
zmin, zmax = min(z_vals), max(z_vals)

# Ajout d'une marge de sécurité (1mm) pour éviter les problèmes aux frontières
margin = 0.001
xmin -= margin
xmax += margin
ymin -= margin
ymax += margin
zmin -= margin
zmax += margin

print(
    f"Limites: X=[{xmin:.4f}, {xmax:.4f}], Y=[{ymin:.4f}, {ymax:.4f}], Z=[{zmin:.4f}, {zmax:.4f}]"
)

# ===================================================================================
# ÉTAPE 2: DISCRÉTISATION SPATIALE
# ===================================================================================
# L'espace 3D est divisé en une grille régulière de nx x ny x nz cellules
# Chaque cellule devient un "état" de la chaîne de Markov

nx, ny, nz = 5, 5, 5  # Nombre de divisions par axe (5x5x5 = 125 états)
n_states = nx * ny * nz  # Total des états spatiaux

dx = (xmax - xmin) / nx  # Largeur d'une cellule en X
dy = (ymax - ymin) / ny  # Largeur d'une cellule en Y
dz = (zmax - zmin) / nz  # Largeur d'une cellule en Z


def compute_states_gpu(x, y, z):
    """
    Convertit les coordonnées continues (x,y,z) en indices d'états discrets.

    Args:
        x, y, z: Arrays numpy des coordonnées

    Returns:
        Tensor GPU contenant les indices d'état pour chaque particule
    """
    # Conversion en tensors GPU et calcul des indices de cellule
    ix = ((torch.tensor(x, device="cuda") - xmin) / dx).long().clamp(0, nx - 1)
    iy = ((torch.tensor(y, device="cuda") - ymin) / dy).long().clamp(0, ny - 1)
    iz = ((torch.tensor(z, device="cuda") - zmin) / dz).long().clamp(0, nz - 1)

    # Index linéaire: ix + iy*nx + iz*nx*ny
    return ix + iy * nx + iz * nx * ny


print(f"📊 Discrétisation: {nx}x{ny}x{nz} = {n_states} états")

# ===================================================================================
# ÉTAPE 3: INITIALISATION DES TENSEURS SUR GPU
# ===================================================================================
# On utilise PyTorch pour effectuer les calculs sur GPU (CUDA)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🖥️  Device: {device}")

# transition_counts[i,j] = nombre de transitions observées de i vers j
transition_counts = torch.zeros((n_states, n_states), dtype=torch.int64, device=device)

# state_counts[i] = nombre total de particules observées dans l'état i
state_counts = torch.zeros(n_states, dtype=torch.int64, device=device)

# ===================================================================================
# ÉTAPE 4: CALCUL DES TRANSITIONS SUR GPU
# ===================================================================================
# Pour chaque paire de fichiers consécutifs, on tracke les particules
# et on compte les transitions d'un état spatial à un autre

print("🔄 Calcul des transitions sur GPU...")

n_files = len(csv_files)
df_prev = None  # DataFrame du timestep précédent
states_prev = None  # États GPU du timestep précédent

for i in tqdm(range(1, n_files), desc="Processing"):
    # Lecture du fichier CSV courant
    with fs.open(csv_files[i], "rb") as f:
        df_curr = pl.read_csv(f)

    # IDs des particules au timestep courant
    ids_curr = df_curr["Particle_ID"].to_numpy()

    # Calcul des états spatiaux sur GPU
    states_curr = compute_states_gpu(
        df_curr["coordinates:0"].to_numpy(),
        df_curr["coordinates:1"].to_numpy(),
        df_curr["coordinates:2"].to_numpy(),
    )

    if df_prev is not None:
        # IDs des particules au timestep précédent
        ids_prev = df_prev["Particle_ID"].to_numpy()

        # Identification des particules communes (trackées dans les 2 timesteps)
        common_ids = np.intersect1d(ids_prev, ids_curr)

        if len(common_ids) > 0:
            # Masques pour extraire les particules communes
            prev_idx = np.isin(ids_prev, common_ids)
            curr_idx = np.isin(ids_curr, common_ids)

            # États avant/après pour les particules communes
            s_prev = states_prev[prev_idx]
            s_curr = states_curr[curr_idx]

            # Comptage vectorisé des transitions uniques
            transitions = torch.stack([s_prev, s_curr], dim=1)
            unique_trans, counts = torch.unique(transitions, dim=0, return_counts=True)

            # Accumulation des comptages de transitions
            for j in range(len(unique_trans)):
                transition_counts[unique_trans[j, 0], unique_trans[j, 1]] += counts[j]

            # Comptage des particules par état d'origine
            unique_prev, _ = torch.unique(s_prev, return_counts=True)
            for u in unique_prev:
                mask = s_prev == u
                state_counts[u] += mask.sum().item()

    # Sauvegarde pour l'itération suivante
    df_prev = df_curr
    states_prev = states_curr

# ===================================================================================
# ÉTAPE 5: NORMALISATION POUR OBTENIR LES PROBABILITÉS P_ij
# ===================================================================================
# Chaque ligne i est divisée par le nombre total de particules parties de i
# pour obtenir des probabilités (somme = 1 par ligne)

print("📐 Normalisation...")

P = torch.zeros((n_states, n_states), dtype=torch.float64, device=device)
for i in range(n_states):
    if state_counts[i] > 0:
        # P[i,j] = transition_counts[i,j] / state_counts[i]
        P[i] = transition_counts[i].float() / state_counts[i].float()

# Vérification de la normalisation
row_sums = P.sum(dim=1)
visited = state_counts > 0
print(
    f"   Somme par ligne (min/max): {row_sums[visited].min().item():.4f} / {row_sums[visited].max().item():.4f}"
)
print(f"   États visités: {visited.sum().item()} / {n_states}")

# ===================================================================================
# ÉTAPE 6: SAUVEGARDE DES RÉSULTATS
# ===================================================================================
# Export en format NumPy pour utilisation ultérieure

# Transfert CPU -> NumPy
P_cpu = P.cpu().numpy()
transition_cpu = transition_counts.cpu().numpy()
state_counts_cpu = state_counts.cpu().numpy()

# Sauvegarde des fichiers
np.save("/kaggle/working/transition_matrix.npy", P_cpu)  # Matrice de transition P_ij
np.save("/kaggle/working/transition_counts.npy", transition_cpu)  # Comptages bruts
np.save("/kaggle/working/state_counts.npy", state_counts_cpu)  # Comptages par état

print(f"\n✅ Matrice sauvegardée!")
print(f"   Shape: {P_cpu.shape}")
print(f"   Fichiers: /kaggle/working/transition_matrix.npy")

# ===================================================================================
# VISUALISATION PRÉLIMINAIRE
# ===================================================================================
# Affiche un extrait de la matrice pour vérification

print("\n📊 Aperçu de la matrice de transition (5x5 coin supérieur):")
print(P_cpu[:5, :5].round(4))