Group internal attributes in canopy into data classes

docs/source/users/demography/canopy.md

@@ -139,7 +139,7 @@ should equal the whole canopy $f_{abs}$ calculated using the simple Beer-Lambert
 equation and the PFT trait values.
 
 ```{code-cell} ipython3
-print(simple_canopy.extinction_profile[-1])
+print(simple_canopy.community_data.extinction_profile[-1])
 ```
 
 ```{code-cell} ipython3
@@ -182,15 +182,15 @@ ax1.set_xlabel("Profile radius (m)")
 ax1.set_ylabel("Vertical height (m)")
 
 # Plot the leaf area between heights for stems
-ax2.plot(simple_canopy.stem_leaf_area, hghts, color="red")
+ax2.plot(simple_canopy.cohort_data.stem_leaf_area, hghts, color="red")
 ax2.set_xlabel("Leaf area (m2)")
 
 # Plot the fraction of light absorbed at different heights
-ax3.plot(simple_canopy.f_abs, hghts, color="red")
+ax3.plot(simple_canopy.cohort_data.f_abs, hghts, color="red")
 ax3.set_xlabel("Light absorption fraction (-)")
 
 # Plot the light extinction profile through the canopy.
-ax4.plot(simple_canopy.extinction_profile, hghts, color="red")
+ax4.plot(simple_canopy.community_data.extinction_profile, hghts, color="red")
 ax4.set_xlabel("Cumulative light\nabsorption fraction (-)")
 ```
 
@@ -312,7 +312,7 @@ print(1 - np.exp(np.sum(-community.stem_traits.par_ext * cohort_lai)))
 ```
 
 ```{code-cell} ipython3
-print(canopy.extinction_profile[-1])
+print(canopy.community_data.extinction_profile[-1])
 ```
 
 ```{code-cell} ipython3
@@ -349,16 +349,16 @@ ax1.set_xlabel("Profile radius (m)")
 ax1.set_ylabel("Vertical height (m)")
 
 # Plot the leaf area between heights for stems
-ax2.plot(canopy.stem_leaf_area, hghts)
+ax2.plot(canopy.cohort_data.stem_leaf_area, hghts)
 ax2.set_xlabel("Leaf area per stem (m2)")
 
 # Plot the fraction of light absorbed at different heights
-ax3.plot(canopy.f_abs, hghts, color="grey")
-ax3.plot(1 - canopy.cohort_f_trans, hghts)
+ax3.plot(canopy.cohort_data.f_abs, hghts, color="grey")
+ax3.plot(1 - canopy.cohort_data.f_trans, hghts)
 ax3.set_xlabel("Light absorption fraction (-)")
 
 # Plot the light extinction profile through the canopy.
-ax4.plot(canopy.extinction_profile, hghts, color="grey")
+ax4.plot(canopy.community_data.extinction_profile, hghts, color="grey")
 _ = ax4.set_xlabel("Cumulative light\nabsorption fraction (-)")
 ```
 
@@ -416,7 +416,7 @@ l_m = \left\lceil \frac{\sum_1^{N_s}{A_c}}{ A(1 - f_G)}\right\rceil
 $$
 
 ```{code-cell} ipython3
-canopy_ppa = Canopy(community=community, canopy_gap_fraction=2 / 32, fit_ppa=True)
+canopy_ppa = Canopy(community=community, canopy_gap_fraction=0 / 32, fit_ppa=True)
 ```
 
 The `canopy_ppa.heights` attribute now contains the heights at which the PPA
@@ -430,7 +430,7 @@ And the final value in the canopy extinction profile still matches the expectati
 above:
 
 ```{code-cell} ipython3
-print(canopy_ppa.extinction_profile[-1])
+print(canopy_ppa.community_data.extinction_profile[-1])
 ```
 
 ### Visualizing layer closure heights and areas
@@ -503,18 +503,18 @@ ax1.legend(frameon=False)
 for val in canopy_ppa.heights:
     ax2.axhline(val, color="red", linewidth=0.5, zorder=0)
 
-for val in canopy_ppa.extinction_profile:
+for val in canopy_ppa.community_data.extinction_profile:
     ax2.axvline(val, color="red", linewidth=0.5, zorder=0)
 
-ax2.plot(canopy.extinction_profile, hghts)
+ax2.plot(canopy.community_data.extinction_profile, hghts)
 
 ax2_top = ax2.twiny()
 ax2_top.set_xlim(ax2.get_xlim())
 extinction_labels = [
     f"$f_{{abs{l + 1}}}$ = {z:.3f}"
-    for l, z in enumerate(np.nditer(canopy_ppa.extinction_profile))
+    for l, z in enumerate(np.nditer(canopy_ppa.community_data.extinction_profile))
 ]
-ax2_top.set_xticks(canopy_ppa.extinction_profile)
+ax2_top.set_xticks(canopy_ppa.community_data.extinction_profile)
 ax2_top.set_xticklabels(extinction_labels, rotation=90)
 
 ax2.set_xlabel("Light extinction (-)")
@@ -548,11 +548,11 @@ The steps below show this partitioning process for the PPA layers calculated abo
 
 ```{code-cell} ipython3
 print("k = \n", community.stem_traits.par_ext, "\n")
