Skip to content

This is a Python tool to employ stratified sampling or treatment randomization with uneven numbers in some strata using pandas. Mainly thought with RCTs in mind, it also works for any other scenario in where you would like to randomly allocate treatment within blocks or strata. The tool also supports having multiple treatments with different pro…

License

Notifications You must be signed in to change notification settings

manmartgarc/stochatreat

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Stochatreat

Build Main Branch Tests codecov
PyPI pypi pypi-downloads
conda-forge Conda conda-downloads
Meta Hatch project linting - Ruff types - Mypy License - MIT

Introduction

This is a Python tool to employ stratified randomization or sampling with uneven numbers in some strata using pandas. Mainly thought with randomized controlled trials (RCTs) in mind, it also works for any other scenario in where you would like to randomly allocate treatment within blocks or strata. The tool also supports having multiple treatments with different probability of assignment within each block or stratum.

Installation

PyPI

You can install this package via pip:

pip install stochatreat

Conda

You can also install this package with conda:

conda install -c conda-forge stochatreat

Usage

Single cluster:

from stochatreat import stochatreat
import numpy as np
import pandas as pd

# make 1000 households in 5 different neighborhoods.
np.random.seed(42)
df = pd.DataFrame(
    data={"id": list(range(1000)), "nhood": np.random.randint(1, 6, size=1000)}
)

# randomly assign treatments by neighborhoods.
treats = stochatreat(
    data=df,  # your dataframe
    stratum_cols="nhood",  # the blocking variable
    treats=2,  # including control
    idx_col="id",  # the unique id column
    random_state=42,  # random seed
    misfit_strategy="stratum",
)  # the misfit strategy to use
# merge back with original data
df = df.merge(treats, how="left", on="id")

# check for allocations
df.groupby("nhood")["treat"].value_counts().unstack()

# previous code should return this
treat    0    1
nhood
1      105  105
2       95   95
3       95   95
4      103  103
5      102  102

Multiple clusters and treatment probabilities:

from stochatreat import stochatreat
import numpy as np
import pandas as pd

# make 1000 households in 5 different neighborhoods, with a dummy indicator
np.random.seed(42)
df = pd.DataFrame(
    data={
        "id": list(range(1000)),
        "nhood": np.random.randint(1, 6, size=1000),
        "dummy": np.random.randint(0, 2, size=1000),
    }
)

# randomly assign treatments by neighborhoods and dummy status.
treats = stochatreat(
    data=df,
    stratum_cols=["nhood", "dummy"],
    treats=2,
    probs=[1 / 3, 2 / 3],
    idx_col="id",
    random_state=42,
    misfit_strategy="global",
)
# merge back with original data
df = df.merge(treats, how="left", on="id")

# check for allocations
df.groupby(["nhood", "dummy"])["treat"].value_counts().unstack()

# previous code should return this
treat         0   1
nhood dummy
1     0      37  75
      1      33  65
2     0      35  69
      1      29  57
3     0      30  58
      1      34  68
4     0      36  72
      1      32  66
5     0      33  68
      1      35  68

Contributing

If you'd like to contribute to the package, make sure you read the contributing guide.

References

About

This is a Python tool to employ stratified sampling or treatment randomization with uneven numbers in some strata using pandas. Mainly thought with RCTs in mind, it also works for any other scenario in where you would like to randomly allocate treatment within blocks or strata. The tool also supports having multiple treatments with different pro…

Topics

Resources

License

Stars

Watchers

Forks

Packages

No packages published

Languages