From c8a228e54953a958dbc7f47ad5643841c53eb9f9 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Fri, 18 Oct 2024 17:30:18 +0100
Subject: [PATCH 01/31] Initial implementation of PandasExporter and Cohorts
 ABCs and test

---
 pyrealm/demography/core.py         |  66 ++++++++++++++++
 tests/unit/demography/test_core.py | 121 +++++++++++++++++++++++++++++
 2 files changed, 187 insertions(+)
 create mode 100644 pyrealm/demography/core.py
 create mode 100644 tests/unit/demography/test_core.py

diff --git a/pyrealm/demography/core.py b/pyrealm/demography/core.py
new file mode 100644
index 00000000..f628da79
--- /dev/null
+++ b/pyrealm/demography/core.py
@@ -0,0 +1,66 @@
+"""Core shared functionality for the {mod}`~pyrealm.demography` module."""
+
+from __future__ import annotations
+
+from abc import ABC
+from typing import ClassVar
+
+import numpy as np
+import pandas as pd
+from numpy.typing import NDArray
+
+
+class PandasExporter(ABC):
+    """Abstract base class implementing pandas export.
+
+    Classes inheriting from this ABC must define a class attribute ``array_attrs`` that
+    names a set of class attributes that are all numpy arrays of equal length. The
+    classes then inherit the `to_pandas` method that exports those attributes to a
+    {class}`pandas.DataFrame`.
+    """
+
+    array_attrs: ClassVar[tuple[str, ...]]
+
+    def to_pandas(self) -> pd.DataFrame:
+        """Convert the instance array attributes into a {class}`pandas.DataFrame."""
+        return pd.DataFrame({k: getattr(self, k) for k in self.array_attrs})
+
+
+class Cohorts(ABC):
+    """Abstract base class implementing cohort manipulation functionality.
+
+    Classes inheriting from this ABC must define a class attribute ``array_attrs`` that
+    names a set of class attributes that are all numpy arrays of equal length. The class
+    then inherit:
+
+    * The `add_cohorts` method, which allows a second instance of the same class to be
+      joined to the calling instance, concatenting each of the array attributes from
+      the second instance onto the calling instance.
+    * The `drop_cohorts` method, which takes a set of indices onto the array attributes
+      and drops the values from those indices for each array attribute.
+    """
+
+    array_attrs: ClassVar[tuple[str, ...]]
+
+    def add_cohorts(self, add: Cohorts) -> None:
+        """Add array attributes from a second instance.
+
+        Args:
+            add: A second instance from which to add array attribute values.
+        """
+        for trait in self.array_attrs:
+            setattr(
+                self,
+                trait,
+                np.concatenate([getattr(self, trait), getattr(add, trait)]),
+            )
+
+    def drop_cohorts(self, drop_indices: NDArray[np.int_]) -> None:
+        """Drop array attribute values from an instance.
+
+        Args:
+            drop_indices: An array of integer indices to drop from each array attribute.
+        """
+
+        for trait in self.array_attrs:
+            setattr(self, trait, np.delete(getattr(self, trait), drop_indices))
diff --git a/tests/unit/demography/test_core.py b/tests/unit/demography/test_core.py
new file mode 100644
index 00000000..890622e1
--- /dev/null
+++ b/tests/unit/demography/test_core.py
@@ -0,0 +1,121 @@
+"""Tests of core demography objects."""
+
+from __future__ import annotations
+
+from dataclasses import InitVar, dataclass, field
+from typing import ClassVar
+
+import numpy as np
+import pandas as pd
+from numpy.typing import NDArray
+
+
+def test_PandasExporter():
+    """Test the PandasExporter abstract base class."""
+
+    from pyrealm.demography.core import PandasExporter
+
+    @dataclass
+    class TestClass(PandasExporter):
+        """Simple test dataclass implementing the ABC."""
+
+        array_attrs: ClassVar[tuple[str, ...]] = ("c", "d", "e")
+        c: NDArray[np.float64]
+        d: NDArray[np.float64]
+        e: NDArray[np.float64]
+
+    # create instance and run method
+    instance = TestClass(c=np.arange(5), d=np.arange(5), e=np.arange(5))
+    pandas_out = instance.to_pandas()
+
+    # simple checks of output class and behaviour
+    assert isinstance(pandas_out, pd.DataFrame)
+    assert pandas_out.shape == (5, 3)
+    assert np.allclose(pandas_out.sum(axis=1), np.arange(5) * 3)
+    assert np.allclose(pandas_out.sum(axis=0), np.repeat(10, 3))
+
+
+def test_Cohorts():
+    """Test the Cohorts abstract base class."""
+
+    from pyrealm.demography.core import Cohorts
+
+    @dataclass
+    class TestClass(Cohorts):
+        """Simple test class implementing the ABC."""
+
+        array_attrs: ClassVar[tuple[str, ...]] = ("a", "b")
+
+        a: NDArray[np.float64]
+        b: NDArray[np.float64]
+
+    # Create instances
+    t1 = TestClass(a=np.array([1, 2, 3]), b=np.array([4, 5, 6]))
+    t2 = TestClass(a=np.array([4, 5, 6]), b=np.array([7, 8, 9]))
+
+    # Add the t2 data into t1 and check the a and b attributes are extended
+    t1.add_cohorts(t2)
+    assert np.allclose(t1.a, np.arange(1, 7))
+    assert np.allclose(t1.b, np.arange(4, 10))
+
+    # Drop some indices and check the a and b attributes are truncated
+    t1.drop_cohorts(np.array([0, 5]))
+    assert np.allclose(t1.a, np.arange(2, 6))
+    assert np.allclose(t1.b, np.arange(5, 9))
+
+
+def test_PandasExporter_Cohorts_multiple_inheritance():
+    """Test the behaviour of a class inheriting both core ABCs."""
+
+    from pyrealm.demography.core import Cohorts, PandasExporter
+
+    @dataclass
+    class TestClass(Cohorts, PandasExporter):
+        """Test class with multiple inheritance."""
+
+        array_attrs: ClassVar[tuple[str, ...]] = ("c", "d", "e")
+
+        n: InitVar[int]
+        start_vals: InitVar[NDArray[np.int_]]
+
+        c: NDArray[np.float64] = field(init=False)
+        d: NDArray[np.int_] = field(init=False)
+        e: NDArray[np.float64] = field(init=False)
+
+        def __post_init__(self, n: int, start_vals: NDArray[np.int_]) -> None:
+            self.c = np.arange(start_vals[0], start_vals[0] + n)
+            self.d = np.arange(start_vals[1], start_vals[1] + n)
+            self.e = np.arange(start_vals[2], start_vals[2] + n)
+
+    # Create instances
+    t1 = TestClass(n=5, start_vals=np.array([1, 2, 3]))
+    t2 = TestClass(n=3, start_vals=np.array([6, 7, 8]))
+
+    # Test dataframe properties
+    t1_out = t1.to_pandas()
+
+    # simple checks of output class and behaviour
+    assert isinstance(t1_out, pd.DataFrame)
+    assert t1_out.shape == (5, 3)
+    assert np.allclose(t1_out.sum(axis=1), np.array([6, 9, 12, 15, 18]))
+    assert np.allclose(t1_out.sum(axis=0), np.repeat(15, 3) + np.array([0, 5, 10]))
+
+    # Add the second set and check the results via pandas
+    t1.add_cohorts(t2)
+    t1_out_add = t1.to_pandas()
+
+    # simple checks of output class and behaviour
+    assert isinstance(t1_out_add, pd.DataFrame)
+    assert t1_out_add.shape == (8, 3)
+    assert np.allclose(t1_out_add.sum(axis=1), np.array([6, 9, 12, 15, 18, 21, 24, 27]))
+    assert np.allclose(t1_out_add.sum(axis=0), np.repeat(36, 3) + np.array([0, 8, 16]))
+
+    # Drop some entries and recheck
+    t1.drop_cohorts(np.array([0, 7]))
+    t1_out_drop = t1.to_pandas()
+
+    # simple checks of output class and behaviour
+    assert isinstance(t1_out_drop, pd.DataFrame)
+    assert t1_out_drop.shape == (6, 3)
+    assert np.allclose(t1_out_drop.sum(axis=1), np.array([9, 12, 15, 18, 21, 24]))
+    assert np.allclose(t1_out_drop.sum(axis=0), np.repeat(27, 3) + np.array([0, 6, 12]))

From b625b82da7af56c26fa580240ee3de4ac0e3984a Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Fri, 18 Oct 2024 18:29:01 +0100
Subject: [PATCH 02/31] Rolling out ABCs to flora classes

---
 docs/source/users/demography/flora.md | 14 +++++++++++---
 pyrealm/demography/flora.py           | 20 +++++++++++---------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/docs/source/users/demography/flora.md b/docs/source/users/demography/flora.md
index 8db0615d..c32729c1 100644
--- a/docs/source/users/demography/flora.md
+++ b/docs/source/users/demography/flora.md
@@ -146,8 +146,12 @@ flora = Flora([short_pft, medium_pft, tall_pft])
 flora
 ```
 
+The class provides a {meth}`~pyrealm.demography.StemTraits.to_pandas()` method to export
+the trait data as a {class}`pandas.DataFrame`, making it easier to use for plotting or
+calculations outside of `pyrealm`.
+
 ```{code-cell} ipython3
-pd.DataFrame({k: getattr(flora, k) for k in flora.trait_attrs})
+flora.to_pandas()
 ```
 
 You can also create `Flora` instances using PFT data stored TOML, JSON and CSV file
@@ -168,7 +172,11 @@ more easily created from a `Flora` object by providing a list of PFT names:
 # Get stem traits for a range of stems
 stem_pfts = ["short", "short", "short", "medium", "medium", "tall"]
 stem_traits = flora.get_stem_traits(pft_names=stem_pfts)
+```
 
-# Show the repeated values
-stem_traits.h_max
+Again, the class provides the {meth}`~pyrealm.demography.StemTraits.to_pandas()` method
+to extract the data:
+
+```{code-cell} ipython3
+stem_traits.to_pandas()
 ```
diff --git a/pyrealm/demography/flora.py b/pyrealm/demography/flora.py
index 41767d5f..3e2384a5 100644
--- a/pyrealm/demography/flora.py
+++ b/pyrealm/demography/flora.py
@@ -37,6 +37,8 @@
 from marshmallow.exceptions import ValidationError
 from numpy.typing import NDArray
 
+from pyrealm.demography.core import PandasExporter
+
 if sys.version_info[:2] >= (3, 11):
     import tomllib
     from tomllib import TOMLDecodeError
@@ -236,7 +238,7 @@ class PlantFunctionalType(PlantFunctionalTypeStrict):
 
 
 @dataclass(frozen=True)
-class Flora:
+class Flora(PandasExporter):
     """A dataclass providing trait data on collection of plant functional types.
 
     A flora provides trait data on the complete collection of plant functional types
@@ -258,10 +260,10 @@ class Flora:
     pfts: InitVar[Sequence[type[PlantFunctionalTypeStrict]]]
     r"""A sequence of plant functional type instances to include in the Flora."""
 
-    # A class variable setting the attribute names of traits.
-    trait_attrs: ClassVar[list[str]] = [
+    # A class variable setting the names of PFT traits held as arrays.
+    array_attrs: ClassVar[tuple[str, ...]] = tuple(
         f.name for f in fields(PlantFunctionalTypeStrict)
-    ]
+    )
 
     # Populated post init
     # - trait arrays
@@ -342,7 +344,7 @@ def __post_init__(self, pfts: Sequence[type[PlantFunctionalTypeStrict]]) -> None
         object.__setattr__(self, "_n_stems", len(pfts))
 
         # Populate the trait attributes with arrays
-        for pft_field in self.trait_attrs:
+        for pft_field in self.array_attrs:
             object.__setattr__(
                 self, pft_field, np.array([getattr(pft, pft_field) for pft in pfts])
             )
@@ -436,12 +438,12 @@ def get_stem_traits(self, pft_names: NDArray[np.str_]) -> StemTraits:
         # matching the pft_names and pass that into the StemTraits constructor.
 
         return StemTraits(
-            **{trt: getattr(self, trt)[pft_index] for trt in self.trait_attrs}
+            **{trt: getattr(self, trt)[pft_index] for trt in self.array_attrs}
         )
 
 
 @dataclass()
-class StemTraits:
+class StemTraits(PandasExporter):
     """A dataclass for stem traits.
 
     This dataclass is used to provide arrays of plant functional type (PFT) traits
