clean up docstrings throughout (#476)

ihmeuw · Aug 28, 2024 · ce781dc · ce781dc
1 parent c83ac74
commit ce781dc
Show file tree

Hide file tree

Showing 18 changed files with 291 additions and 259 deletions.
diff --git a/src/vivarium_public_health/disease/state.py b/src/vivarium_public_health/disease/state.py
@@ -103,9 +103,8 @@ def transition_side_effect(self, index: pd.Index, event_time: pd.Timestamp) -> N
         ----------
         index
             An iterable of integer labels for the simulants.
-        event_time : pandas.Timestamp
+        event_time
             The time at which this transition occurs.
-
         """
         pop = self.population_view.get(index)
         pop[self.event_time_column] = event_time
@@ -141,13 +140,13 @@ def add_rate_transition(
             The end state after the transition.
 
         get_data_functions
-            map from transition type to the function to pull that transition's data
+            Map from transition type to the function to pull that transition's data.
         triggered
+            The trigger for the transition
 
 
         Returns
         -------
-        RateTransition
             The created transition object.
         """
         transition = RateTransition(self, output, get_data_functions, triggered)
@@ -157,7 +156,7 @@ def add_rate_transition(
     def add_proportion_transition(
         self,
         output: "BaseDiseaseState",
-        get_data_functions: Dict[str, Callable] = None,
+        get_data_functions: Optional[Dict[str, Callable]] = None,
         triggered=Trigger.NOT_TRIGGERED,
     ) -> ProportionTransition:
         """Builds a ProportionTransition from this state to the given state.
@@ -166,13 +165,13 @@ def add_proportion_transition(
         ----------
         output
             The end state after the transition.
-
         get_data_functions
-            map from transition type to the function to pull that transition's data
+            Map from transition type to the function to pull that transition's data.
+        triggered
+            The trigger for the transition.
 
         Returns
         -------
-        RateTransition
             The created transition object.
         """
         if "proportion" not in get_data_functions:
@@ -302,8 +301,8 @@ def __init__(
         allow_self_transition: bool = False,
         side_effect_function: Optional[Callable] = None,
         cause_type: str = "cause",
-        get_data_functions: Dict[str, Callable] = None,
-        cleanup_function: Callable = None,
+        get_data_functions: Optional[Dict[str, Callable]] = None,
+        cleanup_function: Optional[Callable] = None,
     ):
         """
         Parameters
@@ -312,14 +311,16 @@ def __init__(
             The name of this state.
         allow_self_transition
             Whether this state allows simulants to remain in the state for
-            multiple time-steps
+            multiple time-steps.
         side_effect_function
             A function to be called when this state is entered.
         cause_type
             The type of cause represented by this state. Either "cause" or "sequela".
         get_data_functions
             A dictionary containing a mapping to functions to retrieve data for
-            various state attributes
+            various state attributes.
+        cleanup_function
+            The cleanup function.
         """
         super().__init__(
             state_id,
@@ -344,7 +345,7 @@ def setup(self, builder: Builder) -> None:
 
         Parameters
         ----------
-        builder : `engine.Builder`
+        builder
             Interface to several simulation tools.
         """
         super().setup(builder)
@@ -506,9 +507,9 @@ def next_state(
         ----------
         index
             An iterable of integer labels for the simulants.
-        event_time:
+        event_time
             The time at which this transition occurs.
-        population_view:
+        population_view
             A view of the internal state of the simulation.
         """
         eligible_index = self._filter_for_transition_eligibility(index, event_time)
@@ -528,7 +529,6 @@ def compute_disability_weight(self, index: pd.Index) -> pd.Series:
 
         Returns
         -------
-        `pandas.Series`
             An iterable of disability weights indexed by the provided `index`.
         """
         disability_weight = pd.Series(0.0, index=index)
@@ -555,8 +555,12 @@ def adjust_mortality_rate(self, index: pd.Index, rates_df: pd.DataFrame) -> pd.D
         ----------
         index
             An iterable of integer labels for the simulants.
-        rates_df : `pandas.DataFrame`
+        rates_df
+            A DataFrame of mortality rates.
 
+        Returns
+        -------
+            The modified DataFrame of mortality rates.
         """
         rate = self.excess_mortality_rate(index, skip_post_processor=True)
         rates_df[self.state_id] = rate
@@ -599,17 +603,18 @@ def _assign_event_time_for_prevalent_cases(
         infected_at = current_time - pd.to_timedelta(infected_at, unit="D")
         return infected_at
 
-    def _filter_for_transition_eligibility(self, index, event_time):
+    def _filter_for_transition_eligibility(self, index, event_time) -> pd.Index:
         """Filter out all simulants who haven't been in the state for the prescribed dwell time.
 
         Parameters
         ----------
         index
             An iterable of integer labels for the simulants.
+        event_time
+            The time at which this transition occurs.
 
         Returns
         -------
-        pd.Index
             A filtered index of the simulants.
         """
         population = self.population_view.get(index, query='alive == "alive"')

diff --git a/src/vivarium_public_health/mslt/delay.py b/src/vivarium_public_health/mslt/delay.py
@@ -19,8 +19,7 @@
 
 
 class DelayedRisk(Component):
-    """
-    A delayed risk represents an exposure whose impact takes time to come into
+    """A delayed risk represents an exposure whose impact takes time to come into
     effect (e.g., smoking uptake and cessation).
 
     The data required by this component are:
@@ -320,8 +319,7 @@ def register_modifier(self, builder: Builder, disease: str) -> None:
     ########################
 
     def on_initialize_simulants(self, pop_data: SimulantData) -> None:
-        """
-        Define the initial distribution of the population across the bins, in
+        """Define the initial distribution of the population across the bins, in
         both the BAU and the intervention scenario.
         """
         # Set all bins to zero, in order to create the required columns.
@@ -364,7 +362,6 @@ def on_time_step_prepare(self, event: Event) -> None:
         - New exposures
         - Cessation of exposure
         - Increased duration of time since exposure
-
         """
         if self.clock().year == self.start_year:
             return
@@ -540,6 +537,17 @@ def pivot_load(builder: Builder, entity_key: str) -> pd.DataFrame:
     Performs a long to wide conversion if dataframe has an index column
     named 'measure'.
 
+    Parameters
+    ----------
+    builder
+        The builder object for the simulation.
+    entity_key
+        The key for the entity to be loaded.
+
+    Returns
+    -------
+    pd.DataFrame
+        The loaded data and potentially pivoted data.
     """
     data = builder.data.load(entity_key)
 

diff --git a/src/vivarium_public_health/mslt/disease.py b/src/vivarium_public_health/mslt/disease.py
@@ -19,7 +19,8 @@
 
 
 class AcuteDisease(Component):
-    """
+    """This component characterises an acute disease.
+
     An acute disease has a sufficiently short duration, relative to the
     time-step size, that it is not meaningful to talk about prevalence.
     Instead, it simply contributes an excess mortality rate, and/or a
@@ -32,10 +33,18 @@ class AcuteDisease(Component):
 
     where `<disease>` is the name as provided to the constructor.
 
-    Parameters
+    Attributes
     ----------
     disease
         The disease name (referred to as `<disease>` here).
+    excess_mortality
+        The excess mortality rate for the disease.
+    int_excess_mortality
+        The excess mortality rate for the disease in the intervention scenario.
+    disability_rate
+        The years lost due to disability (YLD) rate for the disease.
+    int_disability_rate
+        The YLD rate for the disease in the intervention scenario.
 
     """
 
@@ -77,16 +86,14 @@ def setup(self, builder: Builder) -> None:
     ##################################
 
     def mortality_adjustment(self, index, mortality_rate):
-        """
-        Adjust the all-cause mortality rate in the intervention scenario, to
+        """Adjust the all-cause mortality rate in the intervention scenario, to
         account for any change in prevalence (relative to the BAU scenario).
         """
         delta = self.int_excess_mortality(index) - self.excess_mortality(index)
         return mortality_rate + delta
 
     def disability_adjustment(self, index, yld_rate):
-        """
-        Adjust the years lost due to disability (YLD) rate in the intervention
+        """Adjust the years lost due to disability (YLD) rate in the intervention
         scenario, to account for any change in prevalence (relative to the BAU
         scenario).
         """
@@ -106,10 +113,28 @@ class Disease(Component):
 
     where `<disease>` is the name as provided to the constructor.
 
-    Parameters
+    Attributes
     ----------
     disease
         The disease name (referred to as `<disease>` here).
+    clock
+        The simulation clock.
+    start_year
+        The simulation start year.
+    simplified_equations
+        Whether to use simplified equations for the disease model.
+    incidence
+        The incidence rate for the disease.
+    incidence_intervention
+        The incidence rate for the disease in the intervention scenario.
+    remission
+        The remission rate for the disease.
+    excess_mortality
+        The excess mortality rate for the disease.
+    disability_rate
+        The years lost due to disability (YLD) rate for the disease.
+    initial_prevalence
+        The initial prevalence of the disease.
 
     """
 
@@ -231,9 +256,7 @@ def on_initialize_simulants(self, pop_data: SimulantData) -> None:
         self.population_view.update(pop)
 
     def on_time_step_prepare(self, event: Event) -> None:
-        """
-        Update the disease status for both the BAU and intervention scenarios.
-        """
+        """Update the disease status for both the BAU and intervention scenarios."""
         # Do not update the disease status in the first year, the initial data
         # describe the disease state at the end of the year.
         if self.clock().year == self.start_year:
@@ -356,8 +379,7 @@ def on_time_step_prepare(self, event: Event) -> None:
     ##################################
 
     def mortality_adjustment(self, index, mortality_rate):
-        """
-        Adjust the all-cause mortality rate in the intervention scenario, to
+        """Adjust the all-cause mortality rate in the intervention scenario, to
         account for any change in disease prevalence (relative to the BAU
         scenario).
         """
@@ -387,8 +409,7 @@ def mortality_adjustment(self, index, mortality_rate):
         return mortality_rate + delta
 
     def disability_adjustment(self, index, yld_rate):
-        """
-        Adjust the years lost due to disability (YLD) rate in the intervention
+        """Adjust the years lost due to disability (YLD) rate in the intervention
         scenario, to account for any change in disease prevalence (relative to
         the BAU scenario).
         """

diff --git a/src/vivarium_public_health/mslt/intervention.py b/src/vivarium_public_health/mslt/intervention.py
@@ -100,40 +100,43 @@ def adjust_rate(self, index, rates):
 
 
 class ModifyDiseaseIncidence(ModifyDiseaseRate):
-    """
-    Interventions that modify a disease incidence rate, based on a PIF lookup
+    """Interventions that modify a disease incidence rate, based on a PIF lookup
     table.
+
     """
 
     def __init__(self, intervention: str, disease: str):
         super().__init__(intervention=intervention, disease=disease, rate="incidence")
 
 
 class ModifyDiseaseMortality(ModifyDiseaseRate):
-    """
-    Interventions that modify a disease fatality rate, based on a PIF lookup
+    """Interventions that modify a disease fatality rate, based on a PIF lookup
     table.
+
     """
 
     def __init__(self, intervention: str, disease: str):
         super().__init__(intervention=intervention, disease=disease, rate="excess_mortality")
 
 
 class ModifyDiseaseMorbidity(ModifyDiseaseRate):
-    """
-    Interventions that modify a disease disability rate, based on a PIF lookup
+    """Interventions that modify a disease disability rate, based on a PIF lookup
     table.
+
     """
 
     def __init__(self, intervention: str, disease: str):
         super().__init__(intervention=intervention, disease=disease, rate="yld_rate")
 
 
 class ModifyAcuteDiseaseIncidence(Component):
-    """
-    Interventions that modify an acute disease incidence rate.
-    Note that this intervention will simply modify both the disability rate
+    """Interventions that modify an acute disease incidence rate.
+
+    Notes
+    -----
+    This intervention will simply modify both the disability rate
     and the mortality rate for the chosen acute disease.
+
     """
 
     ##############