Initial plants model draft

tests/core/data/all_config.toml

@@ -24,11 +24,11 @@ a_plant_integer = 12
 
 [[plants.ftypes]]
 pft_name = "shrub"
-maxh = 1.0
+max_height = 1.0
 
 [[plants.ftypes]]
 pft_name = "broadleaf"
-maxh = 50.0
+max_height = 50.0
 
 [[animals.functional_groups]]
 name = "carnivorous_bird"

tests/core/data/all_config_bad.toml

@@ -21,11 +21,11 @@ a_plant_integer = 12
 
 [[plants.ftypes]]
 pft_name = "shrub"
-maxh = 1.0
+max_height = 1.0
 
 [[plants.ftypes]]
 pft_name = "broadleaf"
-maxh = 50.0
+max_height = 50.0
 
 [soil]
 model_time_step = "5 days"

tests/models/plants/conftest.py

@@ -0,0 +1,39 @@
+"""Fixtures for plants model testing."""
+
+import numpy as np
+import pytest
+from xarray import DataArray
+
+
+@pytest.fixture()
+def plants_config(shared_datadir):
+    """Simple configuration fixture for use in tests."""
+
+    from virtual_rainforest.core.config import Config
+
+    return Config(shared_datadir / "all_config.toml")
+
+
+@pytest.fixture()
+def flora(plants_config):
+    """Construct a minimal Flora object."""
+    from virtual_rainforest.models.plants.functional_types import Flora
+
+    flora = Flora.from_config(plants_config)
+
+    return flora
+
+
+@pytest.fixture()
+def plants_data():
+    """Construct a minimal data object with plant cohort data."""
+    from virtual_rainforest.core.data import Data
+    from virtual_rainforest.core.grid import Grid
+
+    data = Data(grid=Grid(cell_ny=2, cell_nx=2))
+    data["plant_cohorts_n"] = DataArray(np.array([5] * 4))
+    data["plant_cohorts_pft"] = DataArray(np.array(["broadleaf"] * 4))
+    data["plant_cohorts_cell_id"] = DataArray(np.arange(4))
+    data["plant_cohorts_dbh"] = DataArray(np.array([0.1] * 4))
+
+    return data

tests/models/plants/data/all_config.toml

@@ -0,0 +1,51 @@
+# This file contains every tag required to build the config for the plants and soil
+# modules. Each value has been chosen to be different from the default value, so that
+# this file can be used to test that providing non-default values works.
+[core]
+modules = ["plants", "soil", "animals"]
+
+[core.grid]
+cell_nx = 10
+cell_ny = 10
+
+[core.timing]
+start_date = "2020-01-01"
+update_interval = "2 weeks"
+run_length = "50 years"
+
+[core.data_output_options]
+save_initial_state = true
+save_final_state = true
+out_initial_file_name = "model_at_start.nc"
+out_final_file_name = "model_at_end.nc"
+
+[plants]
+a_plant_integer = 12
+
+[[plants.ftypes]]
+pft_name = "shrub"
+max_height = 1.0
+
+[[plants.ftypes]]
+pft_name = "broadleaf"
+max_height = 50.0
+
+[[animals.functional_groups]]
+name = "carnivorous_bird"
+taxa = "bird"
+diet = "carnivore"
+
+[[animals.functional_groups]]
+name = "herbivorous_bird"
+taxa = "bird"
+diet = "herbivore"
+
+[[animals.functional_groups]]
+name = "carnivorous_mammal"
+taxa = "mammal"
+diet = "carnivore"
+
+[[animals.functional_groups]]
+name = "herbivorous_mammal"
+taxa = "mammal"
+diet = "herbivore"

tests/models/plants/test_community.py

