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.camels_us import CamelsUs
from hydrodataset import SETTING

# Initialize
ds = CamelsUs(SETTING["local_data_path"]["datasets-origin"])

# 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

For unit, we use Pint package's unit system -- unit registry

Parameters

ABC : type description

Source code in hydrodataset/hydro_dataset.py

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

    For unit, we use Pint package's unit system -- unit registry

    Parameters
    ----------
    ABC : _type_
        _description_
    """

    # 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, data_path, cache_path=None):
        self.data_source_dir = Path(ROOT_DIR, data_path)
        if not self.data_source_dir.is_dir():
            self.data_source_dir.mkdir(parents=True)
        if cache_path is None:
            self.cache_dir = Path(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 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_cache_file = self.cache_dir.joinpath(self._attributes_cache_filename)
        try:
            attr_ds = xr.open_dataset(attr_cache_file)
        except FileNotFoundError:
            self.cache_attributes_xrdataset()
            attr_ds = xr.open_dataset(attr_cache_file)

        # 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]
            ds_selected = ds_subset.sel(basin=gage_id_lst)
        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

    def _load_ts_dataset(self, **kwargs):
        """
        Loads the time series dataset from cache.

        This method can be overridden by subclasses to implement different loading
        strategies (e.g., loading multiple files).

        Args:
            **kwargs: Additional keyword arguments for loading.

        Returns:
            xarray.Dataset: The loaded time series dataset.
        """
        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]
        ds_selected = ds_subset.sel(
            basin=gage_id_lst, 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 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__(data_path, cache_path=None)

Source code in hydrodataset/hydro_dataset.py

def __init__(self, data_path, cache_path=None):
    self.data_source_dir = Path(ROOT_DIR, data_path)
    if not self.data_source_dir.is_dir():
        self.data_source_dir.mkdir(parents=True)
    if cache_path is None:
        self.cache_dir = Path(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]
    ds_selected = ds_subset.sel(
        basin=gage_id_lst, 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_cache_file = self.cache_dir.joinpath(self._attributes_cache_filename)
    try:
        attr_ds = xr.open_dataset(attr_cache_file)
    except FileNotFoundError:
        self.cache_attributes_xrdataset()
        attr_ds = xr.open_dataset(attr_cache_file)

    # 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]
        ds_selected = ds_subset.sel(basin=gage_id_lst)
    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