support covid for sarix

reichlab · Nov 21, 2024 · a09605f · a09605f
1 parent c7982f9
commit a09605f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/idmodels/sarix.py b/src/idmodels/sarix.py
@@ -1,7 +1,7 @@
 
 import numpy as np
 import pandas as pd
-from iddata.loader import FluDataLoader
+from iddata.loader import DiseaseDataLoader
 from iddata.utils import get_holidays
 from sarix import sarix
 
@@ -13,8 +13,8 @@ def __init__(self, model_config):
         self.model_config = model_config
 
     def run(self, run_config):
-        fdl = FluDataLoader()
-        df = fdl.load_data(nhsn_kwargs={"as_of": run_config.ref_date},
+        fdl = DiseaseDataLoader()
+        df = fdl.load_data(nhsn_kwargs={"as_of": run_config.ref_date, "disease": run_config.disease},
                            sources=self.model_config.sources,
                            power_transform=self.model_config.power_transform)
         if run_config.locations is not None: