Add StemAllometry class

pyrealm/demography/canopy.py

@@ -79,10 +79,10 @@ def _calculate_canopy(self, community: Community) -> None:
         # Calculate community wide properties: total crown area, maximum height, crown
         # area required to fill a layer and total number of canopy layers
         self.total_community_crown_area = (
-            community.cohort_data["crown_area"] * community.cohort_data["n_individuals"]
+            community.stem_allometry.crown_area * community.cohort_data["n_individuals"]
         ).sum()
 
-        self.max_stem_height = community.cohort_data["stem_height"].max()
+        self.max_stem_height = community.stem_allometry.stem_height.max()
 
         self.crown_area_per_layer = community.cell_area * (1 - self.canopy_gap_fraction)
 
@@ -109,12 +109,12 @@ def _calculate_canopy(self, community: Community) -> None:
             solution = root_scalar(
                 solve_community_projected_canopy_area,
                 args=(
-                    community.cohort_data["stem_height"],
-                    community.cohort_data["crown_area"],
+                    community.stem_allometry.stem_height,
+                    community.stem_allometry.crown_area,
                     community.stem_traits.m,
                     community.stem_traits.n,
                     community.stem_traits.q_m,
-                    community.cohort_data["canopy_z_max"],
+                    community.stem_allometry.canopy_z_max,
                     community.cohort_data["n_individuals"],
                     target_area,
                     False,  # validate
@@ -136,7 +136,7 @@ def _calculate_canopy(self, community: Community) -> None:
         # turning off the validation internally should simply speed up the code.
         self.stem_relative_radius = calculate_relative_canopy_radius_at_z(
             z=self.layer_heights,
-            stem_height=community.cohort_data["stem_height"],
+            stem_height=community.stem_allometry.stem_height,
             m=community.stem_traits.m,
             n=community.stem_traits.n,
             validate=False,
@@ -146,21 +146,21 @@ def _calculate_canopy(self, community: Community) -> None:
         self.stem_crown_area = calculate_stem_projected_crown_area_at_z(
             z=self.layer_heights,
             q_z=self.stem_relative_radius,
-            crown_area=community.cohort_data["crown_area"],
-            stem_height=community.cohort_data["stem_height"],
+            crown_area=community.stem_allometry.crown_area,
+            stem_height=community.stem_allometry.stem_height,
             q_m=community.stem_traits.q_m,
-            z_max=community.cohort_data["canopy_z_max"],
+            z_max=community.stem_allometry.canopy_z_max,
             validate=False,
         )
 
         # Find the projected leaf area of a cohort stem at canopy closure heights.
         self.stem_leaf_area = calculate_stem_projected_leaf_area_at_z(
             z=self.layer_heights,
             q_z=self.stem_relative_radius,
-            crown_area=community.cohort_data["crown_area"],
-            stem_height=community.cohort_data["stem_height"],
+            crown_area=community.stem_allometry.crown_area,
+            stem_height=community.stem_allometry.stem_height,
             f_g=community.stem_traits.f_g,
             q_m=community.stem_traits.q_m,
-            z_max=community.cohort_data["canopy_z_max"],
+            z_max=community.stem_allometry.canopy_z_max,
             validate=False,
         )

pyrealm/demography/community.py

@@ -89,12 +89,17 @@
 ...     cohort_pft_names=cohort_pft_names
 ... )
 
-Convert the community cohort data to a :class:`pandas.DataFrame` for nicer display and
-show some of the calculated T Model predictions:
-
->>> pd.DataFrame(community.cohort_data)[
-...    ['name', 'dbh', 'n_individuals', 'stem_height', 'crown_area', 'stem_mass']
-... ]
+Convert some of the data to a :class:`pandas.DataFrame` for nicer display and show some
+of the calculated T Model predictions:
+
+>>> pd.DataFrame({
+...    'name': community.stem_traits.name,
+...    'dbh': community.stem_allometry.dbh,
+...    'n_individuals': community.cohort_data["n_individuals"],
+...    'stem_height': community.stem_allometry.stem_height,
+...    'crown_area': community.stem_allometry.crown_area,
+...    'stem_mass': community.stem_allometry.stem_mass,
+... })
               name    dbh  n_individuals  stem_height  crown_area  stem_mass
 0   Evergreen Tree  0.100            100     9.890399    2.459835   8.156296
 1  Deciduous Shrub  0.030            200     2.110534    0.174049   0.134266
@@ -117,9 +122,8 @@
 from numpy.typing import NDArray
 
 from pyrealm.core.utilities import check_input_shapes
-from pyrealm.demography import canopy_functions
-from pyrealm.demography import t_model_functions as t_model
 from pyrealm.demography.flora import Flora, StemTraits
+from pyrealm.demography.t_model_functions import StemAllometry
 
 if sys.version_info[:2] >= (3, 11):
     import tomllib
@@ -356,6 +360,7 @@ class Community:
     # Post init properties
     number_of_cohorts: int = field(init=False)
     stem_traits: StemTraits = field(init=False)
+    stem_allometry: StemAllometry = field(init=False)
     cohort_data: dict[str, NDArray] = field(init=False)
 
     def __post_init__(
@@ -405,66 +410,9 @@ def __post_init__(
 
         self.number_of_cohorts = len(cohort_pft_names)
 
-        # Populate the T model fields
-        self._calculate_t_model()
-
-    def _calculate_t_model(self) -> None:
-        """Calculate T Model predictions across cohort data.
-
-        This method populates or updates the community attributes predicted by the T
-        Model :cite:`Li:2014bc` and by the canopy shape extensions to the T Model
-        implemented in PlantFate :cite:`joshi:2022a`.
-        """
-
-        # Add data to cohort dataframes capturing the T Model geometry
-        # - Classic T Model scaling
-        self.cohort_data["stem_height"] = t_model.calculate_heights(
-            h_max=self.stem_traits.h_max,
-            a_hd=self.stem_traits.a_hd,
-            dbh=self.cohort_data["dbh"],
-        )
-
-        self.cohort_data["crown_area"] = t_model.calculate_crown_areas(
-            ca_ratio=self.stem_traits.ca_ratio,
-            a_hd=self.stem_traits.a_hd,
-            dbh=self.cohort_data["dbh"],
-            stem_height=self.cohort_data["stem_height"],
-        )
-
-        self.cohort_data["crown_fraction"] = t_model.calculate_crown_fractions(
-            a_hd=self.stem_traits.a_hd,
-            dbh=self.cohort_data["dbh"],
-            stem_height=self.cohort_data["stem_height"],
-        )
-
-        self.cohort_data["stem_mass"] = t_model.calculate_stem_masses(
-            rho_s=self.stem_traits.rho_s,
-            dbh=self.cohort_data["dbh"],
-            stem_height=self.cohort_data["stem_height"],
-        )
-
-        self.cohort_data["foliage_mass"] = t_model.calculate_foliage_masses(
-            sla=self.stem_traits.sla,
-            lai=self.stem_traits.lai,
-            crown_area=self.cohort_data["crown_area"],
-        )
-
-        self.cohort_data["sapwood_mass"] = t_model.calculate_sapwood_masses(
-            rho_s=self.stem_traits.rho_s,
-            ca_ratio=self.stem_traits.ca_ratio,
-            stem_height=self.cohort_data["stem_height"],
-            crown_area=self.cohort_data["crown_area"],
-            crown_fraction=self.cohort_data["crown_fraction"],
-        )
-
-        # Canopy shape extension to T Model from PlantFATE
-        self.cohort_data["canopy_z_max"] = canopy_functions.calculate_canopy_z_max(
-            z_max_prop=self.stem_traits.z_max_prop,
-            stem_height=self.cohort_data["stem_height"],
-        )
-        self.cohort_data["canopy_r0"] = canopy_functions.calculate_canopy_r0(
-            q_m=self.stem_traits.q_m,
-            crown_area=self.cohort_data["crown_area"],
+        # Populate the stem allometry
+        self.stem_allometry = StemAllometry(
+            stem_traits=self.stem_traits, at_dbh=cohort_dbh_values
         )
 
     @classmethod

pyrealm/demography/t_model_functions.py

@@ -5,10 +5,18 @@
 calculate stem growth given net primary productivity.
 """  # noqa: D205
 
+from dataclasses import InitVar, dataclass, field
+from typing import ClassVar
+
 import numpy as np
 from numpy.typing import NDArray
 
 from pyrealm.core.utilities import check_input_shapes
+from pyrealm.demography.canopy_functions import (
+    calculate_canopy_r0,
+    calculate_canopy_z_max,
+)
+from pyrealm.demography.flora import Flora, StemTraits
 
 
 def _validate_t_model_args(pft_args: list[NDArray], size_args: list[NDArray]) -> None:
@@ -600,3 +608,119 @@ def calculate_growth_increments(
     delta_d = (npp - turnover) / (dWsdt + dWfdt)
 
     return (delta_d, dWsdt * delta_d, dWfdt * delta_d)
+
+
+@dataclass
+class StemAllometry:
+    """Calculate T Model allometric predictions across a set of stems.
+
+    This method calculate predictions of stem allometries for stem height, crown area,
+    crown fraction, stem mass, foliage mass and sapwood mass under the T Model
+    :cite:`Li:2014bc`, given diameters at breast height for a set of plant functional
+    traits.
+
+    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.
+        at_dbh: An array of diameter at breast height values at which to predict stem
+            allometry values.
+    """
+
+    allometry_attrs: ClassVar[tuple[str, ...]] = (
+        "dbh",
+        "stem_height",
+        "crown_area",
+        "crown_fraction",
+        "stem_mass",
+        "foliage_mass",
+        "sapwood_mass",
+        "canopy_r0",
+        "canopy_z_max",
+    )
+
+    # 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."""
+    at_dbh: InitVar[NDArray[np.float32]]
+    """An array of diameter at breast height values at which to predict stem allometry 
+    values."""
+
+    # Post init allometry attributes
+    dbh: NDArray[np.float32] = field(init=False)
+    """Diameter at breast height (metres)"""
+    stem_height: NDArray[np.float32] = field(init=False)
+    """Stem height (metres)"""
+    crown_area: NDArray[np.float32] = field(init=False)
+    """Crown area (square metres)"""
+    crown_fraction: NDArray[np.float32] = field(init=False)
+    """Vertical fraction of the stem covered by the crown (-)"""
+    stem_mass: NDArray[np.float32] = field(init=False)
+    """Stem mass (kg)"""
+    foliage_mass: NDArray[np.float32] = field(init=False)
+    """Foliage mass (kg)"""
+    sapwood_mass: NDArray[np.float32] = field(init=False)
+    """Sapwood mass (kg)"""
+    canopy_r0: NDArray[np.float32] = field(init=False)
+    """Canopy radius scaling factor (-)"""
+    canopy_z_max: NDArray[np.float32] = field(init=False)
+    """Height of maximum crown radius (metres)"""
+
+    def __post_init__(
+        self, stem_traits: Flora | StemTraits, at_dbh: NDArray[np.float32]
+    ) -> None:
+        """Populate the stem allometry attributes from the traits and size data."""
+
+        self.stem_height = calculate_heights(
+            h_max=stem_traits.h_max,
+            a_hd=stem_traits.a_hd,
+            dbh=at_dbh,
+        )
+
+        # Broadcast at_dbh to shape of stem height to get congruent shapes
+        self.dbh = np.broadcast_to(at_dbh, self.stem_height.shape)
+
+        self.crown_area = calculate_crown_areas(
+            ca_ratio=stem_traits.ca_ratio,
+            a_hd=stem_traits.a_hd,
+            dbh=self.dbh,
+            stem_height=self.stem_height,
+        )
+
+        self.crown_fraction = calculate_crown_fractions(
+            a_hd=stem_traits.a_hd,
+            dbh=self.dbh,
+            stem_height=self.stem_height,
+        )
+
+        self.stem_mass = calculate_stem_masses(
+            rho_s=stem_traits.rho_s,
+            dbh=self.dbh,
+            stem_height=self.stem_height,
+        )
+
+        self.foliage_mass = calculate_foliage_masses(
+            sla=stem_traits.sla,
+            lai=stem_traits.lai,
+            crown_area=self.crown_area,
+        )
+
+        self.sapwood_mass = calculate_sapwood_masses(
+            rho_s=stem_traits.rho_s,
+            ca_ratio=stem_traits.ca_ratio,
+            stem_height=self.stem_height,
+            crown_area=self.crown_area,
+            crown_fraction=self.crown_fraction,
+        )
+
+        self.canopy_r0 = calculate_canopy_r0(
+            q_m=stem_traits.q_m,
+            crown_area=self.crown_area,
+        )
+
+        self.canopy_z_max = calculate_canopy_z_max(
+            z_max_prop=stem_traits.z_max_prop,
+            stem_height=self.stem_height,
+        )

tests/unit/demography/test_canopy.py

@@ -38,7 +38,7 @@ def test_Canopy__init__():
         np.ceil(
             (
                 (
-                    community.cohort_data["crown_area"]
+                    community.stem_allometry.crown_area
                     * community.cohort_data["n_individuals"]
                 ).sum()
                 * (1 + canopy_gap_fraction)