Refactor setting of PROPERTIES columns across Modules

src/tlo/core.py

@@ -8,14 +8,18 @@
 
 import json
 from enum import Enum, auto
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Dict, FrozenSet, Optional, Set, Type
 
 import numpy as np
 import pandas as pd
 
 if TYPE_CHECKING:
+    from pathlib import Path
     from typing import Optional
 
+    from tlo.methods import Metadata
+    from tlo.methods.causes import Cause
+    from tlo.population import Population
     from tlo.simulation import Simulation
 
 class Types(Enum):
@@ -76,7 +80,7 @@ class Specifiable:
         Types.BITSET: int,
     }
 
-    def __init__(self, type_, description, categories=None):
+    def __init__(self, type_: Types, description: str, categories: Set[Any] = None):
         """Create a new Specifiable.
 
         :param type_: an instance of Types giving the type of allowed values
@@ -94,16 +98,16 @@ def __init__(self, type_, description, categories=None):
             self.categories = categories
 
     @property
-    def python_type(self):
+    def python_type(self) -> Type[Any]:
         """Return the Python type corresponding to this Specifiable."""
         return self.PYTHON_TYPE_MAP[self.type_]
 
     @property
-    def pandas_type(self):
+    def pandas_type(self) -> Type[Any]:
         """Return the Pandas type corresponding to this Specifiable."""
         return self.PANDAS_TYPE_MAP[self.type_]
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return detailed description of Specifiable."""
 
         delimiter = " === "
@@ -131,8 +135,19 @@ class Property(Specifiable):
         object: float("nan"),
         np.uint32: 0,
     }
-
-    def __init__(self, type_, description, categories=None, *, ordered=False):
+    # An overwrite for the default value map above, if a specific property needs
+    # to instantiate at a different default value.
+    _default_property_value: Any
+
+    def __init__(
+        self,
+        type_: Types,
+        description: str,
+        categories: Set[Any] = None,
+        *,
+        ordered: bool = False,
+        default_property_value: Optional[Any] = None,
+    ) -> None:
         """Create a new property specification.
 
         :param type_: An instance of ``Types`` giving the type of allowed values of this
@@ -142,17 +157,39 @@ def __init__(self, type_, description, categories=None, *, ordered=False):
             ``Types.CATEGORICAL``.
         :param ordered: Whether categories are ordered  if ``type_`` is
             ``Types.CATEGORICAL``.
+        :param default_property_value: The default value to assign to elements in a series
+            of this ``Property``. If ``type_`` is ``Types.CATEGORICAL``, this value must
+            be one of the supplied categories, otherwise it must be of type ``type_``.
         """
         if type_ in [Types.SERIES, Types.DATA_FRAME]:
             raise TypeError("Property cannot be of type SERIES or DATA_FRAME.")
+
         super().__init__(type_, description, categories)
         self.ordered = ordered
 
+        # Set supplied default value, if appropriate
+        self._default_property_value = (
+            default_property_value
+            if default_property_value is not None
+            and (
+                (
+                    self.type_ is Types.CATEGORICAL
+                    and default_property_value in categories
+                )
+                or isinstance(default_property_value, self.python_type)
+            )
+            else None
+        )
+
     @property
-    def _default_value(self):
-        return self.PANDAS_TYPE_DEFAULT_VALUE_MAP[self.pandas_type]
+    def _default_value(self) -> Type[Any]:
+        return (
+            self.PANDAS_TYPE_DEFAULT_VALUE_MAP[self.pandas_type]
+            if self._default_property_value is None
+            else self._default_property_value
+        )
 
