vivarium-public-health 3.0.3__py3-none-any.whl → 3.0.4__py3-none-any.whl

vivarium_public_health/plugins/parser.py

@@ -7,6 +7,7 @@ Component Configuration Parsers in this module are specialized implementations o
 :class:`ComponentConfigurationParser <vivarium.framework.components.parser.ComponentConfigurationParser>`
 that can parse configurations of components specific to the Vivarium Public
 Health package.
+
 """
 
 from importlib import import_module
@@ -34,34 +35,32 @@ from vivarium_public_health.utilities import TargetString
 
 
 class CausesParsingErrors(ParsingError):
-    """
-    Error raised when there are any errors parsing a cause model configuration.
-    """
+    """Error raised when there are any errors parsing a cause model configuration."""
 
     def __init__(self, messages: List[str]):
         super().__init__("\n - " + "\n - ".join(messages))
 
 
 class CausesConfigurationParser(ComponentConfigurationParser):
-    """
+    """Parser for cause model configurations.
+
     Component configuration parser that acts the same as the standard vivarium
     `ComponentConfigurationParser` but adds the additional ability to parse a
     configuration to create `DiseaseModel` components. These DiseaseModel
     configurations can either be specified directly in the configuration in a
     `causes` key or in external configuration files that are specified in the
     `external_configuration` key.
+
     """
 
     DEFAULT_MODEL_CONFIG = {
         "model_type": f"{DiseaseModel.__module__}.{DiseaseModel.__name__}",
         "initial_state": None,
     }
-    """
-    If a cause model configuration does not specify a model type or initial 
-    state, these default values will be used. The default model type is 
-    `DiseaseModel` and the
-    default initial state is `None`. If the initial state is not specified,
-    the cause model must have a state named 'susceptible'.
+    """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'.
     """
 
     DEFAULT_STATE_CONFIG = {
@@ -72,23 +71,16 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         "cleanup_function": None,
         "state_type": None,
     }
-    """
-    If a state configuration does not specify cause_type, transient, 
-    allow_self_transition, side_effect, cleanup_function, or state_type,
-    these default values will be used. The default cause type is 'cause', the
-    default transient value is False, and the default allow_self_transition
-    value is True.
-    """
+    """Default state configuration if it's not explicitly specified."""
 
     DEFAULT_TRANSITION_CONFIG = {"triggered": "NOT_TRIGGERED"}
-    """
-    If a transition configuration does not specify a triggered value, this
-    default value will be used. The default triggered value is 'NOT_TRIGGERED'.
+    """Default triggered value.
+    
+    This value is used if the transition configuration does not explicity specify it.
     """
 
     def parse_component_config(self, component_config: LayeredConfigTree) -> List[Component]:
-        """
-        Parses the component configuration and returns a list of components.
+        """Parses the component configuration and returns a list of components.
 
         In particular, this method looks for an `external_configuration` key
         and/or a `causes` key.
@@ -143,7 +135,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        List
             A list of initialized components.
 
         Raises
@@ -188,18 +179,14 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     #########################
 
     def _add_default_config_layer(self, causes_config: LayeredConfigTree) -> None:
-        """
-        Adds a default layer to the provided configuration that specifies
-        default values for the cause model configuration.
+        """Adds a default layer to the provided configuration.
+
+        This default layer specifies values for the cause model configuration.
 
         Parameters
         ----------
         causes_config
             A LayeredConfigTree defining the cause model configurations
-
-        Returns
-        -------
-        None
         """
         default_config = {}
         for cause_name, cause_config in causes_config.items():
@@ -228,9 +215,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     def _get_cause_model_components(
         self, causes_config: LayeredConfigTree
     ) -> List[Component]:
-        """
-        Parses the cause model configuration and returns a list of
-        `DiseaseModel` components.
+        """Parses the cause model configuration and returns the `DiseaseModel` components.
 
         Parameters
         ----------
@@ -239,7 +224,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        List[Component]
             A list of initialized `DiseaseModel` components
         """
         cause_models = []
@@ -277,9 +261,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     def _get_state(
         self, state_name: str, state_config: LayeredConfigTree, cause_name: str
     ) -> BaseDiseaseState:
-        """
-        Parses a state configuration and returns an initialized `BaseDiseaseState`
-        object.
+        """Parses a state configuration and returns an initialized `BaseDiseaseState` object.
 
         Parameters
         ----------
@@ -292,7 +274,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        BaseDiseaseState
             An initialized `BaseDiseaseState` object
         """
         state_id = cause_name if state_name in ["susceptible", "recovered"] else state_name