@@ -457,9 +459,9 @@ class StemTraits:
     """
 
     # A class variable setting the attribute names of traits.
-    trait_attrs: ClassVar[list[str]] = [
+    array_attrs: ClassVar[tuple[str, ...]] = tuple(
         f.name for f in fields(PlantFunctionalTypeStrict)
-    ]
+    )
 
     # Instance trait attributes
     name: NDArray[np.str_]

From 94b2d6f253d7d150be1106e4aac7f55753477d0a Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Fri, 18 Oct 2024 19:20:31 +0100
Subject: [PATCH 03/31] Adding PandasExporter to tmodel functions classes,
 extending to_pandas() to unstack 2D arrays

---
 docs/source/users/demography/t_model.md | 35 +++++++++++++++++--------
 pyrealm/demography/core.py              | 26 ++++++++++++++++--
 pyrealm/demography/t_model_functions.py |  9 ++++---
 3 files changed, 53 insertions(+), 17 deletions(-)

diff --git a/docs/source/users/demography/t_model.md b/docs/source/users/demography/t_model.md
index cd38532c..2d47e886 100644
--- a/docs/source/users/demography/t_model.md
+++ b/docs/source/users/demography/t_model.md
@@ -40,7 +40,7 @@ from pyrealm.demography.t_model_functions import StemAllocation, StemAllometry
 ```
 
 To generate predictions under the T Model, we need a Flora object providing the
-[trait values](./flora.md) for each of the PFTsto be modelled:
+[trait values](./flora.md) for each of the PFTs to be modelled:
 
 ```{code-cell} ipython3
 # Three PFTS
@@ -77,12 +77,12 @@ for each PFT. This then calculates a single estimate at the given size for each
 single_allometry = StemAllometry(stem_traits=flora, at_dbh=np.array([0.1, 0.1, 0.1]))
 ```
 
-We can display those predictions as a `pandas.DataFrame`:
+The class provides a
+{meth}`~pyrealm.demography.t_model_functions.StemAllometry.to_pandas()` method
+to export the stem data for data exploration.
 
 ```{code-cell} ipython3
-pd.DataFrame(
-    {k: getattr(single_allometry, k) for k in single_allometry.allometry_attrs}
-)
+single_allometry.to_pandas()
 ```
 
 However, the DBH values can also be a column array (an `N` x 1 array). In this case, the
@@ -123,11 +123,19 @@ for ax, (var, ylab) in zip(axes.flatten(), plot_details):
         ax.legend(frameon=False)
 ```
 
+The {meth}`~pyrealm.demography.t_model_functions.StemAllometry.to_pandas()` method
+can still be used, but the values are stacked into columns along with a column index
+
+```{code-cell} ipython3
+allometries.to_pandas()
+```
+
 ## Productivity allocation
 
 The T Model also predicts how potential GPP will be allocated to respiration, turnover
-and growth for stems with a given PFT and allometry. Again, a single value can be
-provided to get a single estimate of the allocation model for each stem:
+and growth for stems with a given PFT and allometry using the
+{meth}`~pyrealm.demography.t_model_functions.StemAllometry` class. Again, a single
+value can be provided to get a single estimate of the allocation model for each stem:
 
 ```{code-cell} ipython3
 single_allocation = StemAllocation(
@@ -136,10 +144,12 @@ single_allocation = StemAllocation(
 single_allocation
 ```
 
+The class provides the
+{meth}`~pyrealm.demography.t_model_functions.StemAllocation.to_pandas()` method to
+export data for exploration.
+
 ```{code-cell} ipython3
-pd.DataFrame(
-    {k: getattr(single_allocation, k) for k in single_allocation.allocation_attrs}
-)
+single_allocation.to_pandas()
 ```
 
 Using a column array of potential GPP values can be used to predict multiple estimates of
@@ -217,6 +227,9 @@ for ax, (var, ylab) in zip(axes, plot_details):
 fig.delaxes(axes[-1])
 ```
 
-```{code-cell} ipython3
+As before, the {meth}`~pyrealm.demography.t_model_functions.StemAllometry.to_pandas()`
+method can be used to export the data for each stem:
 
+```{code-cell} ipython3
+allocation.to_pandas()
 ```
diff --git a/pyrealm/demography/core.py b/pyrealm/demography/core.py
index f628da79..f1d58091 100644
--- a/pyrealm/demography/core.py
+++ b/pyrealm/demography/core.py
@@ -22,8 +22,30 @@ class PandasExporter(ABC):
     array_attrs: ClassVar[tuple[str, ...]]
 
     def to_pandas(self) -> pd.DataFrame:
-        """Convert the instance array attributes into a {class}`pandas.DataFrame."""
-        return pd.DataFrame({k: getattr(self, k) for k in self.array_attrs})
+        """Convert the instance array attributes into a {class}`pandas.DataFrame.
+
+        If the array values are two-dimensional (e.g. stems by heights), the data are
+        stacked and an index field is added.
+        """
+
+        # Extract the attributes into a dictionary
+        data = {k: getattr(self, k) for k in self.array_attrs}
+
+        # Check the first attribute array to see if the values are two dimensional
+        data_shape = data[self.array_attrs[0]].shape
+
+        if len(data_shape) == 2:
+            # create an index entry to show the column of each value
+            stacked_data = {
+                "column_index": np.repeat(np.arange(data_shape[1]), data_shape[0])
+            }
+            # Ravel the attribute data using column-major Fortan style
+            for ky, vl in data.items():
+                stacked_data[ky] = np.ravel(vl, order="F")
+
+            return pd.DataFrame(stacked_data)
+
+        return pd.DataFrame(data)
 
 
 class Cohorts(ABC):
diff --git a/pyrealm/demography/t_model_functions.py b/pyrealm/demography/t_model_functions.py
index dbedb70f..8fd578db 100644
--- a/pyrealm/demography/t_model_functions.py
+++ b/pyrealm/demography/t_model_functions.py
@@ -12,6 +12,7 @@
 from numpy.typing import NDArray
 
 from pyrealm.core.utilities import check_input_shapes
+from pyrealm.demography.core import PandasExporter
 from pyrealm.demography.flora import Flora, StemTraits
 
 
@@ -659,7 +660,7 @@ def calculate_growth_increments(
 
 
 @dataclass
-class StemAllometry:
+class StemAllometry(PandasExporter):
     """Calculate T Model allometric predictions across a set of stems.
 
     This method calculate predictions of stem allometries for stem height, crown area,
@@ -675,7 +676,7 @@ class StemAllometry:
             allometry values.
     """
 
-    allometry_attrs: ClassVar[tuple[str, ...]] = (
+    array_attrs: ClassVar[tuple[str, ...]] = (
         "dbh",
         "stem_height",
         "crown_area",
@@ -796,7 +797,7 @@ def __repr__(self) -> str:
 
 
 @dataclass()
-class StemAllocation:
+class StemAllocation(PandasExporter):
     """Calculate T Model allocation predictions across a set of stems.
 
     This method calculate predictions of allocation of potential GPP for stems under the
@@ -814,7 +815,7 @@ class StemAllocation:
             predict stem allometry values.
     """
 
-    allocation_attrs: ClassVar[tuple[str, ...]] = (
+    array_attrs: ClassVar[tuple[str, ...]] = (
         "potential_gpp",
         "whole_crown_gpp",
         "sapwood_respiration",

From e75c29aa074bfb6079e72e3968aaf07a31cc6559 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Sat, 19 Oct 2024 11:40:57 +0100
Subject: [PATCH 04/31] Adding PandasExporter to CrownProfile class

---
 docs/source/users/demography/crown.md | 24 ++++++++++++------------
 pyrealm/demography/crown.py           |  7 ++++---
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/docs/source/users/demography/crown.md b/docs/source/users/demography/crown.md
index 69e64018..a2f156ab 100644
--- a/docs/source/users/demography/crown.md
+++ b/docs/source/users/demography/crown.md
@@ -218,12 +218,10 @@ flora = Flora([narrow_pft, medium_pft, wide_pft])
 flora
 ```
 
-The Flora object can also be used to show a table of canopy variables:
+The key canopy variables from the flora are:
 
 ```{code-cell} ipython3
-# TODO - add a Flora.to_pandas() method
-flora_data = pd.DataFrame({k: getattr(flora, k) for k in flora.trait_attrs})
-flora_data[["name", "ca_ratio", "m", "n", "f_g", "q_m", "z_max_prop"]]
+flora.to_pandas()[["name", "h_max", "ca_ratio", "m", "n", "f_g", "q_m", "z_max_prop"]]
 ```
 
 The next section of code generates the `StemAllometry` to use for the profiles.
@@ -249,14 +247,12 @@ allometry = StemAllometry(stem_traits=flora, at_dbh=stem_dbh)
 We can again use {mod}`pandas` to get a table of those allometric predictions:
 
 ```{code-cell} ipython3
-pd.DataFrame({k: getattr(allometry, k) for k in allometry.allometry_attrs})
+allometry.to_pandas()
 ```
 
 Finally, we can define a set of vertical heights. In order to calculate the
 variables for each PFT at each height, this needs to be provided as a column array,
-that is with a shape `(N, 1)`.
-
-We can then calculate the crown profiles.
+that is with a shape `(N, 1)`. We can then calculate the crown profiles:
 
 ```{code-cell} ipython3
 # Create a set of vertical heights as a column array.
@@ -278,6 +274,14 @@ above calculated at each height $z$:
 crown_profiles
 ```
 
+The {meth}`~pyrealm.demography.crown.CrownProfile.to_pandas` method can be used to
+extract the data into a table, with the separate stems identified by the column index
+field.
+
+```{code-cell} ipython3
+crown_profiles.to_pandas()
+```
+
 ### Visualising crown profiles
 
 The code below generates a plot of the vertical shape profiles of the crowns for each
@@ -532,7 +536,3 @@ plt.legend(
     bbox_to_anchor=(0.5, 1.15),
 )
 ```
-
-```{code-cell} ipython3
-
-```
diff --git a/pyrealm/demography/crown.py b/pyrealm/demography/crown.py
index 0f1d43ed..fa1f2ed8 100644
--- a/pyrealm/demography/crown.py
+++ b/pyrealm/demography/crown.py
@@ -9,6 +9,7 @@
 from numpy.typing import NDArray
 
 from pyrealm.core.utilities import check_input_shapes
+from pyrealm.demography.core import PandasExporter
 from pyrealm.demography.flora import Flora, StemTraits
 from pyrealm.demography.t_model_functions import StemAllometry
 
@@ -282,7 +283,7 @@ def calculate_stem_projected_leaf_area_at_z(
 
 
 @dataclass
-class CrownProfile:
+class CrownProfile(PandasExporter):
     """Calculate vertical crown profiles for stems.
 
     This method calculates crown profile predictions, given an array of vertical
@@ -315,7 +316,7 @@ class CrownProfile:
         z_max: A row array providing expected z_max height for each PFT.
     """
 
-    var_attr_names: ClassVar[tuple[str, ...]] = (
+    array_attrs: ClassVar[tuple[str, ...]] = (
         "relative_crown_radius",
         "crown_radius",
         "projected_crown_area",
@@ -447,7 +448,7 @@ def get_crown_xy(
     """
 
     # Input validation
-    if attr not in crown_profile.var_attr_names:
+    if attr not in crown_profile.array_attrs:
         raise ValueError(f"Unknown crown profile attribute: {attr}")
 
     # TODO

From f5b0b6a4542fbd9897221a2cd172c87d65cb3fc1 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 21 Oct 2024 17:52:25 +0000
Subject: [PATCH 05/31] [pre-commit.ci] pre-commit autoupdate
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

