Add/test combined arg, revert method order

Author: Max Joseph

pymc_experimental/model_builder.py

@@ -245,13 +245,13 @@ def fit(
         self.idata.add_groups(fit_data=self.data.to_xarray())
         return self.idata
 
-    def predict_posterior(
+    def predict(
         self,
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
         extend_idata: bool = True,
     ) -> xr.Dataset:
         """
-        Generate posterior predictive samples on unseen data.
+        Uses model to predict on unseen data and return point prediction of all the samples
 
         Parameters
         ---------
@@ -262,61 +262,66 @@ def predict_posterior(
 
         Returns
         -------
-        returns posterior predictive samples
+        returns posterior mean of predictive samples
 
         Examples
         --------
         >>> data, model_config, sampler_config = LinearModel.create_sample_input()
         >>> model = LinearModel(model_config, sampler_config)
         >>> idata = model.fit(data)
         >>> x_pred = []
-        >>> prediction_data = pd.DataFrame({'input': x_pred})
-        >>> pred_samples = model.predict_posterior(prediction_data)
+        >>> prediction_data = pd.DataFrame({'input':x_pred})
+        >>> pred_mean = model.predict(prediction_data)
         """
+        posterior_predictive_samples = self.predict_posterior(data_prediction, extend_idata)
+        posterior_means = posterior_predictive_samples.mean(dim=["chain", "draw"], keep_attrs=True)
+        return posterior_means
 
-        if data_prediction is not None:  # set new input data
-            self._data_setter(data_prediction)
-
-        with self.model:  # sample with new input data
-            post_pred = pm.sample_posterior_predictive(self.idata)
-            if extend_idata:
-                self.idata.extend(post_pred)
-
-        posterior_predictive_samples = az.extract(post_pred, "posterior_predictive", combined=False)
-
-        return posterior_predictive_samples
-
-    def predict(
+    def predict_posterior(
         self,
         data_prediction: Dict[str, Union[np.ndarray, pd.DataFrame, pd.Series]] = None,
         extend_idata: bool = True,
+        combined: bool = False,
     ) -> xr.Dataset:
         """
-        Uses model to predict on unseen data and return point prediction of all the samples
+        Generate posterior predictive samples on unseen data.
 
         Parameters
         ---------
         data_prediction : Dictionary of string and either of numpy array, pandas dataframe or pandas Series
             It is the data we need to make prediction on using the model.
         extend_idata : Boolean determining whether the predictions should be added to inference data object.
             Defaults to True.
+        combined: Combine chain and draw dims into sample. Won’t work if a dim named sample already exists.
+            Defaults to False.
 
         Returns
         -------
-        returns posterior mean of predictive samples
+        returns posterior predictive samples
 
         Examples
         --------
         >>> data, model_config, sampler_config = LinearModel.create_sample_input()
         >>> model = LinearModel(model_config, sampler_config)
         >>> idata = model.fit(data)
         >>> x_pred = []
-        >>> prediction_data = pd.DataFrame({'input':x_pred})
-        >>> pred_mean = model.predict(prediction_data)
+        >>> prediction_data = pd.DataFrame({'input': x_pred})
+        >>> pred_samples = model.predict_posterior(prediction_data)
         """
-        posterior_predictive_samples = self.predict_posterior(data_prediction, extend_idata)
-        posterior_means = posterior_predictive_samples.mean(dim=["chain", "draw"], keep_attrs=True)
-        return posterior_means
+
+        if data_prediction is not None:  # set new input data
+            self._data_setter(data_prediction)
+
+        with self.model:  # sample with new input data
+            post_pred = pm.sample_posterior_predictive(self.idata)
+            if extend_idata:
+                self.idata.extend(post_pred)
+
+        posterior_predictive_samples = az.extract(
+            post_pred, "posterior_predictive", combined=combined
+        )
+
+        return posterior_predictive_samples
 
     @property
     def id(self) -> str:

pymc_experimental/tests/test_model_builder.py

@@ -151,16 +151,17 @@ def test_predict():
     assert np.issubdtype(pred["y_model"].dtype, np.floating)
 
 
-def test_predict_posterior():
+@pytest.mark.parametrize("combined", [True, False])
+def test_predict_posterior(combined):
     model = test_ModelBuilder.initial_build_and_fit()
     n_pred = 100
     x_pred = np.random.uniform(low=0, high=1, size=n_pred)
     prediction_data = pd.DataFrame({"input": x_pred})
-    pred = model.predict_posterior(prediction_data)
+    pred = model.predict_posterior(prediction_data, combined=combined)
     chains = model.idata.sample_stats.dims["chain"]
     draws = model.idata.sample_stats.dims["draw"]
-    assert "y_model" in pred
-    assert pred["y_model"].shape == (chains, draws, n_pred)
+    expected_shape = (n_pred, chains * draws) if combined else (chains, draws, n_pred)
+    assert pred["y_model"].shape == expected_shape
     assert np.issubdtype(pred["y_model"].dtype, np.floating)