transition_matrix_article.py

"""
===================================================================================
MODÉLISATION MARKOVIENNE - FORMULE DE L'ARTICLE
===================================================================================

Implémentation stricte de la formule de l'article:
    P_ij = (1/N_LT) * Sum_n [ T_ij(n) / phi(i, t_{n-1}) ]

où:
    - N_LT = nombre de pas de temps pour l'apprentissage (Learning Time)
    - T_ij(n) = nombre de transitions i→j au pas n
    - phi(i, t_{n-1}) = nombre de particules dans l'état i au temps t_{n-1}

Différence avec l'approche GPU:
    - Version GPU: Compte TOUTES les transitions, puis divise par le total
    - Version article: Calcule une matrice normalisée à chaque pas de temps, puis moyenne

===================================================================================
"""

import polars as pl
from huggingface_hub import HfFileSystem
import numpy as np
from tqdm import tqdm

# ===================================================================================
# PARAMÈTRES UTILISATEUR
# ===================================================================================
N_LT = 100  # Nombre de pas de temps pour l'apprentissage (Learning Time)
nx, ny, nz = 5, 5, 5  # Discrétisation spatiale
# ===================================================================================

# ===================================================================================
# CONNEXION AU DÉPÔT HUGGINGFACE
# ===================================================================================
fs = HfFileSystem()
folder_path = "hf://buckets/ktongue/DEM_MCM/Output Telem"
files = sorted(fs.glob(f"{folder_path}/*.csv"))

print(f"📁 {len(files)} fichiers disponibles")
print(f"⚙️  Utilisation de N_LT = {N_LT} pas de temps pour l'apprentissage")

# ===================================================================================
# ÉTAPE 1: CALCUL DES LIMITES SPATIALES
# ===================================================================================
# Échantillonnage rapide pour déterminer les frontières du domaine

print("🔍 Calcul des limites spatiales...")
sample_files = files[::50]  # 1 fichier sur 50 pour rapidité
x_vals, y_vals, z_vals = [], [], []

for f in sample_files:
    with fs.open(f, "rb") as file:
        df = pl.read_csv(file)
    x_vals.extend(df["x"].to_list())
    y_vals.extend(df["y"].to_list())
    z_vals.extend(df["z"].to_list())

# Calcul des min/max avec marge de sécurité
xmin, xmax = min(x_vals) - 0.001, max(x_vals) + 0.001
ymin, ymax = min(y_vals) - 0.001, max(y_vals) + 0.001
zmin, zmax = min(z_vals) - 0.001, max(z_vals) + 0.001

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

# ===================================================================================
# ÉTAPE 2: INITIALISATION DE LA DISCRÉTISATION
# ===================================================================================
n_states = nx * ny * nz  # Total = 125 états
dx = (xmax - xmin) / nx
dy = (ymax - ymin) / ny
dz = (zmax - zmin) / nz

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


def get_state(x, y, z):
    """
    Convertit les coordonnées continues (x,y,z) en indice d'état discret.

    Formule:
        ix = clamp(floor((x - xmin) / dx), 0, nx-1)
        iy = clamp(floor((y - ymin) / dy), 0, ny-1)
        iz = clamp(floor((z - zmin) / dz), 0, nz-1)

        state = ix + iy * nx + iz * nx * ny

    Args:
        x, y, z: Coordonnées spatiales

    Returns:
        Indice d'état (entier entre 0 et n_states-1)
    """
    ix = min(max(int((x - xmin) / dx), 0), nx - 1)
    iy = min(max(int((y - ymin) / dy), 0), ny - 1)
    iz = min(max(int((z - zmin) / dz), 0), nz - 1)
    return ix + iy * nx + iz * nx * ny


# ===================================================================================
# ÉTAPE 3: CALCUL DE LA MATRICE SELON LA FORMULE DE L'ARTICLE
# ===================================================================================
# Formule: P_ij = (1/N_LT) * Sum_n [ T_ij(n) / phi(i, t_{n-1}) ]
#
# Pour chaque pas de temps n:
#   1. Lire les fichiers à t_{n-1} et t_n
#   2. Tracker les particules communes
#   3. Calculer T_ij(n) et phi(i, t_{n-1})
#   4. Accumuler P_ij(n) = T_ij(n) / phi(i, t_{n-1})
#
# Après N_LT pas: P = (1/N_LT) * Sum_n P_ij(n)

print(f"📊 Calcul de la matrice sur {N_LT} pas de temps...")

# Accumulateur pour la somme des matrices normalisées
# P_sum[i,j] = Sum_n [ T_ij(n) / phi(i, t_{n-1}) ]
P_sum = np.zeros((n_states, n_states))

# On prend les N_LT premiers pas de temps
# (modifier start_index pour ignorer le régime transitoire)
start_index = 0
files_to_process = files[start_index : start_index + N_LT + 1]

