Adding allocation traits class var, simple test

pyrealm/demography/t_model_functions.py

@@ -724,3 +724,125 @@ def __post_init__(
             z_max_prop=stem_traits.z_max_prop,
             stem_height=self.stem_height,
         )
+
+
+@dataclass()
+class StemAllocation:
+    """Calculate T Model allocation predictions across a set of stems.
+
+    This method calculate predictions of allocation of potential GPP for stems under the
+    T Model :cite:`Li:2014bc`, given a set of traits for those stems and the stem
+    allometries given the stem size.
+
+    Args:
+        stem_traits: An instance of :class:`~pyrealm.demography.flora.Flora` or
+            :class:`~pyrealm.demography.flora.StemTraits`, providing plant functional
+            trait data for a set of stems.
+        stem_allometry: An instance of
+            :class:`~pyrealm.demography.t_model_functions.StemAllometry`
+            providing the stem size data for which to calculate allocation.
+        at_potential_gpp: An array of diameter at breast height values at which to
+            predict stem allometry values.
+    """
+
+    allocation_attrs: ClassVar[tuple[str, ...]] = (
+        "potential_gpp",
+        "whole_crown_gpp",
+        "sapwood_respiration",
+        "foliar_respiration",
+        "fine_root_respiration",
+        "npp",
+        "turnover",
+        "delta_dbh",
+        "delta_stem_mass",
+        "delta_foliage_mass",
+    )
+
+    # Init vars
+    stem_traits: InitVar[Flora | StemTraits]
+    """An instance of :class:`~pyrealm.demography.flora.Flora` or 
+    :class:`~pyrealm.demography.flora.StemTraits`, providing plant functional trait data
+    for a set of stems."""
+    stem_allometry: InitVar[StemAllometry]
+    """An instance of :class:`~pyrealm.demography.t_model_functions.StemAllometry`
+    providing the stem size data for which to calculate allocation."""
+    at_potential_gpp: InitVar[NDArray[np.float32]]
+    """An array of potential gross primary productivity for each stem that should be
+    allocated to respiration, turnover and growth."""
+
+    # Post init allometry attributes
+    potential_gpp: NDArray[np.float32] = field(init=False)
+    whole_crown_gpp: NDArray[np.float32] = field(init=False)
+    sapwood_respiration: NDArray[np.float32] = field(init=False)
+    foliar_respiration: NDArray[np.float32] = field(init=False)
+    fine_root_respiration: NDArray[np.float32] = field(init=False)
+    npp: NDArray[np.float32] = field(init=False)
+    turnover: NDArray[np.float32] = field(init=False)
+    delta_dbh: NDArray[np.float32] = field(init=False)
+    delta_stem_mass: NDArray[np.float32] = field(init=False)
+    delta_foliage_mass: NDArray[np.float32] = field(init=False)
+
+    def __post_init__(
+        self,
+        stem_traits: Flora | StemTraits,
+        stem_allometry: StemAllometry,
+        at_potential_gpp: NDArray[np.float32],
+    ) -> None:
+        """Populate stem allocation attributes from the traits, allometry and GPP."""
+
+        # Broadcast potential GPP to match stem data outputs
+        self.potential_gpp = np.broadcast_to(at_potential_gpp, stem_allometry.dbh.shape)
+
+        self.whole_crown_gpp = calculate_whole_crown_gpp(
+            potential_gpp=self.potential_gpp,
+            crown_area=stem_allometry.crown_area,
+            par_ext=stem_traits.par_ext,
+            lai=stem_traits.lai,
+        )
+
+        self.sapwood_respiration = calculate_sapwood_respiration(
+            resp_s=stem_traits.resp_s, sapwood_mass=stem_allometry.sapwood_mass
+        )
+
+        self.foliar_respiration = calculate_foliar_respiration(
+            resp_f=stem_traits.resp_f, whole_crown_gpp=self.whole_crown_gpp
+        )
+
+        self.fine_root_respiration = calculate_fine_root_respiration(
+            zeta=stem_traits.zeta,
+            sla=stem_traits.sla,
+            resp_r=stem_traits.resp_r,
+            foliage_mass=stem_allometry.foliage_mass,
+        )
+
+        self.npp = calculate_net_primary_productivity(
+            yld=stem_traits.yld,
+            whole_crown_gpp=self.whole_crown_gpp,
+            foliar_respiration=self.foliar_respiration,
+            fine_root_respiration=self.fine_root_respiration,
+            sapwood_respiration=self.sapwood_respiration,
+        )
+
+        self.turnover = calculate_foliage_and_fine_root_turnover(
+            sla=stem_traits.sla,
+            zeta=stem_traits.zeta,
+            tau_f=stem_traits.tau_f,
+            tau_r=stem_traits.tau_r,
+            foliage_mass=stem_allometry.foliage_mass,
+        )
+
+        (self.delta_dbh, self.delta_stem_mass, self.delta_foliage_mass) = (
+            calculate_growth_increments(
+                rho_s=stem_traits.rho_s,
+                a_hd=stem_traits.a_hd,
+                h_max=stem_traits.h_max,
+                lai=stem_traits.lai,
+                ca_ratio=stem_traits.ca_ratio,
+                sla=stem_traits.sla,
+                zeta=stem_traits.zeta,
+                npp=self.npp,
+                turnover=self.turnover,
+                dbh=stem_allometry.dbh,
+                stem_height=stem_allometry.stem_height,
+            )
+        )

tests/unit/demography/conftest.py

@@ -51,6 +51,10 @@ def rtmodel_flora():
 
     pft_definitions = pft_definitions.drop(columns=["d"])
 
+    # Set foliar respiration to zero to avoid issues with this being applied before
+    # estimating whole crown gpp in rtmodel
+    pft_definitions["resp_f"] = 0
+
     return Flora(
         pfts=[
             PlantFunctionalType(**args)
@@ -70,7 +74,6 @@ def rtmodel_data():
 
     rdata = rdata.rename(
         columns={
-            "dD": "delta_d",
             "D": "dbh",
             "H": "stem_height",
             "fc": "crown_fraction",
@@ -79,18 +82,20 @@ def rtmodel_data():
             "Ws": "stem_mass",
             "Wss": "sapwood_mass",
             "P0": "potential_gpp",
-            "GPP": "crown_gpp",
-            "Rm1": "resp_swd",
-            "Rm2": "resp_frt",
-            "dWs": "delta_mass_stm",
-            "dWfr": "delta_mass_frt",
+            "GPP": "whole_crown_gpp",
+            "Rm1": "sapwood_respiration",
+            "Rm2": "fine_root_respiration",
+            "NPP": "npp",
+            "dD": "delta_dbh",
+            "dWs": "delta_stem_mass",
+            "dWfr": "delta_foliage_mass",
         }
     )
 
     # Fix some scaling differences:
     # The R tmodel implementation rescales reported delta_d as a radial increase in
     # millimetres, not diameter increase in metres
-    rdata["delta_d"] = rdata["delta_d"] / 500
+    rdata["delta_dbh"] = rdata["delta_dbh"] / 500
 
     # Wrap the return data into arrays with PFT as columns and diameter values as rows
     rdata_arrays = {k: np.reshape(v, (3, 6)).T for k, v in rdata.items()}

tests/unit/demography/test_t_model_functions.py

@@ -471,7 +471,7 @@ def test_calculate_whole_crown_gpp(
             )
 
             assert result.shape == exp_shape
-            assert np.allclose(result, rtmodel_data["crown_gpp"][out_idx])
+            assert np.allclose(result, rtmodel_data["whole_crown_gpp"][out_idx])
             return
 
         assert str(excep.value).startswith(excep_msg)
@@ -498,7 +498,7 @@ def test_calculate_sapwood_respiration(
             )
 
             assert result.shape == exp_shape
-            assert np.allclose(result, rtmodel_data["resp_swd"][out_idx])
+            assert np.allclose(result, rtmodel_data["sapwood_respiration"][out_idx])
             return
 
         assert str(excep.value).startswith(excep_msg)
@@ -527,13 +527,14 @@ def test_calculate_foliar_respiration(
         with outcome as excep:
             result = calculate_foliar_respiration(
                 resp_f=rtmodel_flora.resp_f[pft_idx],
-                whole_crown_gpp=rtmodel_data["crown_gpp"][data_idx],
+                whole_crown_gpp=rtmodel_data["whole_crown_gpp"][data_idx],
             )
 
             assert result.shape == exp_shape
             assert np.allclose(
                 result,
-                rtmodel_data["crown_gpp"][data_idx] * rtmodel_flora.resp_f[pft_idx],
+                rtmodel_data["whole_crown_gpp"][data_idx]
+                * rtmodel_flora.resp_f[pft_idx],
             )
             return
 
@@ -563,7 +564,7 @@ def test_calculate_fine_root_respiration(
             )
 
             assert result.shape == exp_shape
-            assert np.allclose(result, rtmodel_data["resp_frt"][out_idx])
+            assert np.allclose(result, rtmodel_data["fine_root_respiration"][out_idx])
             return
 
         assert str(excep.value).startswith(excep_msg)
@@ -588,14 +589,14 @@ def test_calculate_net_primary_productivity(
         with outcome as excep:
             result = calculate_net_primary_productivity(
                 yld=rtmodel_flora.yld[pft_idx],
-                whole_crown_gpp=rtmodel_data["crown_gpp"][data_idx],
+                whole_crown_gpp=rtmodel_data["whole_crown_gpp"][data_idx],
                 foliar_respiration=0,  # Not included here in the R implementation
-                fine_root_respiration=rtmodel_data["resp_frt"][data_idx],
-                sapwood_respiration=rtmodel_data["resp_swd"][data_idx],
+                fine_root_respiration=rtmodel_data["fine_root_respiration"][data_idx],
+                sapwood_respiration=rtmodel_data["sapwood_respiration"][data_idx],
             )
 
             assert result.shape == exp_shape
-            assert np.allclose(result, rtmodel_data["NPP"][out_idx])
+            assert np.allclose(result, rtmodel_data["npp"][out_idx])
             return
 
         assert str(excep.value).startswith(excep_msg)
@@ -658,20 +659,22 @@ def test_calculate_calculate_growth_increments(
                 ca_ratio=rtmodel_flora.ca_ratio[pft_idx],
                 sla=rtmodel_flora.sla[pft_idx],
                 zeta=rtmodel_flora.zeta[pft_idx],
-                npp=rtmodel_data["NPP"][data_idx],
+                npp=rtmodel_data["npp"][data_idx],
                 turnover=rtmodel_data["turnover"][data_idx],
                 dbh=rtmodel_data["dbh"][data_idx],
                 stem_height=rtmodel_data["stem_height"][data_idx],
             )
 
             assert delta_d.shape == exp_shape
-            assert np.allclose(delta_d, rtmodel_data["delta_d"][out_idx])
+            assert np.allclose(delta_d, rtmodel_data["delta_dbh"][out_idx])
 
             assert delta_mass_stm.shape == exp_shape
-            assert np.allclose(delta_mass_stm, rtmodel_data["delta_mass_stm"][out_idx])
+            assert np.allclose(delta_mass_stm, rtmodel_data["delta_stem_mass"][out_idx])
 
             assert delta_mass_frt.shape == exp_shape
-            assert np.allclose(delta_mass_frt, rtmodel_data["delta_mass_frt"][out_idx])
+            assert np.allclose(
+                delta_mass_frt, rtmodel_data["delta_foliage_mass"][out_idx]
+            )
             return
 
         assert str(excep.value).startswith(excep_msg)
@@ -716,6 +719,33 @@ def test_StemAllometry(rtmodel_flora, rtmodel_data):
     )
 
     # Check the variables provided by the rtmodel implementation
-    to_check = set(stem_allometry.allometry_attrs) - set(["canopy_r0", "canopy_z_max"])
-    for var in to_check:
+    vars_to_check = (
+        v
+        for v in stem_allometry.allometry_attrs
+        if v not in ["canopy_r0", "canopy_z_max"]
+    )
+    for var in vars_to_check:
         assert np.allclose(getattr(stem_allometry, var), rtmodel_data[var])
+
+
+def test_StemAllocation(rtmodel_flora, rtmodel_data):
+    """Test the StemAllometry class."""
+
+    from pyrealm.demography.t_model_functions import StemAllocation, StemAllometry
+
+    stem_allometry = StemAllometry(
+        stem_traits=rtmodel_flora, at_dbh=rtmodel_data["dbh"][:, [0]]
+    )
+
+    stem_allocation = StemAllocation(
+        stem_traits=rtmodel_flora,
+        stem_allometry=stem_allometry,
+        at_potential_gpp=rtmodel_data["potential_gpp"],
+    )
+
+    # Check the variables provided by the rtmodel implementation
+    vars_to_check = (
+        v for v in stem_allocation.allocation_attrs if v not in ["foliar_respiration"]
+    )
+    for var in vars_to_check:
+        assert np.allclose(getattr(stem_allocation, var), rtmodel_data[var])