Merge pull request #430 from pymc-labs/ruff-doctest

Enable ruff formatting of docstrings

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])

causalpy/experiments/diff_in_diff.py

@@ -70,8 +70,8 @@ class DifferenceInDifferences(BaseExperiment):
     ...             "random_seed": seed,
     ...             "progressbar": False,
     ...         }
-    ...     )
-    ...  )
+    ...     ),
+    ... )
     """
 
     supports_ols = True

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),
     ... )
     """
 

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={

causalpy/experiments/prepostfit.py

@@ -337,7 +337,7 @@ class InterruptedTimeSeries(PrePostFit):
     ...             "random_seed": seed,
     ...             "progressbar": False,
     ...         }
-    ...     )
+    ...     ),
     ... )
     """
 

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
     <BLANKLINE>

causalpy/pymc_models.py

@@ -50,20 +50,20 @@ 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...
     >>> model.score(X, y)  # doctest: +ELLIPSIS
     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...
     """
 

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