vivarium-public-health 3.1.0__py3-none-any.whl → 3.1.2__py3-none-any.whl

vivarium_public_health/_version.py

@@ -1 +1 @@
-__version__ = "3.1.0"
+__version__ = "3.1.2"

vivarium_public_health/disease/model.py

@@ -8,8 +8,9 @@ function is to provide coordination across a set of disease states and
 transitions at simulation initialization and during transitions.
 
 """
-
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import warnings
+from collections.abc import Callable, Iterable
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -34,25 +35,25 @@ class DiseaseModel(Machine):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             f"{self.name}": {
                 "data_sources": {
-                    "cause_specific_mortality_rate": "self::load_cause_specific_mortality_rate"
+                    "cause_specific_mortality_rate": self.load_cause_specific_mortality_rate,
                 },
             },
         }
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         return [self.state_column]
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return ["age", "sex"]
 
     @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
+    def initialization_requirements(self) -> dict[str, list[str]]:
         return {
             "requires_columns": ["age", "sex"],
             "requires_values": [],
@@ -60,11 +61,11 @@ class DiseaseModel(Machine):
         }
 
     @property
-    def state_names(self) -> List[str]:
+    def state_names(self) -> list[str]:
         return [s.state_id for s in self.states]
 
     @property
-    def transition_names(self) -> List[TransitionString]:
+    def transition_names(self) -> list[TransitionString]:
         return [
             state_name for state in self.states for state_name in state.get_transition_names()
         ]
@@ -76,19 +77,38 @@ class DiseaseModel(Machine):
     def __init__(
         self,
         cause: str,
-        initial_state: Optional[BaseDiseaseState] = None,
-        get_data_functions: Optional[Dict[str, Callable]] = None,
+        initial_state: BaseDiseaseState | None = None,
+        get_data_functions: dict[str, Callable] | None = None,
         cause_type: str = "cause",
+        states: Iterable[BaseDiseaseState] = (),
+        residual_state: BaseDiseaseState | None = None,
         **kwargs,
     ):
-        super().__init__(cause, **kwargs)
+        super().__init__(cause, states=states, **kwargs)
         self.cause = cause
         self.cause_type = cause_type
 
         if initial_state is not None:
-            self.initial_state = initial_state.state_id
+            warnings.warn(
+                "In the future, the 'initial_state' argument to DiseaseModel"
+                " will be used to initialize all simulants into that state. To"
+                " retain the current behavior of defining a residual state, use"
+                " the 'residual_state' argument.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+            if residual_state:
+                raise DiseaseModelError(
+                    "A DiseaseModel cannot be initialized with both"
+                    " 'initial_state and 'residual_state'."
+                )
+
+            self.residual_state = initial_state.state_id
+        elif residual_state is not None:
+            self.residual_state = residual_state.state_id
         else:
-            self.initial_state = self._get_default_initial_state()
+            self.residual_state = self._get_default_residual_state()
 
         self._get_data_functions = (
             get_data_functions if get_data_functions is not None else {}
@@ -111,9 +131,7 @@ class DiseaseModel(Machine):
     # Setup methods #
     #################
 
-    def load_cause_specific_mortality_rate(
-        self, builder: Builder
-    ) -> Union[float, pd.DataFrame]:
+    def load_cause_specific_mortality_rate(self, builder: Builder) -> float | pd.DataFrame:
         if "cause_specific_mortality_rate" not in self._get_data_functions:
             only_morbid = builder.data.load(f"cause.{self.cause}.restrictions")["yld_only"]
             if only_morbid:
@@ -135,7 +153,7 @@ class DiseaseModel(Machine):
     def on_initialize_simulants(self, pop_data: SimulantData) -> None:
         population = self.population_view.subview(["age", "sex"]).get(pop_data.index)
 
-        assert self.initial_state in {s.state_id for s in self.states}
+        assert self.residual_state in {s.state_id for s in self.states}
 
         if pop_data.user_data["sim_state"] == "setup":  # simulation start
             if self.configuration_age_start != self.configuration_age_end != 0:
@@ -176,16 +194,10 @@ class DiseaseModel(Machine):
             )
         else:
             condition_column = pd.Series(
-                self.initial_state, index=population.index, name=self.state_column
+                self.residual_state, index=population.index, name=self.state_column
             )
         self.population_view.update(condition_column)
 
-    def on_time_step(self, event: Event) -> None:
-        self.transition(event.index, event.time)
-
-    def on_time_step_cleanup(self, event: Event) -> None:
-        self.cleanup(event.index, event.time)
-
     ##################################
     # Pipeline sources and modifiers #
     ##################################
@@ -197,7 +209,7 @@ class DiseaseModel(Machine):
     # Helper functions #
     ####################
 
-    def _get_default_initial_state(self):
+    def _get_default_residual_state(self):
         susceptible_states = [s for s in self.states if isinstance(s, SusceptibleState)]
         if len(susceptible_states) != 1:
             raise DiseaseModelError("Disease model must have exactly one SusceptibleState.")
@@ -205,7 +217,7 @@ class DiseaseModel(Machine):
 
     def get_state_weights(
         self, pop_index: pd.Index, prevalence_type: str
-    ) -> Tuple[List[str], Union[np.ndarray, None]]:
+    ) -> tuple[list[str], np.ndarray | None]:
         states = [state for state in self.states if state.lookup_tables.get(prevalence_type)]
 
         if not states:
@@ -219,7 +231,7 @@ class DiseaseModel(Machine):
         weights = np.array(weights).T
         weights_bins = np.cumsum(weights, axis=1)
 
-        state_names = [s.state_id for s in states] + [self.initial_state]
+        state_names = [s.state_id for s in states] + [self.residual_state]
 
         return state_names, weights_bins
 

vivarium_public_health/disease/special_disease.py

@@ -10,11 +10,10 @@ This module contains frequently used, but non-standard disease models.
 import re
 from collections import namedtuple
 from operator import gt, lt
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import pandas as pd
 from vivarium import Component
-from vivarium.framework.engine import Builder
 from vivarium.framework.event import Event
 from vivarium.framework.population import SimulantData
 from vivarium.framework.values import list_combiner, union_post_processor
@@ -96,13 +95,13 @@ class RiskAttributableDisease(Component):
         return f"risk_attributable_disease.{self.cause.name}"
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             self.name: {
                 "data_sources": {
                     "raw_disability_weight": f"{self.cause}.disability_weight",
-                    "cause_specific_mortality_rate": "self::load_cause_specific_mortality_rate_data",
-                    "excess_mortality_rate": "self::load_excess_mortality_rate_data",
+                    "cause_specific_mortality_rate": self.load_cause_specific_mortality_rate_data,
+                    "excess_mortality_rate": self.load_excess_mortality_rate_data,
                     "population_attributable_fraction": 0,
                 },
                 "threshold": None,
@@ -112,7 +111,7 @@ class RiskAttributableDisease(Component):
         }
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         return [
             self.cause.name,
             self.diseased_event_time_column,
@@ -120,11 +119,11 @@ class RiskAttributableDisease(Component):
         ]
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return ["alive"]
 
     @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
+    def initialization_requirements(self) -> dict[str, list[str]]:
         return {
             "requires_columns": [],
             "requires_values": [f"{self.risk.name}.exposure"],

vivarium_public_health/disease/state.py

@@ -7,7 +7,8 @@ This module contains tools to manage standard disease states.
 
 """
 