-print("L_H = \n", canopy_ppa.cohort_lai)
+print("L_H = \n", canopy_ppa.cohort_data.lai)
 ```
 
 ```{code-cell} ipython3
-layer_cohort_f_tr = np.exp(-community.stem_traits.par_ext * canopy_ppa.cohort_lai)
+layer_cohort_f_tr = np.exp(-community.stem_traits.par_ext * canopy_ppa.cohort_data.lai)
 print(layer_cohort_f_tr)
 ```
 

docs/source/users/demography/crown.md

@@ -532,7 +532,3 @@ plt.legend(
     bbox_to_anchor=(0.5, 1.15),
 )
 ```
-
-```{code-cell} ipython3
-
-```

pyrealm/demography/canopy.py

@@ -1,5 +1,7 @@
 """Functionality for canopy modelling."""
 
+from dataclasses import InitVar, dataclass, field
+
 import numpy as np
 from numpy.typing import NDArray
 from scipy.optimize import root_scalar  # type: ignore [import-untyped]
@@ -163,6 +165,172 @@ def fit_perfect_plasticity_approximation(
     return layer_heights[:, None]
 
 
+@dataclass
+class CohortCanopyData:
+    """Dataclass holding canopy data across cohorts.
+
+    The cohort canopy data consists of a set of attributes represented as two
+    dimensional arrays. Each row is different height at which canopy properties are
+    required and the columns represent the different cohorts or the identical stem
+    properties of individuals within cohorts.
+
+    The data class takes the projected leaf area at the required heights and then
+    partitions this into the actual leaf area within each layer, the leaf area index
+    across the whole cohort and then then light absorbtion and transmission fractions of
+    each cohort at each level.
+
+    Args:
+        projected_leaf_area: A two dimensional array providing projected leaf area for a
+            set of cohorts (columns) at a set of required heights (rows), as for example
+            calculated using the :class:`~pyrealm.demography.crown.CrownProfile` class.
+        n_individuals: A one-dimensional array of the number of individuals in each
+            cohort.
+        pft_lai: A one-dimensional array giving the leaf area index trait for the plant
+            functional type of each cohort.
+        pft_par_ext: A one-dimensional array giving the light extinction coefficient for
+            the plant functional type of each cohort.
+        cell_area: A float setting the total canopy area available to the cohorts.
+    """
+
+    # Init vars
+    projected_leaf_area: InitVar[NDArray[np.float64]]
+    """An array of the stem projected leaf area for each cohort at each of the required
+    heights."""
+    n_individuals: InitVar[NDArray[np.int_]]
+    """The number of individuals for each cohort."""
+    pft_lai: InitVar[NDArray[np.float64]]
+    """The leaf area index of the plant functional type for each cohort."""
+    pft_par_ext: InitVar[NDArray[np.float64]]
+    """The extinction coefficient of the plant functional type for each cohort."""
+    cell_area: InitVar[float]
+    """The area available to the community."""
+
+    # Computed variables
+    stem_leaf_area: NDArray[np.float64] = field(init=False)
+    """The leaf area of the crown model for each cohort by layer."""
+    lai: NDArray[np.float64] = field(init=False)
+    """The leaf area index for each cohort by layer."""
+    f_trans: NDArray[np.float64] = field(init=False)
+    """The fraction of light transmitted by each cohort by layer."""
+    f_abs: NDArray[np.float64] = field(init=False)
+    """The fraction of light absorbed by each cohort by layer."""
+    cohort_fapar: NDArray[np.float64] = field(init=False)
+    """The fraction of absorbed radiation for each cohort by layer."""
+    stem_fapar: NDArray[np.float64] = field(init=False)
+    """The fraction of absorbed radiation for each stem by layer."""
+
+    def __post_init__(
+        self,
+        projected_leaf_area: NDArray[np.float64],
+        n_individuals: NDArray[np.int_],
+        pft_lai: NDArray[np.float64],
+        pft_par_ext: NDArray[np.float64],
+        cell_area: float,
+    ) -> None:
+        """Calculates cohort canopy attributes from the input data."""
+        # Partition the projected leaf area into the leaf area in each layer for each
+        # stem and then scale up to the cohort leaf area in each layer.
+        self.stem_leaf_area = np.diff(projected_leaf_area, axis=0, prepend=0)
+
+        # Calculate the leaf area index per layer per stem, using the stem
+        # specific leaf area index values. LAI is a value per m2, so scale back down by
+        # the available community area.
+        self.lai = (
+            self.stem_leaf_area * n_individuals * pft_lai
+        ) / cell_area  # self.filled_community_area
+
+        # Calculate the Beer-Lambert light transmission and absorption components per
+        # layer and cohort
+        self.f_trans = np.exp(-pft_par_ext * self.lai)
+        self.f_abs = 1 - self.f_trans
+
+    def allocate_fapar(
+        self, community_fapar: NDArray[np.float64], n_individuals: NDArray[np.int_]
+    ) -> None:
+        """Allocate community-wide absorption across cohorts.
+
+        The total fraction of light absorbed across layers is a community-wide property
+        - each cohort contributes to the cumulative light absorption. Once the light
+        absorbed within a layer of the community is known, this can then be partitioned
+        back to cohorts and individual stems to give the fraction of canopy top
+        radiation intercepted by each stem within each layer.
+
+        Args:
+            community_fapar: The community wide fraction of light absorbed across all
+                layers and cohorts.
+            n_individuals: The number of individuals within each cohort.
+        """
+
+        # Calculate the fapar profile across cohorts and layers
+        # * The first part of the equation is calculating the relative absorption of
+        #   each cohort within each layer
+        # * Each layer is then multiplied by fraction of the total light absorbed in the
+        #   layer
+        # * The resulting matrix can be multiplied by a canopy top PPFD to generate the
+        #   flux absorbed within each layer for each cohort.
+
+        # Divide the community wide f_APAR among the cohorts, based on their relative
+        # f_abs values
+        self.cohort_fapar = (
+            self.f_abs / self.f_abs.sum(axis=1)[:, None]
+        ) * community_fapar[:, None]
+        # Partition cohort f_APAR between the number of stems
+        self.stem_fapar = self.cohort_fapar / n_individuals
+
+
+@dataclass
+class CommunityCanopyData:
+    """Dataclass holding community-wide canopy data.
+
+    The community canopy data consists of a set of attributes represented as one
+    dimensional arrays, with each entry representing a different vertical height at
+    which canopy properties are required.
+
+    The data class takes the per cohort light transmission at the required heights and
+    calculates the aggregate transmission and absorption fractions within layers across
+    the whole community. It then calculates the cumulative extinction and transmission
+    profiles across layers and the hence the actual fraction of canopy top radiation
+    intercepted across layers (:math:`f_{APAR}`).
+
+    Args:
+        cohort_transmissivity: The per cohort light transmissivity across the required
+            heights, as calculated as
+            :attr:`CohortCanopyData.f_trans<pyrealm.demography.canopy.CohortCanopyData.f_trans>`.
+    """
+
+    # Init vars
+    cohort_transmissivity: InitVar[NDArray[np.float64]]
+    """An array providing the per cohort light transmissivity at each of the required
+    heights."""
+
+    # Calculated variables
+    f_trans: NDArray[np.float64] = field(init=False)
+    """The fraction of light transmitted by the whole community within a layer."""
+    f_abs: NDArray[np.float64] = field(init=False)
+    """The fraction of light absorbed by the whole community within a layer."""
+    transmission_profile: NDArray[np.float64] = field(init=False)
+    """The light transmission profile for the whole community by layer."""
+    extinction_profile: NDArray[np.float64] = field(init=False)
+    """The light extinction profile for the whole community by layer."""
+    fapar: NDArray[np.float64] = field(init=False)
+    """The fraction of absorbed radiation for the whole community across layers."""
+
+    def __post_init__(self, cohort_transmissivity: NDArray[np.float64]) -> None:
+        """Calculates community-wide canopy attributes from the input data."""
+
+        # Aggregate across cohorts into a layer wide transmissivity
+        self.f_trans = cohort_transmissivity.prod(axis=1)
+
+        # Calculate the canopy wide light extinction per layer
+        self.f_abs = 1 - self.f_trans
+
+        # Calculate cumulative light transmission and extinction profiles
+        self.transmission_profile = np.cumprod(self.f_trans)
+        self.extinction_profile = 1 - self.transmission_profile
+
+        self.fapar = -np.diff(self.transmission_profile, prepend=1)
+
+
 class Canopy:
     """Calculate canopy characteristics for a plant community.
 
@@ -214,28 +382,10 @@ def __init__(
 
         self.crown_profile: CrownProfile
         """The crown profiles of the community stems at the provided layer heights."""
-        self.stem_leaf_area: NDArray[np.float64]
-        """The leaf area of the crown model for each cohort by layer."""
-        self.cohort_lai: NDArray[np.float64]
-        """The leaf area index for each cohort by layer."""
-        self.cohort_f_trans: NDArray[np.float64]
-        """The fraction of light transmitted by each cohort by layer."""
-        self.cohort_f_abs: NDArray[np.float64]
-        """The fraction of light absorbed by each cohort by layer."""
-        self.f_trans: NDArray[np.float64]
-        """The fraction of light transmitted by the whole community by layer."""
-        self.f_abs: NDArray[np.float64]
-        """The fraction of light absorbed by the whole community by layer."""
-        self.transmission_profile: NDArray[np.float64]
-        """The light transmission profile for the whole community by layer."""
-        self.extinction_profile: NDArray[np.float64]
-        """The light extinction profile for the whole community by layer."""
-        self.fapar: NDArray[np.float64]
-        """The fraction of absorbed radiation for the whole community by layer."""
-        self.cohort_fapar: NDArray[np.float64]
-        """The fraction of absorbed radiation for each cohort by layer."""
-        self.stem_fapar: NDArray[np.float64]
-        """The fraction of absorbed radiation for each stem by layer."""
+        self.cohort_data: CohortCanopyData
+        """The per-cohort canopy data."""
+        self.community_data: CommunityCanopyData
+        """The community-wide canopy data."""
         self.filled_community_area: float
         """The area filled by crown after accounting for the crown gap fraction."""
 
@@ -282,45 +432,23 @@ def _calculate_canopy(self, community: Community) -> None:
             z=self.heights,
         )
 
-        # Partition the projected leaf area into the leaf area in each layer for each
-        # stem and then scale up to the cohort leaf area in each layer.
-        self.stem_leaf_area = np.diff(
-            self.crown_profile.projected_leaf_area, axis=0, prepend=0
+        # Calculate the per cohort canopy components (LAI, f_trans, f_abs) from the
+        # projected leaf area for each stem at the layer heights
+        self.cohort_data = CohortCanopyData(
+            projected_leaf_area=self.crown_profile.projected_leaf_area,
+            n_individuals=community.cohorts.n_individuals,
+            pft_lai=community.stem_traits.lai,
+            pft_par_ext=community.stem_traits.par_ext,
+            cell_area=community.cell_area,
         )
 
-        # Calculate the leaf area index per layer per stem, using the stem
-        # specific leaf area index values. LAI is a value per m2, so scale back down by
-        # the available community area.
-        self.cohort_lai = (
-            self.stem_leaf_area
-            * community.cohorts.n_individuals
-            * community.stem_traits.lai
-        ) / community.cell_area  # self.filled_community_area
-
-        # Calculate the Beer-Lambert light transmission and absorption components per
-        # layer and cohort
-        self.cohort_f_trans = np.exp(-community.stem_traits.par_ext * self.cohort_lai)
-        self.cohort_f_abs = 1 - self.cohort_f_trans
-
-        # Aggregate across cohorts into a layer wide transimissivity
-        self.f_trans = self.cohort_f_trans.prod(axis=1)
-
-        # Calculate the canopy wide light extinction per layer
-        self.f_abs = 1 - self.f_trans
-
-        # Calculate cumulative light transmission and extinction profiles
-        self.transmission_profile = np.cumprod(self.f_trans)
-        self.extinction_profile = 1 - self.transmission_profile
+        # Calculate the community wide canopy componennts at each layer_height
+        self.community_data = CommunityCanopyData(
+            cohort_transmissivity=self.cohort_data.f_trans
+        )
 
-        # Calculate the fapar profile across cohorts and layers
-        # * The first part of the equation is calculating the relative absorption of
-        #   each cohort within each layer
-        # * Each layer is then multiplied by fraction of the total light absorbed in the
-        #   layer
-        # * The resulting matrix can be multiplied by a canopy top PPFD to generate the
-        #   flux absorbed within each layer for each cohort.
-        self.fapar = -np.diff(self.transmission_profile, prepend=1)
-        self.cohort_fapar = (
-            self.cohort_f_abs / self.cohort_f_abs.sum(axis=1)[:, None]
-        ) * self.fapar[:, None]
-        self.stem_fapar = self.cohort_fapar / community.cohorts.n_individuals
+        # Calculate the proportion of absorbed light intercepted by each cohort and stem
+        self.cohort_data.allocate_fapar(
+            community_fapar=self.community_data.fapar,
+            n_individuals=community.cohorts.n_individuals,
+        )