Revise cohort structure inside community

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)
 ```

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
-
-```

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

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
@@ -95,7 +97,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,
@@ -111,7 +113,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 +135,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,26 +259,27 @@ class CommunityStructuredDataSchema(Schema):
     )
 
     @post_load
-    def cohort_objects_to_arrays(self, data: dict, **kwargs: Any) -> dict:
+    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
         """
 
-        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 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"]]),
+            ),
         )
-        data["cohort_pft_names"] = np.array([c["pft_name"] for c in data["cohorts"]])
-
-        del data["cohorts"]
-
-        return data
 
 
 class CommunityCSVDataSchema(Schema):
@@ -300,21 +336,23 @@ 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
-        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"])
-
-        return data
+        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"]),
+            ),
+        )
 
 
 @dataclass
@@ -353,21 +391,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 +414,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

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]),
+        ),
     )