Feature/add radiation input

tests/conftest.py

@@ -480,6 +480,13 @@ def dummy_climate_data(fixture_core_components):
         coords=full_coordinates,
         name="leaf_area_index",
     )
+    canopy_absorption = np.repeat(a=[np.nan, 1.0, np.nan], repeats=[1, 3, 11])
+    data["canopy_absorption"] = DataArray(
+        np.broadcast_to(canopy_absorption, (3, 15)).T,
+        dims=["layers", "cell_id"],
+        coords=full_coordinates,
+        name="canopy_absorption",
+    )
 
     layer_heights = np.repeat(
         a=[32.0, 30.0, 20.0, 10.0, np.nan, 1.5, 0.1, -0.5, -1.0],

tests/models/abiotic/test_abiotic_model.py

@@ -369,6 +369,33 @@ def test_update_abiotic_model(dummy_climate_data, cfg_string):
 
     model.update(time_index=0)
 
+    # Check that updated vars are in data object
+    for var in [
+        "air_temperature",
+        "canopy_temperature",
+        "soil_temperature",
+        "vapour_pressure",
+        "vapour_pressure_deficit",
+        "air_conductivity",
+        "conductivity_from_ref_height",
+        "leaf_air_heat_conductivity",
+        "leaf_vapour_conductivity",
+        "wind_speed",
+        "friction_velocity",
+        "diabatic_correction_heat_above",
+        "diabatic_correction_momentum_above",
+        "diabatic_correction_heat_canopy",
+        "diabatic_correction_momentum_canopy",
+        "sensible_heat_flux",
+        "latent_heat_flux",
+        "ground_heat_flux",
+        "soil_absorption",
+        "longwave_emission_soil",
+        "molar_density_air",
+        "specific_heat_air",
+    ]:
+        assert var in model.data
+
     friction_velocity_exp = np.repeat(0.161295, 3)
     np.testing.assert_allclose(
         model.data["friction_velocity"],
@@ -377,15 +404,10 @@ def test_update_abiotic_model(dummy_climate_data, cfg_string):
         atol=1e-3,
     )
 
-    wind_speed_exp = np.full((15, 3), np.nan)
-    wind_speed_exp[[0, 1, 2, 3, 11, 12], :] = [
-        [0.727122, 0.727122, 0.727122],
-        [0.615474, 0.615474, 0.615474],
-        [0.587838, 0.587838, 0.587838],
-        [0.537028, 0.537028, 0.537028],
-        [0.506793, 0.506793, 0.506793],
-        [0.50198, 0.50198, 0.50198],
-    ]
+    wind_speed_exp = DataArray(np.full((15, 3), np.nan), dims=["layers", "cell_id"])
+    wind_vals = [0.727122, 0.615474, 0.587838, 0.537028, 0.506793, 0.50198]
+    wind_speed_exp.T[..., [0, 1, 2, 3, 11, 12]] = wind_vals
+
     np.testing.assert_allclose(
         model.data["wind_speed"],
         DataArray(wind_speed_exp),
@@ -397,33 +419,24 @@ def test_update_abiotic_model(dummy_climate_data, cfg_string):
         np.concatenate(
             [
                 [[np.nan, np.nan, np.nan]] * 13,
-                [[20.713125, 20.712525, 20.712458], [20.0] * 3],
+                [[20.713167, 20.708367, 20.707833], [20.0] * 3],
             ],
             axis=0,
         ),
         dims=["layers", "cell_id"],
     )
+
     np.testing.assert_allclose(
         model.data["soil_temperature"],
         exp_new_soiltemp,
         rtol=1e-04,
         atol=1e-04,
     )
 
-    exp_gv = DataArray(
-        np.concatenate(
-            [
-                [[np.nan, np.nan, np.nan]],
-                [
-                    [0.495047, 0.495047, 0.495047],
-                    [0.483498, 0.483498, 0.483498],
-                    [0.46169, 0.46169, 0.46169],
-                ],
-                [[np.nan, np.nan, np.nan]] * 11,
-            ],
-        ),
-        dims=["layers", "cell_id"],
-    )
+    exp_gv = DataArray(np.full((15, 3), np.nan), dims=["layers", "cell_id"])
+    gv_vals = [0.496563, 0.485763, 0.465142]
+    exp_gv.T[..., [1, 2, 3]] = gv_vals
+
     np.testing.assert_allclose(
         model.data["leaf_vapour_conductivity"], exp_gv, rtol=1e-03, atol=1e-03
     )
@@ -445,17 +458,14 @@ def test_update_abiotic_model(dummy_climate_data, cfg_string):
     )
 
     # TODO fix fluxes from soil