@@ -0,0 +1,105 @@
+"""Tests the plant community model code."""
+
+from contextlib import nullcontext as does_not_raise
+from logging import CRITICAL, INFO
+
+import numpy as np
+import pytest
+from xarray import DataArray
+
+from tests.conftest import log_check
+
+
+@pytest.mark.parametrize(
+    argnames="vars,raises,exp_log",
+    argvalues=[
+        pytest.param(
+            (("plant_cohorts_n", DataArray(np.array([5] * 4))),),
+            pytest.raises(ValueError),
+            ((CRITICAL, "Missing plant cohort variables"),),
+            id="missing var",
+        ),
+        pytest.param(
+            (
+                ("plant_cohorts_n", DataArray(np.array([5] * 9), dims="toolong")),
+                ("plant_cohorts_pft", DataArray(np.array(["shrub"] * 4))),
+                ("plant_cohorts_cell_id", DataArray(np.arange(4))),
+                ("plant_cohorts_dbh", DataArray(np.array([0.1] * 4))),
+            ),
+            pytest.raises(ValueError),
+            ((CRITICAL, "Unequal plant cohort variable dimensions"),),
+            id="unequal sizes",
+        ),
+        pytest.param(
+            (
+                ("plant_cohorts_n", DataArray(np.array([5] * 4).reshape(2, 2))),
+                ("plant_cohorts_pft", DataArray(np.array(["shrub"] * 4).reshape(2, 2))),
+                ("plant_cohorts_cell_id", DataArray(np.arange(4).reshape(2, 2))),
+                ("plant_cohorts_dbh", DataArray(np.array([0.1] * 4).reshape(2, 2))),
+            ),
+            pytest.raises(ValueError),
+            ((CRITICAL, "Plant cohort variable data is not one dimensional"),),
+            id="not 1D",
+        ),
+        pytest.param(
+            (
+                ("plant_cohorts_n", DataArray(np.array([5] * 4))),
+                ("plant_cohorts_pft", DataArray(np.array(["shrub"] * 4))),
+                ("plant_cohorts_cell_id", DataArray(DataArray(np.arange(2, 6)))),
+                ("plant_cohorts_dbh", DataArray(np.array([0.1] * 4))),
+            ),
+            pytest.raises(ValueError),
+            ((CRITICAL, "Plant cohort cell ids not in grid cell ids"),),
+            id="bad cell ids",
+        ),
+        pytest.param(
+            (
+                ("plant_cohorts_n", DataArray(np.array([5] * 4))),
+                ("plant_cohorts_pft", DataArray(np.array(["tree"] * 4))),
+                ("plant_cohorts_cell_id", DataArray(DataArray(np.arange(4)))),
+                ("plant_cohorts_dbh", DataArray(np.array([0.1] * 4))),
+            ),
+            pytest.raises(ValueError),
+            ((CRITICAL, "Plant cohort PFTs ids not in configured PFTs"),),
+            id="bad pfts",
+        ),
+        pytest.param(
+            (
+                ("plant_cohorts_n", DataArray(np.array([5] * 4))),
+                ("plant_cohorts_pft", DataArray(np.array(["shrub"] * 4))),
+                ("plant_cohorts_cell_id", DataArray(DataArray(np.arange(4)))),
+                ("plant_cohorts_dbh", DataArray(np.array([0.1] * 4))),
+            ),
+            does_not_raise(),
+            ((INFO, "Plant cohort data loaded"),),
+            id="all good",
+        ),
+    ],
+)
+def test_PlantCommunities__init__(caplog, flora, vars, raises, exp_log):
+    """Test the data handling of the plants __init__."""
+
+    from virtual_rainforest.core.data import Data
+    from virtual_rainforest.core.grid import Grid
+    from virtual_rainforest.models.plants.community import PlantCommunities
+
+    data = Data(grid=Grid(cell_ny=2, cell_nx=2))
+
+    for var, value in vars:
+        data[var] = value
+
+    # Clear any data loading log entries
+    caplog.clear()
+
+    with raises:
+        plants_obj = PlantCommunities(data, flora=flora)
+
+        if isinstance(raises, does_not_raise):
+            # Check the expected contents of plants_obj
+            assert len(plants_obj.communities) == 4
+            cids = set([0, 1, 2, 3])
+            assert set(plants_obj.communities.keys()) == cids
+            for cid in cids:
+                assert len(plants_obj[cid]) == 1
+
+    log_check(caplog, expected_log=exp_log)