-from typing import Any, Callable, Dict, List, Optional
+from collections.abc import Callable
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -36,11 +37,11 @@ class BaseDiseaseState(State):
         return [self.event_time_column, self.event_count_column]
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return [self.model, "alive"]
 
     @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
+    def initialization_requirements(self) -> dict[str, list[str]]:
         return {
             "requires_columns": [self.model],
             "requires_values": [],
@@ -55,7 +56,7 @@ class BaseDiseaseState(State):
         self,
         state_id: str,
         allow_self_transition: bool = False,
-        side_effect_function: Optional[Callable] = None,
+        side_effect_function: Callable | None = None,
         cause_type: str = "cause",
     ):
         super().__init__(state_id, allow_self_transition)  # becomes state_id
@@ -83,7 +84,7 @@ class BaseDiseaseState(State):
     # Helper methods #
     ##################
 
-    def get_initialization_parameters(self) -> Dict[str, Any]:
+    def get_initialization_parameters(self) -> dict[str, Any]:
         """Exclude side effect function and cause type from name and __repr__."""
         initialization_parameters = super().get_initialization_parameters()
         for key in ["side_effect_function", "cause_type"]:
@@ -118,7 +119,7 @@ class BaseDiseaseState(State):
     # Public methods #
     ##################
 
-    def get_transition_names(self) -> List[str]:
+    def get_transition_names(self) -> list[str]:
         transitions = []
         for trans in self.transition_set.transitions:
             init_state = trans.input_state.name.split(".")[1]
@@ -129,7 +130,7 @@ class BaseDiseaseState(State):
     def add_rate_transition(
         self,
         output: "BaseDiseaseState",
-        get_data_functions: Dict[str, Callable] = None,
+        get_data_functions: dict[str, Callable] = None,
         triggered=Trigger.NOT_TRIGGERED,
     ) -> RateTransition:
         """Builds a RateTransition from this state to the given state.
@@ -156,7 +157,7 @@ class BaseDiseaseState(State):
     def add_proportion_transition(
         self,
         output: "BaseDiseaseState",
-        get_data_functions: Optional[Dict[str, Callable]] = None,
+        get_data_functions: dict[str, Callable] | None = None,
         triggered=Trigger.NOT_TRIGGERED,
     ) -> ProportionTransition:
         """Builds a ProportionTransition from this state to the given state.
