making y optional, adding fitted model fixture to make tests more efficient

Author: michaelraczycki

pymc_experimental/model_builder.py

@@ -424,7 +424,7 @@ def load(cls, fname: str):
     def fit(
         self,
         X: pd.DataFrame,
-        y: pd.Series,
+        y: Optional[pd.Series] = None,
         progressbar: bool = True,
         predictor_names: List[str] = None,
         random_seed: RandomState = None,
@@ -464,7 +464,8 @@ def fit(
         """
         if predictor_names is None:
             predictor_names = []
-
+        if y is None:
+            y = np.zeros(X.shape[0])
         y = pd.DataFrame({self.output_var: y})
         self.generate_and_preprocess_model_data(X, y.values.flatten())
         self.build_model(self.X, self.y)

pymc_experimental/tests/test_model_builder.py

@@ -40,6 +40,24 @@ def toy_y(toy_X):
     return y
 
 
+@pytest.fixture(scope="module")
+def fitted_model_instance(toy_X, toy_y):
+    sampler_config = {
+        "draws": 500,
+        "tune": 300,
+        "chains": 2,
+        "target_accept": 0.95,
+    }
+    model_config = {
+        "a": {"loc": 0, "scale": 10},
+        "b": {"loc": 0, "scale": 10},
+        "obs_error": 2,
+    }
+    model = test_ModelBuilder(model_config=model_config, sampler_config=sampler_config)
+    model.fit(toy_X)
+    return model
+
+
 class test_ModelBuilder(ModelBuilder):
 
     _model_type = "LinearModel"
@@ -103,16 +121,11 @@ def default_sampler_config(self) -> Dict:
             "target_accept": 0.95,
         }
 
-    @staticmethod
-    def initial_build_and_fit(toy_X, toy_y, check_idata=True) -> ModelBuilder:
-        model_builder = test_ModelBuilder()
-        model_builder.idata = model_builder.fit(
-            toy_X, toy_y, predictor_names=["input"], random_seed=1234
-        )
-        if check_idata:
-            assert model_builder.idata is not None
-            assert "posterior" in model_builder.idata.groups()
-        return model_builder
+
+def test_initial_build_and_fit(fitted_model_instance, check_idata=True) -> ModelBuilder:
+    if check_idata:
+        assert fitted_model_instance.idata is not None
+        assert "posterior" in fitted_model_instance.idata.groups()
 
 
 def test_save_without_fit_raises_runtime_error():
@@ -129,59 +142,62 @@ def test_empty_sampler_config_fit(toy_X, toy_y):
     assert "posterior" in model_builder.idata.groups()
 
 
-def test_fit(toy_X, toy_y):
-    model_builder = test_ModelBuilder.initial_build_and_fit(toy_X, toy_y)
-    x_pred = np.random.uniform(low=0, high=1, size=100)
-    prediction_data = pd.DataFrame({"input": x_pred})
-    pred = model_builder.predict(prediction_data["input"])
-    post_pred = model_builder.sample_posterior_predictive(
+def test_fit(fitted_model_instance):
+    prediction_data = pd.DataFrame({"input": np.random.uniform(low=0, high=1, size=100)})
+    pred = fitted_model_instance.predict(prediction_data["input"])
+    post_pred = fitted_model_instance.sample_posterior_predictive(
         prediction_data["input"], extend_idata=True, combined=True
     )
-    post_pred[model_builder.output_var].shape[0] == prediction_data.input.shape
+    post_pred[fitted_model_instance.output_var].shape[0] == prediction_data.input.shape
+
+
+def test_fit_no_y(toy_X):
+    model_builder = test_ModelBuilder()
+    model_builder.idata = model_builder.fit(X=toy_X)
+    assert model_builder.model is not None
+    assert model_builder.idata is not None
+    assert "posterior" in model_builder.idata.groups()
 
 
 @pytest.mark.skipif(
     sys.platform == "win32", reason="Permissions for temp files not granted on windows CI."
 )
-def test_save_load(toy_X, toy_y):
-    test_builder = test_ModelBuilder.initial_build_and_fit(toy_X, toy_y)
+def test_save_load(fitted_model_instance):
     temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
-    test_builder.save(temp.name)
+    fitted_model_instance.save(temp.name)
     test_builder2 = test_ModelBuilder.load(temp.name)
-    assert test_builder.idata.groups() == test_builder2.idata.groups()
+    assert fitted_model_instance.idata.groups() == test_builder2.idata.groups()
 
     x_pred = np.random.uniform(low=0, high=1, size=100)
     prediction_data = pd.DataFrame({"input": x_pred})
-    pred1 = test_builder.predict(prediction_data["input"])
+    pred1 = fitted_model_instance.predict(prediction_data["input"])
     pred2 = test_builder2.predict(prediction_data["input"])
     assert pred1.shape == pred2.shape
     temp.close()
 
 
-def test_predict(toy_X, toy_y):
-    model_builder = test_ModelBuilder.initial_build_and_fit(toy_X, toy_y)
+def test_predict(fitted_model_instance):
     x_pred = np.random.uniform(low=0, high=1, size=100)
     prediction_data = pd.DataFrame({"input": x_pred})
-    pred = model_builder.predict(prediction_data["input"])
+    pred = fitted_model_instance.predict(prediction_data["input"])
     # Perform elementwise comparison using numpy
     assert type(pred) == np.ndarray
     assert len(pred) > 0
 
 
 @pytest.mark.parametrize("combined", [True, False])
-def test_sample_posterior_predictive(toy_X, toy_y, combined):
-    model_builder = test_ModelBuilder.initial_build_and_fit(toy_X, toy_y)
+def test_sample_posterior_predictive(fitted_model_instance, combined):
     n_pred = 100
     x_pred = np.random.uniform(low=0, high=1, size=n_pred)
     prediction_data = pd.DataFrame({"input": x_pred})
-    pred = model_builder.sample_posterior_predictive(
+    pred = fitted_model_instance.sample_posterior_predictive(
         prediction_data["input"], combined=combined, extend_idata=True
     )
-    chains = model_builder.idata.sample_stats.dims["chain"]
-    draws = model_builder.idata.sample_stats.dims["draw"]
+    chains = fitted_model_instance.idata.sample_stats.dims["chain"]
+    draws = fitted_model_instance.idata.sample_stats.dims["draw"]
     expected_shape = (n_pred, chains * draws) if combined else (chains, draws, n_pred)
-    assert pred[model_builder.output_var].shape == expected_shape
-    assert np.issubdtype(pred[model_builder.output_var].dtype, np.floating)
+    assert pred[fitted_model_instance.output_var].shape == expected_shape
+    assert np.issubdtype(pred[fitted_model_instance.output_var].dtype, np.floating)
 
 
 def test_id():