From 1fe164bfc43445f5f3a805607b7ca6c70c461196 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 15:31:11 +0100
Subject: [PATCH 1/9] Adding new data classes within Canopy

---
 pyrealm/demography/canopy.py | 203 ++++++++++++++++++++++++-----------
 1 file changed, 142 insertions(+), 61 deletions(-)

diff --git a/pyrealm/demography/canopy.py b/pyrealm/demography/canopy.py
index 11614cea..053c22c6 100644
--- a/pyrealm/demography/canopy.py
+++ b/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,125 @@ def fit_perfect_plasticity_approximation(
     return layer_heights[:, None]
 
 
+@dataclass
+class CohortCanopyData:
+    """TBD.
+
+    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.
+    """
+
+    # 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:
+        """TBD."""
+        # 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:
+        """TBD."""
+
+        # 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:
+    """TBD."""
+
+    # 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 by layer."""
+    f_abs: NDArray[np.float64] = field(init=False)
+    """The fraction of light absorbed by the whole community by 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 by layer."""
+
+    def __post_init__(self, cohort_transmissivity: NDArray[np.float64]) -> None:
+        """TBD."""
+
+        # 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 +335,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 +385,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,
+        )

From bac13d91a2814850545f7ca3ec725e76e5b4f0f7 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 15:31:37 +0100
Subject: [PATCH 2/9] Simple test of CohortCanopyData class

---
 tests/unit/demography/test_canopy.py | 46 +++++++++++++++++++++++++++-
 1 file changed, 45 insertions(+), 1 deletion(-)

diff --git a/tests/unit/demography/test_canopy.py b/tests/unit/demography/test_canopy.py
index 95e04bf2..2f42817c 100644
--- a/tests/unit/demography/test_canopy.py
+++ b/tests/unit/demography/test_canopy.py
@@ -4,6 +4,50 @@
 import pytest
 
 
+@pytest.mark.parametrize(
+    argnames="args,expected",
+    argvalues=(
+        [
+            pytest.param(
+                {
+                    "projected_leaf_area": np.array([[2, 2, 2]]),
+                    "n_individuals": np.array([2, 2, 2]),
+                    "pft_lai": np.array([2, 2, 2]),
+                    "pft_par_ext": np.array([0.5, 0.5, 0.5]),
+                    "cell_area": 8,
+                },
+                (np.full((3,), 2), np.full((3,), 1), np.full((3,), np.exp(-0.5))),
+                id="single layer",
+            ),
+            pytest.param(
+                {
+                    "projected_leaf_area": np.tile([[2], [4], [6]], 3),
+                    "n_individuals": np.array([2, 2, 2]),
+                    "pft_lai": np.array([2, 2, 2]),
+                    "pft_par_ext": np.array([0.5, 0.5, 0.5]),
+                    "cell_area": 8,
+                },
+                (np.full((3, 3), 2), np.full((3, 3), 1), np.full((3, 3), np.exp(-0.5))),
+                id="two layers",
+            ),
+        ]
+    ),
+)
+def test_CohortCanopyData__init__(args, expected):
+    """Test creation of the cohort canopy data."""
+
+    from pyrealm.demography.canopy import CohortCanopyData
+
+    # Calculate canopy components
+    instance = CohortCanopyData(**args)
+
+    # Unpack and test expectations
+    exp_stem_leaf_area, exp_lai, exp_f_trans = expected
+    assert np.allclose(instance.stem_leaf_area, exp_stem_leaf_area)
+    assert np.allclose(instance.lai, exp_lai)
+    assert np.allclose(instance.f_trans, exp_f_trans)
+
+
 def test_Canopy__init__():
     """Test happy path for initialisation.
 
@@ -49,7 +93,7 @@ def test_Canopy__init__():
             / community.cell_area
         )
     )
-    assert canopy.stem_leaf_area.shape == (
+    assert canopy.cohort_data.stem_leaf_area.shape == (
         n_layers_from_crown_area,
         canopy.n_cohorts,
     )

From 74b9eed71518c483ee35cd015782e0919d4ac51c Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 16:05:41 +0100
Subject: [PATCH 3/9] Updating cabopy class testing

---
 tests/unit/demography/test_canopy.py | 62 ++++++++++++++++++++++------
 1 file changed, 50 insertions(+), 12 deletions(-)

diff --git a/tests/unit/demography/test_canopy.py b/tests/unit/demography/test_canopy.py
index 2f42817c..dfa42bd5 100644
--- a/tests/unit/demography/test_canopy.py
+++ b/tests/unit/demography/test_canopy.py
@@ -5,7 +5,7 @@
 
 
 @pytest.mark.parametrize(
-    argnames="args,expected",
+    argnames="cohort_args, cohort_expected, community_expected",
     argvalues=(
         [
             pytest.param(
@@ -17,6 +17,10 @@
                     "cell_area": 8,
                 },
                 (np.full((3,), 2), np.full((3,), 1), np.full((3,), np.exp(-0.5))),
+                (
+                    np.full((1,), np.exp(-0.5)) ** 3,
+                    np.full((1,), np.exp(-0.5)) ** 3,
+                ),
                 id="single layer",
             ),
             pytest.param(
@@ -28,24 +32,58 @@
                     "cell_area": 8,
                 },
                 (np.full((3, 3), 2), np.full((3, 3), 1), np.full((3, 3), np.exp(-0.5))),
-                id="two layers",
+                (
+                    np.full((3,), np.exp(-0.5)) ** 3,
+                    np.power(np.exp(-0.5), np.array([3, 6, 9])),
+                ),
+                id="three layers",
             ),
         ]
     ),
 )
-def test_CohortCanopyData__init__(args, expected):
-    """Test creation of the cohort canopy data."""
+class TestCanopyData:
+    """Shared testing of the cohort and community canopy dataclasses.
+
+    Simple cohort tests:
+    - LAI = (2 leaf area * 2 individuals * 2 LAI) / 8 area = 1
+    - trans = e ^ {-k L}, and since L = 1, trans = e^{-k}
+
+    Simple community tests
+    - Three identical cohorts so community trans = (e{-k})^3 for each layer
+    - Transmission profile (e{-k})^3, e{-k})^6, e{-k})^9)
+    """
+
+    def test_CohortCanopyData__init__(
+        self, cohort_args, cohort_expected, community_expected
+    ):
+        """Test creation of the cohort canopy data."""
+
+        from pyrealm.demography.canopy import CohortCanopyData
+
+        # Calculate canopy components
+        instance = CohortCanopyData(**cohort_args)
+
+        # Unpack and test expectations
+        exp_stem_leaf_area, exp_lai, exp_f_trans = cohort_expected
+        assert np.allclose(instance.stem_leaf_area, exp_stem_leaf_area)
+        assert np.allclose(instance.lai, exp_lai)
+        assert np.allclose(instance.f_trans, exp_f_trans)
+
+    def test_CommunityCanopyData__init__(
+        self, cohort_args, cohort_expected, community_expected
+    ):
+        """Test creation of the community canopy data."""
+
+        from pyrealm.demography.canopy import CohortCanopyData, CommunityCanopyData
 
-    from pyrealm.demography.canopy import CohortCanopyData
+        cohort_data = CohortCanopyData(**cohort_args)
 
-    # Calculate canopy components
-    instance = CohortCanopyData(**args)
+        instance = CommunityCanopyData(cohort_transmissivity=cohort_data.f_trans)
 
-    # Unpack and test expectations
-    exp_stem_leaf_area, exp_lai, exp_f_trans = expected
-    assert np.allclose(instance.stem_leaf_area, exp_stem_leaf_area)
-    assert np.allclose(instance.lai, exp_lai)
-    assert np.allclose(instance.f_trans, exp_f_trans)
+        # Unpack and test expectations
+        exp_f_trans, exp_trans_prof = community_expected
+        assert np.allclose(instance.f_trans, exp_f_trans)
+        assert np.allclose(instance.transmission_profile, exp_trans_prof)
 
 
 def test_Canopy__init__():

From d55f553f7a4ae416fc37fb19e3db50287cd8da94 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 17:40:23 +0100
Subject: [PATCH 4/9] Adding simple test of CohortCanopyData.allocate_fapar

---
 tests/unit/demography/test_canopy.py | 29 ++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/unit/demography/test_canopy.py b/tests/unit/demography/test_canopy.py
index dfa42bd5..cbfbbefc 100644
--- a/tests/unit/demography/test_canopy.py
+++ b/tests/unit/demography/test_canopy.py
@@ -51,6 +51,10 @@ class TestCanopyData:
     Simple community tests
     - Three identical cohorts so community trans = (e{-k})^3 for each layer
     - Transmission profile (e{-k})^3, e{-k})^6, e{-k})^9)
+
+    Allocate fapar
+     - share fapar equally among 3 cohorts and then equally between the two stems in
+       each cohort.
     """
 
     def test_CohortCanopyData__init__(
@@ -85,6 +89,31 @@ def test_CommunityCanopyData__init__(
         assert np.allclose(instance.f_trans, exp_f_trans)
         assert np.allclose(instance.transmission_profile, exp_trans_prof)
 
+    def test_CohortCanopyData_allocate_fapar(
+        self, cohort_args, cohort_expected, community_expected
+    ):
+        """Test creation of the cohort canopy data."""
+
+        from pyrealm.demography.canopy import CohortCanopyData, CommunityCanopyData
+
+        cohort_data = CohortCanopyData(**cohort_args)
+        community_data = CommunityCanopyData(cohort_transmissivity=cohort_data.f_trans)
+
+        cohort_data.allocate_fapar(
+            community_fapar=community_data.fapar,
+            n_individuals=cohort_args["n_individuals"],
+        )
+
+        # Unpack and test expectations
+        exp_f_trans, exp_trans_prof = community_expected
+        expected_fapar = -np.diff(exp_trans_prof, prepend=1)
+        assert np.allclose(
+            cohort_data.cohort_fapar, np.tile((expected_fapar / 3)[:, None], 3)
+        )
+        assert np.allclose(
+            cohort_data.stem_fapar, np.tile((expected_fapar / 6)[:, None], 3)
+        )
+
 
 def test_Canopy__init__():
     """Test happy path for initialisation.

From d2857ce935fbc2c0df7542acce41d02014c92cb6 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 18:23:51 +0100
Subject: [PATCH 5/9] Doc fixes

---
 docs/source/users/demography/canopy.md | 34 +++++++++++++-------------
 docs/source/users/demography/crown.md  |  4 ---
 2 files changed, 17 insertions(+), 21 deletions(-)

diff --git a/docs/source/users/demography/canopy.md b/docs/source/users/demography/canopy.md
index da837e9c..8e0e55b7 100644
--- a/docs/source/users/demography/canopy.md
+++ b/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)
 ```
 
diff --git a/docs/source/users/demography/crown.md b/docs/source/users/demography/crown.md
index 69e64018..7e6b39ba 100644
--- a/docs/source/users/demography/crown.md
+++ b/docs/source/users/demography/crown.md
@@ -532,7 +532,3 @@ plt.legend(
     bbox_to_anchor=(0.5, 1.15),
 )
 ```
-
-```{code-cell} ipython3
-
-```

From ecb36a4899c70c22778de469d7a7c617b3dd9d71 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Wed, 23 Oct 2024 10:52:28 +0100
Subject: [PATCH 6/9] Forgot to fill in docstrings

---
 pyrealm/demography/canopy.py | 67 ++++++++++++++++++++++++++++++------
 1 file changed, 57 insertions(+), 10 deletions(-)

diff --git a/pyrealm/demography/canopy.py b/pyrealm/demography/canopy.py
index 053c22c6..5756214b 100644
--- a/pyrealm/demography/canopy.py
+++ b/pyrealm/demography/canopy.py
@@ -167,10 +167,29 @@ def fit_perfect_plasticity_approximation(
 
 @dataclass
 class CohortCanopyData:
-    """TBD.
+    """Dataclass holding canopy data across cohorts.
 
-    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.
+    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
@@ -208,7 +227,7 @@ def __post_init__(
         pft_par_ext: NDArray[np.float64],
         cell_area: float,
     ) -> None:
-        """TBD."""
+        """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)
@@ -228,7 +247,19 @@ def __post_init__(
     def allocate_fapar(
         self, community_fapar: NDArray[np.float64], n_individuals: NDArray[np.int_]
     ) -> None:
-        """TBD."""
+        """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
@@ -249,7 +280,23 @@ def allocate_fapar(
 
 @dataclass
 class CommunityCanopyData:
-    """TBD."""
+    """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]]
@@ -258,18 +305,18 @@ class CommunityCanopyData:
 
     # Calculated variables
     f_trans: NDArray[np.float64] = field(init=False)
-    """The fraction of light transmitted by the whole community by layer."""
+    """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 by layer."""
+    """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 by layer."""
+    """The fraction of absorbed radiation for the whole community across layers."""
 
     def __post_init__(self, cohort_transmissivity: NDArray[np.float64]) -> None:
-        """TBD."""
+        """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)

From bf33706a142614cdb68815068fb554e372abfe41 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Wed, 23 Oct 2024 11:09:01 +0100
Subject: [PATCH 7/9] Fixing RST error

---
 pyrealm/demography/canopy.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyrealm/demography/canopy.py b/pyrealm/demography/canopy.py
index 5756214b..40181ea0 100644
--- a/pyrealm/demography/canopy.py
+++ b/pyrealm/demography/canopy.py
@@ -182,7 +182,7 @@ class CohortCanopyData:
     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.
+            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

From 0ba5f1144c15520213648243196cd85e14021d3e Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Wed, 23 Oct 2024 16:05:20 +0100
Subject: [PATCH 8/9] Updating calculation structure in cohort and community
 canopy data

---
 pyrealm/demography/canopy.py         | 59 +++++++++++++---------------
 tests/unit/demography/test_canopy.py | 35 +++++------------
 2 files changed, 38 insertions(+), 56 deletions(-)

diff --git a/pyrealm/demography/canopy.py b/pyrealm/demography/canopy.py
index 40181ea0..94650704 100644
--- a/pyrealm/demography/canopy.py
+++ b/pyrealm/demography/canopy.py
@@ -1,5 +1,7 @@
 """Functionality for canopy modelling."""
 
+from __future__ import annotations
+
 from dataclasses import InitVar, dataclass, field
 
 import numpy as np
@@ -174,10 +176,24 @@ class CohortCanopyData:
     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.
+    The data class:
+
+    1. 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 absorption and transmission fractions of each cohort
+       at each level.
+
+    2. Calculates the community-wide transmission and absorption profiles. These are
+       generated as an instance of the class
+       :class:`~pyrealm.demography.canopy.CommunityCanopyData` and stored in the
+       ``community_data`` attribute.
+
+    3. Allocates the 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:
         projected_leaf_area: A two dimensional array providing projected leaf area for a
@@ -219,6 +235,9 @@ class CohortCanopyData:
     stem_fapar: NDArray[np.float64] = field(init=False)
     """The fraction of absorbed radiation for each stem by layer."""
 
+    # Community wide attributes in their own class
+    community_data: CommunityCanopyData = field(init=False)
+
     def __post_init__(
         self,
         projected_leaf_area: NDArray[np.float64],
@@ -244,22 +263,8 @@ def __post_init__(
         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 community wide properties
+        self.community_data = CommunityCanopyData(cohort_transmissivity=self.f_trans)
 
         # Calculate the fapar profile across cohorts and layers
         # * The first part of the equation is calculating the relative absorption of
@@ -273,7 +278,7 @@ def allocate_fapar(
         # f_abs values
         self.cohort_fapar = (
             self.f_abs / self.f_abs.sum(axis=1)[:, None]
-        ) * community_fapar[:, None]
+        ) * self.community_data.fapar[:, None]
         # Partition cohort f_APAR between the number of stems
         self.stem_fapar = self.cohort_fapar / n_individuals
 
@@ -442,13 +447,5 @@ def _calculate_canopy(self, community: Community) -> None:
             cell_area=community.cell_area,
         )
 
-        # Calculate the community wide canopy componennts at each layer_height
-        self.community_data = CommunityCanopyData(
-            cohort_transmissivity=self.cohort_data.f_trans
-        )
-
-        # 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,
-        )
+        # Create a shorter reference to the community data
+        self.community_data = self.cohort_data.community_data
diff --git a/tests/unit/demography/test_canopy.py b/tests/unit/demography/test_canopy.py
index cbfbbefc..b8b03105 100644
--- a/tests/unit/demography/test_canopy.py
+++ b/tests/unit/demography/test_canopy.py
@@ -73,6 +73,16 @@ def test_CohortCanopyData__init__(
         assert np.allclose(instance.lai, exp_lai)
         assert np.allclose(instance.f_trans, exp_f_trans)
 
+        # Unpack and test expectations
+        exp_f_trans, exp_trans_prof = community_expected
+        expected_fapar = -np.diff(exp_trans_prof, prepend=1)
+        assert np.allclose(
+            instance.cohort_fapar, np.tile((expected_fapar / 3)[:, None], 3)
+        )
+        assert np.allclose(
+            instance.stem_fapar, np.tile((expected_fapar / 6)[:, None], 3)
+        )
+
     def test_CommunityCanopyData__init__(
         self, cohort_args, cohort_expected, community_expected
     ):
@@ -89,31 +99,6 @@ def test_CommunityCanopyData__init__(
         assert np.allclose(instance.f_trans, exp_f_trans)
         assert np.allclose(instance.transmission_profile, exp_trans_prof)
 
-    def test_CohortCanopyData_allocate_fapar(
-        self, cohort_args, cohort_expected, community_expected
-    ):
-        """Test creation of the cohort canopy data."""
-
-        from pyrealm.demography.canopy import CohortCanopyData, CommunityCanopyData
-
-        cohort_data = CohortCanopyData(**cohort_args)
-        community_data = CommunityCanopyData(cohort_transmissivity=cohort_data.f_trans)
-
-        cohort_data.allocate_fapar(
-            community_fapar=community_data.fapar,
-            n_individuals=cohort_args["n_individuals"],
-        )
-
-        # Unpack and test expectations
-        exp_f_trans, exp_trans_prof = community_expected
-        expected_fapar = -np.diff(exp_trans_prof, prepend=1)
-        assert np.allclose(
-            cohort_data.cohort_fapar, np.tile((expected_fapar / 3)[:, None], 3)
-        )
-        assert np.allclose(
-            cohort_data.stem_fapar, np.tile((expected_fapar / 6)[:, None], 3)
-        )
-
 
 def test_Canopy__init__():
     """Test happy path for initialisation.

From 34509eb073e18d6f43075c5644543d719577d102 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Wed, 23 Oct 2024 16:11:57 +0100
Subject: [PATCH 9/9] Spaces in RST indents again

---
 pyrealm/demography/canopy.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pyrealm/demography/canopy.py b/pyrealm/demography/canopy.py
index 94650704..60a99109 100644
--- a/pyrealm/demography/canopy.py
+++ b/pyrealm/demography/canopy.py
@@ -190,10 +190,10 @@ class CohortCanopyData:
 
     3. Allocates the 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.
+       - 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:
         projected_leaf_area: A two dimensional array providing projected leaf area for a