from cluster_experiments.random_splitter import *

BalancedClusteredSplitter

Bases: ClusteredSplitter

Like ClusteredSplitter, but ensures that treatments are balanced among clusters. That is, if we have 25 clusters and 2 treatments, 13 clusters should have treatment A and 12 clusters should have treatment B.

Source code in cluster_experiments/random_splitter.py

class BalancedClusteredSplitter(ClusteredSplitter):
    """Like ClusteredSplitter, but ensures that treatments are balanced among clusters. That is, if we have
    25 clusters and 2 treatments, 13 clusters should have treatment A and 12 clusters should have treatment B."""

    def sample_treatment(
        self,
        cluster_df: pd.DataFrame,
    ) -> List[str]:
        """
        Samples treatments for each cluster

        Arguments:
            cluster_df: dataframe to assign treatments to
        """
        n_clusters = len(cluster_df)
        n_treatments = len(self.treatments)
        n_per_treatment = n_clusters // n_treatments
        n_extra = n_clusters % n_treatments
        treatments = []
        for i in range(n_treatments):
            treatments += [self.treatments[i]] * (n_per_treatment + (i < n_extra))
        random.shuffle(treatments)
        return treatments

sample_treatment(cluster_df)

Samples treatments for each cluster

Parameters:

Name	Type	Description	Default
cluster_df	DataFrame	dataframe to assign treatments to	required

Source code in cluster_experiments/random_splitter.py

def sample_treatment(
    self,
    cluster_df: pd.DataFrame,
) -> List[str]:
    """
    Samples treatments for each cluster

    Arguments:
        cluster_df: dataframe to assign treatments to
    """
    n_clusters = len(cluster_df)
    n_treatments = len(self.treatments)
    n_per_treatment = n_clusters // n_treatments
    n_extra = n_clusters % n_treatments
    treatments = []
    for i in range(n_treatments):
        treatments += [self.treatments[i]] * (n_per_treatment + (i < n_extra))
    random.shuffle(treatments)
    return treatments

BalancedSwitchbackSplitter

Bases: BalancedClusteredSplitter, SwitchbackSplitter

Like SwitchbackSplitter, but ensures that treatments are balanced among clusters. That is, if we have 25 clusters and 2 treatments, 13 clusters should have treatment A and 12 clusters should have treatment B.

Source code in cluster_experiments/random_splitter.py

class BalancedSwitchbackSplitter(BalancedClusteredSplitter, SwitchbackSplitter):
    """
    Like SwitchbackSplitter, but ensures that treatments are balanced among clusters. That is, if we have
    25 clusters and 2 treatments, 13 clusters should have treatment A and 12 clusters should have treatment B.
    """

    pass

ClusteredSplitter

Bases: RandomSplitter

Splits randomly using clusters

Parameters:

Name	Type	Description	Default
cluster_cols	List[str]	List of columns to use as clusters	required
treatments	Optional[List[str]]	list of treatments	None
treatment_col	str	Name of the column with the treatment variable.	'treatment'
splitter_weights	Optional[List[float]]	weights to use for the splitter, should have the same length as treatments, each weight should correspond to an element in treatments	None

Usage:

import pandas as pd
from cluster_experiments.random_splitter import ClusteredSplitter
splitter = ClusteredSplitter(cluster_cols=["city"])
df = pd.DataFrame({"city": ["A", "B", "C"]})
df = splitter.assign_treatment_df(df)
print(df)

Source code in cluster_experiments/random_splitter.py

