diff --git a/tests/models/hydrology/test_above_ground.py b/tests/models/hydrology/test_above_ground.py index 894a63927..54a4b10cc 100644 --- a/tests/models/hydrology/test_above_ground.py +++ b/tests/models/hydrology/test_above_ground.py @@ -215,3 +215,18 @@ def test_estimate_interception(): exp_result = np.array([0.0, 1.180619, 5.339031]) np.testing.assert_allclose(result, exp_result) + + +def test_distribute_monthly_rainfall(): + """Test that randomly generated numbers are reproducible.""" + from virtual_rainforest.models.hydrology.above_ground import ( + distribute_monthly_rainfall, + ) + + monthly_rain = np.array([0.0, 20.0, 200.0]) + result = distribute_monthly_rainfall(monthly_rain, 10) + result1 = distribute_monthly_rainfall(monthly_rain, 10) + + assert result.shape == (3, 10) + np.testing.assert_allclose(result.sum(axis=1), monthly_rain) + np.testing.assert_allclose(result, result1) diff --git a/tests/models/hydrology/test_hydrology_model.py b/tests/models/hydrology/test_hydrology_model.py index 0c89f1bb9..6edeeda41 100644 --- a/tests/models/hydrology/test_hydrology_model.py +++ b/tests/models/hydrology/test_hydrology_model.py @@ -401,7 +401,7 @@ def test_setup( dims=["layers", "cell_id"], ), DataArray( - [[0.507389, 0.507389, 0.507389], [0.451635, 0.451635, 0.451635]], + [[0.504361, 0.502679, 0.501345], [0.451635, 0.451635, 0.451635]], dims=["layers", "cell_id"], ), ], @@ -409,7 +409,7 @@ def test_setup( ).assign_coords(model.data["layer_heights"].coords) exp_surf_prec = DataArray( - [177.113493, 177.113493, 177.113493], + [177.561204, 177.822518, 177.591868], dims=["cell_id"], coords={"cell_id": [0, 1, 2]}, ) @@ -419,7 +419,7 @@ def test_setup( coords={"cell_id": [0, 1, 2]}, ) exp_vertical_flow = DataArray( - [55.756815, 55.756815, 55.756815], + [54.432845, 53.850539, 53.01296], dims=["cell_id"], coords={"cell_id": [0, 1, 2]}, ) @@ -429,7 +429,7 @@ def test_setup( coords={"cell_id": [0, 1, 2]}, ) exp_stream_flow = DataArray( - [117.182581, 117.182581, 117.182581], + [117.629926, 117.896289, 117.669184], dims=["cell_id"], coords={"cell_id": [0, 1, 2]}, ) diff --git a/virtual_rainforest/models/hydrology/above_ground.py b/virtual_rainforest/models/hydrology/above_ground.py index b484b4eaa..1eaf3747c 100644 --- a/virtual_rainforest/models/hydrology/above_ground.py +++ b/virtual_rainforest/models/hydrology/above_ground.py @@ -272,3 +272,36 @@ def estimate_interception( * (1 - np.exp(-canopy_density_factor * precipitation / max_capacity)), nan=0.0, ) + + +def distribute_monthly_rainfall( + total_monthly_rainfall: NDArray[np.float32], + num_days: int, +) -> NDArray[np.float32]: + """Distributes total monthly rainfall over the specified number of days. + + At the moment, this function allocates each millimeter of monthly rainfall to a + randomly selected day. In the future, this allocation could be based on observed + rainfall patterns. + + Args: + total_monthly_rainfall: Total monthly rainfall, [mm] + num_days (int, optional): Number of days to distribute the rainfall over. + + Returns: + An array containing the daily rainfall amounts [mm]. + """ + np.random.seed(42) # TODO move this to core for all models to be the same + + daily_rainfall_data = [] + for total_monthly_rainfall in total_monthly_rainfall: + daily_rainfall = np.zeros(num_days) + + for _ in range(int(total_monthly_rainfall)): + day = np.random.randint(0, num_days) # Randomly select a day + daily_rainfall[day] += 1.0 # Add 1.0 mm of rainfall to the selected day + + daily_rainfall *= total_monthly_rainfall / np.sum(daily_rainfall) + daily_rainfall_data.append(daily_rainfall) + + return np.nan_to_num(np.array(daily_rainfall_data), nan=0.0) diff --git a/virtual_rainforest/models/hydrology/hydrology_model.py b/virtual_rainforest/models/hydrology/hydrology_model.py index 0cfc472c1..8fa3d3924 100644 --- a/virtual_rainforest/models/hydrology/hydrology_model.py +++ b/virtual_rainforest/models/hydrology/hydrology_model.py @@ -364,9 +364,10 @@ def update(self, time_index: int) -> None: days: int = 30 # Select variables at relevant heights for current time step - current_precipitation = ( - self.data["precipitation"].isel(time_index=time_index) / days - ).to_numpy() + current_precipitation = above_ground.distribute_monthly_rainfall( + (self.data["precipitation"].isel(time_index=time_index)).to_numpy(), + days, + ) leaf_area_index_sum = self.data["leaf_area_index"].sum(dim="layers").to_numpy() evapotranspiration = ( self.data["evapotranspiration"].sum(dim="layers") / days @@ -420,7 +421,7 @@ def update(self, time_index: int) -> None: # Interception of water in canopy, [mm] interception = above_ground.estimate_interception( leaf_area_index=leaf_area_index_sum, - precipitation=current_precipitation, + precipitation=current_precipitation[:, day], intercept_param_1=self.constants.intercept_param_1, intercept_param_2=self.constants.intercept_param_2, intercept_param_3=self.constants.intercept_param_3, @@ -428,7 +429,7 @@ def update(self, time_index: int) -> None: ) # Precipitation that reaches the surface per day, [mm] - precipitation_surface = current_precipitation - interception + precipitation_surface = current_precipitation[:, day] - interception daily_lists["precipitation_surface"].append(precipitation_surface) # Calculate how much water can be added to soil before capacity is reached,