"""
=========
Component
=========
A base Component class to be used to create components for use in ``vivarium``
simulations.
"""
from __future__ import annotations
import re
from abc import ABC
from collections.abc import Sequence
from importlib import import_module
from inspect import signature
from typing import TYPE_CHECKING, Any, overload
import pandas as pd
from vivarium.artifact import ArtifactException
from vivarium.config_tree import ConfigTree, ConfigurationError
from vivarium.engine.framework.lifecycle import LifeCycleError, lifecycle_states
from vivarium.engine.types import LookupTableData
if TYPE_CHECKING:
import loguru
from vivarium.engine.framework.engine import Builder
from vivarium.engine.framework.event import Event
from vivarium.engine.framework.lookup import LookupTable
from vivarium.engine.framework.population import PopulationView
from vivarium.engine.types import DataInput
DEFAULT_EVENT_PRIORITY = 5
"""The default priority at which events will be triggered."""
[docs]
class Component(ABC):
"""The base class for all components used in a Vivarium simulation.
A :class:`Component` in a Vivarium simulation represents a distinct feature or
aspect of the model. It encapsulates the logic and data needed for that
feature. Components commonly interact with the rest of the simulation by
creating and updating columns in the state table, registering pipelines,
and registering modifiers on pipelines created by other components. Observer
components might also register observations. All components within a
simulation must have a unique name, which is generated by default from the
component's class and the argument passed to its constructor.
The :meth:`setup_component` method is run by Vivarium during the setup phase
and performs a series of operations to prepare the component for the simulation.
These operations include setting the logger for the component, calling the
component's custom :meth:`setup` method, setting the population view if the
component needs one, and registering listeners for each lifecycle event if
the component has defined a method to be triggered on that event.
Subclasses of :class:`Component` should override these properties as needed:
- :attr:`sub_components`
- :attr:`configuration_defaults`
- :attr:`post_setup_priority`
- :attr:`time_step_prepare_priority`
- :attr:`time_step_priority`
- :attr:`time_step_cleanup_priority`
- :attr:`collect_metrics_priority`
- :attr:`simulation_end_priority`
Subclasses of :class:`Component` should override these methods in order to have
operations occur during the appropriate lifecycle phase of a simulation:
- :meth:`setup`
- :meth:`on_post_setup`
- :meth:`on_time_step_prepare`
- :meth:`on_time_step`
- :meth:`on_time_step_cleanup`
- :meth:`on_collect_metrics`
- :meth:`on_simulation_end`
"""
CONFIGURATION_DEFAULTS: dict[str, Any] = {}
"""A dictionary containing the defaults for any configurations managed by this
component. An empty dictionary indicates no managed configurations. Components will
look for a ``data_sources`` block in this dictionary to build lookup tables automatically.
"""
def __init__(self) -> None:
"""Initializes a new instance of the Component class.
This method is the initializer for the Component class. It initializes
logger of type Logger and population_view of type PopulationView to None.
These attributes will be fully initialized in the setup_component method
of this class.
"""
self._repr: str = ""
self._name: str = ""
self._sub_components: Sequence["Component"] = []
self._logger: loguru.Logger | None = None
self.configuration: ConfigTree = ConfigTree()
self._population_view: PopulationView | None = None
def __repr__(self) -> str:
"""Returns a string representation of the :meth:`__init__` call made to create
this object.
The representation is built by retrieving the initialization parameters
and their values. If a value is an instance of :class:`Component`, its own
:meth:`__repr__` is called. The resulting string is stored in the ``_repr``
attribute and returned.
IMPORTANT: this method must not be called within the :meth:`__init__`
functions of this component or its subclasses or its value may not be
initialized correctly.
Returns
-------
A string representation of the __init__ call made to create this object.
"""
if not self._repr:
args = ", ".join(
[
f"{name}={value.__repr__() if isinstance(value, Component) else value}"
for name, value in self.get_initialization_parameters().items()
]
)
self._repr = f"{type(self).__name__}({args})"
return self._repr
def __str__(self) -> str:
return self._repr
##############
# Properties #
##############
@property
def name(self) -> str:
"""The name of the component.
By convention, these are in snake case with arguments of the :meth:`__init__`
appended and separated by ``.``.
Names must be unique within a simulation.
The name is created by first converting the name of the class to snake
case. Then, the names of the initialization parameters are appended,
separated by ``.``. If a parameter is an instance of `Component`, its
:attr:`name` property is used; otherwise, the string representation of the
parameter is used. The resulting string is stored in the ``_name`` attribute
and returned.
IMPORTANT: this property must not be accessed within the :meth:`__init__`
functions of this component or its subclasses or its value may not be
initialized correctly.
"""
if not self._name:
base_name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", type(self).__name__)
base_name = re.sub("([a-z0-9])([A-Z])", r"\1_\2", base_name).lower()
args = [
f"'{value.name}'" if isinstance(value, Component) else str(value)
for value in self.get_initialization_parameters().values()
]
self._name = ".".join([base_name] + args)
return self._name
@property
def logger(self) -> loguru.Logger:
"""The logger for this component.
Raises
------
LifeCycleError
If the logger has not been initialized.
"""
if self._logger is None:
raise LifeCycleError(
f"Logger for component '{self.name}' has not been initialized. "
"This is likely due to having called this prior to simulation setup."
)
return self._logger
@property
def population_view(self) -> PopulationView:
"""The :class:`~vivarium.engine.framework.population.PopulationView` for this component.
Raises
------
PopulationError
If the component does not have access to the state table.
"""
if self._population_view is None:
from vivarium.engine.framework.population.exceptions import PopulationError
raise PopulationError(
f"Component '{self.name}' does not have access to the state table. "
"This is likely due to having called this prior to simulation setup."
)
return self._population_view
@property
def private_columns(self) -> list[str]:
"""The list of private columns created by this component."""
return self.population_view.private_columns
@property
def sub_components(self) -> Sequence["Component"]:
"""The components managed by this component."""
return self._sub_components
@property
def configuration_defaults(self) -> dict[str, Any]:
"""The dictionary containing the defaults for any configurations managed
by this component.
These default values will be stored at the ``component_configs`` layer of the
simulation's :class:`~vivarium.config_tree.main.ConfigTree`.
"""
return self.CONFIGURATION_DEFAULTS
@property
def lookup_table_value_columns(self) -> dict[str, str | list[str]]:
"""A mapping of lookup table names to their value columns."""
return {}
@property
def post_setup_priority(self) -> int:
"""The priority of this component's ``post_setup`` listener."""
return DEFAULT_EVENT_PRIORITY
@property
def time_step_prepare_priority(self) -> int:
"""The priority of this component's ``time_step__prepare`` listener."""
return DEFAULT_EVENT_PRIORITY
@property
def time_step_priority(self) -> int:
"""The priority of this component's ``time_step`` listener."""
return DEFAULT_EVENT_PRIORITY
@property
def time_step_cleanup_priority(self) -> int:
"""The priority of this component's ``time_step__cleanup`` listener."""
return DEFAULT_EVENT_PRIORITY
@property
def collect_metrics_priority(self) -> int:
"""The priority of this component's ``collect_metrics`` listener."""
return DEFAULT_EVENT_PRIORITY
@property
def simulation_end_priority(self) -> int:
"""The priority of this component's ``simulation_end`` listener."""
return DEFAULT_EVENT_PRIORITY
#####################
# Lifecycle methods #
#####################
[docs]
def setup_component(self, builder: Builder) -> None:
"""Sets up the component for a Vivarium simulation.
This method is run by Vivarium during the setup phase. It performs a series
of operations to prepare the component for the simulation.
It sets the :attr:`logger` for the component, sets up the component, sets the
population view, and registers various listeners including ``post_setup``,
``simulant_initializer``, ``time_step__prepare``, ``time_step``, ``time_step__cleanup``,
``collect_metrics``, and ``simulation_end`` listeners.
Parameters
----------
builder
The builder object used to set up the component.
"""
with builder.components._tracking_setup(self):
self._logger = builder.logging.get_logger(self.name)
self.configuration = self.get_configuration(builder)
self.setup(builder)
self._set_population_view(builder)
self._register_post_setup_listener(builder)
self._register_time_step_prepare_listener(builder)
self._register_time_step_listener(builder)
self._register_time_step_cleanup_listener(builder)
self._register_collect_metrics_listener(builder)
self._register_simulation_end_listener(builder)
#######################
# Methods to override #
#######################
[docs]
def setup(self, builder: Builder) -> None:
"""Defines custom actions this component needs to run during the setup
lifecycle phase.
This method is intended to be overridden by subclasses to perform any
necessary setup operations specific to the component. By default, it
does nothing.
Parameters
----------
builder
The builder object used to set up the component.
"""
pass
[docs]
def on_post_setup(self, event: Event) -> None:
"""Method that vivarium will run during the ``post_setup`` event.
This method is intended to be overridden by subclasses if there are
operations they need to perform specifically during the ``post_setup`` event.
Notes
-----
This method is not commonly used functionality.
Parameters
----------
event
The event object associated with the ``post_setup`` event.
"""
pass
[docs]
def on_time_step_prepare(self, event: Event) -> None:
"""Method that vivarium will run during the ``time_step__prepare`` event.
This method is intended to be overridden by subclasses if there are
operations they need to perform specifically during the
``time_step__prepare`` event.
Parameters
----------
event
The event object associated with the ``time_step__prepare`` event.
"""
pass
[docs]
def on_time_step(self, event: Event) -> None:
"""Method that vivarium will run during the ``time_step`` event.
This method is intended to be overridden by subclasses if there are
operations they need to perform specifically during the ``time_step`` event.
Parameters
----------
event
The event object associated with the ``time_step`` event.
"""
pass
[docs]
def on_time_step_cleanup(self, event: Event) -> None:
"""Method that vivarium will run during the ``time_step__cleanup`` event.
This method is intended to be overridden by subclasses if there are
operations they need to perform specifically during the
``time_step__cleanup`` event.
Parameters
----------
event
The event object associated with the ``time_step__cleanup`` event.
"""
pass
[docs]
def on_collect_metrics(self, event: Event) -> None:
"""Method that vivarium will run during the ``collect_metrics`` event.
This method is intended to be overridden by subclasses if there are
operations they need to perform specifically during the ``collect_metrics``
event.
Parameters
----------
event
The event object associated with the ``collect_metrics`` event.
"""
pass
[docs]
def on_simulation_end(self, event: Event) -> None:
"""Method that vivarium will run during the ``simulation_end`` event.
This method is intended to be overridden by subclasses if there are
operations they need to perform specifically during the ``simulation_end``
event.
Parameters
----------
event
The event object associated with the ``simulation_end`` event.
"""
pass
##################
# Helper methods #
##################
[docs]
def get_initialization_parameters(self) -> dict[str, Any]:
"""Retrieves the values of all parameters specified in the :meth:`__init__` that
have an attribute with the same name.
Notes
-----
This retrieves the value of the attribute at the time of calling,
which is not guaranteed to be the same as the original value.
Returns
-------
A dictionary where the keys are the names of the parameters used in
the :meth:`__init__` method and the values are their current values.
"""
return {
parameter_name: getattr(self, parameter_name)
for parameter_name in signature(self.__init__).parameters # type: ignore[misc]
if hasattr(self, parameter_name)
}
[docs]
def get_configuration(self, builder: Builder) -> ConfigTree:
"""Retrieves the configuration for this component from the builder.
This method retrieves the configuration for this component from the
simulation's overall configuration. The configuration is retrieved using
the name of the component as the key.
Parameters
----------
builder
The simulation's builder object.
Returns
-------
The configuration for this component, or a default empty configuration.
"""
if self.name in builder.configuration:
return builder.configuration.get_tree(self.name)
return ConfigTree({"data_sources": {}})
@overload
def build_lookup_table(
self,
builder: Builder,
name: str,
data_source: DataInput | None = None,
value_columns: str | None = None,
) -> LookupTable[pd.Series[Any]]:
...
@overload
def build_lookup_table(
self,
builder: Builder,
name: str,
data_source: DataInput | None = None,
value_columns: list[str] | tuple[str, ...] = ...,
) -> LookupTable[pd.DataFrame]:
...
[docs]
def build_lookup_table(
self,
builder: Builder,
name: str,
data_source: DataInput | None = None,
value_columns: list[str] | tuple[str, ...] | str | None = None,
) -> LookupTable[pd.Series[Any]] | LookupTable[pd.DataFrame]:
"""Builds a LookupTable.
If a data_source is not provided, the method will look for a data source
in the component's configuration under the key "data_sources" with the
provided name.
If value_columns provided is a list or tuple, a LookupTable returning a
DataFrame will be built. If it is a string or None, a LookupTable
returning a Series will be built. If value_columns is None, the name of the
returned Series will be "value".
Parameters
----------
builder
The builder object used to set up the component.
data_source
The data source to build the LookupTable from. If None, the data source
will be retrieved from the component's configuration.
name
The name of the lookup table, used to retrieve the data source from
the configuration if data_source is None.
value_columns
The columns to include in the LookupTable.
Returns
-------
The LookupTable built from the data source.
Raises
------
vivarium.config_tree.exceptions.ConfigurationError
If the data source is invalid.
"""
if data_source is None:
data_source = self.configuration.get(["data_sources", name])
if data_source is None:
raise ConfigurationError(
f"No data source provided for lookup table '{name}', "
"and no data source found in configuration."
)
try:
data = self.get_data(builder, data_source)
return builder.lookup.build_table(
data=data, name=name, value_columns=value_columns
)
except ConfigurationError as e:
raise ConfigurationError(f"Error building lookup table '{name}': {e}")
[docs]
def get_data(self, builder: Builder, data_source: DataInput) -> LookupTableData:
"""Retrieves data from a data source.
If the data source is a float or a DataFrame, it is treated as the data
itself. If the data source is a string, containing the substring '::',
it is treated as a function to call to retrieve the data. The string to
the left of '::' is the module to import, and the string to the right is
the function to call. 'self' can be provided to the left of '::' to call
a method on the component itself. If the data source is a string without
the substring '::', it is treated as a key in the artifact.
Parameters
----------
builder
The builder object used to set up the component.
data_source
The data source to retrieve data from.
Returns
-------
The data retrieved from the data source.
Raises
------
vivarium.config_tree.exceptions.ConfigurationError
If the data source is invalid.
"""
if isinstance(data_source, str):
if "::" in data_source:
module, method = data_source.split("::")
try:
if module == "self":
data_source_callable = getattr(self, method)
else:
data_source_callable = getattr(import_module(module), method)
except ModuleNotFoundError:
raise ConfigurationError(f"Unable to find module '{module}'.")
except AttributeError:
module_string = (
f"component {self.name}" if module == "self" else f"module '{module}'"
)
raise ConfigurationError(
f"There is no method '{method}' for the {module_string}."
)
data: LookupTableData = data_source_callable(builder)
else:
try:
data = builder.data.load(data_source)
except ArtifactException:
raise ConfigurationError(
f"Failed to find key '{data_source}' in artifact."
)
elif callable(data_source):
data = data_source(builder)
else:
data = data_source
return data
def _set_population_view(self, builder: Builder) -> None:
"""Creates the PopulationView for this component.
Parameters
----------
builder
The builder object used to set up the component.
"""
self._population_view = builder.population.get_view(self)
def _register_post_setup_listener(self, builder: Builder) -> None:
"""Registers a ``post_setup`` listener if this component has defined one.
This method allows the component to respond to ``post_setup`` events if it
has its own :meth:`on_post_setup` method. The listener will be registered with
the component's ``post_setup`` priority, allowing control over the order of
operations when multiple components are listening to the same event.
Parameters
----------
builder
The builder with which to register the listener.
"""
if type(self).on_post_setup != Component.on_post_setup:
builder.event.register_listener(
lifecycle_states.POST_SETUP,
self.on_post_setup,
self.post_setup_priority,
)
def _register_time_step_prepare_listener(self, builder: Builder) -> None:
"""Registers a ``time_step__prepare`` listener if this component has defined one.
This method allows the component to respond to ``time_step__prepare`` events
if it has its own :meth:`on_time_step_prepare` method. The listener will be
registered with the component's ``time_step__prepare`` priority.
Parameters
----------
builder
The builder with which to register the listener.
"""
if type(self).on_time_step_prepare != Component.on_time_step_prepare:
builder.event.register_listener(
lifecycle_states.TIME_STEP_PREPARE,
self.on_time_step_prepare,
self.time_step_prepare_priority,
)
def _register_time_step_listener(self, builder: Builder) -> None:
"""Registers a ``time_step`` listener if this component has defined one.
This method allows the component to respond to ``time_step`` events
if it has its own :meth:`on_time_step` method. The listener will be
registered with the component's ``time_step`` priority.
Parameters
----------
builder
The builder with which to register the listener.
"""
if type(self).on_time_step != Component.on_time_step:
builder.event.register_listener(
lifecycle_states.TIME_STEP,
self.on_time_step,
self.time_step_priority,
)
def _register_time_step_cleanup_listener(self, builder: Builder) -> None:
"""Registers a ``time_step__cleanup`` listener if this component has defined one.
This method allows the component to respond to ``time_step__cleanup`` events
if it has its own :meth:`on_time_step_cleanup` method. The listener will be
registered with the component's ``time_step__cleanup`` priority.
Parameters
----------
builder
The builder with which to register the listener.
"""
if type(self).on_time_step_cleanup != Component.on_time_step_cleanup:
builder.event.register_listener(
lifecycle_states.TIME_STEP_CLEANUP,
self.on_time_step_cleanup,
self.time_step_cleanup_priority,
)
def _register_collect_metrics_listener(self, builder: Builder) -> None:
"""Registers a ``collect_metrics`` listener if this component has defined one.
This method allows the component to respond to ``collect_metrics`` events
if it has its own :meth:`on_collect_metrics` method. The listener will be
registered with the component's ``collect_metrics`` priority.
Parameters
----------
builder
The builder with which to register the listener.
"""
if type(self).on_collect_metrics != Component.on_collect_metrics:
builder.event.register_listener(
lifecycle_states.COLLECT_METRICS,
self.on_collect_metrics,
self.collect_metrics_priority,
)
def _register_simulation_end_listener(self, builder: Builder) -> None:
"""Registers a ``simulation_end`` listener if this component has defined one.
This method allows the component to respond to ``simulation_end`` events
if it has its own :meth:`on_simulation_end` method. The listener will be
registered with the component's ``simulation_end`` priority.
Parameters
----------
builder
The builder with which to register the listener.
"""
if type(self).on_simulation_end != Component.on_simulation_end:
builder.event.register_listener(
lifecycle_states.SIMULATION_END,
self.on_simulation_end,
self.simulation_end_priority,
)