print(f"   Traitement des fichiers {start_index} à {start_index + N_LT}...")

# Boucle principale: pour chaque pas de temps
for i in tqdm(range(1, len(files_to_process)), desc="Learning"):
    # -----------------------------------------------------------------
    # Étape 3.1: Lecture de deux fichiers consécutifs
    # -----------------------------------------------------------------
    # df_prev = données à t_{n-1}
    # df_curr = données à t_n
    with fs.open(files_to_process[i - 1], "rb") as f:
        df_prev = pl.read_csv(f)
    with fs.open(files_to_process[i], "rb") as f:
        df_curr = pl.read_csv(f)

    # -----------------------------------------------------------------
    # Étape 3.2: Calcul des états pour chaque particule
    # -----------------------------------------------------------------
    # Utilisation de Polars pour appliquer get_state() à chaque ligne

    # df_prev.with_columns(): ajoute une nouvelle colonne sans modifier l'original
    # pl.struct(['x', 'y', 'z']): combine les 3 colonnes en une structure
    # .map_elements(lambda s: ...): applique la fonction à chaque structure
    # .alias('state_prev'): nomme la nouvelle colonne

    df_prev = df_prev.with_columns(
        [
            pl.struct(["x", "y", "z"])
            .map_elements(
                lambda s: get_state(s["x"], s["y"], s["z"]), return_dtype=pl.Int64
            )
            .alias("state_prev")
        ]
    )

    df_curr = df_curr.with_columns(
        [
            pl.struct(["x", "y", "z"])
            .map_elements(
                lambda s: get_state(s["x"], s["y"], s["z"]), return_dtype=pl.Int64
            )
            .alias("state_curr")
        ]
    )

    # -----------------------------------------------------------------
    # Étape 3.3: Jointure pour tracker les particules
    # -----------------------------------------------------------------
    # On garde seulement les particules présentes AUX DEUX instants
    # merged = [(ID, state_{n-1}, state_n), ...]

    merged = df_prev.select(["ID", "state_prev"]).join(
        df_curr.select(["ID", "state_curr"]), on="ID"
    )

    # -----------------------------------------------------------------
    # Étape 3.4: Comptage des transitions pour ce pas de temps
    # -----------------------------------------------------------------
    # phi_counts: phi(i, t_{n-1}) = nombre de particules dans chaque état source
    # Formula: phi(i, t_{n-1}) = Sum_j T_ij(n)
    phi_counts = merged.group_by("state_prev").len().rename({"len": "total_source"})

    # transitions: T_ij(n) = nombre de transitions i->j observées au pas n
    transitions = merged.group_by(["state_prev", "state_curr"]).len()

    # -----------------------------------------------------------------
    # Étape 3.5: Calcul de la matrice instantanée et accumulation
    # -----------------------------------------------------------------
    # Pour ce pas de temps n, on calcule:
    #   P_ij(n) = T_ij(n) / phi(i, t_{n-1})
    # Et on l'ajoute à P_sum

    trans_with_total = transitions.join(phi_counts, on="state_prev")

    for row in trans_with_total.iter_rows():
        state_from, state_to, count, total = row
        if total > 0:
            # P_sum[i,j] += T_ij(n) / phi(i, t_{n-1})
            P_sum[state_from, state_to] += count / total

# ===================================================================================
# ÉTAPE 4: MOYENNE FINALE
# ===================================================================================
# Formule: P_ij = (1/N_LT) * Sum_n P_ij(n)

P = P_sum / N_LT

print(f"\n✅ Matrice calculée avec N_LT = {N_LT}")
print(f"   Shape: {P.shape}")

# Vérification de la stochasticité
row_sums = P.sum(axis=1)
print(f"   Somme par ligne (min/max): {row_sums.min():.4f} / {row_sums.max():.4f}")

# ===================================================================================
# ÉTAPE 5: SAUVEGARDE
# ===================================================================================
output_file = f"/kaggle/working/transition_matrix_NLT_{N_LT}.npy"
np.save(output_file, P)
print(f"💾 Sauvegardé: {output_file}")

# ===================================================================================
# ANALYSE COMPLÉMENTAIRE
# ===================================================================================
print(f"\n📈 Analyse de la matrice:")

# Probabilité de rester dans la même cellule (diagonale)
diagonal = np.diag(P)
print(f"   Probabilité de rester (diagonale):")
print(f"      Moyenne: {diagonal.mean():.4f}")
print(f"      Écart-type: {diagonal.std():.4f}")

# Transitions les plus probables (hors diagonale)
P_nodiag = P.copy()
np.fill_diagonal(P_nodiag, 0)
max_idx = np.unravel_index(P_nodiag.argmax(), P_nodiag.shape)
print(
    f"   Transition la plus probable: {max_idx[0]} -> {max_idx[1]} (P={P_nodiag.max():.4f})"
)

print(f"\n✨ Terminé!")