Affiner les heatmaps de corrélation et l'annotation des lags

docs/04 - Corrélations binaires/scripts/plot_correlation_heatmap.py

@@ -4,6 +4,9 @@ from __future__ import annotations
 from dataclasses import dataclass
 from pathlib import Path
 import sys
+import argparse
+import pandas as pd
+import numpy as np
 
 
 PROJECT_ROOT = Path(__file__).resolve().parents[3]
@@ -20,7 +23,6 @@ CSV_PATH = Path("data/weather_minutely.csv")
 DOC_DIR = Path(__file__).resolve().parent.parent
 
 CORRELATION_METHODS: tuple[str, ...] = ("pearson", "spearman")
-CORRELATION_TRANSFORM = "square"
 
 
 @dataclass(frozen=True)
@@ -36,13 +38,19 @@ class HeatmapConfig:
 HEATMAP_CONFIGS: dict[str, HeatmapConfig] = {
     "pearson": HeatmapConfig(
         filename="correlation_heatmap.png",
-        title="Corrélations R² (coef. de Pearson)",
-        colorbar_label="Coefficient de corrélation R²",
+        title="Corrélations (coef. de Pearson)",
+        colorbar_label="Coefficient de corrélation",
+        cmap="viridis",
+        vmin=0.0,
+        vmax=1.0,
     ),
     "spearman": HeatmapConfig(
         filename="correlation_heatmap_spearman.png",
-        title="Corrélations R² (coef. de Spearman)",
-        colorbar_label="Coefficient de corrélation R²",
+        title="Corrélations (coef. de Spearman)",
+        colorbar_label="Coefficient de corrélation",
+        cmap="viridis",
+        vmin=0.0,
+        vmax=1.0,
     ),
 }
 
@@ -63,36 +71,92 @@ def _get_heatmap_config(method: str) -> HeatmapConfig:
 
 
 def main() -> None:
+    parser = argparse.ArgumentParser(description="Trace des matrices de corrélation instantanées (signées, absolues ou r²).")
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=DOC_DIR / "figures",
+        help="Dossier de sortie pour les heatmaps.",
+    )
+    parser.add_argument(
+        "--transform",
+        choices=["identity", "absolute", "square"],
+        default="absolute",
+        help="Transformation de la matrice (signée, |r| ou r²). Par défaut : |r|.",
+    )
+    parser.add_argument(
+        "--upper-only",
+        action="store_true",
+        help="Masque la partie inférieure de la matrice pour alléger la lecture.",
+    )
+    args = parser.parse_args()
+
     if not CSV_PATH.exists():
         print(f"⚠ Fichier introuvable : {CSV_PATH}")
         print("   Assurez-vous d'avoir généré le dataset minuté.")
         return
 
     df = load_raw_csv(CSV_PATH)
+    df = df[[v.column for v in VARIABLES]]
     print(f"Dataset minuté chargé : {CSV_PATH}")
     print(f"  Lignes   : {len(df)}")
     print(f"  Colonnes : {list(df.columns)}")
     print()
 
+    transform = args.transform
     matrices = compute_correlation_matrices_for_methods(
         df=df,
         variables=VARIABLES,
         methods=CORRELATION_METHODS,
-        transform=CORRELATION_TRANSFORM,
+        transform=transform,
     )
 
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
     for method, corr in matrices.items():
-        print(f"Matrice de corrélation (méthode={method}, transform={CORRELATION_TRANSFORM}) :")
+        if args.upper_only:
+            mask = np.tril(np.ones_like(corr, dtype=bool), k=-1)
+            corr = corr.mask(mask)
+
+        print(f"Matrice de corrélation (méthode={method}, transform={transform}) :")
         print(corr)
         print()
 
         config = _get_heatmap_config(method)
+        filename = config.filename
+        title = config.title
+        if transform == "absolute":
+            title = f"{title} (|r|)"
+            stem, suffix = filename.rsplit(".", 1)
+            filename = f"{stem}_abs.{suffix}"
+        elif transform == "square":
+            title = f"{title} (r²)"
+            stem, suffix = filename.rsplit(".", 1)
+            filename = f"{stem}_r2.{suffix}"
+            config = HeatmapConfig(
+                filename=filename,
+                title=title,
+                colorbar_label="Coefficient de corrélation r²",
+                cmap="viridis",
+                vmin=0.0,
+                vmax=1.0,
+            )
+        elif transform == "identity":
+            config = HeatmapConfig(
+                filename=filename,
+                title=title,
+                colorbar_label="Coefficient de corrélation r",
+                cmap="coolwarm",
+                vmin=-1.0,
+                vmax=1.0,
+            )
+
         output_path = plot_correlation_heatmap(
             corr=corr,
             variables=VARIABLES,
-            output_path=DOC_DIR / "figures" / config.filename,
+            output_path=args.output_dir / filename,
             annotate=True,
-            title=config.title,
+            title=title,
             cmap=config.cmap,
             vmin=config.vmin,
             vmax=config.vmax,