balance: a python package for balancing biased data samples

[!NOTE] balance is currently in beta and is actively supported. Follow us on github.

What is balance?

balance is a Python package offering a simple workflow and methods for dealing with biased data samples when looking to infer from them to some population of interest.

Biased samples often occur in survey statistics when respondents present non-response bias or survey suffers from sampling bias (that are not missing completely at random). A similar issue arises in observational studies when comparing the treated vs untreated groups, and in any data that suffers from selection bias.

Under the missing at random assumption (MAR), bias in samples could sometimes be (at least partially) mitigated by relying on auxiliary information (a.k.a.: "covariates" or "features") that is present for all items in the sample, as well as present in a sample of items from the population. For example, if we want to infer from a sample of respondents to some survey, we may wish to adjust for non-response using demographic information such as age, gender, education, etc. This can be done by weighing the sample to the population using auxiliary information.

The package is intended for researchers who are interested in balancing biased samples, such as the ones coming from surveys, using a Python package. This need may arise by survey methodologists, demographers, UX researchers, market researchers, and generally data scientists, statisticians, and machine learners.

More about the methodological background can be found in Sarig, T., Galili, T., & Eilat, R. (2023). balance – a Python package for balancing biased data samples.

Installation

Requirements

You need Python 3.9, 3.10, 3.11, 3.12, 3.13, or 3.14 to run balance. balance can be built and run from Linux, OSX, and Windows.

The required Python dependencies are:

REQUIRES = [
    # Numpy and pandas: carefully versioned for binary compatibility
    "numpy>=1.21.0,<2.0; python_version<'3.12'",
    "numpy>=1.24.0; python_version>='3.12'",
    "pandas>=1.5.0,<4.0.0; python_version<'3.12'",
    "pandas>=2.0.0,<4.0.0; python_version>='3.12'",
    # Scientific stack
    "scipy>=1.7.0,<1.14.0; python_version<'3.12'",
    "scipy>=1.11.0; python_version>='3.12'",
    "scikit-learn>=1.0.0,<1.4.0; python_version<'3.12'",
    "scikit-learn>=1.3.0; python_version>='3.12'",
    "ipython",
    "patsy",
    "seaborn",
    "plotly",
    "matplotlib",
    "statsmodels",
    "session-info",
]

See pyproject.toml for more details.

Installing balance

Installing via PyPi

We recommend installing balance from PyPi via pip for the latest stable version:

python -m pip install balance

Installation will use Python wheels from PyPI, available for OSX, Linux, and Windows.

Installing from Source/Git

You can install the latest (bleeding edge) version from Git:

python -m pip install git+https://github.com/facebookresearch/balance.git

Alternatively, if you have a local clone of the repo:

cd balance
python -m pip install .

Or using dev-dependencies:

cd balance
python -m pip install .[dev]

Getting started

balance's workflow in high-level

The core workflow in balance deals with fitting and evaluating weights to a sample. For each unit in the sample (such as a respondent to a survey), balance fits a weight that can be (loosely) interpreted as the number of people from the target population that this respondent represents. This aims to help mitigate the coverage and non-response biases, as illustrated in the following figure.

The weighting of survey data through balance is done in the following main steps:

1. Loading data of the respondents of the survey.
2. Loading data about the target population we would like to correct for.
3. Diagnostics of the sample covariates so to evaluate whether weighting is needed.
4. Adjusting the sample to the target.
5. Evaluation of the results.
6. Use the weights for producing population level estimations.
7. Saving the output weights.

You can see a step-by-step description (with code) of the above steps in the General Framework page.

Code example of using balance

You may run the following code to play with balance's basic workflow (these are snippets taken from the quickstart tutorial):

We start by loading data, and adjusting it:

from balance import load_data, Sample

# load simulated example data
target_df, sample_df = load_data()

# Import sample and target data into a Sample object
sample = Sample.from_frame(sample_df, outcome_columns=["happiness"])
target = Sample.from_frame(target_df)

# Set the target to be the target of sample
sample_with_target = sample.set_target(target)

# Check basic diagnostics of sample vs target before adjusting:
# sample_with_target.covars().plot()

You can read more on evaluation of the pre-adjusted data in the Pre-Adjustment Diagnostics page.

Next, we adjust the sample to the population by fitting balancing survey weights:

# Using ipw to fit survey weights
adjusted = sample_with_target.adjust()

You can read more on adjustment process in the Adjusting Sample to Population page.

The above code gets us an adjusted object with weights. We can evaluate the benefit of the weights to the covariate balance, for example by running:

print(adjusted.summary())
    # Covar ASMD reduction: 62.3%, design effect: 2.249
    # Covar ASMD (7 variables):0.335 -> 0.126
    # Model performance: Model proportion deviance explained: 0.174

