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

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",

vivarium_public_health/plugins/parser.py

@@ -9,9 +9,10 @@ that can parse configurations of components specific to the Vivarium Public
 Health package.
 
 """
-
+import warnings
+from collections.abc import Callable
 from importlib import import_module
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any
 
 import pandas as pd
 from layered_config_tree import LayeredConfigTree
@@ -37,7 +38,7 @@ from vivarium_public_health.utilities import TargetString
 class CausesParsingErrors(ParsingError):
     """Error raised when there are any errors parsing a cause model configuration."""
 
-    def __init__(self, messages: List[str]):
+    def __init__(self, messages: list[str]):
         super().__init__("\n - " + "\n - ".join(messages))
 
 
@@ -56,11 +57,13 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     DEFAULT_MODEL_CONFIG = {
         "model_type": f"{DiseaseModel.__module__}.{DiseaseModel.__name__}",
         "initial_state": None,
+        "residual_state": None,
     }
     """Default cause model configuration if it's not explicitly specified.
     
-    If the initial state is not specified, the cause  model must have a state 
-    named 'susceptible'.
+    Initial state and residual state cannot both be provided. If neither initial
+    state nor residual state has been specified, the cause  model must have a
+    state named 'susceptible'.
     """
 
     DEFAULT_STATE_CONFIG = {
@@ -79,7 +82,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     This value is used if the transition configuration does not explicity specify it.
     """
 
-    def parse_component_config(self, component_config: LayeredConfigTree) -> List[Component]:
+    def parse_component_config(self, component_config: LayeredConfigTree) -> list[Component]:
         """Parses the component configuration and returns a list of components.
 
         In particular, this method looks for an `external_configuration` key
@@ -104,7 +107,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
             causes:
                 cause_1:
                     model_type: vivarium_public_health.disease.DiseaseModel
-                    initial_state: susceptible
+                    residual_state: susceptible
                     states:
                         susceptible:
                             cause_type: cause
@@ -214,7 +217,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
     def _get_cause_model_components(
         self, causes_config: LayeredConfigTree
-    ) -> List[Component]:
+    ) -> list[Component]:
         """Parses the cause model configuration and returns the `DiseaseModel` components.
 
         Parameters
@@ -234,7 +237,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
                 data_sources_config = cause_config.data_sources
                 data_sources = self._get_data_sources(data_sources_config)
 
-            states: Dict[str, BaseDiseaseState] = {
+            states: dict[str, BaseDiseaseState] = {
                 state_name: self._get_state(state_name, state_config, cause_name)
                 for state_name, state_config in cause_config.states.items()
             }
@@ -247,11 +250,13 @@ class CausesConfigurationParser(ComponentConfigurationParser):
                 )
 
             model_type = import_by_path(cause_config.model_type)
-            initial_state = states.get(cause_config.initial_state, None)
+            residual_state = states.get(
+                cause_config.residual_state, states.get(cause_config.initial_state, None)
+            )
             model = model_type(
                 cause_name,
-                initial_state=initial_state,
                 states=list(states.values()),
+                residual_state=residual_state,
                 get_data_functions=data_sources,
             )
             cause_models.append(model)
@@ -347,7 +352,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
     def _get_data_sources(
         self, config: LayeredConfigTree
-    ) -> Dict[str, Callable[[Builder, Any], Any]]:
+    ) -> dict[str, Callable[[Builder, Any], Any]]:
         """Parses a data sources configuration and returns the data sources.
 
         Parameters
@@ -362,9 +367,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         return {name: self._get_data_source(name, config[name]) for name in config.keys()}
 
     @staticmethod
-    def _get_data_source(
-        name: str, source: Union[str, float]
-    ) -> Callable[[Builder, Any], Any]:
+    def _get_data_source(name: str, source: str | float) -> Callable[[Builder, Any], Any]:
         """Parses a data source and returns a callable that can be used to retrieve the data.
 
         Parameters
@@ -403,7 +406,14 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     # Validation methods #
     ######################
 
