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] 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.