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

vivarium_public_health/_version.py

@@ -1 +1 @@
-__version__ = "3.0.2"
+__version__ = "3.0.4"

vivarium_public_health/disease/state.py

@@ -103,9 +103,8 @@ class BaseDiseaseState(State):
         ----------
         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 @@ class BaseDiseaseState(State):
             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 @@ class BaseDiseaseState(State):
     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 @@ class BaseDiseaseState(State):
         ----------
         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 @@ class DiseaseState(BaseDiseaseState):
         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 @@ class DiseaseState(BaseDiseaseState):
             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 @@ class DiseaseState(BaseDiseaseState):
 
         Parameters
         ----------
-        builder : `engine.Builder`
+        builder
             Interface to several simulation tools.
         """
         super().setup(builder)
@@ -506,9 +507,9 @@ class DiseaseState(BaseDiseaseState):
         ----------
         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 @@ class DiseaseState(BaseDiseaseState):
 
         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 @@ class DiseaseState(BaseDiseaseState):
         ----------
         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 @@ class DiseaseState(BaseDiseaseState):
         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"')

vivarium_public_health/mslt/delay.py

@@ -19,8 +19,7 @@ from vivarium.framework.population import SimulantData
 
 
 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 @@ class DelayedRisk(Component):
     ########################
 
     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 @@ class DelayedRisk(Component):
         - 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)
 

vivarium_public_health/mslt/disease.py

@@ -19,7 +19,8 @@ from vivarium.framework.population import SimulantData
 
 
 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 @@ class AcuteDisease(Component):
     ##################################
 
     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 @@ class Disease(Component):
         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 @@ class Disease(Component):
     ##################################
 
     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 @@ class Disease(Component):
         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).
         """

vivarium_public_health/mslt/intervention.py

@@ -100,9 +100,9 @@ class ModifyDiseaseRate(Component):
 
 
 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):
@@ -110,9 +110,9 @@ class ModifyDiseaseIncidence(ModifyDiseaseRate):
 
 
 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):
@@ -120,9 +120,9 @@ class ModifyDiseaseMortality(ModifyDiseaseRate):
 
 
 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):
@@ -130,10 +130,13 @@ class ModifyDiseaseMorbidity(ModifyDiseaseRate):
 
 
 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.
+
     """
 
     ##############

vivarium_public_health/mslt/observer.py

@@ -16,9 +16,8 @@ from vivarium.framework.engine import Builder
 from vivarium.framework.event import Event
 
 
-def output_file(config, suffix, sep="_", ext="csv"):
-    """
-    Determine the output file name for an observer, based on the prefix
+def output_file(config, suffix, sep="_", ext="csv") -> str:
+    """Determine the output file name for an observer, based on the prefix
     defined in ``config.observer.output_prefix`` and the (optional)
     ``config.input_data.input_draw_number``.
 
@@ -33,6 +32,9 @@ def output_file(config, suffix, sep="_", ext="csv"):
     ext
         The output file extension.
 
+    Returns
+    -------
+        The output file name for the observer.
     """
     if "observer" not in config:
         raise ValueError("observer.output_prefix not defined")
@@ -51,16 +53,22 @@ def output_file(config, suffix, sep="_", ext="csv"):
 
 
 class MorbidityMortality(Component):
-    """
-    This class records the all-cause morbidity and mortality rates for each
+    """This class records the all-cause morbidity and mortality rates for each
     cohort at each year of the simulation.
 
-    Parameters
+    Attributes
     ----------
     output_suffix
         The suffix for the CSV file in which to record the
         morbidity and mortality data.
-
+    clock
+        The simulation clock.
+    tables
+        The tables of morbidity and mortality data.
+    table_cols
+        The columns in the tables.
+    output_file
+        The output file for the morbidity and mortality data.
     """
 
     ##############
@@ -148,7 +156,7 @@ class MorbidityMortality(Component):
     # Helper methods #
     ##################
 
-    def calculate_LE(self, table, py_col, denom_col):
+    def calculate_LE(self, table, py_col, denom_col) -> pd.Series:
         """Calculate the life expectancy for each cohort at each time-step.
 
         Parameters
