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

meteo/plots/__init__.py

@@ -13,7 +13,9 @@ from .calendar_overview import (
 from .correlations import (
     plot_correlation_heatmap,
     plot_lagged_correlation,
+    plot_lagged_correlation_multi,
     plot_rolling_correlation_heatmap,
+    CorrelationBand,
 )
 from .rain import plot_daily_rainfall_hyetograph, plot_rainfall_by_season
 from .relationships import (
@@ -55,7 +57,9 @@ __all__ = [
     "rainfall_daily_total_series",
     "plot_correlation_heatmap",
     "plot_lagged_correlation",
+    "plot_lagged_correlation_multi",
     "plot_rolling_correlation_heatmap",
+    "CorrelationBand",
     "plot_daily_rainfall_hyetograph",
     "plot_rainfall_by_season",
     "plot_event_composite",

meteo/plots/correlations.py

@@ -3,16 +3,23 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Sequence
+from typing import Iterable, Sequence
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 
+from meteo.correlation_presets import CorrelationBand
 from .base import export_plot_dataset
 from meteo.variables import Variable
 
-__all__ = ['plot_lagged_correlation', 'plot_correlation_heatmap', 'plot_rolling_correlation_heatmap']
+__all__ = [
+    'plot_lagged_correlation',
+    'plot_lagged_correlation_multi',
+    'plot_correlation_heatmap',
+    'plot_rolling_correlation_heatmap',
+    'CorrelationBand',
+]
 
 
 def plot_lagged_correlation(
@@ -42,13 +49,89 @@ def plot_lagged_correlation(
 
     return output_path.resolve()
 
+
+def plot_lagged_correlation_multi(
+    lag_series: dict[str, pd.Series],
+    var_x: Variable,
+    var_y: Variable,
+    output_path: str | Path,
+    *,
+    title_suffix: str | None = None,
+    ylabel: str = "Corrélation",
+    y_limits: tuple[float, float] | None = None,
+    thresholds: Sequence[float] | None = None,
+    bands: Iterable["CorrelationBand"] | None = None,
+) -> Path:
+    """
+    Trace plusieurs courbes de corrélation en fonction du lag (ex. Pearson/Spearman).
+    """
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    df = pd.concat(lag_series, axis=1)
+    export_plot_dataset(df, output_path)
+
+    plt.figure()
+    colors = ["#1f77b4", "#d1495b", "#2ca02c", "#9467bd"]
+    for idx, (label, series) in enumerate(df.items()):
+        plt.plot(series.index, series, label=label, color=colors[idx % len(colors)], linewidth=1.6)
+
+    ax = plt.gca()
+    if bands:
+        xmin, xmax = df.index.min(), df.index.max()
+        for band in bands:
+            ax.axhspan(band.min_value, band.max_value, color=band.color, alpha=0.25, zorder=0)
+            ax.text(
+                xmax,
+                (band.min_value + band.max_value) / 2.0,
+                band.label,
+                ha="right",
+                va="center",
+                fontsize=8,
+                color="#444444",
+                bbox=dict(facecolor="white", edgecolor="none", alpha=0.6, pad=1.5),
+            )
+
+    plt.axvline(0, linestyle="--", color="#666666", linewidth=1.0)
+    plt.xlabel("Décalage (minutes)\n(lag > 0 : X précède Y)")
+    plt.ylabel(ylabel)
+    title = f"Corrélation décalée : {var_x.label} → {var_y.label}"
+    if title_suffix:
+        title = f"{title} ({title_suffix})"
+    plt.title(title)
+    if thresholds:
+        xmin, xmax = plt.xlim()
+        for thr in thresholds:
+            plt.axhline(thr, color="#999999", linestyle="--", linewidth=1.0, alpha=0.85)
+            plt.text(
+                xmax,
+                thr,
+                f"{thr:.2f}",
+                ha="right",
+                va="center",
+                fontsize=8,
+                color="#555555",
+                bbox=dict(facecolor="white", edgecolor="none", alpha=0.7, pad=1.5),
+            )
+    if y_limits is not None:
+        plt.ylim(*y_limits)
+    plt.grid(True, alpha=0.7)
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150)
+    plt.close()
+
+    return output_path.resolve()
+
 def plot_correlation_heatmap(
     corr: pd.DataFrame,
     variables: Sequence[Variable],
     output_path: str | Path,
     *,
     annotate: bool = True,
+    annotate_values: "pd.DataFrame | None" = None,
     title: str | None = None,
+    figsize: tuple[float, float] | None = None,
     cmap: str | None = None,
     vmin: float | None = None,
     vmax: float | None = None,
@@ -89,7 +172,13 @@ def plot_correlation_heatmap(
 
     data = corr.to_numpy()
 
-    fig, ax = plt.subplots()
+    if figsize is None:
+        n = len(variables)
+        # Augmente la taille pour laisser respirer les annotations
+        side = max(6.0, n * 0.9)
+        figsize = (side, side)
+
+    fig, ax = plt.subplots(figsize=figsize)
     if vmin is None:
         vmin = -1.0
     if vmax is None:
@@ -117,6 +206,11 @@ def plot_correlation_heatmap(
     # Annotation des cases
     if annotate:
         n = data.shape[0]
+        annot_data = (
+            annotate_values.loc[columns, columns].to_numpy()
+            if annotate_values is not None
+            else data
+        )
         norm = im.norm
         cmap_obj = im.cmap
 
@@ -128,18 +222,23 @@ def plot_correlation_heatmap(
 
         for i in range(n):
             for j in range(n):
-                val = data[i, j]
+                val_corr = data[i, j]
+                val_annot = annot_data[i, j]
                 if i == j:
                     text = "—"
-                elif np.isnan(val):
+                elif isinstance(val_annot, (float, int, np.floating)) and np.isnan(val_annot):
                     text = ""
                 else:
-                    text = f"{val:.2f}"
+                    # si annotate_values est fourni, on affiche la valeur annotée brute
+                    if annotate_values is not None:
+                        text = str(val_annot)
+                    else:
+                        text = f"{val_corr:.2f}"
 
                 if not text:
                     continue
 
-                color = _text_color(0.0 if np.isnan(val) else val)
+                color = _text_color(0.0 if np.isnan(val_corr) else val_corr)
                 ax.text(
                     j,
                     i,