docs(examples): repair runnable tail examples

Author: brycewang-stanford

src/statspai/bcf/bcf.py

@@ -81,12 +81,18 @@ def bcf(
 
     Examples
     --------
-    >>> import statspai as sp
-    >>> result = sp.bcf(df, y='outcome', treat='treatment',
-    ...                 covariates=['x1', 'x2', 'x3'])
+    >>> import statspai as sp, numpy as np, pandas as pd
+    >>> rng = np.random.default_rng(0)
+    >>> n = 400
+    >>> X = rng.normal(size=(n, 3))
+    >>> treatment = rng.binomial(1, 0.5, n)
+    >>> outcome = 2.0 * treatment + X @ np.array([1.0, -0.5, 0.3]) + rng.normal(size=n)
+    >>> df = pd.DataFrame(X, columns=["x1", "x2", "x3"])
+    >>> df["outcome"], df["treatment"] = outcome, treatment
+    >>> result = sp.bcf(df, y="outcome", treat="treatment",
+    ...                 covariates=["x1", "x2", "x3"], n_bootstrap=10)
     >>> print(result.summary())
-    >>> cate = result.model_info['cate']  # individual effects
-    >>> cate_sd = result.model_info['cate_sd']  # posterior SD
+    >>> cate = result.model_info["cate"]        # individual effects
     """
     est = BayesianCausalForest(
         data=data, y=y, treat=treat, covariates=covariates,

src/statspai/bunching/bunching.py

@@ -94,11 +94,12 @@ def bunching(
 
     Examples
     --------
-    >>> import statspai as sp
-    >>> result = sp.bunching(df, running_var='income',
-    ...                      threshold=50000, dt=0.10)
+    >>> import statspai as sp, numpy as np, pandas as pd
+    >>> rng = np.random.default_rng(0)
+    >>> income = rng.exponential(50000, 2000)   # taxable income, kink at 50k
+    >>> df = pd.DataFrame({"income": income})
+    >>> result = sp.bunching(df, running_var="income", threshold=50000)
     >>> print(result.summary())
-    >>> print(f"Elasticity: {result.model_info['elasticity']:.4f}")
     """
     est = BunchingEstimator(
         data=data, running_var=running_var, threshold=threshold,

src/statspai/causal_discovery/pcmci.py

@@ -275,9 +275,12 @@ def pcmci(
 
     Examples
     --------
-    >>> import statspai as sp
-    >>> res = sp.pcmci(df_ts, variables=['gdp', 'inflation', 'rates'],
-    ...                tau_max=4, pc_alpha=0.01)
+    >>> import statspai as sp, numpy as np, pandas as pd
+    >>> rng = np.random.default_rng(0)
+    >>> df_ts = pd.DataFrame(rng.normal(size=(120, 3)),
+    ...                      columns=["gdp", "inflation", "rates"])
+    >>> res = sp.pcmci(df_ts, variables=["gdp", "inflation", "rates"],
+    ...                tau_max=2, pc_alpha=0.1)
     >>> res.discovered_links()        # DataFrame of all significant links
     """
     if variables is None:

src/statspai/causal_text/text_treatment.py

@@ -165,14 +165,15 @@ def text_treatment_effect(
 
     Examples
     --------
-    >>> import statspai as sp, pandas as pd
+    >>> import statspai as sp, numpy as np, pandas as pd
+    >>> rng = np.random.default_rng(0)
     >>> df = pd.DataFrame({
-    ...     'text': ['great product', 'terrible bug', 'okay tool', 'love it'],
-    ...     'treatment': [1, 0, 0, 1],
-    ...     'outcome': [4.5, 1.2, 2.8, 4.7],
+    ...     "text": ["great product love it"] * 15 + ["terrible buggy slow"] * 15,
+    ...     "treatment": [1] * 15 + [0] * 15,
+    ...     "outcome": list(rng.normal(4.0, 0.3, 15)) + list(rng.normal(2.0, 0.3, 15)),
     ... })
-    >>> r = sp.text_treatment_effect(df, text_col='text', outcome='outcome',
-    ...                              treatment='treatment', n_components=4)
+    >>> r = sp.text_treatment_effect(df, text_col="text", outcome="outcome",
+    ...                              treatment="treatment", n_components=2)
     >>> r.estimate
     """
     for col in (text_col, outcome, treatment):

src/statspai/deepiv/deep_iv.py

@@ -169,15 +169,20 @@ def deepiv(
 
     Examples
     --------
-    >>> result = sp.deepiv(
+    Requires the optional ``deepiv`` extra (PyTorch); the calls below are
+    illustrative and are skipped by the doctest runner.
+
+    >>> import statspai as sp                              # doctest: +SKIP
+    >>> result = sp.deepiv(                                # doctest: +SKIP
     ...     df, y='lwage', treat='educ',
     ...     instruments=['nearc4'],
     ...     covariates=['exper', 'expersq'],
     ... )
-    >>> print(result.summary())
+    >>> print(result.summary())                           # doctest: +SKIP
+
+    Custom architecture with the unbiased paired-sample gradient:
 
-    >>> # Custom architecture with unbiased gradient
-    >>> result = sp.deepiv(
+    >>> result = sp.deepiv(                                # doctest: +SKIP
     ...     df, y='sales', treat='price',
     ...     instruments=['cost_shifter'],
     ...     covariates=['demand_controls'],

src/statspai/did/report.py

@@ -697,11 +697,18 @@ def cs_report(
 
     Examples
     --------
-    >>> import statspai as sp
-    >>> rpt = sp.did.cs_report(
-    ...     df, y='y', g='g', t='t', i='id', random_state=42)
+    >>> import statspai as sp, numpy as np, pandas as pd
+    >>> rng = np.random.default_rng(0)
+    >>> rows = []
+    >>> for unit in range(80):
+    ...     g = 4 if unit < 40 else 0             # cohort; 0 = never treated
+    ...     for t in range(8):
+    ...         post = g > 0 and t >= g
+    ...         y = 0.3 * t + (2.0 if post else 0.0) + (unit % 5) + rng.normal()
+    ...         rows.append((unit, t, g, y))
+    >>> df = pd.DataFrame(rows, columns=["id", "t", "g", "y"])
+    >>> rpt = sp.cs_report(df, y="y", g="g", t="t", i="id", random_state=42)
     >>> rpt.dynamic           # event-study DataFrame w/ uniform bands
-    >>> rpt.breakdown         # R-R breakdown M* per post event time
     """
     if isinstance(data_or_result, CausalResult):
         cs = data_or_result

src/statspai/dose_response/gps.py

@@ -93,10 +93,18 @@ def dose_response(
 
     Examples
     --------
-    >>> import statspai as sp
-    >>> result = sp.dose_response(df, y='outcome', treat='dosage',
-    ...                           covariates=['age', 'weight'])
-    >>> print(result.detail)  # dose-response curve
+    >>> import statspai as sp, numpy as np, pandas as pd
+    >>> rng = np.random.default_rng(0)
+    >>> n = 400
+    >>> age = rng.normal(40, 10, n)
+    >>> weight = rng.normal(70, 12, n)
+    >>> dosage = 0.5 * age + rng.normal(0, 5, n)      # confounded dose
+    >>> outcome = 0.2 * dosage + 0.1 * weight + rng.normal(0, 1, n)
+    >>> df = pd.DataFrame({"outcome": outcome, "dosage": dosage,
+    ...                    "age": age, "weight": weight})
+    >>> result = sp.dose_response(df, y="outcome", treat="dosage",
+    ...                           covariates=["age", "weight"])
+    >>> result.detail            # dose-response curve
     """
     est = DoseResponse(
         data=data, y=y, treat=treat, covariates=covariates,

src/statspai/matching/match.py

@@ -1149,9 +1149,13 @@ def psplot(
 
     Examples
     --------
-    >>> fig, ax = sp.psplot(df, treat='D', covariates=['x1', 'x2'])
-    >>> fig, ax = sp.psplot(df, treat='D', covariates=['x1', 'x2'],
-    ...                      trim=0.1)
+    >>> import statspai as sp, numpy as np, pandas as pd
+    >>> rng = np.random.default_rng(0)
+    >>> n = 400
+    >>> x1, x2 = rng.normal(size=n), rng.normal(size=n)
+    >>> D = rng.binomial(1, 1 / (1 + np.exp(-(x1 + 0.5 * x2))))
+    >>> df = pd.DataFrame({"D": D, "x1": x1, "x2": x2})
+    >>> fig, ax = sp.psplot(df, treat="D", covariates=["x1", "x2"])
     """
     try:
         import matplotlib.pyplot as plt

src/statspai/tmle/tmle.py

@@ -98,16 +98,17 @@ def tmle(
 
     Examples
     --------
-    >>> import statspai as sp
-    >>> result = sp.tmle(df, y='outcome', treat='treatment',
-    ...                  covariates=['x1', 'x2', 'x3'])
+    >>> import statspai as sp, numpy as np, pandas as pd
+    >>> rng = np.random.default_rng(0)
+    >>> n = 400
+    >>> X = rng.normal(size=(n, 3))
+    >>> treatment = rng.binomial(1, 1 / (1 + np.exp(-X[:, 0])))
+    >>> outcome = 2.0 * treatment + X @ np.array([1.0, -0.5, 0.3]) + rng.normal(size=n)
+    >>> df = pd.DataFrame(X, columns=["x1", "x2", "x3"])
+    >>> df["outcome"], df["treatment"] = outcome, treatment
+    >>> result = sp.tmle(df, y="outcome", treat="treatment",
+    ...                  covariates=["x1", "x2", "x3"])
     >>> print(result.summary())
-
-    >>> # Custom learner library
-    >>> from sklearn.ensemble import RandomForestRegressor
-    >>> result = sp.tmle(df, y='outcome', treat='treatment',
-    ...                  covariates=['x1', 'x2'],
-    ...                  outcome_library=[RandomForestRegressor()])
     """
     est = TMLE(
         data=data, y=y, treat=treat, covariates=covariates,