fix code style

Author: juanitorduz

causalpy/data/simulate_data.py

@@ -28,7 +28,9 @@ def generate_synthetic_control_data(
     """
     Example:
     >> import pathlib
-    >> df, weightings_true = generate_synthetic_control_data(treatment_time=treatment_time)
+    >> df, weightings_true = generate_synthetic_control_data(
+                                treatment_time=treatment_time
+                            )
     >> df.to_csv(pathlib.Path.cwd() / 'synthetic_control.csv', index=False)
     """
 
@@ -45,15 +47,17 @@ def generate_synthetic_control_data(
         }
     )
 
-    # 2. Generate counterfactual, based on weighted sum of non-treated variables. This is the counterfactual with NO treatment.
+    # 2. Generate counterfactual, based on weighted sum of non-treated variables. This
+    # is the counterfactual with NO treatment.
     weightings_true = dirichlet(np.ones(7)).rvs(1)
     df["counterfactual"] = np.dot(df.to_numpy(), weightings_true.T)
 
     # 3. Generate the causal effect
     causal_effect = gamma(10).pdf(np.arange(0, N, 1) - treatment_time)
     df["causal effect"] = causal_effect * -50
 
-    # 4. Generate the actually observed data, ie the treated with the causal effect applied
+    # 4. Generate the actually observed data, ie the treated with the causal effect
+    # applied
     df["actual"] = df["counterfactual"] + df["causal effect"]
 
     # 5. apply observation noise to all relevant variables
@@ -126,7 +130,9 @@ def generate_time_series_data(treatment_time):
 
 
 def generate_time_series_data_simple(treatment_time, slope=0.0):
-    """Generate simple interrupted time series data, with no seasonality or temporal structure"""
+    """Generate simple interrupted time series data, with no seasonality or temporal
+    structure.
+    """
     dates = pd.date_range(
         start=pd.to_datetime("2010-01-01"), end=pd.to_datetime("2020-01-01"), freq="M"
     )

causalpy/pymc_experiments.py

@@ -28,19 +28,28 @@ def print_coefficients(self):
         """Prints the model coefficients"""
         print("Model coefficients:")
         coeffs = az.extract(self.prediction_model.idata.posterior, var_names="beta")
-        # Note: f"{name: <30}" pads the name with spaces so that we have alignment of the stats despite variable names of different lengths
+        # Note: f"{name: <30}" pads the name with spaces so that we have alignment of
+        # the stats despite variable names of different lengths
         for name in self.labels:
             coeff_samples = coeffs.sel(coeffs=name)
             print(
-                f"  {name: <30}{coeff_samples.mean().data:.2f}, 94% HDI [{coeff_samples.quantile(0.03).data:.2f}, {coeff_samples.quantile(1-0.03).data:.2f}]"
+                f"""
+                {name: <30}{coeff_samples.mean().data:.2f},
+                94% HDI [{coeff_samples.quantile(0.03).data:.2f},
+                {coeff_samples.quantile(1-0.03).data:.2f}]
+                """
             )
         # add coeff for measurement std
         coeff_samples = az.extract(
             self.prediction_model.idata.posterior, var_names="sigma"
         )
         name = "sigma"
         print(
-            f"  {name: <30}{coeff_samples.mean().data:.2f}, 94% HDI [{coeff_samples.quantile(0.03).data:.2f}, {coeff_samples.quantile(1-0.03).data:.2f}]"
+            f"""
+            {name: <30}{coeff_samples.mean().data:.2f},
+            94% HDI [{coeff_samples.quantile(0.03).data:.2f},
+            {coeff_samples.quantile(1-0.03).data:.2f}]
+            """
         )
 
 
@@ -120,8 +129,12 @@ def plot(self):
             include_label=False,
         )
         ax[0].plot(self.datapost.index, self.post_y, "k.")
+
         ax[0].set(
-            title=f"Pre-intervention Bayesian $R^2$: {self.score.r2:.3f} (std = {self.score.r2_std:.3f})"
+            title=f"""
+            Pre-intervention Bayesian $R^2$: {self.score.r2:.3f}
+            (std = {self.score.r2_std:.3f})
+            """
         )
 
         plot_xY(self.datapre.index, self.pre_impact, ax=ax[1])
@@ -200,7 +213,8 @@ class DifferenceInDifferences(ExperimentalDesign):
 
     .. note::
 
-       There is no pre/post intervention data distinction for DiD, we fit all the data available.
+        There is no pre/post intervention data distinction for DiD, we fit all the
+        data available.
 
     """
 
@@ -224,7 +238,8 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
-        # TODO: `treated` is a deterministic function of group and time, so this should be a function rather than supplied data
+        # TODO: `treated` is a deterministic function of group and time, so this should
+        # be a function rather than supplied data
 
         # DEVIATION FROM SKL EXPERIMENT CODE =============================
         # fit the model to the observed (pre-intervention) data
@@ -309,7 +324,8 @@ def plot(self):
             showmedians=False,
             widths=0.2,
         )
-        # Plot counterfactual - post-test for treatment group IF no treatment had occurred.
+        # Plot counterfactual - post-test for treatment group IF no treatment
+        # had occurred.
         parts = ax.violinplot(
             az.extract(
                 self.y_pred_counterfactual,
@@ -381,13 +397,16 @@ class RegressionDiscontinuity(ExperimentalDesign):
 
     :param data: A pandas dataframe
     :param formula: A statistical model formula
-    :param treatment_threshold: A scalar threshold value at which the treatment is applied
+    :param treatment_threshold: A scalar threshold value at which the treatment
+                                is applied
     :param prediction_model: A PyMC model
-    :param running_variable_name: The name of the predictor variable that the treatment threshold is based upon
+    :param running_variable_name: The name of the predictor variable that the treatment
+                                  threshold is based upon
 
     .. note::
 
-       There is no pre/post intervention data distinction for the regression discontinuity design, we fit all the data available.
+        There is no pre/post intervention data distinction for the regression
+        discontinuity design, we fit all the data available.
     """
 
     def __init__(
@@ -412,7 +431,8 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
-        # TODO: `treated` is a deterministic function of x and treatment_threshold, so this could be a function rather than supplied data
+        # TODO: `treated` is a deterministic function of x and treatment_threshold, so
+        # this could be a function rather than supplied data
 
         # DEVIATION FROM SKL EXPERIMENT CODE =============================
         # fit the model to the observed (pre-intervention) data
@@ -435,8 +455,10 @@ def __init__(
         (new_x,) = build_design_matrices([self._x_design_info], self.x_pred)
         self.pred = self.prediction_model.predict(X=np.asarray(new_x))
 
-        # calculate discontinuity by evaluating the difference in model expectation on either side of the discontinuity
-        # NOTE: `"treated": np.array([0, 1])`` assumes treatment is applied above (not below) the threshold
+        # calculate discontinuity by evaluating the difference in model expectation on
+        # either side of the discontinuity
+        # NOTE: `"treated": np.array([0, 1])`` assumes treatment is applied above
+        # (not below) the threshold
         self.x_discon = pd.DataFrame(
             {
                 self.running_variable_name: np.array(
@@ -457,7 +479,7 @@ def _is_treated(self, x):
 
         .. warning::
 
-           Assumes treatment is given to those ABOVE the treatment threshold.
+            Assumes treatment is given to those ABOVE the treatment threshold.
         """
         return np.greater_equal(x, self.treatment_threshold)
 
@@ -479,10 +501,13 @@ def plot(self):
             ax=ax,
         )
         # create strings to compose title
-        r2 = f"Bayesian $R^2$ on all data = {self.score.r2:.3f} (std = {self.score.r2_std:.3f})"
+        title_info = f"{self.score.r2:.3f} (std = {self.score.r2_std:.3f})"
+        r2 = f"Bayesian $R^2$ on all data = {title_info}"
         percentiles = self.discontinuity_at_threshold.quantile([0.03, 1 - 0.03]).values
         ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]"
-        discon = f"Discontinuity at threshold = {self.discontinuity_at_threshold.mean():.2f}, "
+        discon = f"""
+            Discontinuity at threshold = {self.discontinuity_at_threshold.mean():.2f},
+            """
         ax.set(title=r2 + "\n" + discon + ci)
         # Intervention line
         ax.axvline(
@@ -502,7 +527,7 @@ def summary(self):
         print(f"Formula: {self.formula}")
         print(f"Running variable: {self.running_variable_name}")
         print(f"Threshold on running variable: {self.treatment_threshold}")
-        print(f"\nResults:")
+        print("\nResults:")
         print(
             f"Discontinuity at threshold = {self.discontinuity_at_threshold.mean():.2f}"
         )

causalpy/pymc_models.py

@@ -21,7 +21,9 @@ def _data_setter(self, X):
             pm.set_data({"X": X})
 
     def fit(self, X, y, coords):
-        """Draw samples from posterior, prior predictive, and posterior predictive distributions."""
+        """Draw samples from posterior, prior predictive, and posterior predictive
+        distributions.
+        """
         self.build_model(X, y, coords)
         with self.model:
             self.idata = pm.sample()
@@ -43,7 +45,8 @@ def score(self, X, y):
 
         .. caution::
 
-           The Bayesian :math:`R^2` is not the same as the traditional coefficient of determination, https://en.wikipedia.org/wiki/Coefficient_of_determination.
+            The Bayesian :math:`R^2` is not the same as the traditional coefficient of
+            determination, https://en.wikipedia.org/wiki/Coefficient_of_determination.
 
         """
         yhat = self.predict(X)

causalpy/skl_experiments.py

@@ -141,9 +141,10 @@ def plot_coeffs(self):
         )
 
 
-# InterruptedTimeSeries and SyntheticControl are basically the same thing but with different
-# predictor variables. So we just have a TimeSeriesExperiment class and InterruptedTimeSeries
-# and SyntheticControl are both equal to the TimeSeriesExperiment class
+# InterruptedTimeSeries and SyntheticControl are basically the same thing but with
+# different predictor variables. So we just have a TimeSeriesExperiment class and
+# InterruptedTimeSeries and SyntheticControl are both equal to the TimeSeriesExperiment
+# class
 
 
 class InterruptedTimeSeries(TimeSeriesExperiment):
@@ -168,7 +169,8 @@ class DifferenceInDifferences(ExperimentalDesign):
     """
     .. note::
 
-       There is no pre/post intervention data distinction for DiD, we fit all the data available.
+        There is no pre/post intervention data distinction for DiD, we fit all the data
+        available.
     """
 
     def __init__(
@@ -190,7 +192,8 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
-        # TODO: `treated` is a deterministic function of group and time, so this should be a function rather than supplied data
+        # TODO: `treated` is a deterministic function of group and time, so this should
+        # be a function rather than supplied data
 
         # fit the model to all the data
         self.prediction_model.fit(X=self.X, y=self.y)
@@ -254,7 +257,8 @@ def plot(self):
             markersize=10,
             label="model fit (treament group)",
         )
-        # Plot counterfactual - post-test for treatment group IF no treatment had occurred.
+        # Plot counterfactual - post-test for treatment group IF no treatment
+        # had occurred.
         ax.plot(
             self.x_pred_counterfactual[self.time_variable_name],
             self.y_pred_counterfactual,
@@ -297,7 +301,8 @@ class RegressionDiscontinuity(ExperimentalDesign):
 
     .. note::
 
-       There is no pre/post intervention data distinction for the regression discontinuity design, we fit all the data available.
+        There is no pre/post intervention data distinction for the regression
+        discontinuity design, we fit all the data available.
 
     """
 
@@ -322,7 +327,8 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
-        # TODO: `treated` is a deterministic function of x and treatment_threshold, so this could be a function rather than supplied data
+        # TODO: `treated` is a deterministic function of x and treatment_threshold, so
+        # this could be a function rather than supplied data
 
         # fit the model to all the data
         self.prediction_model.fit(X=self.X, y=self.y)
@@ -342,8 +348,10 @@ def __init__(
         (new_x,) = build_design_matrices([self._x_design_info], self.x_pred)
         self.pred = self.prediction_model.predict(X=np.asarray(new_x))
 
-        # calculate discontinuity by evaluating the difference in model expectation on either side of the discontinuity
-        # NOTE: `"treated": np.array([0, 1])`` assumes treatment is applied above (not below) the threshold
+        # calculate discontinuity by evaluating the difference in model expectation on
+        # either side of the discontinuity
+        # NOTE: `"treated": np.array([0, 1])`` assumes treatment is applied above
+        # (not below) the threshold
         self.x_discon = pd.DataFrame(
             {
                 self.running_variable_name: np.array(
@@ -359,11 +367,12 @@ def __init__(
         )
 
     def _is_treated(self, x):
-        """Returns ``True`` if ``x`` is greater than or equal to the treatment threshold.
+        """Returns ``True`` if ``x`` is greater than or equal to the treatment
+        threshold.
 
         .. warning::
 
-           Assumes treatment is given to those ABOVE the treatment threshold.
+            Assumes treatment is given to those ABOVE the treatment threshold.
         """
         return np.greater_equal(x, self.treatment_threshold)
 
@@ -406,7 +415,7 @@ def summary(self):
         print(f"Formula: {self.formula}")
         print(f"Running variable: {self.running_variable_name}")
         print(f"Threshold on running variable: {self.treatment_threshold}")
-        print(f"\nResults:")
+        print("\nResults:")
         print(f"Discontinuity at threshold = {self.discontinuity_at_threshold:.2f}")
         print("Model coefficients:")
         for name, val in zip(self.labels, self.prediction_model.coef_[0]):

docs/conf.py

@@ -12,11 +12,13 @@
 import os
 import sys
 
+from causalpy.version import __version__
+
 sys.path.insert(0, os.path.abspath("../"))
 
 # autodoc_mock_imports
 # This avoids autodoc breaking when it can't find packages imported in the code.
-# https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#confval-autodoc_mock_imports
+# https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#confval-autodoc_mock_imports # noqa: E501
 autodoc_mock_imports = [
     "arviz",
     "matplotlib",
@@ -37,7 +39,6 @@
 copyright = "2022, Benjamin T. Vincent"
 author = "Benjamin T. Vincent"
 
-from causalpy.version import __version__
 
 release = __version__
 
@@ -57,10 +58,13 @@
 exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 # -- nbsphinx config ----------------------------------------------------------
-# Opt out of executing the notebooks remotely. This will save time in the remote build process on readthedocs. The notebooks in /docs/notebooks will be parsed/converted, but not re-executed.
+# Opt out of executing the notebooks remotely. This will save time in the remote build
+# process on readthedocs. The notebooks in /docs/notebooks will be parsed/converted,
+# but not re-executed.
 nbsphinx_execute = "never"
 
-# MyST options for working with markdown files. Info about extensions here https://myst-parser.readthedocs.io/en/latest/syntax/optional.html?highlight=math#admonition-directives
+# MyST options for working with markdown files.
+# Info about extensions here https://myst-parser.readthedocs.io/en/latest/syntax/optional.html?highlight=math#admonition-directives # noqa: E501
 myst_enable_extensions = ["dollarmath", "amsmath", "colon_fence", "linkify"]
 
 # -- Options for HTML output -------------------------------------------------

setup.py

@@ -3,6 +3,8 @@
 
 from setuptools import find_packages, setup
 
+from causalpy.version import __version__
+
 PROJECT_ROOT = os.path.dirname(os.path.realpath(__file__))
 README_FILE = os.path.join(PROJECT_ROOT, "README.md")
 VERSION_FILE = os.path.join(PROJECT_ROOT, "bambi", "version.py")
@@ -16,7 +18,6 @@ def get_long_description():
 
 # get version
 sys.path.insert(0, os.path.abspath("../"))
-from causalpy.version import __version__
 
 with open(REQUIREMENTS_FILE) as f:
     install_reqs = f.read().splitlines()