Steppedweddge implementation

cluster_experiments/power_config.py

@@ -79,6 +79,7 @@ class PowerConfig:
 
     # optional mappings
     cupac_model: str = ""
+    comparison_col: str = ""
 
     # Shared
     target_col: str = "target"

cluster_experiments/random_splitter.py

@@ -481,3 +481,43 @@ def from_config(cls, config) -> "StratifiedSwitchbackSplitter":
             splitter_weights=config.splitter_weights,
             washover=washover_cls.from_config(config),
         )
+
+
+class ClusteredSteppedWedge(ClusteredSplitter):
+    """Normally used for gradual rollouts, where the treatment is rolled out in batches."""
+
+    def __init__(self, rollout_df: pd.DataFrame, cluster_cols: List[str]):
+        """rollout df should have a column 'rollout_date' and a column 'perc_cities'"""
+        super().__init__(cluster_cols)
+        self.rollout_df = rollout_df
+
+    def assign_treatment_df(
+        self,
+        df: pd.DataFrame,
+    ) -> pd.DataFrame:
+        """
+        Takes a df, randomizes treatments and adds the treatment column to the dataframe
+
+        Arguments:
+            df: dataframe to assign treatments to
+        """
+        df = df.copy()
+
+        # raise error if any nulls in cluster_cols
+        if df[self.cluster_cols].isnull().values.any():
+            raise ValueError(
+                f"Null values found in cluster_cols: {self.cluster_cols}. "
+                "Please remove nulls before running the splitter."
+            )
+
+        clusters_df = df.loc[:, self.cluster_cols].drop_duplicates()
+        clusters_df[self.treatment_col] = self.sample_treatment(clusters_df)
+        df = df.merge(clusters_df, on=self.cluster_cols, how="left")
+
+        # once we know the treatment cities, we need to stablish a roll out strategy. 2 inputs, rollout date and n cities.
+
+        df = df.merge(self.rollout_df, on=["cluster"], how="left")
+        # replace all 1 by treatment
+
+        df["treatment"] = np.where(df["date"] >= df["rollout_date"], 1, 0)
+        return df

mkdocs.yml

@@ -11,7 +11,10 @@ nav:
     - Index: index.md
   - Cupac example: cupac_example.ipynb
   - Custom classes: create_custom_classes.ipynb
-  - Stratified switchback: switchback.ipynb
+  - Switchback:
+      - Stratified switchback: switchback.ipynb
+      - Switchback calendar visualization: plot_calendars.ipynb
+      - Visualization - 4-hour switches: plot_calendars_hours.ipynb
   - Multiple treatments: multivariate.ipynb
   - AA test clustered: aa_test.ipynb
   - Paired T test: paired_ttest.ipynb
@@ -22,6 +25,7 @@ nav:
       - Pre experiment outcome model: api/cupac_model.md
       - Power config: api/power_config.md
       - Power analysis: api/power_analysis.md
+      - Washover: api/washover.md
 plugins:
   - mkdocstrings:
       watch:

setup.py

@@ -38,6 +38,8 @@
     "mkdocstrings==0.18.0",
     "jinja2<3.1.0",
     "mkdocs-jupyter==0.22.0",
+    "plotnine==0.8.0",
+    "matplotlib==3.4.3",
 ]
 
 dev_packages = test_packages + util_packages + docs_packages