Disease States

This module contains tools to manage standard disease states.

class vivarium_public_health.disease.state.BaseDiseaseState(state_id, allow_self_transition=True, side_effect_function=None, cause_type='cause')[source]

Base class for disease states in a state machine model.

Provides shared infrastructure for tracking state event times, prevalence-based initialization weights, and transitions.

Parameters:

state_id (str)
allow_self_transition (bool)
side_effect_function (Callable | None)
cause_type (str)

property configuration_defaults: dict[str, Any]

Provides default configuration values for this state.

Extends the parent State’s configuration with disease-specific data sources for prevalence.

Configuration structure:

{component_name}:
    data_sources:
        prevalence:
            Source for prevalence data used to initialize simulants
            into this state. Default is the value set on the instance
            (typically 0.0).
        birth_prevalence:
            Source for birth prevalence data used to initialize
            newborn simulants. Default is the value set on the
            instance (typically 0.0).

property has_dwell_time: bool: Whether this state has a non-zero dwell time.

setup(builder)[source]

Perform this component’s setup.

Parameters:: builder (Builder) – Access point for utilizing framework interfaces during setup.
Return type:: None

get_dwell_time(builder)[source]

Load the dwell time for this state from configuration.

Parameters:: builder (Builder) – Access point for utilizing framework interfaces during setup.
Return type:: SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]]
Returns:: The dwell time data, converted from a Timedelta to days if applicable.

initialize_event_time_and_count(pop_data)[source]

Add this state’s columns to the simulation state table.

Parameters:: pop_data (SimulantData) – Metadata about the simulants being initialized.
Return type:: None

get_initialization_parameters()[source]

Exclude side effect function and cause type from name and __repr__.

Return type:: dict[str, Any]

get_initial_event_times(pop_data)[source]

Get initial event times and counts for new simulants.

Parameters:: pop_data (SimulantData) – Metadata about the simulants being initialized.
Return type:: DataFrame
Returns:: A DataFrame with event time and count columns for new simulants.

transition_side_effect(index, event_time)[source]

Updates the simulation state and triggers any side effects associated with this state.

Parameters:

index – An iterable of integer labels for the simulants.
event_time – The time at which this transition occurs.

Return type:

None

get_transition_names()[source]

Get the names of all transitions from this state.

Return type:: list[str]
Returns:: The transition names formatted as {from_state}_TO_{to_state}.

add_rate_transition(output, triggered=Trigger.NOT_TRIGGERED, transition_rate=None, rate_type='transition_rate')[source]

Builds a RateTransition from this state to the given state.

Parameters:

output (BaseDiseaseState) – The end state after the transition.
triggered (Trigger) – The trigger for the transition.
transition_rate (SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]] | None) – The transition rate source. Can be the data itself, a function to retrieve the data, or the artifact key containing the data.
rate_type (str) – The type of rate. Can be “incidence_rate”, “transition_rate”, or “remission_rate”.

Return type:

RateTransition

Returns:

The created transition object.

add_proportion_transition(output, triggered=Trigger.NOT_TRIGGERED, proportion=None)[source]

Builds a ProportionTransition from this state to the given state.

Parameters:

output (BaseDiseaseState) – The end state after the transition.
triggered (Trigger) – The trigger for the transition.
proportion (SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]] | None) – The proportion source. Can be the data itself, a function to retrieve the data, or the artifact key containing the data.

Return type:

ProportionTransition

Returns:

The created transition object.

add_dwell_time_transition(output, triggered=Trigger.NOT_TRIGGERED)[source]

Build a dwell time transition from this state to the given state.

Parameters:

output (BaseDiseaseState) – The end state after the transition.
triggered (Trigger) – The trigger for the transition.

Return type:

Transition

Returns:

The created transition object.

class vivarium_public_health.disease.state.NonDiseasedState(state_id, allow_self_transition=True, side_effect_function=None, cause_type='cause', name_prefix='')[source]

Base class for states representing the absence of a disease condition.

Provides a name prefix mechanism for creating properly named disease states (e.g., susceptible_to_ or recovered_from_).

Parameters:

state_id (str)
allow_self_transition (bool)
side_effect_function (Callable | None)
cause_type (str)
name_prefix (str)

