uv pip install microimpute

# With image export (for plots)
uv pip install microimpute[images]

from microimpute import Imputer
import pandas as pd

# Donor data (complete)
donor = pd.DataFrame({
    'income': [50000, 60000, 70000],
    'age': [30, 40, 50],
    'capital_gains': [5000, 8000, 12000]  # Variable to impute
})

# Recipient data (missing capital_gains)
recipient = pd.DataFrame({
    'income': [55000, 65000],
    'age': [35, 45],
    # capital_gains is missing
})

# Impute using quantile forest
imputer = Imputer(method='quantile_forest')
imputer.fit(
    donor=donor,
    donor_target='capital_gains',
    common_vars=['income', 'age']
)

recipient_imputed = imputer.predict(recipient)
# Now recipient has predicted capital_gains

from microimpute import compare_methods

# Compare different imputation methods
results = compare_methods(
    donor=donor,
    recipient=recipient,
    target_var='capital_gains',
    common_vars=['income', 'age'],
    methods=['linear', 'random_forest', 'quantile_forest']
)

# Shows quantile loss for each method
print(results)

# Quantile loss around 0.1 = good
# Quantile loss around 0.5 = poor
# Compare across methods to choose best

git clone https://github.com/PolicyEngine/microimpute
cd microimpute

tree microimpute/

# Key modules:
ls microimpute/
# - imputer.py - Main Imputer class
# - methods/ - Different imputation methods
# - comparison.py - Method benchmarking
# - utils/ - Utilities

# Quantile forest implementation
cat microimpute/methods/quantile_forest.py

# Random forest
cat microimpute/methods/random_forest.py

# Linear regression
cat microimpute/methods/linear.py

cat pyproject.toml

# microimpute/methods/my_method.py

class MyMethodImputer:
    def fit(self, X_train, y_train):
        """Train on donor data."""
        # Fit your model
        pass

    def predict(self, X_test):
        """Impute on recipient data."""
        # Return predictions
        pass

    def get_quantile_loss(self, X_val, y_val):
        """Compute validation loss."""
        # Evaluate quality
        pass

cd ../policyengine-us-data

# Find usage
grep -r "microimpute" policyengine_us_data/
grep -r "Imputer" policyengine_us_data/

make test

# Or
pytest tests/ -v --cov=microimpute

cat tests/test_imputer.py
cat tests/test_methods.py

from microimpute import Imputer

# Create imputer
imputer = Imputer(method='quantile_forest')

# Fit on donor (complete data)
imputer.fit(
    donor=donor_df,
    donor_target='target_variable',
    common_vars=['age', 'income', 'state']
)

# Predict on recipient (missing target_variable)
recipient_imputed = imputer.predict(recipient_df)

from microimpute import compare_methods

# Test multiple methods
methods = ['linear', 'random_forest', 'quantile_forest', 'xgboost']

results = compare_methods(
    donor=donor,
    recipient=recipient,
    target_var='target',
    common_vars=common_vars,
    methods=methods
)

# Use method with lowest quantile loss
best_method = results.sort_values('quantile_loss').iloc[0]['method']

# Impute several variables
variables_to_impute = [
    'short_term_capital_gains',
    'long_term_capital_gains',
    'qualified_dividends'
]

for var in variables_to_impute:
    imputer = Imputer(method='quantile_forest')
    imputer.fit(donor=irs_puf, donor_target=var, common_vars=common_vars)
    cps[var] = imputer.predict(cps)

from microimpute import tune_hyperparameters

# Automatically find best hyperparameters
best_params, study = tune_hyperparameters(
    donor=donor,
    target_var='target',
    common_vars=common_vars,
    method='quantile_forest',
    n_trials=100
)

# Use tuned parameters
imputer = Imputer(method='quantile_forest', **best_params)

from sklearn.model_selection import cross_val_score

# Split donor for validation
# Impute on validation set
# Measure accuracy

import plotly.express as px

# Compare imputed vs actual (on donor validation set)
fig = px.scatter(
    x=actual_values,
    y=imputed_values,
    labels={'x': 'Actual', 'y': 'Imputed'}
)
fig.add_trace(px.line(x=[min, max], y=[min, max]))  # 45-degree line

1. Load CPS survey data
   ↓
2. microimpute: Fill missing variables from IRS PUF
   ↓
3. microcalibrate: Adjust weights to match benchmarks
   ↓
4. Validation: Check against administrative totals
   ↓
5. Package: Distribute enhanced dataset
   ↓
6. PolicyEngine: Use for population simulations

# Use random forest (faster than quantile forest)
imputer = Imputer(method='random_forest')

# Or subsample donor
donor_sample = donor.sample(n=10000, random_state=42)
imputer.fit(donor=donor_sample, ...)

# Use quantile forest with tuning
best_params, _ = tune_hyperparameters(...)
imputer = Imputer(method='quantile_forest', **best_params)