vivarium-public-health 3.1.3__py3-none-any.whl → 3.1.4__py3-none-any.whl

vivarium_public_health/_version.py

@@ -1 +1 @@
-__version__ = "3.1.3"
+__version__ = "3.1.4"

vivarium_public_health/disease/exceptions.py

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

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

vivarium_public_health/disease/state.py

@@ -6,19 +6,20 @@ Disease States
 This module contains tools to manage standard disease states.
 
 """
-
+import warnings
 from collections.abc import Callable
 from typing import Any
 
 import numpy as np
 import pandas as pd
 from vivarium.framework.engine import Builder
-from vivarium.framework.lookup import LookupTableData
 from vivarium.framework.population import PopulationView, SimulantData
 from vivarium.framework.randomness import RandomnessStream
 from vivarium.framework.state_machine import State, Transient, Transition, Trigger
 from vivarium.framework.values import Pipeline, list_combiner, union_post_processor
+from vivarium.types import DataInput, LookupTableData
 
+from vivarium_public_health.disease.exceptions import DiseaseModelError
 from vivarium_public_health.disease.transition import (
     ProportionTransition,
     RateTransition,
@@ -28,10 +29,25 @@ from vivarium_public_health.utilities import get_lookup_columns, is_non_zero
 
 
 class BaseDiseaseState(State):
+
     ##############
     # Properties #
     ##############
 
+    @property
+    def configuration_defaults(self) -> dict[str, Any]:
+        configuration_defaults = super().configuration_defaults
+        additional_defaults = {
+            "prevalence": self.prevalence,
+            "birth_prevalence": self.birth_prevalence,
+        }
+        data_sources = {
+            **configuration_defaults[self.name]["data_sources"],
+            **additional_defaults,
+        }
+        configuration_defaults[self.name]["data_sources"] = data_sources
+        return configuration_defaults
+
     @property
     def columns_created(self):
         return [self.event_time_column, self.event_count_column]
@@ -59,13 +75,15 @@ class BaseDiseaseState(State):
         side_effect_function: Callable | None = None,
         cause_type: str = "cause",
     ):
-        super().__init__(state_id, allow_self_transition)  # becomes state_id
+        super().__init__(state_id, allow_self_transition)
         self.cause_type = cause_type
 
         self.side_effect_function = side_effect_function
 
         self.event_time_column = self.state_id + "_event_time"
         self.event_count_column = self.state_id + "_event_count"
+        self.prevalence = 0.0
+        self.birth_prevalence = 0.0
 
     ########################
     # Event-driven methods #
@@ -87,10 +105,7 @@ class BaseDiseaseState(State):
     def get_initialization_parameters(self) -> dict[str, Any]:
         """Exclude side effect function and cause type from name and __repr__."""
         initialization_parameters = super().get_initialization_parameters()
-        for key in ["side_effect_function", "cause_type"]:
-            if key in initialization_parameters.keys():
-                del initialization_parameters[key]
-        return initialization_parameters
+        return {"state_id": initialization_parameters["state_id"]}
 
     def get_initial_event_times(self, pop_data: SimulantData) -> pd.DataFrame:
         return pd.DataFrame(
@@ -132,6 +147,8 @@ class BaseDiseaseState(State):
         output: "BaseDiseaseState",
         get_data_functions: dict[str, Callable] = None,
         triggered=Trigger.NOT_TRIGGERED,
+        transition_rate: DataInput | None = None,
+        rate_type: str = "transition_rate",
     ) -> RateTransition:
         """Builds a RateTransition from this state to the given state.
 
@@ -139,18 +156,29 @@ 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.
         triggered
             The trigger for the transition
-
+        transition_rate
+            The transition rate source. Can be the data itself, a function to
+            retrieve the data, or the artifact key containing the data.
+        rate_type
+            The type of rate. Can be "incidence_rate", "transition_rate", or
+            "remission_rate".
 
         Returns
         -------
             The created transition object.
         """
-        transition = RateTransition(self, output, get_data_functions, triggered)
+        transition = RateTransition(
+            input_state=self,
+            output_state=output,
+            get_data_functions=get_data_functions,
+            triggered=triggered,
+            transition_rate=transition_rate,
+            rate_type=rate_type,
+        )
         self.add_transition(transition)
         return transition
 
@@ -159,6 +187,7 @@ class BaseDiseaseState(State):
         output: "BaseDiseaseState",
         get_data_functions: dict[str, Callable] | None = None,
         triggered=Trigger.NOT_TRIGGERED,
+        proportion: DataInput | None = None,
     ) -> ProportionTransition:
         """Builds a ProportionTransition from this state to the given state.
 
