1

128 lines
3.4 KiB
Python

"""Fonctions utilitaires pour générer des hexbins à partir de scénarios."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence
import pandas as pd
from meteo.correlation_presets import HexbinScenario
from meteo.variables import VARIABLES_BY_KEY, Variable
from .relationships import plot_hexbin_with_third_variable
__all__ = [
"HexbinPlotResult",
"prepare_hexbin_dataframe",
"generate_hexbin_for_scenario",
"generate_hexbin_scenarios",
]
@dataclass(frozen=True)
class HexbinPlotResult:
"""Résultat d'un hexbin généré à partir d'un scénario prédéfini."""
scenario: HexbinScenario
var_x: Variable
var_y: Variable
var_color: Variable
output_path: Path
point_count: int
def _get_variables_for_scenario(scenario: HexbinScenario) -> tuple[Variable, Variable, Variable]:
try:
return (
VARIABLES_BY_KEY[scenario.key_x],
VARIABLES_BY_KEY[scenario.key_y],
VARIABLES_BY_KEY[scenario.key_color],
)
except KeyError as exc:
raise KeyError(f"Variable inconnue dans le scénario hexbin : {exc}") from exc
def prepare_hexbin_dataframe(
df: pd.DataFrame,
var_x: Variable,
var_y: Variable,
var_color: Variable,
) -> pd.DataFrame:
"""Sélectionne et nettoie les colonnes utiles à un hexbin."""
columns = [var_x.column, var_y.column, var_color.column]
missing = [col for col in columns if col not in df.columns]
if missing:
raise KeyError(f"Colonnes absentes dans le DataFrame : {missing}")
return df[columns].dropna()
def generate_hexbin_for_scenario(
df: pd.DataFrame,
scenario: HexbinScenario,
*,
base_output_dir: str | Path,
cmap: str = "viridis",
) -> HexbinPlotResult:
"""Construit un hexbin complet à partir d'un scénario et retourne ses métadonnées."""
var_x, var_y, var_color = _get_variables_for_scenario(scenario)
cleaned_df = prepare_hexbin_dataframe(df, var_x, var_y, var_color)
output_dir = Path(base_output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / scenario.filename
reduce_func = scenario.get_reduce_func()
reduce_label = scenario.get_reduce_label()
plot_hexbin_with_third_variable(
df=cleaned_df,
var_x=var_x,
var_y=var_y,
var_color=var_color,
output_path=output_path,
gridsize=scenario.gridsize,
mincnt=scenario.mincnt,
reduce_func=reduce_func,
reduce_func_label=reduce_label,
cmap=cmap,
)
return HexbinPlotResult(
scenario=scenario,
var_x=var_x,
var_y=var_y,
var_color=var_color,
output_path=output_path.resolve(),
point_count=len(cleaned_df),
)
def generate_hexbin_scenarios(
df: pd.DataFrame,
scenarios: Sequence[HexbinScenario],
*,
base_output_dir: str | Path,
cmap: str = "viridis",
) -> list[HexbinPlotResult]:
"""Génère en série plusieurs hexbins prédéfinis et retourne les résultats."""
if not scenarios:
raise ValueError("Aucun scénario hexbin fourni.")
results: list[HexbinPlotResult] = []
for scenario in scenarios:
result = generate_hexbin_for_scenario(
df=df,
scenario=scenario,
base_output_dir=base_output_dir,
cmap=cmap,
)
results.append(result)
return results