adjusted.covars().plot(library = "seaborn", dist_type = "kde")

And get:

We can also check the impact of the weights on the outcome using:

# For the outcome:
print(adjusted.outcomes().summary())
    # 1 outcomes: ['happiness']
    # Mean outcomes:
    #             happiness
    # source
    # self        54.221388
    # unadjusted  48.392784
    #
    # Response rates (relative to number of respondents in sample):
    #    happiness
    # n     1000.0
    # %      100.0
adjusted.outcomes().plot()

You can read more on evaluation of the post-adjusted data in the Evaluating and using the adjustment weights page.

Finally, the adjusted data can be downloaded using:

adjusted.to_download()  # Or:
# adjusted.to_csv()

To see a more detailed step-by-step code example with code output prints and plots (both static and interactive), please go over to the tutorials section.

Implemented methods for adjustments

balance currently implements various adjustment methods. Click the links to learn more about each:

1. Logistic regression using L1 (LASSO) penalization.
2. Covariate Balancing Propensity Score (CBPS).
3. Post-stratification.
4. Raking.

Implemented methods for diagnostics/evaluation

For diagnostics the main tools (comparing before, after applying weights, and the target population) are:

1. Plots
   1. barplots
   2. density plots (for weights and covariances)
   3. qq-plots
   4. love plot — per-covariate ASMD before-vs-after on a sorted scatter with a +0.1 reference cutoff (ASMD is non-negative, so only the positive threshold line is drawn), in the spirit of R's cobalt::love.plot (added in v0.21 via BalanceDFCovars.love_plot())
2. Statistical summaries
   1. Weights distributions
      1. Kish's design effect
      2. Main summaries (mean, median, variances, quantiles)
   2. Covariate distributions
      1. Absolute Standardized Mean Difference (ASMD). For continuous variables, it is Cohen's d. Categorical variables are one-hot encoded, Cohen's d is calculated for each category and ASMD for a categorical variable is defined as Cohen's d, average across all categories.

You can read more on evaluation of the post-adjusted data in the Evaluating and using the adjustment weights page.

Estimating a population outcome: IPW vs. outcome-model estimators

Once a sample is reweighted to a target, balance offers more than one way to estimate the target-population mean of an outcome. They use the same covariates differently, so comparing them is a useful robustness check:

• μ̂_IPW — inverse-propensity (Hájek) weighted mean: outcomes().mean() (the adjusted self row). The weighted average of the responders' observed outcome using the adjusted weights. Consistent when the weights (propensity model) are correct.
• μ̂_OM — outcome-model / g-computation estimate: outcomes_hat().mean() (the target-population row). Fit a learner ĝ(X) ≈ E[Y|X] on the responders, apply it to the target covariates, and average the predicted outcomes ŷ_T with the target weights. Consistent when the outcome model is correct. Added in v0.23.

The two estimators are complementary — agreement between μ̂_IPW and μ̂_OM is reassuring; disagreement points at a misspecified weighting or outcome model. A doubly-robust / AIPW estimator that combines both is available as bf.aipw() (point estimate only; variance/CI and cross-fitting remain deferred).

For a runnable end-to-end walkthrough, see the outcome-model tutorial.

from balance import Sample

bf = (
    Sample.from_frame(sample_df, outcome_columns=["happiness"])
    .set_target(Sample.from_frame(target_df))
    .adjust(method="ipw")
)

# μ̂_IPW — weighted mean of the observed outcome (`self` row):
bf.outcomes().mean()

# μ̂_OM — fit ĝ(X) on responders, apply to the target, average ŷ_T (target row):
bf.fit_outcome_model()
bf.predict_outcomes(on="target")
bf.outcomes_hat().mean()

# An honest confidence interval for μ̂_OM via a nonparametric bootstrap
# (resample responders, refit ĝ*, predict on the fixed target, re-average):
bf.outcomes_hat().mean_with_ci(ci_method="bootstrap", n_bootstrap=200, random_seed=2020)

# outcomes_hat().summary() reports the estimator and *scopes* any doubly-robust
# claim to the fit weights — it never prints a blanket "doubly robust" (a plain
# g-computation for the non-linear default; DR only for a weighted linear+intercept fit):
bf.outcomes_hat().summary()

Train/holdout transfer. You can fit the outcome model on one frame (a train split) and apply it to a different frame (a holdout / scoring split) with the same covariate schema, mirroring how set_fitted_model transfers a fitted IPW model. set_fitted_outcome_model copies the already-fitted model onto the holdout without re-fitting (the fitted learner is shared by identity), so predicting on the holdout target gives μ̂_OM computed with the train model:

# train_bf carries a fitted outcome model (train_bf.fit_outcome_model(...));
# holdout_bf has the same covariates plus its own target to score:
scored = holdout_bf.set_fitted_outcome_model(train_bf, inplace=False)
scored.predict_outcomes(on="target")
scored.outcomes_hat().mean()      # μ̂_OM on the holdout target via train_bf's model

Separately, bf.outcomes().weights_impact_on_outcome_ss() is a diagnostic (not a third estimator): it reports how much applying the weights moves the outcome mean (weighted vs. unweighted), with a significance test — useful for judging whether the reweighting materially changed the outcome estimate. See the outcome-model design doc for the estimator theory and the AIPW roadmap.

Design-based inference

balance complements an adjacent library that handles the design-based inference step balance does not — covering the non-probability → DiD half of a survey-based campaign or policy evaluation in Python.

Survey-weighted Difference-in-Differences (diff-diff)

balance pairs naturally with diff-diff, the open-source Python package for modern Difference-in-Differences (Callaway & Sant'Anna 2021, Sun & Abraham 2021, Borusyak-Jaravel-Spiess 2024, Synthetic DiD, Continuous DiD, Triple Difference) with built-in survey-design variance. The two libraries solve adjacent halves of a survey-based campaign or policy evaluation:

• balance produces non-probability weights for a sample against a target population frame (IPW, CBPS, rake, post-stratification).
• diff-diff consumes those weights via its SurveyDesign and returns design-consistent ATT(g, t) with HonestDiD sensitivity.

The thin adapter balance.interop.diff_diff (added in the upcoming v0.21 release; see CHANGELOG.md) turns the handoff into a single import. Install the optional extra:

pip install "balance[did]"

Then the canonical workflow is:

from balance import Sample
from balance.interop.diff_diff import fit_did

# `set_target(...)` requires a Sample/SampleFrame, not a raw DataFrame --
# wrap your target DataFrame with Sample.from_frame(...) first.
s = (
    Sample.from_frame(df, ...)
    .set_target(Sample.from_frame(target_df))
    .adjust(method="ipw")
)
results = fit_did(s, estimator="CallawaySantAnna", outcome="y", time="t",
                  unit="state", treatment_first="first_treat",
                  estimation_method="dr")

fit_did builds the SurveyDesign from the active balance weight column, strips balance's history columns, and routes kwargs into the chosen diff-diff estimator. For a complete walk-through on a BRFSS-shaped public-health panel (including HonestDiD sensitivity), see the tutorial: tutorials/balance_diff_diff_brfss.ipynb. Upstream project: github.com/igerber/diff-diff.

Developer and AI assistant resources

• ARCHITECTURE.md — internal architecture documentation (class hierarchy, data flow, design decisions)
• .github/copilot-instructions.md — instructions for LLM/AI coding assistants working on balance

Other resources

• Presentation: "Balancing biased data samples with the 'balance' Python package" - presented in the Israeli Statistical Association (ISA) conference on June 1st 2023.

More details

Getting help, submitting bug reports and contributing code

You are welcome to:

• Learn more in the balance website.
• Ask for help on: https://github.com/facebookresearch/balance/issues/new?template=support_question.md
• Submit bug-reports and features' suggestions at: https://github.com/facebookresearch/balance/issues
• Send a pull request on: https://github.com/facebookresearch/balance. See the CONTRIBUTING file for how to help out. And our CODE OF CONDUCT for our expectations from contributors.

Citing balance

Sarig, T., Galili, T., & Eilat, R. (2023). balance – a Python package for balancing biased data samples. https://arxiv.org/abs/2307.06024

@misc{sarig2023balance,
      title={balance - a Python package for balancing biased data samples},
      author={Tal Sarig and Tal Galili and Roee Eilat},
      year={2023},
      eprint={2307.06024},
      archivePrefix={arXiv},
      primaryClass={stat.CO}
}

License

The balance package is licensed under the MIT license, and all the documentation on the site (including text and images) is under CC-BY.

News

You can follow updates on our:

• Blog
• Changelog

Acknowledgements / People

The balance package is actively maintained by people from the Central Applied Science team (in Menlo Park and Tel Aviv), by Wesley Lee, Tal Sarig, and Tal Galili.

The balance package was (and is) developed by many people, including: Roee Eilat, Tal Galili, Daniel Haimovich, Kevin Liou, Steve Mandala, Adam Obeng (author of the initial internal Meta version), Tal Sarig, Luke Sonnet, Sean Taylor, Barak Yair Reif, Soumyadip Sarkar, and others. If you worked on balance in the past, please email us to be added to this list.

The balance package was open-sourced by Tal Sarig, Tal Galili and Steve Mandala in late 2022.

Branding created by Dana Beaty, from the Meta AI Design and Marketing Team. For logo files, see here.