-    _CAUSE_KEYS = {"model_type", "initial_state", "states", "transitions", "data_sources"}
+    _CAUSE_KEYS = {
+        "model_type",
+        "initial_state",
+        "states",
+        "transitions",
+        "data_sources",
+        "residual_state",
+    }
     _STATE_KEYS = {
         "state_type",
         "cause_type",
@@ -489,7 +499,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         if error_messages:
             raise CausesParsingErrors(error_messages)
 
-    def _validate_cause(self, cause_name: str, cause_config: Dict[str, Any]) -> List[str]:
+    def _validate_cause(self, cause_name: str, cause_config: dict[str, Any]) -> list[str]:
         """Validates a cause configuration and returns a list of error messages.
 
         Parameters
@@ -535,11 +545,28 @@ class CausesConfigurationParser(ComponentConfigurationParser):
             )
         else:
             initial_state = cause_config.get("initial_state", None)
-            if initial_state is not None and initial_state not in states_config:
+            residual_state = cause_config.get("residual_state", None)
+            if initial_state is not None:
+                warnings.warn(
+                    "In the future, the 'initial_state' cause configuration will"
+                    " be used to initialize all simulants into that state. To"
+                    " retain the current behavior of defining a residual state,"
+                    " use the 'residual_state' cause configuration.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+                if residual_state is None:
+                    residual_state = initial_state
+                else:
+                    error_messages.append(
+                        "A cause may not have both 'initial_state and"
+                        " 'residual_state' configurations."
+                    )
+
+            if residual_state is not None and residual_state not in states_config:
                 error_messages.append(
-                    f"Initial state '{cause_config['initial_state']}' for cause "
-                    f"'{cause_name}' must be present in the states for cause "
-                    f"'{cause_name}."
+                    f"Residual state '{residual_state}' for cause '{cause_name}'"
+                    f" must be present in the states for cause '{cause_name}."
                 )
             for state_name, state_config in states_config.items():
                 error_messages += self._validate_state(cause_name, state_name, state_config)
@@ -563,8 +590,8 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         return error_messages
 
     def _validate_state(
-        self, cause_name: str, state_name: str, state_config: Dict[str, Any]
-    ) -> List[str]:
+        self, cause_name: str, state_name: str, state_config: dict[str, Any]
+    ) -> list[str]:
         """Validates a state configuration and returns a list of error messages.
 
         Parameters
@@ -655,9 +682,9 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         self,
         cause_name: str,
         transition_name: str,
-        transition_config: Dict[str, Any],
-        states_config: Dict[str, Any],
-    ) -> List[str]:
+        transition_config: dict[str, Any],
+        states_config: dict[str, Any],
+    ) -> list[str]:
         """Validates a transition configuration and returns a list of error messages.
 
         Parameters
@@ -755,8 +782,8 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
     @staticmethod
     def _validate_imported_type(
-        import_path: str, cause_name: str, entity_type: str, entity_name: Optional[str] = None
-    ) -> List[str]:
+        import_path: str, cause_name: str, entity_type: str, entity_name: str | None = None
+    ) -> list[str]:
         """Validates an imported type and returns a list of error messages.
 
         Parameters
@@ -797,8 +824,8 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         return error_messages
 
     def _validate_data_sources(
-        self, config: Dict[str, Any], cause_name: str, config_type: str, config_name: str
-    ) -> List[str]:
+        self, config: dict[str, Any], cause_name: str, config_type: str, config_name: str
+    ) -> list[str]:
         """Validates the data sources in a configuration and returns any error messages.
 
         Parameters

vivarium_public_health/population/add_new_birth_cohorts.py

@@ -7,8 +7,6 @@ This module contains several different models of fertility.
 
 """
 
-from typing import Dict, List, Optional
-
 import numpy as np
 import pandas as pd
 from vivarium import Component
@@ -158,15 +156,15 @@ class FertilityAgeSpecificRates(Component):
     ##############
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         return ["last_birth_time", "parent_id"]
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return ["sex"]
 
     @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
