fix a couple bugs with naming

ihmeuw · Dec 13, 2023 · d00c301 · d00c301
1 parent d33a4ee
commit d00c301
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git 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
@@ -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