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

vivarium_public_health/plugins/parser.py

@@ -10,8 +10,9 @@ 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))
 
 
@@ -81,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
@@ -216,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
@@ -236,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()
             }
@@ -351,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
@@ -366,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
@@ -500,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
@@ -591,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
@@ -683,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
@@ -783,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
@@ -825,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,7 +100,7 @@ class Mortality(Component):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         return {
             "mortality": {
                 "data_sources": {
@@ -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": [],

vivarium_public_health/risks/distributions.py

@@ -9,7 +9,7 @@ exposure distributions.
 """
 
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, List, Optional, Union
+from collections.abc import Callable
 
 import numpy as np
 import pandas as pd
@@ -38,7 +38,7 @@ class RiskExposureDistribution(Component, ABC):
         self,
         risk: EntityString,
         distribution_type: str,
-        exposure_data: Optional[Union[int, float, pd.DataFrame]] = None,
+        exposure_data: int | float | pd.DataFrame | None = None,
     ) -> None:
         super().__init__()
         self.risk = risk
@@ -51,14 +51,14 @@ class RiskExposureDistribution(Component, ABC):
     # Setup methods #
     #################
 
-    def get_configuration(self, builder: "Builder") -> Optional[LayeredConfigTree]:
+    def get_configuration(self, builder: "Builder") -> LayeredConfigTree | None:
         return builder.configuration[self.risk]
 
     @abstractmethod
     def build_all_lookup_tables(self, builder: "Builder") -> None:
         raise NotImplementedError
 
-    def get_exposure_data(self, builder: Builder) -> Union[int, float, pd.DataFrame]:
+    def get_exposure_data(self, builder: Builder) -> int | float | pd.DataFrame:
         if self._exposure_data is not None:
             return self._exposure_data
         return self.get_data(builder, self.configuration["data_sources"]["exposure"])
@@ -92,11 +92,11 @@ class EnsembleDistribution(RiskExposureDistribution):
     ##############
 
     @property
-    def columns_created(self) -> List[str]:
+    def columns_created(self) -> list[str]:
         return [self._propensity]
 
     @property
-    def initialization_requirements(self) -> Dict[str, List[str]]:
+    def initialization_requirements(self) -> dict[str, list[str]]:
         return {
             "requires_columns": [],
             "requires_values": [],
@@ -265,7 +265,7 @@ class ContinuousDistribution(RiskExposureDistribution):
 
 class PolytomousDistribution(RiskExposureDistribution):
     @property
-    def categories(self) -> List[str]:
+    def categories(self) -> list[str]:
         # These need to be sorted so the cumulative sum is in the ocrrect order of categories
         # and results are therefore reproducible and correct
         return sorted(self.lookup_tables["exposure"].value_columns)
@@ -286,8 +286,8 @@ class PolytomousDistribution(RiskExposureDistribution):
         )
 
     def get_exposure_value_columns(
-        self, exposure_data: Union[int, float, pd.DataFrame]
-    ) -> Optional[List[str]]:
+        self, exposure_data: int | float | pd.DataFrame
+    ) -> list[str] | None:
         if isinstance(exposure_data, pd.DataFrame):
             return list(exposure_data["parameter"].unique())
         return None
@@ -342,7 +342,7 @@ class DichotomousDistribution(RiskExposureDistribution):
         )
         self.lookup_tables["paf"] = self.build_lookup_table(builder, 0.0)
 
-    def get_exposure_data(self, builder: Builder) -> Union[int, float, pd.DataFrame]:
+    def get_exposure_data(self, builder: Builder) -> int | float | pd.DataFrame:
         exposure_data = super().get_exposure_data(builder)
 
         if isinstance(exposure_data, (int, float)):
@@ -373,8 +373,8 @@ class DichotomousDistribution(RiskExposureDistribution):
         return exposure_data
 
     def get_exposure_value_columns(
-        self, exposure_data: Union[int, float, pd.DataFrame]
-    ) -> Optional[List[str]]:
+        self, exposure_data: int | float | pd.DataFrame
+    ) -> list[str] | None:
         if isinstance(exposure_data, pd.DataFrame):
             return self.get_value_columns(exposure_data)
         return None
@@ -470,9 +470,9 @@ def clip(q):
 
 
 def get_risk_distribution_parameter(
-    value_columns_getter: Callable[[Union[pd.DataFrame]], List[str]],
-    data: Union[float, pd.DataFrame],
-) -> Union[float, pd.Series]:
+    value_columns_getter: Callable[[pd.DataFrame], list[str]],
+    data: float | pd.DataFrame,
+) -> float | pd.Series:
     if isinstance(data, pd.DataFrame):
         value_columns = value_columns_getter(data)
         if len(value_columns) > 1:

vivarium_public_health/risks/effect.py

@@ -7,9 +7,9 @@ This module contains tools for modeling the relationship between risk
 exposure models and disease models.
 
 """
-
+from collections.abc import Callable
 from importlib import import_module
-from typing import Any, Callable, Dict, List, Tuple, Union
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -58,7 +58,7 @@ class RiskEffect(Component):
         return f"risk_effect.{risk.name}_on_{target}"
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         """Default values for any configurations managed by this component."""
         return {
             self.name: {
@@ -150,7 +150,7 @@ class RiskEffect(Component):
         self,
         builder: Builder,
         configuration=None,
-    ) -> Union[str, float, pd.DataFrame]:
+    ) -> str | float | pd.DataFrame:
         if configuration is None:
             configuration = self.configuration
 
@@ -173,8 +173,8 @@ class RiskEffect(Component):
         return rr_data
 
     def get_filtered_data(
-        self, builder: "Builder", data_source: Union[str, float, pd.DataFrame]
-    ) -> Union[float, pd.DataFrame]:
+        self, builder: "Builder", data_source: str | float | pd.DataFrame
+    ) -> float | pd.DataFrame:
         data = super().get_data(builder, data_source)
 
         if isinstance(data, pd.DataFrame):
@@ -191,8 +191,8 @@ class RiskEffect(Component):
         return data
 
     def process_categorical_data(
-        self, builder: Builder, rr_data: Union[str, float, pd.DataFrame]
-    ) -> Tuple[Union[str, float, pd.DataFrame], List[str]]:
+        self, builder: Builder, rr_data: str | float | pd.DataFrame
+    ) -> tuple[str | float | pd.DataFrame, list[str]]:
         if not isinstance(rr_data, pd.DataFrame):
             cat1 = builder.data.load("population.demographic_dimensions")
             cat1["parameter"] = "cat1"
@@ -347,7 +347,7 @@ class NonLogLinearRiskEffect(RiskEffect):
     ##############
 
     @property
-    def configuration_defaults(self) -> Dict[str, Any]:
+    def configuration_defaults(self) -> dict[str, Any]:
         """Default values for any configurations managed by this component."""
         return {
             self.name: {
@@ -419,7 +419,7 @@ class NonLogLinearRiskEffect(RiskEffect):
         self,
         builder: Builder,
         configuration=None,
-    ) -> Union[str, float, pd.DataFrame]:
+    ) -> str | float | pd.DataFrame:
         if configuration is None:
             configuration = self.configuration