add support for fourier pooling option

Author: nickreich

src/idmodels/sarix.py

@@ -116,12 +116,14 @@ class SARIXFourierModel(SARIXModel):
     SARIX model with Fourier seasonality terms.
 
     Adds annual seasonal patterns using Fourier harmonics to the base SARIX model.
-    Requires fourier_K parameter in model_config to specify number of harmonic pairs.
+
+    Required model_config parameters:
+    - fourier_K: Number of Fourier harmonic pairs (int)
+    - fourier_pooling: How to share Fourier coefficients across locations ('none' or 'shared')
     """
     def _get_sarix_module(self):
-        """Return the sarix_fourier module for Fourier-enhanced fitting."""
-        from sarixfourier import sarix_fourier
-        return sarix_fourier
+        """Return the sarix module (same module, but with Fourier parameters)."""
+        return sarix
 
     def _get_extra_sarix_params(self, df):
         """Return Fourier-specific parameters for SARIX constructor."""
@@ -131,7 +133,8 @@ def _get_extra_sarix_params(self, df):
 
         return {
             "day_of_year": day_of_year,
-            "fourier_K": self.model_config.fourier_K
+            "fourier_K": self.model_config.fourier_K,
+            "fourier_pooling": self.model_config.fourier_pooling
         }
 
 

tests/integration/test_sarix.py

@@ -7,7 +7,7 @@
 import pandas as pd
 from pandas.testing import assert_frame_equal
 
-from idmodels.sarix import SARIXModel
+from idmodels.sarix import SARIXModel, SARIXFourierModel
 
 
 def test_sarix(tmp_path):
@@ -145,6 +145,194 @@ def test_sarix_shared_sigma_pooling_multiple_batches(tmp_path):
         "All predictions should be non-negative"
 
 
+def test_sarix_fourier_none_pooling(tmp_path):
+    """Test SARIXFourierModel with fourier_pooling='none' (unpooled)."""
+    model_config = SimpleNamespace(
+        model_class="sarix_fourier",
+        model_name="sarix_p2_fourier_K2_none",
+
+        # data sources
+        sources=["nhsn"],
+
+        # fit locations separately or jointly
+        fit_locations_separately=False,
+
+        # SARIX model parameters
+        p=2,
+        P=0,
+        d=0,
+        D=0,
+        season_period=1,
+
+        # power transform
+        power_transform="4rt",
+
+        # parameter pooling
+        theta_pooling="shared",
+        sigma_pooling="shared",
+
+        # Fourier parameters
+        fourier_K=2,
+        fourier_pooling="none",  # Unpooled Fourier coefficients
+
+        # covariates
+        x=[]
+    )
+
+    # Use subset of locations for faster testing
+    run_config = SimpleNamespace(
+        disease="flu",
+        ref_date=datetime.date.fromisoformat("2024-01-06"),
+        output_root=tmp_path / "model-output",
+        artifact_store_root=tmp_path / "artifact-store",
+        save_feat_importance=False,
+        locations=["US", "01", "02", "04", "05"],
+        max_horizon=2,  # Reduced for faster testing
+        q_levels=[0.025, 0.50, 0.975],
+        q_labels=["0.025", "0.5", "0.975"],
+        num_warmup=50,  # Reduced for faster testing
+        num_samples=50,
+        num_chains=1
+    )
+
+    model = SARIXFourierModel(model_config)
+    model.run(run_config)
+
+    # Verify output structure
+    actual_df = pd.read_csv(
+        run_config.output_root / "UMass-sarix_p2_fourier_K2_none" /
+        "2024-01-06-UMass-sarix_p2_fourier_K2_none.csv"
+    )
+
+    # Assertions
+    assert len(actual_df) > 0, "Output dataframe should not be empty"
+    assert set(actual_df["location"].unique()) == set(run_config.locations), \
+        "Output should contain predictions for all input locations"
+    assert all(actual_df["output_type"] == "quantile"), \
+        "All outputs should be quantiles"
+    assert set(actual_df["output_type_id"].astype(str).unique()) == set(run_config.q_labels), \
+        "Output should contain all specified quantile levels"
+    assert actual_df["value"].notna().all(), \
+        "All predictions should be non-null"
+    assert (actual_df["value"] >= 0).all(), \
+        "All predictions should be non-negative"
+
+
+def test_sarix_fourier_shared_pooling(tmp_path):
+    """Test SARIXFourierModel with fourier_pooling='shared' (pooled across locations)."""
+    model_config = SimpleNamespace(
+        model_class="sarix_fourier",
+        model_name="sarix_p2_fourier_K2_shared",
+
+        # data sources
+        sources=["nhsn"],
+
+        # fit locations separately or jointly
+        fit_locations_separately=False,
+
+        # SARIX model parameters
+        p=2,
+        P=0,
+        d=0,
+        D=0,
+        season_period=1,
+
+        # power transform
+        power_transform="4rt",
+
+        # parameter pooling
+        theta_pooling="shared",
+        sigma_pooling="shared",
+
+        # Fourier parameters
+        fourier_K=2,
+        fourier_pooling="shared",  # Shared Fourier coefficients
+
+        # covariates
+        x=[]
+    )
+
+    # Use subset of locations for faster testing
+    run_config = SimpleNamespace(
+        disease="flu",
+        ref_date=datetime.date.fromisoformat("2024-01-06"),
+        output_root=tmp_path / "model-output",
+        artifact_store_root=tmp_path / "artifact-store",
+        save_feat_importance=False,
+        locations=["US", "01", "02", "04", "05"],
+        max_horizon=2,
+        q_levels=[0.025, 0.50, 0.975],
+        q_labels=["0.025", "0.5", "0.975"],
+        num_warmup=50,
+        num_samples=50,
+        num_chains=1
+    )
+
+    model = SARIXFourierModel(model_config)
+    model.run(run_config)
+
+    # Verify output structure
+    actual_df = pd.read_csv(
+        run_config.output_root / "UMass-sarix_p2_fourier_K2_shared" /
+        "2024-01-06-UMass-sarix_p2_fourier_K2_shared.csv"
+    )
+
+    # Assertions
+    assert len(actual_df) > 0, "Output dataframe should not be empty"
+    assert set(actual_df["location"].unique()) == set(run_config.locations), \
+        "Output should contain predictions for all input locations"
+    assert all(actual_df["output_type"] == "quantile"), \
+        "All outputs should be quantiles"
+    assert set(actual_df["output_type_id"].astype(str).unique()) == set(run_config.q_labels), \
+        "Output should contain all specified quantile levels"
+    assert actual_df["value"].notna().all(), \
+        "All predictions should be non-null"
+    assert (actual_df["value"] >= 0).all(), \
+        "All predictions should be non-negative"
+
+
+def test_sarix_fourier_missing_pooling_parameter():
+    """Test that SARIXFourierModel raises error when fourier_pooling is missing."""
+    model_config = SimpleNamespace(
+        model_class="sarix_fourier",
+        model_name="sarix_p2_fourier_K2_nopooling",
+        sources=["nhsn"],
+        fit_locations_separately=False,
+        p=2, P=0, d=0, D=0, season_period=1,
+        power_transform="4rt",
+        theta_pooling="shared",
+        sigma_pooling="shared",
+        fourier_K=2,
+        # fourier_pooling is MISSING - should cause error
+        x=[]
+    )
+
+    run_config = SimpleNamespace(
+        disease="flu",
+        ref_date=datetime.date.fromisoformat("2024-01-06"),
+        output_root=Path("/tmp") / "model-output",
+        artifact_store_root=Path("/tmp") / "artifact-store",
+        save_feat_importance=False,
+        locations=["US"],
+        max_horizon=1,
+        q_levels=[0.5],
+        q_labels=["0.5"],
+        num_warmup=10,
+        num_samples=10,
+        num_chains=1
+    )
+
+    model = SARIXFourierModel(model_config)
+
+    # Should raise AttributeError when trying to access missing fourier_pooling
+    try:
+        model.run(run_config)
+        assert False, "Should have raised AttributeError for missing fourier_pooling"
+    except AttributeError as e:
+        assert "fourier_pooling" in str(e), \
+            f"Error should mention fourier_pooling, got: {str(e)}"
+
+
 def _np_percentile_val():
     return numpy.array(
         [[[2.22541624e-01, 1.82324940e-01, 1.27709944e-01],