balance Quickstart (raking): Analyzing and adjusting the bias on a simulated toy dataset¶

The raking method is an advanced technique that extends post-stratification. It is well-suited for situations where we have marginal distributions of multiple covariates and we don't know the joint distribution. Raking works by applying post-stratification to the data based on the first covariate, using the resulting output weights as input for adjustment based on the second covariate, and so forth. Once all covariates have been utilized for adjustment, the process is repeated until a specified level of convergence is attained

One of the main advantages of raking is its ability to work with user-level data while also utilizing marginal distributions that lack user-level granularity. Another benefit is its capacity to closely fit these distributions, depending on the convergence achieved. This is in contrast to techniques such as inverse probability weighting (IPW) and covariate balancing propensity score (CBPS), which may only approximate the data and potentially fail to fit them even at marginal levels.

This notebook demonstrates how to use the raking method and showcases the high degree of fit it can provide.

Load the data¶

In [1]:
%matplotlib inline

import plotly.offline as offline
offline.init_notebook_mode()

from balance import load_data

INFO (2026-08-19 01:11:01,756) [__init__/<module> (line 77)]: Using balance version 0.23.0

INFO (2026-08-19 01:11:01,757) [__init__/<module> (line 82)]: 
balance (Version 0.23.0) loaded:
    📖 Documentation: https://import-balance.org/
    🛠️ Help / Issues: https://github.com/facebookresearch/balance/issues/
    📄 Citation:
        Sarig, T., Galili, T., & Eilat, R. (2023).
        balance - a Python package for balancing biased data samples.
        https://arxiv.org/abs/2307.06024

    Tip: You can view this message anytime with balance.help()
In [2]:
target_df, sample_df = load_data()

print("target_df: \n", target_df.head())
print("sample_df: \n", sample_df.head())

target_df: 
        id gender age_group     income  happiness
0  100000   Male       45+  10.183951  61.706333
1  100001   Male       45+   6.036858  79.123670
2  100002   Male     35-44   5.226629  44.206949
3  100003    NaN       45+   5.752147  83.985716
4  100004    NaN     25-34   4.837484  49.339713
sample_df: 
   id  gender age_group     income  happiness
0  0    Male     25-34   6.428659  26.043029
1  1  Female     18-24   9.940280  66.885485
2  2    Male     18-24   2.673623  37.091922
3  3     NaN     18-24  10.550308  49.394050
4  4     NaN     18-24   2.689994  72.304208
In [3]:
from balance import Sample

Raking can work with numerical variables since the variable is automatically bucketed. But for the simplicity of the discussion, we'll focus only on age and gender.

In [4]:
sample = Sample.from_frame(sample_df[['id', 'gender', 'age_group', 'happiness']], outcome_columns=["happiness"])
target = Sample.from_frame(target_df[['id', 'gender', 'age_group', 'happiness']], outcome_columns=["happiness"])
sample_with_target = sample.set_target(target)

WARNING (2026-08-19 01:11:01,793) [input_validation/guess_id_column (line 336)]: Guessed id column name id for the data

WARNING (2026-08-19 01:11:01,805) [sample_frame/from_frame (line 380)]: No weights passed. Adding a 'weight' column and setting all values to 1

WARNING (2026-08-19 01:11:01,808) [input_validation/guess_id_column (line 336)]: Guessed id column name id for the data

WARNING (2026-08-19 01:11:01,823) [sample_frame/from_frame (line 380)]: No weights passed. Adding a 'weight' column and setting all values to 1

Fit models using ipw and rake¶

Fit an ipw model:

In [5]:
adjusted_ipw = sample_with_target.adjust(method = "ipw")

INFO (2026-08-19 01:11:01,837) [ipw/ipw (line 735)]: Starting ipw function

INFO (2026-08-19 01:11:01,842) [adjustment/apply_transformations (line 435)]: Adding the variables: []

INFO (2026-08-19 01:11:01,842) [adjustment/apply_transformations (line 436)]: Transforming the variables: ['gender', 'age_group']

INFO (2026-08-19 01:11:01,848) [adjustment/apply_transformations (line 472)]: Final variables in output: ['gender', 'age_group']

INFO (2026-08-19 01:11:01,931) [ipw/ipw (line 812)]: Building model matrix

INFO (2026-08-19 01:11:01,932) [ipw/ipw (line 813)]: The formula used to build the model matrix: ['gender + age_group + _is_na_gender']

INFO (2026-08-19 01:11:01,932) [ipw/ipw (line 814)]: The number of columns in the model matrix: 7

INFO (2026-08-19 01:11:01,933) [ipw/ipw (line 815)]: The number of rows in the model matrix: 11000

INFO (2026-08-19 01:11:20,111) [ipw/ipw (line 1011)]: Done with sklearn

INFO (2026-08-19 01:11:20,112) [ipw/ipw (line 1013)]: max_de: None

INFO (2026-08-19 01:11:20,113) [ipw/ipw (line 1035)]: Starting model selection

INFO (2026-08-19 01:11:20,115) [ipw/ipw (line 1091)]: Chosen lambda: 0.041158338186664825

INFO (2026-08-19 01:11:20,116) [ipw/ipw (line 1108)]: Proportion null deviance explained 0.11579381555381918

Fit a raking model (on the user level data as input):

In [6]:
adjusted_rake = sample_with_target.adjust(method = "rake")

INFO (2026-08-19 01:11:20,136) [adjustment/apply_transformations (line 435)]: Adding the variables: []

INFO (2026-08-19 01:11:20,137) [adjustment/apply_transformations (line 436)]: Transforming the variables: ['gender', 'age_group']

INFO (2026-08-19 01:11:20,144) [adjustment/apply_transformations (line 472)]: Final variables in output: ['gender', 'age_group']

INFO (2026-08-19 01:11:20,154) [rake/rake (line 623)]: Final covariates and levels that will be used in raking: {'age_group': ['18-24', '25-34', '35-44', '45+'], 'gender': ['Female', 'Male', '__NaN__']}.

Rake model-glance diagnostics¶

The regular diagnostics table also includes compact model_glance rows for a fitted rake model. To persist all replay metadata (including the number of variables), fit the same data through BalanceFrame.fit(method="rake") and inspect the diagnostics table.

In [7]:
from balance.balance_frame import BalanceFrame
from balance.sample_frame import SampleFrame

rake_resp = SampleFrame.from_frame(
    sample_df[["id", "gender", "age_group", "happiness"]],
    id_column="id",
    outcome_columns=["happiness"],
)
rake_tgt = SampleFrame.from_frame(
    target_df[["id", "gender", "age_group", "happiness"]],
    id_column="id",
    outcome_columns=["happiness"],
)
adjusted_rake_bf = BalanceFrame(sample=rake_resp, target=rake_tgt).fit(
    method="rake",
    variables=["gender", "age_group"],
)
rake_glance = adjusted_rake_bf.diagnostics().query("metric == 'model_glance'")
rake_glance[["var", "val"]]

WARNING (2026-08-19 01:11:20,195) [sample_frame/from_frame (line 380)]: No weights passed. Adding a 'weight' column and setting all values to 1

WARNING (2026-08-19 01:11:20,217) [sample_frame/from_frame (line 380)]: No weights passed. Adding a 'weight' column and setting all values to 1

WARNING (2026-08-19 01:11:20,227) [adjustment/_warn_default_transformations_not_transferable (line 649)]: rake(store_fit_metadata=True) is being used together with transformations='default'. The fitted model can be replayed in-place via BalanceFrame.predict_weights(), but transfer scoring via predict_weights(data=...) will raise. Pass deterministic transformations at fit time to enable transfer.

INFO (2026-08-19 01:11:20,228) [adjustment/apply_transformations (line 435)]: Adding the variables: []

INFO (2026-08-19 01:11:20,229) [adjustment/apply_transformations (line 436)]: Transforming the variables: ['gender', 'age_group']

INFO (2026-08-19 01:11:20,234) [adjustment/apply_transformations (line 472)]: Final variables in output: ['gender', 'age_group']

INFO (2026-08-19 01:11:20,242) [rake/rake (line 623)]: Final covariates and levels that will be used in raking: {'age_group': ['18-24', '25-34', '35-44', '45+'], 'gender': ['Female', 'Male', '__NaN__']}.

INFO (2026-08-19 01:11:20,256) [balance_frame/diagnostics (line 3904)]: Starting computation of diagnostics of the fitting

INFO (2026-08-19 01:11:20,539) [balance_frame/diagnostics (line 3931)]: Done computing diagnostics
Out[7]:
var val
37 converged 1.000000
38 iterations 3.000000
39 final_conv 0.000059
40 n_variables 2.000000

When comparing the results of ipw and rake, we can see that rake has a larger design effect, and that it provides a perfect fit. In contrast, ipw gives only a partial fit.

We can see it in the ASMD and also the bar plots.

In [8]:
print(adjusted_ipw.summary())

Adjustment details:
    method: ipw
    weight trimming mean ratio: 20
Covariate diagnostics:
    Covar ASMD reduction: 77.6%
    Covar ASMD (6 variables): 0.243 -> 0.054
    Covar mean KLD reduction: 92.2%
    Covar mean KLD (2 variables): 0.179 -> 0.014
Weight diagnostics:
    design effect (Deff): 1.527
    effective sample size proportion (ESSP): 0.655
    effective sample size (ESS): 654.8
Outcome weighted means:
            happiness
source               
self           53.889
target         56.278
unadjusted     48.559
Model performance: Model proportion deviance explained: 0.116
In [9]:
print(adjusted_rake.summary())

Adjustment details:
    method: rake
Covariate diagnostics:
    Covar ASMD reduction: 100.0%
    Covar ASMD (6 variables): 0.243 -> 0.000
    Covar mean KLD reduction: 100.0%
    Covar mean KLD (2 variables): 0.179 -> 0.000
Weight diagnostics:
    design effect (Deff): 2.103
    effective sample size proportion (ESSP): 0.476
    effective sample size (ESS): 475.6
Outcome weighted means:
            happiness
source               
self           55.484
target         56.278
unadjusted     48.559
In [10]:
adjusted_ipw.covars().plot()
In [11]:
adjusted_rake.covars().plot()

Outcome analysis¶

In [12]:
print(adjusted_ipw.outcomes().summary())
adjusted_ipw.outcomes().plot()

1 outcomes: ['happiness']
Mean outcomes (with 95% confidence intervals):
source       self  target  unadjusted           self_ci         target_ci     unadjusted_ci
happiness  53.889  56.278      48.559  (52.736, 55.042)  (55.961, 56.595)  (47.669, 49.449)

Weights impact on outcomes (t_test):
           mean_yw0  mean_yw1  mean_diff  diff_ci_lower  diff_ci_upper  t_stat  p_value       n
outcome                                                                                        
happiness    48.559    53.889       5.33           2.58          8.081   3.803      0.0  1000.0

Response rates (relative to number of respondents in sample):
   happiness
n     1000.0
%      100.0
Response rates (relative to notnull rows in the target):
    happiness
n     1000.0
%       10.0
Response rates (in the target):
    happiness
n    10000.0
%      100.0

The above shows the estimated mean happiness for our sample: unadjusted, IPW-adjusted, and target values. The following shows the corresponding happiness outcomes after raking:

In [13]:
print(adjusted_rake.outcomes().summary())
adjusted_rake.outcomes().plot()

1 outcomes: ['happiness']
Mean outcomes (with 95% confidence intervals):
source       self  target  unadjusted           self_ci         target_ci     unadjusted_ci
happiness  55.484  56.278      48.559  (54.173, 56.796)  (55.961, 56.595)  (47.669, 49.449)

Weights impact on outcomes (t_test):
           mean_yw0  mean_yw1  mean_diff  diff_ci_lower  diff_ci_upper  t_stat  p_value       n
outcome                                                                                        
happiness    48.559    55.484      6.926          2.827         11.024   3.316    0.001  1000.0

Response rates (relative to number of respondents in sample):
   happiness
n     1000.0
%      100.0
Response rates (relative to notnull rows in the target):
    happiness
n     1000.0
%       10.0
Response rates (in the target):
    happiness
n    10000.0
%      100.0

As we can see, both IPW and raking impact the outcome estimate. Raking achieves exact balance on the target marginals (as shown by the perfect ASMD scores earlier), which can be important when precise matching to known population distributions is required.

Using target margins with rake¶

The benefit of rake is that we can fit directly to known marginal target totals without first constructing a row-level target DataFrame. Pass those totals to rake() with target_margins. Each variable must sum to the same positive target total.

In order to demonstrate this point, let us assume we have another target population in mind, with different marginal totals for gender and age group.

In [14]:
from balance.weighting_methods.rake import rake
import numpy as np

target_margins = {
    "gender": {"Female": 100.0, "Male": 850.0, np.nan: 50.0},
    "age_group": {"18-24": 250.0, "25-34": 250.0, "35-44": 250.0, "45+": 250.0},
}

target_margins
Out[14]:
{'gender': {'Female': 100.0, 'Male': 850.0, nan: 50.0},
 'age_group': {'18-24': 250.0, '25-34': 250.0, '35-44': 250.0, '45+': 250.0}}

Now fit rake() from the sample covariates and weights. Because target_margins supplies the target population, we pass None for target_df and target_weights.

In [15]:
marginal_rake_result = rake(
    sample.covars().df,
    sample.weight_series,
    target_df=None,
    target_weights=None,
    target_margins=target_margins,
)

adjusted_rake_2_weights = marginal_rake_result["weight"]
adjusted_rake_2_weights.head()

INFO (2026-08-19 01:11:21,927) [adjustment/apply_transformations (line 435)]: Adding the variables: []

INFO (2026-08-19 01:11:21,928) [adjustment/apply_transformations (line 436)]: Transforming the variables: ['gender', 'age_group']

INFO (2026-08-19 01:11:21,932) [adjustment/apply_transformations (line 472)]: Final variables in output: ['gender', 'age_group']

INFO (2026-08-19 01:11:21,938) [rake/rake (line 623)]: Final covariates and levels that will be used in raking: {'age_group': ['18-24', '25-34', '35-44', '45+'], 'gender': ['Female', 'Male', '__NaN__']}.
Out[15]:
index
0    1.105660
1    0.196483
2    0.656030
3    0.324401
4    0.324401
Name: rake_weight, dtype: float64

The fitted weights are normalized to the common total supplied by target_margins:

In [16]:
adjusted_rake_2_weights.sum()
Out[16]:
np.float64(1000.0)

As the following code shows, the fitted weights recover the requested marginal totals for age and gender.

In [17]:
def weighted_totals_by_column(df, weights, column):
    values = df[column].astype(object).where(df[column].notna(), "__NaN__")
    return weights.groupby(values).sum().sort_index()

print("Gender weighted totals:")
print(weighted_totals_by_column(sample.covars().df, adjusted_rake_2_weights, "gender"))
print("\nAge-group weighted totals:")
print(weighted_totals_by_column(sample.covars().df, adjusted_rake_2_weights, "age_group"))

Gender weighted totals:
gender
Female     100.0
Male       850.0
__NaN__     50.0
Name: rake_weight, dtype: float64

Age-group weighted totals:
age_group
18-24    249.999834
25-34    249.999227
35-44    249.998940
45+      250.001999
Name: rake_weight, dtype: float64
In [ ]: