ClimateGraph 0.4.0__py3-none-any.whl

ClimateGraph/plot/primitive.py

@@ -0,0 +1,99 @@
+"""Drawable units for composed figures.
+
+A primitive draws one layer onto shared axes, which is what a ``custom`` plot
+composes.
+"""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, ClassVar
+
+from pydantic import BaseModel
+
+from ClimateGraph.utils.registry import RegistryMixin
+
+if TYPE_CHECKING:
+    import matplotlib as mpl
+    import xarray as xr
+
+
+class Primitive(RegistryMixin, ABC):
+    """The Primitive abstract class.
+
+    A minimal rendering unit — the atomic building block of a ``type: custom``
+    plot. A primitive does no data processing of its own: the plot orchestrator
+    resolves the dataset/var, applies the domain, filters time, reduces dims and
+    converts units, then hands the prepared ``xr.DataArray`` to ``render`` along
+    with the coordinate names bound to each axis. The primitive only draws and
+    returns the matplotlib artist, so legends and colorbars stay a figure-level
+    concern of the orchestrator.
+
+    Mirrors the ``Plot`` / ``Domain`` ABC + registry shape via ``RegistryMixin``:
+    concrete primitives self-register by class name and ``aliases``, and their
+    ``config`` classes are gathered into a discriminated union by
+    ``build_config_union``.
+    """
+
+    registry: ClassVar[dict[str, type["Primitive"]]] = {}
+    aliases: ClassVar[list[str]] = []
+    config: ClassVar[type[BaseModel] | None] = None
+
+    required_coords: ClassVar[frozenset[str]] = frozenset()
+    wants_colorbar: ClassVar[bool] = False
+    wants_legend: ClassVar[bool] = False
+
+    def __init__(self, name: str, subplot_config: BaseModel, **kwargs):
+        """Primitive initialization.
+
+        Parameters
+        ----------
+        name : str
+            Name of the owning plot, used for reference / error messages.
+        subplot_config : BaseModel
+            The concrete subplot config (a ``SubplotConfig`` subclass) driving
+            this primitive.
+        """
+        self.name = name
+        self.subplot_config = subplot_config
+        self.style = dict(getattr(subplot_config, "model_extra", None) or {})
+
+    @classmethod
+    def get_primitive_class(cls, name: str):
+        """Look up a registered primitive subclass by name."""
+        return cls.get_class(name)
+
+    @classmethod
+    def check_primitive_class(cls, name: str) -> bool:
+        """Return True if ``type`` names a registered primitive."""
+        return cls.check_class(name)
+
+    @abstractmethod
+    def render(
+        self,
+        ax: "mpl.axes.Axes",
+        data: "xr.DataArray",
+        *,
+        x: str | None,
+        y: str | None,
+        label: str | None,
+    ) -> "mpl.artist.Artist":
+        """Draw the prepared data onto ``ax`` and return the artist.
+
+        Parameters
+        ----------
+        ax : matplotlib Axes
+            The shared axes to draw on.
+        data : xr.DataArray
+            Fully prepared data (domain/time/dim-reduce/unit already applied).
+        x : str | None
+            Coordinate name bound to the x axis.
+        y : str | None
+            Coordinate name bound to the y axis (None for 1-D primitives).
+        label : str | None
+            Legend / colorbar label for this artist.
+
+        Returns
+        -------
+        matplotlib Artist
+            The artist created, so the orchestrator can build legends/colorbars.
+        """
+        pass

ClimateGraph/plot/primitives.py

