HydroDataset Base Class

Overview

HydroDataset is the abstract base class that defines the unified interface for all hydrological dataset implementations in hydrodataset. It provides common functionality for data reading, caching, variable standardization, and unit management.

Key Responsibilities

Data Reading

• Read basin/station IDs
• Read timeseries data as xarray.Dataset
• Read attribute/static data as xarray.Dataset
• Support for selecting specific variables and time ranges

NetCDF Caching

• Generate and cache timeseries data in .nc format
• Generate and cache attribute data in .nc format
• Fast subsequent reads from cached files
• Cache location configured in hydro_setting.yml

Variable Standardization

• Map dataset-specific variable names to standard names
• Support multiple data sources for same variable
• Automatic unit conversion and tracking

Feature Management

• List available static features (attributes)
• List available dynamic features (timeseries)
• Clean feature names for consistency

Architecture

Required Properties (Subclass Implementation)

Subclasses must implement these properties:

• _attributes_cache_filename: NetCDF filename for cached attributes (e.g., "camels_us_attributes.nc")
• _timeseries_cache_filename: NetCDF filename for cached timeseries (e.g., "camels_us_timeseries.nc")
• default_t_range: Default time range as ["YYYY-MM-DD", "YYYY-MM-DD"]

Optional Properties (Subclass Override)

• _subclass_static_definitions: Dictionary mapping standard static variable names to dataset-specific names and units
• _dynamic_variable_mapping: Dictionary mapping StandardVariable constants to dataset-specific timeseries variables

Usage Pattern

Creating a New Dataset Class

from hydrodataset import HydroDataset, StandardVariable
from aqua_fetch import DatasetClass

class MyDataset(HydroDataset):
    def __init__(self, data_path, region=None, download=False):
        super().__init__(data_path)
        self.region = region
        self.download = download
        self.aqua_fetch = DatasetClass(data_path)

    @property
    def _attributes_cache_filename(self):
        return "mydataset_attributes.nc"

    @property
    def _timeseries_cache_filename(self):
        return "mydataset_timeseries.nc"

    @property
    def default_t_range(self):
        return ["1980-01-01", "2020-12-31"]

    # Define static variable mappings
    _subclass_static_definitions = {
        "area": {"specific_name": "area_km2", "unit": "km^2"},
        "p_mean": {"specific_name": "p_mean", "unit": "mm/day"},
    }

    # Define dynamic variable mappings
    _dynamic_variable_mapping = {
        StandardVariable.STREAMFLOW: {
            "default_source": "observations",
            "sources": {
                "observations": {"specific_name": "q_cms", "unit": "m^3/s"},
            },
        },
        StandardVariable.PRECIPITATION: {
            "default_source": "gauge",
            "sources": {
                "gauge": {"specific_name": "precip_mm", "unit": "mm/day"},
                "era5": {"specific_name": "tp", "unit": "mm/day"},
            },
        },
    }

Using a Dataset Instance

from hydrodataset import open_dataset

# Initialize via the resolver — source="local" or source="cloud"
ds = open_dataset("camels_us", source="local")

# Get available features
print(ds.available_static_features)
print(ds.available_dynamic_features)

# Read data
basin_ids = ds.read_object_ids()
timeseries = ds.read_ts_xrdataset(
    gage_id_lst=basin_ids[:5],
    t_range=["1990-01-01", "1995-12-31"],
    var_lst=["streamflow", "precipitation"]
)
attributes = ds.read_attr_xrdataset(
    gage_id_lst=basin_ids[:5],
    var_lst=["area", "p_mean"]
)

# Access convenience methods
areas = ds.read_area(gage_id_lst=basin_ids[:5])
mean_precip = ds.read_mean_prcp(gage_id_lst=basin_ids[:5])

API Reference

hydrodataset.hydro_dataset.HydroDataset

Bases: ABC

An interface for Hydrological Dataset.