@@ -170,15 +199,26 @@ class BaseDiseaseState(State):
             Map from transition type to the function to pull that transition's data.
         triggered
             The trigger for the transition.
+        proportion
+            The proportion source.  Can be the data itself, a function to
+            retrieve the data, or the artifact key containing the data.
 
         Returns
         -------
             The created transition object.
         """
-        if "proportion" not in get_data_functions:
+        if (
+            get_data_functions is None or "proportion" not in get_data_functions
+        ) and proportion is None:
             raise ValueError("You must supply a proportion function.")
 
-        transition = ProportionTransition(self, output, get_data_functions, triggered)
+        transition = ProportionTransition(
+            input_state=self,
+            output_state=output,
+            get_data_functions=get_data_functions,
+            triggered=triggered,
+            proportion=proportion,
+        )
         self.add_transition(transition)
         return transition
 
@@ -220,17 +260,19 @@ class NonDiseasedState(BaseDiseaseState):
         self,
         output: BaseDiseaseState,
         get_data_functions: dict[str, Callable] = None,
-        **kwargs,
+        triggered=Trigger.NOT_TRIGGERED,
+        transition_rate: DataInput | None = None,
+        **_kwargs,
     ) -> RateTransition:
-        if get_data_functions is None:
-            get_data_functions = {
-                "incidence_rate": lambda builder, cause: builder.data.load(
-                    f"{self.cause_type}.{cause}.incidence_rate"
-                )
-            }
-        elif "incidence_rate" not in get_data_functions:
-            raise ValueError("You must supply an incidence rate function.")
-        return super().add_rate_transition(output, get_data_functions, **kwargs)
+        if get_data_functions is None and transition_rate is None:
+            transition_rate = f"{self.cause_type}.{output.state_id}.incidence_rate"
+        return super().add_rate_transition(
+            output=output,
+            get_data_functions=get_data_functions,
+            triggered=triggered,
+            transition_rate=transition_rate,
+            rate_type="incidence_rate",
+        )
 
 
 class SusceptibleState(NonDiseasedState):
@@ -253,6 +295,13 @@ class SusceptibleState(NonDiseasedState):
             name_prefix="susceptible_to_",
         )
 
+    ##################
+    # Public methods #
+    ##################
+
+    def has_initialization_weights(self) -> bool:
+        return True
+
 
 class RecoveredState(NonDiseasedState):
     def __init__(
@@ -280,17 +329,20 @@ class DiseaseState(BaseDiseaseState):
 
     @property
     def configuration_defaults(self) -> dict[str, Any]:
-        return {
-            f"{self.name}": {
-                "data_sources": {
-                    "prevalence": self.load_prevalence,
-                    "birth_prevalence": self.load_birth_prevalence,
-                    "dwell_time": self.load_dwell_time,
-                    "disability_weight": self.load_disability_weight,
-                    "excess_mortality_rate": self.load_excess_mortality_rate,
-                },
-            },
+        configuration_defaults = super().configuration_defaults
+        additional_defaults = {
+            "prevalence": self._prevalence_source,
+            "birth_prevalence": self._birth_prevalence_source,
+            "dwell_time": self._dwell_time_source,
+            "disability_weight": self._disability_weight_source,
+            "excess_mortality_rate": self._excess_modality_rate_source,
         }
+        data_sources = {
+            **configuration_defaults[self.name]["data_sources"],
+            **additional_defaults,
+        }
+        configuration_defaults[self.name]["data_sources"] = data_sources
+        return configuration_defaults
 
     #####################
     # Lifecycle methods #
@@ -304,6 +356,11 @@ class DiseaseState(BaseDiseaseState):
         cause_type: str = "cause",
         get_data_functions: dict[str, Callable] | None = None,
         cleanup_function: Callable | None = None,
+        prevalence: DataInput | None = None,
+        birth_prevalence: DataInput | None = None,
+        dwell_time: DataInput | None = None,
+        disability_weight: DataInput | None = None,
+        excess_mortality_rate: DataInput | None = None,
     ):
         """
         Parameters