class ClusteredSplitter(RandomSplitter):
    """
    Splits randomly using clusters

    Arguments:
        cluster_cols: List of columns to use as clusters
        treatments: list of treatments
        treatment_col: Name of the column with the treatment variable.
        splitter_weights: weights to use for the splitter, should have the same length as treatments, each weight should correspond to an element in treatments

    Usage:
    ```python
    import pandas as pd
    from cluster_experiments.random_splitter import ClusteredSplitter
    splitter = ClusteredSplitter(cluster_cols=["city"])
    df = pd.DataFrame({"city": ["A", "B", "C"]})
    df = splitter.assign_treatment_df(df)
    print(df)
    ```
    """

    def __init__(
        self,
        cluster_cols: List[str],
        treatments: Optional[List[str]] = None,
        treatment_col: str = "treatment",
        splitter_weights: Optional[List[float]] = None,
    ) -> None:
        self.treatments = treatments or ["A", "B"]
        self.cluster_cols = cluster_cols
        self.treatment_col = treatment_col
        self.splitter_weights = splitter_weights

    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")
        return df

    def sample_treatment(
        self,
        cluster_df: pd.DataFrame,
    ) -> List[str]:
        """
        Samples treatments for each cluster

        Arguments:
            cluster_df: dataframe to assign treatments to
        """
        return random.choices(
            self.treatments, k=len(cluster_df), weights=self.splitter_weights
        )

assign_treatment_df(df)

Takes a df, randomizes treatments and adds the treatment column to the dataframe

Parameters:

Name	Type	Description	Default
df	DataFrame	dataframe to assign treatments to	required

Source code in cluster_experiments/random_splitter.py

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")
    return df

sample_treatment(cluster_df)

Samples treatments for each cluster

Parameters:

Name	Type	Description	Default
cluster_df	DataFrame	dataframe to assign treatments to	required

Source code in cluster_experiments/random_splitter.py

def sample_treatment(
    self,
    cluster_df: pd.DataFrame,
) -> List[str]:
    """
    Samples treatments for each cluster

    Arguments:
        cluster_df: dataframe to assign treatments to
    """
    return random.choices(
        self.treatments, k=len(cluster_df), weights=self.splitter_weights
    )

FixedSizeClusteredSplitter

Bases: ClusteredSplitter

This class represents a splitter that splits clusters into treatment groups with a predefined number of treatment clusters. This is particularly useful for synthetic control analysis, where we only want 1 cluster ( unit) to be in treatment group and the rest in control The cluster that receives treatment remains random.

Attributes:

Name	Type	Description
cluster_cols	List[str]	List of columns to use as clusters.
n_treatment_clusters	int	The predefined number of treatment clusters.

Source code in cluster_experiments/random_splitter.py

class FixedSizeClusteredSplitter(ClusteredSplitter):
    """
    This class  represents a splitter that splits clusters into treatment groups with a predefined number of
    treatment clusters. This is particularly useful for synthetic control analysis, where we only want 1 cluster (
    unit) to be in treatment group and the rest in control The cluster that receives treatment remains random.

    Attributes:
        cluster_cols (List[str]): List of columns to use as clusters.
        n_treatment_clusters (int): The predefined number of treatment clusters.

    """

    def __init__(self, cluster_cols: List[str], n_treatment_clusters: int):
        super().__init__(cluster_cols=cluster_cols)
        self.n_treatment_clusters = n_treatment_clusters

    def sample_treatment(
        self,
        cluster_df: pd.DataFrame,
    ) -> List[str]:
        """
        Samples treatments for each cluster.

        Args:
            cluster_df (pd.DataFrame): Dataframe to assign treatments to.

        Returns:
            List[str]: A list of treatments for each cluster.
        """
        n_control_treatment = [
            len(cluster_df) - self.n_treatment_clusters,
            self.n_treatment_clusters,
        ]

        sample_treatment = [
            treatment
            for treatment, count in zip(self.treatments, n_control_treatment)
            for _ in range(count)
        ]
        random.shuffle(sample_treatment)
        return sample_treatment

sample_treatment(cluster_df)

Samples treatments for each cluster.

Parameters:

Name	Type	Description	Default
cluster_df	DataFrame	Dataframe to assign treatments to.	required

Returns:

Type	Description
List[str]	List[str]: A list of treatments for each cluster.

Source code in cluster_experiments/random_splitter.py