Subclasses wrap an aqua_fetch reader (self.aqua_fetch) and standardize its data through the ADR 0001 resolver: resolve_data_path returns an absolute URI (local path or s3://) which is passed as uri. Data is cached locally as NetCDF (or as Zarr on cloud storage) and exposed through the standard read_attr_xrdataset / read_ts_xrdataset APIs using standardized variable names.

Parameters

uri : str Absolute path (local) or s3:// URI (cloud) of the dataset root. cache_path : str or Path, optional Cache directory for the NetCDF/Zarr caches. When omitted the default cache directory from ~/hydro_setting.yml is used.

Source code in hydrodataset/hydro_dataset.py

class HydroDataset(ABC):
    """An interface for Hydrological Dataset.

    Subclasses wrap an aqua_fetch reader (``self.aqua_fetch``) and standardize
    its data through the ADR 0001 resolver: ``resolve_data_path`` returns an
    absolute URI (local path or ``s3://``) which is passed as ``uri``. Data is
    cached locally as NetCDF (or as Zarr on cloud storage) and exposed through
    the standard ``read_attr_xrdataset`` / ``read_ts_xrdataset`` APIs using
    standardized variable names.

    Parameters
    ----------
    uri : str
        Absolute path (local) or ``s3://`` URI (cloud) of the dataset root.
    cache_path : str or Path, optional
        Cache directory for the NetCDF/Zarr caches. When omitted the default
        cache directory from ``~/hydro_setting.yml`` is used.
    """

    # A unified definition for static variables, including name mapping and units
    _base_static_definitions = {
        "area": {"specific_name": "area_km2", "unit": "km^2"},
        "p_mean": {"specific_name": "p_mean", "unit": "mm/day"},
    }
    # variable name map for timeseries
    _dynamic_variable_mapping = {}

    def __init__(self, uri, cache_path=None):
        from pathlib import PurePosixPath

        uri_str = str(uri)
        if uri_str.startswith("s3://"):
            # S3 URI — skip local path validation
            self.data_source_dir = uri_str
        else:
            data_path = Path(uri_str)
            # Accept both POSIX (/data/hydro) and Windows (C:\data\hydro) absolute forms
            if not (
                PurePosixPath(str(data_path)).is_absolute()
                or data_path.is_absolute()
            ):
                raise ValueError(
                    f"uri must be an absolute path. "
                    f"Use hydrodataset.configs.data_resolver.resolve_data_path() "
                    f"to resolve from a dataset id. "
                    f"Got: {data_path}"
                )
            self.data_source_dir = data_path
            if not self.data_source_dir.is_dir():
                self.data_source_dir.mkdir(parents=True)
        if cache_path is None:
            from hydrodataset.configs.settings import get_cache_dir

            self.cache_dir = get_cache_dir()
        else:
            self.cache_dir = Path(cache_path)
        if not self.cache_dir.is_dir():
            self.cache_dir.mkdir(parents=True)

        # Merge static variable definitions
        self._static_variable_definitions = self._base_static_definitions.copy()
        if hasattr(self.__class__, "_subclass_static_definitions"):
            self._static_variable_definitions.update(self._subclass_static_definitions)

    def get_name(self):
        raise NotImplementedError

    def set_data_source_describe(self):
        raise NotImplementedError

    def download_data_source(self):
        raise NotImplementedError

    def is_data_ready(self):
        raise NotImplementedError

    def read_object_ids(self) -> np.ndarray:
        """Read watershed station ID list."""
        if hasattr(self, "aqua_fetch"):
            stations_list = self.aqua_fetch.stations()
            return np.sort(np.array(stations_list))
        raise NotImplementedError

    def read_target_cols(
        self, gage_id_lst=None, t_range=None, target_cols=None, **kwargs
    ) -> np.ndarray:
        raise NotImplementedError

    def read_relevant_cols(
        self, gage_id_lst=None, t_range=None, var_lst=None, forcing_type=None, **kwargs
    ) -> np.ndarray:
        """3d data (site_num * time_length * var_num), time-series data"""
        raise NotImplementedError

    def read_constant_cols(
        self, gage_id_lst=None, var_lst=None, **kwargs
    ) -> np.ndarray:
        """2d data (site_num * var_num), non-time-series data"""
        raise NotImplementedError

    def read_other_cols(
        self, object_ids=None, other_cols: dict = None, **kwargs
    ) -> dict:
        """some data which cannot be easily treated as constant vars or time-series with same length as relevant vars
        CONVENTION: other_cols is a dict, where each item is also a dict with all params in it
        """
        raise NotImplementedError

    def get_constant_cols(self) -> np.ndarray:
        """the constant cols in this data_source"""
        raise NotImplementedError

    def get_relevant_cols(self) -> np.ndarray:
        """the relevant cols in this data_source"""
        raise NotImplementedError

    def get_target_cols(self) -> np.ndarray:
        """the target cols in this data_source"""
        raise NotImplementedError

    def get_other_cols(self) -> dict:
        """the other cols in this data_source"""
        raise NotImplementedError

    def _dynamic_features(self) -> list:
        """the dynamic features in this data_source"""
        if hasattr(self, "aqua_fetch"):
            original_features = self.aqua_fetch.dynamic_features
            return self._clean_feature_names(original_features)
        raise NotImplementedError

    @staticmethod
    def _clean_feature_names(feature_names):
        """Clean feature names to be compatible with NetCDF format and our internal standard.

        The cleaning process follows these steps:
        1. Remove units in parentheses (along with any preceding whitespace)
           e.g., 'Prcp(mm/day)' -> 'Prcp' or 'Temp (°C)' -> 'Temp'
        2. Convert all characters to lowercase
           e.g., 'Prcp' -> 'prcp'
        3. Remove any remaining invalid characters (only keep a-z, 0-9, and _)
           This ensures NetCDF variable naming compliance

        Args:
            feature_names (list or pd.Index): Original feature names that may contain
                units and special characters

        Returns:
            list: Cleaned feature names with only lowercase letters, numbers, and underscores

        Examples:
            >>> _clean_feature_names(['Prcp(mm/day)_daymet', 'Temp (°C)'])
            ['prcp_daymet', 'temp']
        """
        if not isinstance(feature_names, pd.Index):
            feature_names = pd.Index(feature_names)

        # Remove units in parentheses, then convert to lowercase
        cleaned_names = feature_names.str.replace(
            r"\s*\([^)]*\)", "", regex=True
        ).str.lower()
        # Replace any remaining invalid characters
        cleaned_names = cleaned_names.str.replace(r"""[^a-z0-9_]""", "", regex=True)
        return cleaned_names.tolist()

    def _static_features(self) -> list:
        """the static features in this data_source"""
        if hasattr(self, "aqua_fetch"):
            original_features = self.aqua_fetch.static_features
            return self._clean_feature_names(original_features)
        raise NotImplementedError

    @property
    @abstractmethod
    def _attributes_cache_filename(self):
        pass

    @property
    @abstractmethod
    def _timeseries_cache_filename(self):
        pass

    @property
    @abstractmethod
    def default_t_range(self):
        pass

    def cache_timeseries_xrdataset(self):
        if hasattr(self, "aqua_fetch"):
            # Build a lookup map from specific name to unit
            unit_lookup = {}
            if hasattr(self, "_dynamic_variable_mapping"):
                for (
                    std_name,
                    mapping_info,
                ) in self._dynamic_variable_mapping.items():
                    for source, source_info in mapping_info["sources"].items():
                        unit_lookup[source_info["specific_name"]] = source_info["unit"]

            gage_id_lst = self.read_object_ids().tolist()
            original_var_lst = self.aqua_fetch.dynamic_features
            cleaned_var_lst = self._clean_feature_names(original_var_lst)
            # Create a mapping from original variable names to cleaned names
            # to ensure correct correspondence even if list order changes
            var_name_mapping = dict(zip(original_var_lst, cleaned_var_lst))

            batch_data = self.aqua_fetch.fetch_stations_features(
                stations=gage_id_lst,
                dynamic_features=original_var_lst,
                static_features=None,
                st=self.default_t_range[0],
                en=self.default_t_range[1],
                as_dataframe=False,
            )

            dynamic_data = (
                batch_data[1] if isinstance(batch_data, tuple) else batch_data
            )

            new_data_vars = {}
            time_coord = dynamic_data.coords["time"]

            # Process only the variables that exist in the data source
            # Subclasses can add additional variables in their override methods
            for original_var in tqdm(
                original_var_lst,
                desc="Processing variables",
                total=len(original_var_lst),
            ):
                cleaned_var = var_name_mapping[original_var]
                var_data = []
                for station in gage_id_lst:
                    if station in dynamic_data.data_vars:
                        station_data = dynamic_data[station].sel(
                            dynamic_features=original_var
                        )
                        if "dynamic_features" in station_data.coords:
                            station_data = station_data.drop("dynamic_features")
                        var_data.append(station_data)

                if var_data:
                    combined = xr.concat(var_data, dim="basin")
                    combined["basin"] = gage_id_lst
                    combined.attrs["units"] = unit_lookup.get(cleaned_var, "unknown")
                    new_data_vars[cleaned_var] = combined

            new_ds = xr.Dataset(
                data_vars=new_data_vars,
                coords={
                    "basin": gage_id_lst,
                    "time": time_coord,
                },
            )

            batch_filepath = self.cache_dir.joinpath(self._timeseries_cache_filename)
            batch_filepath.parent.mkdir(parents=True, exist_ok=True)
            new_ds.to_netcdf(batch_filepath)
            print(f"成功保存到: {batch_filepath}")
        else:
            raise NotImplementedError

    def _assign_units_to_dataset(self, ds, units_map):
        def get_unit_by_prefix(var_name):
            for prefix, unit in units_map.items():
                if var_name.startswith(prefix):
                    return unit
            return None

        def get_unit(var_name):
            prefix_unit = get_unit_by_prefix(var_name)
            if prefix_unit:
                return prefix_unit
            return "undefined"

        for var in ds.data_vars:
            unit = get_unit(var)
            ds[var].attrs["units"] = unit
            if unit == "class":
                ds[var].attrs["description"] = "Classification code"
        return ds

        return ds

    def _get_attribute_units(self) -> dict:
        """Builds a unit dictionary from the static variable definitions."""
        return {
            info["specific_name"]: info["unit"]
            for std_name, info in self._static_variable_definitions.items()
        }

    def _write_zarr_units(self, root, kind: str) -> None:
        """Attach a ``units`` attr to each data array of a cloud zarr group.

        Mirrors the units the local NetCDF cache writes, so cloud zarr is
        self-describing and consistent with local. Keys are cleaned with
        ``_clean_feature_names`` to match the zarr array names exactly.

        Args:
            root: an open ``zarr`` group (the timeseries or attributes store).
            kind: ``"dynamic"`` for timeseries variables (unit source is
                ``_dynamic_variable_mapping``, default ``"unknown"``) or
                ``"static"`` for attributes (unit source is
                ``_static_variable_definitions``, default ``"undefined"``).

        Coordinate/bookkeeping arrays (``basin``/``time``/``_progress``) are
        skipped.
        """
        clean = self._clean_feature_names
        if kind == "dynamic":
            lookup = {}
            for info in getattr(self, "_dynamic_variable_mapping", {}).values():
                for src in info.get("sources", {}).values():
                    lookup[clean([src["specific_name"]])[0]] = src.get(
                        "unit", "unknown"
                    )
            default = "unknown"
        else:
            lookup = {
                clean([info["specific_name"]])[0]: info["unit"]
                for info in self._static_variable_definitions.values()
            }
            default = "undefined"
        for name in root.array_keys():
            if name in ("basin", "time", "_progress"):
                continue
            root[name].attrs["units"] = lookup.get(name, default)

    def cache_attributes_xrdataset(self):
        if hasattr(self, "aqua_fetch"):
            df_attr = self.aqua_fetch.fetch_static_features()
            print(df_attr.columns)
            # Clean column names using the unified method
            df_attr.columns = self._clean_feature_names(df_attr.columns)
            # Remove duplicate columns if any (keep first occurrence)
            if df_attr.columns.duplicated().any():
                df_attr = df_attr.loc[:, ~df_attr.columns.duplicated()]
            # Ensure index is string type for basin IDs
            df_attr.index = df_attr.index.astype(str)
            ds_attr = df_attr.to_xarray()
            # Check if the coordinate is named 'basin', if not rename it
            coord_names = list(ds_attr.dims.keys())
            if len(coord_names) > 0 and coord_names[0] != "basin":
                ds_attr = ds_attr.rename({coord_names[0]: "basin"})
            units_map = self._get_attribute_units()
            ds_attr = self._assign_units_to_dataset(ds_attr, units_map)
            ds_attr.to_netcdf(self.cache_dir.joinpath(self._attributes_cache_filename))
        else:
            raise NotImplementedError

    def read_attr_xrdataset(
        self,
        gage_id_lst: list = None,
        var_lst: list = None,
        to_numeric: bool = True,
        **kwargs,
    ) -> xr.Dataset:
        """Reads attribute data for a list of basins using standardized variable names.

        Args:
            gage_id_lst: A list of basin identifiers.
            var_lst: A list of **standard** attribute names to retrieve.
                If None, nothing will be returned.
            to_numeric: If True, converts all non-numeric variables to numeric codes
                and stores the original labels in the variable's attributes.
                Defaults to True.

        Returns:
            An xarray Dataset containing the attribute data for the requested basins,
            with variables named using the standard names.
        """
        if not var_lst:
            return None

        # 1. Translate standard names to dataset-specific names
        target_vars_to_fetch = []
        rename_map = {}
        for std_name in var_lst:
            if std_name not in self._static_variable_definitions:
                raise ValueError(
                    f"'{std_name}' is not a recognized standard static variable."
                )
            actual_var_name = self._static_variable_definitions[std_name][
                "specific_name"
            ]
            target_vars_to_fetch.append(actual_var_name)
            rename_map[actual_var_name] = std_name

        # 2. Read data from cache using actual variable names
        attr_ds = self._load_attr_dataset()

        # 3. Select variables and basins
        ds_subset = attr_ds[target_vars_to_fetch]
        if gage_id_lst is not None:
            gage_id_lst = [str(gid) for gid in gage_id_lst]

            # Intersect with available basins
            available_basins = set(ds_subset["basin"].values.astype(str))
            valid_gages = [g for g in gage_id_lst if str(g) in available_basins]
            missing_gages = set(gage_id_lst) - set(valid_gages)

            if missing_gages:
                print(
                    f"Warning: {len(missing_gages)} requested basins are missing from the attribute cache and will be skipped."
                )

            ds_selected = ds_subset.sel(basin=valid_gages)
        else:
            ds_selected = ds_subset

        # 4. Rename to standard names
        final_ds = ds_selected.rename(rename_map)

        if not to_numeric:
            return final_ds

        # 5. If to_numeric is True, perform conversion
        converted_ds = xr.Dataset(coords=final_ds.coords)
        for var_name, da in final_ds.data_vars.items():
            if np.issubdtype(da.dtype, np.number):
                converted_ds[var_name] = da
            else:
                # Assumes string-like array that needs factorizing
                numeric_vals, labels = pd.factorize(da.values, sort=True)
                new_da = xr.DataArray(
                    numeric_vals,
                    coords=da.coords,
                    dims=da.dims,
                    name=da.name,
                    attrs=da.attrs,  # Preserve original attributes
                )
                new_da.attrs["labels"] = labels.tolist()
                converted_ds[var_name] = new_da
        return converted_ds

    # ── Cloud utilities ───────────────────────────────────────────────────────

    def _is_cloud(self) -> bool:
        return str(self.data_source_dir).startswith("s3://")

    def _make_s3fs(self):
        """Create an s3fs filesystem from storage.s3 in hydro_setting.yml."""
        import s3fs
        from hydrodataset.configs.settings import get_s3_config
        cfg = get_s3_config()
        return s3fs.S3FileSystem(
            key=cfg.get("access_key_id"),
            secret=cfg.get("secret_access_key"),
            endpoint_url=cfg.get("endpoint_url"),
            config_kwargs={"s3": {"addressing_style": "virtual"}},
        )

    def _zarr_path_and_opts(self, zarr_name: str):
        """Return (s3_uri, storage_options) for a zarr store on OSS.

        zarr 3 manages its own async fs internally; passing storage_options
        avoids the event-loop mismatch caused by mixing a sync s3fs with
        zarr's internal async loop.
        """
        from hydrodataset.configs.settings import get_s3_config
        cfg = get_s3_config()
        bucket = cfg["bucket"]
        out = f"s3://{bucket}/zarr/{zarr_name}"
        opts = {
            "key": cfg.get("access_key_id"),
            "secret": cfg.get("secret_access_key"),
            "endpoint_url": cfg.get("endpoint_url"),
            "config_kwargs": {"s3": {"addressing_style": "virtual"}},
        }
        return out, opts

    def _open_zarr_pref_consolidated(self, out: str, opts: dict) -> xr.Dataset:
        """Open a cloud zarr, preferring consolidated metadata (one GET) and
        falling back to unconsolidated for stores written before B. Both paths
        keep mask_and_scale=False so real 0/1970 values survive.
        """
        try:
            return xr.open_zarr(
                out, storage_options=opts, consolidated=True, mask_and_scale=False
            )
        except Exception:
            return xr.open_zarr(
                out, storage_options=opts, consolidated=False, mask_and_scale=False
            )

    # ── Cache loading hooks ───────────────────────────────────────────────────

    def _load_attr_dataset(self) -> xr.Dataset:
        """Load the attributes dataset. Cloud → zarr on OSS; local → NC cache."""
        if self._is_cloud():
            import zarr as _zarr

            zarr_name = self._attributes_cache_filename.replace(".nc", ".zarr")
            out, opts = self._zarr_path_and_opts(zarr_name)
            try:
                return self._open_zarr_pref_consolidated(out, opts)
            except _zarr.errors.GroupNotFoundError:
                print(f"Attributes zarr not found at {out}, generating...")
                import s3fs as _s3fs
                _fs = self._make_s3fs()
                _raw = out.removeprefix("s3://")
                if _fs.exists(_raw):
                    _fs.rm(_raw, recursive=True)
                self.cache_attributes_to_zarr()
                return self._open_zarr_pref_consolidated(out, opts)
        attr_cache_file = self.cache_dir.joinpath(self._attributes_cache_filename)
        try:
            return xr.open_dataset(attr_cache_file)
        except FileNotFoundError:
            self.cache_attributes_xrdataset()
            return xr.open_dataset(attr_cache_file)

    def _load_ts_dataset(self, **kwargs):
        """Load the timeseries dataset. Cloud → zarr on OSS; local → NC cache."""
        if self._is_cloud():
            import zarr as _zarr

            zarr_name = self._timeseries_cache_filename.replace(".nc", ".zarr")
            out, opts = self._zarr_path_and_opts(zarr_name)
            try:
                return self._open_zarr_pref_consolidated(out, opts)
            except _zarr.errors.GroupNotFoundError:
                print(f"Timeseries zarr not found at {out}, generating...")
                import s3fs as _s3fs
                _fs = self._make_s3fs()
                _raw = out.removeprefix("s3://")
                if _fs.exists(_raw):
                    _fs.rm(_raw, recursive=True)
                self.cache_timeseries_to_zarr()
                return self._open_zarr_pref_consolidated(out, opts)
        ts_cache_file = self.cache_dir.joinpath(self._timeseries_cache_filename)
        if not os.path.isfile(ts_cache_file):
            self.cache_timeseries_xrdataset()
        return xr.open_dataset(ts_cache_file)

    def read_ts_xrdataset(
        self,
        gage_id_lst: list = None,
        t_range: list = None,
        var_lst: list = None,
        sources: dict = None,
        **kwargs,
    ):
        if (
            not hasattr(self, "_dynamic_variable_mapping")
            or not self._dynamic_variable_mapping
        ):
            raise NotImplementedError(
                "This dataset does not support the standardized variable mapping."
            )

        if var_lst is None:
            var_lst = list(self._dynamic_variable_mapping.keys())

        if t_range is None:
            t_range = self.default_t_range

        target_vars_to_fetch = []
        rename_map = {}

        for std_name in var_lst:
            if std_name not in self._dynamic_variable_mapping:
                raise ValueError(
                    f"'{std_name}' is not a recognized standard variable for this dataset."
                )

            mapping_info = self._dynamic_variable_mapping[std_name]

            # Determine which source(s) to use and if they were explicitly requested
            is_explicit_source = sources and std_name in sources
            sources_to_use = []
            if is_explicit_source:
                provided_sources = sources[std_name]
                if isinstance(provided_sources, list):
                    sources_to_use.extend(provided_sources)
                else:
                    sources_to_use.append(provided_sources)
            else:
                sources_to_use.append(mapping_info["default_source"])

            # A suffix is only needed if the user explicitly requested multiple sources
            needs_suffix = is_explicit_source and len(sources_to_use) > 1
            for source in sources_to_use:
                if source not in mapping_info["sources"]:
                    raise ValueError(
                        f"Source '{source}' is not available for variable '{std_name}'."
                    )

                actual_var_name = mapping_info["sources"][source]["specific_name"]
                target_vars_to_fetch.append(actual_var_name)
                output_name = f"{std_name}_{source}" if needs_suffix else std_name
                rename_map[actual_var_name] = output_name

        # Read data from cache using actual variable names
        ts = self._load_ts_dataset(**kwargs)
        missing_vars = [v for v in target_vars_to_fetch if v not in ts.data_vars]
        if missing_vars:
            # To provide a better error message, map back to standard names
            reverse_rename_map = {v: k for k, v in rename_map.items()}
            missing_std_vars = [reverse_rename_map.get(v, v) for v in missing_vars]
            raise ValueError(
                f"The following variables are missing from the cache file: {missing_std_vars}"
            )

        ds_subset = ts[target_vars_to_fetch]

        # Intersect requested gage_id_lst with available basins in the netcdf
        available_basins = set(ds_subset["basin"].values.astype(str))
        valid_gages = [g for g in gage_id_lst if str(g) in available_basins]
        missing_gages = set(gage_id_lst) - set(valid_gages)

        if missing_gages:
            print(
                f"Warning: {len(missing_gages)} requested basins are missing from the dataset cache and will be skipped."
            )
            # Optional: print first few missing IDs
            # print(f"Missing IDs: {list(missing_gages)[:10]}...")

        ds_selected = ds_subset.sel(
            basin=valid_gages, time=slice(t_range[0], t_range[1])
        )
        final_ds = ds_selected.rename(rename_map)
        return final_ds

    def get_available_dynamic_features(self) -> dict:
        """
        Returns a dictionary of available standard dynamic feature names
        and their possible sources.
        """
        if (
            not hasattr(self, "_dynamic_variable_mapping")
            or not self._dynamic_variable_mapping
        ):
            return {}

        feature_info = {}
        for std_name, mapping_info in self._dynamic_variable_mapping.items():
            feature_info[std_name] = {
                "default_source": mapping_info.get("default_source"),
                "available_sources": list(mapping_info.get("sources", {}).keys()),
            }
        return feature_info

    def get_available_static_features(self) -> list:
        """Returns a list of available standard static feature names."""
        return list(self._static_variable_definitions.keys())

    @property
    def available_static_features(self) -> list:
        """Returns a list of available static attribute names."""
        return self.get_available_static_features()

    @property
    def available_dynamic_features(self) -> dict:
        """Returns a dictionary of available dynamic feature names and their possible sources."""
        return self.get_available_dynamic_features()

    def read_area(self, gage_id_lst: list[str]) -> xr.Dataset:
        """Reads the catchment area for a list of basins.

        Args:
            gage_id_lst: A list of basin identifiers for which to retrieve the area.

        Returns:
            An xarray Dataset containing the area data for the requested basins.
        """
        data_ds = self.read_attr_xrdataset(gage_id_lst=gage_id_lst, var_lst=["area"])
        return data_ds

    def _p_mean_from_precip(self, basin_ids) -> list:
        """Per-basin mean of the precipitation timeseries (mm/day).

        For datasets whose raw attributes lack a single ``p_mean`` column, the
        local ``cache_attributes_xrdataset`` override computes it from the
        precipitation timeseries. The cloud ``cache_attributes_to_zarr`` path
        uses this helper to do the same so cloud matches the local NC cache.
        """
        basin_ids = [str(b) for b in basin_ids]
        prcp = self.read_ts_xrdataset(
            gage_id_lst=basin_ids,
            t_range=self.default_t_range,
            var_lst=["precipitation"],
        )["precipitation"]
        # scale the per-timestep mean to mm/day (no-op for daily data, x24 for
        # hourly) so p_mean is consistent regardless of timeseries resolution
        times = pd.DatetimeIndex(prcp["time"].values)
        if len(times) > 1:
            step_hours = (times[1] - times[0]).total_seconds() / 3600.0
            steps_per_day = 24.0 / step_hours if step_hours else 1.0
        else:
            steps_per_day = 1.0
        series = (prcp.mean(dim="time") * steps_per_day).to_series()
        return [float(series.get(b, float("nan"))) for b in basin_ids]

    def read_mean_prcp(self, gage_id_lst: list[str], unit: str = "mm/d") -> xr.Dataset:
        """Reads the mean daily precipitation for a list of basins, with unit conversion.

        Args:
            gage_id_lst: A list of basin identifiers.
            unit: The desired unit for the output precipitation. Defaults to "mm/d".
                Supported units: ['mm/d', 'mm/day', 'mm/h', 'mm/hour', 'mm/3h',
                'mm/3hour', 'mm/8d', 'mm/8day'].

        Returns:
            An xarray Dataset containing the mean precipitation data in the specified units.

        Raises:
            ValueError: If an unsupported unit is provided.
        """
        prcp_var_name = "p_mean"
        data_ds = self.read_attr_xrdataset(
            gage_id_lst=gage_id_lst, var_lst=[prcp_var_name]
        )
        # No conversion needed
        if unit in ["mm/d", "mm/day"]:
            return data_ds

        # Conversion needed, create a new dataset
        converted_ds = data_ds.copy()
        # After renaming, the variable in the dataset is now the standard name
        if unit in ["mm/h", "mm/hour"]:
            converted_ds[prcp_var_name] = data_ds[prcp_var_name] / 24
        elif unit in ["mm/3h", "mm/3hour"]:
            converted_ds[prcp_var_name] = data_ds[prcp_var_name] / 8
        elif unit in ["mm/8d", "mm/8day"]:
            converted_ds[prcp_var_name] = data_ds[prcp_var_name] * 8
        else:
            raise ValueError(
                "unit must be one of ['mm/d', 'mm/day', 'mm/h', 'mm/hour', 'mm/3h', 'mm/3hour', 'mm/8d', 'mm/8day']"
            )
        return converted_ds

available_static_features property

Returns a list of available static attribute names.

available_dynamic_features property

Returns a dictionary of available dynamic feature names and their possible sources.

default_t_range abstractmethod property

__init__(uri, cache_path=None)

Source code in hydrodataset/hydro_dataset.py

def __init__(self, uri, cache_path=None):
    from pathlib import PurePosixPath

    uri_str = str(uri)
    if uri_str.startswith("s3://"):
        # S3 URI — skip local path validation
        self.data_source_dir = uri_str
    else:
        data_path = Path(uri_str)
        # Accept both POSIX (/data/hydro) and Windows (C:\data\hydro) absolute forms
        if not (
            PurePosixPath(str(data_path)).is_absolute()
            or data_path.is_absolute()
        ):
            raise ValueError(
                f"uri must be an absolute path. "
                f"Use hydrodataset.configs.data_resolver.resolve_data_path() "
                f"to resolve from a dataset id. "
                f"Got: {data_path}"
            )
        self.data_source_dir = data_path
        if not self.data_source_dir.is_dir():
            self.data_source_dir.mkdir(parents=True)
    if cache_path is None:
        from hydrodataset.configs.settings import get_cache_dir

        self.cache_dir = get_cache_dir()
    else:
        self.cache_dir = Path(cache_path)
    if not self.cache_dir.is_dir():
        self.cache_dir.mkdir(parents=True)

    # Merge static variable definitions
    self._static_variable_definitions = self._base_static_definitions.copy()
    if hasattr(self.__class__, "_subclass_static_definitions"):
        self._static_variable_definitions.update(self._subclass_static_definitions)

read_object_ids()

Read watershed station ID list.

Source code in hydrodataset/hydro_dataset.py

def read_object_ids(self) -> np.ndarray:
    """Read watershed station ID list."""
    if hasattr(self, "aqua_fetch"):
        stations_list = self.aqua_fetch.stations()
        return np.sort(np.array(stations_list))
    raise NotImplementedError

read_ts_xrdataset(gage_id_lst=None, t_range=None, var_lst=None, sources=None, **kwargs)

Source code in hydrodataset/hydro_dataset.py

def read_ts_xrdataset(
    self,
    gage_id_lst: list = None,
    t_range: list = None,
    var_lst: list = None,
    sources: dict = None,
    **kwargs,
):
    if (
        not hasattr(self, "_dynamic_variable_mapping")
        or not self._dynamic_variable_mapping
    ):
        raise NotImplementedError(
            "This dataset does not support the standardized variable mapping."
        )

    if var_lst is None:
        var_lst = list(self._dynamic_variable_mapping.keys())

    if t_range is None:
        t_range = self.default_t_range

    target_vars_to_fetch = []
    rename_map = {}

    for std_name in var_lst:
        if std_name not in self._dynamic_variable_mapping:
            raise ValueError(
                f"'{std_name}' is not a recognized standard variable for this dataset."
            )

        mapping_info = self._dynamic_variable_mapping[std_name]

        # Determine which source(s) to use and if they were explicitly requested
        is_explicit_source = sources and std_name in sources
        sources_to_use = []
        if is_explicit_source:
            provided_sources = sources[std_name]
            if isinstance(provided_sources, list):
                sources_to_use.extend(provided_sources)
            else:
                sources_to_use.append(provided_sources)
        else:
            sources_to_use.append(mapping_info["default_source"])

        # A suffix is only needed if the user explicitly requested multiple sources
        needs_suffix = is_explicit_source and len(sources_to_use) > 1
        for source in sources_to_use:
            if source not in mapping_info["sources"]:
                raise ValueError(
                    f"Source '{source}' is not available for variable '{std_name}'."
                )

            actual_var_name = mapping_info["sources"][source]["specific_name"]
            target_vars_to_fetch.append(actual_var_name)
            output_name = f"{std_name}_{source}" if needs_suffix else std_name
            rename_map[actual_var_name] = output_name

    # Read data from cache using actual variable names
    ts = self._load_ts_dataset(**kwargs)
    missing_vars = [v for v in target_vars_to_fetch if v not in ts.data_vars]
    if missing_vars:
        # To provide a better error message, map back to standard names
        reverse_rename_map = {v: k for k, v in rename_map.items()}
        missing_std_vars = [reverse_rename_map.get(v, v) for v in missing_vars]
        raise ValueError(
            f"The following variables are missing from the cache file: {missing_std_vars}"
        )

    ds_subset = ts[target_vars_to_fetch]

    # Intersect requested gage_id_lst with available basins in the netcdf
    available_basins = set(ds_subset["basin"].values.astype(str))
    valid_gages = [g for g in gage_id_lst if str(g) in available_basins]
    missing_gages = set(gage_id_lst) - set(valid_gages)

    if missing_gages:
        print(
            f"Warning: {len(missing_gages)} requested basins are missing from the dataset cache and will be skipped."
        )
        # Optional: print first few missing IDs
        # print(f"Missing IDs: {list(missing_gages)[:10]}...")

    ds_selected = ds_subset.sel(
        basin=valid_gages, time=slice(t_range[0], t_range[1])
    )
    final_ds = ds_selected.rename(rename_map)
    return final_ds

read_attr_xrdataset(gage_id_lst=None, var_lst=None, to_numeric=True, **kwargs)

Reads attribute data for a list of basins using standardized variable names.

Parameters:

Name	Type	Description	Default
gage_id_lst	list	A list of basin identifiers.	None
var_lst	list	A list of standard attribute names to retrieve. If None, nothing will be returned.	None
to_numeric	bool	If True, converts all non-numeric variables to numeric codes and stores the original labels in the variable's attributes. Defaults to True.	True

Returns:

Type	Description
Dataset	An xarray Dataset containing the attribute data for the requested basins,
Dataset	with variables named using the standard names.

Source code in hydrodataset/hydro_dataset.py

def read_attr_xrdataset(
    self,
    gage_id_lst: list = None,
    var_lst: list = None,
    to_numeric: bool = True,
    **kwargs,
) -> xr.Dataset:
    """Reads attribute data for a list of basins using standardized variable names.

    Args:
        gage_id_lst: A list of basin identifiers.
        var_lst: A list of **standard** attribute names to retrieve.
            If None, nothing will be returned.
        to_numeric: If True, converts all non-numeric variables to numeric codes
            and stores the original labels in the variable's attributes.
            Defaults to True.

    Returns:
        An xarray Dataset containing the attribute data for the requested basins,
        with variables named using the standard names.
    """
    if not var_lst:
        return None

    # 1. Translate standard names to dataset-specific names
    target_vars_to_fetch = []
    rename_map = {}
    for std_name in var_lst:
        if std_name not in self._static_variable_definitions:
            raise ValueError(
                f"'{std_name}' is not a recognized standard static variable."
            )
        actual_var_name = self._static_variable_definitions[std_name][
            "specific_name"
        ]
        target_vars_to_fetch.append(actual_var_name)
        rename_map[actual_var_name] = std_name

    # 2. Read data from cache using actual variable names
    attr_ds = self._load_attr_dataset()

    # 3. Select variables and basins
    ds_subset = attr_ds[target_vars_to_fetch]
    if gage_id_lst is not None:
        gage_id_lst = [str(gid) for gid in gage_id_lst]

        # Intersect with available basins
        available_basins = set(ds_subset["basin"].values.astype(str))
        valid_gages = [g for g in gage_id_lst if str(g) in available_basins]
        missing_gages = set(gage_id_lst) - set(valid_gages)

        if missing_gages:
            print(
                f"Warning: {len(missing_gages)} requested basins are missing from the attribute cache and will be skipped."
            )

        ds_selected = ds_subset.sel(basin=valid_gages)
    else:
        ds_selected = ds_subset

    # 4. Rename to standard names
    final_ds = ds_selected.rename(rename_map)

    if not to_numeric:
        return final_ds

    # 5. If to_numeric is True, perform conversion
    converted_ds = xr.Dataset(coords=final_ds.coords)
    for var_name, da in final_ds.data_vars.items():
        if np.issubdtype(da.dtype, np.number):
            converted_ds[var_name] = da
        else:
            # Assumes string-like array that needs factorizing
            numeric_vals, labels = pd.factorize(da.values, sort=True)
            new_da = xr.DataArray(
                numeric_vals,
                coords=da.coords,
                dims=da.dims,
                name=da.name,
                attrs=da.attrs,  # Preserve original attributes
            )
            new_da.attrs["labels"] = labels.tolist()
            converted_ds[var_name] = new_da
    return converted_ds

cache_timeseries_xrdataset()

Source code in hydrodataset/hydro_dataset.py

def cache_timeseries_xrdataset(self):
    if hasattr(self, "aqua_fetch"):
        # Build a lookup map from specific name to unit
        unit_lookup = {}
        if hasattr(self, "_dynamic_variable_mapping"):
            for (
                std_name,
                mapping_info,
            ) in self._dynamic_variable_mapping.items():
                for source, source_info in mapping_info["sources"].items():
                    unit_lookup[source_info["specific_name"]] = source_info["unit"]

        gage_id_lst = self.read_object_ids().tolist()
        original_var_lst = self.aqua_fetch.dynamic_features
        cleaned_var_lst = self._clean_feature_names(original_var_lst)
        # Create a mapping from original variable names to cleaned names
        # to ensure correct correspondence even if list order changes
        var_name_mapping = dict(zip(original_var_lst, cleaned_var_lst))

        batch_data = self.aqua_fetch.fetch_stations_features(
            stations=gage_id_lst,
            dynamic_features=original_var_lst,
            static_features=None,
            st=self.default_t_range[0],
            en=self.default_t_range[1],
            as_dataframe=False,
        )

        dynamic_data = (
            batch_data[1] if isinstance(batch_data, tuple) else batch_data
        )

        new_data_vars = {}
        time_coord = dynamic_data.coords["time"]

        # Process only the variables that exist in the data source
        # Subclasses can add additional variables in their override methods
        for original_var in tqdm(
            original_var_lst,
            desc="Processing variables",
            total=len(original_var_lst),
        ):
            cleaned_var = var_name_mapping[original_var]
            var_data = []
            for station in gage_id_lst:
                if station in dynamic_data.data_vars:
                    station_data = dynamic_data[station].sel(
                        dynamic_features=original_var
                    )
                    if "dynamic_features" in station_data.coords:
                        station_data = station_data.drop("dynamic_features")
                    var_data.append(station_data)

            if var_data:
                combined = xr.concat(var_data, dim="basin")
                combined["basin"] = gage_id_lst
                combined.attrs["units"] = unit_lookup.get(cleaned_var, "unknown")
                new_data_vars[cleaned_var] = combined

        new_ds = xr.Dataset(
            data_vars=new_data_vars,
            coords={
                "basin": gage_id_lst,
                "time": time_coord,
            },
        )

        batch_filepath = self.cache_dir.joinpath(self._timeseries_cache_filename)
        batch_filepath.parent.mkdir(parents=True, exist_ok=True)
        new_ds.to_netcdf(batch_filepath)
        print(f"成功保存到: {batch_filepath}")
    else:
        raise NotImplementedError

cache_attributes_xrdataset()

Source code in hydrodataset/hydro_dataset.py

def cache_attributes_xrdataset(self):
    if hasattr(self, "aqua_fetch"):
        df_attr = self.aqua_fetch.fetch_static_features()
        print(df_attr.columns)
        # Clean column names using the unified method
        df_attr.columns = self._clean_feature_names(df_attr.columns)
        # Remove duplicate columns if any (keep first occurrence)
        if df_attr.columns.duplicated().any():
            df_attr = df_attr.loc[:, ~df_attr.columns.duplicated()]
        # Ensure index is string type for basin IDs
        df_attr.index = df_attr.index.astype(str)
        ds_attr = df_attr.to_xarray()
        # Check if the coordinate is named 'basin', if not rename it
        coord_names = list(ds_attr.dims.keys())
        if len(coord_names) > 0 and coord_names[0] != "basin":
            ds_attr = ds_attr.rename({coord_names[0]: "basin"})
        units_map = self._get_attribute_units()
        ds_attr = self._assign_units_to_dataset(ds_attr, units_map)
        ds_attr.to_netcdf(self.cache_dir.joinpath(self._attributes_cache_filename))
    else:
        raise NotImplementedError

read_area(gage_id_lst)

Reads the catchment area for a list of basins.

Parameters:

Name	Type	Description	Default
gage_id_lst	list[str]	A list of basin identifiers for which to retrieve the area.	required

Returns:

Type	Description
Dataset	An xarray Dataset containing the area data for the requested basins.

Source code in hydrodataset/hydro_dataset.py

def read_area(self, gage_id_lst: list[str]) -> xr.Dataset:
    """Reads the catchment area for a list of basins.

    Args:
        gage_id_lst: A list of basin identifiers for which to retrieve the area.

    Returns:
        An xarray Dataset containing the area data for the requested basins.
    """
    data_ds = self.read_attr_xrdataset(gage_id_lst=gage_id_lst, var_lst=["area"])
    return data_ds

read_mean_prcp(gage_id_lst, unit='mm/d')

Reads the mean daily precipitation for a list of basins, with unit conversion.

Parameters:

Name	Type	Description	Default
gage_id_lst	list[str]	A list of basin identifiers.	required
unit	str	The desired unit for the output precipitation. Defaults to "mm/d". Supported units: ['mm/d', 'mm/day', 'mm/h', 'mm/hour', 'mm/3h', 'mm/3hour', 'mm/8d', 'mm/8day'].	'mm/d'

Returns:

Type	Description
Dataset	An xarray Dataset containing the mean precipitation data in the specified units.

Raises:

Type	Description
ValueError	If an unsupported unit is provided.

Source code in hydrodataset/hydro_dataset.py

def read_mean_prcp(self, gage_id_lst: list[str], unit: str = "mm/d") -> xr.Dataset:
    """Reads the mean daily precipitation for a list of basins, with unit conversion.

    Args:
        gage_id_lst: A list of basin identifiers.
        unit: The desired unit for the output precipitation. Defaults to "mm/d".
            Supported units: ['mm/d', 'mm/day', 'mm/h', 'mm/hour', 'mm/3h',
            'mm/3hour', 'mm/8d', 'mm/8day'].

    Returns:
        An xarray Dataset containing the mean precipitation data in the specified units.

    Raises:
        ValueError: If an unsupported unit is provided.
    """
    prcp_var_name = "p_mean"
    data_ds = self.read_attr_xrdataset(
        gage_id_lst=gage_id_lst, var_lst=[prcp_var_name]
    )
    # No conversion needed
    if unit in ["mm/d", "mm/day"]:
        return data_ds

    # Conversion needed, create a new dataset
    converted_ds = data_ds.copy()
    # After renaming, the variable in the dataset is now the standard name
    if unit in ["mm/h", "mm/hour"]:
        converted_ds[prcp_var_name] = data_ds[prcp_var_name] / 24
    elif unit in ["mm/3h", "mm/3hour"]:
        converted_ds[prcp_var_name] = data_ds[prcp_var_name] / 8
    elif unit in ["mm/8d", "mm/8day"]:
        converted_ds[prcp_var_name] = data_ds[prcp_var_name] * 8
    else:
        raise ValueError(
            "unit must be one of ['mm/d', 'mm/day', 'mm/h', 'mm/hour', 'mm/3h', 'mm/3hour', 'mm/8d', 'mm/8day']"
        )
    return converted_ds