Skip to content

Commit 87a5002

Browse files
drbenvincentdrbenvincent
committed
fix doctests
1 parent b139865 commit 87a5002

File tree

1 file changed

+72
-64
lines changed

1 file changed

+72
-64
lines changed

‎causalpy/pymc_models.py‎

Lines changed: 72 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -977,11 +977,11 @@ class TransferFunctionLinearRegression(PyMCModel):
977977
\epsilon[t] = \rho \cdot \epsilon[t-1] + \nu[t]
978978
979979
But ``\epsilon[t] = y[t] - \mu[t]`` depends on ``\mu`` (which depends on ``\beta``,
980-
``\theta``, ``half_life``, etc.), so this fails:
980+
``\theta``, ``half_life``, etc.), so this fails::
981981
982-
>>> # This FAILS with TypeError
983-
>>> residuals = y - mu # depends on parameters!
984-
>>> pm.AR("epsilon", rho, observed=residuals) # ❌
982+
# This FAILS with TypeError
983+
residuals = y - mu # depends on parameters!
984+
pm.AR("epsilon", rho, observed=residuals) # ❌
985985
986986
*Implemented Solution (TransferFunctionARRegression):*
987987
@@ -1003,12 +1003,12 @@ class TransferFunctionLinearRegression(PyMCModel):
10031003
ar[t] = \rho \cdot ar[t-1] + \eta[t]
10041004
\epsilon_{obs}[t] \sim N(0, \sigma^2_{obs})
10051005
1006-
This could use PyMC's built-in ``pm.AR`` since ``ar[t]`` is latent (unobserved):
1006+
This could use PyMC's built-in ``pm.AR`` since ``ar[t]`` is latent (unobserved)::
10071007
1008-
>>> # This WOULD work (but changes model interpretation)
1009-
>>> ar = pm.AR("ar", rho=[rho_param], sigma=sigma_ar, shape=n_obs)
1010-
>>> mu_total = mu_baseline + ar
1011-
>>> pm.Normal("y", mu=mu_total, sigma=sigma_obs, observed=y_data) # ✅
1008+
# This WOULD work (but changes model interpretation)
1009+
ar = pm.AR("ar", rho=[rho_param], sigma=sigma_ar, shape=n_obs)
1010+
mu_total = mu_baseline + ar
1011+
pm.Normal("y", mu=mu_total, sigma=sigma_obs, observed=y_data) # ✅
10121012
10131013
**Why We Don't Use This Approach:**
10141014
@@ -1037,19 +1037,20 @@ class TransferFunctionLinearRegression(PyMCModel):
10371037
**Prior Customization**:
10381038
10391039
Priors are managed using the ``Prior`` class from ``pymc_extras`` and can be
1040-
customized via the ``priors`` parameter:
1041-
1042-
>>> from pymc_extras.prior import Prior
1043-
>>> model = cp.pymc_models.TransferFunctionLinearRegression(
1044-
... saturation_type=None,
1045-
... adstock_config={...},
1046-
... priors={
1047-
... "beta": Prior(
1048-
... "Normal", mu=0, sigma=100, dims=["treated_units", "coeffs"]
1049-
... ),
1050-
... "sigma": Prior("HalfNormal", sigma=50, dims=["treated_units"]),
1051-
... },
1052-
... )
1040+
customized via the ``priors`` parameter::
1041+
1042+
from pymc_extras.prior import Prior
1043+
1044+
model = cp.pymc_models.TransferFunctionLinearRegression(
1045+
saturation_type=None,
1046+
adstock_config={...},
1047+
priors={
1048+
"beta": Prior(
1049+
"Normal", mu=0, sigma=100, dims=["treated_units", "coeffs"]
1050+
),
1051+
"sigma": Prior("HalfNormal", sigma=50, dims=["treated_units"]),
1052+
},
1053+
)
10531054
10541055
By default, data-informed priors are set automatically via ``priors_from_data()``:
10551056
@@ -1061,16 +1062,19 @@ class TransferFunctionLinearRegression(PyMCModel):
10611062
10621063
Examples
10631064
--------
1064-
>>> import causalpy as cp
1065-
>>> model = cp.pymc_models.TransferFunctionLinearRegression(
1066-
... saturation_type=None,
1067-
... adstock_config={
1068-
... "half_life_prior": {"dist": "Gamma", "alpha": 4, "beta": 2},
1069-
... "l_max": 8,
1070-
... "normalize": True,
1071-
... },
1072-
... sample_kwargs={"chains": 4, "draws": 2000, "tune": 1000},
1073-
... )
1065+
Basic usage::
1066+
1067+
import causalpy as cp
1068+
1069+
model = cp.pymc_models.TransferFunctionLinearRegression(
1070+
saturation_type=None,
1071+
adstock_config={
1072+
"half_life_prior": {"dist": "Gamma", "alpha": 4, "beta": 2},
1073+
"l_max": 8,
1074+
"normalize": True,
1075+
},
1076+
sample_kwargs={"chains": 4, "draws": 2000, "tune": 1000},
1077+
)
10741078
"""
10751079

10761080
def __init__(
@@ -1518,21 +1522,22 @@ class TransferFunctionARRegression(PyMCModel):
15181522
**Prior Customization**:
15191523
15201524
Priors are managed using the ``Prior`` class from ``pymc_extras`` and can be
1521-
customized via the ``priors`` parameter:
1522-
1523-
>>> from pymc_extras.prior import Prior
1524-
>>> model = cp.pymc_models.TransferFunctionARRegression(
1525-
... saturation_type=None,
1526-
... adstock_config={...},
1527-
... priors={
1528-
... "beta": Prior(
1529-
... "Normal", mu=0, sigma=100, dims=["treated_units", "coeffs"]
1530-
... ),
1531-
... "rho": Prior(
1532-
... "Uniform", lower=-0.95, upper=0.95, dims=["treated_units"]
1533-
... ),
1534-
... },
1535-
... )
1525+
customized via the ``priors`` parameter::
1526+
1527+
from pymc_extras.prior import Prior
1528+
1529+
model = cp.pymc_models.TransferFunctionARRegression(
1530+
saturation_type=None,
1531+
adstock_config={...},
1532+
priors={
1533+
"beta": Prior(
1534+
"Normal", mu=0, sigma=100, dims=["treated_units", "coeffs"]
1535+
),
1536+
"rho": Prior(
1537+
"Uniform", lower=-0.95, upper=0.95, dims=["treated_units"]
1538+
),
1539+
},
1540+
)
15361541
15371542
By default, data-informed priors are set automatically via ``priors_from_data()``:
15381543
@@ -1545,23 +1550,26 @@ class TransferFunctionARRegression(PyMCModel):
15451550
15461551
Examples
15471552
--------
1548-
>>> import causalpy as cp
1549-
>>> model = cp.pymc_models.TransferFunctionARRegression(
1550-
... saturation_type=None,
1551-
... adstock_config={
1552-
... "half_life_prior": {"dist": "Gamma", "alpha": 4, "beta": 2},
1553-
... "l_max": 8,
1554-
... "normalize": True,
1555-
... },
1556-
... sample_kwargs={"chains": 4, "draws": 2000, "tune": 1000},
1557-
... )
1558-
>>> result = cp.GradedInterventionTimeSeries(
1559-
... data=df,
1560-
... y_column="outcome",
1561-
... treatment_names=["treatment"],
1562-
... base_formula="1 + time + covariate",
1563-
... model=model,
1564-
... )
1553+
Basic usage::
1554+
1555+
import causalpy as cp
1556+
1557+
model = cp.pymc_models.TransferFunctionARRegression(
1558+
saturation_type=None,
1559+
adstock_config={
1560+
"half_life_prior": {"dist": "Gamma", "alpha": 4, "beta": 2},
1561+
"l_max": 8,
1562+
"normalize": True,
1563+
},
1564+
sample_kwargs={"chains": 4, "draws": 2000, "tune": 1000},
1565+
)
1566+
result = cp.GradedInterventionTimeSeries(
1567+
data=df,
1568+
y_column="outcome",
1569+
treatment_names=["treatment"],
1570+
base_formula="1 + time + covariate",
1571+
model=model,
1572+
)
15651573
15661574
References
15671575
----------

0 commit comments

Comments
 (0)