add a param unit for read_mean_prcp but not finished yet

hydrodataset/camels.py

@@ -1430,7 +1430,7 @@ 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) -> np.ndarray:
+    def read_mean_prcp(self, gage_id_lst, unit="mm/d"):
         if self.region in ["US", "AUS", "BR", "GB"]:
             if self.region == "US":
                 return self.read_attr_xrdataset(gage_id_lst, ["p_mean"])

hydrodataset/caravan.py

@@ -846,7 +846,7 @@ def streamflow_unit(self):
     def read_area(self, gage_id_lst=None):
         return self.read_attr_xrdataset(gage_id_lst, ["area"])
 
-    def read_mean_prcp(self, gage_id_lst=None):
+    def read_mean_prcp(self, gage_id_lst=None, unit="mm/d"):
         return self.read_attr_xrdataset(gage_id_lst, ["p_mean"])
 
 

hydrodataset/hydro_dataset.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2022-09-05 23:20:24
-LastEditTime: 2024-09-10 18:26:45
+LastEditTime: 2024-11-04 20:06:05
 LastEditors: Wenyu Ouyang
 Description: main modules for hydrodataset
 FilePath: \hydrodataset\hydrodataset\hydro_dataset.py
@@ -113,6 +113,20 @@ def read_area(self, gage_id_lst):
         """read area of each basin/unit"""
         raise NotImplementedError
 
-    def read_mean_prcp(self, gage_id_lst):
-        """read mean precipitation of each basin/unit"""
+    def read_mean_prcp(self, gage_id_lst, unit="mm/d"):
+        """read mean precipitation of each basin
+        default unit is mm/d, but one can chose other units and we will convert the unit to the specified unit
+
+        Parameters
+        ----------
+        gage_id_lst : list, optional
+            the list of gage ids, by default None
+        unit : str, optional
+            the unit of precipitation, by default "mm/d"
+
+        Returns
+        -------
+        xr.Dataset
+            the mean precipitation of each basin
+        """
         raise NotImplementedError

tests/test_camels.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2022-09-05 23:20:24
-LastEditTime: 2024-09-10 19:21:25
+LastEditTime: 2024-11-04 20:15:11
 LastEditors: Wenyu Ouyang
 Description: Tests for `hydrodataset` package
 FilePath: \hydrodataset\tests\test_camels.py
@@ -279,3 +279,64 @@ def test_read_camels_us_model_output_data_no_data():
         result = camels.read_camels_us_model_output_data(gage_id_lst, t_range, var_lst)
         assert result.shape == (1, 3653, 2)
         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"
+    camels.read_constant_cols = MagicMock(return_value=np.array([1.0, 2.0, 3.0]))
+
+    gage_id_lst = ["12345678", "12345679", "12345680"]
+    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_br():
+    camels = Camels()
+    camels.region = "BR"
+    camels.read_constant_cols = MagicMock(return_value=np.array([1.0, 2.0, 3.0]))
+
+    gage_id_lst = ["12345678", "12345679", "12345680"]
+    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_gb():
+    camels = Camels()
+    camels.region = "GB"
+    camels.read_constant_cols = MagicMock(return_value=np.array([1.0, 2.0, 3.0]))
+
+    gage_id_lst = ["12345678", "12345679", "12345680"]
+    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_cl():
+    camels = Camels()
+    camels.region = "CL"
+    camels.read_constant_cols = MagicMock(return_value=np.array([1.0, 2.0, 3.0]))
+
+    gage_id_lst = ["12345678", "12345679", "12345680"]
+    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_invalid_region():
+    camels = Camels()
+    camels.region = "INVALID"
+
+    gage_id_lst = ["12345678", "12345679", "12345680"]
+    try:
+        camels.read_mean_prcp(gage_id_lst)
+    except NotImplementedError as e:
+        assert str(e) == CAMELS_NO_DATASET_ERROR_LOG

tests/test_caravan.py

@@ -7,6 +7,7 @@
 FilePath: /hydrodataset/tests/test_caravan.py
 Copyright (c) 2023-2024 Wenyu Ouyang. All rights reserved.
 """
+
 import os
 import numpy as np
 import pytest
@@ -182,6 +183,7 @@ def test_read_area(caravan):
 
 
 def test_read_prcp_mean(caravan):
+    # be careful for the unit!
     caravan_ids = caravan.read_object_ids()
     prcp_mean = caravan.read_mean_prcp(
         caravan_ids[:3].tolist() + caravan_ids[-2:].tolist()