@@ -198,7 +199,7 @@ class NonDiseasedState(BaseDiseaseState):
         self,
         state_id: str,
         allow_self_transition: bool = False,
-        side_effect_function: Optional[Callable] = None,
+        side_effect_function: Callable | None = None,
         cause_type: str = "cause",
         name_prefix: str = "",
     ):
@@ -218,7 +219,7 @@ class NonDiseasedState(BaseDiseaseState):
     def add_rate_transition(
         self,
         output: BaseDiseaseState,
-        get_data_functions: Dict[str, Callable] = None,
+        get_data_functions: dict[str, Callable] = None,
         **kwargs,
     ) -> RateTransition:
         if get_data_functions is None:
@@ -241,7 +242,7 @@ class SusceptibleState(NonDiseasedState):
         self,
         state_id: str,
         allow_self_transition: bool = False,
-        side_effect_function: Optional[Callable] = None,
+        side_effect_function: Callable | None = None,
         cause_type: str = "cause",
     ):
         super().__init__(
@@ -258,7 +259,7 @@ class RecoveredState(NonDiseasedState):
         self,
         state_id: str,
         allow_self_transition: bool = False,
-        side_effect_function: Optional[Callable] = None,
+        side_effect_function: Callable | None = None,
         cause_type: str = "cause",
     ):
         super().__init__(
@@ -278,15 +279,15 @@ class DiseaseState(BaseDiseaseState):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             f"{self.name}": {
                 "data_sources": {
-                    "prevalence": "self::load_prevalence",
-                    "birth_prevalence": "self::load_birth_prevalence",
-                    "dwell_time": "self::load_dwell_time",
-                    "disability_weight": "self::load_disability_weight",
-                    "excess_mortality_rate": "self::load_excess_mortality_rate",
+                    "prevalence": self.load_prevalence,
+                    "birth_prevalence": self.load_birth_prevalence,
+                    "dwell_time": self.load_dwell_time,
+                    "disability_weight": self.load_disability_weight,
+                    "excess_mortality_rate": self.load_excess_mortality_rate,
                 },
             },
         }
@@ -299,10 +300,10 @@ class DiseaseState(BaseDiseaseState):
         self,
         state_id: str,
         allow_self_transition: bool = False,
-        side_effect_function: Optional[Callable] = None,
+        side_effect_function: Callable | None = None,
         cause_type: str = "cause",
-        get_data_functions: Optional[Dict[str, Callable]] = None,
-        cleanup_function: Optional[Callable] = None,
+        get_data_functions: dict[str, Callable] | None = None,
+        cleanup_function: Callable | None = None,
     ):
         """
         Parameters
@@ -472,7 +473,7 @@ class DiseaseState(BaseDiseaseState):
     def add_rate_transition(
         self,
         output: BaseDiseaseState,
-        get_data_functions: Dict[str, Callable] = None,
+        get_data_functions: dict[str, Callable] = None,
         **kwargs,
     ) -> RateTransition:
         if get_data_functions is None:

vivarium_public_health/disease/transition.py

@@ -7,7 +7,8 @@ This module contains tools to model transitions between disease states.
 
 """
 
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
 
 import pandas as pd
 from vivarium.framework.engine import Builder
@@ -39,17 +40,17 @@ class RateTransition(Transition):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             f"{self.name}": {
                 "data_sources": {
-                    "transition_rate": "self::load_transition_rate",
+                    "transition_rate": self.load_transition_rate,
                 },
             },
         }
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return ["alive"]
 
     @property
@@ -77,7 +78,7 @@ class RateTransition(Transition):
         self,
         input_state: "BaseDiseaseState",
         output_state: "BaseDiseaseState",
