"""
==========================
Vivarium Interactive Tools
==========================
This module provides an interface for interactive simulation usage. The main
part is the :class:`InteractiveContext`, a sub-class of the main simulation
object in ``vivarium`` that has been extended to include convenience
methods for running and exploring the simulation in an interactive setting.
See the associated tutorials for :ref:`running <interactive_tutorial>` and
:ref:`exploring <exploration_tutorial>` for more information.
"""
from math import ceil
from typing import Any, Callable, Dict, List
import pandas as pd
from vivarium.framework.engine import SimulationContext
from vivarium.framework.time import Time, Timedelta
from vivarium.framework.values import Pipeline
from .utilities import log_progress, run_from_ipython
[docs]
class InteractiveContext(SimulationContext):
"""A simulation context with helper methods for running simulations interactively."""
def __init__(self, *args, setup=True, **kwargs):
super().__init__(*args, **kwargs)
if setup:
self.setup()
@property
def current_time(self) -> Time:
"""Returns the current simulation time."""
return self._clock.time
[docs]
def setup(self):
super().setup()
self.initialize_simulants()
[docs]
def step(self, step_size: Timedelta = None):
"""Advance the simulation one step.
Parameters
----------
step_size
An optional size of step to take. Must be compatible with the
simulation clock's step size (usually a pandas.Timedelta).
"""
old_step_size = self._clock._clock_step_size
if step_size is not None:
if not (
isinstance(step_size, type(self._clock.step_size))
or isinstance(self._clock.step_size, type(step_size))
):
raise ValueError(
f"Provided time must be compatible with {type(self._clock.step_size)}"
)
self._clock._clock_step_size = step_size
super().step()
self._clock._clock_step_size = old_step_size
[docs]
def run(self, with_logging: bool = True) -> int:
"""Run the simulation for the duration specified in the configuration.
Parameters
----------
with_logging
Whether or not to log the simulation steps. Only works in an ipython
environment.
Returns
-------
int
The number of steps the simulation took.
"""
return self.run_until(self._clock.stop_time, with_logging=with_logging)
[docs]
def run_for(self, duration: Timedelta, with_logging: bool = True) -> int:
"""Run the simulation for the given time duration.
Parameters
----------
duration
The length of time to run the simulation for. Should be compatible
with the simulation clock's step size (usually a pandas
Timedelta).
with_logging
Whether or not to log the simulation steps. Only works in an ipython
environment.
Returns
-------
int
The number of steps the simulation took.
"""
return self.run_until(self._clock.time + duration, with_logging=with_logging)
[docs]
def run_until(self, end_time: Time, with_logging: bool = True) -> int:
"""Run the simulation until the provided end time.
Parameters
----------
end_time
The time to run the simulation until. The simulation will run until
its clock is greater than or equal to the provided end time. Must be
compatible with the simulation clock's step size (usually a pandas.Timestamp)
with_logging
Whether or not to log the simulation steps. Only works in an ipython
environment.
Returns
-------
int
The number of steps the simulation took.
"""
if not (
isinstance(end_time, type(self._clock.time))
or isinstance(self._clock.time, type(end_time))
):
raise ValueError(
f"Provided time must be compatible with {type(self._clock.time)}"
)
iterations = int(ceil((end_time - self._clock.time) / self._clock.step_size))
self.take_steps(number_of_steps=iterations, with_logging=with_logging)
assert self._clock.time - self._clock.step_size < end_time <= self._clock.time
return iterations
[docs]
def take_steps(
self, number_of_steps: int = 1, step_size: Timedelta = None, with_logging: bool = True
):
"""Run the simulation for the given number of steps.
Parameters
----------
number_of_steps
The number of steps to take.
step_size
An optional size of step to take. Must be compatible with the
simulation clock's step size (usually a pandas.Timedelta).
with_logging
Whether or not to log the simulation steps. Only works in an ipython
environment.
"""
if not isinstance(number_of_steps, int):
raise ValueError("Number of steps must be an integer.")
if run_from_ipython() and with_logging:
for _ in log_progress(range(number_of_steps), name="Step"):
self.step(step_size)
else:
for _ in range(number_of_steps):
self.step(step_size)
[docs]
def get_population(self, untracked: bool = False) -> pd.DataFrame:
"""Get a copy of the population state table.
Parameters
----------
untracked
Whether or not to return simulants who are no longer being tracked
by the simulation.
"""
return self._population.get_population(untracked)
[docs]
def list_values(self) -> List[str]:
"""List the names of all pipelines in the simulation."""
return list(self._values.keys())
[docs]
def get_value(self, value_pipeline_name: str) -> Pipeline:
"""Get the value pipeline associated with the given name."""
return self._values.get_value(value_pipeline_name)
[docs]
def list_events(self) -> List[str]:
"""List all event types registered with the simulation."""
return self._events.list_events()
[docs]
def get_listeners(self, event_type: str) -> Dict[int, List[Callable]]:
"""Get all listeners of a particular type of event.
Available event types can be found by calling
:func:`InteractiveContext.list_events`.
Parameters
----------
event_type
The type of event to grab the listeners for.
"""
if event_type not in self._events:
raise ValueError(f"No event {event_type} in system.")
return self._events.get_listeners(event_type)
[docs]
def get_emitter(self, event_type: str) -> Callable:
"""Get the callable that emits the given type of events.
Available event types can be found by calling
:func:`InteractiveContext.list_events`.
Parameters
----------
event_type
The type of event to grab the listeners for.
"""
if event_type not in self._events:
raise ValueError(f"No event {event_type} in system.")
return self._events.get_emitter(event_type)
[docs]
def list_components(self) -> Dict[str, Any]:
"""Get a mapping of component names to components currently in the simulation.
Returns
-------
Dict[str, Any]
A dictionary mapping component names to components.
"""
return self._component_manager.list_components()
[docs]
def get_component(self, name: str) -> Any:
"""Get the component in the simulation that has ``name``, if present.
Names are guaranteed to be unique.
Parameters
----------
name
A component name.
Returns
-------
A component that has the name ``name`` else None.
"""
return self._component_manager.get_component(name)
[docs]
def print_initializer_order(self):
"""Print the order in which population initializers are called."""
initializers = []
for r in self._resource:
name = r.__name__
if hasattr(r, "__self__"):
obj = r.__self__
initializers.append(f"{obj.__class__.__name__}({obj.name}).{name}")
else:
initializers.append(f"Unbound function {name}")
print("\n".join(initializers))
[docs]
def print_lifecycle_order(self):
"""Print the order of lifecycle events (including user event handlers)."""
print(self._lifecycle)
def __repr__(self):
return "InteractiveContext()"