@@ -330,8 +311,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         sink_state: BaseDiseaseState,
         transition_config: LayeredConfigTree,
     ) -> None:
-        """
-        Adds a transition between two states.
+        """Adds a transition between two states.
 
         Parameters
         ----------
@@ -341,10 +321,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
             The state the transition ends at
         transition_config
             A `LayeredConfigTree` defining the transition to add
-
-        Returns
-        -------
-        None
         """
         triggered = Trigger[transition_config.triggered]
         if "data_sources" in transition_config:
@@ -372,9 +348,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     def _get_data_sources(
         self, config: LayeredConfigTree
     ) -> Dict[str, Callable[[Builder, Any], Any]]:
-        """
-        Parses a data sources configuration and returns a dictionary of data
-        sources.
+        """Parses a data sources configuration and returns the data sources.
 
         Parameters
         ----------
@@ -383,7 +357,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        Dict[str, Callable[[Builder, Any], Any]]
             A dictionary of data source getters
         """
         return {name: self._get_data_source(name, config[name]) for name in config.keys()}
@@ -392,9 +365,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     def _get_data_source(
         name: str, source: Union[str, float]
     ) -> Callable[[Builder, Any], Any]:
-        """
-        Parses a data source and returns a callable that can be used to retrieve
-        the data.
+        """Parses a data source and returns a callable that can be used to retrieve the data.
 
         Parameters
         ----------
@@ -405,7 +376,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        Callable[[Builder, Any], Any]
             A callable that can be used to retrieve the data
         """
         if isinstance(source, float):
@@ -465,18 +435,13 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
     @staticmethod
     def _validate_external_configuration(external_configuration: LayeredConfigTree) -> None:
-        """
-        Validates the external configuration.
+        """Validates the external configuration.
 
         Parameters
         ----------
         external_configuration
             A LayeredConfigTree defining the external configuration
 
-        Returns
-        -------
-        None
-
         Raises
         ------
         CausesParsingErrors
@@ -504,18 +469,13 @@ class CausesConfigurationParser(ComponentConfigurationParser):
             raise CausesParsingErrors(error_messages)
 
     def _validate_causes_config(self, causes_config: LayeredConfigTree) -> None:
-        """
-        Validates the cause model configuration.
+        """Validates the cause model configuration.
 
         Parameters
         ----------
         causes_config
             A LayeredConfigTree defining the cause model configurations
 
-        Returns
-        -------
-        None
-
         Raises
         ------
         CausesParsingErrors
@@ -530,8 +490,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
             raise CausesParsingErrors(error_messages)
 
     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.
+        """Validates a cause configuration and returns a list of error messages.
 
         Parameters
         ----------
@@ -542,7 +501,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        List[str]
             A list of error messages
         """
         error_messages = []
@@ -607,8 +565,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     def _validate_state(
         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.
+        """Validates a state configuration and returns a list of error messages.
 
         Parameters
         ----------
@@ -621,7 +578,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        List[str]
             A list of error messages
         """
         error_messages = []
@@ -702,8 +658,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
         transition_config: Dict[str, Any],
         states_config: Dict[str, Any],
     ) -> List[str]:
-        """
-        Validates a transition configuration and returns a list of error messages.
+        """Validates a transition configuration and returns a list of error messages.
 
         Parameters
         ----------
@@ -718,7 +673,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        List[str]
             A list of error messages
         """
         error_messages = []
@@ -803,8 +757,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     def _validate_imported_type(
         import_path: str, cause_name: str, entity_type: str, entity_name: Optional[str] = None
     ) -> List[str]:
-        """
-        Validates an imported type and returns a list of error messages.
+        """Validates an imported type and returns a list of error messages.
 
         Parameters
         ----------
@@ -820,7 +773,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        List[str]
             A list of error messages
         """
         expected_type = {"model": DiseaseModel, "state": BaseDiseaseState}[entity_type]
@@ -847,9 +799,7 @@ class CausesConfigurationParser(ComponentConfigurationParser):
     def _validate_data_sources(
         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 a list of
-        error messages.
+        """Validates the data sources in a configuration and returns any error messages.
 
         Parameters
         ----------
