In [1]:
from pathlib import Path
from survey_kit.utilities.random import RandomData
from survey_kit.utilities.formula_builder import FormulaBuilder
from survey_kit.calibration.moment import Moment
from survey_kit.calibration.calibration import Calibration
from survey_kit.utilities.dataframe import summary
import narwhals as nw
from survey_kit import logger
In [2]:
logger.info("Generating data for weighting")
n_rows = 100_000
df_population = (
RandomData(n_rows=n_rows, seed=12332151)
.index("index")
.integer("v_1", 1, 10)
.np_distribution("v_f_continuous_0", "normal", loc=10, scale=2)
.np_distribution("v_f_continuous_1", "normal", loc=10, scale=2)
.np_distribution("v_f_continuous_2", "normal", loc=10, scale=2)
.float("v_extra", -1, 2)
.np_distribution("weight_0", "normal", loc=10, scale=1)
.np_distribution("weight_1", "normal", loc=10, scale=1)
.integer("year", 2016, 2021)
.integer("month", 1, 12)
.to_df()
.lazy()
)
df_treatment = (
RandomData(n_rows=n_rows, seed=894654)
.index("index")
.integer("v_1", 1, 10)
# Intentionally set the loc/scale as different than above
.np_distribution("v_f_continuous_0", "normal", loc=11, scale=4)
.np_distribution("v_f_continuous_1", "normal", loc=11, scale=4)
.np_distribution("v_f_continuous_2", "normal", loc=11, scale=4)
.float("v_extra", -1, 2)
.np_distribution("weight_0", "normal", loc=10, scale=1)
.np_distribution("weight_1", "normal", loc=10, scale=1)
.integer("year", 2016, 2021)
.integer("month", 1, 12)
.to_df()
.lazy()
)
# print(df.describe())
Generating data for weighting
In [3]:
logger.info("Weighting 'function'")
f = FormulaBuilder(df=df_population, constant=False)
f.continuous(columns=["v_1", "v_f_continuous_*", "v_f_p2_*"])
# f.simple_interaction(columns=["v_1","v_f_continuous_0"])
logger.info("Define the target moments that the weighting will match")
logger.info(" This can be a dataset or a single row of pop controls")
m = Moment(
df=df_population,
formula=f.formula,
weight="weight_0",
index="index",
by=["year"],
equalize_by=True,
rescale=True,
)
logger.info("You can save/reload moments if you want")
# m.save("/my/path/moment")
# m_loaded = Moment.load("/my/path/moment")
Weighting 'function'
Define the target moments that the weighting will match
This can be a dataset or a single row of pop controls
You can save/reload moments if you want
In [4]:
# Calibrate the data in df_treatment to the moment above
c = Calibration(
df=df_treatment, moments=m, weight="weight_1", final_weight="weight_final"
)
c.run(
# Drop a moment if there are too few observations
min_obs=5,
# If it fails to converge, set bounds on the weights
# final weights = (base*ratio) where the bounds are on the ratio
# for "best possible" weights
bounds=(0.001, 1000),
)
# Merge the final weights back on the treatment data
df_treatment = c.get_final_weights(df_treatment)
Aggregating any sub_moments
Calibrating weights using aebw
min obs = 5
Calibration using combined moments
Entropy Balance Rewighting, Sanders (2024)
Input matrix is sparse? False
Problem Size: 100000 rows, 31 moments
# Criterion ||Eq. Const.|| ||FOC Lagr.|| PrimalStepSize DualStepSize Opt. Violation.
SparseLinearSolver(size=31, spd=True) using: pypardiso (with PARDISO analysis reuse)
0 0.000000 42897.345673 0.0000 inf inf 42897.34567303
1 2758.741585 18718.889431 56.2768 134.44454489 0.87644069 18718.97402661
2 9511.441231 12.601388 10.3729 71.07822924 0.02816004 16.32154598
3 9484.568346 0.266563 1.2351 12.11729118 0.02869139 1.26357097
4 9483.867173 0.001644 0.3279 0.88765227 0.00072279 0.32789490
5 9483.851573 0.000007 0.0330 0.04234383 0.00000673 0.03298252
6 9483.851465 0.000000 0.0004 0.00497025 0.00000105 0.00037928
7 9483.851465 0.000000 0.0000 0.00005853 0.00000001 0.00000005
Optimality converged.
Time elapsed: 0.6667935000150464
Optimization completed, success?: True
Converged: True
Maximum Difference: 4.551914400963142e-15
shape: (30, 9) ┌────────────────────────────────┬──────────┬─────────┬───────────┬────────────┬─────────┬────────────────┬─────────────┬─────────────┐ │ Variable ┆ Initial ┆ Targets ┆ Estimates ┆ Calibrated ┆ NonZero ┆ NonZero_Target ┆ diff ┆ percent │ ╞════════════════════════════════╪══════════╪═════════╪═══════════╪════════════╪═════════╪════════════════╪═════════════╪═════════════╡ │ m0_year==2021:v_1 ┆ 1.006379 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16681 ┆ 16506 ┆ -4.5519e-15 ┆ -4.5519e-13 │ │ m0_year==2017:v_f_continuous_0 ┆ 1.092225 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16569 ┆ 16887 ┆ -2.6645e-15 ┆ -2.6645e-13 │ │ m0_year==2017:_in ┆ 0.993348 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16569 ┆ 16887 ┆ -2.6645e-15 ┆ -2.6645e-13 │ │ m0_year==2016:v_f_continuous_1 ┆ 1.10321 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16676 ┆ 16676 ┆ -2.3315e-15 ┆ -2.3315e-13 │ │ m0_year==2017:v_f_continuous_2 ┆ 1.096367 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16569 ┆ 16887 ┆ 2.2204e-15 ┆ 2.2204e-13 │ │ m0_year==2019:v_f_continuous_2 ┆ 1.09556 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16624 ┆ 16763 ┆ -1.9984e-15 ┆ -1.9984e-13 │ │ m0_year==2020:v_f_continuous_1 ┆ 1.097323 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16646 ┆ 16618 ┆ -1.8874e-15 ┆ -1.8874e-13 │ │ m0_year==2020:v_f_continuous_0 ┆ 1.092051 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16646 ┆ 16618 ┆ -1.8874e-15 ┆ -1.8874e-13 │ │ m0_year==2021:v_f_continuous_1 ┆ 1.09959 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16681 ┆ 16506 ┆ -1.7764e-15 ┆ -1.7764e-13 │ │ m0_year==2020:v_1 ┆ 1.001671 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16646 ┆ 16618 ┆ -1.7764e-15 ┆ -1.7764e-13 │ │ m0_year==2019:_in ┆ 0.997436 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16624 ┆ 16763 ┆ 1.7764e-15 ┆ 1.7764e-13 │ │ m0_year==2016:_in ┆ 1.000639 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16676 ┆ 16676 ┆ 1.5543e-15 ┆ 1.5543e-13 │ │ m0_year==2018:_in ┆ 1.008746 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16804 ┆ 16550 ┆ -1.4433e-15 ┆ -1.4433e-13 │ │ m0_year==2016:v_f_continuous_2 ┆ 1.09744 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16676 ┆ 16676 ┆ -1.4433e-15 ┆ -1.4433e-13 │ │ m0_year==2017:v_1 ┆ 1.002875 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16569 ┆ 16887 ┆ -1.4433e-15 ┆ -1.4433e-13 │ │ m0_year==2020:_in ┆ 0.998146 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16646 ┆ 16618 ┆ 1.1102e-15 ┆ 1.1102e-13 │ │ m0_year==2019:v_f_continuous_0 ┆ 1.102635 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16624 ┆ 16763 ┆ -1.1102e-15 ┆ -1.1102e-13 │ │ m0_year==2019:v_f_continuous_1 ┆ 1.101078 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16624 ┆ 16763 ┆ 1.1102e-15 ┆ 1.1102e-13 │ │ m0_year==2018:v_f_continuous_0 ┆ 1.112906 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16804 ┆ 16550 ┆ -9.9920e-16 ┆ -9.9920e-14 │ │ m0_year==2018:v_f_continuous_1 ┆ 1.11314 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16804 ┆ 16550 ┆ 8.8818e-16 ┆ 8.8818e-14 │ │ m0_year==2016:v_f_continuous_0 ┆ 1.101391 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16676 ┆ 16676 ┆ -8.8818e-16 ┆ -8.8818e-14 │ │ m0_year==2020:v_f_continuous_2 ┆ 1.100242 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16646 ┆ 16618 ┆ -8.8818e-16 ┆ -8.8818e-14 │ │ m0_year==2019:v_1 ┆ 1.008309 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16624 ┆ 16763 ┆ -8.8818e-16 ┆ -8.8818e-14 │ │ m0_year==2018:v_1 ┆ 1.01166 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16804 ┆ 16550 ┆ -7.7716e-16 ┆ -7.7716e-14 │ │ m0_year==2018:v_f_continuous_2 ┆ 1.113311 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16804 ┆ 16550 ┆ -5.5511e-16 ┆ -5.5511e-14 │ │ m0_year==2021:v_f_continuous_2 ┆ 1.09976 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16681 ┆ 16506 ┆ -5.5511e-16 ┆ -5.5511e-14 │ │ m0_year==2021:_in ┆ 1.001684 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16681 ┆ 16506 ┆ -5.5511e-16 ┆ -5.5511e-14 │ │ m0_year==2021:v_f_continuous_0 ┆ 1.103274 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16681 ┆ 16506 ┆ -4.4409e-16 ┆ -4.4409e-14 │ │ m0_year==2016:v_1 ┆ 0.999911 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16676 ┆ 16676 ┆ -3.3307e-16 ┆ -3.3307e-14 │ │ m0_year==2017:v_f_continuous_1 ┆ 1.092783 ┆ 1 ┆ 1.0 ┆ 1 ┆ 16569 ┆ 16887 ┆ 2.2204e-16 ┆ 2.2204e-14 │ └────────────────────────────────┴──────────┴─────────┴───────────┴────────────┴─────────┴────────────────┴─────────────┴─────────────┘
In [5]:
logger.info("'Population' estimates")
_ = summary(df_population, weight="weight_0")
'Population' estimates
┌──────────────────┬─────────┬─────────────┬───────────────┬───────────────┬───────────┬───────────┐ │ Variable ┆ n ┆ n (missing) ┆ mean ┆ std ┆ min ┆ max │ ╞══════════════════╪═════════╪═════════════╪═══════════════╪═══════════════╪═══════════╪═══════════╡ │ index ┆ 100,000 ┆ 0 ┆ 49,987.162827 ┆ 28,870.556017 ┆ 0.0 ┆ 99,999.0 │ │ v_1 ┆ 100,000 ┆ 0 ┆ 5.484124 ┆ 2.875588 ┆ 1.0 ┆ 10.0 │ │ v_f_continuous_0 ┆ 100,000 ┆ 0 ┆ 9.989362 ┆ 1.996633 ┆ 1.491748 ┆ 18.835062 │ │ v_f_continuous_1 ┆ 100,000 ┆ 0 ┆ 10.002072 ┆ 2.006889 ┆ 1.37638 ┆ 18.82769 │ │ v_f_continuous_2 ┆ 100,000 ┆ 0 ┆ 9.998039 ┆ 2.004505 ┆ 1.166252 ┆ 19.254231 │ │ v_extra ┆ 100,000 ┆ 0 ┆ 0.504297 ┆ 0.867154 ┆ -0.999978 ┆ 1.999996 │ │ weight_1 ┆ 100,000 ┆ 0 ┆ 10.004412 ┆ 1.002786 ┆ 5.35507 ┆ 14.023032 │ │ year ┆ 100,000 ┆ 0 ┆ 2,018.493399 ┆ 1.706467 ┆ 2,016.0 ┆ 2,021.0 │ │ month ┆ 100,000 ┆ 0 ┆ 6.504686 ┆ 3.44999 ┆ 1.0 ┆ 12.0 │ └──────────────────┴─────────┴─────────────┴───────────────┴───────────────┴───────────┴───────────┘
In [6]:
logger.info("\n\n'Treatment', original weights")
_ = summary(df_treatment, weight="weight_1")
'Treatment', original weights
┌──────────────────┬─────────┬─────────────┬───────────────┬───────────────┬───────────┬───────────┐ │ Variable ┆ n ┆ n (missing) ┆ mean ┆ std ┆ min ┆ max │ ╞══════════════════╪═════════╪═════════════╪═══════════════╪═══════════════╪═══════════╪═══════════╡ │ index ┆ 100,000 ┆ 0 ┆ 49,997.423141 ┆ 28,873.734839 ┆ 0.0 ┆ 99,999.0 │ │ v_1 ┆ 100,000 ┆ 0 ┆ 5.51228 ┆ 2.868293 ┆ 1.0 ┆ 10.0 │ │ v_f_continuous_0 ┆ 100,000 ┆ 0 ┆ 10.99576 ┆ 4.013016 ┆ -5.459849 ┆ 30.287625 │ │ v_f_continuous_1 ┆ 100,000 ┆ 0 ┆ 11.014155 ┆ 3.999313 ┆ -7.360353 ┆ 27.415263 │ │ v_f_continuous_2 ┆ 100,000 ┆ 0 ┆ 11.00231 ┆ 4.005346 ┆ -8.478953 ┆ 27.102181 │ │ v_extra ┆ 100,000 ┆ 0 ┆ 0.500786 ┆ 0.867307 ┆ -0.999996 ┆ 1.999988 │ │ weight_0 ┆ 100,000 ┆ 0 ┆ 9.997927 ┆ 1.001265 ┆ 5.26969 ┆ 14.090873 │ │ year ┆ 100,000 ┆ 0 ┆ 2,018.500692 ┆ 1.707684 ┆ 2,016.0 ┆ 2,021.0 │ │ month ┆ 100,000 ┆ 0 ┆ 6.489706 ┆ 3.447587 ┆ 1.0 ┆ 12.0 │ │ weight_final ┆ 100,000 ┆ 0 ┆ 1.010009 ┆ 0.475246 ┆ 0.152202 ┆ 7.668396 │ └──────────────────┴─────────┴─────────────┴───────────────┴───────────────┴───────────┴───────────┘
In [7]:
logger.info("\n\n'Treatment', calibrated")
_ = summary(df_treatment, weight="weight_final")
'Treatment', calibrated
┌──────────────────┬─────────┬─────────────┬──────────────┬───────────────┬───────────┬───────────┐ │ Variable ┆ n ┆ n (missing) ┆ mean ┆ std ┆ min ┆ max │ ╞══════════════════╪═════════╪═════════════╪══════════════╪═══════════════╪═══════════╪═══════════╡ │ index ┆ 100,000 ┆ 0 ┆ 49,956.05108 ┆ 28,879.526059 ┆ 0.0 ┆ 99,999.0 │ │ v_1 ┆ 100,000 ┆ 0 ┆ 5.484124 ┆ 2.868615 ┆ 1.0 ┆ 10.0 │ │ v_f_continuous_0 ┆ 100,000 ┆ 0 ┆ 9.989362 ┆ 4.006812 ┆ -5.459849 ┆ 30.287625 │ │ v_f_continuous_1 ┆ 100,000 ┆ 0 ┆ 10.002072 ┆ 4.006384 ┆ -7.360353 ┆ 27.415263 │ │ v_f_continuous_2 ┆ 100,000 ┆ 0 ┆ 9.998039 ┆ 4.007537 ┆ -8.478953 ┆ 27.102181 │ │ v_extra ┆ 100,000 ┆ 0 ┆ 0.502204 ┆ 0.867469 ┆ -0.999996 ┆ 1.999988 │ │ weight_0 ┆ 100,000 ┆ 0 ┆ 9.999062 ┆ 1.000964 ┆ 5.26969 ┆ 14.090873 │ │ weight_1 ┆ 100,000 ┆ 0 ┆ 10.102167 ┆ 0.998455 ┆ 5.539069 ┆ 14.227871 │ │ year ┆ 100,000 ┆ 0 ┆ 2,018.5 ┆ 1.707834 ┆ 2,016.0 ┆ 2,021.0 │ │ month ┆ 100,000 ┆ 0 ┆ 6.492081 ┆ 3.448134 ┆ 1.0 ┆ 12.0 │ └──────────────────┴─────────┴─────────────┴──────────────┴───────────────┴───────────┴───────────┘