africas2s 0.2.0__py3-none-any.whl

africas2s/__init__.py

@@ -0,0 +1,100 @@
+from .downscale import downscale
+from .optimize import optimize
+from .ensemble import ensemble, pool_ensembles
+from .skill import skill, SkillReport
+from .compare import skill_compare, ComparisonReport
+from .pev import prediction_error_variance
+from .flex_forecast import flex_forecast, FlexForecastResult
+from .pipelines import seasonal_mme, SeasonalMMEResult
+from .indices import Index
+from .calibrate import calibrate, CalibrateResult, LogitConfig
+from .climate import (
+    accumulate,
+    frequency_below,
+    percent_of_normal,
+    percentile_of,
+    rank_of_record,
+    seasonal_reduce,
+    seasonal_stack,
+)
+from .analog import (
+    AnalogSet, analogs_from_years, analogs_from_index, analogs_from_field, analogs_where,
+)
+from .completion import complete, CompletionResult
+from .series import quantile_map, error_bounds, ErrorBounds
+from .methods.smoothed_regression import seasonal_coefficients
+from .io import write_terciles, tercile_mae
+from .terrain import load_terrain
+from .combine import combine_terciles, mask_by_skill, dry_mask
+from .plotting.forecasts import plot_tercile_forecast as plot_terciles
+from .plotting.scenarios import plot_accumulation_scenarios, plot_index_scatter
+from .plotting.maps import natural_earth_borders, plot_field_map, plot_choropleth
+from .plotting.style import TercileStyle
+from .plotting.forecasts import plot_field, plot_tercile_comparison, render_styled_terciles
+from . import methods     # trigger registration
+from . import metrics     # trigger registration
+from .metrics import loo_predict, loo_corr, permutation_test, fdr
+from . import strategies  # trigger registration
+from . import time        # calendar / season-step utilities
+from . import aggregations  # daily-rainfall season timing and dry spells
+from . import training     # train / inference separation (§10.2)
+from .training import train
+
+__all__ = [
+    "downscale",
+    "train",
+    "optimize",
+    "ensemble",
+    "pool_ensembles",
+    "skill",
+    "SkillReport",
+    "skill_compare",
+    "ComparisonReport",
+    "prediction_error_variance",
+    "flex_forecast",
+    "FlexForecastResult",
+    "seasonal_mme",
+    "SeasonalMMEResult",
+    "Index",
+    "calibrate",
+    "CalibrateResult",
+    "LogitConfig",
+    "accumulate",
+    "frequency_below",
+    "percent_of_normal",
+    "percentile_of",
+    "rank_of_record",
+    "seasonal_reduce",
+    "seasonal_stack",
+    "AnalogSet",
+    "analogs_from_years",
+    "analogs_from_index",
+    "analogs_from_field",
+    "analogs_where",
+    "complete",
+    "CompletionResult",
+    "quantile_map",
+    "error_bounds",
+    "ErrorBounds",
+    "seasonal_coefficients",
+    "write_terciles",
+    "tercile_mae",
+    "load_terrain",
+    "combine_terciles",
+    "mask_by_skill",
+    "dry_mask",
+    "loo_predict",
+    "loo_corr",
+    "permutation_test",
+    "fdr",
+    "plot_terciles",
+    "plot_accumulation_scenarios",
+    "plot_index_scatter",
+    "plot_field_map",
+    "plot_choropleth",
+    "natural_earth_borders",
+    "TercileStyle",
+    "plot_field",
+    "render_styled_terciles",
+    "plot_tercile_comparison",
+]

africas2s/_optional.py

@@ -0,0 +1,10 @@
+def require_optional(package_name, install_hint=None):
+    """Import and return a package, raising a clear error if missing."""
+    try:
+        return __import__(package_name)
+    except ImportError:
+        hint = install_hint or f"pip install {package_name}"
+        raise ImportError(
+            f"{package_name} is required for this feature. "
+            f"Install with: {hint}"
+        )

africas2s/_spatial.py

@@ -0,0 +1,30 @@
+"""Shared lat/lon dimension resolution for gridded DataArrays.
+
+A single place to recognize the spatial-axis naming conventions africas2s
+accepts (lat/latitude/Y/y and lon/longitude/X/x), so every entry point —
+indices, ensemble regression, logistic calibration, plotting — agrees on which
+dims are spatial. Adding a new alias here reaches all of them at once.
+"""
+from __future__ import annotations
+
+import xarray as xr
+
+_LAT_ALIASES = ("lat", "latitude", "Y", "y")
+_LON_ALIASES = ("lon", "longitude", "X", "x")
+
+
+def spatial_dims(da: xr.DataArray, *, context: str = "data") -> tuple[str, str]:
+    """Return the ``(lat_dim, lon_dim)`` names on ``da``.
+
+    Resolves the common aliases (lat/latitude/Y/y and lon/longitude/X/x).
+    Raises ``ValueError`` naming ``context`` when either axis can't be found.
+    """
+    lat = next((d for d in _LAT_ALIASES if d in da.dims), None)
+    lon = next((d for d in _LON_ALIASES if d in da.dims), None)
+    if lat is None or lon is None:
+        raise ValueError(
+            f"{context} could not find lat/lon dims on data with dims "
+            f"{tuple(da.dims)}; expected one of {_LAT_ALIASES} and "
+            f"{_LON_ALIASES}."
+        )
+    return lat, lon

africas2s/_warnings.py

@@ -0,0 +1,30 @@
+"""Scoped suppression of numpy's benign "All-NaN slice encountered" warning.
+
+AfricaS2S computes per-cell statistics (tercile boundaries via ``xarray.quantile``,
+climatologies via ``nanmean``, logit labels via ``nanpercentile``) over gridded
+fields whose masked cells — ocean, out-of-region, dry-masked — are legitimately
+all-NaN along the reduction axis. numpy warns "All-NaN slice encountered" on those
+slices even though NaN-in / NaN-out is exactly the intended result (the cell has
+no forecast and is masked downstream).
+
+Decorate a public entry point with :func:`quiet_all_nan_slices` to silence just
+that one message for the duration of the call. It uses ``catch_warnings`` so the
+process-wide warning filters are restored on exit — no global state is mutated,
+and only this specific benign message is affected.
+"""
+from __future__ import annotations
+
+import functools
+import warnings
+
+
+def quiet_all_nan_slices(fn):
+    """Wrap ``fn`` so numpy's "All-NaN slice encountered" RuntimeWarning is
+    suppressed while it runs (masked grid cells are expected to be all-NaN)."""
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore", "All-NaN slice encountered", RuntimeWarning)
+            return fn(*args, **kwargs)
+    return wrapper

africas2s/aggregations/__init__.py

@@ -0,0 +1,21 @@
+"""Daily-rainfall aggregations: season timing and dry spells.
+
+Access is module-qualified, following ``africas2s.time`` and
+``africas2s.tercile``::
+
+    from africas2s import aggregations as agg
+
+    onset = agg.onset(daily, season="MAM")
+"""
+from .spells import DrySpellResult, dry_spell
+from .timing import ONSET_DEFAULTS, TimingResult, cessation, onset, season_length
+
+__all__ = [
+    "onset",
+    "cessation",
+    "season_length",
+    "dry_spell",
+    "TimingResult",
+    "DrySpellResult",
+    "ONSET_DEFAULTS",
+]

africas2s/aggregations/_runs.py

@@ -0,0 +1,73 @@
+"""Rolling totals and consecutive-run lengths over a daily step axis.
+
+Onset's false-start guard, cessation, and dry-spell statistics are one question
+read three ways: where are the runs of consecutive days at or below a
+threshold. This module answers it once, for every cell, year and member
+simultaneously. Nothing here knows what a rainy season is.
+"""
+from __future__ import annotations
+
+import numpy as np
+import xarray as xr
+
+__all__ = ["rolling_total", "dry_run_lengths", "has_data"]
+
+
+def rolling_total(rain: xr.DataArray, window: int, *, dim: str = "step") -> xr.DataArray:
+    """Total over ``window`` steps *starting* at each position.
+
+    ``xarray.rolling`` reports the window *ending* at each position, so the
+    result is shifted back by ``window - 1``. A position whose window runs off
+    the end of the axis, or contains a NaN, is NaN: an accumulation cannot be
+    confirmed from incomplete data.
+    """
+    if window < 1:
+        raise ValueError(f"window must be >= 1, got {window}")
+    rolled = rain.rolling({dim: window}, min_periods=window).sum()
+    return rolled.shift({dim: -(window - 1)})
+
+
+def _run_lengths_1d(dry: np.ndarray) -> np.ndarray:
+    """Length of the run of True ending at each position; 0 where False.
+
+    Walks the axis once. ``maximum.accumulate`` over the indices of False
+    entries gives, at each position, the index of the most recent False, and
+    the distance to it is the run length.
+    """
+    n = dry.shape[-1]
+    idx = np.arange(n)
+    reset = np.where(~dry, idx, -1)
+    last_false = np.maximum.accumulate(reset, axis=-1)
+    return np.where(dry, idx - last_false, 0).astype(np.int64)
+
+
+def dry_run_lengths(rain: xr.DataArray, threshold: float, *,
+                    dim: str = "step") -> xr.DataArray:
+    """Length of the consecutive at-or-below-``threshold`` run ending at each step.
+
+    Wet steps are 0. The threshold is inclusive, matching the "1 mm/day or
+    less" wording of the criterion.
+
+    NaN steps count as not dry, because ``NaN <= threshold`` is False. A gap in
+    the record therefore breaks a run rather than extending it, which is the
+    conservative direction: missing data never manufactures a dry spell.
+    """
+    dry = rain <= threshold
+    return xr.apply_ufunc(
+        _run_lengths_1d,
+        dry,
+        input_core_dims=[[dim]],
+        output_core_dims=[[dim]],
+        dask="parallelized",
+        output_dtypes=[np.int64],
+    ).transpose(*rain.dims)
+
+
+def has_data(rain: xr.DataArray, *, dim: str = "step") -> xr.DataArray:
+    """True where the cell has at least one finite value along ``dim``.
+
+    The three-state ``occurred`` contract depends on this. Rolling sums over an
+    all-NaN cell yield NaN, NaN fails every comparison, and a naive ``any``
+    would report the cell as a failed season rather than as absent data.
+    """
+    return rain.notnull().any(dim)

africas2s/aggregations/spells.py

@@ -0,0 +1,93 @@
+"""Dry-spell statistics within a season.
+
+The same run-length kernel that powers onset's false-start guard, read as
+counting rather than dating. Emmett's observation in the design thread: because
+the onset criterion already needs a dry-spell definition, the dry-spell metric
+falls out of it rather than being built separately.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import xarray as xr
+
+from ..climate import seasonal_stack
+from ._runs import dry_run_lengths, has_data
+from .timing import (
+    _provenance,
+    _require_daily,
+    _stamp,
+    _warn_on_incomplete_record,
+    _warn_on_suspicious_units,
+)
+
+__all__ = ["DrySpellResult", "dry_spell"]
+
+
+@dataclass
+class DrySpellResult:
+    """Dry-spell statistics, one value per season.
+
+    ``longest`` is the length of the longest run at or below the threshold,
+    whatever its length. ``count`` is the number of distinct runs reaching
+    ``dry_spell_days``. Both are NaN for cells with no data, and 0.0 for cells
+    that simply had no dry days.
+    """
+
+    longest: xr.DataArray
+    count: xr.DataArray
+    params: dict = field(default_factory=dict)
+
+
+def dry_spell(
+    daily: xr.DataArray,
+    season,
+    *,
+    dry_spell_mm: float = 1.0,
+    dry_spell_days: int = 7,
+    time_dim: str = "time",
+) -> DrySpellResult:
+    """Longest dry run and count of qualifying dry runs, per season.
+
+    A run is dry when each of its days receives ``dry_spell_mm`` or less. The
+    defaults are the dry-spell half of the default onset criterion.
+
+    No tail is stacked: unlike onset, this question is entirely contained
+    within the season. ``daily`` must be daily; any other cadence raises,
+    because a dry-spell length counted in days has no meaning on a coarser
+    axis.
+    """
+    _require_daily(daily, time_dim)
+
+    # Metres-instead-of-millimetres is as silently wrong here as it is for
+    # onset, in the opposite direction: every day falls under the threshold and
+    # the whole season reads as one dry spell.
+    _warn_on_suspicious_units(daily)
+
+    rain = seasonal_stack(daily, season, time_dim=time_dim)
+
+    # A record that stops mid-season does not return NaN here: seasonal_stack
+    # NaN-pads the gap, NaN counts as not-dry, so every run spanning the cut is
+    # silently truncated and `longest` under-reports. Warn rather than let a
+    # confident wrong number through.
+    _warn_on_incomplete_record(
+        daily, season, rain["year"].values, 0, time_dim,
+        window="full season",
+        consequence=("Dry runs are cut where the record ends, so longest and "
+                     "count under-report rather than reporting no data."),
+    )
+
+    lengths = dry_run_lengths(rain, dry_spell_mm)
+    ok = has_data(rain)
+
+    longest = lengths.max("step").astype(float).where(ok)
+
+    # A run reaches length L at exactly one step, the L-th day of that run, so
+    # equality counts each qualifying run once regardless of how long it grows.
+    count = (lengths == dry_spell_days).sum("step").astype(float).where(ok)
+
+    params = dict(dry_spell_mm=dry_spell_mm, dry_spell_days=dry_spell_days)
+    prov = _provenance(params, season)
+    arrays = {"longest": longest, "count": count}
+    _stamp(arrays, prov)
+    return DrySpellResult(**arrays, params=prov)