diff --git a/causalpy/data/simulate_data.py b/causalpy/data/simulate_data.py index 0dfff360..fe3af1d3 100644 --- a/causalpy/data/simulate_data.py +++ b/causalpy/data/simulate_data.py @@ -71,9 +71,7 @@ def generate_synthetic_control_data( Example -------- >>> from causalpy.data.simulate_data import generate_synthetic_control_data - >>> df, weightings_true = generate_synthetic_control_data( - ... treatment_time=70 - ... ) + >>> df, weightings_true = generate_synthetic_control_data(treatment_time=70) """ # 1. Generate non-treated variables @@ -267,8 +265,9 @@ def generate_regression_discontinuity_data( >>> import pathlib >>> from causalpy.data.simulate_data import generate_regression_discontinuity_data >>> df = generate_regression_discontinuity_data(true_treatment_threshold=0.5) - >>> df.to_csv(pathlib.Path.cwd() / 'regression_discontinuity.csv', - ... index=False) # doctest: +SKIP + >>> df.to_csv( + ... pathlib.Path.cwd() / "regression_discontinuity.csv", index=False + ... ) # doctest: +SKIP """ def is_treated(x): @@ -298,13 +297,9 @@ def generate_ancova_data( >>> import pathlib >>> from causalpy.data.simulate_data import generate_ancova_data >>> df = generate_ancova_data( - ... N=200, - ... pre_treatment_means=np.array([10, 12]), - ... treatment_effect=2, - ... sigma=1 + ... N=200, pre_treatment_means=np.array([10, 12]), treatment_effect=2, sigma=1 ... ) - >>> df.to_csv(pathlib.Path.cwd() / 'ancova_data.csv', - ... index=False) # doctest: +SKIP + >>> df.to_csv(pathlib.Path.cwd() / "ancova_data.csv", index=False) # doctest: +SKIP """ group = np.random.choice(2, size=N) pre = np.random.normal(loc=pre_treatment_means[group]) diff --git a/causalpy/experiments/diff_in_diff.py b/causalpy/experiments/diff_in_diff.py index 26fb8112..daae8e50 100644 --- a/causalpy/experiments/diff_in_diff.py +++ b/causalpy/experiments/diff_in_diff.py @@ -70,8 +70,8 @@ class DifferenceInDifferences(BaseExperiment): ... "random_seed": seed, ... "progressbar": False, ... } - ... ) - ... ) + ... ), + ... ) """ supports_ols = True diff --git a/causalpy/experiments/instrumental_variable.py b/causalpy/experiments/instrumental_variable.py index a272ad3e..bb0adfc7 100644 --- a/causalpy/experiments/instrumental_variable.py +++ b/causalpy/experiments/instrumental_variable.py @@ -74,17 +74,17 @@ class InstrumentalVariable(BaseExperiment): ... "cores": 4, ... "target_accept": 0.95, ... "progressbar": False, - ... } + ... } >>> instruments_formula = "X ~ 1 + Z" >>> formula = "y ~ 1 + X" >>> instruments_data = test_data[["X", "Z"]] >>> data = test_data[["y", "X"]] >>> iv = cp.InstrumentalVariable( - ... instruments_data=instruments_data, - ... data=data, - ... instruments_formula=instruments_formula, - ... formula=formula, - ... model=InstrumentalVariableRegression(sample_kwargs=sample_kwargs), + ... instruments_data=instruments_data, + ... data=data, + ... instruments_formula=instruments_formula, + ... formula=formula, + ... model=InstrumentalVariableRegression(sample_kwargs=sample_kwargs), ... ) """ diff --git a/causalpy/experiments/inverse_propensity_weighting.py b/causalpy/experiments/inverse_propensity_weighting.py index 3e87046c..13f74767 100644 --- a/causalpy/experiments/inverse_propensity_weighting.py +++ b/causalpy/experiments/inverse_propensity_weighting.py @@ -56,7 +56,7 @@ class InversePropensityWeighting(BaseExperiment): >>> result = cp.InversePropensityWeighting( ... df, ... formula="trt ~ 1 + age + race", - ... outcome_variable ="outcome", + ... outcome_variable="outcome", ... weighting_scheme="robust", ... model=cp.pymc_models.PropensityScore( ... sample_kwargs={ diff --git a/causalpy/experiments/prepostfit.py b/causalpy/experiments/prepostfit.py index 127a0799..720d001a 100644 --- a/causalpy/experiments/prepostfit.py +++ b/causalpy/experiments/prepostfit.py @@ -337,7 +337,7 @@ class InterruptedTimeSeries(PrePostFit): ... "random_seed": seed, ... "progressbar": False, ... } - ... ) + ... ), ... ) """ diff --git a/causalpy/experiments/prepostnegd.py b/causalpy/experiments/prepostnegd.py index def51ed2..9cecc49b 100644 --- a/causalpy/experiments/prepostnegd.py +++ b/causalpy/experiments/prepostnegd.py @@ -69,9 +69,9 @@ class PrePostNEGD(BaseExperiment): ... "random_seed": seed, ... "progressbar": False, ... } - ... ) + ... ), ... ) - >>> result.summary(round_to=1) # doctest: +NUMBER + >>> result.summary(round_to=1) # doctest: +NUMBER ==================Pretest/posttest Nonequivalent Group Design=================== Formula: post ~ 1 + C(group) + pre diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index 32470dd9..032da928 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -50,12 +50,12 @@ class PyMCModel(pm.Model): >>> X = rng.normal(loc=0, scale=1, size=(20, 2)) >>> y = rng.normal(loc=0, scale=1, size=(20,)) >>> model = MyToyModel( - ... sample_kwargs={ - ... "chains": 2, - ... "draws": 2000, - ... "progressbar": False, - ... "random_seed": 42, - ... } + ... sample_kwargs={ + ... "chains": 2, + ... "draws": 2000, + ... "progressbar": False, + ... "random_seed": 42, + ... } ... ) >>> model.fit(X, y) Inference data... @@ -63,7 +63,7 @@ class PyMCModel(pm.Model): r2 ... r2_std ... dtype: float64 - >>> X_new = rng.normal(loc=0, scale=1, size=(20,2)) + >>> X_new = rng.normal(loc=0, scale=1, size=(20, 2)) >>> model.predict(X_new) Inference data... """ @@ -286,11 +286,11 @@ class InstrumentalVariableRegression(PyMCModel): >>> ## Ensure the endogeneity of the the treatment variable >>> X = -1 + 4 * Z + e2 + 2 * e1 >>> y = 2 + 3 * X + 3 * e1 - >>> t = X.reshape(10,1) - >>> y = y.reshape(10,1) - >>> Z = np.asarray([[1, Z[i]] for i in range(0,10)]) - >>> X = np.asarray([[1, X[i]] for i in range(0,10)]) - >>> COORDS = {'instruments': ['Intercept', 'Z'], 'covariates': ['Intercept', 'X']} + >>> t = X.reshape(10, 1) + >>> y = y.reshape(10, 1) + >>> Z = np.asarray([[1, Z[i]] for i in range(0, 10)]) + >>> X = np.asarray([[1, X[i]] for i in range(0, 10)]) + >>> COORDS = {"instruments": ["Intercept", "Z"], "covariates": ["Intercept", "X"]} >>> sample_kwargs = { ... "tune": 5, ... "draws": 10, @@ -300,12 +300,20 @@ class InstrumentalVariableRegression(PyMCModel): ... "progressbar": False, ... } >>> iv_reg = InstrumentalVariableRegression(sample_kwargs=sample_kwargs) - >>> iv_reg.fit(X, Z,y, t, COORDS, { - ... "mus": [[-2,4], [0.5, 3]], - ... "sigmas": [1, 1], - ... "eta": 2, - ... "lkj_sd": 1, - ... }, None) + >>> iv_reg.fit( + ... X, + ... Z, + ... y, + ... t, + ... COORDS, + ... { + ... "mus": [[-2, 4], [0.5, 3]], + ... "sigmas": [1, 1], + ... "eta": 2, + ... "lkj_sd": 1, + ... }, + ... None, + ... ) Inference data... """ diff --git a/pyproject.toml b/pyproject.toml index 1024a3f8..2195651e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -118,6 +118,9 @@ omit-covered-files = false generate-badge = "docs/source/_static/" badge-format = "svg" +[tool.ruff.format] +docstring-code-format = true + [tool.ruff.lint] extend-select = [ "I", # isort