@@ -864,7 +814,6 @@ class CausesConfigurationParser(ComponentConfigurationParser):
 
         Returns
         -------
-        List[str]
             A list of error messages
         """
         error_messages = []

vivarium_public_health/population/add_new_birth_cohorts.py

@@ -81,8 +81,7 @@ class FertilityCrudeBirthRate(Component):
 
     new_births = sim_pop_size_t0 * live_births / true_pop_size * step_size
 
-    Where
-
+    Where:
     sim_pop_size_t0 = the initial simulation population size
     live_births = annual number of live births in the true population
     true_pop_size = the true population size
@@ -126,6 +125,7 @@ class FertilityCrudeBirthRate(Component):
     def on_time_step(self, event: Event) -> None:
         """Adds new simulants every time step based on the Crude Birth Rate
         and an assumption that birth is a Poisson process
+
         Parameters
         ----------
         event
@@ -151,9 +151,7 @@ class FertilityCrudeBirthRate(Component):
 
 
 class FertilityAgeSpecificRates(Component):
-    """
-    A simulant-specific model for fertility and pregnancies.
-    """
+    """A simulant-specific model for fertility and pregnancies."""
 
     ##############
     # Properties #
@@ -180,11 +178,11 @@ class FertilityAgeSpecificRates(Component):
     #####################
 
     def setup(self, builder: Builder) -> None:
-        """Setup the common randomness stream and
-        age-specific fertility lookup tables.
+        """Setup the common randomness stream and age-specific fertility lookup tables.
+
         Parameters
         ----------
-        builder : vivarium.engine.Builder
+        builder
             Framework coordination object.
         """
         age_specific_fertility_rate = self.load_age_specific_fertility_rate_data(builder)
@@ -238,9 +236,10 @@ class FertilityAgeSpecificRates(Component):
 
     def on_time_step(self, event: Event) -> None:
         """Produces new children and updates parent status on time steps.
+
         Parameters
         ----------