def sample_treatment(
    self,
    cluster_df: pd.DataFrame,
) -> List[str]:
    """
    Samples treatments for each cluster.

    Args:
        cluster_df (pd.DataFrame): Dataframe to assign treatments to.

    Returns:
        List[str]: A list of treatments for each cluster.
    """
    n_control_treatment = [
        len(cluster_df) - self.n_treatment_clusters,
        self.n_treatment_clusters,
    ]

    sample_treatment = [
        treatment
        for treatment, count in zip(self.treatments, n_control_treatment)
        for _ in range(count)
    ]
    random.shuffle(sample_treatment)
    return sample_treatment

NonClusteredSplitter

Bases: RandomSplitter

Splits randomly without clusters

Parameters:

Name	Type	Description	Default
treatments	Optional[List[str]]	list of treatments	None
treatment_col	str	Name of the column with the treatment variable.	'treatment'

Usage:

import pandas as pd
from cluster_experiments.random_splitter import NonClusteredSplitter
splitter = NonClusteredSplitter(
    treatments=["A", "B"],
)
df = pd.DataFrame({"city": ["A", "B", "C"]})
df = splitter.assign_treatment_df(df)
print(df)

Source code in cluster_experiments/random_splitter.py

class NonClusteredSplitter(RandomSplitter):
    """
    Splits randomly without clusters

    Arguments:
        treatments: list of treatments
        treatment_col: Name of the column with the treatment variable.

    Usage:
    ```python
    import pandas as pd
    from cluster_experiments.random_splitter import NonClusteredSplitter
    splitter = NonClusteredSplitter(
        treatments=["A", "B"],
    )
    df = pd.DataFrame({"city": ["A", "B", "C"]})
    df = splitter.assign_treatment_df(df)
    print(df)
    ```
    """

    def __init__(
        self,
        treatments: Optional[List[str]] = None,
        treatment_col: str = "treatment",
        splitter_weights: Optional[List[float]] = None,
    ) -> None:
        self.treatments = treatments or ["A", "B"]
        self.treatment_col = treatment_col
        self.splitter_weights = splitter_weights

    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()
        df[self.treatment_col] = random.choices(
            self.treatments, k=len(df), weights=self.splitter_weights
        )
        return df

    @classmethod
    def from_config(cls, config):
        """Creates a NonClusteredSplitter from a PowerConfig"""
        return cls(
            treatments=config.treatments,
            treatment_col=config.treatment_col,
            splitter_weights=config.splitter_weights,
        )

assign_treatment_df(df)

Takes a df, randomizes treatments and adds the treatment column to the dataframe

Parameters:

Name	Type	Description	Default
df	DataFrame	dataframe to assign treatments to	required

Source code in cluster_experiments/random_splitter.py

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()
    df[self.treatment_col] = random.choices(
        self.treatments, k=len(df), weights=self.splitter_weights
    )
    return df

from_config(config) classmethod

Creates a NonClusteredSplitter from a PowerConfig

Source code in cluster_experiments/random_splitter.py

@classmethod
def from_config(cls, config):
    """Creates a NonClusteredSplitter from a PowerConfig"""
    return cls(
        treatments=config.treatments,
        treatment_col=config.treatment_col,
        splitter_weights=config.splitter_weights,
    )

RandomSplitter

Bases: ABC

Abstract class to split instances in a switchback or clustered way. It can be used to create a calendar/split of clusters or to run a power analysis.

In order to create your own RandomSplitter, you should write your own assign_treatment_df method, that takes a dataframe as an input and returns the same dataframe with the treatment_col column.

Parameters:

Name	Type	Description	Default
cluster_cols	Optional[List[str]]	List of columns to use as clusters	None
treatments	Optional[List[str]]	list of treatments	None
treatment_col	str	Name of the column with the treatment variable.	'treatment'
splitter_weights	Optional[List[float]]	weights to use for the splitter, should have the same length as treatments, each weight should correspond to an element in treatments	None

Source code in cluster_experiments/random_splitter.py