@@ -322,7 +379,31 @@ class DiseaseState(BaseDiseaseState):
             various state attributes.
         cleanup_function
             The cleanup function.
+        prevalence
+            The prevalence source. This is used to initialize simulants. Can be
+            the data itself, a function to retrieve the data, or the artifact
+            key containing the data.
+        birth_prevalence
+            The birth prevalence source. This is used to initialize newborn
+            simulants. Can be the data itself, a function to retrieve the data,
+            or the artifact key containing the data.
+        dwell_time
+            The dwell time source. This is used to determine how long a simulant
+            must remain in the state before transitioning. Can be the data
+            itself, a function to retrieve the data, or the artifact key
+            containing the data.
+        disability_weight
+            The disability weight source. This is used to calculate the
+            disability weight for simulants in this state. Can be the data
+            itself, a function to retrieve the data, or the artifact key
+            containing the data.
+        excess_mortality_rate
+            The excess mortality rate source. This is used to calculate the
+            excess mortality rate for simulants in this state. Can be the data
+            itself, a function to retrieve the data, or the artifact key
+            containing the data.
         """
+
         super().__init__(
             state_id,
             allow_self_transition=allow_self_transition,
@@ -340,6 +421,34 @@ class DiseaseState(BaseDiseaseState):
         )
         self._cleanup_function = cleanup_function
 
+        if get_data_functions is not None:
+            warnings.warn(
+                "The argument 'get_data_functions' has been deprecated. Use"
+                " cause_specific_mortality_rate instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+            for data_type in self._get_data_functions:
+                try:
+                    data_source = locals()[data_type]
+                except KeyError:
+                    data_source = None
+
+                if locals()[data_type] is not None:
+                    raise DiseaseModelError(
+                        f"It is not allowed to pass '{data_type}' both as a"
+                        " stand-alone argument and as part of get_data_functions."
+                    )
+
+        self._prevalence_source = self.get_prevalence_source(prevalence)
+        self._birth_prevalence_source = self.get_birth_prevalence_source(birth_prevalence)
+        self._dwell_time_source = self.get_dwell_time_source(dwell_time)
+        self._disability_weight_source = self.get_disability_weight_source(disability_weight)
+        self._excess_modality_rate_source = self.get_excess_mortality_rate_source(
+            excess_mortality_rate
+        )
+
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
         """Performs this component's simulation setup.
@@ -377,32 +486,45 @@ class DiseaseState(BaseDiseaseState):
     # Setup methods #
     #################
 
-    def load_prevalence(self, builder: Builder) -> LookupTableData:
+    def _get_data_functions_source(self, data_type: str) -> DataInput:
+        def data_source(builder: Builder) -> LookupTableData:
+            return self._get_data_functions[data_type](builder, self.state_id)
+
+        return data_source
+
+    def get_prevalence_source(self, prevalence: DataInput | None) -> DataInput:
         if "prevalence" in self._get_data_functions:
-            return self._get_data_functions["prevalence"](builder, self.state_id)
+            return self._get_data_functions_source("prevalence")
+        elif prevalence is not None:
+            return prevalence
         else:
-            return builder.data.load(f"{self.cause_type}.{self.state_id}.prevalence")
+            return f"{self.cause_type}.{self.state_id}.prevalence"
 
-    def load_birth_prevalence(self, builder: Builder) -> LookupTableData:
+    def get_birth_prevalence_source(self, birth_prevalence: DataInput | None) -> DataInput:
         if "birth_prevalence" in self._get_data_functions:
-            return self._get_data_functions["birth_prevalence"](builder, self.state_id)
+            return self._get_data_functions_source("birth_prevalence")
+        elif birth_prevalence is not None:
+            return birth_prevalence
         else:
-            return 0
+            return 0.0
 
-    def load_dwell_time(self, builder: Builder) -> LookupTableData:
+    def get_dwell_time_source(self, dwell_time: DataInput | None) -> DataInput:
         if "dwell_time" in self._get_data_functions:
-            dwell_time = self._get_data_functions["dwell_time"](builder, self.state_id)
-        else:
-            dwell_time = 0
-
-        if isinstance(dwell_time, pd.Timedelta):
-            dwell_time = dwell_time.total_seconds() / (60 * 60 * 24)
-        if (
-            isinstance(dwell_time, pd.DataFrame) and np.any(dwell_time.value != 0)
-        ) or dwell_time > 0:
-            self.transition_set.allow_null_transition = True
-
-        return dwell_time
+            dwell_time = self._get_data_functions_source("dwell_time")
+        elif dwell_time is None:
+            dwell_time = 0.0
+
+        def dwell_time_source(builder: Builder) -> LookupTableData:
+            dwell_time_ = self.get_data(builder, dwell_time)
+            if isinstance(dwell_time_, pd.Timedelta):
+                dwell_time_ = dwell_time_.total_seconds() / (60 * 60 * 24)
+            if (
+                isinstance(dwell_time_, pd.DataFrame) and np.any(dwell_time_.value != 0)
+            ) or dwell_time_ > 0:
+                self.transition_set.allow_null_transition = True
+            return dwell_time_
+
+        return dwell_time_source
 
     def get_dwell_time_pipeline(self, builder: Builder) -> Pipeline:
         required_columns = get_lookup_columns([self.lookup_tables["dwell_time"]])
@@ -412,20 +534,22 @@ class DiseaseState(BaseDiseaseState):
             requires_columns=required_columns,
         )
 
-    def load_disability_weight(self, builder: Builder) -> LookupTableData:
+    def get_disability_weight_source(self, disability_weight: DataInput | None) -> DataInput:
         if "disability_weight" in self._get_data_functions:
-            disability_weight = self._get_data_functions["disability_weight"](
-                builder, self.state_id
-            )
+            disability_weight = self._get_data_functions_source("disability_weight")
+        elif disability_weight is not None:
+            disability_weight = disability_weight
         else:
-            disability_weight = builder.data.load(
-                f"{self.cause_type}.{self.state_id}.disability_weight"
-            )
+            disability_weight = f"{self.cause_type}.{self.state_id}.disability_weight"
 
-        if isinstance(disability_weight, pd.DataFrame) and len(disability_weight) == 1:
-            disability_weight = disability_weight.value[0]  # sequela only have single value
+        def disability_weight_source(builder: Builder) -> LookupTableData:
+            disability_weight_ = self.get_data(builder, disability_weight)
+            if isinstance(disability_weight_, pd.DataFrame) and len(disability_weight_) == 1:
+                # sequela only have single value
+                disability_weight_ = disability_weight_.value[0]
+            return disability_weight_
 
-        return disability_weight
+        return disability_weight_source
 
     def get_disability_weight_pipeline(self, builder: Builder) -> Pipeline:
         lookup_columns = get_lookup_columns([self.lookup_tables["disability_weight"]])
@@ -435,16 +559,25 @@ class DiseaseState(BaseDiseaseState):
             requires_columns=lookup_columns + ["alive", self.model],
         )
 
-    def load_excess_mortality_rate(self, builder: Builder) -> LookupTableData:
+    def get_excess_mortality_rate_source(
+        self, excess_mortality_rate: DataInput | None
+    ) -> DataInput:
         if "excess_mortality_rate" in self._get_data_functions:
-            return self._get_data_functions["excess_mortality_rate"](builder, self.state_id)
-        elif builder.data.load(f"cause.{self.model}.restrictions")["yld_only"]:
-            return 0
-        else:
+            excess_mortality_rate = self._get_data_functions_source("excess_mortality_rate")
+        elif excess_mortality_rate is None:
+            excess_mortality_rate = f"{self.cause_type}.{self.state_id}.excess_mortality_rate"
+
+        def excess_mortality_rate_source(builder: Builder) -> LookupTableData:
+            if excess_mortality_rate is not None:
+                return self.get_data(builder, excess_mortality_rate)
+            elif builder.data.load(f"cause.{self.model}.restrictions")["yld_only"]:
+                return 0
             return builder.data.load(
                 f"{self.cause_type}.{self.state_id}.excess_mortality_rate"
             )
 
