vivarium-public-health 3.1.3__tar.gz → 3.1.4__tar.gz

{vivarium_public_health-3.1.3 → vivarium_public_health-3.1.4}/.github/workflows/build.yml

@@ -20,6 +20,7 @@ jobs:
     with:
       dependencies: "layered_config_tree,vivarium"
       skip_mypy: true
+      python_version: ${{ matrix.python-version }}
     secrets:
       NOTIFY_EMAIL: ${{ secrets.NOTIFY_EMAIL }}
       NOTIFY_PASSWORD: ${{ secrets.NOTIFY_PASSWORD }}

{vivarium_public_health-3.1.3 → vivarium_public_health-3.1.4}/CHANGELOG.rst

@@ -1,3 +1,9 @@
+**3.1.4 - 11/22/24**
+
+ - Feature: Enable initializing a population of all newborns
+ - Feature: Enable setting data sources directly on DiseaseModel components
+ - Refactor: Update DiseaseModel to leverage the on_initialize_simulants method of Machine
+
 **3.1.3 - 12/03/24**
 
  - Feature: Create relative risk pipeline

{vivarium_public_health-3.1.3 → vivarium_public_health-3.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vivarium_public_health
-Version: 3.1.3
+Version: 3.1.4
 Summary: Components for modelling diseases, risks, and interventions with ``vivarium``
 Home-page: https://github.com/ihmeuw/vivarium_public_health
 Author: The vivarium developers
@@ -26,7 +26,7 @@ Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
 Classifier: Topic :: Scientific/Engineering :: Physics
 Classifier: Topic :: Software Development :: Libraries
 License-File: LICENSE.txt
-Requires-Dist: vivarium>=3.2.0
+Requires-Dist: vivarium>=3.2.3
 Requires-Dist: layered_config_tree>=1.0.1
 Requires-Dist: loguru
 Requires-Dist: numpy<2.0.0

{vivarium_public_health-3.1.3 → vivarium_public_health-3.1.4}/setup.py

@@ -15,7 +15,7 @@ if __name__ == "__main__":
         long_description = f.read()
 
     install_requirements = [
-        "vivarium>=3.2.0",
+        "vivarium>=3.2.3",
         "layered_config_tree>=1.0.1",
         "loguru",
         "numpy<2.0.0",

vivarium_public_health-3.1.4/src/vivarium_public_health/_version.py

@@ -0,0 +1 @@
+__version__ = "3.1.4"

vivarium_public_health-3.1.4/src/vivarium_public_health/disease/exceptions.py

@@ -0,0 +1,5 @@
+from vivarium.exceptions import VivariumError
+
+
+class DiseaseModelError(VivariumError):
+    pass

{vivarium_public_health-3.1.3 → vivarium_public_health-3.1.4}/src/vivarium_public_health/disease/model.py

@@ -10,24 +10,20 @@ transitions at simulation initialization and during transitions.
 """
 import warnings
 from collections.abc import Callable, Iterable
+from functools import partial
 from typing import Any
 
-import numpy as np
 import pandas as pd
-from vivarium.exceptions import VivariumError
 from vivarium.framework.engine import Builder
-from vivarium.framework.event import Event
 from vivarium.framework.population import SimulantData
 from vivarium.framework.state_machine import Machine
+from vivarium.types import DataInput, LookupTableData
 
+from vivarium_public_health.disease.exceptions import DiseaseModelError
 from vivarium_public_health.disease.state import BaseDiseaseState, SusceptibleState
 from vivarium_public_health.disease.transition import TransitionString
 
 
-class DiseaseModelError(VivariumError):
-    pass
-
-
 class DiseaseModel(Machine):
 
     ##############
@@ -82,57 +78,77 @@ class DiseaseModel(Machine):
         cause_type: str = "cause",
         states: Iterable[BaseDiseaseState] = (),
         residual_state: BaseDiseaseState | None = None,
-        **kwargs,
-    ):
-        super().__init__(cause, states=states, **kwargs)
+        cause_specific_mortality_rate: DataInput | None = None,
+    ) -> None:
+        super().__init__(cause, states=states)
         self.cause = cause
         self.cause_type = cause_type
+        self.residual_state = self._get_residual_state(initial_state, residual_state)
+        self._csmr_source = cause_specific_mortality_rate
+        self._get_data_functions = (
+            get_data_functions if get_data_functions is not None else {}
+        )
 
-        if initial_state is not None:
+        if get_data_functions is not None:
             warnings.warn(
-                "In the future, the 'initial_state' argument to DiseaseModel"
-                " will be used to initialize all simulants into that state. To"
-                " retain the current behavior of defining a residual state, use"
-                " the 'residual_state' argument.",
+                "The argument 'get_data_functions' has been deprecated. Use"
+                " cause_specific_mortality_rate instead.",
                 DeprecationWarning,
                 stacklevel=2,
             )
 
-            if residual_state:
+            if cause_specific_mortality_rate is not None:
                 raise DiseaseModelError(
-                    "A DiseaseModel cannot be initialized with both"
-                    " 'initial_state and 'residual_state'."
+                    "It is not allowed to pass cause_specific_mortality_rate"
+                    " both as a stand-alone argument and as part of"
+                    " get_data_functions."
                 )
 
-            self.residual_state = initial_state.state_id
-        elif residual_state is not None:
-            self.residual_state = residual_state.state_id
-        else:
-            self.residual_state = self._get_default_residual_state()
-
-        self._get_data_functions = (
-            get_data_functions if get_data_functions is not None else {}
-        )
-
     def setup(self, builder: Builder) -> None:
         """Perform this component's setup."""
         super().setup(builder)
 
         self.configuration_age_start = builder.configuration.population.initialization_age_min
         self.configuration_age_end = builder.configuration.population.initialization_age_max
+
         builder.value.register_value_modifier(
             "cause_specific_mortality_rate",
             self.adjust_cause_specific_mortality_rate,
             requires_columns=["age", "sex"],
         )
-        self.randomness = builder.randomness.get_stream(f"{self.state_column}_initial_states")
+
+    def on_initialize_simulants(self, pop_data: SimulantData) -> None:
+        """Initialize the simulants in the population.
+
+        If all simulants are initialized at age 0, birth prevalence is used.
+        Otherwise, prevalence is used.
+
+        Parameters
+        ----------
+        pop_data
+            The population data object.
+        """
+        if pop_data.user_data.get("age_end", self.configuration_age_end) == 0:
+            initialization_table_name = "birth_prevalence"
+        else:
+            initialization_table_name = "prevalence"
+
+        for state in self.states:
+            state.lookup_tables["initialization_weights"] = state.lookup_tables[
+                initialization_table_name
+            ]
+
+        super().on_initialize_simulants(pop_data)
 
     #################
     # Setup methods #
     #################
 
     def load_cause_specific_mortality_rate(self, builder: Builder) -> float | pd.DataFrame:
-        if "cause_specific_mortality_rate" not in self._get_data_functions:
+        if (
+            "cause_specific_mortality_rate" not in self._get_data_functions
+            and self._csmr_source is None
+        ):
             only_morbid = builder.data.load(f"cause.{self.cause}.restrictions")["yld_only"]
             if only_morbid:
                 csmr_data = 0.0
@@ -140,64 +156,14 @@ class DiseaseModel(Machine):
                 csmr_data = builder.data.load(
                     f"{self.cause_type}.{self.cause}.cause_specific_mortality_rate"
                 )
+        elif self._csmr_source is not None:
+            csmr_data = self.get_data(builder, self._csmr_source)
         else:
             csmr_data = self._get_data_functions["cause_specific_mortality_rate"](
                 self.cause, builder
             )
         return csmr_data
 
-    ########################
-    # Event-driven methods #
-    ########################
-
-    def on_initialize_simulants(self, pop_data: SimulantData) -> None:
-        population = self.population_view.subview(["age", "sex"]).get(pop_data.index)
-
-        assert self.residual_state in {s.state_id for s in self.states}
-
-        if pop_data.user_data["sim_state"] == "setup":  # simulation start
-            if self.configuration_age_start != self.configuration_age_end != 0:
-                state_names, weights_bins = self.get_state_weights(
-                    pop_data.index, "prevalence"
-                )
-            else:
-                raise NotImplementedError(
-                    "We do not currently support an age 0 cohort. "
-                    "configuration.population.initialization_age_min and "
-                    "configuration.population.initialization_age_max "
-                    "cannot both be 0."
-                )
-
-        else:  # on time step
-            if pop_data.user_data["age_start"] == pop_data.user_data["age_end"] == 0:
-                state_names, weights_bins = self.get_state_weights(
-                    pop_data.index, "birth_prevalence"
-                )
-            else:
-                state_names, weights_bins = self.get_state_weights(
-                    pop_data.index, "prevalence"
-                )
-
-        if state_names and not population.empty:
-            # only do this if there are states in the model that supply prevalence data
-            population["sex_id"] = population.sex.apply({"Male": 1, "Female": 2}.get)
-
-            condition_column = self.assign_initial_status_to_simulants(
-                population,
-                state_names,
-                weights_bins,
-                self.randomness.get_draw(population.index),
-            )
-
-            condition_column = condition_column.rename(
-                columns={"condition_state": self.state_column}
-            )
-        else:
-            condition_column = pd.Series(
-                self.residual_state, index=population.index, name=self.state_column
-            )
-        self.population_view.update(condition_column)
-
     ##################################
     # Pipeline sources and modifiers #
     ##################################
@@ -209,40 +175,78 @@ class DiseaseModel(Machine):
     # Helper functions #
     ####################
 
-    def _get_default_residual_state(self):
-        susceptible_states = [s for s in self.states if isinstance(s, SusceptibleState)]
-        if len(susceptible_states) != 1:
-            raise DiseaseModelError("Disease model must have exactly one SusceptibleState.")
-        return susceptible_states[0].state_id
-
-    def get_state_weights(
-        self, pop_index: pd.Index, prevalence_type: str
-    ) -> tuple[list[str], np.ndarray | None]:
-        states = [state for state in self.states if state.lookup_tables.get(prevalence_type)]
-
-        if not states:
-            return states, None
+    def _get_residual_state(
+        self, initial_state: BaseDiseaseState, residual_state: BaseDiseaseState
+    ) -> BaseDiseaseState:
+        """Get the residual state for the DiseaseModel.
 
-        weights = [state.lookup_tables.get(prevalence_type)(pop_index) for state in states]
-        for w in weights:
-            w.reset_index(inplace=True, drop=True)
-        weights += ((1 - np.sum(weights, axis=0)),)
+        This will be the residual state if it is provided, otherwise it will be
+        the model's SusceptibleState. This method also calculates the residual
+        state's birth_prevalence and prevalence.
+        """
+        if initial_state is not None:
+            warnings.warn(
+                "In the future, the 'initial_state' argument to DiseaseModel"
+                " will be used to initialize all simulants into that state. To"
+                " retain the current behavior of defining a residual state, use"
+                " the 'residual_state' argument.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
 
-        weights = np.array(weights).T
-        weights_bins = np.cumsum(weights, axis=1)
+            if residual_state:
+                raise DiseaseModelError(
+                    "A DiseaseModel cannot be initialized with both"
+                    " 'initial_state and 'residual_state'."
+                )
 
-        state_names = [s.state_id for s in states] + [self.residual_state]
+            residual_state = initial_state
+        elif residual_state is None:
+            susceptible_states = [s for s in self.states if isinstance(s, SusceptibleState)]
+            if len(susceptible_states) != 1:
+                raise DiseaseModelError(
+                    "Disease model must have exactly one SusceptibleState."
+                )
+            residual_state = susceptible_states[0]
 
-        return state_names, weights_bins
+        if residual_state not in self.states:
+            raise DiseaseModelError(
+                f"Residual state '{self.residual_state}' must be one of the"
+                f" states: {self.states}."
+            )
 
-    @staticmethod
-    def assign_initial_status_to_simulants(
-        simulants_df, state_names, weights_bins, propensities
-    ):
-        simulants = simulants_df[["age", "sex"]].copy()
+        residual_state.birth_prevalence = partial(
+            self._get_residual_state_probabilities, table_name="birth_prevalence"
+        )
+        residual_state.prevalence = partial(
+            self._get_residual_state_probabilities, table_name="prevalence"
+        )
 
-        choice_index = (propensities.values[np.newaxis].T > weights_bins).sum(axis=1)
-        initial_states = pd.Series(np.array(state_names)[choice_index], index=simulants.index)
+        return residual_state
+
+    def _get_residual_state_probabilities(
+        self, builder: Builder, table_name: str
+    ) -> LookupTableData:
+        """Calculate the probabilities of the residual state based on the other states."""
+        non_residual_states = [s for s in self.states if s != self.residual_state]
+        non_residual_probabilities = 0
+        for state in non_residual_states:
+            weights_source = builder.configuration[state.name].data_sources[table_name]
+            weights = state.get_data(builder, weights_source)
+            if isinstance(weights, pd.DataFrame):
+                weights = weights.set_index(
+                    [c for c in weights.columns if c != "value"]
+                ).squeeze()
+            non_residual_probabilities += weights
+
+        residual_probabilities = 1 - non_residual_probabilities
+
+        if pd.Series(residual_probabilities < 0).any():
+            raise ValueError(
+                f"The {table_name} for the states in the DiseaseModel must sum"
+                " to less than 1."
+            )
+        if isinstance(residual_probabilities, pd.Series):
+            residual_probabilities = residual_probabilities.reset_index()
 
-        simulants.loc[:, "condition_state"] = initial_states
-        return simulants
+        return residual_probabilities