+    def initialization_requirements(self) -> dict[str, list[str]]:
         return {
             "requires_columns": ["sex"],
             "requires_values": [],

vivarium_public_health/population/base_population.py

@@ -8,9 +8,7 @@ characteristics to simulants.
 
 """
 
-from __future__ import annotations
-
-from typing import Callable, Dict, Iterable, List
+from collections.abc import Callable, Iterable
 
 import numpy as np
 import pandas as pd
@@ -48,7 +46,7 @@ class BasePopulation(Component):
     ##############
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         return ["age", "sex", "alive", "location", "entrance_time", "exit_time"]
 
     @property
@@ -91,7 +89,7 @@ class BasePopulation(Component):
     #################
 
     @staticmethod
-    def get_randomness_streams(builder: Builder) -> Dict[str, RandomnessStream]:
+    def get_randomness_streams(builder: Builder) -> dict[str, RandomnessStream]:
         return {
             "general_purpose": builder.randomness.get_stream("population_generation"),
             "bin_selection": builder.randomness.get_stream(
@@ -255,7 +253,7 @@ class AgeOutSimulants(Component):
     ##############
 
     @property
-    def columns_required(self) -> List[str]:
+    def columns_required(self) -> list[str]:
         """A list of the columns this component requires that it did not create."""
         return self._columns_required
 
@@ -290,9 +288,9 @@ def generate_population(
     simulant_ids: pd.Index,
     creation_time: pd.Timestamp,
     step_size: pd.Timedelta,
-    age_params: Dict[str, float],
+    age_params: dict[str, float],
     demographic_proportions: pd.DataFrame,
-    randomness_streams: Dict[str, RandomnessStream],
+    randomness_streams: dict[str, RandomnessStream],
     register_simulants: Callable[[pd.DataFrame], None],
     key_columns: Iterable[str] = ("entrance_time", "age"),
 ) -> pd.DataFrame:
@@ -378,7 +376,7 @@ def _assign_demography_with_initial_age(
     pop_data: pd.DataFrame,
     initial_age: float,
     step_size: pd.Timedelta,
-    randomness_streams: Dict[str, RandomnessStream],
+    randomness_streams: dict[str, RandomnessStream],
     register_simulants: Callable[[pd.DataFrame], None],
 ) -> pd.DataFrame:
     """Assigns age, sex, and location information to the provided simulants given a fixed age.
@@ -441,7 +439,7 @@ def _assign_demography_with_age_bounds(
     pop_data: pd.DataFrame,
     age_start: float,
     age_end: float,
-    randomness_streams: Dict[str, RandomnessStream],
+    randomness_streams: dict[str, RandomnessStream],
     register_simulants: Callable[[pd.DataFrame], None],
     key_columns: Iterable[str] = ("entrance_time", "age"),
 ) -> pd.DataFrame:

vivarium_public_health/population/data_transformations.py

@@ -9,7 +9,6 @@ it into different distributions for sampling.
 """
 
 from collections import namedtuple
-from typing import Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -302,7 +301,7 @@ def smooth_ages(
 def _get_bins_and_proportions(
     pop_data: pd.DataFrame,
     age: AgeValues,
-) -> Tuple[EndpointValues, AgeValues]:
+) -> tuple[EndpointValues, AgeValues]:
     """Finds and returns the bin edges and the population proportions in
     the current and neighboring bins.
 
@@ -376,7 +375,7 @@ def _get_bins_and_proportions(
 
 def _construct_sampling_parameters(
     age: AgeValues, endpoint: EndpointValues, proportion: AgeValues
-) -> Tuple[EndpointValues, EndpointValues, float, float]:
+) -> tuple[EndpointValues, EndpointValues, float, float]:
     """Calculates some sampling distribution parameters from known values.
 
     Parameters
@@ -442,12 +441,12 @@ def _construct_sampling_parameters(
 
 
 def _compute_ages(
-    uniform_rv: Union[np.ndarray, float],
+    uniform_rv: np.ndarray | float,
     start: float,
     height: float,
     slope: float,
     normalization: float,
-) -> Union[np.ndarray, float]:
+) -> np.ndarray | float:
     """Produces samples from the local age distribution.
 
     Parameters

vivarium_public_health/population/mortality.py

@@ -45,7 +45,7 @@ back the modified unmodeled csmr.
 
 """
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 import pandas as pd
 from vivarium import Component
@@ -100,12 +100,12 @@ class Mortality(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "mortality": {
                 "data_sources": {
                     "all_cause_mortality_rate": "cause.all_causes.cause_specific_mortality_rate",
-                    "unmodeled_cause_specific_mortality_rate": "self::load_unmodeled_csmr",
+                    "unmodeled_cause_specific_mortality_rate": self.load_unmodeled_csmr,
                     "life_expectancy": "population.theoretical_minimum_risk_life_expectancy",
                 },
                 "unmodeled_causes": [],
@@ -113,18 +113,18 @@ class Mortality(Component):
         }
 
     @property
-    def standard_lookup_tables(self) -> List[str]:
+    def standard_lookup_tables(self) -> list[str]:
         return [
             "all_cause_mortality_rate",
             "life_expectancy",
         ]
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         return [self.cause_of_death_column_name, self.years_of_life_lost_column_name]
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         return ["alive", "exit_time", "age", "sex"]
 
     @property
@@ -184,7 +184,7 @@ class Mortality(Component):
             requires_columns=required_columns,
         )
 
-    def load_unmodeled_csmr(self, builder: Builder) -> Union[float, pd.DataFrame]:
+    def load_unmodeled_csmr(self, builder: Builder) -> float | pd.DataFrame:
         # todo validate that all data have the same columns
         raw_csmr = 0.0
         for idx, cause in enumerate(builder.configuration[self.name].unmodeled_causes):

vivarium_public_health/results/disability.py

@@ -8,8 +8,6 @@ in the simulation.
 
 """
 
-from typing import Union
-
 import pandas as pd
 from layered_config_tree import LayeredConfigTree
 from loguru import logger
@@ -185,7 +183,7 @@ class DisabilityObserver(PublicHealthObserver):
     # Aggregators #
     ###############
 
-    def disability_weight_aggregator(self, dw: pd.DataFrame) -> Union[float, pd.Series]:
+    def disability_weight_aggregator(self, dw: pd.DataFrame) -> float | pd.Series:
         """Aggregate disability weights for the time step.
 
         Parameters

vivarium_public_health/results/disease.py

@@ -8,7 +8,7 @@ in the simulation.
 
 """
 
-from typing import Any, Dict, List
+from typing import Any
 
 import pandas as pd
 from layered_config_tree import LayeredConfigTree
@@ -66,7 +66,7 @@ class DiseaseObserver(PublicHealthObserver):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         """A dictionary containing the defaults for any configurations managed by
         this component.
         """
@@ -79,17 +79,17 @@ class DiseaseObserver(PublicHealthObserver):
         }
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         """Columns created by this observer."""
         return [self.previous_state_column_name]
 
     @property