-        get_data_functions: Dict[str, Callable] = None,
+        get_data_functions: dict[str, Callable] = None,
         triggered=Trigger.NOT_TRIGGERED,
     ):
         super().__init__(
@@ -108,7 +109,7 @@ class RateTransition(Transition):
     # Setup methods #
     #################
 
-    def load_transition_rate(self, builder: Builder) -> Union[float, pd.DataFrame]:
+    def load_transition_rate(self, builder: Builder) -> float | pd.DataFrame:
         if "incidence_rate" in self._get_data_functions:
             rate_data = self._get_data_functions["incidence_rate"](
                 builder, self.output_state.state_id
@@ -152,11 +153,11 @@ class ProportionTransition(Transition):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             f"{self.name}": {
                 "data_sources": {
-                    "proportion": "self::load_proportion",
+                    "proportion": self.load_proportion,
                 },
             },
         }
@@ -169,7 +170,7 @@ class ProportionTransition(Transition):
         self,
         input_state: "BaseDiseaseState",
         output_state: "BaseDiseaseState",
-        get_data_functions: Dict[str, Callable] = None,
+        get_data_functions: dict[str, Callable] = None,
         triggered=Trigger.NOT_TRIGGERED,
     ):
         super().__init__(
@@ -183,7 +184,7 @@ class ProportionTransition(Transition):
     # Setup methods #
     #################
 
-    def load_proportion(self, builder: Builder) -> Union[float, pd.DataFrame]:
+    def load_proportion(self, builder: Builder) -> float | pd.DataFrame:
         if "proportion" not in self._get_data_functions:
             raise ValueError("Must supply a proportion function")
         return self._get_data_functions["proportion"](builder, self.output_state.state_id)

vivarium_public_health/mslt/delay.py

@@ -8,7 +8,7 @@ lifetable simulation.
 
 """
 
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -93,7 +93,7 @@ class DelayedRisk(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             self.risk: {
                 "constant_prevalence": False,
@@ -103,15 +103,15 @@ class DelayedRisk(Component):
         }
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         return self._bin_names
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return ["age", "sex", "population"]
 
     @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
+    def initialization_requirements(self) -> dict[str, list[str]]:
         return {
             "requires_columns": ["age", "sex", "population"],
             "requires_values": [],

vivarium_public_health/mslt/disease.py

@@ -8,7 +8,7 @@ simulations.
 
 """
 
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -143,7 +143,7 @@ class Disease(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             self.disease: {
                 "simplified_no_remission_equations": False,
@@ -151,7 +151,7 @@ class Disease(Component):
         }
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         columns = []
         for scenario in ["", "_intervention"]:
             for rate in ["_S", "_C"]:
@@ -160,11 +160,11 @@ class Disease(Component):
         return columns
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return ["age", "sex"]
 
     @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
+    def initialization_requirements(self) -> dict[str, list[str]]:
         return {
             "requires_columns": ["age", "sex"],
             "requires_values": [],

vivarium_public_health/mslt/intervention.py

@@ -8,7 +8,7 @@ simulations.
 
 """
 
-from typing import Any, Dict
+from typing import Any
 
 from vivarium import Component
 from vivarium.framework.engine import Builder
@@ -22,7 +22,7 @@ class ModifyAllCauseMortality(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "intervention": {
                 self.intervention: {
@@ -62,7 +62,7 @@ class ModifyDiseaseRate(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "intervention": {
                 self.intervention: {
@@ -144,7 +144,7 @@ class ModifyAcuteDiseaseIncidence(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "intervention": {
                 self.intervention: {
@@ -187,7 +187,7 @@ class ModifyAcuteDiseaseMorbidity(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "intervention": {
                 self.intervention: {
@@ -228,7 +228,7 @@ class ModifyAcuteDiseaseMortality(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "intervention": {
                 self.intervention: {
@@ -269,7 +269,7 @@ class TobaccoFreeGeneration(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "tobacco_free_generation": {
                 "year": 2020,
@@ -310,7 +310,7 @@ class TobaccoEradication(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "tobacco_eradication": {
                 "year": 2020,

vivarium_public_health/mslt/magic_wand_components.py

@@ -8,7 +8,7 @@ multi-state lifetable simulations.
 
 """
 
-from typing import Any, Dict
+from typing import Any
 
 from vivarium import Component
 from vivarium.framework.engine import Builder
@@ -73,7 +73,7 @@ class ModifyAcuteDiseaseYLD(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "intervention": {
                 self.disease: {
@@ -112,7 +112,7 @@ class ModifyAcuteDiseaseMortality(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "intervention": {
                 self.disease: {

vivarium_public_health/mslt/observer.py

@@ -8,8 +8,6 @@ multi-state lifetable simulations.
 
 """
 
-from typing import List, Optional
-
 import pandas as pd
 from vivarium import Component
 from vivarium.framework.engine import Builder
@@ -76,7 +74,7 @@ class MorbidityMortality(Component):
     ##############
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return [
             "age",
             "sex",
@@ -224,7 +222,7 @@ class Disease(Component):
     ##############
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return [
             "age",
             "sex",
@@ -238,7 +236,7 @@ class Disease(Component):
     # Lifecycle methods #
     #####################
 
-    def __init__(self, disease: str, output_suffix: Optional[str] = None):
+    def __init__(self, disease: str, output_suffix: str | None = None):
         super().__init__()
         self.disease = disease
         if output_suffix is None:
@@ -335,7 +333,7 @@ class TobaccoPrevalence(Component):
     ##############
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return ["age", "sex", "bau_population", "population"] + self._bin_names
 
     #####################