add_rate_transition(output, triggered=Trigger.NOT_TRIGGERED, transition_rate=None)[source]

Build a rate transition from this state to the given state.

If no transition rate is provided, uses the incidence rate for the output state from the artifact.

Parameters:

output (BaseDiseaseState) – The end state after the transition.
triggered (Trigger) – The trigger for the transition.
transition_rate (SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]] | None) – The transition rate source. Can be the data itself, a function to retrieve the data, or the artifact key containing the data.

Return type:

RateTransition

Returns:

The created transition object.

class vivarium_public_health.disease.state.SusceptibleState(state_id, allow_self_transition=True, side_effect_function=None, cause_type='cause')[source]

State representing susceptibility to a disease.

Automatically prepends susceptible_to_ to the state ID.

Parameters:

state_id (str)
allow_self_transition (bool)
side_effect_function (Callable | None)
cause_type (str)

has_initialization_weights()[source]

Whether this state has initialization weights.

Return type:: bool
Returns:: Always True for susceptible states.

class vivarium_public_health.disease.state.RecoveredState(state_id, allow_self_transition=True, side_effect_function=None, cause_type='cause')[source]

State representing recovery from a disease.

Automatically prepends recovered_from_ to the state ID.

Parameters:

state_id (str)
allow_self_transition (bool)
side_effect_function (Callable | None)
cause_type (str)

class vivarium_public_health.disease.state.ExcessMortalityState(*args, **kwargs)[source]

Mixin for disease states that may have excess mortality.

Parameters:

args (Any)
kwargs (Any)

has_excess_mortality(builder)[source]

Determine whether this state has non-zero excess mortality.

Parameters:: builder (Builder) – Access point for utilizing framework interfaces during setup.
Return type:: bool
Returns:: True if the state has non-zero excess mortality data.

class vivarium_public_health.disease.state.DiseaseState(state_id, allow_self_transition=True, side_effect_function=None, cause_type='cause', prevalence=None, birth_prevalence=0.0, dwell_time=0.0, disability_weight=None, excess_mortality_rate=None)[source]

State representing a disease in a state machine model.

Parameters:

state_id (str)
allow_self_transition (bool)
side_effect_function (Callable | None)
cause_type (str)
prevalence (DataInput | None)
birth_prevalence (DataInput)
dwell_time (DataInput)
disability_weight (DataInput | None)
excess_mortality_rate (DataInput | None)

property configuration_defaults: dict[str, Any]

Provides default configuration values for this disease state.

Extends BaseDiseaseState’s configuration with additional data sources for disease burden metrics.

Configuration structure:

{component_name}:
    data_sources:
        prevalence:
            Source for prevalence data. Defaults to the
            ``prevalence`` constructor argument, or if not
            provided, loads from artifact at
            ``cause.{state_id}.prevalence``.
        birth_prevalence:
            Source for birth prevalence data. Defaults to the
            ``birth_prevalence`` constructor argument, or if not
            provided, loads from artifact at
            ``cause.{state_id}.birth_prevalence``.
        dwell_time:
            Source for dwell time data (minimum time in state
            before transition). Defaults to the ``dwell_time``
            constructor argument, or if not provided, defaults
            to 0 (no minimum dwell time).
        disability_weight:
            Source for disability weight data used to calculate
            years lived with disability (YLDs). Defaults to the
            ``disability_weight`` constructor argument, or if not
            provided, loads from artifact at
            ``cause.{state_id}.disability_weight``.
        excess_mortality_rate:
            Source for excess mortality rate data. Defaults to the
            ``excess_mortality_rate`` constructor argument, or if
            not provided, loads from artifact at
            ``cause.{state_id}.excess_mortality_rate``.

setup(builder)[source]

Performs this component’s simulation setup.

Parameters:: builder (Builder) – Access point for utilizing framework interfaces during setup.
Return type:: None

get_prevalence_source(prevalence)[source]

Resolve the prevalence data source.

Parameters:: prevalence (SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]] | None) – The prevalence source provided at construction, or None to use the default artifact key.
Return type:: SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]]
Returns:: The resolved prevalence data source.

get_disability_weight_source(disability_weight)[source]

Resolve the disability weight data source.

