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

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,16 +207,40 @@ 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.
         """
-
-        # Drop the unused 'value' column and rename the pipeline names to causes
-        results = results.drop(columns=["value"]).rename(
-            columns={col: col.replace(".disability_weight", "") for col in results.columns},
-        )
+        if len(self.causes_of_disability) > 1:
+            # Drop the unused 'value' column and rename the remaining pipeline names to cause names
+            results = results.rename(
+                columns={
+                    col: col.replace(".disability_weight", "") for col in results.columns
+                }
+            )[[col.state_id for col in self.causes_of_disability]]
+        else:
+            # Rename the 'value' column to the single cause of disability
+            results = results.rename(
+                columns={COLUMNS.VALUE: self.causes_of_disability[0].state_id}
+            )
         # Get desired index names prior to stacking
         idx_names = list(results.index.names) + [COLUMNS.SUB_ENTITY]
         results = pd.DataFrame(results.stack(), columns=[COLUMNS.VALUE])
@@ -173,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]

vivarium_public_health/results/disease.py

@@ -41,6 +41,24 @@ class DiseaseObserver(PublicHealthObserver):
                         - "sex"
                     include:
                         - "sample_stratification"
+
+    Attributes
+    ----------
+    disease
+        The name of the disease being observed.
+    previous_state_column_name
+        The name of the column that stores the previous state of the disease.
+    step_size
+        The time step size of the simulation.
+    disease_model
+        The disease model for the disease being observed.
+    entity_type
+        The type of entity being observed.
+    entity
+        The entity being observed.
+    transition_stratification_name
+        The stratification name for transitions between disease states.
+
     """
 
     ##############
@@ -49,6 +67,9 @@ class DiseaseObserver(PublicHealthObserver):
 
     @property
     def configuration_defaults(self) -> Dict[str, Any]:
+        """A dictionary containing the defaults for any configurations managed by
+        this component.
+        """
         return {
             "stratification": {
                 self.disease: super().configuration_defaults["stratification"][
@@ -59,14 +80,17 @@ class DiseaseObserver(PublicHealthObserver):
 
     @property
     def columns_created(self) -> List[str]:
+        """Columns created by this observer."""
         return [self.previous_state_column_name]
 
     @property
     def columns_required(self) -> List[str]:
+        """Columns required by this observer."""
         return [self.disease]
 
     @property
     def initialization_requirements(self) -> Dict[str, List[str]]:
+        """Requirements for observer initialization."""
         return {
             "requires_columns": [self.disease],
         }
@@ -76,6 +100,13 @@ class DiseaseObserver(PublicHealthObserver):
     #####################
 
     def __init__(self, disease: str) -> None:
+        """Constructor for this observer.
+
+        Parameters
+        ----------
+        disease
+            The name of the disease being observed.
+        """
         super().__init__()
         self.disease = disease
         self.previous_state_column_name = f"previous_{self.disease}"
@@ -85,6 +116,7 @@ class DiseaseObserver(PublicHealthObserver):
     #################
 
     def setup(self, builder: Builder) -> None:
+        """Set up the observer."""
         self.step_size = builder.time.step_size()
         self.disease_model = builder.components.get_component(f"disease_model.{self.disease}")
         self.entity_type = self.disease_model.cause_type
@@ -92,10 +124,35 @@ class DiseaseObserver(PublicHealthObserver):
         self.transition_stratification_name = f"transition_{self.disease}"
 
     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[self.disease]
 
     def register_observations(self, builder: Builder) -> None:
-
+        """Register stratifications and observations.
+
+        Notes
+        -----
+        Ideally, each observer registers a single observation. This one, however,
+        registeres two.
+
+        While it's typical for all stratification registrations to be encapsulated
+        in a single class (i.e. the
+        :class:ResultsStratifier <vivarium_public_health.results.stratification.ResultsStratifier),
+        this observer registers two additional stratifications. While they could
+        be registered in the ``ResultsStratifier`` as well, they are specific to
+        this observer and so they are registered here while we have easy access
+        to the required names and categories.
+        """
         self.register_disease_state_stratification(builder)
         self.register_transition_stratification(builder)
 
@@ -104,6 +161,7 @@ class DiseaseObserver(PublicHealthObserver):
         self.register_transition_count_observation(builder, pop_filter)
 
     def register_disease_state_stratification(self, builder: Builder) -> None:
+        """Register the disease state stratification."""
         builder.results.register_stratification(
             self.disease,
             [state.state_id for state in self.disease_model.states],
@@ -111,6 +169,20 @@ class DiseaseObserver(PublicHealthObserver):
         )
 
     def register_transition_stratification(self, builder: Builder) -> None:
+        """Register the transition stratification.
+
+        This stratification is used to track transitions between disease states.
+        It appends 'no_transition' to the list of transition categories and also
+        includes it as an exluded category.
+
+        Notes
+        -----
+        It is important to include 'no_transition' in bith the list of transition
+        categories as well as the list of excluded categories. This is because
+        it must exist as a category for the transition mapping to work correctly,
+        but then we don't want to include it later during the actual stratification
+        process.
+        """
         transitions = [
             str(transition) for transition in self.disease_model.transition_names
         ] + ["no_transition"]
@@ -130,6 +202,7 @@ class DiseaseObserver(PublicHealthObserver):
         )
 
     def register_person_time_observation(self, builder: Builder, pop_filter: str) -> None:
+        """Register a person time observation."""
         self.register_adding_observation(
             builder=builder,
             name=f"person_time_{self.disease}",
@@ -144,6 +217,7 @@ class DiseaseObserver(PublicHealthObserver):
     def register_transition_count_observation(
         self, builder: Builder, pop_filter: str
     ) -> None:
+        """Register a transition count observation."""
         self.register_adding_observation(
             builder=builder,
             name=f"transition_count_{self.disease}",
@@ -158,6 +232,17 @@ class DiseaseObserver(PublicHealthObserver):
         )
 
     def map_transitions(self, df: pd.DataFrame) -> pd.Series:
+        """Map previous and current disease states to transition string.
+
+        Parameters
+        ----------
+        df
+            The DataFrame containing the disease states.
+
+        Returns
+        -------
+            The transitions between disease states.
+        """
         transitions = pd.Series(index=df.index, dtype=str)
         transition_mask = df[self.previous_state_column_name] != df[self.disease]
         transitions[~transition_mask] = "no_transition"
@@ -179,7 +264,10 @@ class DiseaseObserver(PublicHealthObserver):
         self.population_view.update(pop)
 
     def on_time_step_prepare(self, event: Event) -> None:
-        # This enables tracking of transitions between states
+        """Update the previous state column to the current state.
+
+        This enables tracking of transitions between states.
+        """
         prior_state_pop = self.population_view.get(event.index)
         prior_state_pop[self.previous_state_column_name] = prior_state_pop[self.disease]
         self.population_view.update(prior_state_pop)
@@ -189,6 +277,17 @@ class DiseaseObserver(PublicHealthObserver):
     ###############
 
     def aggregate_state_person_time(self, x: pd.DataFrame) -> float:
+        """Aggregate person time for the time step.
+
+        Parameters
+        ----------
+        x
+            The DataFrame containing the population.
+
+        Returns
+        -------
+            The aggregated person time.
+        """
         return len(x) * to_years(self.step_size())
 
     ##############################
@@ -196,6 +295,26 @@ class DiseaseObserver(PublicHealthObserver):
     ##############################
 
     def format(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
+        """Rename the appropriate column to 'sub_entity'.
+
+        The primary thing this method does is rename the appropriate column
+        (either the transition stratification name of the disease name, depending
+        on the measure) to 'sub_entity'. We do this here instead of the
+        'get_sub_entity_column' method simply because we do not want the original
+        column at all. If we keep it here and then return it as the sub-entity
+        column later, the final results would have both.
+
+        Parameters
+        ----------
+        measure
+            The measure.
+        results
+            The results to format.
+
+        Returns
+        -------
+            The formatted results.
+        """
         results = results.reset_index()
         if "transition_count_" in measure:
             sub_entity = self.transition_stratification_name
@@ -205,6 +324,7 @@ class DiseaseObserver(PublicHealthObserver):
         return results
 
     def get_measure_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'measure' column values."""
         if "transition_count_" in measure:
             measure_name = "transition_count"
         if "person_time_" in measure:
@@ -212,11 +332,14 @@ class DiseaseObserver(PublicHealthObserver):
         return pd.Series(measure_name, index=results.index)
 
     def get_entity_type_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'entity_type' column values."""
         return pd.Series(self.entity_type, index=results.index)
 
     def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'entity' column values."""
         return pd.Series(self.entity, index=results.index)
 
     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]

vivarium_public_health/results/mortality.py

@@ -47,6 +47,18 @@ class MortalityObserver(PublicHealthObserver):
     This observer needs to access the has_excess_mortality attribute of the causes
     we're observing, but this attribute gets defined in the setup of the cause models.
     As a result, the model specification should list this observer after causes.
+
+    Attributes
+    ----------
+    required_death_columns
+        Columns required by the deaths observation.
+    required_yll_columns
+        Columns required by the ylls observation.
+    clock
+        The simulation clock.
+    causes_of_death
+        Causes of death to be observed.
+
     """
 
     def __init__(self) -> None:
@@ -65,12 +77,12 @@ class MortalityObserver(PublicHealthObserver):
 
     @property
     def mortality_classes(self) -> list[type]:
+        """The classes to be considered as causes of death."""
         return [DiseaseState, RiskAttributableDisease]
 
     @property
     def configuration_defaults(self) -> Dict[str, Any]:
-        """
-        A dictionary containing the defaults for any configurations managed by
+        """A dictionary containing the defaults for any configurations managed by
         this component.
         """
         config_defaults = super().configuration_defaults
@@ -79,6 +91,7 @@ class MortalityObserver(PublicHealthObserver):
 
     @property
     def columns_required(self) -> List[str]:
+        """Columns required by this observer."""
         return [
             "alive",
             "years_of_life_lost",
@@ -91,10 +104,22 @@ class MortalityObserver(PublicHealthObserver):
     #################
 
     def setup(self, builder: Builder) -> None:
+        """Set up the observer."""
         self.clock = builder.time.clock()
         self.set_causes_of_death(builder)
 
     def set_causes_of_death(self, builder: Builder) -> None:
+        """Set the causes of death to be observed.
+
+        The causes to be observed are any registered components of class types
+        found in the ``mortality_classes`` property.
+
+        Notes
+        -----
+        We do not actually exclude any categories in this method.
+
+        Also note that we add 'not_dead' and 'other_causes' categories here.
+        """
         causes_of_death = [
             cause
             for cause in builder.components.get_components_by_type(
@@ -112,9 +137,35 @@ class MortalityObserver(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[self.get_configuration_name()]
 
     def register_observations(self, builder: Builder) -> None:
+        """Register stratifications and observations.
+
+        Notes
+        -----
+        Ideally, each observer registers a single observation. This one, however,
+        registeres two.
+
+        While it's typical for all stratification registrations to be encapsulated
+        in a single class (i.e. the
+        :class:ResultsStratifier <vivarium_public_health.results.stratification.ResultsStratifier),
+        this observer potentially registers an additional one. While it could
+        be registered in the ``ResultsStratifier`` as well, it is specific to
+        this observer and so it is registered here while we have easy access
+        to the required categories.
+        """
         pop_filter = 'alive == "dead" and tracked == True'
         additional_stratifications = self.configuration.include
         if not self.configuration.aggregate:
@@ -155,10 +206,12 @@ class MortalityObserver(PublicHealthObserver):
     ###############
 
     def count_deaths(self, x: pd.DataFrame) -> float:
+        """Count the number of deaths that occurred during this time step."""
         died_of_cause = x["exit_time"] > self.clock()
         return sum(died_of_cause)
 
     def calculate_ylls(self, x: pd.DataFrame) -> float:
+        """Calculate the years of life lost during this time step."""
         died_of_cause = x["exit_time"] > self.clock()
         return x.loc[died_of_cause, "years_of_life_lost"].sum()
 
@@ -167,6 +220,26 @@ class MortalityObserver(PublicHealthObserver):
     ##############################
 
     def format(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
+        """Rename the appropriate column to 'entity'.
+
+        The primary thing this method does is rename the 'cause_of_death' column
+        to 'entity' (or, it we are aggregating, and there is no 'cause_of_death'
+        column, we simply create a new 'entity' column). We do this here instead
+        of the 'get_entity_column' method simply because we do not want the
+        'cause_of_death' at all. If we keep it here and then return it as the
+        entity column later, the final results would have both.
+
+        Parameters
+        ----------
+        measure
+            The measure.
+        results
+            The results to format.
+
+        Returns
+        -------
+            The formatted results.
+        """
         results = results.reset_index()
         if self.configuration.aggregate:
             results[COLUMNS.ENTITY] = "all_causes"
@@ -175,12 +248,15 @@ class MortalityObserver(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_death}
         return results[COLUMNS.ENTITY].map(entity_type_map).astype(CategoricalDtype())
 
     def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'entity' column values."""
         # The entity col was created in the 'format' method
         return results[COLUMNS.ENTITY]
 
     def get_sub_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'sub_entity' column values."""
         return results[COLUMNS.ENTITY]