-    # exp_latent_heat = DataArray(np.full((15, 3), np.nan), dims=["layers", "cell_id"])
-    # lat_heat_vals = [27.916181, 27.386375, 15.775225, 1]
-    # exp_latent_heat.T[..., [1, 2, 3, 13]] = lat_heat_vals
     exp_latent_heat = DataArray(
         np.concatenate(
             [
                 [[np.nan, np.nan, np.nan]],
                 [
-                    [27.916181, 27.916181, 27.916181],
-                    [27.386375, 27.386375, 27.386375],
-                    [15.775225, 15.775225, 15.775225],
+                    [28.07077, 28.07077, 28.07077],
+                    [27.568713, 27.568715, 27.568716],
+                    [16.006245, 16.006317, 16.006325],
                 ],
                 [[np.nan, np.nan, np.nan]] * 9,
                 [[2.254, 22.54, 225.4]],
@@ -467,10 +477,21 @@ def test_update_abiotic_model(dummy_climate_data, cfg_string):
         model.data["latent_heat_flux"], exp_latent_heat, rtol=1e-04, atol=1e-04
     )
 
-    exp_sens_heat = DataArray(np.full((15, 3), np.nan), dims=["layers", "cell_id"])
-    sens_heat_vals = [-16.814787, -16.29302, -5.416152, -185.669563]
-    exp_sens_heat.T[..., [1, 2, 3, 13]] = sens_heat_vals
-
+    exp_sens_heat = DataArray(
+        np.concatenate(
+            [
+                [[np.nan, np.nan, np.nan]],
+                [
+                    [-16.970825, -16.970825, -16.970825],
+                    [-16.47644, -16.47644, -16.47644],
+                    [-5.637158, -5.637226, -5.637233],
+                ],
+                [[np.nan, np.nan, np.nan]] * 9,
+                [[-192.074608, -192.074608, -192.074608]],
+                [[np.nan, np.nan, np.nan]],
+            ]
+        )
+    )
     np.testing.assert_allclose(
         model.data["sensible_heat_flux"], exp_sens_heat, rtol=1e-04, atol=1e-04
     )

tests/models/abiotic/test_soil_energy_balance.py

@@ -102,9 +102,6 @@ def test_calculate_soil_heat_balance(dummy_climate_data):
     )
 
     data = dummy_climate_data
-    data["shortwave_radiation_surface"] = DataArray(
-        np.array([100, 10, 0]), dims="cell_id"
-    )
     data["soil_evaporation"] = DataArray(np.array([0.001, 0.01, 0.1]), dims="cell_id")
     data["molar_density_air"] = DataArray(
         np.full((15, 3), 38), dims=["layers", "cell_id"]
@@ -119,6 +116,7 @@ def test_calculate_soil_heat_balance(dummy_climate_data):
 
     result = calculate_soil_heat_balance(
         data=data,
+        time_index=0,
         topsoil_layer_index=13,
         update_interval=43200,
         abiotic_consts=AbioticConsts,
@@ -137,7 +135,7 @@ def test_calculate_soil_heat_balance(dummy_climate_data):
     variables = [var for var in result if var not in var_list]
     assert variables
 
-    np.testing.assert_allclose(result["soil_absorption"], np.array([87.5, 8.75, 0.0]))
+    np.testing.assert_allclose(result["soil_absorption"], np.repeat(79.625, 3))
     np.testing.assert_allclose(
         result["longwave_emission_soil"],
         np.repeat(0.007258, 3),
@@ -158,7 +156,7 @@ def test_calculate_soil_heat_balance(dummy_climate_data):
     )
     np.testing.assert_allclose(
         result["ground_heat_flux"],
-        np.array([81.841, -17.195, -228.805]),
+        np.array([73.966007, 53.680007, -149.179993]),
         rtol=1e-04,
         atol=1e-04,
     )

virtual_ecosystem/example_data/generation_scripts/topofcanopy_radiation_example_data.py

@@ -0,0 +1,40 @@
+"""Simple top of canopy shortwave radiation for `ve_run` example data.
+
+This code creates top of canopy shortwave radiation data as input to setup the abiotic
+model. The current values are typical hourly averages for tropical regions.
+
+Once the new netcdf file is created, the final step is to add the grid information to
+the grid config `TOML` to load this data correctly when setting up a Virtual Ecosystem
+Simulation. Here, we can also add the 45 m offset to position the coordinated at the
+centre of the grid cell.
+
+[core.grid]
+cell_nx = 9
+cell_ny = 9
+cell_area = 8100
+xoff = -45.0
+yoff = -45.0
+"""
+
+import numpy as np
+from xarray import DataArray, Dataset
+
+from virtual_ecosystem.example_data.generation_scripts.common import (
+    cell_id,
+    n_cells,
+    n_dates,
+    time,
+    time_index,
+)
+
+data = Dataset()
+
+# Spatio-temporal shortwave radiation flux data [W m-2]
+data["topofcanopy_radiation"] = DataArray(
+    data=np.full((n_cells, n_dates), fill_value=250),
+    coords={"cell_id": cell_id, "time_index": time_index},
+)
+
+data["time"] = DataArray(time, coords={"time_index": time_index})
+
+data.to_netcdf("../data/example_topofcanopy_radiation.nc", format="NETCDF3_64BIT")

virtual_ecosystem/models/abiotic/abiotic_model.py

@@ -61,7 +61,7 @@ class AbioticModel(
         "soil_temperature",
         "vapour_pressure",
         "vapour_pressure_deficit",
-        "air_heat_conductivity",
+        "air_conductivity",
         "conductivity_from_ref_height",
         "leaf_air_heat_conductivity",
         "leaf_vapour_conductivity",
@@ -76,6 +76,8 @@ class AbioticModel(
         "ground_heat_flux",
         "soil_absorption",
         "longwave_emission_soil",
+        "molar_density_air",
+        "specific_heat_air",
     ),
 ):
     """A class describing the abiotic model.
@@ -307,13 +309,14 @@ def update(self, time_index: int, **kwargs: Any) -> None:
         wind_output = {}
 
         # Might make sense to store the shape and use np.full(shape, np.nan)
-        wind_speed_data = np.full_like(self.data["leaf_area_index"], np.nan)
-        wind_speed_data[true_aboveground_rows, :] = wind_update["wind_speed"]
-        wind_output["wind_speed"] = DataArray(
-            wind_speed_data,
-            dims=self.data["layer_heights"].dims,
-            coords=self.data["layer_heights"].coords,
-        )
+        for var in ["wind_speed", "molar_density_air", "specific_heat_air"]:
+            var_out = np.full_like(self.data["leaf_area_index"], np.nan)
+            var_out[true_aboveground_rows, :] = wind_update[var]
+            wind_output[var] = DataArray(
+                var_out,
+                dims=self.data["layer_heights"].dims,
+                coords=self.data["layer_heights"].coords,
+            )
 
         for var in [
             "friction_velocity",
@@ -322,14 +325,15 @@ def update(self, time_index: int, **kwargs: Any) -> None:
             "diabatic_correction_heat_canopy",
             "diabatic_correction_momentum_canopy",
         ]:
-            var_out = DataArray(wind_update[var], dims="cell_id")
-            wind_output[var] = var_out
+            var_out = wind_update[var]
+            wind_output[var] = DataArray(var_out, dims="cell_id")
 
         self.data.add_from_dict(output_dict=wind_output)
 
         # Soil energy balance
         soil_heat_balance = soil_energy_balance.calculate_soil_heat_balance(
             data=self.data,
+            time_index=time_index,
             topsoil_layer_index=topsoil_layer_index,
             update_interval=43200,  # TODO self.model_timing.update_interval
             abiotic_consts=self.model_constants,

virtual_ecosystem/models/abiotic/soil_energy_balance.py

@@ -182,6 +182,7 @@ def update_surface_temperature(
 
 def calculate_soil_heat_balance(
     data: Data,
+    time_index: int,
     topsoil_layer_index: int,
     update_interval: Quantity,
     abiotic_consts: AbioticConsts,
@@ -219,6 +220,7 @@ def calculate_soil_heat_balance(
 
     Args:
         data: instance if a data object
+        time_index: time index
         update_interval: Update interval, [s]
         AbioticConsts: set of constants specific to abiotic model
         CoreConsts: set of constants that are shared across the model
@@ -232,11 +234,15 @@ def calculate_soil_heat_balance(
     output = {}
 
     # Calculate soil absorption of shortwave radiation, [W m-2]
+    shortwave_radiation_surface = data["topofcanopy_radiation"].isel(
+        time_index=time_index
+    ) - (data["canopy_absorption"].sum())
     soil_absorption = calculate_soil_absorption(
-        shortwave_radiation_surface=data["shortwave_radiation_surface"].to_numpy(),
+        shortwave_radiation_surface=shortwave_radiation_surface.to_numpy(),
         surface_albedo=abiotic_consts.surface_albedo,
     )
     output["soil_absorption"] = soil_absorption
+    output["shortwave_radiation_surface"] = shortwave_radiation_surface.to_numpy()
 
     # Calculate longwave emission from topsoil, [W m-2]; note that this is the soil
     # temperature of the previous time step

virtual_ecosystem/models/abiotic/wind.py

@@ -644,13 +644,15 @@ def calculate_wind_profile(
         standard_pressure=core_constants.standard_pressure,
         celsius_to_kelvin=core_constants.zero_Celsius,
     )
+    output["molar_density_air"] = molar_density_air
 
     # Calculate specific heat of air, [J mol-1 K-1]
     specific_heat_air = calculate_specific_heat_air(
         temperature=air_temperature,
         molar_heat_capacity_air=core_constants.molar_heat_capacity_air,
         specific_heat_equ_factors=abiotic_constants.specific_heat_equ_factors,
     )
+    output["specific_heat_air"] = specific_heat_air
 
     # Calculate the total leaf area index, [m2 m-2]
     leaf_area_index_sum = np.nansum(leaf_area_index, axis=0)