weighting.py

from argparse import ArgumentParser

import numpy as np
import pandas as pd
from ruamel.yaml import YAML
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
import torch

from skorch import NeuralNetClassifier

DATASET_PATHSPEC = "./config/data_pathspec.yml"
yaml = YAML(typ='safe')
yaml.default_flow_style = False

class PermutationWeighter(object):
    def __init__(self, classifier, seed=42, **kwargs):
        self.clf = NeuralNetClassifier(
            classifier,
            **kwargs
        )

    def fit_weights(self, df, treatment_col, feature_cols, perm_seed=42, evaluate=True):
        obs_data = df.copy()
        obs_data['perm_label'] = 0
        perm_data = df.copy()
        perm_data['perm_label'] = 1

        np.random.seed(perm_seed)
        perm_data[treatment_col] = np.random.choice(perm_data[treatment_col].values, size=len(perm_data), replace=False)
        weighting_data = pd.concat([obs_data, perm_data], axis=0)
        X_ = torch.from_numpy(weighting_data[feature_cols + [treatment_col]].values.astype(np.float32)).to(self.clf.device)
        self.clf.fit(
            {"X_": X_},
            torch.from_numpy(weighting_data['perm_label'].values).long().to(self.clf.device))
        if evaluate:
            x_dict = {"X_": X_}
            probs = self.clf.predict_proba(x_dict)[:, 1]
            auc = roc_auc_score(weighting_data['perm_label'], probs)
            print("Permuted-vs.-Real AUC:", auc)

    def estimate_weights(self, data, treatment_col=None, feature_cols=None):
        if treatment_col is not None and feature_cols is not None:
            data = torch.from_numpy(data.loc[:, feature_cols + [treatment_col]].values.astype(np.float32)).to(self.clf.device)
        weights = self.clf.predict_proba(data)
        return weights[:, 1] / weights[:, 0]



if __name__ == '__main__':
    psr = ArgumentParser()
    psr.add_argument("--seed", type=int, default=42)
    psr.add_argument("--perm-seed", type=int, default=0)
    psr.add_argument("--dataset", type=str, required=True)
    psr.add_argument("--treatment-col", type=str, default="t")
    psr.add_argument("--feature-cols", type=str, nargs='+', default=[])
    psr.add_argument("--feature-prefixes", type=str, nargs='+', default="x")
    psr.add_argument("--classifier-kwargs", type=str, nargs='+', default=[])

    args = psr.parse_args()

    with open(DATASET_PATHSPEC, "r") as f:
        dataset_cfg = yaml.load(f)
    path = dataset_cfg[args.dataset]["data"]
    print("Loading data from", path) 
    df = pd.read_csv(path, low_memory=False, index_col=0)

    dev_df, test_df = train_test_split(df, test_size=0.3, random_state=args.seed)

    feature_cols = args.feature_cols
    for pref in args.feature_prefixes:
        feature_cols += [c for c in df.columns if c.startswith(pref)]
    print("Fitting permutation weights...")
    classifier_kwargs = dict([x.split("=") for x in args.classifier_kwargs])
    pw = PermutationWeighter(RandomForestClassifier, seed=args.seed, **classifier_kwargs)
    print("Using treatment col:",  args.treatment_col)
    print("Features:", feature_cols)
    dev_pw = pw.fit_weights(dev_df, args.treatment_col, feature_cols, args.perm_seed)
    df["perm_ipsw"] = pw.estimate_weights(df, args.treatment_col, feature_cols)
    df.to_csv(path)
    print("Saved permutation weights to", path)