ClimateGraph 0.4.0__py3-none-any.whl

ClimateGraph/data/data.py

@@ -0,0 +1,508 @@
+"""The dataset abstraction.
+
+Brings datasets onto a common set of dimensions regardless of their format, file
+layout or shape, so the rest of the framework can work with them interchangeably.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+import cartopy.crs as ccrs
+import numpy as np
+import xarray as xr
+
+from ClimateGraph.reader import Reader
+from ClimateGraph.reader.reader.reader import ReadSpec
+from ClimateGraph.utils.dataset_utils import (
+    _record,
+    change_unit,
+    dim_reduction,
+    normalize_vars,
+)
+from ClimateGraph.utils.general_utils import ReductionMethodEnum
+from ClimateGraph.utils.registry import RegistryMixin
+from ClimateGraph.utils.resample_engine import ResampleEngine, get_engine
+
+log = logging.getLogger(__name__)
+
+
+class Data(RegistryMixin, ABC):
+    """The Data abstract class.
+
+    A class that abstracts the core nature of environmental data, independent of the topology of the data.
+    This class contains and implements attributes and methods common to all the Data that ClimateGraph is supposed to handle.
+    """
+
+    registry: dict[str, type["Data"]] = {}
+    aliases: list[str] = list()
+    geom_dims: tuple[str, ...] = ()
+
+    @classmethod
+    def get_data_subclass(cls, name: str):
+        return cls.get_class(name)
+
+    @classmethod
+    def check_topology_type(cls, type: str) -> bool:
+        return cls.check_class(type)
+
+    @classmethod
+    def create(
+        cls,
+        name: str,
+        topology: str,
+        reader: str,
+        path: Path | list[Path],
+        vars: dict[str, dict[str, str]] | list[str] | None,
+        crs: ccrs,
+        reader_kwargs: dict,
+    ):
+        """Build a Data object of the subclass matching ``topology``.
+
+        Parameters
+        ----------
+        name : str
+            Name of the object, used for reference inside ClimateGraph execution.
+        topology : str
+            Type of topology, used for lookup in the Data registry.
+        reader : str
+            Specific reader name used for reading the files that compose the data object. A reader must exist for the corresponding topology and reader name.
+        path : Path | List[Path]
+            Path or list of paths that compose the data object.
+        vars : Dict[str, Dict[str, str]]
+            Dictionary where keys are Variable names, and the value is another Dictionary with "name" (with the name of the variable in the files) and "unit" (with the "pint" unit name for this variable) keys.
+        crs : Cartopy.CRS
+            Cartopy Coordinate Reference System object
+        reader_kwargs : Dict
+            Other keyword arguments in Dict format passed directly to the Reader, allows for further reader specialization.
+
+        Returns
+        -------
+        Data
+            Data object created with the arguments in the adequate Data subclass
+        """
+        data_cls = cls.get_data_subclass(topology)
+        reader_cls = Reader.get_reader_subclass(topology, reader)
+        return data_cls(name, path, vars, reader_cls, crs, reader_kwargs=reader_kwargs)
+
+    def __init__(
+        self,
+        name: str,
+        path: Path | list[Path],
+        vars: dict[str, dict[str, str]] | list[str] | None,
+        reader: Reader,
+        crs: ccrs.CRS,
+        reader_kwargs: dict[str, Any] | None = None,
+    ):
+        """Create a dataset handle. Nothing is read until ``obj`` is accessed.
+
+        Parameters
+        ----------
+        name : str
+            Name for reference inside the ClimateGraph execution.
+        path : Path | list[Path]
+            Path or list of paths that compose the data objects
+        vars : Dict[str, Dict[str, str]]
+            Dictionary where keys are Variable names, and the value is another Dictionary with "name" (with the name of the variable in the files) and "unit" (with the "pint" unit name for this variable) keys.
+        reader : Reader
+            Specific reader object used for reading the files that compose the data object. A reader must exist for the corresponding topology and reader name.
+        crs : ccrs.CRS
+            Cartopy Coordinate Reference System object
+        reader_kwargs : Dict[str, Any] | None, optional
+            Other keyword arguments in Dict format passed directly to the Reader, allows for further reader specialization, by default None
+        """
+        self.reader = reader
+        self.reader_kwargs = {} if reader_kwargs is None else reader_kwargs
+
+        self._obj = None
+        self._geom = None
+        self._resample_cache: dict = {}
+        self._vars = normalize_vars(vars)
+        self._path = None
+        self._bbox = None
+        self._dims = None
+
+        self.name = name
+        self.path = path
+        self.crs = crs
+
+    def copy(self):
+        """Return a copy of this object. Used when a domain produces a filtered dataset.
+
+        Returns
+        -------
+        Data
+            Exact replica of the Data object
+        """
+        new = self.__class__(
+            self.name,
+            self.path,
+            self.vars,
+            self.reader,
+            self.crs,
+            reader_kwargs=self.reader_kwargs,
+        )
+        new.obj = self._obj
+        new._geom = self._geom
+        new._bbox = self._bbox
+        new._dims = self._dims
+
+        return new
+
+    @abstractmethod
+    def _set_geom(self):
+        """Build the pyresample geometry used for resampling.
+
+        The core of the topology abstraction: each topology defines its own
+        spatial geometry here.
+        """
+        pass
+
+    @property
+    def obj(self) -> xr.Dataset:
+        """The underlying dataset, loaded on first access.
+
+        Returns
+        -------
+        xr.Dataset
+            Xarray dataset of the Data object.
+        """
+        if self._obj is None:
+            self.load_obj()
+        return self._obj
+
+    @obj.setter
+    def obj(self, _obj: xr.Dataset):
+        """Replace the underlying dataset, clearing everything derived from the old one.
+
+        Parameters
+        ----------
+        _obj : xr.Dataset
+             Xarray Dataset object meant to be set in the Data object.
+        """
+        self._obj = _obj
+        self._bbox = None
+        self._geom = None
+        self._dims = None
+        self._resample_cache = {}
+
+    @property
+    def geom(self):
+        if self._geom is None:
+            self._set_geom()
+        return self._geom
+
+    @property
+    def vars(self):
+        """The variable names this dataset exposes.
+
+        Returns
+        -------
+        Dict[str, Dict[str, str]] | None
+            Normalized dict mapping each user-facing variable name to a dict with
+            "name" (file-native name) and "unit" (pint unit, possibly None)
+            keys, or None when the data block declared no vars, and vars will be treated
+            by native file names.
+        """
+        return self._vars
+
+    def var_unit(self, var_name: str) -> str | None:
+        """Resolve the declared source unit for ``var_name``.
+
+        Single choke point for reading a variable's declared unit, so callers
+        never index self.vars[...]["unit"]
+
+        Parameters
+        ----------
+        var_name : str
+            User-facing variable name.
+
+        Returns
+        -------
+        str | None
+            The declared pint unit, or ``None`` when vars wasn't declared.
+        """
+        if isinstance(self._vars, dict):
+            entry = self._vars.get(var_name)
+            if entry is not None:
+                return entry.get("unit")
+        return None
+
+    @property
+    def dims(self):
+        """Mapping of dimension name to dimension size.
+
+        Returns
+        -------
+        Dict[str, int]
+            Dictionary mapping dimension names to dimension size.
+        """
+        if self._dims is None:
+            self._dims = dict(self.obj.sizes)
+        return self._dims
+
+    @property
+    def bbox(self):
+        """Bounding box around the spatial extent of the dataset.
+
+        Returns
+        -------
+        Tuple[float, float, float, float]
+            Tuple representing (Minimum Longitude,  Minimum Latitude, Maximum Longitude, Maximum Latitude)
+        """
+        if self._bbox is None:
+            lons, lats = self.get_coordinates(["longitude", "latitude"])
+            self._bbox = (
+                float(lons.min()),
+                float(lats.min()),
+                float(lons.max()),
+                float(lats.max()),
+            )
+        return self._bbox
+
+    def load_obj(self):
+        """Read the dataset into ``obj`` by building
+        a ReadSpec and invoking ``reader.read(spec)``. ``load_mode``,
+        ``cache_dir`` and ``engine`` are pulled out of ``reader_kwargs``
+        if present; the remainder lives on ``spec.extras`` for the
+        subclass to consume.
+
+        Returns
+        -------
+        xr.Dataset
+            Xarray dataset of the Data object.
+        """
+        extras = dict(self.reader_kwargs)
+        load_mode = extras.pop("load_mode", "safe")
+        cache_dir = extras.pop("cache_dir", None)
+        engine = extras.pop("engine", None)
+        spec = ReadSpec(
+            paths=self.path if isinstance(self.path, list) else [self.path],
+            vars=self.vars,
+            load_mode=load_mode,
+            cache_dir=cache_dir,
+            engine=engine,
+            extras=extras,
+        )
+        self._obj = self.reader.read(spec)
+        return self._obj
+
+    def get_var(
+        self,
+        var_name: str,
+        in_unit: str | None = None,
+        reduction_func: str | None = None,
+        keep_dims: str | list[str] | None = None,
+        reduction_dims: str | list[str] | None = None,
+        as_array: bool = False,
+        dim_reduce: dict[str, str | dict] | None = None,
+    ) -> xr.DataArray | np.ndarray:
+        """Get a variable, reduced and unit-converted, ready to draw.
+
+        Parameters
+        ----------
+        var_name : str
+            Variable name.
+        in_unit : str | None, optional
+            Measure of unit in which to convert the variable, by default None
+        reduction_func : str | None, optional
+            name of function with which to reduce the variable dimensions, dimensions in reduction_dims or, that aren't in keep_dims arg. Works only if keep_dims or reduction_dims is also specified. By default None
+        keep_dims : str | list[str] | None, optional
+            Dimensions to keep after reducing the others with the reduction_func. Only works if reduction_func is not None. By default None
+        reduction_dims : str | list[str] | None, optional
+            Dimensions to reduce with the reduction_func. Only works if reduction_func is not None. By default None
+        as_array : bool, optional
+            Boolean representing if the resulting data will be returned in numpy.ndarray format, by default False
+
+        Returns
+        -------
+        xr.DataArray | numpy.ndarray
+            Variable obtained from the obj attribute. Processed if required.
+
+        Raises
+        ------
+        KeyError
+            If the variable can't be found in the obj object.
+        """
+        if var_name not in self.obj:
+            raise KeyError(
+                f"Variable {var_name!r} not found in dataset {self.name!r} "
+                f"(available: {list(self.obj.data_vars)})."
+            )
+        xa = self.obj.data_vars[var_name]
+
+        if dim_reduce is not None:
+            xa = dim_reduction(xa, dim_reduce, name=var_name)
+
+        if reduction_func is not None:
+            reduction_func = (
+                ReductionMethodEnum(reduction_func).func
+                if isinstance(reduction_func, str)
+                else reduction_func
+            )
+            if reduction_dims is None and keep_dims is not None:
+                reduction_dims = list(set(self.dims) - set(keep_dims))
+            xa = xa.reduce(reduction_func, reduction_dims)
+            _record(
+                xa,
+                f"reduced {var_name!r} over {reduction_dims} with {reduction_func.__name__}",
+            )
+
+        if in_unit is not None:
+            xa = change_unit(xa, self.var_unit(var_name), in_unit)
+
+        if as_array:
+            xa = xa.to_numpy()
+
+        return xa
+
+    def get_coordinates(
+        self,
+        coord_names: list[str] | str,
+        as_array: bool = False,
+    ) -> list[xr.DataArray] | xr.DataArray | list[np.ndarray] | np.ndarray:
+        """Get one or more coordinates from the dataset.
+
+        Parameters
+        ----------
+        coord_names : list[str] | str
+            Coordinate name or names.
+        as_array : bool, optional
+            Boolean representing if the resulting data will be returned in numpy.ndarray format, by default False
+
+        Returns
+        -------
+        list[xr.DataArray] | xr.DataArray | list[np.ndarray] | np.ndarray
+            Coordinate(s) obtained from the obj attribute. Processed if required.
+
+        Raises
+        ------
+        KeyError
+            If a coordinate can't be found in the obj object.
+        """
+        obj = self.obj
+        if isinstance(coord_names, str):
+            coord_names = [coord_names]
+        missing_coords = [i for i in coord_names if i not in obj.coords]
+        if len(missing_coords) > 0:
+            raise KeyError(f"Coordinates {missing_coords} not in {self.name} dataset.")
+
+        coords = [obj.coords[coord] for coord in coord_names]
+        if as_array:
+            coords = [array.to_numpy() for array in coords]
+
+        if len(coords) == 1:
+            coords = coords[0]
+        return coords
+
+    def resample_vars(
+        self,
+        other: "Data",
+        vars: str | list[str],
+        radius_of_influence: int = 10000,
+        engine: str | ResampleEngine = "pyresample",
+        engine_kwargs: dict | None = None,
+    ) -> xr.Dataset:
+        """Project ``other``'s vars onto ``self``'s spatial geometry.
+
+        Purely spatial: each variable of ``other`` is reprojected onto ``self``'s
+        geometry with the requested engine, keeping ``other``'s own time axis and any
+        extra dims (z, member, …) untouched, and inheriting ``self``'s spatial
+        coordinates (site / latitude / longitude / region / …). It does NOT align time
+        or convert units — those are the caller's concern. Two sources with different
+        time extents resampled onto the same target never clobber each other.
+
+        The result is memoized on the source (``other._resample_cache``), keyed by
+        ``(self.name, vars, radius, engine, engine_kwargs)``, so repeated
+        ``resample_to`` domains sharing a target recompute the projection once and hit
+        the cache thereafter. The cache is invalidated whenever ``other.obj`` is
+        reassigned; the target's geometry is assumed stable across a run.
+
+        Parameters
+        ----------
+        other : Data
+            Data whose vars are reprojected onto ``self``'s geometry.
+        vars : str | list[str]
+            Variable name or list of names to resample.
+        radius_of_influence : int, optional
+            Radius length in meters to use for resampling. by default 10000
+        engine : str | ResampleEngine, optional
+            Resample backend name or instance. Default ``"pyresample"``.
+        engine_kwargs : dict | None, optional
+            Extra kwargs forwarded to the engine constructor
+            (e.g. ``method``, ``sigmas``). Default None.
+
+        Returns
+        -------
+        xr.Dataset
+            One variable per input var, keyed ``"{var}__{other.name}"``, on
+            ``self``'s geometry and carrying ``other``'s time / extra dims.
+        """
+        if isinstance(vars, str):
+            vars = [vars]
+
+        engine_key = engine if isinstance(engine, str) else type(engine).__name__
+        cache_key = (
+            self.name,
+            tuple(sorted(vars)),
+            radius_of_influence,
+            engine_key,
+            tuple(sorted((engine_kwargs or {}).items())),
+        )
+        if cache_key in other._resample_cache:
+            log.debug("resample cache hit: %s -> %s (%s)", other.name, self.name, vars)
+            return other._resample_cache[cache_key]
+
+        resample_engine = get_engine(engine, **(engine_kwargs or {}))
+        info = resample_engine.prepare(
+            other.geom,
+            self.geom,
+            radius_of_influence=radius_of_influence,
+        )
+        _resample = resample_engine.make_resampler(
+            info, self.geom.shape, other.geom.shape
+        )
+
+        dst_geom_dims = [d for d in self.obj.dims if d in set(self.geom_dims)]
+        dst_sizes = {d: self.obj.sizes[d] for d in dst_geom_dims}
+        dst_coords = {
+            name: coord
+            for name, coord in self.obj.coords.items()
+            if set(coord.dims) <= set(self.geom_dims)
+        }
+
+        resampled_vars = {}
+        for var in vars:
+            var_src = other.get_var(var)
+            src_geom_dims = [d for d in var_src.dims if d in set(other.geom_dims)]
+
+            resampled = xr.apply_ufunc(
+                _resample,
+                var_src,
+                input_core_dims=[src_geom_dims],
+                output_core_dims=[dst_geom_dims],
+                dask="parallelized",
+                output_dtypes=[var_src.dtype],
+                dask_gufunc_kwargs={"output_sizes": dst_sizes},
+            )
+            resampled_vars[f"{var}__{other.name}"] = resampled.assign_coords(dst_coords)
+
+        result = xr.Dataset(resampled_vars)
+        # load so anything downstream doesn't re compute all processed until here
+        result = result.load()
+        _record(
+            result,
+            f"spatially resampled {list(resampled_vars)} from {other.name} "
+            f"({type(other).__name__}) to {self.name} ({type(self).__name__})",
+        )
+
+        save_to = self.reader_kwargs.get("save_resampled_to")
+        if save_to:
+            target = Path(save_to)
+            target.parent.mkdir(parents=True, exist_ok=True)
+            result.to_netcdf(target)
+            _record(result, f"saved resampled data to {target}")
+
+        other._resample_cache[cache_key] = result
+        return result

ClimateGraph/data/point_surface.py

@@ -0,0 +1,91 @@
+"""Station and point-observation topology."""
+
+from collections.abc import Sequence
+from pathlib import Path
+
+import cartopy.crs as ccrs
+import numpy as np
+import xarray as xr
+from pyresample.geometry import SwathDefinition
+
+from .data import Data
+
+
+class _PrebuiltReader:
+    """No-op reader for in-memory PointSurface objects whose ``obj`` is set directly.
+
+    ``from_points`` builds the dataset itself, so ``read`` is never reached; it
+    only exists to satisfy ``Data.__init__``/``copy`` which carry a reader around.
+    """
+
+    @staticmethod
+    def read(spec):
+        raise RuntimeError(
+            "_PrebuiltReader has no files to read; its obj is set in-memory."
+        )
+
+
+class PointSurface(Data):
+    """The PointSurface topology class.
+
+    A class that gives particular representation to Point Surface data. This data should be in ("time", "site") dimensions.
+    """
+
+    aliases = ["pt_sfc", "point_surface", "point"]
+    geom_dims = ("site",)
+
+    def _set_geom(self):
+        """Build the pyresample ``SwathDefinition`` over the one-dimensional site coordinates."""
+        lons, lats = self.get_coordinates(["longitude", "latitude"], as_array=True)
+        self._geom = SwathDefinition(lons=lons, lats=lats, crs=self.crs)
+
+    @classmethod
+    def from_points(
+        cls,
+        name: str,
+        names: Sequence[str],
+        lats: Sequence[float],
+        lons: Sequence[float],
+        crs: ccrs.CRS | None = None,
+    ) -> "PointSurface":
+        """Build an in-memory point target from named coordinates.
+
+        Used as the resample target of a ``Points`` domain: a coords-only
+        PointSurface whose ``site`` dimension holds the provided **names**, with
+        ``latitude``/``longitude`` on ``site`` (so ``_set_geom`` can build the
+        SwathDefinition). No data variables — the resample only needs the geometry
+        and the site-scoped coords. ``obj`` is set directly, so no reader runs.
+
+        Parameters
+        ----------
+        name : str
+            Identity for the target (used in the resample cache key / history).
+        names : Sequence[str]
+            Point names — become the ``site`` index (labels and ``.sel`` keys).
+        lats, lons : Sequence[float]
+            Latitude / longitude per point, in the ``crs`` (default lat/lon).
+        crs : ccrs.CRS | None, optional
+            Coordinate reference system. Defaults to ``PlateCarree`` (lat/lon).
+
+        Returns
+        -------
+        PointSurface
+            An in-memory point target ready to resample onto.
+        """
+        site = np.asarray(list(names), dtype=object)
+        ds = xr.Dataset(
+            coords={
+                "site": site,
+                "latitude": ("site", np.asarray(lats, dtype=float)),
+                "longitude": ("site", np.asarray(lons, dtype=float)),
+            }
+        )
+        instance = cls(
+            name=name,
+            path=Path(f"memory://{name}"),
+            vars=None,
+            reader=_PrebuiltReader,
+            crs=crs or ccrs.PlateCarree(),
+        )
+        instance.obj = ds
+        return instance

ClimateGraph/data/regular_grid.py

@@ -0,0 +1,25 @@
+"""Structured latitude/longitude grid topology."""
+
+from pyresample import SwathDefinition
+
+from .data import Data
+
+
+class RegularGrid(Data):
+    """The RegularGrid topology class.
+
+    A class that implements Regular Grid specific logic. This data should be in ("time", "x", "y", "z") dimensions.
+    """
+
+    aliases = ["regular_grid", "grid", "regulargrid"]
+    geom_dims = ("x", "y")
+
+    def _set_geom(self):
+        """Build the pyresample geometry.
+
+        Uses a ``SwathDefinition`` rather than an ``AreaDefinition``, which does
+        not behave correctly here.
+        """
+        # TODO: change to AreaDefinition
+        lons, lats = self.get_coordinates(["longitude", "latitude"], as_array=True)
+        self._geom = SwathDefinition(lons=lons, lats=lats, crs=self.crs)

ClimateGraph/data/satellite_swath.py

@@ -0,0 +1,21 @@
+"""Satellite swath topology.
+
+not usable yet
+"""
+
+from pyresample.geometry import SwathDefinition
+
+from .data import Data
+
+
+class SatelliteSwath(Data):
+    """The Satellite Swath topology class. Not being used and not recommended.
+
+    A class that gives particular representation to Satellite Swath data.
+    """
+
+    def _set_geom(self):
+        lons, lats = self.get_coordinates(["longitude", "latitude"])
+        swath_def = SwathDefinition(lons=lons, lats=lats, crs=self.crs)
+        self.geom = swath_def.compute_optimal_bb_area(self.crs.proj4_params)
+        return self.geom

ClimateGraph/domain/__init__.py

@@ -0,0 +1,11 @@
+"""Spatial and attribute subsetting applied to datasets."""
+
+import importlib
+import pkgutil
+
+for module in pkgutil.iter_modules(__path__):
+    importlib.import_module(f"{__name__}.{module.name}")
+
+from .domain import Domain
+
+__all__ = ["Domain"]