BaseStateTracker#

class BaseStateTracker(model, verbose=False, params=None)[source]#

Bases: BasePhase

Component for tracking the number in each SEIR state for each time tick.

This class maintains a time series of state counts across nodes in the model. The states are dynamically generated as properties based on model.params.states (e.g., “S”, “E”, “I”, “R”). Each state can be accessed as a property that returns a numpy array containing the time series for that state.

The tracking can be done at different aggregation levels: - aggregation_level = -1: Sum over all patches (default, backward compatible) - aggregation_level >= 0: Group by geographic level and track separately

Parameters:

  • model – The simulation model containing nodes, states, and parameters.

  • verbose (bool) – Whether to print verbose output during simulation. Defaults to False.

  • params (BaseStateTrackerParams | None) – Component-specific parameters. If None, will use default parameters.

Initialize the component.

Parameters:

  • model – The model instance this component belongs to.

  • verbose (bool) – Whether to enable verbose output. Defaults to False.

__call__(model, tick)[source]#

Execute component logic for a given simulation tick.

Parameters:

  • model – The model instance.

  • tick (int) – The current simulation tick.

Return type:

None

plot(fig=None)[source]#

Plots the time series of SEIR state counts across all nodes using subplots.

This function creates a separate subplot for each state, showing how the number of individuals in each state changes over time. Each state gets its own subplot for better visibility.

Parameters:

fig (Figure | None) – A matplotlib Figure object. If None, a new figure will be created.

Yields:

None – This function uses a generator to yield control back to the caller. If used directly (not as a generator), it will show the plot immediately.

Example

# Use as a generator (for model.visualize()): for _ in tracker.plot():

plt.show()

plot_combined(fig=None)[source]#

Plots all SEIR states on a single plot for easy comparison.

Parameters:

fig (Figure | None) – A matplotlib Figure object. If None, a new figure will be created.

Yields:

None – This function uses a generator to yield control back to the caller.

get_dataframe()[source]#

Get a DataFrame of state counts over time.

Returns:

  • tick: Time step

  • state: State name (S, E, I, R, etc.)

  • patch_id: Patch/group identifier. Matches the id column of the scenario DataFrame at the requested aggregation_level, or "all_patches" when aggregation_level=-1 (the default).

  • count: Number of individuals in this state

Return type:

DataFrame