diff --git a/src/vivarium_gates_nutrition_optimization/components/intervention.py b/src/vivarium_gates_nutrition_optimization/components/intervention.py index 6896d78..eee11bd 100644 --- a/src/vivarium_gates_nutrition_optimization/components/intervention.py +++ b/src/vivarium_gates_nutrition_optimization/components/intervention.py @@ -35,7 +35,7 @@ def columns_required(self) -> List[str]: @property def initialization_requirements(self) -> Dict[str, List[str]]: - return {"requires_streams": [self.name], "requires_columns": [self.columns_required]} + return {"requires_streams": [self.name], "requires_columns": self.columns_required} # noinspection PyAttributeOutsideInit def setup(self, builder: Builder) -> None: diff --git a/src/vivarium_gates_nutrition_optimization/components/pregnancy.py b/src/vivarium_gates_nutrition_optimization/components/pregnancy.py index 98ef394..f215b3f 100644 --- a/src/vivarium_gates_nutrition_optimization/components/pregnancy.py +++ b/src/vivarium_gates_nutrition_optimization/components/pregnancy.py @@ -23,8 +23,8 @@ def __init__(self, state_id, *args, **kwargs): class PregnantState(DiseaseState): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, state_id, *args, **kwargs): + super().__init__(state_id, *args, **kwargs) self.new_children = NewChildren() ############## @@ -212,6 +212,10 @@ def get_initial_event_times(self, pop_data: SimulantData) -> pd.DataFrame: class PregnancyModel(DiseaseModel): + def __init__(self, cause: str, **kwargs): + super().__init__(cause, **kwargs) + self._name = f"disease_model.{cause}" + @property def time_step_priority(self) -> int: return 3