vivarium-public-health 3.1.1__tar.gz → 3.1.2__tar.gz

{vivarium_public_health-3.1.1 → vivarium_public_health-3.1.2}/CHANGELOG.rst

@@ -1,3 +1,8 @@
+**3.1.2 - 11/13/24**
+
+ - Modernize type hinting
+ - Refactor: Update DiseaseModel to rely on Machine's time-step and time-step-cleanup methods
+
 **3.1.1 - 11/12/24**
 
  - Feature: Rename DiseaseModel initial_state to residual_state

{vivarium_public_health-3.1.1 → vivarium_public_health-3.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vivarium_public_health
-Version: 3.1.1
+Version: 3.1.2
 Summary: Components for modelling diseases, risks, and interventions with ``vivarium``
 Home-page: https://github.com/ihmeuw/vivarium_public_health
 Author: The vivarium developers

vivarium_public_health-3.1.2/src/vivarium_public_health/_version.py

@@ -0,0 +1 @@
+__version__ = "3.1.2"

{vivarium_public_health-3.1.1 → vivarium_public_health-3.1.2}/src/vivarium_public_health/disease/model.py

@@ -9,8 +9,8 @@ transitions at simulation initialization and during transitions.
 
 """
 import warnings
-from collections.abc import Iterable
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from collections.abc import Callable, Iterable
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -35,7 +35,7 @@ class DiseaseModel(Machine):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             f"{self.name}": {
                 "data_sources": {
@@ -45,15 +45,15 @@ class DiseaseModel(Machine):
         }
 
     @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": [],
@@ -61,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()
         ]
@@ -77,8 +77,8 @@ 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,
@@ -131,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:
@@ -200,12 +198,6 @@ class DiseaseModel(Machine):
             )
         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 #
     ##################################
@@ -225,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:

{vivarium_public_health-3.1.1 → vivarium_public_health-3.1.2}/src/vivarium_public_health/disease/special_disease.py

@@ -10,7 +10,7 @@ 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
@@ -95,7 +95,7 @@ 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": {
@@ -111,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,
@@ -119,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-3.1.1 → vivarium_public_health-3.1.2}/src/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,7 +279,7 @@ class DiseaseState(BaseDiseaseState):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             f"{self.name}": {
                 "data_sources": {
@@ -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-3.1.1 → vivarium_public_health-3.1.2}/src/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,7 +40,7 @@ class RateTransition(Transition):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             f"{self.name}": {
                 "data_sources": {
@@ -49,7 +50,7 @@ class RateTransition(Transition):
         }
 
     @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,7 +153,7 @@ class ProportionTransition(Transition):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             f"{self.name}": {
                 "data_sources": {
@@ -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-3.1.1 → vivarium_public_health-3.1.2}/src/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-3.1.1 → vivarium_public_health-3.1.2}/src/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-3.1.1 → vivarium_public_health-3.1.2}/src/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-3.1.1 → vivarium_public_health-3.1.2}/src/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-3.1.1 → vivarium_public_health-3.1.2}/src/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
 
     #####################

{vivarium_public_health-3.1.1 → vivarium_public_health-3.1.2}/src/vivarium_public_health/mslt/population.py

@@ -8,8 +8,6 @@ multi-state lifetable simulations.
 
 """
 
-from typing import List, Optional
-
 import numpy as np
 import pandas as pd
 from vivarium import Component
@@ -50,7 +48,7 @@ class BasePopulation(Component):
     ##############
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         return [
             "age",
             "sex",
@@ -71,7 +69,7 @@ class BasePopulation(Component):
         ]
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return ["tracked"]
 
     #####################
@@ -122,7 +120,7 @@ class Mortality(Component):
     ##############
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return [
             "population",
             "bau_population",
@@ -190,7 +188,7 @@ class Disability(Component):
     ##############
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return [
             "bau_yld_rate",
             "yld_rate",