read_mean_prcp now only support CAMELS_US

examples/scripts.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2022-09-06 23:42:46
-LastEditTime: 2024-09-14 13:27:17
+LastEditTime: 2024-11-11 16:44:31
 LastEditors: Wenyu Ouyang
 Description: examples for using hydrodataset
 FilePath: \hydrodataset\examples\scripts.py
@@ -18,7 +18,7 @@
 camels_cl_path = os.path.join("camels", "camels_cl")
 camels_gb_path = os.path.join("camels", "camels_gb")
 camels_us_path = os.path.join("camels", "camels_us")
-camels_aus_v2_path=os.path.join("camels","camels_aus_v2")
+camels_aus_v2_path = os.path.join("camels", "camels_aus_v2")
 
 aus_v2_region = "AUS_v2"
 aus_region = "AUS"
@@ -174,7 +174,7 @@
 )
 
 # # ---------------------------- AUS-V2 -------------------------------
-camels_aus_v2=Camels(camels_aus_v2_path, download=False, region=aus_v2_region)
+camels_aus_v2 = Camels(camels_aus_v2_path, download=False, region=aus_v2_region)
 gage_ids = camels_aus_v2.read_object_ids()
 assert gage_ids.size == 561
 attrs = camels_aus_v2.read_constant_cols(

hydrodataset/camels.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2022-01-05 18:01:11
-LastEditTime: 2024-09-10 19:15:33
+LastEditTime: 2024-11-11 17:26:42
 LastEditors: Wenyu Ouyang
 Description: Read Camels Series ("AUStralia", "BRazil", "ChiLe", "GreatBritain", "UnitedStates") datasets
 FilePath: \hydrodataset\hydrodataset\camels.py
@@ -595,7 +595,9 @@ def get_relevant_cols(self) -> np.ndarray:
                 if root == self.data_source_description["CAMELS_FORCING_DIR"]:
                     continue
                 forcing_types.extend(
-                    file[:-4] for file in files if file not in ["ClimaticIndices.csv", "desktop.ini"]
+                    file[:-4]
+                    for file in files
+                    if file not in ["ClimaticIndices.csv", "desktop.ini"]
                 )
             return np.array(forcing_types)
         elif self.region == "BR":
@@ -1295,7 +1297,9 @@ def read_relevant_cols(
                 chosen_data = forcing_data[gage_id_lst].values[ind1, :]
                 x[:, ind2, k] = chosen_data.T
         elif self.region == "AUS_v2":
-            for k in tqdm(range(len(var_lst)), desc="Read forcing data of CAMELS-AUS-V2"):
+            for k in tqdm(
+                range(len(var_lst)), desc="Read forcing data of CAMELS-AUS-V2"
+            ):
                 if "precipitation_" in var_lst[k]:
                     forcing_dir = os.path.join(
                         self.data_source_description["CAMELS_FORCING_DIR"],
@@ -1559,13 +1563,37 @@ def read_area(self, gage_id_lst) -> np.ndarray:
         else:
             raise NotImplementedError(CAMELS_NO_DATASET_ERROR_LOG)
 
-    def read_mean_prcp(self, gage_id_lst, unit="mm/d") -> np.ndarray:
+    def read_mean_prcp(self, gage_id_lst, unit="mm/d") -> xr.Dataset:
+        """Read mean precipitation data
+
+        Parameters
+        ----------
+        gage_id_lst : list
+            station ids
+        unit : str, optional
+            the unit of mean_prcp, by default "mm/d"
+
+        Returns
+        -------
+        xr.Dataset
+            TODO: now only support CAMELS-US
+
+        Raises
+        ------
+        NotImplementedError
+            some regions are not supported
+        ValueError
+            unit must be one of ['mm/d', 'mm/day', 'mm/h', 'mm/hour', 'mm/3h', 'mm/3hour', 'mm/8d', 'mm/8day']
+        """
         if self.region in ["US", "AUS", "AUS_v2", "BR", "GB"]:
             if self.region == "US":
                 data = self.read_attr_xrdataset(gage_id_lst, ["p_mean"])
-            data = self.read_constant_cols(
-                gage_id_lst, ["p_mean"], is_return_dict=False,
-            )
+            else:
+                data = self.read_constant_cols(
+                    gage_id_lst,
+                    ["p_mean"],
+                    is_return_dict=False,
+                )
         elif self.region == "CL":
             # there are different p_mean values for different forcings, here we chose p_mean_cr2met now
             data = self.read_constant_cols(
@@ -1579,9 +1607,11 @@ def read_mean_prcp(self, gage_id_lst, unit="mm/d") -> np.ndarray:
             converted_data = data / 24
         elif unit in ["mm/3h", "mm/3hour"]:
             converted_data = data / 8
+        elif unit in ["mm/8d", "mm/8day"]:
+            converted_data = data * 8
         else:
             raise ValueError(
-                "unit must be one of ['mm/d', 'mm/day', 'mm/h', 'mm/hour', 'mm/3h', 'mm/3hour']"
+                "unit must be one of ['mm/d', 'mm/day', 'mm/h', 'mm/hour', 'mm/3h', 'mm/3hour', 'mm/8d', 'mm/8day']"
             )
         return converted_data
 
@@ -1897,4 +1927,4 @@ def read_attr_xrdataset(self, gage_id_lst=None, var_lst=None, **kwargs):
 
     @property
     def streamflow_unit(self):
-        return "foot^3/s"
+        return "foot^3/s"

tests/test_camels.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2022-09-05 23:20:24
-LastEditTime: 2024-11-04 20:15:11
+LastEditTime: 2024-11-11 17:29:13
 LastEditors: Wenyu Ouyang
 Description: Tests for `hydrodataset` package
 FilePath: \hydrodataset\tests\test_camels.py
@@ -16,6 +16,7 @@
 from hydrodataset import Camels
 import numpy as np
 import pandas as pd
+import xarray as xr
 from unittest.mock import patch, MagicMock
 
 
@@ -76,6 +77,7 @@ def test_read_mean_prcp():
     gage_ids = camels.read_object_ids()
     mean_prcp = camels.read_mean_prcp(gage_ids[:5])
     print(mean_prcp)
+    assert isinstance(mean_prcp, xr.Dataset)
 
 
 def test_read_target_cols_us():
@@ -281,16 +283,6 @@ def test_read_camels_us_model_output_data_no_data():
         assert np.all(np.isnan(result))
 
 
-def test_read_mean_prcp_us():
-    camels = Camels()
-    camels.region = "US"
-    camels.read_attr_xrdataset = MagicMock(return_value=np.array([1.0, 2.0, 3.0]))
-
-    gage_id_lst = ["01013500", "01013501", "01013502"]
-    result = camels.read_mean_prcp(gage_id_lst)
-    assert np.array_equal(result, np.array([1.0, 2.0, 3.0]))
-
-
 def test_read_mean_prcp_aus():
     camels = Camels()
     camels.region = "AUS"