Skip to content

Commit 5c2870e

Browse files
jpreszlerjpreszler
committed
dcotest on pymc_models good except rng caused score differences
1 parent adb94c5 commit 5c2870e

File tree

1 file changed

+82
-87
lines changed

1 file changed

+82
-87
lines changed

‎causalpy/pymc_models.py

Lines changed: 82 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -48,15 +48,17 @@ def build_model(self, X, y, coords) -> None:
4848
4949
Example
5050
-------
51+
>>> import pymc as pm
52+
>>> from causalpy.pymc_models import ModelBuilder
5153
>>> class CausalPyModel(ModelBuilder):
52-
>>> def build_model(self, X, y):
53-
>>> with self:
54-
>>> X_ = pm.MutableData(name="X", value=X)
55-
>>> y_ = pm.MutableData(name="y", value=y)
56-
>>> beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1])
57-
>>> sigma = pm.HalfNormal("sigma", sigma=1)
58-
>>> mu = pm.Deterministic("mu", pm.math.dot(X_, beta))
59-
>>> pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_)
54+
... def build_model(self, X, y):
55+
... with self:
56+
... X_ = pm.MutableData(name="X", value=X)
57+
... y_ = pm.MutableData(name="y", value=y)
58+
... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1])
59+
... sigma = pm.HalfNormal("sigma", sigma=1)
60+
... mu = pm.Deterministic("mu", pm.math.dot(X_, beta))
61+
... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_)
6062
"""
6163
raise NotImplementedError("This method must be implemented by a subclass")
6264

@@ -83,37 +85,22 @@ def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None:
8385
>>> import pymc as pm
8486
>>> from causalpy.pymc_models import ModelBuilder
8587
>>> class MyToyModel(ModelBuilder):
86-
... def build_model(self, X, y, coords):
87-
... with self:
88-
... X_ = pm.MutableData(name="X", value=X)
89-
... y_ = pm.MutableData(name="y", value=y)
90-
... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1])
91-
... sigma = pm.HalfNormal("sigma", sigma=1)
92-
... mu = pm.Deterministic("mu", pm.math.dot(X_, beta))
93-
... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_)
88+
... def build_model(self, X, y, coords):
89+
... with self:
90+
... X_ = pm.MutableData(name="X", value=X)
91+
... y_ = pm.MutableData(name="y", value=y)
92+
... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1])
93+
... sigma = pm.HalfNormal("sigma", sigma=1)
94+
... mu = pm.Deterministic("mu", pm.math.dot(X_, beta))
95+
... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_)
9496
>>> rng = np.random.default_rng(seed=42)
9597
>>> X = rng.normal(loc=0, scale=1, size=(20, 2))
9698
>>> y = rng.normal(loc=0, scale=1, size=(20,))
97-
>>> model = MyToyModel(sample_kwargs={"chains": 2, "draws": 2})
98-
>>> model.fit(X, y)
99-
Only 2 samples in chain.
100-
Auto-assigning NUTS sampler...
101-
Initializing NUTS using jitter+adapt_diag...
102-
Multiprocess sampling (2 chains in 4 jobs)
103-
NUTS: [beta, sigma]
104-
Sampling 2 chains for 1_000 tune and 2 draw iterations (2_000 + 4 draws total)
105-
took 0 seconds.gences]
106-
The number of samples is too small to check convergence reliably.
107-
Sampling: [beta, sigma, y_hat]
108-
Sampling: [y_hat]
109-
Inference data with groups:
110-
> posterior
111-
> posterior_predictive
112-
> sample_stats
113-
> prior
114-
> prior_predictive
115-
> observed_data
116-
> constant_data
99+
>>> model = MyToyModel(
100+
... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False}
101+
... )
102+
>>> model.fit(X, y) # doctest: +ELLIPSIS
103+
Inference ...
117104
"""
118105
self.build_model(X, y, coords)
119106
with self.model:
@@ -133,16 +120,30 @@ def predict(self, X):
133120
134121
Example
135122
-------
136-
Assumes `model` has been initialized and .fit() has been run,
137-
see ModelBuilder().fit() for example.
138-
123+
>>> import causalpy as cp
124+
>>> import numpy as np
125+
>>> import pymc as pm
126+
>>> from causalpy.pymc_models import ModelBuilder
127+
>>> class MyToyModel(ModelBuilder):
128+
... def build_model(self, X, y, coords):
129+
... with self:
130+
... X_ = pm.MutableData(name="X", value=X)
131+
... y_ = pm.MutableData(name="y", value=y)
132+
... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1])
133+
... sigma = pm.HalfNormal("sigma", sigma=1)
134+
... mu = pm.Deterministic("mu", pm.math.dot(X_, beta))
135+
... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_)
136+
>>> rng = np.random.default_rng(seed=42)
137+
>>> X = rng.normal(loc=0, scale=1, size=(20, 2))
138+
>>> y = rng.normal(loc=0, scale=1, size=(20,))
139+
>>> model = MyToyModel(
140+
... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False}
141+
... )
142+
>>> model.fit(X, y) # doctest: +ELLIPSIS
143+
Inference...
139144
>>> X_new = rng.normal(loc=0, scale=1, size=(20,2))
140-
>>> model.predict(X_new)
141-
Sampling: [beta, y_hat]
142-
Inference data with groups:
143-
> posterior_predictive
144-
> observed_data
145-
> constant_data
145+
>>> model.predict(X_new) # doctest: +ELLIPSIS
146+
Inference...
146147
"""
147148

148149
self._data_setter(X)
@@ -162,9 +163,28 @@ def score(self, X, y) -> pd.Series:
162163
163164
Example
164165
--------
165-
Assuming `model` has been fit
166-
167-
>>> model.score(X, y) # X, y are random data here
166+
>>> import causalpy as cp
167+
>>> import numpy as np
168+
>>> import pymc as pm
169+
>>> from causalpy.pymc_models import ModelBuilder
170+
>>> class MyToyModel(ModelBuilder):
171+
... def build_model(self, X, y, coords):
172+
... with self:
173+
... X_ = pm.MutableData(name="X", value=X)
174+
... y_ = pm.MutableData(name="y", value=y)
175+
... beta = pm.Normal("beta", mu=0, sigma=1, shape=X_.shape[1])
176+
... sigma = pm.HalfNormal("sigma", sigma=1)
177+
... mu = pm.Deterministic("mu", pm.math.dot(X_, beta))
178+
... pm.Normal("y_hat", mu=mu, sigma=sigma, observed=y_)
179+
>>> rng = np.random.default_rng(seed=42)
180+
>>> X = rng.normal(loc=0, scale=1, size=(20, 2))
181+
>>> y = rng.normal(loc=0, scale=1, size=(20,))
182+
>>> model = MyToyModel(
183+
... sample_kwargs={"chains": 2, "draws": 2, "progressbar": False}
184+
... )
185+
>>> model.fit(X, y) # doctest: +ELLIPSIS
186+
Inference...
187+
>>> model.score(X, y)
168188
Sampling: [y_hat]
169189
r2 0.352251
170190
r2_std 0.051624
@@ -196,27 +216,14 @@ class WeightedSumFitter(ModelBuilder):
196216
197217
Example
198218
--------
219+
>>> import causalpy as cp
220+
>>> import numpy as np
221+
>>> from causalpy.pymc_models import WeightedSumFitter
199222
>>> sc = cp.load_data("sc")
200223
>>> X = sc[['a', 'b', 'c', 'd', 'e', 'f', 'g']]
201224
>>> y = np.asarray(sc['actual']).reshape((sc.shape[0], 1))
202-
>>> wsf = WeightedSumFitter()
203-
>>> wsf.fit(X,y)
204-
Auto-assigning NUTS sampler...
205-
Initializing NUTS using jitter+adapt_diag...
206-
Multiprocess sampling (4 chains in 4 jobs)
207-
NUTS: [beta, sigma]
208-
Sampling 4 chains for 1_000 tune and 1_000 draw iterations
209-
(4_000 + 4_000 draws total) took 3 seconds.
210-
Sampling: [beta, sigma, y_hat]
211-
Sampling: [y_hat]
212-
Inference data with groups:
213-
> posterior
214-
> posterior_predictive
215-
> sample_stats
216-
> prior
217-
> prior_predictive
218-
> observed_data
219-
> constant_data
225+
>>> wsf = WeightedSumFitter(sample_kwargs={"progressbar": False})
226+
>>> _ = wsf.fit(X,y)
220227
"""
221228

