Demography dataclasses and signatures

[davidorme]

With the current WIP pull requests #296, there is now quite a lot of structure and function in the demography modules. The implementation in #296 has a good chunk of the functionality we want but the internal structures are a little creaky. There is a whole lot of dictionaries of variables and also a lot of symmetries in the data structure that we aren't making use of.

So - I moaned about this in #296. Given what is in that PR, I'd rather aim towards something where we have:

• better documented and separated data structures,
• better clarity about what data is required where
• clearer reuse of code with similar use cases.
• Attributes rather than lots of dictionary names:
  • Clearer to read
  • Deeper typing
  • Autocompletion etc in VSCode

Here's a more thought through sketch of some signatures.

Flora

The Flora class is a dict of PlantFunctionalType instances. I think originally that was primary use mode we envisioned, but actually all of the internal usage is of the Flora.data attribute (which is a dict of arrays). On top of that, the internals of the Community and Canopy make extensive use of arrays of PlantFunctionalType traits. So I think we have basically four objects that have extremely similar internals.

@dataclass(frozen=True)
class PlantFunctionalTypeStrict:
    """PFT with all traits explicitly declared - mostly used for deserialising."""
    name: str
    a_hd: float
    ...
    f_g: float

@dataclass(frozen=True)
class PlantFunctionalType(PlantFunctionalTypeStrict):
    """PFT with default values, used as a programatic convenience."""
    a_hd: float = 116.0
    ...
    f_g: float = 0.05

@dataclass(frozen=True)
class Flora():
    """
    * Is constructed from a list of PFTs.
    * Guarantees unique PFT names
    * Otherwise has all the same trait attributes as PFTS but as row arrays.
    * Also has a bonus extra dictionary onto the original PFT instances
    """
    pfts: InitVar[list[PlantFunctionalTypeStrict]]
    name: NDArray[np.str_] = field(post_init=True)
    a_hd: NDArray[np.float32] = field(post_init=True)
    ...
    f_g: NDArray[np.float32]  = field(post_init=True)
    pft_dict: dict[str, PlantFunctionalTypeStrict] = field(post_init=True)

@dataclass()
class StemTraits():
    """
    * Core data class for Communities - holds the stem traits for each cohort
    * Provides all the same trait attributes as row arrays.
    """
    name: NDArray[np.str_] 
    a_hd: NDArray[np.float32]
    ...
    f_g: NDArray[np.float32]

Allometries

The T Model code can then keep the individual functions with low level signatures but provide a high level API that generates a StemAllometry object:

@dataclass()
class StemAllometry():
    """
    * Data class containing allometric predictions from stem traits.
    * Core data class for Communities - holds actual stem allometries
    """
    stem_height: NDArray[np.float32]
    ...
    canopy_z_max: NDArray[np.float32]


def get_stem_allometries(stems: StemTraits | Flora, dbh: NDArray) -> StemAllometry:
    """
    * Generates predictions of stem allometries at given stem sizes for either:
        * StemTraits for community simulation
        * Flora for trait exploration
    """

Allocation model

The allocation model takes specific stem allometries and calculates the expected allocation of GPP across the stem. It can take a very similar form to the allometry.

@dataclass()
class StemAllocation():
    """
    * Data class containing allocation predictions for stem allometries.
    * Core data class for Community - applying growth.
    """
    whole_crown_gpp: NDArray[np.float32],
    ...
    delta_foliage_mass: NDArray[np.float32]

def get_stem_allocations(stems: StemTraits | Flora, allometry: StemAllometry, potential_gpp: NDArray) -> StemAllocation:
    """
    * Generates predictions of stem allometries at given stem sizes for either:
        * StemTraits for community simulation
        * Flora for trait exploration
    """

[davidorme]

Actually, it's probably cleaner to have the StemAllometry class simply created from those inputs.

@dataclass()
class StemAllometry():
    """
    * Data class containing allometric predictions from stem traits.
    * Core data class for Communities - holds actual stem allometries
    """
    # Init vars
    stems: InitVar[StemTraits | Flora]
    dbh: InitVarNDArray

    # Post init attributes
    stem_height: NDArray[np.float32] = field(post_init=True)
    ...
    canopy_z_max: NDArray[np.float32] = field(post_init=True)