From dd11e1f9474af5927eb2da093af24430285f79c2 Mon Sep 17 00:00:00 2001 From: "ouyang,wenyu" Date: Mon, 4 Nov 2024 20:19:31 +0800 Subject: [PATCH] add a param unit for read_mean_prcp but not finished yet --- hydrodataset/camels.py | 2 +- hydrodataset/caravan.py | 2 +- hydrodataset/hydro_dataset.py | 20 +++++++++-- tests/test_camels.py | 63 ++++++++++++++++++++++++++++++++++- tests/test_caravan.py | 2 ++ 5 files changed, 83 insertions(+), 6 deletions(-) diff --git a/hydrodataset/camels.py b/hydrodataset/camels.py index 650f62d..cde52aa 100644 --- a/hydrodataset/camels.py +++ b/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"]) diff --git a/hydrodataset/caravan.py b/hydrodataset/caravan.py index 9965c62..5b707e8 100644 --- a/hydrodataset/caravan.py +++ b/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"]) diff --git a/hydrodataset/hydro_dataset.py b/hydrodataset/hydro_dataset.py index cad3383..bfed5af 100644 --- a/hydrodataset/hydro_dataset.py +++ b/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 diff --git a/tests/test_camels.py b/tests/test_camels.py index 738e305..957a476 100644 --- a/tests/test_camels.py +++ b/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 diff --git a/tests/test_caravan.py b/tests/test_caravan.py index beab166..af4714d 100644 --- a/tests/test_caravan.py +++ b/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()