fixing doctest

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

Author: NathanielF

causalpy/pymc_experiments.py

@@ -1499,12 +1499,12 @@ class InversePropensityWeighting(ExperimentalDesign, PropensityDataValidator):
     Example
     --------
     >>> import causalpy as cp
-    >>> df = cp.load_data("rd")
+    >>> df = cp.load_data("nhefs")
     >>> seed = 42
     >>> result = cp.pymc_experiments.InversePropensityWeighting(
     ...     df,
-    ...     formula="t ~ 1 + x",
-    ...     outcome_variable ="y",
+    ...     formula="trt ~ 1 + age + race",
+    ...     outcome_variable ="outcome",
     ...     weighting_scheme="robust",
     ...     model=cp.pymc_models.PropensityScore(
     ...         sample_kwargs={

causalpy/pymc_models.py

@@ -390,11 +390,11 @@ class PropensityScore(ModelBuilder):
     >>> import numpy as np
     >>> from causalpy.pymc_models import PropensityScore
     >>> df = cp.load_data('nhefs')
-    >>> X = df[["outcome", "trt", "age", "race"]]
-    >>> t = np.asarray(df["trt"]).reshape((df["trt"].shape[0],1))
+    >>> X = df[["age", "race"]]
+    >>> t = np.asarray(df["trt"])
     >>> ps = PropensityScore(sample_kwargs={"progressbar": False})
-    >>> ps.fit(X, y, coords={
-    ...                 'coeffs': ['trt', 'age', 'race'],
+    >>> ps.fit(X, t, coords={
+    ...                 'coeffs': ['age', 'race'],
     ...                 'obs_indx': np.arange(df.shape[0])
     ...                },
     ... )