class RandomSplitter(ABC):
    """
    Abstract class to split instances in a switchback or clustered way. It can be used to create a calendar/split of clusters
    or to run a power analysis.

    In order to create your own RandomSplitter, you should write your own assign_treatment_df method, that takes a dataframe as an input and returns the same dataframe with the treatment_col column.

    Arguments:
        cluster_cols: List of columns to use as clusters
        treatments: list of treatments
        treatment_col: Name of the column with the treatment variable.
        splitter_weights: weights to use for the splitter, should have the same length as treatments, each weight should correspond to an element in treatments

    """

    def __init__(
        self,
        cluster_cols: Optional[List[str]] = None,
        treatments: Optional[List[str]] = None,
        treatment_col: str = "treatment",
        splitter_weights: Optional[List[float]] = None,
    ) -> None:
        self.treatments = treatments or ["A", "B"]
        self.cluster_cols = cluster_cols or []
        self.treatment_col = treatment_col
        self.splitter_weights = splitter_weights

    @abstractmethod
    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
        """

    @classmethod
    def from_config(cls, config):
        """Creates a RandomSplitter from a PowerConfig"""
        return cls(
            treatments=config.treatments,
            cluster_cols=config.cluster_cols,
            treatment_col=config.treatment_col,
            splitter_weights=config.splitter_weights,
        )

assign_treatment_df(df) abstractmethod

Takes a df, randomizes treatments and adds the treatment column to the dataframe

Parameters:

Name	Type	Description	Default
df	DataFrame	dataframe to assign treatments to	required

Source code in cluster_experiments/random_splitter.py

@abstractmethod
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
    """

from_config(config) classmethod

Creates a RandomSplitter from a PowerConfig

Source code in cluster_experiments/random_splitter.py

@classmethod
def from_config(cls, config):
    """Creates a RandomSplitter from a PowerConfig"""
    return cls(
        treatments=config.treatments,
        cluster_cols=config.cluster_cols,
        treatment_col=config.treatment_col,
        splitter_weights=config.splitter_weights,
    )

RepeatedSampler

Bases: RandomSplitter

Doesn't actually split the data, but repeatedly samples (i.e. duplicates) all rows for all treatments. This is useful for backtesting, where we assume to have access to all counterfactuals.

Parameters:

Name	Type	Description	Default
treatments	Optional[List[str]]	list of treatments	None
treatment_col	str	Name of the column with the treatment variable.	'treatment'

Usage:

import pandas as pd
from cluster_experiments.random_splitter import RepeatedSampler
splitter = RepeatedSampler(
    treatments=["A", "B"],
)
df = pd.DataFrame({"city": ["A", "B", "C"]})
df = splitter.assign_treatment_df(df)
print(df)

Source code in cluster_experiments/random_splitter.py

class RepeatedSampler(RandomSplitter):
    """
    Doesn't actually split the data, but repeatedly samples (i.e. duplicates) all rows for all treatments.
    This is useful for backtesting, where we assume to have access to all counterfactuals.

    Arguments:
        treatments: list of treatments
        treatment_col: Name of the column with the treatment variable.

    Usage:
    ```python
    import pandas as pd
    from cluster_experiments.random_splitter import RepeatedSampler
    splitter = RepeatedSampler(
        treatments=["A", "B"],
    )
    df = pd.DataFrame({"city": ["A", "B", "C"]})
    df = splitter.assign_treatment_df(df)
    print(df)
    ```
    """

    def __init__(
        self,
        treatments: Optional[List[str]] = None,
        treatment_col: str = "treatment",
    ) -> None:
        self.treatments = treatments or ["A", "B"]
        self.treatment_col = treatment_col

    def assign_treatment_df(
        self,
        df: pd.DataFrame,
    ) -> pd.DataFrame:
        df = df.copy()

        dfs = []
        for treatment in self.treatments:
            df_treat = df.copy().assign(**{self.treatment_col: treatment})
            dfs.append(df_treat)

        return pd.concat(dfs).reset_index(drop=True)

    @classmethod
    def from_config(cls, config):
        """Creates a RepeatedSampler from a PowerConfig"""
        return cls(
            treatments=config.treatments,
            treatment_col=config.treatment_col,
        )

from_config(config) classmethod

Creates a RepeatedSampler from a PowerConfig

Source code in cluster_experiments/random_splitter.py

@classmethod
def from_config(cls, config):
    """Creates a RepeatedSampler from a PowerConfig"""
    return cls(
        treatments=config.treatments,
        treatment_col=config.treatment_col,
    )

StratifiedClusteredSplitter

Bases: RandomSplitter

Splits randomly with clusters, ensuring a balanced allocation of treatment groups across clusters and strata. To be used, for example, when having days as clusters and days of the week as stratus. This splitter will make sure that we won't have all Sundays in treatment and no Sundays in control.

Parameters:

Name	Type	Description	Default
cluster_cols	Optional[List[str]]	List of columns to use as clusters	None
treatments	Optional[List[str]]	list of treatments	None
treatment_col	str	Name of the column with the treatment variable.	'treatment'
strata_cols	Optional[List[str]]	List of columns to use as strata	None

Usage:

import pandas as pd
from cluster_experiments.random_splitter import StratifiedClusteredSplitter
splitter = StratifiedClusteredSplitter(cluster_cols=["city"],strata_cols=["country"])
df = pd.DataFrame({"city": ["A", "B", "C","D"], "country":["C1","C2","C2","C1"]})
df = splitter.assign_treatment_df(df)
print(df)

Source code in cluster_experiments/random_splitter.py

