@@ -50,20 +50,20 @@ class PyMCModel(pm.Model):
5050 >>> X = rng.normal(loc=0, scale=1, size=(20, 2))
5151 >>> y = rng.normal(loc=0, scale=1, size=(20,))
5252 >>> model = MyToyModel(
53- ... sample_kwargs={
54- ... "chains": 2,
55- ... "draws": 2000,
56- ... "progressbar": False,
57- ... "random_seed": 42,
58- ... }
53+ ... sample_kwargs={
54+ ... "chains": 2,
55+ ... "draws": 2000,
56+ ... "progressbar": False,
57+ ... "random_seed": 42,
58+ ... }
5959 ... )
6060 >>> model.fit(X, y)
6161 Inference data...
6262 >>> model.score(X, y) # doctest: +ELLIPSIS
6363 r2 ...
6464 r2_std ...
6565 dtype: float64
66- >>> X_new = rng.normal(loc=0, scale=1, size=(20,2))
66+ >>> X_new = rng.normal(loc=0, scale=1, size=(20, 2))
6767 >>> model.predict(X_new)
6868 Inference data...
6969 """
@@ -286,11 +286,11 @@ class InstrumentalVariableRegression(PyMCModel):
286286 >>> ## Ensure the endogeneity of the the treatment variable
287287 >>> X = -1 + 4 * Z + e2 + 2 * e1
288288 >>> y = 2 + 3 * X + 3 * e1
289- >>> t = X.reshape(10,1)
290- >>> y = y.reshape(10,1)
291- >>> Z = np.asarray([[1, Z[i]] for i in range(0,10)])
292- >>> X = np.asarray([[1, X[i]] for i in range(0,10)])
293- >>> COORDS = {' instruments' : [' Intercept', 'Z' ], ' covariates' : [' Intercept', 'X' ]}
289+ >>> t = X.reshape(10, 1)
290+ >>> y = y.reshape(10, 1)
291+ >>> Z = np.asarray([[1, Z[i]] for i in range(0, 10)])
292+ >>> X = np.asarray([[1, X[i]] for i in range(0, 10)])
293+ >>> COORDS = {" instruments" : [" Intercept", "Z" ], " covariates" : [" Intercept", "X" ]}
294294 >>> sample_kwargs = {
295295 ... "tune": 5,
296296 ... "draws": 10,
@@ -300,12 +300,20 @@ class InstrumentalVariableRegression(PyMCModel):
300300 ... "progressbar": False,
301301 ... }
302302 >>> iv_reg = InstrumentalVariableRegression(sample_kwargs=sample_kwargs)
303- >>> iv_reg.fit(X, Z,y, t, COORDS, {
304- ... "mus": [[-2,4], [0.5, 3]],
305- ... "sigmas": [1, 1],
306- ... "eta": 2,
307- ... "lkj_sd": 1,
308- ... }, None)
303+ >>> iv_reg.fit(
304+ ... X,
305+ ... Z,
306+ ... y,
307+ ... t,
308+ ... COORDS,
309+ ... {
310+ ... "mus": [[-2, 4], [0.5, 3]],
311+ ... "sigmas": [1, 1],
312+ ... "eta": 2,
313+ ... "lkj_sd": 1,
314+ ... },
315+ ... None,
316+ ... )
309317 Inference data...
310318 """
311319
0 commit comments