Experiment length runner (#180)

* add timelines

* add notebook

* add notebook

cluster_experiments/cupac.py

@@ -2,10 +2,11 @@
 
 import pandas as pd
 from numpy.typing import ArrayLike
-from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
+from sklearn.base import BaseEstimator
+from sklearn.utils.validation import NotFittedError, check_is_fitted
 
 
-class EmptyRegressor(BaseEstimator, RegressorMixin):
+class EmptyRegressor(BaseEstimator):
     """
     Empty regressor class. It does not do anything, used to glue the code of other estimators and PowerAnalysis
 
@@ -21,7 +22,7 @@ def from_config(cls, config):
         return cls()
 
 
-class TargetAggregation(BaseEstimator, RegressorMixin):
+class TargetAggregation(BaseEstimator):
     """
     Adds average of target using pre-experiment data
 
@@ -117,12 +118,14 @@ def __init__(
         cupac_model: Optional[BaseEstimator] = None,
         target_col: str = "target",
         features_cupac_model: Optional[List[str]] = None,
+        cache_fit: bool = True,
     ):
         self.cupac_model: BaseEstimator = cupac_model or EmptyRegressor()
         self.target_col = target_col
         self.cupac_outcome_name = f"estimate_{target_col}"
         self.features_cupac_model: List[str] = features_cupac_model or []
         self.is_cupac = not isinstance(self.cupac_model, EmptyRegressor)
+        self.cache_fit = cache_fit
 
     def _prep_data_cupac(
         self, df: pd.DataFrame, pre_experiment_df: pd.DataFrame
@@ -165,23 +168,40 @@ def add_covariates(
             df=df, pre_experiment_df=pre_experiment_df
         )
 
-        # Fit model
-        self.cupac_model.fit(pre_experiment_x, pre_experiment_y)
+        # Fit model if it has not been fitted before
+        self._fit_cupac_model(pre_experiment_x, pre_experiment_y)
 
         # Predict
-        if isinstance(self.cupac_model, RegressorMixin):
-            estimated_target = self.cupac_model.predict(df_predict)
-        elif isinstance(self.cupac_model, ClassifierMixin):
-            estimated_target = self.cupac_model.predict_proba(df_predict)[:, 1]
-        else:
-            raise ValueError(
-                "cupac_model should be an instance of RegressorMixin or ClassifierMixin"
-            )
+        estimated_target = self._predict_cupac_model(df_predict)
 
         # Add cupac outcome name to df
         df[self.cupac_outcome_name] = estimated_target
         return df
 
+    def _fit_cupac_model(
+        self, pre_experiment_x: pd.DataFrame, pre_experiment_y: pd.Series
+    ):
+        """Fits the cupac model.
+        Caches the fitted model in the object, so we only fit it once.
+        We can disable this by setting cache_fit to False.
+        """
+        if not self.cache_fit:
+            self.cupac_model.fit(pre_experiment_x, pre_experiment_y)
+            return
+
+        try:
+            check_is_fitted(self.cupac_model)
+        except NotFittedError:
+            self.cupac_model.fit(pre_experiment_x, pre_experiment_y)
+
+    def _predict_cupac_model(self, df_predict: pd.DataFrame) -> ArrayLike:
+        """Predicts the cupac model"""
+        if hasattr(self.cupac_model, "predict_proba"):
+            return self.cupac_model.predict_proba(df_predict)[:, 1]
+        if hasattr(self.cupac_model, "predict"):
+            return self.cupac_model.predict(df_predict)
+        raise ValueError("cupac_model should have predict or predict_proba method.")
+
     def need_covariates(self, pre_experiment_df: Optional[pd.DataFrame] = None) -> bool:
         return pre_experiment_df is not None and self.is_cupac
 

cluster_experiments/power_analysis.py

@@ -424,9 +424,12 @@ def from_config(cls, config: PowerConfig) -> "PowerAnalysis":
             target_col=config.target_col,
             treatment_col=config.treatment_col,
             treatment=config.treatment,
+            control=config.control,
             n_simulations=config.n_simulations,
             alpha=config.alpha,
+            features_cupac_model=config.features_cupac_model,
             seed=config.seed,
+            hypothesis=config.hypothesis,
         )
 
     def check_treatment_col(self):
@@ -606,6 +609,7 @@ def __init__(
         features_cupac_model: Optional[List[str]] = None,
         seed: Optional[int] = None,
         hypothesis: str = "two-sided",
+        time_col: Optional[str] = None,
     ):
         self.splitter = splitter
         self.analysis = analysis
@@ -616,6 +620,7 @@ def __init__(
         self.treatment_col = treatment_col
         self.alpha = alpha
         self.hypothesis = hypothesis
+        self.time_col = time_col
 
         self.cupac_handler = CupacHandler(
             cupac_model=cupac_model,
@@ -798,6 +803,111 @@ def _get_average_standard_error(
 
         return std_error_mean
 
+    def run_average_standard_error(
+        self,
+        df: pd.DataFrame,
+        pre_experiment_df: Optional[pd.DataFrame] = None,
+        verbose: bool = False,
+        n_simulations: Optional[int] = None,
+        experiment_length: Iterable[int] = (),
+    ) -> Generator[Tuple[float, int], None, None]:
+        """
+        Run power analysis by simulation, using standard errors from the analysis.
+
+        Args:
+            df: Dataframe with outcome and treatment variables.
+            pre_experiment_df: Dataframe with pre-experiment data.
+            verbose: Whether to show progress bar.
+            n_simulations: Number of simulations to run.
+            experiment_length: Length of the experiment in days.
+        """
+        n_simulations = self.n_simulations if n_simulations is None else n_simulations
+
+        for n_days in experiment_length:
+            df_time = df.copy()
+            experiment_start = df_time[self.time_col].min()
+            df_time = df_time.loc[
+                df_time[self.time_col] < experiment_start + pd.Timedelta(days=n_days)
+            ]
+            std_error_mean = self._get_average_standard_error(
+                df=df_time,
+                pre_experiment_df=pre_experiment_df,
+                verbose=verbose,
+                n_simulations=n_simulations,
+            )
+            yield std_error_mean, n_days
+
+    def power_time_line(
+        self,
+        df: pd.DataFrame,
+        pre_experiment_df: Optional[pd.DataFrame] = None,
+        verbose: bool = False,
+        average_effects: Iterable[float] = (),
+        experiment_length: Iterable[int] = (),
+        n_simulations: Optional[int] = None,
+        alpha: Optional[float] = None,
+    ) -> List[Dict]:
+        """
+        Run power analysis by simulation, using standard errors from the analysis.
+
+        Args:
+            df: Dataframe with outcome and treatment variables.
+            pre_experiment_df: Dataframe with pre-experiment data.
+            verbose: Whether to show progress bar.
+            average_effects: Average effects to test.
+            experiment_length: Length of the experiment in days.
+            n_simulations: Number of simulations to run.
+            alpha: Significance level.
+        """
+        alpha = self.alpha if alpha is None else alpha
+
+        results = []
+        for std_error_mean, n_days in self.run_average_standard_error(
+            df=df,
+            pre_experiment_df=pre_experiment_df,
+            verbose=verbose,
+            n_simulations=n_simulations,
+            experiment_length=experiment_length,
+        ):
+            for effect in average_effects:
+                power = self._normal_power_calculation(
+                    alpha=alpha, std_error=std_error_mean, average_effect=effect
+                )
+                results.append(
+                    {"effect": effect, "power": power, "experiment_length": n_days}
+                )
+
+        return results
+
+    def mde_time_line(
+        self,
+        df: pd.DataFrame,
+        pre_experiment_df: Optional[pd.DataFrame] = None,
+        verbose: bool = False,
+        powers: Iterable[float] = (),
+        experiment_length: Iterable[int] = (),
+        n_simulations: Optional[int] = None,
+        alpha: Optional[float] = None,
+    ) -> List[Dict]:
+        alpha = self.alpha if alpha is None else alpha
+
+        results = []
+        for std_error_mean, n_days in self.run_average_standard_error(
+            df=df,
+            pre_experiment_df=pre_experiment_df,
+            verbose=verbose,
+            n_simulations=n_simulations,
+            experiment_length=experiment_length,
+        ):
+            for power in powers:
+                mde = self._normal_mde_calculation(
+                    alpha=alpha, std_error=std_error_mean, power=power
+                )
+                results.append(
+                    {"power": power, "mde": mde, "experiment_length": n_days}
+                )
+        return results
+
     def power_line(
         self,
         df: pd.DataFrame,
@@ -888,9 +998,13 @@ def from_config(cls, config: PowerConfig) -> "NormalPowerAnalysis":
             target_col=config.target_col,
             treatment_col=config.treatment_col,
             treatment=config.treatment,
+            control=config.control,
             n_simulations=config.n_simulations,
             alpha=config.alpha,
+            features_cupac_model=config.features_cupac_model,
             seed=config.seed,
+            hypothesis=config.hypothesis,
+            time_col=config.time_col,
         )
 
     def check_treatment_col(self):

cluster_experiments/power_config.py

@@ -152,7 +152,8 @@ def __post_init__(self):
                 self._set_and_log("washover_time_delta", None, "splitter")
             if self._are_different(self.washover, ""):
                 self._set_and_log("washover", "", "splitter")
-            if self._are_different(self.time_col, None):
+            # an exception is made when we have no perturbator (normal power analysis)
+            if self._are_different(self.time_col, None) and self.perturbator != "":
                 self._set_and_log("time_col", None, "splitter")
 
         if self.perturbator not in {"normal", "beta_relative_positive"}:

mkdocs.yml

@@ -20,7 +20,9 @@ nav:
   - Paired T test: paired_ttest.ipynb
   - Different hypotheses tests: analysis_with_different_hypotheses.ipynb
   - Washover: washover_example.ipynb
-  - Normal Power: normal_power.ipynb
+  - Normal Power:
+      - Compare with simulation: normal_power.ipynb
+      - Time-lines: normal_power_lines.ipynb
   - API:
       - Experiment analysis: api/experiment_analysis.md
       - Perturbators: api/perturbator.md

setup.py

@@ -47,7 +47,7 @@
 
 setup(
     name="cluster_experiments",
-    version="0.17.0",
+    version="0.18.0",
     packages=find_packages(),
     extras_require={
         "dev": dev_packages,

tests/power_analysis/test_normal_power_analysis.py

@@ -1,3 +1,4 @@
+import pandas as pd
 import pytest
 
 from cluster_experiments.experiment_analysis import ClusteredOLSAnalysis, OLSAnalysis
@@ -348,3 +349,67 @@ def test_mde_power_line(df):
     # then
     assert mde_power_line[0.9] > mde_power_line[0.8]
     assert mde_power_line[0.8] > mde_power_line[0.7]
+
+
+def test_mde_time_line(df):
+    # given
+    pw_normal = NormalPowerAnalysis.from_dict(
+        {
+            "splitter": "non_clustered",
+            "analysis": "ols",
+            "n_simulations": 5,
+            "hypothesis": "two-sided",
+            "seed": 20240922,
+            "time_col": "date",
+        }
+    )
+    df_cp = df.copy()
+    df_cp["date"] = pd.to_datetime(df_cp["date"])
+
+    # when
+    mde_time_line = pw_normal.mde_time_line(
+        df_cp, experiment_length=[1, 2, 3], powers=[0.8]
+    )
+    mde_df = pd.DataFrame(mde_time_line)
+
+    # then
+    assert (
+        mde_df.query("experiment_length == 1")["mde"].squeeze()
+        > mde_df.query("experiment_length == 2")["mde"].squeeze()
+    )
+    assert (
+        mde_df.query("experiment_length == 2")["mde"].squeeze()
+        > mde_df.query("experiment_length == 3")["mde"].squeeze()
+    )
+
+
+def test_power_time_line(df):
+    # given
+    pw_normal = NormalPowerAnalysis.from_dict(
+        {
+            "splitter": "non_clustered",
+            "analysis": "ols",
+            "n_simulations": 5,
+            "hypothesis": "two-sided",
+            "seed": 20240922,
+            "time_col": "date",
+        }
+    )
+    df_cp = df.copy()
+    df_cp["date"] = pd.to_datetime(df_cp["date"])
+
+    # when
+    power_time_line = pw_normal.power_time_line(
+        df_cp, experiment_length=[1, 2, 3], average_effects=[0.1]
+    )
+    power_df = pd.DataFrame(power_time_line)
+
+    # then
+    assert (
+        power_df.query("experiment_length == 1")["power"].squeeze()
+        < power_df.query("experiment_length == 2")["power"].squeeze()
+    )
+    assert (
+        power_df.query("experiment_length == 2")["power"].squeeze()
+        < power_df.query("experiment_length == 3")["power"].squeeze()
+    )