updates:
- [github.com/astral-sh/ruff-pre-commit: v0.6.9 → v0.7.0](https://github.com/astral-sh/ruff-pre-commit/compare/v0.6.9...v0.7.0)
- [github.com/pre-commit/mirrors-mypy: v1.11.2 → v1.12.1](https://github.com/pre-commit/mirrors-mypy/compare/v1.11.2...v1.12.1)
---
 .pre-commit-config.yaml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 41cf7f95..4ae3bf2a 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -6,7 +6,7 @@ repos:
       - id: debug-statements
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.6.9
+    rev: v0.7.0
     hooks:
       # Run the linter.
       - id: ruff
@@ -14,7 +14,7 @@ repos:
       # Run the formatter.
       - id: ruff-format
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: "v1.11.2"
+    rev: "v1.12.1"
     hooks:
       - id: mypy
         additional_dependencies: [numpy, types-tabulate, pandas-stubs]

From 7b76d60310988259d1b135d761b98a692b9786d5 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 10:24:08 +0100
Subject: [PATCH 06/31] Initial Cohorts dataclass and basic test

---
 pyrealm/demography/community.py         | 99 +++++++++++++------------
 tests/unit/demography/test_community.py | 48 ++++++++++++
 2 files changed, 99 insertions(+), 48 deletions(-)

diff --git a/pyrealm/demography/community.py b/pyrealm/demography/community.py
index 2f593471..02c23f43 100644
--- a/pyrealm/demography/community.py
+++ b/pyrealm/demography/community.py
@@ -111,7 +111,7 @@
 
 import json
 import sys
-from dataclasses import InitVar, dataclass, field
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
 
@@ -133,6 +133,39 @@
     from tomli import TOMLDecodeError
 
 
+@dataclass
+class Cohorts:
+    """A dataclass to hold data for a set of plant cohorts.
+
+    The attributes should be numpy arrays of equal length, containing an entry for each
+    cohort in the data class.
+    """
+
+    dbh_values: NDArray[np.float64]
+    n_individuals: NDArray[np.int_]
+    pft_names: NDArray[np.str_]
+
+    def __post_init__(self) -> None:
+        """Validation of cohorts data."""
+
+        # TODO - validation - maybe make this optional to reduce workload within
+        # simulations
+
+        # Check cohort data types
+        if not (
+            isinstance(self.dbh_values, np.ndarray)
+            and isinstance(self.n_individuals, np.ndarray)
+            and isinstance(self.pft_names, np.ndarray)
+        ):
+            raise ValueError("Cohort data not passed as numpy arrays")
+
+        # Check the cohort inputs are of equal length
+        try:
+            check_input_shapes(self.dbh_values, self.n_individuals, self.dbh_values)
+        except ValueError:
+            raise ValueError("Cohort arrays are of unequal length")
+
+
 class CohortSchema(Schema):
     """A validation schema for Cohort data objects.
 
@@ -224,7 +257,7 @@ class CommunityStructuredDataSchema(Schema):
     )
 
     @post_load
-    def cohort_objects_to_arrays(self, data: dict, **kwargs: Any) -> dict:
+    def cohort_objects_to_arrays(self, data: dict, **kwargs: Any) -> Cohorts:
         """Convert cohorts to arrays.
 
         This post load method converts the cohort objects into arrays, which is the
@@ -235,15 +268,11 @@ def cohort_objects_to_arrays(self, data: dict, **kwargs: Any) -> dict:
             kwargs: Additional keyword arguments passed by marshmallow
         """
 
-        data["cohort_dbh_values"] = np.array([c["dbh_value"] for c in data["cohorts"]])
-        data["cohort_n_individuals"] = np.array(
-            [c["n_individuals"] for c in data["cohorts"]]
+        return Cohorts(
+            dbh_values=np.array([c["dbh_value"] for c in data["cohorts"]]),
+            n_individuals=np.array([c["n_individuals"] for c in data["cohorts"]]),
+            pft_names=np.array([c["pft_name"] for c in data["cohorts"]]),
         )
-        data["cohort_pft_names"] = np.array([c["pft_name"] for c in data["cohorts"]])
-
-        del data["cohorts"]
-
-        return data
 
 
 class CommunityCSVDataSchema(Schema):
@@ -304,15 +333,17 @@ def make_cell_data_scalar(self, data: dict, **kwargs: Any) -> dict:
         """Make cell data scalar.
 
         This post load method reduces the repeated cell id and cell area across CSV data
-        rows into the scalar inputs required to initialise a Community object.
+        rows into the scalar inputs required to initialise a Community object and
+        packages the data on individual cohorts into a Cohorts object.
         """
 
         data["cell_id"] = data["cell_id"][0]
         data["cell_area"] = data["cell_area"][0]
-
-        data["cohort_dbh_values"] = np.array(data["cohort_dbh_values"])
-        data["cohort_n_individuals"] = np.array(data["cohort_n_individuals"])
-        data["cohort_pft_names"] = np.array(data["cohort_pft_names"])
+        data["cohorts"] = Cohorts(
+            dbh_values=np.array(data["cohort_dbh_values"]),
+            n_individuals=np.array(data["cohort_n_individuals"]),
+            pft_names=np.array(data["cohort_pft_names"]),
+        )
 
         return data
 
@@ -353,21 +384,15 @@ class Community:
     flora: Flora
 
     # - arrays representing properties of cohorts
-    cohort_dbh_values: InitVar[NDArray[np.float64]]
-    cohort_n_individuals: InitVar[NDArray[np.int_]]
-    cohort_pft_names: InitVar[NDArray[np.str_]]
+    cohorts: Cohorts
 
     # 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__(
         self,
-        cohort_dbh_values: NDArray[np.float64],
-        cohort_n_individuals: NDArray[np.int_],
-        cohort_pft_names: NDArray[np.str_],
     ) -> None:
         """Validate inputs and populate derived community attributes.
 
@@ -382,37 +407,15 @@ def __post_init__(
         if not (isinstance(self.cell_id, int) and self.cell_id >= 0):
             raise ValueError("Community cell id must be a integer >= 0.")
 
-        # Check cohort data types
-        if not (
-            isinstance(cohort_dbh_values, np.ndarray)
-            and isinstance(cohort_n_individuals, np.ndarray)
-            and isinstance(cohort_pft_names, np.ndarray)
-        ):
-            raise ValueError("Cohort data not passed as numpy arrays.")
-
-        # Check the cohort inputs are of equal length
-        try:
-            check_input_shapes(
-                cohort_dbh_values, cohort_n_individuals, cohort_dbh_values
-            )
-        except ValueError:
-            raise ValueError("Cohort arrays are of unequal length")
-
-        # Store as a dictionary
-        self.cohort_data: dict[str, NDArray] = {
-            "name": cohort_pft_names,
-            "dbh": cohort_dbh_values,
-            "n_individuals": cohort_n_individuals,
-        }
+        # How many cohorts
+        self.number_of_cohorts = len(self.cohorts.dbh_values)
 
         # Get the stem traits for the cohorts
-        self.stem_traits = self.flora.get_stem_traits(cohort_pft_names)
-
-        self.number_of_cohorts = len(cohort_pft_names)
+        self.stem_traits = self.flora.get_stem_traits(self.cohorts.pft_names)
 
         # Populate the stem allometry
         self.stem_allometry = StemAllometry(
-            stem_traits=self.stem_traits, at_dbh=cohort_dbh_values
+            stem_traits=self.stem_traits, at_dbh=self.cohorts.dbh_values
         )
 
     @classmethod
diff --git a/tests/unit/demography/test_community.py b/tests/unit/demography/test_community.py
index 1e6cdea6..4394523e 100644
--- a/tests/unit/demography/test_community.py
+++ b/tests/unit/demography/test_community.py
@@ -47,6 +47,54 @@ def check_expected(community, expected):
     )
 
 
+@pytest.mark.parametrize(
+    argnames="args,outcome,excep_message",
+    argvalues=[
+        pytest.param(
+            {
+                "pft_names": np.array(["broadleaf", "conifer"]),
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
+            },
+            does_not_raise(),
+            None,
+            id="correct",
+        ),
+        pytest.param(
+            {
+                "pft_names": False,
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
+            },
+            pytest.raises(ValueError),
+            "Cohort data not passed as numpy arrays",
+            id="not np array",
+        ),
+        pytest.param(
+            {
+                "pft_names": np.array(["broadleaf", "conifer"]),
+                "n_individuals": np.array([6, 1, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
+            },
+            pytest.raises(ValueError),
+            "Cohort arrays are of unequal length",
+            id="not np array",
+        ),
+    ],
+)
+def test_Cohorts(args, outcome, excep_message):
+    """Test the cohorts data structure."""
+    from pyrealm.demography.community import Cohorts
+
+    with outcome as excep:
+        cohorts = Cohorts(**args)
+        # trivial test of success
+        assert len(cohorts.dbh_values) == 2
+        return
+
+    assert str(excep.value) == excep_message
+
+
 @pytest.mark.parametrize(
     argnames="args,outcome,excep_message",
     argvalues=[

From 84898647bebf9dc8bfc4723f68cb6193d15b5d3f Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 11:07:16 +0100
Subject: [PATCH 07/31] Updating community tests

---
 tests/unit/demography/test_community.py | 121 ++++++++++--------------
 1 file changed, 51 insertions(+), 70 deletions(-)

diff --git a/tests/unit/demography/test_community.py b/tests/unit/demography/test_community.py
index 4394523e..cf7bd5b1 100644
--- a/tests/unit/demography/test_community.py
+++ b/tests/unit/demography/test_community.py
@@ -34,7 +34,7 @@ def check_expected(community, expected):
     """Helper function to provide simple check of returned community objects."""
 
     assert np.allclose(
-        community.cohort_data["n_individuals"],
+        community.cohorts.n_individuals,
         expected["n_individuals"],
     )
     assert np.allclose(
@@ -68,7 +68,17 @@ def check_expected(community, expected):
             },
             pytest.raises(ValueError),
             "Cohort data not passed as numpy arrays",
-            id="not np array",
+            id="not_iterable",
+        ),
+        pytest.param(
+            {
+                "pft_names": ["broadleaf", "conifer"],
+                "n_individuals": [6, 1],
+                "dbh_values": [0.2, 0.5],
+            },
+            pytest.raises(ValueError),
+            "Cohort data not passed as numpy arrays",
+            id="lists_not_arrays",
         ),
         pytest.param(
             {
@@ -96,123 +106,91 @@ def test_Cohorts(args, outcome, excep_message):
 
 
 @pytest.mark.parametrize(
-    argnames="args,outcome,excep_message",
+    argnames="args,cohort_data,outcome,excep_message",
     argvalues=[
         pytest.param(
+            {"cell_id": 1, "cell_area": 100},
             {
-                "cell_id": 1,
-                "cell_area": 100,
-                "cohort_pft_names": np.array(["broadleaf", "conifer"]),
-                "cohort_n_individuals": np.array([6, 1]),
-                "cohort_dbh_values": np.array([0.2, 0.5]),
+                "pft_names": np.array(["broadleaf", "conifer"]),
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
             },
             does_not_raise(),
             None,
             id="correct",
         ),
         pytest.param(
+            {"cell_area": 100},
             {
-                "cell_id": 1,
-                "cell_area": 100,
-                "cohort_pft_names": ["broadleaf", "conifer"],
-                "cohort_n_individuals": [6, 1],
-                "cohort_dbh_values": [0.2, 0.5],
-            },
-            pytest.raises(ValueError),
-            "Cohort data not passed as numpy arrays.",
-            id="lists_not_arrays",
-        ),
-        pytest.param(
-            {
-                "cell_id": 1,
-                "cell_area": 100,
-                "cohort_pft_names": np.array(["broadleaf", "conifer"]),
-                "cohort_n_individuals": np.array([6, 1]),
-                "cohort_dbh_values": np.array([0.2, 0.5, 0.9]),
-            },
-            pytest.raises(ValueError),
-            "Cohort arrays are of unequal length",
-            id="unequal_cohort_arrays",
-        ),
-        pytest.param(
-            {
-                "cell_area": 100,
-                "cohort_pft_names": np.array(["broadleaf", "conifer"]),
-                "cohort_n_individuals": np.array([6, 1]),
-                "cohort_dbh_values": np.array([0.2, 0.5]),
+                "pft_names": np.array(["broadleaf", "conifer"]),
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
             },
             pytest.raises(TypeError),
             "Community.__init__() missing 1 required positional argument: 'cell_id'",
             id="missing_arg",
         ),
         pytest.param(
+            {"cell_id": 1, "cell_area": 100, "cell_elevation": 100},
             {
-                "cell_id": 1,
-                "cell_area": 100,
-                "cell_elevation": 100,
-                "cohort_pft_names": np.array(["broadleaf", "conifer"]),
-                "cohort_n_individuals": np.array([6, 1]),
-                "cohort_dbh_values": np.array([0.2, 0.5]),
+                "pft_names": np.array(["broadleaf", "conifer"]),
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
             },
             pytest.raises(TypeError),
             "Community.__init__() got an unexpected keyword argument 'cell_elevation'",
             id="extra_arg",
         ),
         pytest.param(
+            {"cell_id": 1, "cell_area": "100"},
             {
-                "cell_id": 1,
-                "cell_area": "100",
-                "cohort_pft_names": np.array(["broadleaf", "conifer"]),
-                "cohort_n_individuals": np.array([6, 1]),
-                "cohort_dbh_values": np.array([0.2, 0.5]),
+                "pft_names": np.array(["broadleaf", "conifer"]),
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
             },
             pytest.raises(ValueError),
             "Community cell area must be a positive number.",
             id="cell_area_as_string",
         ),
         pytest.param(
+            {"cell_id": 1, "cell_area": -100},
             {
-                "cell_id": 1,
-                "cell_area": -100,
-                "cohort_pft_names": np.array(["broadleaf", "conifer"]),
-                "cohort_n_individuals": np.array([6, 1]),
-                "cohort_dbh_values": np.array([0.2, 0.5]),
+                "pft_names": np.array(["broadleaf", "conifer"]),
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
             },
             pytest.raises(ValueError),
             "Community cell area must be a positive number.",
             id="cell_area_negative",
         ),
         pytest.param(
+            {"cell_id": "1", "cell_area": 100},
             {
-                "cell_id": "1",
-                "cell_area": 100,
-                "cohort_pft_names": np.array(["broadleaf", "conifer"]),
-                "cohort_n_individuals": np.array([6, 1]),
-                "cohort_dbh_values": np.array([0.2, 0.5]),
+                "pft_names": np.array(["broadleaf", "conifer"]),
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
             },
             pytest.raises(ValueError),
             "Community cell id must be a integer >= 0.",
             id="cell_id_as_string",
         ),
         pytest.param(
+            {"cell_id": -1, "cell_area": 100},
             {
-                "cell_id": -1,
-                "cell_area": 100,
-                "cohort_pft_names": np.array(["broadleaf", "conifer"]),
-                "cohort_n_individuals": np.array([6, 1]),
-                "cohort_dbh_values": np.array([0.2, 0.5]),
+                "pft_names": np.array(["broadleaf", "conifer"]),
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
             },
             pytest.raises(ValueError),
             "Community cell id must be a integer >= 0.",
             id="cell_id_negative",
         ),
         pytest.param(
+            {"cell_id": 1, "cell_area": 100},
             {
-                "cell_id": 1,
-                "cell_area": 100,
-                "cohort_pft_names": np.array(["broadleaf", "juniper"]),
-                "cohort_n_individuals": np.array([6, 1]),
-                "cohort_dbh_values": np.array([0.2, 0.5]),
+                "pft_names": np.array(["broadleaf", "juniper"]),
+                "n_individuals": np.array([6, 1]),
+                "dbh_values": np.array([0.2, 0.5]),
             },
             pytest.raises(ValueError),
             "Plant functional types unknown in flora: juniper",
@@ -221,7 +199,7 @@ def test_Cohorts(args, outcome, excep_message):
     ],
 )
 def test_Community__init__(
-    fixture_flora, fixture_expected, args, outcome, excep_message
+    fixture_flora, fixture_expected, args, cohort_data, outcome, excep_message
 ):
     """Test Community initialisation.
 
@@ -229,10 +207,13 @@ def test_Community__init__(
     properties.
     """
 
-    from pyrealm.demography.community import Community
+    from pyrealm.demography.community import Cohorts, Community
+
+    # Build the cohorts object
+    cohorts = Cohorts(**cohort_data)
 
     with outcome as excep:
-        community = Community(**args, flora=fixture_flora)
+        community = Community(**args, cohorts=cohorts, flora=fixture_flora)
 
         if isinstance(outcome, does_not_raise):
             # Simple test that data is loaded and trait and t model data calculated

From aab0756f6660bff111a02e196c869e8312888d1c Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 11:08:01 +0100
Subject: [PATCH 08/31] Adopting new Cohorts structure through other modules

---
 pyrealm/demography/canopy.py         |  8 +++---
 pyrealm/demography/community.py      | 39 ++++++++++++++++------------
 tests/unit/demography/test_canopy.py |  4 +--
 3 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/pyrealm/demography/canopy.py b/pyrealm/demography/canopy.py
index 282b20d4..11614cea 100644
--- a/pyrealm/demography/canopy.py
+++ b/pyrealm/demography/canopy.py
@@ -118,7 +118,7 @@ def fit_perfect_plasticity_approximation(
 
     # Calculate the number of layers to contain the total community crown area
     total_community_crown_area = (
-        community.stem_allometry.crown_area * community.cohort_data["n_individuals"]
+        community.stem_allometry.crown_area * community.cohorts.n_individuals
     ).sum()
     crown_area_per_layer = community.cell_area * (1 - canopy_gap_fraction)
     n_layers = int(np.ceil(total_community_crown_area / crown_area_per_layer))
@@ -144,7 +144,7 @@ def fit_perfect_plasticity_approximation(
                 community.stem_traits.n,
                 community.stem_traits.q_m,
                 community.stem_allometry.crown_z_max,
-                community.cohort_data["n_individuals"],
+                community.cohorts.n_individuals,
                 target_area,
                 False,  # validate
             ),
@@ -293,7 +293,7 @@ def _calculate_canopy(self, community: Community) -> None:
         # the available community area.
         self.cohort_lai = (
             self.stem_leaf_area
-            * community.cohort_data["n_individuals"]
+            * community.cohorts.n_individuals
             * community.stem_traits.lai
         ) / community.cell_area  # self.filled_community_area
 
@@ -323,4 +323,4 @@ def _calculate_canopy(self, community: Community) -> None:
         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.cohort_data["n_individuals"]
+        self.stem_fapar = self.cohort_fapar / community.cohorts.n_individuals
diff --git a/pyrealm/demography/community.py b/pyrealm/demography/community.py
index 02c23f43..b98d3c0d 100644
--- a/pyrealm/demography/community.py
+++ b/pyrealm/demography/community.py
@@ -95,7 +95,7 @@
 >>> pd.DataFrame({
 ...    'name': community.stem_traits.name,
 ...    'dbh': community.stem_allometry.dbh,
-...    'n_individuals': community.cohort_data["n_individuals"],
+...    'n_individuals': community.cohorts.n_individuals,
 ...    'stem_height': community.stem_allometry.stem_height,
 ...    'crown_area': community.stem_allometry.crown_area,
 ...    'stem_mass': community.stem_allometry.stem_mass,
@@ -257,21 +257,26 @@ class CommunityStructuredDataSchema(Schema):
     )
 
     @post_load
-    def cohort_objects_to_arrays(self, data: dict, **kwargs: Any) -> Cohorts:
+    def convert_to_community_args(self, data: dict, **kwargs: Any) -> dict[str, Any]:
         """Convert cohorts to arrays.
 
-        This post load method converts the cohort objects into arrays, which is the
-        format used to initialise a Community object.
+        This post load method converts the cohort arrays into a Cohorts objects and
+        packages the data up into the required arguments used to initialise a Community
+        object.
 
         Args:
             data: Data passed to the validator
             kwargs: Additional keyword arguments passed by marshmallow
         """
 
-        return Cohorts(
-            dbh_values=np.array([c["dbh_value"] for c in data["cohorts"]]),
-            n_individuals=np.array([c["n_individuals"] for c in data["cohorts"]]),
-            pft_names=np.array([c["pft_name"] for c in data["cohorts"]]),
+        return dict(
+            cell_id=data["cell_id"],
+            cell_area=data["cell_area"],
+            cohorts=Cohorts(
+                dbh_values=np.array([c["dbh_value"] for c in data["cohorts"]]),
+                n_individuals=np.array([c["n_individuals"] for c in data["cohorts"]]),
+                pft_names=np.array([c["pft_name"] for c in data["cohorts"]]),
+            ),
         )
 
 
@@ -329,7 +334,7 @@ def validate_consistent_cell_data(self, data: dict, **kwargs: Any) -> None:
             raise ValueError("Cell area varies in community data")
 
     @post_load
-    def make_cell_data_scalar(self, data: dict, **kwargs: Any) -> dict:
+    def convert_to_community_args(self, data: dict, **kwargs: Any) -> dict[str, Any]:
         """Make cell data scalar.
 
         This post load method reduces the repeated cell id and cell area across CSV data
@@ -337,16 +342,16 @@ def make_cell_data_scalar(self, data: dict, **kwargs: Any) -> dict:
         packages the data on individual cohorts into a Cohorts object.
         """
 
-        data["cell_id"] = data["cell_id"][0]
-        data["cell_area"] = data["cell_area"][0]
-        data["cohorts"] = Cohorts(
-            dbh_values=np.array(data["cohort_dbh_values"]),
-            n_individuals=np.array(data["cohort_n_individuals"]),
-            pft_names=np.array(data["cohort_pft_names"]),
+        return dict(
+            cell_id=data["cell_id"][0],
+            cell_area=data["cell_area"][0],
+            cohorts=Cohorts(
+                dbh_values=np.array(data["cohort_dbh_values"]),
+                n_individuals=np.array(data["cohort_n_individuals"]),
+                pft_names=np.array(data["cohort_pft_names"]),
+            ),
         )
 
-        return data
-
 
 @dataclass
 class Community:
diff --git a/tests/unit/demography/test_canopy.py b/tests/unit/demography/test_canopy.py
index a078fbfe..fbcc81db 100644
--- a/tests/unit/demography/test_canopy.py
+++ b/tests/unit/demography/test_canopy.py
@@ -40,7 +40,7 @@ def test_Canopy__init__():
             (
                 (
                     community.stem_allometry.crown_area
-                    * community.cohort_data["n_individuals"]
+                    * community.cohorts.n_individuals
                 ).sum()
                 * (1 + canopy_gap_fraction)
             )
@@ -78,7 +78,7 @@ def test_solve_canopy_area_filling_height(fixture_community):
             z=this_height,
             stem_height=fixture_community.stem_allometry.stem_height,
             crown_area=fixture_community.stem_allometry.crown_area,
-            n_individuals=fixture_community.cohort_data["n_individuals"],
+            n_individuals=fixture_community.cohorts.n_individuals,
             m=fixture_community.stem_traits.m,
             n=fixture_community.stem_traits.n,
             q_m=fixture_community.stem_traits.q_m,

From fbcb953b18a6cd8a77020408997aba25924d829e Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 11:14:10 +0100
Subject: [PATCH 09/31] Fix broken fixtures and tests

---
 tests/unit/demography/conftest.py    | 10 ++++++----
 tests/unit/demography/test_canopy.py | 10 ++++++----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/tests/unit/demography/conftest.py b/tests/unit/demography/conftest.py
index 735cb4c1..6260c4b9 100644
--- a/tests/unit/demography/conftest.py
+++ b/tests/unit/demography/conftest.py
@@ -66,7 +66,7 @@ def rtmodel_flora():
 @pytest.fixture
 def fixture_community():
     """A fixture providing a simple community."""
-    from pyrealm.demography.community import Community
+    from pyrealm.demography.community import Cohorts, Community
     from pyrealm.demography.flora import Flora, PlantFunctionalType
 
     # A simple community containing one sample stem, with an initial crown gap fraction
@@ -76,9 +76,11 @@ def fixture_community():
         cell_id=1,
         cell_area=100,
         flora=flora,
-        cohort_n_individuals=np.repeat([1], 4),
-        cohort_pft_names=np.repeat(["test"], 4),
-        cohort_dbh_values=np.array([0.2, 0.4, 0.6, 0.8]),
+        cohorts=Cohorts(
+            n_individuals=np.repeat([1], 4),
+            pft_names=np.repeat(["test"], 4),
+            dbh_values=np.array([0.2, 0.4, 0.6, 0.8]),
+        ),
     )
 
 
diff --git a/tests/unit/demography/test_canopy.py b/tests/unit/demography/test_canopy.py
index fbcc81db..95e04bf2 100644
--- a/tests/unit/demography/test_canopy.py
+++ b/tests/unit/demography/test_canopy.py
@@ -12,7 +12,7 @@ def test_Canopy__init__():
     """
 
     from pyrealm.demography.canopy import Canopy
-    from pyrealm.demography.community import Community
+    from pyrealm.demography.community import Cohorts, Community
     from pyrealm.demography.flora import Flora, PlantFunctionalType
 
     flora = Flora(
@@ -25,9 +25,11 @@ def test_Canopy__init__():
     community = Community(
         cell_id=1,
         cell_area=20,
-        cohort_pft_names=np.array(["broadleaf", "conifer"]),
-        cohort_n_individuals=np.array([6, 1]),
-        cohort_dbh_values=np.array([0.2, 0.5]),
+        cohorts=Cohorts(
+            pft_names=np.array(["broadleaf", "conifer"]),
+            n_individuals=np.array([6, 1]),
+            dbh_values=np.array([0.2, 0.5]),
+        ),
         flora=flora,
     )
 

From 5e5d811447e9ab0c286ccb37d519057d8bb54be2 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 11:40:02 +0100
Subject: [PATCH 10/31] Missed a doctest failure

---
 pyrealm/demography/community.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/pyrealm/demography/community.py b/pyrealm/demography/community.py
index b98d3c0d..8e629d6d 100644
--- a/pyrealm/demography/community.py
+++ b/pyrealm/demography/community.py
@@ -84,9 +84,11 @@
 ...     cell_id=1,
 ...     cell_area=1000.0,
 ...     flora=flora,
-...     cohort_dbh_values=cohort_dbh_values,
-...     cohort_n_individuals=cohort_n_individuals,
-...     cohort_pft_names=cohort_pft_names
+...     cohorts=Cohorts(
+...         dbh_values=cohort_dbh_values,
+...         n_individuals=cohort_n_individuals,
+...         pft_names=cohort_pft_names,
+...     ),
 ... )
 
 Convert some of the data to a :class:`pandas.DataFrame` for nicer display and show some

From c43bedf3a5f7f892308e5b4b475165753503e539 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Tue, 22 Oct 2024 11:54:08 +0100
Subject: [PATCH 11/31] Fixing docs

---
 docs/source/users/demography/canopy.md    | 28 +++++++++++++----------
 docs/source/users/demography/community.md | 14 +++++-------
 2 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/docs/source/users/demography/canopy.md b/docs/source/users/demography/canopy.md
index b3c17943..da837e9c 100644
--- a/docs/source/users/demography/canopy.md
+++ b/docs/source/users/demography/canopy.md
@@ -40,7 +40,7 @@ import numpy as np
 import pandas as pd
 
 from pyrealm.demography.flora import PlantFunctionalType, Flora
-from pyrealm.demography.community import Community
+from pyrealm.demography.community import Cohorts, Community
 from pyrealm.demography.crown import CrownProfile, get_crown_xy
 from pyrealm.demography.canopy import Canopy
 from pyrealm.demography.t_model_functions import StemAllometry
@@ -117,9 +117,11 @@ simple_community = Community(
     flora=simple_flora,
     cell_area=total_area,
     cell_id=1,
-    cohort_dbh_values=stem_dbh,
-    cohort_n_individuals=np.array([1]),
-    cohort_pft_names=np.array(["defaults"]),
+    cohorts=Cohorts(
+        dbh_values=stem_dbh,
+        n_individuals=np.array([1]),
+        pft_names=np.array(["defaults"]),
+    ),
 )
 
 # Get the canopy model for the simple case from the canopy top
@@ -255,9 +257,11 @@ community = Community(
     flora=flora,
     cell_area=32,
     cell_id=1,
-    cohort_dbh_values=np.array([0.1, 0.20, 0.5]),
-    cohort_n_individuals=np.array([7, 3, 2]),
-    cohort_pft_names=np.array(["short", "short", "tall"]),
+    cohorts=Cohorts(
+        dbh_values=np.array([0.1, 0.20, 0.5]),
+        n_individuals=np.array([7, 3, 2]),
+        pft_names=np.array(["short", "short", "tall"]),
+    ),
 )
 
 # Calculate the canopy profile across vertical heights
@@ -281,7 +285,7 @@ profiles = get_crown_xy(
 for idx, crown in enumerate(profiles):
 
     # Get spaced but slightly randomized stem locations
-    n_stems = community.cohort_data["n_individuals"][idx]
+    n_stems = community.cohorts.n_individuals[idx]
     stem_locations = np.linspace(0, 10, num=n_stems) + np.random.normal(size=n_stems)
 
     # Plot the crown model for each stem
@@ -298,7 +302,7 @@ vertical profile is equal to the expected value across the whole community.
 ```{code-cell} ipython3
 # Calculate L_h for each cohort
 cohort_lai = (
-    community.cohort_data["n_individuals"]
+    community.cohorts.n_individuals
     * community.stem_traits.lai
     * community.stem_allometry.crown_area
 ) / community.cell_area
@@ -439,13 +443,13 @@ individuals in each cohort.
 ```{code-cell} ipython3
 # Calculate the total projected crown area across the community at each height
 community_crown_area = np.nansum(
-    canopy.crown_profile.projected_crown_area * community.cohort_data["n_individuals"],
+    canopy.crown_profile.projected_crown_area * community.cohorts.n_individuals,
     axis=1,
 )
 
 # Do the same for the projected leaf area
 community_leaf_area = np.nansum(
-    canopy.crown_profile.projected_leaf_area * community.cohort_data["n_individuals"],
+    canopy.crown_profile.projected_leaf_area * community.cohorts.n_individuals,
     axis=1,
 )
 ```
@@ -609,6 +613,6 @@ print(cohort_fapar)
    cohort to given the $f_{APAR}$ for each stem at each height.
 
 ```{code-cell} ipython3
-stem_fapar = cohort_fapar / community.cohort_data["n_individuals"]
+stem_fapar = cohort_fapar / community.cohorts.n_individuals
 print(stem_fapar)
 ```
diff --git a/docs/source/users/demography/community.md b/docs/source/users/demography/community.md
index 8bee4431..2486e3b8 100644
--- a/docs/source/users/demography/community.md
+++ b/docs/source/users/demography/community.md
@@ -36,7 +36,7 @@ import numpy as np
 import pandas as pd
 
 from pyrealm.demography.flora import PlantFunctionalType, Flora
-from pyrealm.demography.community import Community
+from pyrealm.demography.community import Cohorts, Community
 ```
 
 ```{code-cell} ipython3
@@ -57,16 +57,14 @@ community = Community(
     flora=flora,
     cell_area=32,
     cell_id=1,
-    cohort_dbh_values=np.array([0.02, 0.20, 0.5]),
-    cohort_n_individuals=np.array([15, 5, 2]),
-    cohort_pft_names=np.array(["short", "short", "tall"]),
+    cohorts=Cohorts(
+        dbh_values=np.array([0.02, 0.20, 0.5]),
+        n_individuals=np.array([15, 5, 2]),
+        pft_names=np.array(["short", "short", "tall"]),
+    ),
 )
 ```
 
 ```{code-cell} ipython3
 community
 ```
-
-```{code-cell} ipython3
-
-```

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 12/31] 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 13/31] 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 14/31] 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 15/31] 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 16/31] 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 17/31] 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 18/31] 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 ae29444cf855edff1ec8abd3a560355683cfb7fe Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Wed, 23 Oct 2024 13:50:09 +0100
Subject: [PATCH 19/31] Adding PandasExporter to community and canopy classes

---
 pyrealm/demography/canopy.py    | 5 +++--
 pyrealm/demography/community.py | 3 ++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/pyrealm/demography/canopy.py b/pyrealm/demography/canopy.py
index 40181ea0..cbfd88f8 100644
--- a/pyrealm/demography/canopy.py
+++ b/pyrealm/demography/canopy.py
@@ -7,6 +7,7 @@
 from scipy.optimize import root_scalar  # type: ignore [import-untyped]
 
 from pyrealm.demography.community import Community
+from pyrealm.demography.core import PandasExporter
 from pyrealm.demography.crown import (
     CrownProfile,
     _validate_z_qz_args,
@@ -166,7 +167,7 @@ def fit_perfect_plasticity_approximation(
 
 
 @dataclass
-class CohortCanopyData:
+class CohortCanopyData(PandasExporter):
     """Dataclass holding canopy data across cohorts.
 
     The cohort canopy data consists of a set of attributes represented as two
@@ -279,7 +280,7 @@ def allocate_fapar(
 
 
 @dataclass
-class CommunityCanopyData:
+class CommunityCanopyData(PandasExporter):
     """Dataclass holding community-wide canopy data.
 
     The community canopy data consists of a set of attributes represented as one
diff --git a/pyrealm/demography/community.py b/pyrealm/demography/community.py
index 8e629d6d..4c5ce15d 100644
--- a/pyrealm/demography/community.py
+++ b/pyrealm/demography/community.py
@@ -124,6 +124,7 @@
 from numpy.typing import NDArray
 
 from pyrealm.core.utilities import check_input_shapes
+from pyrealm.demography.core import PandasExporter
 from pyrealm.demography.flora import Flora, StemTraits
 from pyrealm.demography.t_model_functions import StemAllometry
 
@@ -136,7 +137,7 @@
 
 
 @dataclass
-class Cohorts:
+class Cohorts(PandasExporter):
     """A dataclass to hold data for a set of plant cohorts.
 
     The attributes should be numpy arrays of equal length, containing an entry for each

From 80a74ced116e02fa2d6843f1e2815dc86162ba54 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Wed, 23 Oct 2024 14:03:38 +0100
Subject: [PATCH 20/31] Adding tests of inherited to_pandas, making unstacking
 field into an index in to_pandas

---
 pyrealm/demography/core.py                    |  8 ++--
 tests/unit/demography/test_flora.py           | 45 +++++++++++++++++++
 .../unit/demography/test_t_model_functions.py | 34 +++++++++++---
 3 files changed, 76 insertions(+), 11 deletions(-)

diff --git a/pyrealm/demography/core.py b/pyrealm/demography/core.py
index f1d58091..1f84eda9 100644
--- a/pyrealm/demography/core.py
+++ b/pyrealm/demography/core.py
@@ -24,8 +24,8 @@ class PandasExporter(ABC):
     def to_pandas(self) -> pd.DataFrame:
         """Convert the instance array attributes into a {class}`pandas.DataFrame.
 
-        If the array values are two-dimensional (e.g. stems by heights), the data are
-        stacked and an index field is added.
+        If the array values are two-dimensional (i.e. stem or cohort data by vertical
+        heights), the data are stacked and an index is added.
         """
 
         # Extract the attributes into a dictionary
@@ -37,13 +37,13 @@ def to_pandas(self) -> pd.DataFrame:
         if len(data_shape) == 2:
             # create an index entry to show the column of each value
             stacked_data = {
-                "column_index": np.repeat(np.arange(data_shape[1]), data_shape[0])
+                "column_stem_index": np.repeat(np.arange(data_shape[1]), data_shape[0])
             }
             # Ravel the attribute data using column-major Fortan style
             for ky, vl in data.items():
                 stacked_data[ky] = np.ravel(vl, order="F")
 
-            return pd.DataFrame(stacked_data)
+            return pd.DataFrame(stacked_data).set_index("column_stem_index")
 
         return pd.DataFrame(data)
 
diff --git a/tests/unit/demography/test_flora.py b/tests/unit/demography/test_flora.py
index d9959350..45d325b4 100644
--- a/tests/unit/demography/test_flora.py
+++ b/tests/unit/demography/test_flora.py
@@ -6,6 +6,7 @@
 from json import JSONDecodeError
 
 import numpy as np
+import pandas as pd
 import pytest
 from marshmallow.exceptions import ValidationError
 from pandas.errors import ParserError
@@ -312,6 +313,11 @@ def test_flora_from_csv(filename, outcome):
                 assert nm in flora.name
 
 
+#
+# Test Flora methods
+#
+
+
 @pytest.mark.parametrize(
     argnames="pft_names,outcome",
     argvalues=[
@@ -334,3 +340,42 @@ def test_flora_get_stem_traits(fixture_flora, pft_names, outcome):
 
         for trt in stem_traits.trait_attrs:
             assert len(getattr(stem_traits, trt)) == len(pft_names)
+
+
+def test_Flora_to_pandas(fixture_flora):
+    """Test the inherited to_pandas method as applied to a Flora object."""
+
+    df = fixture_flora.to_pandas()
+
+    assert isinstance(df, pd.DataFrame)
+    assert df.shape == (fixture_flora.n_pfts, len(fixture_flora.array_attrs))
+    assert set(fixture_flora.array_attrs) == set(df.columns)
+
+
+#
+# Direct constructor for StemTraits
+#
+
+
+def test_StemTraits(fixture_flora):
+    """Basic test of StemTraits constructor and inherited to_pandas method."""
+    from pyrealm.demography.flora import StemTraits
+
+    # Construct some input data from the fixture
+    flora_df = fixture_flora.to_pandas()
+    args = {ky: np.concatenate([val, val]) for ky, val in flora_df.items()}
+
+    instance = StemTraits(**args)
+
+    # Very basic check that the result is as expected
+    assert len(instance.a_hd) == 2 * fixture_flora.n_pfts
+
+    # Test the to_pandas method here too
+    stem_traits_df = instance.to_pandas()
+
+    assert stem_traits_df.shape == (
+        2 * fixture_flora.n_pfts,
+        len(fixture_flora.array_attrs),
+    )
+
+    assert set(instance.array_attrs) == set(stem_traits_df.columns)
diff --git a/tests/unit/demography/test_t_model_functions.py b/tests/unit/demography/test_t_model_functions.py
index 661a8bf3..0279b141 100644
--- a/tests/unit/demography/test_t_model_functions.py
+++ b/tests/unit/demography/test_t_model_functions.py
@@ -728,7 +728,7 @@ def test_calculate_dbh_from_height_edge_cases():
 
 
 def test_StemAllometry(rtmodel_flora, rtmodel_data):
-    """Test the StemAllometry class."""
+    """Test the StemAllometry class and inherited methods."""
 
     from pyrealm.demography.t_model_functions import StemAllometry
 
@@ -736,15 +736,24 @@ def test_StemAllometry(rtmodel_flora, rtmodel_data):
         stem_traits=rtmodel_flora, at_dbh=rtmodel_data["dbh"][:, [0]]
     )
 
-    # Check the variables provided by the rtmodel implementation
+    # Check the values of the variables calculated against the expectations from the
+    # rtmodel implementation
     vars_to_check = (
-        v
-        for v in stem_allometry.allometry_attrs
-        if v not in ["crown_r0", "crown_z_max"]
+        v for v in stem_allometry.array_attrs if v not in ["crown_r0", "crown_z_max"]
     )
     for var in vars_to_check:
         assert np.allclose(getattr(stem_allometry, var), rtmodel_data[var])
 
+    # Test the inherited to_pandas method
+    df = stem_allometry.to_pandas()
+
+    assert df.shape == (
+        stem_allometry._n_stems * stem_allometry._n_pred,
+        len(stem_allometry.array_attrs),
+    )
+
+    assert set(stem_allometry.array_attrs) == set(df.columns)
+
 
 def test_StemAllocation(rtmodel_flora, rtmodel_data):
     """Test the StemAllometry class."""
@@ -761,9 +770,20 @@ def test_StemAllocation(rtmodel_flora, rtmodel_data):
         at_potential_gpp=rtmodel_data["potential_gpp"],
     )
 
-    # Check the variables provided by the rtmodel implementation
+    # Check the values of the variables calculated against the expectations from the
+    # rtmodel implementation
     vars_to_check = (
-        v for v in stem_allocation.allocation_attrs if v not in ["foliar_respiration"]
+        v for v in stem_allocation.array_attrs if v not in ["foliar_respiration"]
     )
     for var in vars_to_check:
         assert np.allclose(getattr(stem_allocation, var), rtmodel_data[var])
+
+    # Test the inherited to_pandas method
+    df = stem_allocation.to_pandas()
+
+    assert df.shape == (
+        stem_allocation._n_stems * stem_allocation._n_pred,
+        len(stem_allocation.array_attrs),
+    )
+
+    assert set(stem_allocation.array_attrs) == set(df.columns)

From b6a75939190ef462577c19f9f9f77da44b3f42e2 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Wed, 23 Oct 2024 15:38:20 +0100
Subject: [PATCH 21/31] Updating community docstrings to use to_pandas

---
 pyrealm/demography/community.py         | 55 +++++++++++++++++--------
 tests/unit/demography/test_community.py | 15 +++++--
 tests/unit/demography/test_crown.py     | 13 +++++-
 tests/unit/demography/test_flora.py     |  2 +-
 4 files changed, 62 insertions(+), 23 deletions(-)

diff --git a/pyrealm/demography/community.py b/pyrealm/demography/community.py
index 4c5ce15d..8a611c5e 100644
--- a/pyrealm/demography/community.py
+++ b/pyrealm/demography/community.py
@@ -91,22 +91,34 @@
 ...     ),
 ... )
 
-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.cohorts.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
-2   Evergreen Tree  0.120            150    11.436498    3.413238  13.581094
-3  Deciduous Shrub  0.025            180     1.858954    0.127752   0.082126
+The data in the Community class is stored under three attributes, each of which stores
+an instance of a dataclass holding related parts of the community data. All have a
+``to_pandas`` method that can be used to visualise and explore the data:
+
+>>> community.cohorts.to_pandas()
+   dbh_values  n_individuals        pft_names
+0       0.100            100   Evergreen Tree
+1       0.030            200  Deciduous Shrub
+2       0.120            150   Evergreen Tree
+3       0.025            180  Deciduous Shrub
+
+>>> community.stem_allometry.to_pandas()[
+...     ["stem_height", "crown_area", "stem_mass", "crown_r0", "crown_z_max"]
+... ]
+   stem_height  crown_area  stem_mass  crown_r0  crown_z_max
+0     9.890399    2.459835   8.156296  0.339477     7.789552
+1     2.110534    0.174049   0.134266  0.083788     1.642777
+2    11.436498    3.413238  13.581094  0.399890     9.007241
+3     1.858954    0.127752   0.082126  0.071784     1.446955
+
+>>> community.stem_traits.to_pandas()[
+...     ["name", "a_hd", "ca_ratio", "sla", "par_ext", "q_m",  "z_max_prop"]
+... ]
+              name   a_hd  ca_ratio   sla  par_ext       q_m  z_max_prop
+0   Evergreen Tree  120.0     380.0  12.0      0.6  2.606561    0.787587
+1  Deciduous Shrub  100.0     350.0  15.0      0.4  2.809188    0.778371
+2   Evergreen Tree  120.0     380.0  12.0      0.6  2.606561    0.787587
+3  Deciduous Shrub  100.0     350.0  15.0      0.4  2.809188    0.778371
 """  # noqa: D205
 
 from __future__ import annotations
@@ -115,7 +127,7 @@
 import sys
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import pandas as pd
@@ -144,9 +156,16 @@ class Cohorts(PandasExporter):
     cohort in the data class.
     """
 
+    # A class variable setting the attribute names of traits.
+    array_attrs: ClassVar[tuple[str, ...]] = tuple(
+        ["dbh_values", "n_individuals", "pft_names"]
+    )
+
+    # Instance attributes
     dbh_values: NDArray[np.float64]
     n_individuals: NDArray[np.int_]
     pft_names: NDArray[np.str_]
+    n_cohorts: int = field(init=False)
 
     def __post_init__(self) -> None:
         """Validation of cohorts data."""
@@ -168,6 +187,8 @@ def __post_init__(self) -> None:
         except ValueError:
             raise ValueError("Cohort arrays are of unequal length")
 
+        self.n_cohorts = len(self.dbh_values)
+
 
 class CohortSchema(Schema):
     """A validation schema for Cohort data objects.
diff --git a/tests/unit/demography/test_community.py b/tests/unit/demography/test_community.py
index cf7bd5b1..af894291 100644
--- a/tests/unit/demography/test_community.py
+++ b/tests/unit/demography/test_community.py
@@ -100,6 +100,13 @@ def test_Cohorts(args, outcome, excep_message):
         cohorts = Cohorts(**args)
         # trivial test of success
         assert len(cohorts.dbh_values) == 2
+
+        # test the to_pandas method
+        df = cohorts.to_pandas()
+
+        assert df.shape == (cohorts.n_cohorts, len(cohorts.array_attrs))
+        assert set(cohorts.array_attrs) == set(df.columns)
+
         return
 
     assert str(excep.value) == excep_message
@@ -215,10 +222,10 @@ def test_Community__init__(
     with outcome as excep:
         community = Community(**args, cohorts=cohorts, flora=fixture_flora)
 
-        if isinstance(outcome, does_not_raise):
-            # Simple test that data is loaded and trait and t model data calculated
-            check_expected(community=community, expected=fixture_expected)
-            return
+        # Simple test that data is loaded and trait and t model data calculated
+        check_expected(community=community, expected=fixture_expected)
+
+        return
 
     # Check exception message
     assert str(excep.value) == excep_message
diff --git a/tests/unit/demography/test_crown.py b/tests/unit/demography/test_crown.py
index a7cd10ec..0ab6612d 100644
--- a/tests/unit/demography/test_crown.py
+++ b/tests/unit/demography/test_crown.py
@@ -586,7 +586,8 @@ def test_CrownProfile(fixture_community):
     """Test the CrownProfile class.
 
     This implements a subset of the tests in the more detailed function checks above to
-    validate that this wrapper class works as intended.
+    validate that this wrapper class works as intended. It also tests the inherited
+    to_pandas method.
     """
 
     from pyrealm.demography.crown import CrownProfile
@@ -615,3 +616,13 @@ def test_CrownProfile(fixture_community):
         np.diag(crown_profile.projected_leaf_area),
         fixture_community.stem_allometry.crown_area,
     )
+
+    # Test the inherited to_pandas method
+    df = crown_profile.to_pandas()
+
+    assert df.shape == (
+        crown_profile._n_stems * crown_profile._n_pred,
+        len(crown_profile.array_attrs),
+    )
+
+    assert set(crown_profile.array_attrs) == set(df.columns)
diff --git a/tests/unit/demography/test_flora.py b/tests/unit/demography/test_flora.py
index 45d325b4..de70a5fb 100644
--- a/tests/unit/demography/test_flora.py
+++ b/tests/unit/demography/test_flora.py
@@ -338,7 +338,7 @@ def test_flora_get_stem_traits(fixture_flora, pft_names, outcome):
     with outcome:
         stem_traits = fixture_flora.get_stem_traits(pft_names=pft_names)
 
-        for trt in stem_traits.trait_attrs:
+        for trt in stem_traits.array_attrs:
             assert len(getattr(stem_traits, trt)) == len(pft_names)
 
 

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 22/31] 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 23/31] 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

From 3e3dbd8c3ebac5fdcbf3a414eaa8d60191673ac9 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Wed, 23 Oct 2024 17:41:48 +0100
Subject: [PATCH 24/31] Making the canopy dataclasses PandasExporter compliant
 and adding simple tests

---
 pyrealm/demography/canopy.py         | 18 ++++++++++++++++++
 tests/unit/demography/test_canopy.py | 22 +++++++++++++++++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/pyrealm/demography/canopy.py b/pyrealm/demography/canopy.py
index 006db57e..503a60fe 100644
--- a/pyrealm/demography/canopy.py
+++ b/pyrealm/demography/canopy.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import InitVar, dataclass, field
+from typing import ClassVar
 
 import numpy as np
 from numpy.typing import NDArray
@@ -209,6 +210,15 @@ class CohortCanopyData(PandasExporter):
         cell_area: A float setting the total canopy area available to the cohorts.
     """
 
+    array_attrs: ClassVar[tuple[str, ...]] = (
+        "stem_leaf_area",
+        "lai",
+        "f_trans",
+        "f_abs",
+        "cohort_fapar",
+        "stem_fapar",
+    )
+
     # 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
@@ -304,6 +314,14 @@ class CommunityCanopyData(PandasExporter):
             :attr:`CohortCanopyData.f_trans<pyrealm.demography.canopy.CohortCanopyData.f_trans>`.
     """
 
+    array_attrs: ClassVar[tuple[str, ...]] = (
+        "f_trans",
+        "f_abs",
+        "transmission_profile",
+        "extinction_profile",
+        "fapar",
+    )
+
     # Init vars
     cohort_transmissivity: InitVar[NDArray[np.float64]]
     """An array providing the per cohort light transmissivity at each of the required
diff --git a/tests/unit/demography/test_canopy.py b/tests/unit/demography/test_canopy.py
index b8b03105..3f78f65a 100644
--- a/tests/unit/demography/test_canopy.py
+++ b/tests/unit/demography/test_canopy.py
@@ -73,7 +73,7 @@ def test_CohortCanopyData__init__(
         assert np.allclose(instance.lai, exp_lai)
         assert np.allclose(instance.f_trans, exp_f_trans)
 
-        # Unpack and test expectations
+        # Unpack and test expectations for cohort and stem fapar
         exp_f_trans, exp_trans_prof = community_expected
         expected_fapar = -np.diff(exp_trans_prof, prepend=1)
         assert np.allclose(
@@ -83,6 +83,16 @@ def test_CohortCanopyData__init__(
             instance.stem_fapar, np.tile((expected_fapar / 6)[:, None], 3)
         )
 
+        # Test the inherited to_pandas method
+        df = instance.to_pandas()
+
+        assert df.shape == (
+            np.prod(exp_stem_leaf_area.shape),
+            len(instance.array_attrs),
+        )
+
+        assert set(instance.array_attrs) == set(df.columns)
+
     def test_CommunityCanopyData__init__(
         self, cohort_args, cohort_expected, community_expected
     ):
@@ -99,6 +109,16 @@ def test_CommunityCanopyData__init__(
         assert np.allclose(instance.f_trans, exp_f_trans)
         assert np.allclose(instance.transmission_profile, exp_trans_prof)
 
+        # Test the inherited to_pandas method
+        df = instance.to_pandas()
+
+        assert df.shape == (
+            len(exp_f_trans),
+            len(instance.array_attrs),
+        )
+
+        assert set(instance.array_attrs) == set(df.columns)
+
 
 def test_Canopy__init__():
     """Test happy path for initialisation.

From 0fa203b7e56071e09c39cdceff087c55c20e9db0 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Thu, 24 Oct 2024 10:31:00 +0100
Subject: [PATCH 25/31] Working around irritating sphinx bug on inherited
 methods links not being supported

---
 docs/source/users/demography/crown.md   |  6 +++---
 docs/source/users/demography/flora.md   | 12 +++++++-----
 docs/source/users/demography/t_model.md | 20 +++++++++++---------
 3 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/docs/source/users/demography/crown.md b/docs/source/users/demography/crown.md
index a2f156ab..4f648f38 100644
--- a/docs/source/users/demography/crown.md
+++ b/docs/source/users/demography/crown.md
@@ -274,9 +274,9 @@ above calculated at each height $z$:
 crown_profiles
 ```
 
-The {meth}`~pyrealm.demography.crown.CrownProfile.to_pandas` method can be used to
-extract the data into a table, with the separate stems identified by the column index
-field.
+The {meth}`~pyrealm.demography.core.PandasExporter.to_pandas()` method of the
+{meth}`~pyrealm.demography.crown.CrownProfile` class can be used to extract the data
+into a table, with the separate stems identified by the column index field.
 
 ```{code-cell} ipython3
 crown_profiles.to_pandas()
diff --git a/docs/source/users/demography/flora.md b/docs/source/users/demography/flora.md
index c32729c1..7c589049 100644
--- a/docs/source/users/demography/flora.md
+++ b/docs/source/users/demography/flora.md
@@ -146,9 +146,10 @@ flora = Flora([short_pft, medium_pft, tall_pft])
 flora
 ```
 
-The class provides a {meth}`~pyrealm.demography.StemTraits.to_pandas()` method to export
-the trait data as a {class}`pandas.DataFrame`, making it easier to use for plotting or
-calculations outside of `pyrealm`.
+The {meth}`~pyrealm.demography.core.PandasExporter.to_pandas()` method of the
+{meth}`~pyrealm.demography.t_model_functions.StemTraits` class exports the trait data as
+a {class}`pandas.DataFrame`, making it easier to use for plotting or calculations
+outside of `pyrealm`.
 
 ```{code-cell} ipython3
 flora.to_pandas()
@@ -174,8 +175,9 @@ stem_pfts = ["short", "short", "short", "medium", "medium", "tall"]
 stem_traits = flora.get_stem_traits(pft_names=stem_pfts)
 ```
 
-Again, the class provides the {meth}`~pyrealm.demography.StemTraits.to_pandas()` method
-to extract the data:
+Again, the {meth}`~pyrealm.demography.core.PandasExporter.to_pandas()` method of the
+{meth}`~pyrealm.demography.t_model_functions.StemTraits` class can be use to extract
+the data:
 
 ```{code-cell} ipython3
 stem_traits.to_pandas()
diff --git a/docs/source/users/demography/t_model.md b/docs/source/users/demography/t_model.md
index 2d47e886..44c0fd73 100644
--- a/docs/source/users/demography/t_model.md
+++ b/docs/source/users/demography/t_model.md
@@ -77,9 +77,9 @@ for each PFT. This then calculates a single estimate at the given size for each
 single_allometry = StemAllometry(stem_traits=flora, at_dbh=np.array([0.1, 0.1, 0.1]))
 ```
 
-The class provides a
-{meth}`~pyrealm.demography.t_model_functions.StemAllometry.to_pandas()` method
-to export the stem data for data exploration.
+The {meth}`~pyrealm.demography.t_model_functions.StemAllometry` class provides the
+{meth}`~pyrealm.demography.core.PandasExporter.to_pandas()` method to export the stem
+data for data exploration.
 
 ```{code-cell} ipython3
 single_allometry.to_pandas()
@@ -123,8 +123,9 @@ for ax, (var, ylab) in zip(axes.flatten(), plot_details):
         ax.legend(frameon=False)
 ```
 
-The {meth}`~pyrealm.demography.t_model_functions.StemAllometry.to_pandas()` method
-can still be used, but the values are stacked into columns along with a column index
+The {meth}`~pyrealm.demography.core.PandasExporter.to_pandas()` method of the
+{meth}`~pyrealm.demography.t_model_functions.StemAllometry` class can still be used, but
+the values are stacked into columns along with a index showing the different cohorts.
 
 ```{code-cell} ipython3
 allometries.to_pandas()
@@ -144,8 +145,8 @@ single_allocation = StemAllocation(
 single_allocation
 ```
 
-The class provides the
-{meth}`~pyrealm.demography.t_model_functions.StemAllocation.to_pandas()` method to
+The {meth}`~pyrealm.demography.core.PandasExporter.to_pandas()` method of the
+{meth}`~pyrealm.demography.t_model_functions.StemAllocation` class can be used to
 export data for exploration.
 
 ```{code-cell} ipython3
@@ -227,8 +228,9 @@ for ax, (var, ylab) in zip(axes, plot_details):
 fig.delaxes(axes[-1])
 ```
 
-As before, the {meth}`~pyrealm.demography.t_model_functions.StemAllometry.to_pandas()`
-method can be used to export the data for each stem:
+As before, the {meth}`~pyrealm.demography.core.PandasExporter.to_pandas()` method of the
+{meth}`~pyrealm.demography.t_model_functions.StemAllometry` classs can be used to export
+the data for each stem:
 
 ```{code-cell} ipython3
 allocation.to_pandas()

From aee8976b846e7309aff6df63a9010f31ebc2e0f2 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Thu, 24 Oct 2024 10:44:49 +0100
Subject: [PATCH 26/31] Add check for adding two CohortMethods subclasses of
 different type, test of check

---
 pyrealm/demography/core.py         | 12 +++++++---
 tests/unit/demography/test_core.py | 35 ++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/pyrealm/demography/core.py b/pyrealm/demography/core.py
index 1f84eda9..01ea9da1 100644
--- a/pyrealm/demography/core.py
+++ b/pyrealm/demography/core.py
@@ -48,7 +48,7 @@ def to_pandas(self) -> pd.DataFrame:
         return pd.DataFrame(data)
 
 
-class Cohorts(ABC):
+class CohortMethods(ABC):
     """Abstract base class implementing cohort manipulation functionality.
 
     Classes inheriting from this ABC must define a class attribute ``array_attrs`` that
@@ -64,12 +64,18 @@ class Cohorts(ABC):
 
     array_attrs: ClassVar[tuple[str, ...]]
 
-    def add_cohorts(self, add: Cohorts) -> None:
-        """Add array attributes from a second instance.
+    def add_cohorts(self, add: CohortMethods) -> None:
+        """Add array attributes from a second instance implementing the base class.
 
         Args:
             add: A second instance from which to add array attribute values.
         """
+
+        if not isinstance(add, self.__class__):
+            raise ValueError(
+                f"Cannot add {type(add).__name__} instance to {type(self).__name__}"
+            )
+
         for trait in self.array_attrs:
             setattr(
                 self,
diff --git a/tests/unit/demography/test_core.py b/tests/unit/demography/test_core.py
index 890622e1..d15f8491 100644
--- a/tests/unit/demography/test_core.py
+++ b/tests/unit/demography/test_core.py
@@ -7,6 +7,7 @@
 
 import numpy as np
 import pandas as pd
+import pytest
 from numpy.typing import NDArray
 
 
@@ -64,6 +65,40 @@ class TestClass(Cohorts):
     assert np.allclose(t1.b, np.arange(5, 9))
 
 
+def test_Cohorts_failure():
+    """Test the Cohorts abstract base class failure mode."""
+
+    from pyrealm.demography.core import CohortMethods
+
+    @dataclass
+    class TestClass(CohortMethods):
+        """Simple test class implementing the ABC."""
+
+        array_attrs: ClassVar[tuple[str, ...]] = ("a", "b")
+
+        a: NDArray[np.float64]
+        b: NDArray[np.float64]
+
+    @dataclass
+    class NotTheSameClass(CohortMethods):
+        """A different simple test class implementing the ABC."""
+
+        array_attrs: ClassVar[tuple[str, ...]] = ("c", "d")
+
+        c: NDArray[np.float64]
+        d: NDArray[np.float64]
+
+    # Create instances
+    t1 = TestClass(a=np.array([1, 2, 3]), b=np.array([4, 5, 6]))
+    t2 = NotTheSameClass(c=np.array([4, 5, 6]), d=np.array([7, 8, 9]))
+
+    # Check that adding a different
+    with pytest.raises(ValueError) as excep:
+        t1.add_cohorts(t2)
+
+    assert str(excep.value) == "Cannot add NotTheSameClass instance to TestClass"
+
+
 def test_PandasExporter_Cohorts_multiple_inheritance():
     """Test the behaviour of a class inheriting both core ABCs."""
 

From 375fe228154c6deac1e38bcbcf236ca32693c01b Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Thu, 24 Oct 2024 10:47:53 +0100
Subject: [PATCH 27/31] Updated demography.core.Cohorts to CohortMethods,
 forgot to update tests

---
 tests/unit/demography/test_core.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/unit/demography/test_core.py b/tests/unit/demography/test_core.py
index d15f8491..0ac76546 100644
--- a/tests/unit/demography/test_core.py
+++ b/tests/unit/demography/test_core.py
@@ -39,10 +39,10 @@ class TestClass(PandasExporter):
 def test_Cohorts():
     """Test the Cohorts abstract base class."""
 
-    from pyrealm.demography.core import Cohorts
+    from pyrealm.demography.core import CohortMethods
 
     @dataclass
-    class TestClass(Cohorts):
+    class TestClass(CohortMethods):
         """Simple test class implementing the ABC."""
 
         array_attrs: ClassVar[tuple[str, ...]] = ("a", "b")
@@ -102,10 +102,10 @@ class NotTheSameClass(CohortMethods):
 def test_PandasExporter_Cohorts_multiple_inheritance():
     """Test the behaviour of a class inheriting both core ABCs."""
 
-    from pyrealm.demography.core import Cohorts, PandasExporter
+    from pyrealm.demography.core import CohortMethods, PandasExporter
 
     @dataclass
-    class TestClass(Cohorts, PandasExporter):
+    class TestClass(CohortMethods, PandasExporter):
         """Test class with multiple inheritance."""
 
         array_attrs: ClassVar[tuple[str, ...]] = ("c", "d", "e")

From 6980eb8ad554696c7a48040ce16ba9808897aac7 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Thu, 24 Oct 2024 12:27:53 +0100
Subject: [PATCH 28/31] Adding CohortMethods base class to Cohorts and testing,
 renaming methods, test typing

---
 pyrealm/demography/community.py         | 21 ++++++++++++++++--
 pyrealm/demography/core.py              | 17 ++++++++++-----
 tests/unit/demography/test_community.py | 25 +++++++++++++++++++++
 tests/unit/demography/test_core.py      | 29 +++++++++++++++----------
 4 files changed, 73 insertions(+), 19 deletions(-)

diff --git a/pyrealm/demography/community.py b/pyrealm/demography/community.py
index 8a611c5e..12872ff8 100644
--- a/pyrealm/demography/community.py
+++ b/pyrealm/demography/community.py
@@ -136,7 +136,7 @@
 from numpy.typing import NDArray
 
 from pyrealm.core.utilities import check_input_shapes
-from pyrealm.demography.core import PandasExporter
+from pyrealm.demography.core import CohortMethods, PandasExporter
 from pyrealm.demography.flora import Flora, StemTraits
 from pyrealm.demography.t_model_functions import StemAllometry
 
@@ -149,7 +149,7 @@
 
 
 @dataclass
-class Cohorts(PandasExporter):
+class Cohorts(PandasExporter, CohortMethods):
     """A dataclass to hold data for a set of plant cohorts.
 
     The attributes should be numpy arrays of equal length, containing an entry for each
@@ -538,6 +538,23 @@ def from_toml(cls, path: Path, flora: Flora) -> Community:
 
         return cls(**file_data, flora=flora)
 
+    def add_cohorts(self, new_cohorts: Cohorts) -> None:
+        """Add a new set of cohorts to the community.
+
+        This method extends the ``cohorts`` attribute with the new cohort data and then
+        also extends the ``stem_traits`` and ``stem_allometry`` to match.
+        """
+
+        self.cohorts.add_cohort_data(new_cohorts)
+
+        new_stem_traits = self.flora.get_stem_traits(new_cohorts.pft_names)
+        self.stem_traits.add_cohort_data(new_stem_traits)
+
+        new_stem_allometry = StemAllometry(
+            stem_traits=new_stem_traits, at_dbh=new_cohorts.dbh_values
+        )
+        self.stem_allometry.add_cohort_data(new_stem_allometry)
+
     # @classmethod
     # def load_communities_from_csv(
     #     cls, cell_area: float, csv_path: str, flora: Flora
diff --git a/pyrealm/demography/core.py b/pyrealm/demography/core.py
index 01ea9da1..d8caf1bd 100644
--- a/pyrealm/demography/core.py
+++ b/pyrealm/demography/core.py
@@ -64,26 +64,31 @@ class CohortMethods(ABC):
 
     array_attrs: ClassVar[tuple[str, ...]]
 
-    def add_cohorts(self, add: CohortMethods) -> None:
+    def add_cohort_data(self, new_data: CohortMethods) -> None:
         """Add array attributes from a second instance implementing the base class.
 
         Args:
-            add: A second instance from which to add array attribute values.
+            new_data: A second instance from which to add cohort data to array attribute
+                values.
         """
 
-        if not isinstance(add, self.__class__):
+        # Check the incoming dataclass matches the calling instance.
+        if not isinstance(new_data, self.__class__):
             raise ValueError(
-                f"Cannot add {type(add).__name__} instance to {type(self).__name__}"
+                f"Cannot add cohort data from an {type(new_data).__name__} "
+                f"instance to {type(self).__name__}"
             )
 
+        # Concatenate the array attributes from the incoming instance to the calling
+        # instance.
         for trait in self.array_attrs:
             setattr(
                 self,
                 trait,
-                np.concatenate([getattr(self, trait), getattr(add, trait)]),
+                np.concatenate([getattr(self, trait), getattr(new_data, trait)]),
             )
 
-    def drop_cohorts(self, drop_indices: NDArray[np.int_]) -> None:
+    def drop_cohort_data(self, drop_indices: NDArray[np.int_]) -> None:
         """Drop array attribute values from an instance.
 
         Args:
diff --git a/tests/unit/demography/test_community.py b/tests/unit/demography/test_community.py
index af894291..6c6eeab6 100644
--- a/tests/unit/demography/test_community.py
+++ b/tests/unit/demography/test_community.py
@@ -112,6 +112,31 @@ def test_Cohorts(args, outcome, excep_message):
     assert str(excep.value) == excep_message
 
 
+def test_Cohorts_CohortMethods():
+    """Test the inherited CohortMethods methods."""
+
+    from pyrealm.demography.community import Cohorts
+
+    # Create and instance to modify using methods
+    cohorts = Cohorts(
+        pft_names=np.array(["broadleaf", "conifer"]),
+        n_individuals=np.array([6, 1]),
+        dbh_values=np.array([0.2, 0.5]),
+    )
+
+    # Check failure mode
+    with pytest.raises(ValueError) as excep:
+        cohorts.add_cohort_data(new_data=dict(a=1))
+
+    assert str(excep.value) == "Cannot add cohort data from an dict instance to Cohorts"
+
+    # Check success of adding and dropping data
+    cohorts.add_cohort_data(new_data=cohorts)
+    assert np.allclose(cohorts.dbh_values, np.array([0.2, 0.5, 0.2, 0.5]))
+    cohorts.drop_cohort_data(drop_indices=np.array([0, 2]))
+    assert np.allclose(cohorts.dbh_values, np.array([0.5, 0.5]))
+
+
 @pytest.mark.parametrize(
     argnames="args,cohort_data,outcome,excep_message",
     argvalues=[
diff --git a/tests/unit/demography/test_core.py b/tests/unit/demography/test_core.py
index 0ac76546..b9c743da 100644
--- a/tests/unit/demography/test_core.py
+++ b/tests/unit/demography/test_core.py
@@ -11,7 +11,7 @@
 from numpy.typing import NDArray
 
 
-def test_PandasExporter():
+def test_PandasExporter() -> None:
     """Test the PandasExporter abstract base class."""
 
     from pyrealm.demography.core import PandasExporter
@@ -26,7 +26,11 @@ class TestClass(PandasExporter):
         e: NDArray[np.float64]
 
     # create instance and run method
-    instance = TestClass(c=np.arange(5), d=np.arange(5), e=np.arange(5))
+    instance = TestClass(
+        c=np.arange(5, dtype=np.float64),
+        d=np.arange(5, dtype=np.float64),
+        e=np.arange(5, dtype=np.float64),
+    )
     pandas_out = instance.to_pandas()
 
     # simple checks of output class and behaviour
@@ -36,7 +40,7 @@ class TestClass(PandasExporter):
     assert np.allclose(pandas_out.sum(axis=0), np.repeat(10, 3))
 
 
-def test_Cohorts():
+def test_Cohorts() -> None:
     """Test the Cohorts abstract base class."""
 
     from pyrealm.demography.core import CohortMethods
@@ -55,17 +59,17 @@ class TestClass(CohortMethods):
     t2 = TestClass(a=np.array([4, 5, 6]), b=np.array([7, 8, 9]))
 
     # Add the t2 data into t1 and check the a and b attributes are extended
-    t1.add_cohorts(t2)
+    t1.add_cohort_data(t2)
     assert np.allclose(t1.a, np.arange(1, 7))
     assert np.allclose(t1.b, np.arange(4, 10))
 
     # Drop some indices and check the a and b attributes are truncated
-    t1.drop_cohorts(np.array([0, 5]))
+    t1.drop_cohort_data(np.array([0, 5]))
     assert np.allclose(t1.a, np.arange(2, 6))
     assert np.allclose(t1.b, np.arange(5, 9))
 
 
-def test_Cohorts_failure():
+def test_Cohorts_add_cohort_data_failure() -> None:
     """Test the Cohorts abstract base class failure mode."""
 
     from pyrealm.demography.core import CohortMethods
@@ -94,12 +98,15 @@ class NotTheSameClass(CohortMethods):
 
     # Check that adding a different
     with pytest.raises(ValueError) as excep:
-        t1.add_cohorts(t2)
+        t1.add_cohort_data(t2)
 
-    assert str(excep.value) == "Cannot add NotTheSameClass instance to TestClass"
+    assert (
+        str(excep.value)
+        == "Cannot add cohort data from an NotTheSameClass instance to TestClass"
+    )
 
 
-def test_PandasExporter_Cohorts_multiple_inheritance():
+def test_PandasExporter_Cohorts_multiple_inheritance() -> None:
     """Test the behaviour of a class inheriting both core ABCs."""
 
     from pyrealm.demography.core import CohortMethods, PandasExporter
@@ -136,7 +143,7 @@ def __post_init__(self, n: int, start_vals: NDArray[np.int_]) -> None:
     assert np.allclose(t1_out.sum(axis=0), np.repeat(15, 3) + np.array([0, 5, 10]))
 
     # Add the second set and check the results via pandas
-    t1.add_cohorts(t2)
+    t1.add_cohort_data(t2)
     t1_out_add = t1.to_pandas()
 
     # simple checks of output class and behaviour
@@ -146,7 +153,7 @@ def __post_init__(self, n: int, start_vals: NDArray[np.int_]) -> None:
     assert np.allclose(t1_out_add.sum(axis=0), np.repeat(36, 3) + np.array([0, 8, 16]))
 
     # Drop some entries and recheck
-    t1.drop_cohorts(np.array([0, 7]))
+    t1.drop_cohort_data(np.array([0, 7]))
     t1_out_drop = t1.to_pandas()
 
     # simple checks of output class and behaviour

From dbffac0c5bf438e318f9476ad97731d37086c67e Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Thu, 24 Oct 2024 13:31:44 +0100
Subject: [PATCH 29/31] Adding CohortMethods to StemTraits and StemAllometry

---
 pyrealm/demography/community.py               | 28 ++++++++++++----
 pyrealm/demography/core.py                    | 12 ++++++-
 pyrealm/demography/flora.py                   |  4 +--
 pyrealm/demography/t_model_functions.py       |  4 +--
 tests/unit/demography/test_flora.py           | 31 ++++++++++++++++++
 .../unit/demography/test_t_model_functions.py | 32 +++++++++++++++++++
 6 files changed, 100 insertions(+), 11 deletions(-)

diff --git a/pyrealm/demography/community.py b/pyrealm/demography/community.py
index 12872ff8..8b4eff7b 100644
--- a/pyrealm/demography/community.py
+++ b/pyrealm/demography/community.py
@@ -538,22 +538,38 @@ def from_toml(cls, path: Path, flora: Flora) -> Community:
 
         return cls(**file_data, flora=flora)
 
-    def add_cohorts(self, new_cohorts: Cohorts) -> None:
+    def drop_cohorts(self, drop_indices: NDArray[np.int_]) -> None:
+        """Drop cohorts from the community.
+
+        This method drops the identified cohorts from the ``cohorts`` attribute and then
+        removes their data from  the ``stem_traits`` and ``stem_allometry`` attributes
+        to match.
+        """
+
+        self.cohorts.drop_cohort_data(drop_indices=drop_indices)
+        self.stem_traits.drop_cohort_data(drop_indices=drop_indices)
+        self.stem_allometry.drop_cohort_data(drop_indices=drop_indices)
+
+    def add_cohorts(self, new_data: Cohorts) -> None:
         """Add a new set of cohorts to the community.
 
         This method extends the ``cohorts`` attribute with the new cohort data and then
         also extends the ``stem_traits`` and ``stem_allometry`` to match.
+
+        Args:
+            new_data: An instance of :class:`~pyrealm.demography.community.Cohorts`
+                containing cohort data to add to the community.
         """
 
-        self.cohorts.add_cohort_data(new_cohorts)
+        self.cohorts.add_cohort_data(new_data=new_data)
 
-        new_stem_traits = self.flora.get_stem_traits(new_cohorts.pft_names)
-        self.stem_traits.add_cohort_data(new_stem_traits)
+        new_stem_traits = self.flora.get_stem_traits(pft_names=new_data.pft_names)
+        self.stem_traits.add_cohort_data(new_data=new_stem_traits)
 
         new_stem_allometry = StemAllometry(
-            stem_traits=new_stem_traits, at_dbh=new_cohorts.dbh_values
+            stem_traits=new_stem_traits, at_dbh=new_data.dbh_values
         )
-        self.stem_allometry.add_cohort_data(new_stem_allometry)
+        self.stem_allometry.add_cohort_data(new_data=new_stem_allometry)
 
     # @classmethod
     # def load_communities_from_csv(
diff --git a/pyrealm/demography/core.py b/pyrealm/demography/core.py
index d8caf1bd..6b31abb2 100644
--- a/pyrealm/demography/core.py
+++ b/pyrealm/demography/core.py
@@ -95,5 +95,15 @@ def drop_cohort_data(self, drop_indices: NDArray[np.int_]) -> None:
             drop_indices: An array of integer indices to drop from each array attribute.
         """
 
+        # TODO - Probably part of tackling #317
+        #        The delete axis=0 here is tied to the case of dropping rows from 2D
+        #        arrays, but then I'm thinking it makes more sense to _only_ support 2D
+        #        arrays rather than the current mixed bag of getting a 1D array when a
+        #        single height is provided. Promoting that kind of input to 2D and then
+        #        enforcing an identical internal structure seems better.
+        #      - But! Trait data does not have 2 dimensions!
+        #      - Also to check here - this can lead to empty instances, which probably
+        #        are a thing we want, if mortality removes all cohorts.
+
         for trait in self.array_attrs:
-            setattr(self, trait, np.delete(getattr(self, trait), drop_indices))
+            setattr(self, trait, np.delete(getattr(self, trait), drop_indices, axis=0))
diff --git a/pyrealm/demography/flora.py b/pyrealm/demography/flora.py
index 35927a06..481c512f 100644
--- a/pyrealm/demography/flora.py
+++ b/pyrealm/demography/flora.py
@@ -37,7 +37,7 @@
 from marshmallow.exceptions import ValidationError
 from numpy.typing import NDArray
 
-from pyrealm.demography.core import PandasExporter
+from pyrealm.demography.core import CohortMethods, PandasExporter
 
 if sys.version_info[:2] >= (3, 11):
     import tomllib
@@ -443,7 +443,7 @@ def get_stem_traits(self, pft_names: NDArray[np.str_]) -> StemTraits:
 
 
 @dataclass()
-class StemTraits(PandasExporter):
+class StemTraits(PandasExporter, CohortMethods):
     """A dataclass for stem traits.
 
     This dataclass is used to provide arrays of plant functional type (PFT) traits
diff --git a/pyrealm/demography/t_model_functions.py b/pyrealm/demography/t_model_functions.py
index 0b51becb..b885dd36 100644
--- a/pyrealm/demography/t_model_functions.py
+++ b/pyrealm/demography/t_model_functions.py
@@ -12,7 +12,7 @@
 from numpy.typing import NDArray
 
 from pyrealm.core.utilities import check_input_shapes
-from pyrealm.demography.core import PandasExporter
+from pyrealm.demography.core import CohortMethods, PandasExporter
 from pyrealm.demography.flora import Flora, StemTraits
 
 
@@ -660,7 +660,7 @@ def calculate_growth_increments(
 
 
 @dataclass
-class StemAllometry(PandasExporter):
+class StemAllometry(PandasExporter, CohortMethods):
     """Calculate T Model allometric predictions across a set of stems.
 
     This method calculate predictions of stem allometries for stem height, crown area,
diff --git a/tests/unit/demography/test_flora.py b/tests/unit/demography/test_flora.py
index de70a5fb..989bcd55 100644
--- a/tests/unit/demography/test_flora.py
+++ b/tests/unit/demography/test_flora.py
@@ -379,3 +379,34 @@ def test_StemTraits(fixture_flora):
     )
 
     assert set(instance.array_attrs) == set(stem_traits_df.columns)
+
+
+def test_StemTraits_CohortMethods(fixture_flora):
+    """Test the StemTraits inherited cohort methods."""
+
+    from pyrealm.demography.t_model_functions import StemTraits
+
+    # Construct some input data with duplicate PFTs by doubling the fixture_flora data
+    flora_df = fixture_flora.to_pandas()
+    args = {ky: np.concatenate([val, val]) for ky, val in flora_df.items()}
+
+    stem_traits = StemTraits(**args)
+
+    # Check failure mode
+    with pytest.raises(ValueError) as excep:
+        stem_traits.add_cohort_data(new_data=dict(a=1))
+
+    assert (
+        str(excep.value) == "Cannot add cohort data from an dict instance to StemTraits"
+    )
+
+    # Check success of adding and dropping data
+    # Add a copy of itself as new cohort data and check the shape
+    stem_traits.add_cohort_data(new_data=stem_traits)
+    assert stem_traits.h_max.shape == (4 * fixture_flora.n_pfts,)
+    assert stem_traits.h_max.sum() == 4 * flora_df["h_max"].sum()
+
+    # Remove all but the first two rows and what's left should be aligned with the
+    # original data
+    stem_traits.drop_cohort_data(drop_indices=np.arange(2, 8))
+    assert np.allclose(stem_traits.h_max, flora_df["h_max"])
diff --git a/tests/unit/demography/test_t_model_functions.py b/tests/unit/demography/test_t_model_functions.py
index 0279b141..aa97cdd6 100644
--- a/tests/unit/demography/test_t_model_functions.py
+++ b/tests/unit/demography/test_t_model_functions.py
@@ -755,6 +755,38 @@ def test_StemAllometry(rtmodel_flora, rtmodel_data):
     assert set(stem_allometry.array_attrs) == set(df.columns)
 
 
+def test_StemAllometry_CohortMethods(rtmodel_flora, rtmodel_data):
+    """Test the StemAllometry inherited cohort methods."""
+
+    from pyrealm.demography.t_model_functions import StemAllometry
+
+    stem_allometry = StemAllometry(
+        stem_traits=rtmodel_flora, at_dbh=rtmodel_data["dbh"][:, [0]]
+    )
+    check_data = stem_allometry.crown_fraction.copy()
+
+    # Check failure mode
+    with pytest.raises(ValueError) as excep:
+        stem_allometry.add_cohort_data(new_data=dict(a=1))
+
+    assert (
+        str(excep.value)
+        == "Cannot add cohort data from an dict instance to StemAllometry"
+    )
+
+    # Check success of adding and dropping data
+    n_entries = len(rtmodel_data["dbh"])
+    # Add a copy of itself as new cohort data and check the shape
+    stem_allometry.add_cohort_data(new_data=stem_allometry)
+    assert stem_allometry.crown_fraction.shape == (2 * n_entries, rtmodel_flora.n_pfts)
+    assert stem_allometry.crown_fraction.sum() == 2 * check_data.sum()
+
+    # Remove the rows from the first copy and what's left should be aligned with the
+    # original data
+    stem_allometry.drop_cohort_data(drop_indices=np.arange(n_entries))
+    assert np.allclose(stem_allometry.crown_fraction, check_data)
+
+
 def test_StemAllocation(rtmodel_flora, rtmodel_data):
     """Test the StemAllometry class."""
 

From 1c30d7b0189e08cbc981547285c619952eb36e22 Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Thu, 24 Oct 2024 14:27:05 +0100
Subject: [PATCH 30/31] Testing the Community level add and drop methods

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

diff --git a/tests/unit/demography/test_community.py b/tests/unit/demography/test_community.py
index 6c6eeab6..bc6620a7 100644
--- a/tests/unit/demography/test_community.py
+++ b/tests/unit/demography/test_community.py
@@ -256,6 +256,49 @@ def test_Community__init__(
     assert str(excep.value) == excep_message
 
 
+def test_Community_add_and_drop(fixture_flora):
+    """Tests the add and drop cohort methods."""
+
+    from pyrealm.demography.community import Cohorts, Community
+
+    # Build the cohorts object, with two cohorts in the same order as the two PFTs in
+    # the fixture flora.
+    cohorts = Cohorts(
+        pft_names=fixture_flora.name,
+        n_individuals=np.array([6, 1]),
+        dbh_values=np.array([0.2, 0.5]),
+    )
+    community = Community(cell_id=1, cell_area=32, flora=fixture_flora, cohorts=cohorts)
+
+    # Check the initial state of the three attributes that should be modified
+    assert np.allclose(community.cohorts.n_individuals, np.array([6, 1]))
+    assert np.allclose(community.stem_traits.h_max, fixture_flora.h_max)
+    assert np.allclose(community.stem_allometry.dbh, np.array([0.2, 0.5]))
+
+    # Add a new set of cohorts
+    new_cohorts = Cohorts(
+        pft_names=fixture_flora.name,
+        n_individuals=np.array([8, 2]),
+        dbh_values=np.array([0.3, 0.6]),
+    )
+    community.add_cohorts(new_cohorts)
+
+    # Test the three attributes again to check they've all been doubled.
+    assert np.allclose(community.cohorts.n_individuals, np.array([6, 1, 8, 2]))
+    assert np.allclose(community.stem_traits.h_max, np.tile(fixture_flora.h_max, 2))
+    assert np.allclose(community.stem_allometry.dbh, np.array([0.2, 0.5, 0.3, 0.6]))
+
+    # Drop some rows
+    community.drop_cohorts(drop_indices=np.array([1, 3]))
+
+    # Test the three attributes again to check they've all been reduced.
+    assert np.allclose(community.cohorts.n_individuals, np.array([6, 8]))
+    assert np.allclose(
+        community.stem_traits.h_max, np.repeat(fixture_flora.h_max[0], 2)
+    )
+    assert np.allclose(community.stem_allometry.dbh, np.array([0.2, 0.3]))
+
+
 @pytest.mark.parametrize(
     argnames="file_data,outcome,excep_message",
     argvalues=[

From 416ac44cdb8a960622fdc822c62548a52a5bd9fd Mon Sep 17 00:00:00 2001
From: David Orme <davidorme@users.noreply.github.com>
Date: Thu, 24 Oct 2024 14:39:00 +0100
Subject: [PATCH 31/31] Adding demography.core module into API docs, fixing
 broken docs references

---
 docs/source/api/demography_api.md     |  8 ++++++++
 docs/source/users/demography/flora.md |  9 ++++-----
 pyrealm/demography/core.py            | 14 +++++++++++++-
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/docs/source/api/demography_api.md b/docs/source/api/demography_api.md
index db9dd62e..0630e5a9 100644
--- a/docs/source/api/demography_api.md
+++ b/docs/source/api/demography_api.md
@@ -30,6 +30,14 @@ language_info:
     :members:
 ```
 
+## The {mod}`~pyrealm.demography.core` module
+
+```{eval-rst}
+.. automodule:: pyrealm.demography.core
+    :autosummary:
+    :members:
+```
+
 ## The {mod}`~pyrealm.demography.flora` module
 
 ```{eval-rst}
diff --git a/docs/source/users/demography/flora.md b/docs/source/users/demography/flora.md
index 7c589049..e1de1b3b 100644
--- a/docs/source/users/demography/flora.md
+++ b/docs/source/users/demography/flora.md
@@ -147,9 +147,9 @@ flora
 ```
 
 The {meth}`~pyrealm.demography.core.PandasExporter.to_pandas()` method of the
-{meth}`~pyrealm.demography.t_model_functions.StemTraits` class exports the trait data as
-a {class}`pandas.DataFrame`, making it easier to use for plotting or calculations
-outside of `pyrealm`.
+{meth}`~pyrealm.demography.flora.StemTraits` class exports the trait data as a
+{class}`pandas.DataFrame`, making it easier to use for plotting or calculations outside
+of `pyrealm`.
 
 ```{code-cell} ipython3
 flora.to_pandas()
@@ -176,8 +176,7 @@ stem_traits = flora.get_stem_traits(pft_names=stem_pfts)
 ```
 
 Again, the {meth}`~pyrealm.demography.core.PandasExporter.to_pandas()` method of the
-{meth}`~pyrealm.demography.t_model_functions.StemTraits` class can be use to extract
-the data:
+{meth}`~pyrealm.demography.flora.StemTraits` class can be use to extract the data:
 
 ```{code-cell} ipython3
 stem_traits.to_pandas()
diff --git a/pyrealm/demography/core.py b/pyrealm/demography/core.py
index 6b31abb2..3d149be6 100644
--- a/pyrealm/demography/core.py
+++ b/pyrealm/demography/core.py
@@ -1,4 +1,16 @@
-"""Core shared functionality for the {mod}`~pyrealm.demography` module."""
+"""Core shared functionality for the {mod}`~pyrealm.demography` module.
+
+This module implements two abstract base classes that are used to share core methods
+across demography classes:
+
+* {class}`~pyrealm.demography.core.PandasExporter` provides the utility
+  {meth}`~pyrealm.demography.core.PandasExporter.to_pandas` method for extracting data
+  from demography classes for plotting and exploring data.
+* {class}`~pyrealm.demography.core.CohortMethods` provides the utility
+  {meth}`~pyrealm.demography.core.CohortMethods.add_cohort_data` and
+  {meth}`~pyrealm.demography.core.CohortMethods.drop_cohort_data` methods that are used
+  to append new cohort data across some demography dataclasses.
+"""
 
 from __future__ import annotations