tests/models/plants/test_functional_types.py

@@ -0,0 +1,43 @@
+"""Test module for plants.functional_types.py.
+
+This module tests the functionality of the plant functional types submodule.
+"""
+
+
+def test_plant_functional_type():
+    """Simple test of PlantFunctionalType dataclass."""
+    from virtual_rainforest.models.plants.functional_types import PlantFunctionalType
+
+    pft = PlantFunctionalType(pft_name="tree", max_height=12.0)
+
+    assert pft.pft_name == "tree"
+    assert pft.max_height == 12.0
+
+
+def test_flora__init__():
+    """Simple test of Flora __init__."""
+    from virtual_rainforest.models.plants.functional_types import (
+        Flora,
+        PlantFunctionalType,
+    )
+
+    flora = Flora(
+        [
+            PlantFunctionalType(pft_name="shrub", max_height=1.0),
+            PlantFunctionalType(pft_name="broadleaf", max_height=50.0),
+        ]
+    )
+
+    assert len(flora) == 2
+    assert tuple(flora.keys()) == ("shrub", "broadleaf")
+
+
+def test_plant_functional_types_from_config(plants_config):
+    """Simple test of Flora from_config factory method."""
+
+    from virtual_rainforest.models.plants.functional_types import Flora
+
+    flora = Flora.from_config(plants_config)
+
+    assert len(flora) == 2
+    assert tuple(flora.keys()) == ("shrub", "broadleaf")

tests/models/plants/test_plants_model.py

@@ -0,0 +1,29 @@
+"""Tests for the model.plants.plants_model submodule."""
+
+
+def test_PlantsModel__init__(plants_data, flora):
+    """Test the PlantsModel.__init__ method."""
+    from pint import Quantity
+
+    from virtual_rainforest.models.plants.plants_model import PlantsModel
+
+    plants_model = PlantsModel(
+        data=plants_data, update_interval=Quantity("1 month"), flora=flora
+    )
+
+    # Currently trivial test.
+    assert len(plants_model.flora) == len(flora)
+
+
+def test_PlantsModel_from_config(plants_data, plants_config):
+    """Test the PlantsModel.from_config factory method."""
+    from pint import Quantity
+
+    from virtual_rainforest.models.plants.plants_model import PlantsModel
+
+    plants_model = PlantsModel.from_config(
+        data=plants_data, config=plants_config, update_interval=Quantity("1 month")
+    )
+
+    # Currently trivial test.
+    assert isinstance(plants_model, PlantsModel)

virtual_rainforest/models/plants/__init__.py

@@ -1,8 +1,16 @@
-from importlib import resources
+"""The :mod:`~virtual_rainforest.models.plants` module provides a plants model for use
+in the Virtual Rainforest. The submodules provide:
 
-from virtual_rainforest.core.config import register_schema
+* The :mod:`~virtual_rainforest.models.plants.plants_model` submodule provides the
+  PlantsModel class as the main API for interacting with the plants model.
+* The :mod:`~virtual_rainforest.models.plants.constants` submodule provides dataclasses
+  containing constants used in the model.
+* The :mod:`~virtual_rainforest.models.plants.community` submodule provides the
+  PlantCohort dataclass that records the details of an individual cohort and the
+  PlantCommunities class that records list of plant cohorts by grid cell.
+"""  # noqa: D205, D415
 
-with resources.path(
-    "virtual_rainforest.models.plants", "model_schema.json"
-) as schema_file_path:
-    register_schema(module_name="plants", schema_file_path=schema_file_path)
+from virtual_rainforest.core.base_model import register_model
+from virtual_rainforest.models.plants.plants_model import PlantsModel
+
+register_model(__name__, PlantsModel)