@@ -0,0 +1,160 @@
+"""The concrete primitives and their subplot config models."""
+
+from typing import ClassVar, Literal
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict, model_validator
+
+from ClimateGraph.utils.general_utils import ReductionMethodEnum
+
+from .primitive import Primitive
+
+
+def _drop_nan_points(
+    x: np.ndarray, y: np.ndarray, vals: np.ndarray
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Drop entries whose value is NaN, keeping x/y/value aligned."""
+    keep = ~np.isnan(vals)
+    return x[keep], y[keep], vals[keep]
+
+
+class SubplotConfig(BaseModel):
+    """Shared config for every primitive subplot.
+
+    Flattened (not nested) so a ``subplots:`` entry reads naturally in YAML —
+    ``type: contourf`` sits next to ``dataset:``/``x:``/``y:`` directly. Extra
+    keys are allowed so matplotlib styling (color, linewidth, alpha, ...) flows
+    straight through to the primitive's draw call via ``model_extra``.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    # Coordinate roles the concrete primitive requires (overridden per subclass).
+    required_coords: ClassVar[frozenset[str]] = frozenset()
+
+    dataset: str
+    var: str | None = None
+    x: str | None = None
+    y: str | None = None
+    # Optional per-subplot time override. When set it wins over the plot-level
+    # `time` for this subplot only, so heterogeneous times can overlay on one
+    # figure (e.g. the same date across two years). When None the subplot follows
+    # the plot's time iteration.
+    time: str | None = None
+    dim_reduce: dict[str, str | dict] | None = None
+    reduction_method: ReductionMethodEnum = ReductionMethodEnum.mean
+    unit: str | None = None
+    label: str | None = None
+
+    @model_validator(mode="after")
+    def _check_required_coords(self):
+        """Check the subplot declares the coordinates this primitive needs."""
+        missing = sorted(c for c in self.required_coords if getattr(self, c) is None)
+        if missing:
+            raise ValueError(
+                f"{self.type!r} subplot requires coordinate role(s) {missing} "
+                f"(set them under the subplot, e.g. x:/y:)."
+            )
+        return self
+
+
+class SeriesConfig(SubplotConfig):
+    """A line along one dimension (``x``)."""
+
+    required_coords: ClassVar[frozenset[str]] = frozenset({"x"})
+    type: Literal["series", "line"]
+
+
+class Series(Primitive):
+    """Primitive: a line of the variable along the ``x`` coordinate."""
+
+    config = SeriesConfig
+    aliases = ["series", "line"]
+    required_coords = frozenset({"x"})
+    wants_legend = True
+
+    def render(self, ax, data, *, x, y, label):
+        """Draw a line along one dimension."""
+        xvals = (
+            data[x].values
+            if x in data.coords
+            else np.arange(data.sizes.get(x, data.size))
+        )
+        (line,) = ax.plot(xvals, data.values, label=label, **self.style)
+        return line
+
+
+class ContourFillConfig(SubplotConfig):
+    """A filled contour over a 2-D (``x``, ``y``) grid."""
+
+    required_coords: ClassVar[frozenset[str]] = frozenset({"x", "y"})
+    type: Literal["contourf", "contour_fill", "contourfill"]
+    levels: int = 10
+    cmap: str = "viridis"
+
+
+class ContourFill(Primitive):
+    """Primitive: a filled contour of the variable on x/y.
+
+    No CRS / projection awareness (unlike ``SpatialOverlay``) — plain axes,
+    generic by design. Projection kwargs, if ever needed, flow through
+    ``model_extra``.
+    """
+
+    config = ContourFillConfig
+    aliases = ["contourf", "contour_fill", "contourfill"]
+    required_coords = frozenset({"x", "y"})
+    wants_colorbar = True
+
+    def render(self, ax, data, *, x, y, label):
+        """Draw filled contours over two coordinates."""
+        cfg = self.subplot_config
+        return data.plot.contourf(
+            x=x,
+            y=y,
+            ax=ax,
+            levels=cfg.levels,
+            cmap=cfg.cmap,
+            add_colorbar=False,
+            add_labels=False,
+            **self.style,
+        )
+
+
+class PointsConfig(SubplotConfig):
+    """A scatter of points at (``x``, ``y``) coloured by value."""
+
+    required_coords: ClassVar[frozenset[str]] = frozenset({"x", "y"})
+    type: Literal["points", "scatter"]
+    markersize: float = 40.0
+    cmap: str = "viridis"
+    edgecolor: str | None = "k"
+    drop_nans: bool = False
+
+
+class Points(Primitive):
+    """Primitive: a coordinate scatter coloured by the variable value."""
+
+    config = PointsConfig
+    aliases = ["points", "scatter"]
+    required_coords = frozenset({"x", "y"})
+    wants_colorbar = True
+
+    def render(self, ax, data, *, x, y, label):
+        """Draw coloured points at two coordinates."""
+        cfg = self.subplot_config
+        xv = np.asarray(data[x].values).ravel()
+        yv = np.asarray(data[y].values).ravel()
+        cv = np.asarray(data.values).ravel()
+        if cfg.drop_nans:
+            xv, yv, cv = _drop_nan_points(xv, yv, cv)
+        return ax.scatter(
+            xv,
+            yv,
+            c=cv,
+            s=cfg.markersize,
+            cmap=cfg.cmap,
+            edgecolor=cfg.edgecolor,
+            label=label,
+            **self.style,
+        )

ClimateGraph/reader/__init__.py

@@ -0,0 +1,15 @@
+"""Readers: turning files into datasets."""
+
+import importlib
+import pkgutil
+
+for module in pkgutil.iter_modules(__path__):
+    importlib.import_module(f"{__name__}.{module.name}")
+
+from .reader import Reader
+
+__all__ = [
+    "Reader",
+]
+
+# Eventually add downloaders

ClimateGraph/reader/reader/__init__.py

@@ -0,0 +1,14 @@
+"""The reader base class and the per-topology reader families."""
+
+import importlib
+import pkgutil
+
+from .reader import Reader
+
+__all__ = [
+    "Reader",
+    #    "Downloader",
+]
+
+for module in pkgutil.iter_modules(__path__):
+    importlib.import_module(f"{__name__}.{module.name}")

ClimateGraph/reader/reader/point_surface/__init__.py

@@ -0,0 +1,7 @@
+"""Readers for station and point-observation data."""
+
+import importlib
+import pkgutil
+
+for module in pkgutil.iter_modules(__path__):
+    importlib.import_module(f"{__name__}.{module.name}")

ClimateGraph/reader/reader/point_surface/csv_base.py

@@ -0,0 +1,135 @@
+"""Base reader for station data stored as tables.
+
+Drives the per-file loop itself, since tabular inputs cannot go through
+``open_mfdataset``, and attaches station coordinates from a metadata file.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pandas as pd
+import xarray as xr
+
+from ..reader import ReadSpec
+from .default import DefaultPointSurfaceReader
+
+
+class CSVPointSurfaceReader(DefaultPointSurfaceReader):
+    """Shared base for CSV/Excel point-surface readers.
+
+    The layout-specific readers (station-per-file, single-file,
+    variable-per-file) live in CSV/Excel rather than NetCDF and share three
+    concerns this base centralises:
+
+    1. **Per-file open is not NetCDF.** The default ``_open_many`` uses
+       ``xr.open_mfdataset``; here it routes safe-mode through the exact same
+       per-file loop the unsafe path runs, so the two modes are identical for
+       CSV input (subclasses only implement ``_open_one`` / ``_to_xarray``).
+    2. **A station-metadata join.** ``_postprocess`` attaches every column of a
+       sidecar metadata table onto the ``site`` dimension as a coordinate,
+       including the ``latitude``/``longitude`` that ``PointSurface._set_geom``
+       requires.
+    3. **The canonical-name rename** is reused from ``DefaultPointSurfaceReader``
+       (variable-per-file layouts set ``restrict_rename_to_present``).
+
+    Subclasses set ``default_station_key``, ``default_time_col`` and
+    ``latlon_aliases`` and implement ``_open_one`` / ``_to_xarray`` / ``_join``.
+    """
+
+    # Pieces hold one (variable-per-file) or a disjoint set (station-per-file)
+    # of vars; filter the rename map to what's actually present in each piece.
+    restrict_rename_to_present = True
+
+    # Subclass-supplied defaults; overridable per data block via ``spec.extras``.
+    default_station_key: str = "station_id"
+    default_time_col: str = "time"
+    # Maps the metadata file's lat/lon column names onto the canonical
+    # ``latitude``/``longitude`` the rest of the pipeline expects.
+    latlon_aliases: dict[str, str] = {}
+
+    # ----- safe == unsafe for CSV ----------------------------------------
+
+    @classmethod
+    def _open_many(cls, paths: list[Path], spec: ReadSpec) -> xr.Dataset:
+        """Run the per-file hooks in a plain loop (mirrors the unsafe path).
+
+        ``xr.open_mfdataset`` only understands NetCDF-like engines, so the
+        default safe-mode implementation can't be used. Routing both modes
+        through the same loop also makes ``safe`` and ``unsafe`` produce
+        byte-identical datasets for CSV input.
+        """
+        pieces = [
+            cls._preprocess(cls._to_xarray(cls._open_one(p, spec), spec), spec)
+            for p in paths
+        ]
+        return cls._join(pieces, spec)
+
+    # ----- metadata join --------------------------------------------------
+
+    @classmethod
+    def _postprocess(cls, ds: xr.Dataset, spec: ReadSpec) -> xr.Dataset:
+        """Attach the station metadata once the values are assembled."""
+        ds = cls._attach_metadata(ds, spec)
+        return super()._postprocess(ds, spec)
+
+    @classmethod
+    def _attach_metadata(cls, ds: xr.Dataset, spec: ReadSpec) -> xr.Dataset:
+        """Join a station-metadata table onto the ``site`` dimension.
+
+        Every metadata column becomes a ``(site,)`` coordinate. Stations with
+        no matching metadata row get NaN-filled coordinates rather than being
+        dropped. The join keys are normalised on both sides via
+        ``_norm_site_key`` so e.g. integer ids in the metadata match string
+        filename-stem ids on the data, or accented station names match their
+        ASCII-uppercased CSV column headers (see subclass overrides).
+        """
+        meta_path = spec.extras.get("metadata")
+        if meta_path is None:
+            raise ValueError(
+                f"{cls.__name__} requires a 'metadata' path in the data block."
+            )
+        meta = cls._read_metadata(Path(meta_path))
+
+        key = spec.extras.get("station_key", cls.default_station_key)
+        if key not in meta.columns:
+            raise ValueError(
+                f"{cls.__name__}: station key {key!r} not in metadata columns "
+                f"{list(meta.columns)}."
+            )
+        meta = meta.set_index(key).rename(columns=cls.latlon_aliases)
+
+        # Normalise the metadata index to the canonical join key and drop any
+        # duplicate stations (keep first) so reindex can align unambiguously.
+        meta.index = [cls._norm_site_key(v) for v in meta.index]
+        meta = meta[~meta.index.duplicated(keep="first")]
+
+        site_keys = [cls._norm_site_key(v) for v in ds["site"].values]
+        meta = meta.reindex(site_keys)
+
+        for col in meta.columns:
+            ds = ds.assign_coords({col: ("site", meta[col].to_numpy())})
+        return ds
+
+    @staticmethod
+    def _read_metadata(path: Path) -> pd.DataFrame:
+        """Read the metadata sidecar — Excel by suffix, else CSV."""
+        if path.suffix.lower() in (".xlsx", ".xls"):
+            return pd.read_excel(path, engine="openpyxl")
+        return pd.read_csv(path)
+
+    @staticmethod
+    def _norm_site_key(value) -> str:
+        """Canonicalise a station id for the metadata join.
+
+        Numeric ids (whether stored as ``int``, ``float`` like ``109.0``, or
+        the string ``"104"``) collapse to their integer string form so that a
+        filename stem matches an ``int`` metadata id. Everything else is
+        stringified and stripped. Subclasses with text station names (e.g.
+        variable-per-file with accented headers) override this.
+        """
+        try:
+            f = float(value)
+        except (ValueError, TypeError):
+            return str(value).strip()
+        return str(int(f)) if f == int(f) else str(f)

ClimateGraph/reader/reader/point_surface/default.py

@@ -0,0 +1,36 @@
+"""Base reader for station data stored as NetCDF."""
+
+import xarray as xr
+
+from ..reader import Reader, ReadSpec
+
+
+class DefaultPointSurfaceReader(Reader):
+    type_aliases = ["DefaultPointSurface"]
+    topology = "PointSurface"
+
+    rename: dict[str, str] = {}
+
+    # When True, drop rename keys that aren't present in the piece before
+    # renaming. NetCDF readers leave this False so a missing variable still
+    # raises (a useful error). Variable-per-file readers (CSV/Excel) set it
+    # True: each piece holds only one variable, so the full rename map would
+    # otherwise reference names absent from that piece and ``rename`` would
+    # raise.
+    restrict_rename_to_present: bool = False
+
+    @classmethod
+    def _preprocess(cls, ds: xr.Dataset, spec: ReadSpec) -> xr.Dataset:
+        """Apply the variable renaming and selection declared in the spec."""
+        rename = dict(cls.rename)
+        rename.update(spec.extras.get("rename", {}))
+        if spec.vars is not None:
+            # Skip composed vars with no file name
+            rename.update(
+                {d["name"]: name for name, d in spec.vars.items() if d.get("name")}
+            )
+        if "x" in ds.coords and "site" not in ds.coords:
+            rename["x"] = "site"
+        if cls.restrict_rename_to_present:
+            rename = {k: v for k, v in rename.items() if k in ds.variables}
+        return ds.rename(rename)

ClimateGraph/reader/reader/point_surface/dmc.py

@@ -0,0 +1,10 @@
+from .default import DefaultPointSurfaceReader
+
+
+class DMC(DefaultPointSurfaceReader):
+    """Dirección Meteorológica de Chile point-surface reader.
+
+    Currently a thin subclass of the default — kept distinct so DMC-specific
+    download/preprocessing (e.g. fetching from the DMC API) can land here
+    without touching SINCA.
+    """

ClimateGraph/reader/reader/point_surface/sinca.py

@@ -0,0 +1,10 @@
+from .default import DefaultPointSurfaceReader
+
+
+class SINCA(DefaultPointSurfaceReader):
+    """SINCA (Sistema de Información Nacional de Calidad del Aire) reader.
+
+    Currently a thin subclass of the default — kept distinct so SINCA-specific
+    download/preprocessing (e.g. fetching from the SINCA API) can land here
+    without touching DMC.
+    """

ClimateGraph/reader/reader/point_surface/single_file.py

@@ -0,0 +1,60 @@
+"""CSV layout: one file containing every station."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pandas as pd
+import xarray as xr
+
+from ..reader import ReadSpec
+from .csv_base import CSVPointSurfaceReader
+
+
+class SingleFileReader(CSVPointSurfaceReader):
+    """Reader for the single-file CSV layout (one file holds everything).
+
+    Layout: one CSV holding every station's record, keyed by a station-code
+    column, plus a ``stations.csv`` metadata table. Originally written for São
+    Paulo (CETESB/QUALAR), whose defaults this keeps; its timestamps in
+    ``local_date`` are tz-aware (``-03:00``) and stored as naive local
+    wall-clock to stay consistent with the other in-situ readers.
+    """
+
+    type_aliases = ["single-file"]
+    default_station_key = "code"
+    default_time_col = "local_date"
+    latlon_aliases = {"lat": "latitude", "lon": "longitude"}
+
+    # Non-numeric columns that ride along as metadata (per ``code``) rather
+    # than as data variables.
+    drop_cols = ["type"]
+
+    @classmethod
+    def _open_one(cls, path: Path, spec: ReadSpec) -> pd.DataFrame:
+        """Read the one CSV holding every station into a frame."""
+        time_col = spec.extras.get("time_col", cls.default_time_col)
+        df = pd.read_csv(path)
+        # Parse the tz-aware timestamp, then drop the tz to keep naive local
+        # wall-clock (xarray cannot hold tz-aware datetimes anyway).
+        df[time_col] = pd.to_datetime(df[time_col]).dt.tz_localize(None)
+        return df
+
+    @classmethod
+    def _to_xarray(cls, raw: pd.DataFrame, spec: ReadSpec) -> xr.Dataset:
+        """Pivot the frame to a dataset indexed by time and site."""
+        time_col = spec.extras.get("time_col", cls.default_time_col)
+        code_col = spec.extras.get("station_key", cls.default_station_key)
+
+        df = raw.drop(columns=cls.drop_cols, errors="ignore")
+        # Collapse to one row per (time, code), taking the first non-null per
+        # column so the partial rows merge rather than colliding.
+        df = df.groupby([time_col, code_col]).first()
+        ds = df.to_xarray()
+        return ds.rename({time_col: "time", code_col: "site"})
+
+    @classmethod
+    def _join(cls, pieces: list[xr.Dataset], spec: ReadSpec) -> xr.Dataset:
+        # Single file in practice; base default handles len == 1.
+        """Concatenate the per-file datasets along time."""
+        return super()._join(pieces, spec)

ClimateGraph/reader/reader/point_surface/station_per_file.py

@@ -0,0 +1,55 @@
+"""CSV layout: one file per station."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pandas as pd
+import xarray as xr
+
+from ..reader import ReadSpec
+from .csv_base import CSVPointSurfaceReader
+
+
+class StationPerFileReader(CSVPointSurfaceReader):
+    """Reader for the station-per-file CSV layout (one station per file).
+
+    Layout: one CSV per station (``{station_id}.csv``) holding a time column
+    plus one column per variable, accompanied by a ``stations.csv`` metadata
+    table. The station id is the file's stem. Originally written for the Chile
+    SINCA per-station export, whose defaults this keeps.
+    """
+
+    type_aliases = ["station-per-file"]
+    default_station_key = "station_id"
+    default_time_col = "timestamp"
+    latlon_aliases = {"latitud": "latitude", "longitud": "longitude"}
+
+    @classmethod
+    def _open_one(cls, path: Path, spec: ReadSpec) -> pd.DataFrame:
+        """Read one station's CSV into a frame."""
+        time_col = spec.extras.get("time_col", cls.default_time_col)
+        df = pd.read_csv(path, parse_dates=[time_col])
+        # Stash the station id (filename stem) for _to_xarray, which has no
+        # access to the path.
+        df.attrs["station_id"] = path.stem
+        return df
+
+    @classmethod
+    def _to_xarray(cls, raw: pd.DataFrame, spec: ReadSpec) -> xr.Dataset:
+        """Convert one station's frame into a single-site dataset."""
+        time_col = spec.extras.get("time_col", cls.default_time_col)
+        df = raw.set_index(time_col)
+        df.index.name = "time"
+        ds = df.to_xarray()
+        ds = ds.expand_dims(site=[raw.attrs["station_id"]])
+        return ds.transpose("time", "site")
+
+    @classmethod
+    def _join(cls, pieces: list[xr.Dataset], spec: ReadSpec) -> xr.Dataset:
+        """Concatenate the per-station datasets along the site dimension."""
+        if len(pieces) == 1:
+            return pieces[0]
+        # One station per file: stack along site (NOT the default time-concat),
+        # outer-joining the time axes so stations with differing coverage align.
+        return xr.concat(pieces, dim="site", join="outer")