Review validation and shapes in demography

docs/source/users/demography/canopy.md

@@ -110,7 +110,7 @@ stem_dbh = np.array([0.5])
 simple_stem = StemAllometry(stem_traits=simple_flora, at_dbh=stem_dbh)
 
 # The total area is exactly the crown area
-total_area = simple_stem.crown_area[0]
+total_area = simple_stem.crown_area[0][0]
 
 # Define a simple community
 simple_community = Community(
@@ -126,7 +126,7 @@ simple_community = Community(
 
 # Get the canopy model for the simple case from the canopy top
 # to the ground
-hghts = np.linspace(simple_stem.stem_height[0], 0, num=101)[:, None]
+hghts = np.linspace(simple_stem.stem_height[0][0], 0, num=101)[:, None]
 simple_canopy = Canopy(
     community=simple_community,
     layer_heights=hghts,

docs/source/users/demography/crown.md

@@ -333,7 +333,7 @@ for pft_idx, offset, colour in zip((0, 1, 2), (0, 5, 12), ("r", "g", "b")):
 
     ax.plot(
         [offset - stem_rz_max, offset + stem_rz_max],
-        [allometry.crown_z_max[pft_idx]] * 2,
+        [allometry.crown_z_max[:, pft_idx]] * 2,
         color=colour,
         linewidth=1,
         linestyle=":",
@@ -452,8 +452,8 @@ for f_g in np.linspace(0, 1, num=11):
 
 # Add a horizontal line for z_max
 ax.plot(
-    [-1, allometry_f_g.crown_area[0] + 1],
-    [allometry_f_g.crown_z_max, allometry_f_g.crown_z_max],
+    [-1, allometry_f_g.crown_area[0][0] + 1],
+    [allometry_f_g.crown_z_max[0][0], allometry_f_g.crown_z_max[0][0]],
     linestyle="--",
     color="black",
     label="$z_{max}$",
@@ -462,7 +462,7 @@ ax.plot(
 
 ax.set_ylabel(r"Vertical height ($z$, m)")
 ax.set_xlabel(r"Projected leaf area ($\tilde{A}_{cp}(z)$, m2)")
-ax.legend(frameon=False)
+_ = ax.legend(frameon=False)
 ```
 
 ## Plotting tools for crown shapes
@@ -525,7 +525,7 @@ for cr_xy, (ch, cpr), (lh, lpr) in zip(
     ax.plot(lpr, lh, color="red", linewidth=1)
 
 ax.set_aspect(0.5)
-plt.legend(
+_ = plt.legend(
     handles=[
         Patch(color="lightgrey", label="Crown profile"),
         Line2D([0], [0], label="Projected crown", color="0.4", linewidth=2),
@@ -534,5 +534,6 @@ plt.legend(
     ncols=3,
     loc="upper center",
     bbox_to_anchor=(0.5, 1.15),
+    frameon=False,
 )
 ```

pyrealm/demography/canopy.py

@@ -10,10 +10,9 @@
 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.core import PandasExporter, _validate_demography_array_arguments
 from pyrealm.demography.crown import (
     CrownProfile,
-    _validate_z_qz_args,
     calculate_relative_crown_radius_at_z,
     calculate_stem_projected_crown_area_at_z,
 )
@@ -66,9 +65,14 @@ def solve_canopy_area_filling_height(
     z_arr = np.array(z)
 
     if validate:
-        _validate_z_qz_args(
-            z=z_arr,
-            stem_properties=[n_individuals, crown_area, stem_height, m, n, q_m, z_max],
+        _validate_demography_array_arguments(
+            trait_args={"m": m, "n": n, "q_m": q_m, "n_individuals": n_individuals},
+            size_args={
+                "z": z_arr,
+                "crown_area": crown_area,
+                "stem_height": stem_height,
+                "z_max": z_max,
+            },
         )
 
     q_z = calculate_relative_crown_radius_at_z(

pyrealm/demography/community.py

@@ -104,12 +104,13 @@
 
 >>> 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
+... ]  # doctest: +NORMALIZE_WHITESPACE
+                   stem_height  crown_area  stem_mass  crown_r0  crown_z_max
+column_stem_index
+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"]

pyrealm/demography/core.py

@@ -92,12 +92,17 @@ def add_cohort_data(self, new_data: CohortMethods) -> None:
             )
 
         # Concatenate the array attributes from the incoming instance to the calling
-        # instance.
+        # instance. This need to respect the attribute array dimensions. If the
+        # attribute is one dimensional (e.g. traits), then concatenate on axis=0, but
+        # for 2 traits, need to concatenate on axis 1 to extend the trait axis.
         for trait in self.array_attrs:
+            current = getattr(self, trait)
             setattr(
                 self,
                 trait,
-                np.concatenate([getattr(self, trait), getattr(new_data, trait)]),
+                np.concatenate(
+                    [current, getattr(new_data, trait)], axis=(current.ndim - 1)
+                ),
             )
 
     def drop_cohort_data(self, drop_indices: NDArray[np.int_]) -> None:
@@ -107,15 +112,152 @@ 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.
+        # Drop from each trait along the last dimension - handles 2D height x stem and
+        # 1D stem traits.
 
         for trait in self.array_attrs:
-            setattr(self, trait, np.delete(getattr(self, trait), drop_indices, axis=0))
+            current = getattr(self, trait)
+            setattr(
+                self, trait, np.delete(current, drop_indices, axis=(current.ndim - 1))
+            )
+
+
+def _validate_demography_array_arguments(
+    trait_args: dict[str, NDArray],
+    size_args: dict[str, NDArray] = {},
+    at_size_args: dict[str, NDArray] = {},
+) -> None:
+    """Shared validation for demography inputs.
+
+    Trait arguments of functions should always be 1 dimensional arrays (row arrays) or
+    arrays of size 1 ('scalar'), representing trait values that are constant within a
+    cohort or stem. If multiple traits are being validated, then they should all have
+    have the same shape or be scalar.
+
+    In addition to simple validation of trait inputs, many demographic functions make
+    predictions of stem allometries, allocation and crown profile at a range of sizes.
+    These size arguments provide values at which to estimate predictions across stems
+    with different traits and  could represent different stem sizes (e.g. dbh) or
+    different stem heights at which to evaluate crown profiles. If size arguments are
+    provided, then they are also validated. All size arguments must have the same shape
+    or be scalar. Given that the stem traits provide N values, the shape of the size
+    arguments can then be:
+
+    * A scalar array (i.e. with size 1, such as np.array(1) or np.array([1])), that
+      provides a single size at which to calculate predictions for all N stems.
+    * A one dimensional array with identical shape to the stem properties (``(N,)``)
+      that provides individual single values for each stem.
+    * A two dimensional column vector (i.e. with shape ``(M, 1)``) that provides a set
+      of ``M`` values at which to calculate multiple predictions for all stem traits.
+    * A two-dimensional array ``(M, N)`` that provides individual predictions for each
+      stem.
+
+    Lastly, some functions generate predictions for stems at a particular size, given a
+    third input. For example, the stem allocation process evaluates how potential GPP is
+    allocated for a stem of a given size. (TBD - add crown example). If provided, these
+    ``at_size_args`` values are validated as follows:
+
+    * if ``z`` is a row array, ``q_z`` must then have identical shape, or
+    * if ``z`` is a column array ``(N, 1)``, ``q_z`` must then have shape ``(N,
+        n_stem_properties``).
+
+    Args:
+        trait_args: A dictionary of row arrays representing trait values, keyed by the
+            trait names.
+        size_args: A dictionary of arrays representing size values for stem allometry or
+            canopy height at which to evaluate functions, keyed by the value name.
+        at_size_args: A dictionary of arrays providing values that are to be evaluated
+            given predictions for a set of traits at a given size and which must be
+            congruent with both trait and size arguments. The dictionary should be keyed
+            with the value name.
+    """
+
+    # NOTE - this validation deliberately does not use check_input_shapes because it is
+    # insisting on _identical_ shapes or scalar arrays is too restrictive. See
+    # discussion in https://github.com/ImperialCollegeLondon/pyrealm/pull/342
+
+    # Check PFT inputs are all equal sized 1D row arrays or a mix of 1D rows and scalar
+    # values
+    try:
+        trait_args_shape = np.broadcast_shapes(
+            *[arr.shape for arr in trait_args.values()]
+        )
+    except ValueError:
+        raise ValueError(
+            f"Trait arguments are not equal shaped or "
+            f"scalar: {','.join(trait_args.keys())}"
+        )
+
+    if len(trait_args_shape) > 1:
+        raise ValueError(
+            f"Trait arguments are not 1D arrays: {','.join(trait_args.keys())}"
+        )
+
+    # Check the size inputs are broadcastable if they are provided.
+    if size_args:
+        try:
+            size_args_shape = np.broadcast_shapes(
+                *[arr.shape for arr in size_args.values()]
+            )
+        except ValueError:
+            raise ValueError(
+                f"Size arguments are not equal shaped or "
+                f"scalar: {','.join(size_args.keys())}"
+            )
+
+        # Test whether the shape of the size args are compatible with the traits args
+        try:
+            trait_size_shape = np.broadcast_shapes(trait_args_shape, size_args_shape)
+        except ValueError:
+            raise ValueError(
+                f"The array shapes of the trait {trait_args_shape} and "
+                f"size {size_args_shape} arguments are not congruent."
+            )
+
+    # Now check at size args if provided
+    if at_size_args and not size_args:
+        raise ValueError("Only provide `at_size_args` when `size_args` also provided.")
+
+    if at_size_args:
+        # Are the at_size values broadcastable?
+        try:
+            at_size_args_shape = np.broadcast_shapes(
+                *[arr.shape for arr in at_size_args.values()]
+            )
+        except ValueError:
+            raise ValueError(
+                f"At size arguments are not equal shaped or "
+                f"scalar: {','.join(at_size_args.keys())}"
+            )
+
+        # Are they congruent with the trait and size values.
+        try:
+            _ = np.broadcast_shapes(trait_size_shape, at_size_args_shape)
+        except ValueError:
+            raise ValueError(
+                f"The broadcast shapes of the trait and size arguments "
+                f"{trait_size_shape} are not congruent with the shape of the at_size "
+                f"arguments {at_size_args_shape}."
+            )
+
+
+def _enforce_2D(array: NDArray) -> NDArray:
+    """Utility conversion to force two dimensional outputs.
+
+    Depending on the input dimensions, the calculations in the T Model and othe
+    demography models can return scalar, one dimensional or two dimensional arrays. This
+    utility function is used to give a consistent output array dimensionality.
+
+    Args:
+        array: A numpy array
+    """
+
+    match array.ndim:
+        case 0:
+            return array[None, None]
+        case 1:
+            return array[None, :]
+        case 2:
+            return array
+        case _:
+            raise ValueError("Demography array of more than 2 dimensions.")