adding some tests and a conclusion

Signed-off-by: Nathaniel <[email protected]>

Author: NathanielF

causalpy/pymc_models.py

@@ -527,6 +527,69 @@ def fit_outcome_model(
         normal_outcome=True,
         spline_component=False,
     ):
+        """
+        Fit a Bayesian outcome model using covariates and previously estimated propensity scores.
+
+        This function implements the second stage of a modular two-step causal inference procedure.
+        It uses propensity scores extracted from a prior treatment model (via `self.fit()`) to adjust
+        for confounding when estimating treatment effects on an outcome variable `y`.
+
+        Parameters
+        ----------
+        X_outcome : array-like, shape (n_samples, n_covariates)
+            Covariate matrix for the outcome model.
+
+        y : array-like, shape (n_samples,)
+            Observed outcome variable.
+
+        coords : dict
+            Coordinate dictionary for named dimensions in the PyMC model. Should include
+            a key "outcome_coeffs" for `X_outcome`.
+
+        priors : dict, optional
+            Dictionary specifying priors for outcome model parameters:
+                - "b_outcome": list [mean, std] for regression coefficients.
+                - "a_outcome": list [mean, std] for the intercept.
+                - "sigma": standard deviation of the outcome noise (default 1).
+
+        noncentred : bool, default True
+            If True, use a non-centred parameterization for the outcome coefficients.
+
+        normal_outcome : bool, default True
+            If True, assume a Normal likelihood for the outcome.
+            If False, use a Student-t likelihood with unknown degrees of freedom.
+
+        spline_component : bool, default False
+            If True, include a spline basis expansion on the propensity score to allow
+            flexible (nonlinear) adjustment. Uses B-splines with 30 internal knots.
+
+        Returns
+        -------
+        idata_outcome : arviz.InferenceData
+            The posterior and prior predictive samples from the outcome model.
+
+        model_outcome : pm.Model
+            The PyMC model object.
+
+        Raises
+        ------
+        AttributeError
+            If the `self.idata` attribute is not available, which indicates that
+            `fit()` (i.e., the treatment model) has not been called yet.
+
+        Notes
+        -----
+        - This model uses a sampled version of the propensity score (`p`) from the
+        posterior of the treatment model, randomly selecting one posterior draw
+        per call.
+        - The term `beta_ps[0] * p + beta_ps[1] * (p * treatment)` captures both
+        main and interaction effects of the propensity score.
+        - Including spline adjustment enables modeling nonlinear relationships
+        between the propensity score and the outcome.
+        - Compatible with IPW-style estimation when combined with weighted loss or
+        diagnostics outside this function.
+
+        """
         if not hasattr(self, "idata"):
             raise AttributeError("""Object is missing required attribute 'idata'
                                  so cannot proceed. Call fit() first""")
@@ -551,7 +614,6 @@ def fit_outcome_model(
                     dims="outcome_coeffs",
                 )
 
-            beta_ps_spline = pm.Normal("beta_ps_spline", 0, 1, size=34)
             beta_ps = pm.Normal("beta_ps", 0, 1, size=2)
 
             chosen = np.random.choice(range(propensity_scores.shape[1]))

causalpy/tests/test_integration_pymc_examples.py

@@ -14,6 +14,7 @@
 import arviz as az
 import numpy as np
 import pandas as pd
+import pymc as pm
 import pytest
 from matplotlib import pyplot as plt
 
@@ -723,5 +724,29 @@ def test_inverse_prop(mock_pymc_sample):
     assert isinstance(fig, plt.Figure)
     assert isinstance(axs, list)
     assert all(isinstance(ax, plt.Axes) for ax in axs)
+    plt.close()
     with pytest.raises(NotImplementedError):
         result.get_plot_data()
+
+    ### testing outcome model
+    idata_normal, model_normal = result.model.fit_outcome_model(
+        X_outcome=result.X_outcome,
+        y=result.y,
+        coords=result.coords,
+        normal_outcome=True,
+        spline_component=False,
+    )
+    assert isinstance(idata_normal, az.InferenceData)
+    assert isinstance(model_normal, pm.Model)
+    assert "beta_" in idata_normal.posterior
+    assert "beta_ps" in idata_normal.posterior
+
+    # Test spline model
+    idata_spline, _ = result.model.fit_outcome_model(
+        X_outcome=result.X_outcome,
+        y=result.y,
+        coords=result.coords,
+        normal_outcome=True,
+        spline_component=True,
+    )
+    assert "spline_features" in idata_spline.posterior