Parameters:: disability_weight (SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]] | None) – The disability weight source provided at construction, or None to use the default artifact key.
Return type:: SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]]
Returns:: The resolved disability weight data source.

register_disability_weight_pipeline(builder)[source]

Register the disability weight pipeline with the simulation.

Parameters:: builder (Builder) – Access point for utilizing framework interfaces during setup.
Return type:: None

get_excess_mortality_rate_source(excess_mortality_rate)[source]

Resolve the excess mortality rate data source.

Parameters:: excess_mortality_rate (SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]] | None) – The excess mortality rate source provided at construction, or None to use the default artifact key.
Return type:: SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]]
Returns:: The resolved excess mortality rate data source.

register_excess_mortality_rate_pipeline(builder)[source]

Register the excess mortality rate pipeline with the simulation.

Parameters:: builder (Builder) – Access point for utilizing framework interfaces during setup.
Return type:: None

get_randomness_prevalence(builder)[source]

Get a randomness stream for assigning prevalent cases.

Parameters:: builder (Builder) – Access point for utilizing framework interfaces during setup.
Return type:: RandomnessStream
Returns:: A randomness stream for prevalent case assignment.

has_initialization_weights()[source]

Whether this state has initialization weights.

Return type:: bool
Returns:: Always True for disease states.

add_rate_transition(output, triggered=Trigger.NOT_TRIGGERED, transition_rate=None, rate_type='transition_rate')[source]

Build a rate transition from this state to the given state.

If no transition rate is provided, uses the remission rate for this state from the artifact.

Parameters:

output (BaseDiseaseState) – The end state after the transition.
triggered (Trigger) – The trigger for the transition.
transition_rate (SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]] | str | Callable[[Builder], SupportsFloat | Timedelta | timedelta | Timestamp | datetime | DataFrame | list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | tuple[SupportsFloat | Timedelta | timedelta | Timestamp | datetime, ...] | Mapping[str, list[SupportsFloat | Timedelta | timedelta | Timestamp | datetime] | list[str]]] | None) – The transition rate source. Can be the data itself, a function to retrieve the data, or the artifact key containing the data.
rate_type (str) – The type of rate. Can be “incidence_rate”, “transition_rate”, or “remission_rate”.

Return type:

RateTransition

Returns:

The created transition object.

next_state(index, event_time, population_view)[source]

Moves a population among different disease states.

Parameters:

index – An iterable of integer labels for the simulants.
event_time – The time at which this transition occurs.
population_view – A view of the internal state of the simulation.

Return type:

None

compute_disability_weight(index)[source]

Gets the disability weight associated with this state.

Parameters:: index – An iterable of integer labels for the simulants.
Returns:: An iterable of disability weights indexed by the provided index.
Return type:: pd.Series[float]

compute_excess_mortality_rate(index)[source]

Get the excess mortality rate associated with this state.

Parameters:: index – An iterable of integer labels for the simulants.
Returns:: An iterable of excess mortality rates indexed by the provided index.
Return type:: pd.Series[float]

adjust_mortality_rate(index, rates_df)[source]

Modifies the baseline mortality rate for a simulant if they are in this state.

Parameters:

index – An iterable of integer labels for the simulants.
rates_df – A DataFrame of mortality rates.

Returns:

The modified DataFrame of mortality rates.

Return type:

pd.DataFrame

get_initial_event_times(pop_data)[source]

Get initial event times for new simulants, including prevalent cases.

Parameters:: pop_data (SimulantData) – Metadata about the simulants being initialized.
Return type:: DataFrame
Returns:: A DataFrame with event time and count columns for new simulants.

with_condition(index)[source]

Get the subset of simulants who are in this disease state.

Parameters:: index – An iterable of integer labels for the simulants.
Returns:: The subset of simulants who are alive and in this state.
Return type:: pd.Index[int]

class vivarium_public_health.disease.state.TransientDiseaseState(state_id, allow_self_transition=True, side_effect_function=None, cause_type='cause')[source]

A disease state that simulants pass through instantaneously.

Simulants do not remain in this state; they transition to another state within the same time step.

Parameters:

state_id (str)
allow_self_transition (bool)
side_effect_function (Callable | None)
cause_type (str)