class StratifiedClusteredSplitter(RandomSplitter):
    """
    Splits randomly with clusters, ensuring a balanced allocation of treatment groups across clusters and strata.
    To be used, for example, when having days as clusters and days of the week as stratus. This splitter will make sure
    that we won't have all Sundays in treatment and no Sundays in control.

    Arguments:
        cluster_cols: List of columns to use as clusters
        treatments: list of treatments
        treatment_col: Name of the column with the treatment variable.
        strata_cols: List of columns to use as strata

    Usage:
    ```python
    import pandas as pd
    from cluster_experiments.random_splitter import StratifiedClusteredSplitter
    splitter = StratifiedClusteredSplitter(cluster_cols=["city"],strata_cols=["country"])
    df = pd.DataFrame({"city": ["A", "B", "C","D"], "country":["C1","C2","C2","C1"]})
    df = splitter.assign_treatment_df(df)
    print(df)
    ```
    """

    def __init__(
        self,
        cluster_cols: Optional[List[str]] = None,
        treatments: Optional[List[str]] = None,
        treatment_col: str = "treatment",
        strata_cols: Optional[List[str]] = None,
    ) -> None:
        super().__init__(
            cluster_cols=cluster_cols,
            treatments=treatments,
            treatment_col=treatment_col,
        )
        if not strata_cols or strata_cols == [""]:
            raise ValueError(
                f"Splitter {self.__class__.__name__} requires strata_cols,"
                f" got {strata_cols = }"
            )
        self.strata_cols = strata_cols

    def assign_treatment_df(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        df_unique_shuffled = (
            df.loc[:, list(set(self.cluster_cols + self.strata_cols))]
            .drop_duplicates()
            .sample(frac=1)
            .reset_index(drop=True)
        )

        # check that, for a given cluster, there is only 1 strata
        for strata_col in self.strata_cols:
            if (
                df_unique_shuffled.groupby(self.cluster_cols)[strata_col]
                .nunique()
                .max()
                > 1
            ):
                raise ValueError(
                    f"There are multiple values in {strata_col} for the same cluster item \n"
                    "You cannot stratify on this column",
                )

        # random shuffling
        random_sorted_treatments = list(np.random.permutation(self.treatments))

        df_unique_shuffled[self.treatment_col] = (
            df_unique_shuffled.groupby(self.strata_cols, as_index=False)
            .cumcount()
            .mod(len(random_sorted_treatments))
            .map(dict(enumerate(random_sorted_treatments)))
        )

        df = df.merge(
            df_unique_shuffled, on=self.cluster_cols + self.strata_cols, how="left"
        )

        return df

    @classmethod
    def from_config(cls, config):
        """Creates a StratifiedClusteredSplitter from a PowerConfig"""
        return cls(
            treatments=config.treatments,
            cluster_cols=config.cluster_cols,
            strata_cols=config.strata_cols,
            treatment_col=config.treatment_col,
        )

from_config(config) classmethod

Creates a StratifiedClusteredSplitter from a PowerConfig

Source code in cluster_experiments/random_splitter.py

@classmethod
def from_config(cls, config):
    """Creates a StratifiedClusteredSplitter from a PowerConfig"""
    return cls(
        treatments=config.treatments,
        cluster_cols=config.cluster_cols,
        strata_cols=config.strata_cols,
        treatment_col=config.treatment_col,
    )

StratifiedSwitchbackSplitter

Bases: StratifiedClusteredSplitter, SwitchbackSplitter

Splits randomly with clusters, ensuring a balanced allocation of treatment groups across clusters and strata. To be used, for example, when having days as clusters and days of the week as stratus. This splitter will make sure that we won't have all Sundays in treatment and no Sundays in control.

It can be created using the time_col and switch_frequency arguments, just like the SwitchbackSplitter.

Parameters:

Name	Type	Description	Default
time_col	str	Name of the column with the time variable.	'date'
switch_frequency	str	Frequency of the switchback. Must be a string (e.g. "1D")	'1D'
cluster_cols	Optional[List[str]]	List of columns to use as clusters	None
treatments	Optional[List[str]]	list of treatments	None
treatment_col	str	Name of the column with the treatment variable.	'treatment'
splitter_weights	Optional[List[float]]	List of weights for the treatments. If None, all treatments will have the same weight.	None
strata_cols	Optional[List[str]]	List of columns to use as strata	None

Usage:

import pandas as pd
from cluster_experiments.random_splitter import StratifiedSwitchbackSplitter
splitter = StratifiedSwitchbackSplitter(time_col="date",switch_frequency="1D",strata_cols=["country"], cluster_cols=["country", "date"])
df = pd.DataFrame({"date": ["2020-01-01", "2020-01-02", "2020-01-03","2020-01-04"], "country":["C1","C2","C2","C1"]})
df = splitter.assign_treatment_df(df)
print(df)

Source code in cluster_experiments/random_splitter.py

class StratifiedSwitchbackSplitter(StratifiedClusteredSplitter, SwitchbackSplitter):
    """
    Splits randomly with clusters, ensuring a balanced allocation of treatment groups across clusters and strata.
    To be used, for example, when having days as clusters and days of the week as stratus. This splitter will make sure
    that we won't have all Sundays in treatment and no Sundays in control.

    It can be created using the time_col and switch_frequency arguments, just like the SwitchbackSplitter.

    Arguments:
        time_col: Name of the column with the time variable.
        switch_frequency: Frequency of the switchback. Must be a string (e.g. "1D")
        cluster_cols: List of columns to use as clusters
        treatments: list of treatments
        treatment_col: Name of the column with the treatment variable.
        splitter_weights: List of weights for the treatments. If None, all treatments will have the same weight.
        strata_cols: List of columns to use as strata

    Usage:
    ```python
    import pandas as pd
    from cluster_experiments.random_splitter import StratifiedSwitchbackSplitter
    splitter = StratifiedSwitchbackSplitter(time_col="date",switch_frequency="1D",strata_cols=["country"], cluster_cols=["country", "date"])
    df = pd.DataFrame({"date": ["2020-01-01", "2020-01-02", "2020-01-03","2020-01-04"], "country":["C1","C2","C2","C1"]})
    df = splitter.assign_treatment_df(df)
    print(df)
    ```
    """

    def __init__(
        self,
        time_col: str = "date",
        switch_frequency: str = "1D",
        cluster_cols: Optional[List[str]] = None,
        treatments: Optional[List[str]] = None,
        treatment_col: str = "treatment",
        splitter_weights: Optional[List[float]] = None,
        washover: Optional[Washover] = None,
        strata_cols: Optional[List[str]] = None,
    ) -> None:
        # Inherit init from SwitchbackSplitter
        SwitchbackSplitter.__init__(
            self,
            time_col=time_col,
            switch_frequency=switch_frequency,
            cluster_cols=cluster_cols,
            treatments=treatments,
            treatment_col=treatment_col,
            splitter_weights=splitter_weights,
            washover=washover,
        )
        self.strata_cols = strata_cols or ["strata"]

    def assign_treatment_df(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        df = self._prepare_switchback_df(df)
        df = StratifiedClusteredSplitter.assign_treatment_df(self, df)
        return self.washover.washover(
            df=df,
            treatment_col=self.treatment_col,
            truncated_time_col=self.time_col,
            cluster_cols=self.cluster_cols,
        )

    @classmethod
    def from_config(cls, config) -> "StratifiedSwitchbackSplitter":
        """Creates a StratifiedSwitchbackSplitter from a PowerConfig"""
        washover_cls = _get_mapping_key(washover_mapping, config.washover)
        return cls(
            treatments=config.treatments,
            cluster_cols=config.cluster_cols,
            strata_cols=config.strata_cols,
            treatment_col=config.treatment_col,
            time_col=config.time_col,
            switch_frequency=config.switch_frequency,
            splitter_weights=config.splitter_weights,
            washover=washover_cls.from_config(config),
        )

from_config(config) classmethod

Creates a StratifiedSwitchbackSplitter from a PowerConfig

Source code in cluster_experiments/random_splitter.py

@classmethod
def from_config(cls, config) -> "StratifiedSwitchbackSplitter":
    """Creates a StratifiedSwitchbackSplitter from a PowerConfig"""
    washover_cls = _get_mapping_key(washover_mapping, config.washover)
    return cls(
        treatments=config.treatments,
        cluster_cols=config.cluster_cols,
        strata_cols=config.strata_cols,
        treatment_col=config.treatment_col,
        time_col=config.time_col,
        switch_frequency=config.switch_frequency,
        splitter_weights=config.splitter_weights,
        washover=washover_cls.from_config(config),
    )

SwitchbackSplitter

Bases: ClusteredSplitter

Splits randomly using clusters and time column

It is a clustered splitter but one of the cluster columns is obtained by truncating the time column to the switch frequency.

Parameters:

Name	Type	Description	Default
time_col	Optional[str]	Name of the column with the time variable.	None
switch_frequency	Optional[str]	Frequency to switch treatments. Uses pandas frequency aliases (https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases)	None
cluster_cols	Optional[List[str]]	List of columns to use as clusters	None
treatments	Optional[List[str]]	list of treatments	None
treatment_col	str	Name of the column with the treatment variable.	'treatment'
splitter_weights	Optional[List[float]]	weights to use for the splitter, should have the same length as treatments, each weight should correspond to an element in treatments	None

Usage:

import pandas as pd
from cluster_experiments.random_splitter import SwitchbackSplitter
splitter = SwitchbackSplitter(time_col="date", switch_frequency="1D", cluster_cols=["date"])
df = pd.DataFrame({"date": pd.date_range("2020-01-01", "2020-01-03")})
df = splitter.assign_treatment_df(df)
print(df)

Source code in cluster_experiments/random_splitter.py

class SwitchbackSplitter(ClusteredSplitter):
    """
    Splits randomly using clusters and time column

    It is a clustered splitter but one of the cluster columns is obtained by truncating the time column to the switch frequency.

    Arguments:
        time_col: Name of the column with the time variable.
        switch_frequency: Frequency to switch treatments. Uses pandas frequency aliases (https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases)
        cluster_cols: List of columns to use as clusters
        treatments: list of treatments
        treatment_col: Name of the column with the treatment variable.
        splitter_weights: weights to use for the splitter, should have the same length as treatments, each weight should correspond to an element in treatments

    Usage:
    ```python
    import pandas as pd
    from cluster_experiments.random_splitter import SwitchbackSplitter
    splitter = SwitchbackSplitter(time_col="date", switch_frequency="1D", cluster_cols=["date"])
    df = pd.DataFrame({"date": pd.date_range("2020-01-01", "2020-01-03")})
    df = splitter.assign_treatment_df(df)
    print(df)
    ```
    """

    def __init__(
        self,
        time_col: Optional[str] = None,
        switch_frequency: Optional[str] = None,
        cluster_cols: Optional[List[str]] = None,
        treatments: Optional[List[str]] = None,
        treatment_col: str = "treatment",
        splitter_weights: Optional[List[float]] = None,
        washover: Optional[Washover] = None,
    ) -> None:
        self.time_col = time_col or "date"
        self.switch_frequency = switch_frequency or "1D"
        self.cluster_cols = cluster_cols or []
        self.treatments = treatments or ["A", "B"]
        self.treatment_col = treatment_col
        self.splitter_weights = splitter_weights
        self.washover = washover or EmptyWashover()
        self._check_clusters()

    def _check_clusters(self):
        """Check if time_col is in cluster_cols"""
        assert (
            self.time_col in self.cluster_cols
        ), "in switchback splitters, time_col must be in cluster_cols"

    def _get_time_col_cluster(self, df: pd.DataFrame) -> pd.Series:
        df = df.copy()
        df[self.time_col] = pd.to_datetime(df[self.time_col])
        # Given the switch frequency, truncate the time column to the switch frequency
        # Using pandas frequency aliases: https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases
        if "W" in self.switch_frequency or "M" in self.switch_frequency:
            return df[self.time_col].dt.to_period(self.switch_frequency).dt.start_time
        return df[self.time_col].dt.floor(self.switch_frequency)

    def _prepare_switchback_df(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        # Build time_col switchback column
        # Overwriting column, this is the worst! If we use the column as a covariate, we're screwed. Needs improvement
        df[_original_time_column(self.time_col)] = df[self.time_col]
        df[self.time_col] = self._get_time_col_cluster(df)
        return df

    def assign_treatment_df(
        self,
        df: pd.DataFrame,
    ) -> pd.DataFrame:
        """
        Creates the switchback column, adds it to cluster_cols and then calls ClusteredSplitter assign_treatment_df

        Arguments:
            df: dataframe to assign treatments to
        """
        df = df.copy()
        df = self._prepare_switchback_df(df)
        df = super().assign_treatment_df(df)
        df = self.washover.washover(
            df,
            truncated_time_col=self.time_col,
            treatment_col=self.treatment_col,
            cluster_cols=self.cluster_cols,
        )
        return df

    @classmethod
    def from_config(cls, config) -> "SwitchbackSplitter":
        washover_cls = _get_mapping_key(washover_mapping, config.washover)
        return cls(
            time_col=config.time_col,
            switch_frequency=config.switch_frequency,
            cluster_cols=config.cluster_cols,
            treatments=config.treatments,
            treatment_col=config.treatment_col,
            splitter_weights=config.splitter_weights,
            washover=washover_cls.from_config(config),
        )

assign_treatment_df(df)

Creates the switchback column, adds it to cluster_cols and then calls ClusteredSplitter assign_treatment_df

Parameters:

Name	Type	Description	Default
df	DataFrame	dataframe to assign treatments to	required

Source code in cluster_experiments/random_splitter.py

def assign_treatment_df(
    self,
    df: pd.DataFrame,
) -> pd.DataFrame:
    """
    Creates the switchback column, adds it to cluster_cols and then calls ClusteredSplitter assign_treatment_df

    Arguments:
        df: dataframe to assign treatments to
    """
    df = df.copy()
    df = self._prepare_switchback_df(df)
    df = super().assign_treatment_df(df)
    df = self.washover.washover(
        df,
        truncated_time_col=self.time_col,
        treatment_col=self.treatment_col,
        cluster_cols=self.cluster_cols,
    )
    return df