+        return excess_mortality_rate_source
+
     def get_excess_mortality_rate_pipeline(self, builder: Builder) -> Pipeline:
         lookup_columns = get_lookup_columns([self.lookup_tables["excess_mortality_rate"]])
         return builder.value.register_rate_producer(
@@ -470,24 +603,27 @@ class DiseaseState(BaseDiseaseState):
     # Public methods #
     ##################
 
+    def has_initialization_weights(self) -> bool:
+        return True
+
     def add_rate_transition(
         self,
         output: BaseDiseaseState,
         get_data_functions: dict[str, Callable] = None,
-        **kwargs,
+        triggered=Trigger.NOT_TRIGGERED,
+        transition_rate: DataInput | None = None,
+        rate_type: str = "transition_rate",
     ) -> RateTransition:
-        if get_data_functions is None:
-            get_data_functions = {
-                "remission_rate": lambda builder, cause: builder.data.load(
-                    f"{self.cause_type}.{cause}.remission_rate"
-                )
-            }
-        elif (
-            "remission_rate" not in get_data_functions
-            and "transition_rate" not in get_data_functions
-        ):
-            raise ValueError("You must supply a transition rate or remission rate function.")
-        return super().add_rate_transition(output, get_data_functions, **kwargs)
+        if get_data_functions is None and transition_rate is None:
+            transition_rate = f"{self.cause_type}.{self.state_id}.remission_rate"
+            rate_type = "remission_rate"
+        return super().add_rate_transition(
+            output=output,
+            get_data_functions=get_data_functions,
+            triggered=triggered,
+            transition_rate=transition_rate,
+            rate_type=rate_type,
+        )
 
     def add_dwell_time_transition(
         self,

vivarium_public_health/disease/transition.py

@@ -6,7 +6,7 @@ Disease Transitions
 This module contains tools to model transitions between disease states.
 
 """
-
+import warnings
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
@@ -15,7 +15,9 @@ from vivarium.framework.engine import Builder
 from vivarium.framework.state_machine import Transition, Trigger
 from vivarium.framework.utilities import rate_to_probability
 from vivarium.framework.values import list_combiner, union_post_processor
+from vivarium.types import DataInput
 
+from vivarium_public_health.disease.exceptions import DiseaseModelError
 from vivarium_public_health.utilities import get_lookup_columns
 
 if TYPE_CHECKING:
@@ -44,7 +46,7 @@ class RateTransition(Transition):
         return {
             f"{self.name}": {
                 "data_sources": {
-                    "transition_rate": self.load_transition_rate,
+                    "transition_rate": self._rate_source,
                 },
             },
         }
@@ -55,19 +57,37 @@ class RateTransition(Transition):
 
     @property
     def transition_rate_pipeline_name(self) -> str:
-        if "incidence_rate" in self._get_data_functions:
-            pipeline_name = f"{self.output_state.state_id}.incidence_rate"
-        elif "remission_rate" in self._get_data_functions:
-            pipeline_name = f"{self.input_state.state_id}.remission_rate"
-        elif "transition_rate" in self._get_data_functions:
-            pipeline_name = (
-                f"{self.input_state.state_id}_to_{self.output_state.state_id}.transition_rate"
-            )
+        if self._get_data_functions:
+            if "incidence_rate" in self._get_data_functions:
+                pipeline_name = f"{self.output_state.state_id}.incidence_rate"
+            elif "remission_rate" in self._get_data_functions:
+                pipeline_name = f"{self.input_state.state_id}.remission_rate"
+            elif "transition_rate" in self._get_data_functions:
+                pipeline_name = (
+                    f"{self.input_state.state_id}_to_{self.output_state.state_id}"
+                    ".transition_rate"
+                )
+            else:
+                raise DiseaseModelError(
+                    "Cannot determine rate_transition pipeline name: "
+                    "no valid data functions supplied."
+                )
         else:
-            raise ValueError(
-                "Cannot determine rate_transition pipeline name: "
-                "no valid data functions supplied."
-            )
+            if self.rate_type == "incidence_rate":
+                pipeline_name = f"{self.output_state.state_id}.incidence_rate"
+            elif self.rate_type == "remission_rate":
+                pipeline_name = f"{self.input_state.state_id}.remission_rate"
+            elif self.rate_type == "transition_rate":
+                pipeline_name = (
+                    f"{self.input_state.state_id}_to_{self.output_state.state_id}"
+                    ".transition_rate"
+                )
+            else:
+                raise DiseaseModelError(
+                    "Cannot determine rate_transition pipeline name: invalid"
+                    f" rate_type '{self.rate_type} supplied."
+                )
+
         return pipeline_name
 
     #####################
@@ -80,6 +100,8 @@ class RateTransition(Transition):
         output_state: "BaseDiseaseState",
         get_data_functions: dict[str, Callable] = None,
         triggered=Trigger.NOT_TRIGGERED,
+        transition_rate: DataInput | None = None,
+        rate_type: str = "transition_rate",
     ):
         super().__init__(
             input_state, output_state, probability_func=self._probability, triggered=triggered
@@ -87,6 +109,22 @@ class RateTransition(Transition):
         self._get_data_functions = (
             get_data_functions if get_data_functions is not None else {}
         )
+        self._rate_source = self._get_rate_source(transition_rate)
+        self.rate_type = rate_type
+
+        if get_data_functions is not None:
+            warnings.warn(
+                "The argument 'get_data_functions' has been deprecated. Use"
+                " 'transition_rate' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if transition_rate is not None:
+                raise DiseaseModelError(
+                    "It is not allowed to pass a transition rate"
+                    " both as a stand-alone argument and as part of"
+                    " get_data_functions."
+                )
 
     # noinspection PyAttributeOutsideInit
     def setup(self, builder: Builder) -> None:
@@ -109,17 +147,19 @@ class RateTransition(Transition):
     # Setup methods #
     #################
 
-    def load_transition_rate(self, builder: Builder) -> float | pd.DataFrame:
-        if "incidence_rate" in self._get_data_functions:
-            rate_data = self._get_data_functions["incidence_rate"](
+    def _get_rate_source(self, transition_rate: DataInput | None) -> DataInput:
+        if transition_rate is not None:
+            rate_data = transition_rate
+        elif "incidence_rate" in self._get_data_functions:
+            rate_data = lambda builder: self._get_data_functions["incidence_rate"](
                 builder, self.output_state.state_id
             )
         elif "remission_rate" in self._get_data_functions:
-            rate_data = self._get_data_functions["remission_rate"](
+            rate_data = lambda builder: self._get_data_functions["remission_rate"](
                 builder, self.input_state.state_id
             )
         elif "transition_rate" in self._get_data_functions:
-            rate_data = self._get_data_functions["transition_rate"](
+            rate_data = lambda builder: self._get_data_functions["transition_rate"](
                 builder, self.input_state.state_id, self.output_state.state_id
             )
         else:
@@ -172,21 +212,38 @@ class ProportionTransition(Transition):
         output_state: "BaseDiseaseState",
         get_data_functions: dict[str, Callable] = None,
         triggered=Trigger.NOT_TRIGGERED,
+        proportion: DataInput | None = None,
     ):
         super().__init__(
             input_state, output_state, probability_func=self._probability, triggered=triggered
         )
+        self._proportion_source = proportion
         self._get_data_functions = (
             get_data_functions if get_data_functions is not None else {}
         )
 
+        if get_data_functions is not None:
+            warnings.warn(
+                "The argument 'get_data_functions' has been deprecated. Use"
+                " 'proportion' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if proportion is not None:
+                raise DiseaseModelError(
+                    "It is not allowed to pass a proportion both as a"
+                    " stand-alone argument and as part of get_data_functions."
+                )
+
     #################
     # Setup methods #
     #################
 
-    def load_proportion(self, builder: Builder) -> float | pd.DataFrame:
+    def load_proportion(self, builder: Builder) -> DataInput:
+        if self._proportion_source is not None:
+            return self._proportion_source
         if "proportion" not in self._get_data_functions:
-            raise ValueError("Must supply a proportion function")
+            raise DiseaseModelError("Must supply a proportion function")
         return self._get_data_functions["proportion"](builder, self.output_state.state_id)
 
     def _probability(self, index):

{vivarium_public_health-3.1.3.dist-info → vivarium_public_health-3.1.4.dist-info}/METADATA

@@ -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.dist-info → vivarium_public_health-3.1.4.dist-info}/RECORD

@@ -1,13 +1,14 @@
 vivarium_public_health/__about__.py,sha256=RgWycPypKZS80TpSX7o41cREnG8PfguNHDHLuLyl820,487
 vivarium_public_health/__init__.py,sha256=GDeeP-7OlCBwPuv_xQoB1wNmvCaFsqfTB7qnnYApm0w,1343
-vivarium_public_health/_version.py,sha256=xoSqNkNOCK0xzFXnsO80epHc1vbiRC8Nbn4Cy-2wBX4,22
+vivarium_public_health/_version.py,sha256=2uixSkocUHf2KiY1oTfzz_5AQGmlrHnypVxGgr4mV9c,22
 vivarium_public_health/utilities.py,sha256=QNXQ6fhAr1HcV-GwKw7wQLz6QyuNxqNvMA-XujKjTgs,3035
 vivarium_public_health/disease/__init__.py,sha256=VUJHDLlE6ngo2qHNQUtZ8OWH5H_T7_ao-xsYKDkRmHw,443
-vivarium_public_health/disease/model.py,sha256=ZwhhQCc8jj_QeJZO2zLtp_yWzqRxvLjuzW7iDUmmBGA,8852
+vivarium_public_health/disease/exceptions.py,sha256=vb30IIV82OiDf2cNZCs_E2rF6mdDDHbnZSND60no5CU,97
+vivarium_public_health/disease/model.py,sha256=qA0mhbIL0rfPtJ1Csmvxvf757vUdrjcbpxYmL_dsniI,9175
 vivarium_public_health/disease/models.py,sha256=01UK7yB2zGPFzmlIpvhd-XnGe6vSCMDza3QTidgY7Nc,3479
 vivarium_public_health/disease/special_disease.py,sha256=kTVuE5rQjUK62ysComG8nB2f61aCKdca9trRB1zsDCQ,14537
-vivarium_public_health/disease/state.py,sha256=NwTnxB_i05magsJuEc1_Z_0xMMKsqSPXIpibUdxPrWY,22333
-vivarium_public_health/disease/transition.py,sha256=4g_F8L3Godb4yQjHRr42n95QXo6m0ldI6c8_vu6cTfo,6429
+vivarium_public_health/disease/state.py,sha256=l1iBZl8IpNOegQjZIY2xq09SBhsn5jSoYD89CymJAUs,28309
+vivarium_public_health/disease/transition.py,sha256=qRpHzq29S_I_KzyDncAahg27bnHNdrA5tsUWuJ4Zty4,8964
 vivarium_public_health/mslt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vivarium_public_health/mslt/delay.py,sha256=TrFp9nmv-PSu6xdQXfyeM4X8RrXTLUnQMxOuhdvD-Wk,22787
 vivarium_public_health/mslt/disease.py,sha256=DONSItBiOUk1qBE6Msw0vrV0XLW4BPzWVxwVK90GK1I,16638
@@ -42,8 +43,8 @@ vivarium_public_health/treatment/__init__.py,sha256=wONElu9aJbBYwpYIovYPYaN_GYfV
 vivarium_public_health/treatment/magic_wand.py,sha256=zGIhrNgB9q6JD7fHlvbDQb3H5e_N_QsROO4Y0kl_JQM,1955
 vivarium_public_health/treatment/scale_up.py,sha256=hVz0ELXDqlpcExI31rKdepxqcW_hy2hZSa6qCzv6udU,7020
 vivarium_public_health/treatment/therapeutic_inertia.py,sha256=8Z97s7GfcpfLu1U1ESJSqeEk4L__a3M0GbBV21MFg2s,2346
-vivarium_public_health-3.1.3.dist-info/LICENSE.txt,sha256=mN4bNLUQNcN9njYRc_3jCZkfPySVpmM6MRps104FxA4,1548
-vivarium_public_health-3.1.3.dist-info/METADATA,sha256=csuhuXUl84A5QlA7S10sjMvafQXonRvAlwm2XSj-ck4,4028
-vivarium_public_health-3.1.3.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-vivarium_public_health-3.1.3.dist-info/top_level.txt,sha256=VVInlpzCFD0UNNhjOq_j-a29odzjwUwYFTGfvqbi4dY,23
-vivarium_public_health-3.1.3.dist-info/RECORD,,
+vivarium_public_health-3.1.4.dist-info/LICENSE.txt,sha256=mN4bNLUQNcN9njYRc_3jCZkfPySVpmM6MRps104FxA4,1548
+vivarium_public_health-3.1.4.dist-info/METADATA,sha256=7h8TxTjuThUizK3wr0mzLB31wwS4Kqr9SVgJA9Gf7gA,4028
+vivarium_public_health-3.1.4.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+vivarium_public_health-3.1.4.dist-info/top_level.txt,sha256=VVInlpzCFD0UNNhjOq_j-a29odzjwUwYFTGfvqbi4dY,23
+vivarium_public_health-3.1.4.dist-info/RECORD,,