-    def create_series(self, name, size):
+    def create_series(self, name: str, size: int) -> pd.Series:
         """Create a Pandas Series for this property.
 
         The values will be left uninitialised.
@@ -201,48 +238,47 @@ class attribute on a subclass.
     # Subclasses can override this to declare the set of initialisation dependencies
     # Declares modules that need to be registered in simulation and initialised before
     # this module
-    INIT_DEPENDENCIES = frozenset()
+    INIT_DEPENDENCIES: FrozenSet[str] = frozenset()
 
     # Subclasses can override this to declare the set of optional init. dependencies
     # Declares modules that need to be registered in simulation and initialised before
     # this module if they are present, but are not required otherwise
-    OPTIONAL_INIT_DEPENDENCIES = frozenset()
+    OPTIONAL_INIT_DEPENDENCIES: FrozenSet[str] = frozenset()
 
     # Subclasses can override this to declare the set of additional dependencies
     # Declares any modules that need to be registered in simulation in addition to those
     # in INIT_DEPENDENCIES to allow running simulation
-    ADDITIONAL_DEPENDENCIES = frozenset()
+    ADDITIONAL_DEPENDENCIES: FrozenSet[str] = frozenset()
 
     # Subclasses can override this to declare the set of modules that this module can be
     # used in place of as a dependency
-    ALTERNATIVE_TO = frozenset()
+    ALTERNATIVE_TO: FrozenSet[str] = frozenset()
 
     # Subclasses can override this set to add metadata tags to their class
     # See tlo.methods.Metadata class
-    METADATA = {}
+    METADATA: FrozenSet[Metadata] = frozenset()
 
-    # Subclasses can override this set to declare the causes death that this module contributes to
+    # Subclasses can override this dict to declare the causes death that this module contributes to
     # This is a dict of the form {<name_used_by_the_module : Cause()}: see core.Cause
-    CAUSES_OF_DEATH = {}
+    CAUSES_OF_DEATH: Dict[str, Cause] = {}
 
     # Subclasses can override this set to declare the causes disability that this module contributes to
     # This is a dict of the form {<name_used_by_the_module : Cause()}: see core.Cause
-    CAUSES_OF_DISABILITY = {}
+    CAUSES_OF_DISABILITY: Dict[str, Cause] = {}
 
     # Subclasses may declare this dictionary to specify module-level parameters.
     # We give an empty definition here as default.
-    PARAMETERS = {}
+    PARAMETERS: Dict[str, Parameter] = {}
 
     # Subclasses may declare this dictionary to specify properties of individuals.
     # We give an empty definition here as default.
-    PROPERTIES = {}
+    PROPERTIES: Dict[str, Property] = {}
 
     # The explicit attributes of the module. We list these to distinguish dynamic
     # parameters created from the PARAMETERS specification.
     __slots__ = ('name', 'parameters', 'rng', 'sim')
 
-
-    def __init__(self, name=None):
+    def __init__(self, name: str = None) -> None:
         """Construct a new disease module ready to be included in a simulation.
 
         Initialises an empty parameters dictionary and module-specific random number
@@ -255,7 +291,7 @@ def __init__(self, name=None):
         self.name = name or self.__class__.__name__
         self.sim: Optional[Simulation] = None
 
-    def load_parameters_from_dataframe(self, resource: pd.DataFrame):
+    def load_parameters_from_dataframe(self, resource: pd.DataFrame) -> None:
         """Automatically load parameters from resource dataframe, updating the class parameter dictionary
 
         Goes through parameters dict self.PARAMETERS and updates the self.parameters with values
@@ -316,7 +352,7 @@ def load_parameters_from_dataframe(self, resource: pd.DataFrame):
             # Save the values to the parameters
             self.parameters[parameter_name] = parameter_value
 
-    def read_parameters(self, data_folder):
+    def read_parameters(self, data_folder: str | Path) -> None:
         """Read parameter values from file, if required.
 
         Must be implemented by subclasses.
@@ -326,23 +362,34 @@ def read_parameters(self, data_folder):
         """
         raise NotImplementedError
 
-    def initialise_population(self, population):
+    def initialise_population(self, population: Population) -> None:
         """Set our property values for the initial population.
 
-        Must be implemented by subclasses.
-
         This method is called by the simulation when creating the initial population, and is
         responsible for assigning initial values, for every individual, of those properties
         'owned' by this module, i.e. those declared in its PROPERTIES dictionary.
 
+        By default, all ``Property``s in ``self.PROPERTIES`` will have
+        their columns set to the default value for their dtype in the population dataframe.
+        Modules that inherit from this class do not have to redefine this method to obtain
+        this behaviour by default, however if you want to add additional steps or initialise
+        your properties in another way, you may need to define this method explicitly and
+        call super().initialise_population. See the template docstring in
+        ``tlo.methods.skeleton.initialise_population`` for more information.
+
         TODO: We probably need to declare somehow which properties we 'read' here, so the
         simulation knows what order to initialise modules in!
 
-        :param population: the population of individuals
+        :param population: The population of individuals in the simulation.
         """
-        raise NotImplementedError
+        df = population.props
+
+        for property_name, property in self.PROPERTIES.items():
+            df.loc[df.is_alive, property_name] = (
+                property._default_value
+            )
 
-    def initialise_simulation(self, sim):
+    def initialise_simulation(self, sim: Simulation) -> None:
         """Get ready for simulation start.
 
         Must be implemented by subclasses.
@@ -353,15 +400,15 @@ def initialise_simulation(self, sim):
         """
         raise NotImplementedError
 
-    def pre_initialise_population(self):
+    def pre_initialise_population(self) -> None:
         """Carry out any work before any populations have been initialised
 
         This optional method allows access to all other registered modules, before any of
         the modules have initialised a population. This is expected to be useful for
         when a module's properties rely upon information from other modules.
         """
 
-    def on_birth(self, mother_id, child_id):
+    def on_birth(self, mother_id: int, child_id: int) -> None:
         """Initialise our properties for a newborn individual.
 
         Must be implemented by subclasses.
@@ -373,6 +420,6 @@ def on_birth(self, mother_id, child_id):
         """
         raise NotImplementedError
 
-    def on_simulation_end(self):
+    def on_simulation_end(self) -> None:
         """This is called after the simulation has ended.
         Modules do not need to declare this."""