Skip to content

Commit 61646c1

Browse files
committed
simplify doctest outputs
1 parent 70de921 commit 61646c1

File tree

1 file changed

+13
-17
lines changed

1 file changed

+13
-17
lines changed

causalpy/pymc_experiments.py

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -128,10 +128,10 @@ class PrePostFit(ExperimentalDesign):
128128
>>> result = cp.pymc_experiments.PrePostFit(
129129
... sc,
130130
... treatment_time,
131-
... formula="actual ~ 0 + a + b + c + d + e + f + g",
131+
... formula="actual ~ 0 + a + g",
132132
... model=cp.pymc_models.WeightedSumFitter(
133133
... sample_kwargs={
134-
... "draws": 2000,
134+
... "draws": 200,
135135
... "target_accept": 0.95,
136136
... "random_seed": seed,
137137
... "progressbar": False
@@ -140,16 +140,11 @@ class PrePostFit(ExperimentalDesign):
140140
... )
141141
>>> result.summary() # doctest: +NUMBER
142142
==================================Pre-Post Fit==================================
143-
Formula: actual ~ 0 + a + b + c + d + e + f + g
143+
Formula: actual ~ 0 + a + g
144144
Model coefficients:
145-
a 0.3, 94% HDI [0.3, 0.3]
146-
b 0.0, 94% HDI [0.0, 0.0]
147-
c 0.3, 94% HDI [0.2, 0.3]
148-
d 0.0, 94% HDI [0.0, 0.1]
149-
e 0.0, 94% HDI [0.0, 0.0]
150-
f 0.1, 94% HDI [0.1, 0.2]
151-
g 0.0, 94% HDI [0.0, 0.0]
152-
sigma 0.2, 94% HDI [0.2, 0.3]
145+
a 0.62, 94% HDI [0.61, 0.64]
146+
g 0.38, 94% HDI [0.36, 0.39]
147+
sigma 0.76, 94% HDI [0.64, 0.90]
153148
"""
154149

155150
def __init__(
@@ -771,13 +766,14 @@ class RegressionDiscontinuity(ExperimentalDesign):
771766
--------
772767
>>> import causalpy as cp
773768
>>> df = cp.load_data("rd")
769+
>>> df['y'] = df['y'] - 1 # added for doctest stability
774770
>>> seed = 42
775771
>>> result = cp.pymc_experiments.RegressionDiscontinuity(
776772
... df,
777773
... formula="y ~ 1 + x + treated + x:treated",
778774
... model=cp.pymc_models.LinearRegression(
779775
... sample_kwargs={
780-
... "draws": 2000,
776+
... "draws": 200,
781777
... "target_accept": 0.95,
782778
... "random_seed": seed,
783779
... "progressbar": False,
@@ -792,12 +788,12 @@ class RegressionDiscontinuity(ExperimentalDesign):
792788
Threshold on running variable: 0.5
793789
<BLANKLINE>
794790
Results:
795-
Discontinuity at threshold = 0.91
791+
Discontinuity at threshold = 0.92
796792
Model coefficients:
797-
Intercept 0.0, 94% HDI [0.0, 0.1]
798-
treated[T.True] 2.4, 94% HDI [1.6, 3.2]
799-
x 1.3, 94% HDI [1.1, 1.5]
800-
x:treated[T.True] -3.0, 94% HDI [-4.1, -2.0]
793+
Intercept -0.9, 94% HDI [-1.0, -0.8]
794+
treated[T.True] 2.4, 94% HDI [1.7, 3.2]
795+
x 1.3, 94% HDI [1.1, 1.4]
796+
x:treated[T.True] -3.1, 94% HDI [-4.1, -2.0]
801797
sigma 0.3, 94% HDI [0.3, 0.4]
802798
"""
803799

0 commit comments

Comments
 (0)