-        event : vivarium.population.PopulationEvent
+        event
             The event that triggered the function call.
         """
         # Get a view on all living women who haven't had a child in at least nine months.

vivarium_public_health/population/base_population.py

@@ -126,7 +126,6 @@ class BasePopulation(Component):
         respectively.  Here we are agnostic to the methods of entrance and exit (e.g., birth,
         migration, death, etc.) as these characteristics can be inferred from this column and
         other information about the simulant and the simulation parameters.
-
         """
 
         age_params = {
@@ -282,7 +281,6 @@ def generate_population(
 
     Returns
     -------
-    pandas.DataFrame
         Table with columns
             'entrance_time'
                 The `pandas.Timestamp` describing when the simulant entered
@@ -299,7 +297,6 @@ def generate_population(
                 The location indicating where the simulant resides.
             'sex'
                 Either 'Male' or 'Female'.  The sex of the simulant.
-
     """
     simulants = pd.DataFrame(
         {
@@ -361,7 +358,6 @@ def _assign_demography_with_initial_age(
 
     Returns
     -------
-    pandas.DataFrame
         Table with same columns as `simulants` and with the additional
         columns 'age', 'sex',  and 'location'.
     """
@@ -426,7 +422,6 @@ def _assign_demography_with_age_bounds(
 
     Returns
     -------
-    pandas.DataFrame
         Table with same columns as `simulants` and with the additional columns
         'age', 'sex',  and 'location'.
 

vivarium_public_health/population/data_transformations.py

@@ -33,7 +33,6 @@ def assign_demographic_proportions(
 
     Returns
     -------
-    pandas.DataFrame
         Table with columns
             'age' : Midpoint of the age group,
             'age_start' : Lower bound of the age group,
@@ -101,7 +100,6 @@ def rescale_binned_proportions(
 
     Returns
     -------
-    pandas.DataFrame
         Table with the same columns as `pop_data` where all bins outside the range
         (age_start, age_end) have been discarded.  If age_start and age_end
         don't fall cleanly on age boundaries, the bins in which they lie are clipped and
@@ -174,7 +172,6 @@ def rescale_binned_proportions(
 def _add_edge_age_groups(pop_data: pd.DataFrame) -> pd.DataFrame:
     """Pads the population data with age groups that enforce constant
     left interpolation and interpolation to zero on the right.
-
     """
     index_cols = ["location", "year_start", "year_end", "sex"]
     age_cols = ["age", "age_start", "age_end"]
@@ -250,7 +247,6 @@ def smooth_ages(
 
     Returns
     -------
-    pandas.DataFrame
         Table with same columns as `simulants` with ages smoothed out within the age bins.
     """
     simulants = simulants.copy()
@@ -324,7 +320,7 @@ def _get_bins_and_proportions(
 
     Returns
     -------
-    Tuple[EndpointValues, AgeValues]
+        A tuple of endpoints tuples and ages tuples.
         The `EndpointValues` tuple has values (
             age at left edge of bin,
             age at right edge of bin,
@@ -334,7 +330,6 @@ def _get_bins_and_proportions(
             proportion of pop in previous bin,
             proportion of pop in next bin,
         )
-
     """
     left = float(pop_data.loc[pop_data["age"] == age.current, "age_start"].iloc[0])
     right = float(pop_data.loc[pop_data["age"] == age.current, "age_end"].iloc[0])
@@ -405,7 +400,6 @@ def _construct_sampling_parameters(
 
     Returns
     -------
-    Tuple[EndpointValues, EndpointValues, float, float]
         A tuple of (pdf, slope, area, cdf_inflection_point) where
             pdf is a tuple with values (
                 pdf evaluated at left bin edge,
@@ -474,7 +468,6 @@ def _compute_ages(
 
     Returns
     -------
-    Union[np.ndarray, float]
         Smoothed ages from one half of the age bin distribution.
     """
     if abs(slope) < np.finfo(np.float32).eps:

vivarium_public_health/population/mortality.py

@@ -59,9 +59,9 @@ from vivarium_public_health.utilities import get_lookup_columns
 
 
 class Mortality(Component):
-    """
-    This is the mortality component which models sources of mortality for a model.
-    THe component models all cause mortality and allows for disease models to contribute
+    """This is the mortality component which models of mortality in a population.
+
+    The component models all cause mortality and allows for disease models to contribute
     cause specific mortality. Data used by this class should be supplied in the artifact
     and is configurable in the configuration to build lookup tables. For instance, let's
     say we want to use sex and hair color to build a lookup table for all cause mortality.

vivarium_public_health/results/columns.py

@@ -4,7 +4,7 @@ from vivarium.framework.results import VALUE_COLUMN
 
 
 class __Columns(NamedTuple):
-    """column names"""
+    """Container class for column names used in results dataframes."""
 
     VALUE: str = VALUE_COLUMN
     MEASURE: str = "measure"

vivarium_public_health/results/disability.py

@@ -8,7 +8,7 @@ in the simulation.
 
 """
 
-from typing import Any, List, Union
+from typing import Union
 
 import pandas as pd
 from layered_config_tree import LayeredConfigTree
@@ -42,6 +42,17 @@ class DisabilityObserver(PublicHealthObserver):
                         - "sex"
                     include:
                         - "sample_stratification"
+    Attributes
+    ----------
+    disability_weight_pipeline_name
+        The name of the pipeline that produces disability weights.
+    step_size
+        The time step size of the simulation.
+    disability_weight
+        The pipeline that produces disability weights.
+    causes_of_disability
+        The causes of disability to be observed.
+
     """
 
     ##############
@@ -50,7 +61,7 @@ class DisabilityObserver(PublicHealthObserver):
 
     @property
     def disability_classes(self) -> list[type]:
-        """The classes to be considered for causes of disability."""
+        """The classes to be considered as causes of disability."""
         return [DiseaseState, RiskAttributableDisease]
 
     #####################
@@ -66,16 +77,25 @@ class DisabilityObserver(PublicHealthObserver):
     #################
 
     def setup(self, builder: Builder) -> None:
+        """Set up the observer."""
         self.step_size = pd.Timedelta(days=builder.configuration.time.step_size)
         self.disability_weight = self.get_disability_weight_pipeline(builder)
         self.set_causes_of_disability(builder)
 
     def set_causes_of_disability(self, builder: Builder) -> None:
-        """Set the causes of disability to be observed by removing any excluded
-        via the model spec from the list of all disability class causes. We implement
-        exclusions here because disabilities are unique in that they are not
-        registered stratifications and so cannot be excluded during the stratification
-        call like other categories.
+        """Set the causes of disability to be observed.
+
+        The causes to be observed are any registered components of class types
+        found in the ``disability_classes`` property *excluding* any listed in
+        the model spec as ``excluded_categories``.
+
+        Notes
+        -----
+        We implement exclusions here instead of during the stratification call
+        like most other categories because disabilities are unique in that they are
+        *not* actually registered stratifications.
+
+        Also note that we add an 'all_causes' category here.
         """
         causes_of_disability = builder.components.get_components_by_type(
             self.disability_classes
@@ -111,9 +131,21 @@ class DisabilityObserver(PublicHealthObserver):
         ]
 
     def get_configuration(self, builder: Builder) -> LayeredConfigTree:
+        """Get the stratification configuration for this observer.
+
+        Parameters
+        ----------
+        builder
+            The builder object for the simulation.
+
+        Returns
+        -------
+            The stratification configuration for this observer.
+        """
         return builder.configuration.stratification.disability
 
     def register_observations(self, builder: Builder) -> None:
+        """Register an observation for years lived with disability."""
         cause_pipelines = [
             f"{cause.state_id}.disability_weight" for cause in self.causes_of_disability
         ]
@@ -131,6 +163,17 @@ class DisabilityObserver(PublicHealthObserver):
         )
 
     def get_disability_weight_pipeline(self, builder: Builder) -> Pipeline:
+        """Register (and return) the pipeline that produces disability weights.
+
+        Parameters
+        ----------
+        builder
+            The builder object for the simulation.
+
+        Returns
+        -------
+            The pipeline that produces disability weights.
+        """
         return builder.value.register_value_producer(
             self.disability_weight_pipeline_name,
             source=lambda index: [pd.Series(0.0, index=index)],
@@ -143,6 +186,17 @@ class DisabilityObserver(PublicHealthObserver):
     ###############
 
     def disability_weight_aggregator(self, dw: pd.DataFrame) -> Union[float, pd.Series]:
+        """Aggregate disability weights for the time step.
+
+        Parameters
+        ----------
+        dw
+            The disability weights to aggregate.
+
+        Returns
+        -------
+            The aggregated disability weights.
+        """
         aggregated_dw = (dw * to_years(self.step_size)).sum().squeeze()
         if isinstance(aggregated_dw, pd.Series):
             aggregated_dw.index.name = "cause_of_disability"
@@ -153,10 +207,27 @@ class DisabilityObserver(PublicHealthObserver):
     ##############################
 
     def format(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
-        """Format results. Note that ylds are unique in that we
-        can't stratify by cause of disability (because there can be multiple at
-        once), and so the results here are actually wide by disability weight
-        pipeline name.
+        """Format wide YLD results to match typical/long stratified results.
+
+        YLDs are unique in that we can't stratify by cause of disability (because
+        there can be multiple at once), and so the results here are actually wide
+        by disability weight pipeline name. This method formats the results to be
+        long by cause of disability.
+
+        Parameters
+        ----------
+        measure
+            The measure.
+        results
+            The wide results to format.
+
+        Returns
+        -------
+            The results stacked by causes of disability.
+
+        Notes
+        -----
+        This method also adds the 'sub_entity' column to the results.
         """
         if len(self.causes_of_disability) > 1:
             # Drop the unused 'value' column and rename the remaining pipeline names to cause names
@@ -180,15 +251,18 @@ class DisabilityObserver(PublicHealthObserver):
         return results
 
     def get_entity_type_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'entity_type' column values."""
         entity_type_map = {
             cause.state_id: cause.cause_type for cause in self.causes_of_disability
         }
         return results[COLUMNS.SUB_ENTITY].map(entity_type_map).astype(CategoricalDtype())
 
     def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'entity' column values."""
         entity_map = {cause.state_id: cause.model for cause in self.causes_of_disability}
         return results[COLUMNS.SUB_ENTITY].map(entity_map).astype(CategoricalDtype())
 
     def get_sub_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'sub_entity' column values."""
         # The sub-entity col was created in the 'format' method
         return results[COLUMNS.SUB_ENTITY]