@@ -162,9 +170,7 @@ class MorbidityMortality(Component):
 
         Returns
         -------
-            The life expectancy for each table row, represented as a
-            pandas.Series object.
-
+            The life expectancy for each table row.
         """
         # Group the person-years by cohort.
         group_cols = ["year_of_birth", "sex"]
@@ -180,17 +186,36 @@ class MorbidityMortality(Component):
 
 
 class Disease(Component):
-    """
-    This class records the disease incidence rate and disease prevalence for
+    """This class records the disease incidence rate and disease prevalence for
     each cohort at each year of the simulation.
 
-    Parameters
+    Attributes
     ----------
     disease
         The name of the chronic disease.
     output_suffix
         The suffix for the CSV file in which to record the
         disease data.
+    bau_S_col
+        The name of the BAU susceptible column.
+    bau_C_col
+        The name of the BAU chronic column.
+    int_S_col
+        The name of the intervention susceptible column.
+    int_C_col
+        The name of the intervention chronic column.
+    bau_incidence
+        The incidence rate for the BAU scenario.
+    int_incidence
+        The incidence rate for the intervention scenario.
+    tables
+        The tables of disease data.
+    table_cols
+        The columns in the tables.
+    clock
+        The simulation clock.
+    output_file
+        The output file for the disease data.
 
     """
 
@@ -286,12 +311,23 @@ class Disease(Component):
 class TobaccoPrevalence(Component):
     """This class records the prevalence of tobacco use in the population.
 
-    Parameters
+    Attributes
     ----------
     output_suffix
         The suffix for the CSV file in which to record the
         prevalence data.
-
+    config
+        The builder configuration object.
+    clock
+        The simulation clock.
+    bin_years
+        The number of years post-exposure to consider.
+    tables
+        The tables of tobacco prevalence data.
+    table_cols
+        The columns in the tables.
+    output_file
+        The output file for the tobacco prevalence data.
     """
 
     ##############
@@ -339,7 +375,7 @@ class TobaccoPrevalence(Component):
     # Setup methods #
     #################
 
-    def get_bin_names(self):
+    def get_bin_names(self) -> list[str]:
         """Return the bin names for both the BAU and the intervention scenario.
 
         These names take the following forms:
@@ -356,6 +392,9 @@ class TobaccoPrevalence(Component):
 
         The intervention bin names take the form ``"name_intervention.X"``.
 
+        Returns
+        -------
+            The bin names for tobacco use.
         """
         if self.bin_years == 0:
             delay_bins = [str(0)]

vivarium_public_health/mslt/population.py

@@ -19,8 +19,7 @@ from vivarium.framework.population import SimulantData
 
 
 class BasePopulation(Component):
-    """
-    This component implements the core population demographics: age, sex,
+    """This component implements the core population demographics: age, sex,
     population size.
 
     The configuration options for this component are:
@@ -113,9 +112,9 @@ class BasePopulation(Component):
 
 
 class Mortality(Component):
-    """
-    This component reduces the population size of each cohort over time,
+    """This component reduces the population size of each cohort over time,
     according to the all-cause mortality rate.
+
     """
 
     ##############
@@ -156,8 +155,7 @@ class Mortality(Component):
     ########################
 
     def on_time_step(self, event: Event) -> None:
-        """
-        Calculate the number of deaths and survivors at each time-step, for
+        """Calculate the number of deaths and survivors at each time-step, for
         both the BAU and intervention scenarios.
         """
         pop = self.population_view.get(event.index)
@@ -181,10 +179,10 @@ class Mortality(Component):
 
 
 class Disability(Component):
-    """
-    This component calculates the health-adjusted life years (HALYs) for each
+    """This component calculates the health-adjusted life years (HALYs) for each
     cohort over time, according to the years lost due to disability (YLD)
     rate.
+
     """
 
     ##############
@@ -219,8 +217,7 @@ class Disability(Component):
     ########################
 
     def on_time_step(self, event: Event) -> None:
-        """
-        Calculate the HALYs for each cohort at each time-step, for both the
+        """Calculate the HALYs for each cohort at each time-step, for both the
         BAU and intervention scenarios.
         """
         pop = self.population_view.get(event.index)