85 lines
2.2 KiB
Python
85 lines
2.2 KiB
Python
# meteo/season.py
|
|
from __future__ import annotations
|
|
|
|
from typing import Iterable, Sequence
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
|
|
SEASON_LABELS = np.array(["winter", "spring", "summer", "autumn"])
|
|
MONTH_TO_SEASON_INDEX = {
|
|
12: 0,
|
|
1: 0,
|
|
2: 0,
|
|
3: 1,
|
|
4: 1,
|
|
5: 1,
|
|
6: 2,
|
|
7: 2,
|
|
8: 2,
|
|
9: 3,
|
|
10: 3,
|
|
11: 3,
|
|
}
|
|
|
|
|
|
def _ensure_datetime_index(index: pd.Index) -> pd.DatetimeIndex:
|
|
if not isinstance(index, pd.DatetimeIndex):
|
|
raise TypeError("Cette fonction nécessite un DatetimeIndex.")
|
|
return index
|
|
|
|
|
|
def _season_indices_for_month(months: np.ndarray, hemisphere: str) -> np.ndarray:
|
|
base_indices = np.vectorize(MONTH_TO_SEASON_INDEX.get)(months)
|
|
if hemisphere == "south":
|
|
return (base_indices + 2) % len(SEASON_LABELS)
|
|
return base_indices
|
|
|
|
|
|
def compute_season_series(
|
|
index: pd.Index,
|
|
*,
|
|
hemisphere: str = "north",
|
|
column_name: str = "season",
|
|
) -> pd.Series:
|
|
"""
|
|
Retourne une série catégorielle indiquant la saison météorologique pour chaque timestamp.
|
|
"""
|
|
hemisphere = hemisphere.lower()
|
|
if hemisphere not in {"north", "south"}:
|
|
raise ValueError("hemisphere doit valoir 'north' ou 'south'.")
|
|
|
|
dt_index = _ensure_datetime_index(index)
|
|
month_array = dt_index.month.to_numpy()
|
|
season_indices = _season_indices_for_month(month_array, hemisphere)
|
|
labels = SEASON_LABELS[season_indices]
|
|
return pd.Series(labels, index=dt_index, name=column_name)
|
|
|
|
|
|
def add_season_column(
|
|
df: pd.DataFrame,
|
|
*,
|
|
hemisphere: str = "north",
|
|
column_name: str = "season",
|
|
) -> pd.DataFrame:
|
|
"""
|
|
Ajoute une colonne 'season' (winter/spring/summer/autumn) au DataFrame.
|
|
"""
|
|
series = compute_season_series(df.index, hemisphere=hemisphere, column_name=column_name)
|
|
df[column_name] = series
|
|
return df
|
|
|
|
|
|
def sort_season_labels(
|
|
labels: Iterable[str],
|
|
*,
|
|
order: Sequence[str] | None = None,
|
|
) -> list[str]:
|
|
"""
|
|
Trie la liste fournie en respectant l'ordre saisonnier par défaut.
|
|
"""
|
|
reference = [str(season) for season in (order if order is not None else SEASON_LABELS)]
|
|
label_set = {str(label) for label in labels if label}
|
|
return [season for season in reference if season in label_set]
|