222229
def build_model(self, X, y, coords):
@@ -261,31 +268,19 @@ class LinearRegression(ModelBuilder):
261268
262269
Example
263270
--------
271+
>>> import causalpy as cp
272+
>>> import numpy as np
273+
>>> from causalpy.pymc_models import LinearRegression
264274
>>> rd = cp.load_data("rd")
265275
>>> X = rd[["x", "treated"]]
266276
>>> y = np.asarray(rd["y"]).reshape((rd["y"].shape[0],1))
267-
>>> lr = LinearRegression()
277+
>>> lr = LinearRegression(sample_kwargs={"progressbar": False})
268278
>>> lr.fit(X, y, coords={
269-
'coeffs': ['x', 'treated'],
270-
'obs_indx': np.arange(rd.shape[0])
271-
}
272-
)
273-
Auto-assigning NUTS sampler...
274-
Initializing NUTS using jitter+adapt_diag...
275-
Multiprocess sampling (4 chains in 4 jobs)
276-
NUTS: [beta, sigma]
277-
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (
278-
4_000 + 4_000 draws total) took 1 seconds.
279-
Sampling: [beta, sigma, y_hat]
280-
Sampling: [y_hat]
281-
Inference data with groups:
282-
> posterior
283-
> posterior_predictive
284-
> sample_stats
285-
> prior
286-
> prior_predictive
287-
> observed_data
288-
> constant_data
279+
... 'coeffs': ['x', 'treated'],
280+
... 'obs_indx': np.arange(rd.shape[0])
281+
... },
282+
... ) # doctest: +ELLIPSIS
283+
Inference...
289284
"""
290285

291286
def build_model(self, X, y, coords):

0 commit comments

Comments
 (0)