-    def columns_required(self) -> List[str]:
+    def columns_required(self) -> list[str]:
         """Columns required by this observer."""
         return [self.disease]
 
     @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
+    def initialization_requirements(self) -> dict[str, list[str]]:
         """Requirements for observer initialization."""
         return {
             "requires_columns": [self.disease],

vivarium_public_health/results/mortality.py

@@ -8,7 +8,7 @@ excess mortality in the simulation, including "other causes".
 
 """
 
-from typing import Any, Dict, List
+from typing import Any
 
 import pandas as pd
 from layered_config_tree import LayeredConfigTree
@@ -81,7 +81,7 @@ class MortalityObserver(PublicHealthObserver):
         return [DiseaseState, RiskAttributableDisease]
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         """A dictionary containing the defaults for any configurations managed by
         this component.
         """
@@ -90,7 +90,7 @@ class MortalityObserver(PublicHealthObserver):
         return config_defaults
 
     @property
-    def columns_required(self) -> List[str]:
+    def columns_required(self) -> list[str]:
         """Columns required by this observer."""
         return [
             "alive",

vivarium_public_health/results/observer.py

@@ -8,7 +8,7 @@ public health models.
 
 """
 
-from typing import Callable, List, Optional, Union
+from collections.abc import Callable
 
 import pandas as pd
 from vivarium.framework.engine import Builder
@@ -32,12 +32,12 @@ class PublicHealthObserver(Observer):
         name: str,
         pop_filter: str,
         when: str = "collect_metrics",
-        requires_columns: List[str] = [],
-        requires_values: List[str] = [],
-        additional_stratifications: List[str] = [],
-        excluded_stratifications: List[str] = [],
-        aggregator_sources: Optional[List[str]] = None,
-        aggregator: Callable[[pd.DataFrame], Union[float, pd.Series]] = len,
+        requires_columns: list[str] = [],
+        requires_values: list[str] = [],
+        additional_stratifications: list[str] = [],
+        excluded_stratifications: list[str] = [],
+        aggregator_sources: list[str] | None = None,
+        aggregator: Callable[[pd.DataFrame], float | pd.Series] = len,
     ) -> None:
         """Registers an adding observation to the results system.
 

vivarium_public_health/results/risk.py

@@ -7,7 +7,7 @@ This module contains tools for observing risk exposure during the simulation.
 
 """
 
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import pandas as pd
 from layered_config_tree import LayeredConfigTree
@@ -59,7 +59,7 @@ class CategoricalRiskObserver(PublicHealthObserver):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         """A dictionary containing the defaults for any configurations managed by
         this component.
         """
@@ -72,7 +72,7 @@ class CategoricalRiskObserver(PublicHealthObserver):
         }
 
     @property
-    def columns_required(self) -> Optional[List[str]]:
+    def columns_required(self) -> list[str] | None:
         """The columns required by this observer."""
         return ["alive"]
 

vivarium_public_health/risks/base_risk.py

@@ -8,7 +8,7 @@ exposure.
 
 """
 
-from typing import Any, Dict, List
+from typing import Any
 
 import pandas as pd
 from vivarium import Component
@@ -106,7 +106,7 @@ class Risk(Component):
         return self.risk
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             self.name: {
                 "data_sources": {
@@ -122,14 +122,14 @@ class Risk(Component):
         }
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         columns_to_create = [self.propensity_column_name]
         if self.create_exposure_column:
             columns_to_create.append(self.exposure_column_name)
         return columns_to_create
 
     @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
+    def initialization_requirements(self) -> dict[str, list[str]]:
         return {
             "requires_columns": [],
             "requires_values": [],