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

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]

vivarium_public_health/results/observer.py

@@ -8,16 +8,19 @@ from vivarium_public_health.results.columns import COLUMNS
 
 
 class PublicHealthObserver(Observer):
-    """A convenience class for typical public health observers. It provides
-    an entry point for registering the most common observation type
-    as well as standardized results formatting methods to overwrite as necessary.
+    """A convenience class for typical public health observers.
+
+    It exposes a method for registering the most common observation type
+    (adding observation) as well methods for formatting public health results
+    in a standardized way (to be overwritten as necessary).
+
     """
 
     def register_adding_observation(
         self,
         builder: Builder,
-        name,
-        pop_filter,
+        name: str,
+        pop_filter: str,
         when: str = "collect_metrics",
         requires_columns: List[str] = [],
         requires_values: List[str] = [],
@@ -25,7 +28,42 @@ class PublicHealthObserver(Observer):
         excluded_stratifications: List[str] = [],
         aggregator_sources: Optional[List[str]] = None,
         aggregator: Callable[[pd.DataFrame], Union[float, pd.Series]] = len,
-    ):
+    ) -> None:
+        """Registers an adding observation to the results system.
+
+        An "adding" observation is one that adds/sums new results to existing
+        result values. It is the most common type of observation used in public
+        health models.
+
+        Parameters
+        ----------
+        builder
+            The builder object.
+        name
+            Name of the observation. It will also be the name of the output results
+            file for this particular observation.
+        pop_filter
+            A Pandas query filter string to filter the population down to the
+            simulants who should be considered for the observation.
+        when
+            Name of the lifecycle phase the observation should happen. Valid values are:
+            "time_step__prepare", "time_step", "time_step__cleanup", or "collect_metrics".
+        requires_columns
+            List of the state table columns that are required by either the `pop_filter`
+            or the `aggregator`.
+        requires_values
+            List of the value pipelines that are required by either the `pop_filter`
+            or the `aggregator`.
+        additional_stratifications
+            List of additional stratification names by which to stratify this
+            observation by.
+        excluded_stratifications
+            List of default stratification names to remove from this observation.
+        aggregator_sources
+            List of population view columns to be used in the `aggregator`.
+        aggregator
+            Function that computes the quantity for this observation.
+        """
         builder.results.register_adding_observation(
             name=name,
             pop_filter=pop_filter,
@@ -42,6 +80,27 @@ class PublicHealthObserver(Observer):
     def format_results(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
         """Top-level results formatter that calls standard sub-methods to be
         overwritten as necessary.
+
+        Public health observations typically require four columns in addition to
+        any stratifications and results columns: 'measure', 'entity_type', 'entity',
+        and 'sub_entity'. This method provides a standardized way to format
+        results by providing five sub-methods to be overwritten as necessary:
+        - format()
+        - get_measure_column()
+        - get_entity_type_column()
+        - get_entity_column()
+        - get_sub_entity_column()
+
+        Parameters
+        ----------
+        measure
+            The measure name.
+        results
+            The raw results.
+
+        Returns
+        -------
+            The formatted results.
         """
 
         results = self.format(measure, results)
@@ -63,16 +122,92 @@ class PublicHealthObserver(Observer):
         return results[ordered_columns]
 
     def format(self, measure: str, results: pd.DataFrame) -> pd.DataFrame:
+        """Format results.
+
+        This method should be overwritten in subclasses to provide custom formatting
+        for the results.
+
+        Parameters
+        ----------
+        measure
+            The measure name.
+        results
+            The raw results.
+
+        Returns
+        -------
+            The formatted results.
+        """
         return results
 
     def get_measure_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'measure' column.
+
+        This method should be overwritten in subclasses to provide the 'measure' column.
+
+        Parameters
+        ----------
+        measure
+            The measure name.
+        results
+            The raw results.
+
+        Returns
+        -------
+            The 'measure' column values.
+        """
         return pd.Series(measure, index=results.index)
 
     def get_entity_type_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'entity_type' column.
+
+        This method should be overwritten in subclasses to provide the 'entity_type' column.
+
+        Parameters
+        ----------
+        measure
+            The measure name.
+        results
+            The raw results.
+
+        Returns
+        -------
+            The 'entity_type' column values.
+        """
         return pd.Series(None, index=results.index)
 
     def get_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'entity' column.
+
+        This method should be overwritten in subclasses to provide the 'entity' column.
+
+        Parameters
+        ----------
+        measure
+            The measure name.
+        results
+            The raw results.
+
+        Returns
+        -------
+            The 'entity' column values.
+        """
         return pd.Series(None, index=results.index)
 
     def get_sub_entity_column(self, measure: str, results: pd.DataFrame) -> pd.Series:
+        """Get the 'sub_entity' column.
+
+        This method should be overwritten in subclasses to provide the 'sub_entity' column.
+
+        Parameters
+        ----------
+        measure
+            The measure name.
+        results
+            The raw results.
+
+        Returns
+        -------
+            The 'sub_entity' column values.
+        """
         return pd.Series(None, index=results.index)