Ajout des matrices de corrélation + Refactoring

meteo/analysis/__init__.py

@@ -6,9 +6,11 @@ from .core import BinnedStatistics, DiurnalCycleStats, MONTH_ORDER
 from .correlations import (
     compute_correlation_matrix,
     compute_correlation_matrix_for_variables,
+    compute_correlation_matrices_for_methods,
     compute_lagged_correlation,
     compute_rolling_correlation_series,
     compute_rolling_correlations_for_pairs,
+    transform_correlation_matrix,
 )
 from .events import build_event_aligned_segments, detect_threshold_events
 from .filters import filter_by_condition
@@ -28,9 +30,11 @@ __all__ = [
     "MONTH_ORDER",
     "compute_correlation_matrix",
     "compute_correlation_matrix_for_variables",
+    "compute_correlation_matrices_for_methods",
     "compute_lagged_correlation",
     "compute_rolling_correlation_series",
     "compute_rolling_correlations_for_pairs",
+    "transform_correlation_matrix",
     "build_event_aligned_segments",
     "detect_threshold_events",
     "filter_by_condition",

meteo/analysis/correlations.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Literal, Sequence
+from typing import Callable, Literal, Sequence
 
 import numpy as np
 import pandas as pd
@@ -11,13 +11,16 @@ from meteo.variables import Variable
 
 from .core import _ensure_datetime_index
 
-__all__ = ['compute_correlation_matrix', 'compute_correlation_matrix_for_variables', 'compute_lagged_correlation', 'compute_rolling_correlation_series', 'compute_rolling_correlations_for_pairs']
+__all__ = ['compute_correlation_matrix', 'compute_correlation_matrix_for_variables', 'compute_correlation_matrices_for_methods', 'compute_lagged_correlation', 'compute_rolling_correlation_series', 'compute_rolling_correlations_for_pairs', 'transform_correlation_matrix']
+
+CorrelationMethod = Literal["pearson", "spearman", "kendall"]
+CorrelationTransform = Literal["identity", "absolute", "square"]
 
 
 def compute_correlation_matrix(
     df: pd.DataFrame,
     *,
-    method: Literal["pearson", "spearman"] = "pearson",
+    method: CorrelationMethod = "pearson",
 ) -> pd.DataFrame:
     """
     Calcule la matrice de corrélation entre toutes les colonnes numériques
@@ -36,7 +39,7 @@ def compute_correlation_matrix_for_variables(
     df: pd.DataFrame,
     variables: Sequence[Variable],
     *,
-    method: Literal["pearson", "spearman"] = "pearson",
+    method: CorrelationMethod = "pearson",
 ) -> pd.DataFrame:
     """
     Calcule la matrice de corrélation pour un sous-ensemble de variables,
@@ -70,6 +73,46 @@ def compute_correlation_matrix_for_variables(
     corr = corr.loc[columns, columns]
     return corr
 
+
+def transform_correlation_matrix(
+    corr: pd.DataFrame,
+    *,
+    transform: CorrelationTransform | Callable[[pd.DataFrame], pd.DataFrame] = "identity",
+) -> pd.DataFrame:
+    """Applique une transformation générique sur une matrice de corrélation."""
+
+    if callable(transform):
+        return transform(corr)
+
+    if transform == "identity":
+        return corr
+    if transform == "absolute":
+        return corr.abs()
+    if transform == "square":
+        return corr.pow(2)
+
+    raise ValueError(f"Transformation de corrélation inconnue : {transform!r}")
+
+
+def compute_correlation_matrices_for_methods(
+    df: pd.DataFrame,
+    variables: Sequence[Variable],
+    *,
+    methods: Sequence[CorrelationMethod],
+    transform: CorrelationTransform | Callable[[pd.DataFrame], pd.DataFrame] = "identity",
+) -> dict[str, pd.DataFrame]:
+    """Calcule plusieurs matrices de corrélation en une seule passe."""
+
+    if not methods:
+        raise ValueError("La liste des méthodes de corrélation est vide.")
+
+    matrices: dict[str, pd.DataFrame] = {}
+    for method in methods:
+        corr = compute_correlation_matrix_for_variables(df, variables, method=method)
+        matrices[method] = transform_correlation_matrix(corr, transform=transform)
+
+    return matrices
+
 def compute_lagged_correlation(
     df: pd.DataFrame,
     var_x: Variable,

meteo/plots/correlations.py

@@ -48,6 +48,11 @@ def plot_correlation_heatmap(
     output_path: str | Path,
     *,
     annotate: bool = True,
+    title: str | None = None,
+    cmap: str | None = None,
+    vmin: float | None = None,
+    vmax: float | None = None,
+    colorbar_label: str | None = None,
 ) -> Path:
     """
     Trace une heatmap de la matrice de corrélation.
@@ -63,6 +68,14 @@ def plot_correlation_heatmap(
         Chemin du fichier image à écrire.
     annotate :
         Si True, affiche la valeur numérique dans chaque case.
+    title :
+        Titre personalisé (par défaut, libellé générique).
+    cmap :
+        Nom de la palette matplotlib à utiliser (par défaut, palette standard).
+    vmin / vmax :
+        Borne d'échelle de couleurs. Si None, valeurs classiques [-1, 1].
+    colorbar_label :
+        Libellé pour la barre de couleur (par défaut "Corrélation").
     """
     output_path = Path(output_path)
     output_path.parent.mkdir(parents=True, exist_ok=True)
@@ -77,7 +90,16 @@ def plot_correlation_heatmap(
     data = corr.to_numpy()
 
     fig, ax = plt.subplots()
-    im = ax.imshow(data, vmin=-1.0, vmax=1.0)
+    if vmin is None:
+        vmin = -1.0
+    if vmax is None:
+        vmax = 1.0
+
+    im_kwargs = {"vmin": vmin, "vmax": vmax}
+    if cmap is not None:
+        im_kwargs["cmap"] = cmap
+
+    im = ax.imshow(data, **im_kwargs)
 
     # Ticks et labels
     ax.set_xticks(np.arange(len(labels)))
@@ -86,31 +108,45 @@ def plot_correlation_heatmap(
     ax.set_yticklabels(labels)
 
     # Axe en haut/bas selon préférence (ici on laisse en bas)
-    ax.set_title("Matrice de corrélation (coef. de Pearson)")
+    ax.set_title(title or "Matrice de corrélation")
 
     # Barre de couleur
     cbar = plt.colorbar(im, ax=ax)
-    cbar.set_label("Corrélation")
+    cbar.set_label(colorbar_label or "Corrélation")
 
     # Annotation des cases
     if annotate:
         n = data.shape[0]
+        norm = im.norm
+        cmap_obj = im.cmap
+
+        def _text_color(value: float) -> str:
+            rgba = cmap_obj(norm(value))
+            r, g, b, _ = rgba
+            luminance = 0.2126 * r + 0.7152 * g + 0.0722 * b
+            return "white" if luminance < 0.5 else "black"
+
         for i in range(n):
             for j in range(n):
+                val = data[i, j]
                 if i == j:
                     text = "—"
+                elif np.isnan(val):
+                    text = ""
                 else:
-                    val = data[i, j]
-                    if np.isnan(val):
-                        text = ""
-                    else:
-                        text = f"{val:.2f}"
+                    text = f"{val:.2f}"
+
+                if not text:
+                    continue
+
+                color = _text_color(0.0 if np.isnan(val) else val)
                 ax.text(
                     j,
                     i,
                     text,
                     ha="center",
                     va="center",
+                    color=color,
                 )
 
     plt.tight_layout()