code: from pathlib import Path
  from typing import Dict, List, Tuple, Union
  
  import numpy as np
  import pandas as pd
  import xarray
  
  from neuralhydrology.datasetzoo.basedataset import BaseDataset
  from neuralhydrology.utils.config import Config
  
  
  class CamelsUS(BaseDataset):
      """Data set class for the CAMELS US data set by [#]_ and [#]_.
      
      Parameters
      ----------
      cfg : Config
          The run configuration.
      is_train : bool 
          Defines if the dataset is used for training or evaluating. If True (training), means/stds for each feature
          are computed and stored to the run directory. If one-hot encoding is used, the mapping for the one-hot encoding 
          is created and also stored to disk. If False, a `scaler` input is expected and similarly the `id_to_int` input
          if one-hot encoding is used. 
      period : {'train', 'validation', 'test'}
          Defines the period for which the data will be loaded
      basin : str, optional
          If passed, the data for only this basin will be loaded. Otherwise the basin(s) are read from the appropriate
          basin file, corresponding to the `period`.
      additional_features : List[Dict[str, pd.DataFrame]], optional
          List of dictionaries, mapping from a basin id to a pandas DataFrame. This DataFrame will be added to the data
          loaded from the dataset and all columns are available as 'dynamic_inputs', 'static_inputs' and 
          'target_variables'
      id_to_int : Dict[str, int], optional
          If the config argument 'use_basin_id_encoding' is True in the config and period is either 'validation' or 
          'test', this input is required. It is a dictionary, mapping from basin id to an integer (the one-hot encoding).
      scaler : Dict[str, Union[pd.Series, xarray.DataArray]], optional
          If period is either 'validation' or 'test', this input is required. It contains the means and standard 
          deviations for each feature and is stored to the run directory during training (train_data/train_data_scaler.p)
          
      References
      ----------
      .. [#] A. J. Newman, M. P. Clark, K. Sampson, A. Wood, L. E. Hay, A. Bock, R. J. Viger, D. Blodgett, 
          L. Brekke, J. R. Arnold, T. Hopson, and Q. Duan: Development of a large-sample watershed-scale 
          hydrometeorological dataset for the contiguous USA: dataset characteristics and assessment of regional 
          variability in hydrologic model performance. Hydrol. Earth Syst. Sci., 19, 209-223, 
          doi:10.5194/hess-19-209-2015, 2015
      .. [#] Addor, N., Newman, A. J., Mizukami, N. and Clark, M. P.: The CAMELS data set: catchment attributes and 
          meteorology for large-sample studies, Hydrol. Earth Syst. Sci., 21, 5293-5313, doi:10.5194/hess-21-5293-2017,
          2017.
      """
  
      def __init__(self,
                   cfg: Config,
                   is_train: bool,
                   period: str,
                   basin: str = None,
                   additional_features: List[Dict[str, pd.DataFrame]] = [],
                   id_to_int: Dict[str, int] = {},
                   scaler: Dict[str, Union[pd.Series, xarray.DataArray]] = {}):
          super(CamelsUS, self).__init__(cfg=cfg,
                                         is_train=is_train,
                                         period=period,
                                         basin=basin,
                                         additional_features=additional_features,
                                         id_to_int=id_to_int,
                                         scaler=scaler)
  
      def _load_basin_data(self, basin: str) -> pd.DataFrame:
          """Load input and output data from text files."""
          # get forcings
          dfs = []
          for forcing in self.cfg.forcings:
              df, area = load_camels_us_forcings(self.cfg.data_dir, basin, forcing)
  
              # rename columns
              if len(self.cfg.forcings) > 1:
                  df = df.rename(columns={col: f"{col}_{forcing}" for col in df.columns})
              dfs.append(df)
          df = pd.concat(dfs, axis=1)
  
          # add discharge
          df['QObs(mm/d)'] = load_camels_us_discharge(self.cfg.data_dir, basin, area)
  
          # replace invalid discharge values by NaNs
          qobs_cols = [col for col in df.columns if "qobs" in col.lower()]
          for col in qobs_cols:
              df.loc[df[col] < 0, col] = np.nan
  
          return df
  
      def _load_attributes(self) -> pd.DataFrame:
          if self.cfg.camels_attributes:
  
              df = load_camels_us_attributes(self.cfg.data_dir, basins=self.basins)
  
              # remove all attributes not defined in the config
              drop_cols = [c for c in df.columns if c not in self.cfg.camels_attributes]
              df = df.drop(drop_cols, axis=1)
  
              return df
  
  
  def load_camels_us_attributes(data_dir: Path, basins: List[str] = []) -> pd.DataFrame:
      """Load CAMELS US attributes from the dataset provided by [#]_
  
      Parameters
      ----------
      data_dir : Path
          Path to the CAMELS US directory. This folder must contain a 'camels_attributes_v2.0' folder (the original 
          data set) containing the corresponding txt files for each attribute group.
      basins : List[str], optional
          If passed, return only attributes for the basins specified in this list. Otherwise, the attributes of all basins
          are returned.
  
      Returns
      -------
      pandas.DataFrame
          Basin-indexed DataFrame, containing the attributes as columns.
  
      References
      ----------
      .. [#] Addor, N., Newman, A. J., Mizukami, N. and Clark, M. P.: The CAMELS data set: catchment attributes and 
          meteorology for large-sample studies, Hydrol. Earth Syst. Sci., 21, 5293-5313, doi:10.5194/hess-21-5293-2017,
          2017.
      """
      attributes_path = Path(data_dir) / 'camels_attributes_v2.0'
  
      if not attributes_path.exists():
          raise RuntimeError(f"Attribute folder not found at {attributes_path}")
  
      txt_files = attributes_path.glob('camels_*.txt')
  
      # Read-in attributes into one big dataframe
      dfs = []
      for txt_file in txt_files:
          df_temp = pd.read_csv(txt_file, sep=';', header=0, dtype={'gauge_id': str})
          df_temp = df_temp.set_index('gauge_id')
  
          dfs.append(df_temp)
  
      df = pd.concat(dfs, axis=1)
      # convert huc column to double digit strings
      df['huc'] = df['huc_02'].apply(lambda x: str(x).zfill(2))
      df = df.drop('huc_02', axis=1)
  
      if basins:
          # drop rows of basins not contained in the passed list
          drop_basins = [b for b in df.index if b not in basins]
          df = df.drop(drop_basins, axis=0)
  
      return df
  
  
  def load_camels_us_forcings(data_dir: Path, basin: str, forcings: str) -> Tuple[pd.DataFrame, int]:
      """Load the forcing data for a basin of the CAMELS US data set.
  
      Parameters
      ----------
      data_dir : Path
          Path to the CAMELS US directory. This folder must contain a 'basin_mean_forcing' folder containing one 
          subdirectory for each forcing. The forcing directories have to contain 18 subdirectories (for the 18 HUCS) as in
          the original CAMELS data set. In each HUC folder are the forcing files (.txt), starting with the 8-digit basin 
          id.
      basin : str
          8-digit USGS identifier of the basin.
      forcings : str
          Can be e.g. 'daymet' or 'nldas', etc. Must match the folder names in the 'basin_mean_forcing' directory. 
  
      Returns
      -------
      pd.DataFrame
          Time-indexed DataFrame, containing the forcing data.
      int
          Catchment area (m2), specified in the header of the forcing file.
      """
      forcing_path = data_dir / 'basin_mean_forcing' / forcings
      if not forcing_path.is_dir():
          raise OSError(f"{forcing_path} does not exist")
  
      files = list(forcing_path.glob('**/*_forcing_leap.txt'))
      file_path = [f for f in files if f.name[:8] == basin]
      if file_path:
          file_path = file_path[0]
      else:
          raise FileNotFoundError(f'No file for Basin {basin} at {file_path}')
  
      df = pd.read_csv(file_path, sep='\s+', header=3)
      df["date"] = pd.to_datetime(df.Year.map(str) + "/" + df.Mnth.map(str) + "/" + df.Day.map(str), format="%Y/%m/%d")
      df = df.set_index("date")
  
      # load area from header
      with open(file_path, 'r') as fp:
          content = fp.readlines()
          area = int(content[2])
  
      return df, area
  
  
  def load_camels_us_discharge(data_dir: Path, basin: str, area: int) -> pd.Series:
      """Load the discharge data for a basin of the CAMELS US data set.
  
      Parameters
      ----------
      data_dir : Path
          Path to the CAMELS US directory. This folder must contain a 'usgs_streamflow' folder with 18
          subdirectories (for the 18 HUCS) as in the original CAMELS data set. In each HUC folder are the discharge files 
          (.txt), starting with the 8-digit basin id.
      basin : str
          8-digit USGS identifier of the basin.
      area : int
          Catchment area (m2), used to normalize the discharge.
  
      Returns
      -------
      pd.Series
          Time-index pandas.Series of the discharge values (mm/day)
      """
  
      discharge_path = data_dir / 'usgs_streamflow'
      files = list(discharge_path.glob('**/*_streamflow_qc.txt'))
      file_path = [f for f in files if f.name[:8] == basin]
      if file_path:
          file_path = file_path[0]
      else:
          raise FileNotFoundError(f'No file for Basin {basin} at {file_path}')
  
      col_names = ['basin', 'Year', 'Mnth', 'Day', 'QObs', 'flag']
      df = pd.read_csv(file_path, sep='\s+', header=None, names=col_names)
      df["date"] = pd.to_datetime(df.Year.map(str) + "/" + df.Mnth.map(str) + "/" + df.Day.map(str), format="%Y/%m/%d")
      df = df.set_index("date")
  
      # normalize discharge from cubic feed per second to mm per day
      df.QObs = 28316846.592 * df.QObs * 86400 / (area * 10**6)
  
      return df.QObs
  
apis: pandas.read_csv pandas.concat