init

Author: juanitorduz

causalpy/data/simulate_data.py

@@ -31,7 +31,7 @@ def _smoothed_gaussian_random_walk(
     gaussian_random_walk_mu: float,
     gaussian_random_walk_sigma: float,
     N: int,
-    lowess_kwargs: dict[str, Any],
+    lowess_kwargs: dict,
 ) -> tuple[np.ndarray, np.ndarray]:
     """
     Generates Gaussian random walk data and applies LOWESS.
@@ -57,7 +57,7 @@ def generate_synthetic_control_data(
     treatment_time: int = 70,
     grw_mu: float = 0.25,
     grw_sigma: float = 1,
-    lowess_kwargs: dict[str, Any] | None = None,
+    lowess_kwargs: dict = default_lowess_kwargs,
 ) -> tuple[pd.DataFrame, np.ndarray]:
     """
     Generates data for synthetic control example.
@@ -78,9 +78,6 @@ def generate_synthetic_control_data(
     >>> from causalpy.data.simulate_data import generate_synthetic_control_data
     >>> df, weightings_true = generate_synthetic_control_data(treatment_time=70)
     """
-    if lowess_kwargs is None:
-        lowess_kwargs = default_lowess_kwargs
-
     # 1. Generate non-treated variables
     df = pd.DataFrame(
         {
@@ -166,7 +163,9 @@ def generate_time_series_data(
     return df
 
 
-def generate_time_series_data_seasonal(treatment_time: pd.Timestamp) -> pd.DataFrame:
+def generate_time_series_data_seasonal(
+    treatment_time: pd.Timestamp,
+) -> pd.DataFrame:
     """
     Generates 10 years of monthly data with seasonality
     """
@@ -184,7 +183,9 @@ def generate_time_series_data_seasonal(treatment_time: pd.Timestamp) -> pd.DataF
 
     N = df.shape[0]
     idx = np.arange(N)[df.index > treatment_time]
-    df["causal effect"] = 100 * gamma(10).pdf(np.arange(0, N, 1) - np.min(idx))
+    df["causal effect"] = 100 * gamma(10).pdf(
+        np.array(np.arange(0, N, 1)) - int(np.min(idx))
+    )
 
     df["y"] += df["causal effect"]
     df["y"] += norm(0, 2).rvs(N)
@@ -310,8 +311,8 @@ def impact(x: np.ndarray) -> np.ndarray:
 def generate_ancova_data(
     N: int = 200,
     pre_treatment_means: np.ndarray = np.array([10, 12]),
-    treatment_effect: float = 2,
-    sigma: float = 1,
+    treatment_effect: int = 2,
+    sigma: int = 1,
 ) -> pd.DataFrame:
     """
     Generate ANCOVA example data
@@ -445,7 +446,7 @@ def generate_multicell_geolift_data() -> pd.DataFrame:
 
 
 def generate_seasonality(
-    n: int = 12, amplitude: float = 1, length_scale: float = 0.5
+    n: int = 12, amplitude: int = 1, length_scale: float = 0.5
 ) -> np.ndarray:
     """Generate monthly seasonality by sampling from a Gaussian process with a
     Gaussian kernel, using numpy code"""
@@ -463,9 +464,9 @@ def generate_seasonality(
 def periodic_kernel(
     x1: np.ndarray,
     x2: np.ndarray,
-    period: float = 1,
-    length_scale: float = 1,
-    amplitude: float = 1,
+    period: int = 1,
+    length_scale: float = 1.0,
+    amplitude: int = 1,
 ) -> np.ndarray:
     """Generate a periodic kernel for gaussian process"""
     return amplitude**2 * np.exp(
@@ -475,10 +476,10 @@ def periodic_kernel(
 
 def create_series(
     n: int = 52,
-    amplitude: float = 1,
-    length_scale: float = 2,
+    amplitude: int = 1,
+    length_scale: int = 2,
     n_years: int = 4,
-    intercept: float = 3,
+    intercept: int = 3,
 ) -> np.ndarray:
     """
     Returns numpy tile with generated seasonality data repeated over

causalpy/experiments/base.py

@@ -16,6 +16,7 @@
 """
 
 from abc import abstractmethod
+from typing import Any, Union
 
 import arviz as az
 import matplotlib.pyplot as plt
@@ -29,10 +30,12 @@
 class BaseExperiment:
     """Base class for quasi experimental designs."""
 
+    labels: list[str]
+
     supports_bayes: bool
     supports_ols: bool
 
-    def __init__(self, model=None):
+    def __init__(self, model: Union[PyMCModel, RegressorMixin] | None = None) -> None:
         # Ensure we've made any provided Scikit Learn model (as identified as being type
         # RegressorMixin) compatible with CausalPy by appending our custom methods.
         if isinstance(model, RegressorMixin):
@@ -50,16 +53,19 @@ def __init__(self, model=None):
         if self.model is None:
             raise ValueError("model not set or passed.")
 
+    def fit(self, *args: Any, **kwargs: Any) -> None:
+        raise NotImplementedError("fit method not implemented")
+
     @property
-    def idata(self):
+    def idata(self) -> az.InferenceData:
         """Return the InferenceData object of the model. Only relevant for PyMC models."""
         return self.model.idata
 
-    def print_coefficients(self, round_to=None):
+    def print_coefficients(self, round_to: int | None = None) -> None:
         """Ask the model to print its coefficients."""
         self.model.print_coefficients(self.labels, round_to)
 
-    def plot(self, *args, **kwargs) -> tuple:
+    def plot(self, *args: Any, **kwargs: Any) -> tuple:
         """Plot the model.
 
         Internally, this function dispatches to either `_bayesian_plot` or `_ols_plot`
@@ -75,16 +81,16 @@ def plot(self, *args, **kwargs) -> tuple:
                 raise ValueError("Unsupported model type")
 
     @abstractmethod
-    def _bayesian_plot(self, *args, **kwargs):
+    def _bayesian_plot(self, *args: Any, **kwargs: Any) -> tuple:
         """Abstract method for plotting the model."""
         raise NotImplementedError("_bayesian_plot method not yet implemented")
 
     @abstractmethod
-    def _ols_plot(self, *args, **kwargs):
+    def _ols_plot(self, *args: Any, **kwargs: Any) -> tuple:
         """Abstract method for plotting the model."""
         raise NotImplementedError("_ols_plot method not yet implemented")
 
-    def get_plot_data(self, *args, **kwargs) -> pd.DataFrame:
+    def get_plot_data(self, *args: Any, **kwargs: Any) -> pd.DataFrame:
         """Recover the data of an experiment along with the prediction and causal impact information.
 
         Internally, this function dispatches to either :func:`get_plot_data_bayesian` or :func:`get_plot_data_ols`
@@ -98,11 +104,11 @@ def get_plot_data(self, *args, **kwargs) -> pd.DataFrame:
             raise ValueError("Unsupported model type")
 
     @abstractmethod
-    def get_plot_data_bayesian(self, *args, **kwargs):
+    def get_plot_data_bayesian(self, *args: Any, **kwargs: Any) -> pd.DataFrame:
         """Abstract method for recovering plot data."""
         raise NotImplementedError("get_plot_data_bayesian method not yet implemented")
 
     @abstractmethod
-    def get_plot_data_ols(self, *args, **kwargs):
+    def get_plot_data_ols(self, *args: Any, **kwargs: Any) -> pd.DataFrame:
         """Abstract method for recovering plot data."""
         raise NotImplementedError("get_plot_data_ols method not yet implemented")

causalpy/experiments/diff_in_diff.py

@@ -15,6 +15,8 @@
 Difference in differences
 """
 
+from typing import Union
+
 import arviz as az
 import numpy as np
 import pandas as pd
@@ -92,8 +94,8 @@ def __init__(
         time_variable_name: str,
         group_variable_name: str,
         post_treatment_variable_name: str = "post_treatment",
-        model=None,
-        **kwargs,
+        model: Union[PyMCModel, RegressorMixin] | None = None,
+        **kwargs: dict,
     ) -> None:
         super().__init__(model=model)
         self.causal_impact: xr.DataArray | float | None
@@ -234,14 +236,14 @@ def __init__(
                 f"{self.group_variable_name}:{self.post_treatment_variable_name}"
             )
             matched_key = next((k for k in coef_map if interaction_term in k), None)
-            att = coef_map.get(matched_key)
+            att = coef_map.get(matched_key) if matched_key is not None else None
             self.causal_impact = att
         else:
             raise ValueError("Model type not recognized")
 
         return
 
-    def input_validation(self):
+    def input_validation(self) -> None:
         # Validate formula structure and interaction interaction terms
         self._validate_formula_interaction_terms()
 
@@ -269,7 +271,7 @@ def input_validation(self):
                 coded. Consisting of 0's and 1's only."""
             )
 
-    def _validate_formula_interaction_terms(self):
+    def _validate_formula_interaction_terms(self) -> None:
         """
         Validate that the formula contains at most one interaction term and no three-way or higher-order interactions.
         Raises FormulaException if more than one interaction term is found or if any interaction term has more than 2 variables.
@@ -299,7 +301,7 @@ def _validate_formula_interaction_terms(self):
                 "Multiple interaction terms are not currently supported as they complicate interpretation of the causal effect."
             )
 
-    def summary(self, round_to=None) -> None:
+    def summary(self, round_to: int | None = 2) -> None:
         """Print summary of main results and model coefficients.
 
         :param round_to:
@@ -311,11 +313,13 @@ def summary(self, round_to=None) -> None:
         print(self._causal_impact_summary_stat(round_to))
         self.print_coefficients(round_to)
 
-    def _causal_impact_summary_stat(self, round_to=None) -> str:
+    def _causal_impact_summary_stat(self, round_to: int | None = None) -> str:
         """Computes the mean and 94% credible interval bounds for the causal impact."""
         return f"Causal impact = {convert_to_string(self.causal_impact, round_to=round_to)}"
 
-    def _bayesian_plot(self, round_to=None, **kwargs) -> tuple[plt.Figure, plt.Axes]:
+    def _bayesian_plot(
+        self, round_to: int | None = None, **kwargs: dict
+    ) -> tuple[plt.Figure, plt.Axes]:
         """
         Plot the results
 
@@ -463,9 +467,10 @@ def _plot_causal_impact_arrow(results, ax):
         )
         return fig, ax
 
-    def _ols_plot(self, round_to=None, **kwargs) -> tuple[plt.Figure, plt.Axes]:
+    def _ols_plot(
+        self, round_to: int | None = 2, **kwargs: dict
+    ) -> tuple[plt.Figure, plt.Axes]:
         """Generate plot for difference-in-differences"""
-        round_to = kwargs.get("round_to")
         fig, ax = plt.subplots()
 
         # Plot raw data
@@ -528,11 +533,15 @@ def _ols_plot(self, round_to=None, **kwargs) -> tuple[plt.Figure, plt.Axes]:
             va="center",
         )
         # formatting
+        # In OLS context, causal_impact should be a float, but mypy doesn't know this
+        causal_impact_value = (
+            float(self.causal_impact) if self.causal_impact is not None else 0.0
+        )
         ax.set(
             xlim=[-0.05, 1.1],
             xticks=[0, 1],
             xticklabels=["pre", "post"],
-            title=f"Causal impact = {round_num(self.causal_impact, round_to)}",
+            title=f"Causal impact = {round_num(causal_impact_value, round_to)}",
         )
         ax.legend(fontsize=LEGEND_FONT_SIZE)
         return fig, ax

causalpy/experiments/instrumental_variable.py

@@ -97,10 +97,10 @@ def __init__(
         data: pd.DataFrame,
         instruments_formula: str,
         formula: str,
-        model=None,
-        priors=None,
-        **kwargs,
-    ):
+        model: BaseExperiment | None = None,
+        priors: dict | None = None,
+        **kwargs: dict,
+    ) -> None:
         super().__init__(model=model)
         self.expt_type = "Instrumental Variable Regression"
         self.data = data
@@ -138,11 +138,11 @@ def __init__(
                 "lkj_sd": 1,
             }
         self.priors = priors
-        self.model.fit(
+        self.model.fit(  # type: ignore[call-arg,union-attr]
             X=self.X, Z=self.Z, y=self.y, t=self.t, coords=COORDS, priors=self.priors
         )
 
-    def input_validation(self):
+    def input_validation(self) -> None:
         """Validate the input data and model formula for correctness"""
         treatment = self.instruments_formula.split("~")[0]
         test = treatment.strip() in self.instruments_data.columns
@@ -165,7 +165,7 @@ def input_validation(self):
                 The coefficients should be interpreted appropriately."""
             )
 
-    def get_2SLS_fit(self):
+    def get_2SLS_fit(self) -> None:
         """
         Two Stage Least Squares Fit
 
@@ -187,7 +187,7 @@ def get_2SLS_fit(self):
         self.first_stage_reg = first_stage_reg
         self.second_stage_reg = second_stage_reg
 
-    def get_naive_OLS_fit(self):
+    def get_naive_OLS_fit(self) -> None:
         """
         Naive Ordinary Least Squares
 
@@ -199,7 +199,7 @@ def get_naive_OLS_fit(self):
         self.ols_beta_params = dict(zip(self._x_design_info.column_names, beta_params))
         self.ols_reg = ols_reg
 
-    def plot(self, round_to=None):
+    def plot(self, *args, **kwargs) -> None:  # type: ignore[override]
         """
         Plot the results
 
@@ -208,7 +208,7 @@ def plot(self, round_to=None):
         """
         raise NotImplementedError("Plot method not implemented.")
 
-    def summary(self, round_to=None) -> None:
+    def summary(self, round_to: int | None = None) -> None:
         """Print summary of main results and model coefficients.
 
         :param round_to: