simulations.py

# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.12.0
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %%
BASE_SEED = 0
DATA_SETUP = dict(
    num_features=100,
    noise=0.1,
    delta=0.5,
    divergent=True,
)

import os
DATA_FILE_FORMAT = os.environ.get("DATA_FILE_FORMAT") or "results/simulations_{}_paper_snapshot.csv"
NUM_SHADOWS = int(os.environ.get("NUM_SHADOWS") or 30)
NUM_EVAL_MODELS = int(os.environ.get("NUM_EVAL_MODELS") or 200)
RESTORE_SAVED_DATA = False

# %%
# %load_ext autoreload
# %autoreload 2

import math
import hashlib
import itertools
import collections

import numpy as np
import pandas as pd
import seaborn as sns
import diffprivlib.models as dp

from IPython.display import display
from tqdm import autonotebook as tqdm
from matplotlib import pyplot as plt
from sklearn.datasets import make_spd_matrix
from sklearn.model_selection import train_test_split
from scipy import stats
from joblib import Parallel, delayed

from mia import run_shadow_model_attack, run_threshold_estimator, apply_estimator_func
from utils import infer_from, normalize_vuln_mean, normalize_vuln_std
from model_zoo import model_zoo, lr_setup, renaming_dict, DpClassifierFactory
from plotting import add_significance_markers

import plot_params

# %%
import logging
logger = logging.Logger(name="default")


# %% [markdown]
# ## Simulating disparate vulnerability

# %% [markdown]
# ### Synthetic data

# %%
def create_multinomial_setup(
        num_features=2,
        noise=0.1,
        delta=0.5,
        divergent=True,
        seed=1
):
    gen = np.random.RandomState(seed)
    weights = np.ones(num_features) / num_features
    if divergent:
        means_a = np.zeros([num_features])
        means_b = np.ones([num_features])
    else:
        means_a = gen.uniform(-1, 1, num_features)
        means_b = gen.uniform(-1, 1, num_features)

    cov = make_spd_matrix(num_features, random_state=seed) * noise

    return means_a, means_b, cov, noise, delta

def gen_sim_data(setup, gen, size_a, size_b):
    means0, means1, cov, noise, delta = setup
    y = [0] * (size_a // 2) + [1] * math.ceil(size_a / 2) + \
        [0] * (size_b // 2) + [1] * math.ceil(size_b / 2)
    a0 = gen.multivariate_normal(means0 - delta, cov, size=size_a // 2)
    a1 = gen.multivariate_normal(means1, cov, size=math.ceil(size_a / 2))
    b0 = gen.multivariate_normal(means0, cov, size=size_b // 2)
    b1 = gen.multivariate_normal(means1 - delta, cov, size=math.ceil(size_b / 2))
    X = np.vstack([a0, a1, b0, b1])

    hi = 2**32
    index = gen.randint(hi, size=size_a + size_b)
    return pd.DataFrame(X, index=index), pd.Series(np.array(y), index=index)

# Demo:
setup = create_multinomial_setup(seed=BASE_SEED, **DATA_SETUP)
gen = np.random.RandomState(BASE_SEED)
size_a = 10000
size_b = 10000
X, y = gen_sim_data(setup, gen, size_a, size_b)
z = np.array(["C"] * size_a + ["T"] * size_b)
X.shape, y.shape, z.shape

sample_df = pd.DataFrame(X)
sample_df = sample_df.assign(y=y).assign(z=z)
sample_df = sample_df.rename(columns={
    "z": "$z$",
    "y": "$y$",
})

g = sns.displot(data=sample_df, x=2, y=3, hue="$y$", col="$z$", kind="kde")
g.set_axis_labels("", "")
g.fig.set_figwidth(12)
g.fig.set_figheight(8)
g.fig.set_tight_layout(tight=True)
# g.fig.savefig("images/synthetic_data_demo.pdf")

# %% [markdown]
# ### Populate Model Zoo

# %%
# Compute max data norm for DP classifiers.
setup = create_multinomial_setup(seed=BASE_SEED, **DATA_SETUP)
gen = np.random.RandomState(BASE_SEED)
size_a = 100000
size_b = 100000
sample, _ = gen_sim_data(setup, gen, size_a, size_b)

sample_norm = np.linalg.norm(sample, axis=1, ord=2).max()
sample_norm

# %%
dp_model_zoo = {
    f"eps{eps}": DpClassifierFactory(eps, data_norm=sample_norm)
    for eps in [0.1, 1, 2, 10]
}
model_zoo = dict(model_zoo, **dp_model_zoo)


# %%
class ThreshClassifier:
    def __init__(self, threshold=0.5):
        self.threshold = threshold

    def fit(self, *args, **kwargs):
        pass

    def predict_proba(self, xs, *args, **kwargs):
        p = xs[:, 0] > self.threshold
        p = np.expand_dims(p, 1)
        return np.hstack([1-p, p])

ThreshClassifier().predict_proba(np.array([[0, 1], [1, 0]]))
model_zoo["threshold"] = lambda: ThreshClassifier()

# %%
renaming_dict = dict(renaming_dict, **{
    "threshold": "Control",
    "eps0.1": r"$\varepsilon=0.1$",
    "eps0.5": r"$\varepsilon=0.5$",
    "eps1": r"$\varepsilon=1$",
    "eps2": r"$\varepsilon=2$",
    "eps2.5": r"$\varepsilon=2.5$",
    "eps5": r"$\varepsilon=5$",
    "eps7.5": r"$\varepsilon=7.5$",
    "eps10": r"$\varepsilon=10$",
    "eps15": r"$\varepsilon=15$",
})


# %%
def eval_sim_vuln(
    data,
    method="shadow_attack_preds",
    num_targets=None,
    **kwargs,
):
    if "shadow" in method:
        shadow_indices = data.y_train.index[:-num_targets]
        target_indices = data.y_train.index[-num_targets:]

        vuln = run_shadow_model_attack(
            data.y_train.loc[shadow_indices].values,
            data.preds_train.loc[shadow_indices].values,
            data.y_test.loc[shadow_indices].values,
            data.preds_test.loc[shadow_indices].values,
            data.y_train.loc[target_indices].values,
            data.preds_train.loc[target_indices].values,
            data.y_test.loc[target_indices].values,
            data.preds_test.loc[target_indices].values,
            method=method,
            **kwargs,
        )

    elif "threshold" in method:
        vuln = apply_estimator_func(
            data.y_train.values,
            data.preds_train.values,
            data.y_test.values,
            data.preds_test.values,
            run_threshold_estimator,
            method=method,
            **kwargs,
        )
    else:
        raise NotImplementedError

    return vuln


# %%
class SimDatasetGen:
    def __init__(self, setup):
        self.setup = setup

    def __call__(self, rng, treatment_group_size, control_group_size):
        X_a_train, y_a_train = gen_sim_data(self.setup, rng, control_group_size, 0)
        X_b_train, y_b_train = gen_sim_data(self.setup, rng, 0, treatment_group_size)
        X_a_test, y_a_test = gen_sim_data(self.setup, rng, control_group_size, 0)
        X_b_test, y_b_test = gen_sim_data(self.setup, rng, 0, treatment_group_size)
        return (X_a_train, y_a_train, X_b_train, y_b_train,
                X_a_test, y_a_test, X_b_test, y_b_test)


# %%
def start_logger_if_necessary():
    logger = logging.getLogger("default")
    if len(logger.handlers) == 0:
        logger.setLevel(logging.INFO)
        sh = logging.StreamHandler()
        sh.setFormatter(logging.Formatter("%(asctime)s %(levelname)-8s %(message)s"))
        fh = logging.FileHandler('simulations.log', mode='w')
        fh.setFormatter(logging.Formatter("%(asctime)s %(levelname)-8s %(message)s"))
        logger.addHandler(sh)
        logger.addHandler(fh)
    return logger


# %%
def get_attack_data(
    model_fn,
    treatment_group_sizes,
    dataset_gen,
    seed=1,
    num_models=100,
    max_size=None,
    parallel=True,
    sensitive_fit=False,
    sensitive_pred=False,
):
    """Generate datasets, train models, and collect their outputs."""

    columns=[
        "group",
        "treatment_size",
        "group_size",
        "rep",
        "baseline_acc",
        "baseline_acc_group",
        "train_acc",
        "test_acc",
        "test_acc_group",
        "preds_train",
        "preds_test",
        "y_train",
        "y_test",
    ]
    if max_size is None:
        max_size = max(treatment_group_sizes)
    it = []
    for size in treatment_group_sizes:
        it.extend(itertools.product(
            [size],
            range(num_models),
        ))

    # Add a PRNG for each iteration.
    rngs = [np.random.RandomState(hash(f"{seed}-{i}") % 2**32) for i, _ in enumerate(it)]
    it = tqdm.tqdm(list(zip(it, rngs)))
    it.set_description("Generating model outputs")

    results = pd.DataFrame()

    def run_iter(treatment_group_size, rep, rng):
        logger = start_logger_if_necessary()

        fit_args = {}

        datasets = dataset_gen(rng, treatment_group_size, max_size)
        (X_a_train, y_a_train, X_b_train, y_b_train,
         X_a_test, y_a_test, X_b_test, y_b_test) = datasets

        test_acc_a_baseline = max(y_a_test.mean(), 1 - y_a_test.mean())
        test_acc_b_baseline = max(y_b_test.mean(), 1 - y_b_test.mean())
        test_acc_baseline = (y_a_test.sum() + y_b_test.sum()) / (len(y_a_test) + len(y_b_test))
        test_acc_baseline = max(test_acc_baseline, 1 - test_acc_baseline)

        X_train = np.vstack([X_a_train, X_b_train])
        y_train = np.concatenate([y_a_train, y_b_train])

        # Special case for the fair classifier
        if sensitive_fit or sensitive_pred:
            train_group_ids = np.array([0] * len(X_a_train) + [1] * len(X_b_train))
            fit_args = dict(sensitive_features=train_group_ids)

        clf = model_fn()
        clf.fit(X_train, y_train, **fit_args)

        preds_a_train = infer_from(clf, X_a_train, sensitive_features=[0] * len(X_a_train)
                                   if sensitive_pred else None)
        preds_a_test = infer_from(clf, X_a_test, sensitive_features=[0] * len(X_a_test)
                                  if sensitive_pred else None)
        preds_b_train = infer_from(clf, X_b_train, sensitive_features=[1] * len(X_b_train)
                                   if sensitive_pred else None)
        preds_b_test = infer_from(clf, X_b_test, sensitive_features=[1] * len(X_b_test)
                                  if sensitive_pred else None)

        # Add indices
        preds_a_train = pd.Series(preds_a_train, index=X_a_train.index)
        preds_a_test = pd.Series(preds_a_test, index=X_a_test.index)
        preds_b_train = pd.Series(preds_b_train, index=X_b_train.index)
        preds_b_test = pd.Series(preds_b_test, index=X_b_test.index)

        # Compute accuracies
        train_acc = np.concatenate([
            (preds_a_train >= 0.5) == y_a_train,
            (preds_b_train >= 0.5) == y_b_train]).mean()

        test_acc_a_vals = (preds_a_test >= 0.5) == y_a_test
        test_acc_b_vals = (preds_b_test >= 0.5) == y_b_test
        test_acc = np.concatenate([test_acc_a_vals, test_acc_b_vals]).mean()

        test_acc_a = test_acc_a_vals.mean()
        test_acc_b = test_acc_b_vals.mean()

        result = pd.DataFrame()
        result = result.append(dict(zip(columns, [
            "a",
            treatment_group_size,
            max_size,
            rep,
            test_acc_baseline,
            test_acc_a_baseline,
            train_acc,
            test_acc,
            test_acc_a,
            preds_a_train,
            preds_a_test,
            y_a_train,
            y_a_test,
        ])), ignore_index=True)
        return result.append(dict(zip(columns, [
            "b",
            treatment_group_size,
            treatment_group_size,
            rep,
            test_acc_baseline,
            test_acc_b_baseline,
            train_acc,
            test_acc,
            test_acc_b,
            preds_b_train,
            preds_b_test,
            y_b_train,
            y_b_test,
        ])), ignore_index=True)

    if parallel:
        data = Parallel(n_jobs=8)(delayed(run_iter)(size, rep, rng)
                                             for (size, rep), rng in it)
    else:
        data = []
        for (size, rep), rng in it:
            data.append(run_iter(size, rep, rng))

    return pd.concat(data, axis=0, ignore_index=True)


# %%
def run_sim_all_sizes(
    model_fn,
    treatment_group_sizes,
    dataset_gen=None,
    data=None,
    method="average_loss_threshold",
    max_size=None,
    num_models=100,
    num_targets=None,
    parallel=True,
    microdata=False,
    sensitive_fit=False,
    sensitive_pred=False,
    seed=1,
):
    if data is None:
        assert dataset_gen is not None
        data = get_attack_data(
            model_fn,
            treatment_group_sizes,
            dataset_gen,
            max_size=max_size,
            num_models=num_models,
            parallel=parallel,
            sensitive_fit=sensitive_fit,
            sensitive_pred=sensitive_pred,
            seed=seed,
        )

    results = pd.DataFrame()
    it = tqdm.tqdm(treatment_group_sizes)
    it.set_description("Evaluating vuln. in group sizes")

    for treatment_group_size in it:
        size_data = data.query(f"treatment_size == {treatment_group_size}")
        group_a_data = size_data.query("group == 'a'").reset_index()
        group_b_data = size_data.query("group == 'b'").reset_index()

        vuln_a = eval_sim_vuln(group_a_data, method,
                           num_targets=num_targets,
                           parallel=parallel,
                           microdata=microdata)
        vuln_b = eval_sim_vuln(group_b_data, method,
                           num_targets=num_targets,
                           parallel=parallel,
                           microdata=microdata)

        if microdata:
            vuln_a_vals = vuln_a.groupby("model_id").vuln.mean()
            vuln_b_vals = vuln_b.groupby("model_id").vuln.mean()

        else:
            vuln_a_vals = vuln_a
            vuln_b_vals = vuln_b

        for vuln_a_val, vuln_b_val in zip(vuln_a_vals, vuln_b_vals):
            results = results.append(
                dict(
                    treatment_size=treatment_group_size,
                    control_size=group_a_data.group_size.mean(),
                    vuln_a=vuln_a_val,
                    vuln_b=vuln_b_val,
                    baseline_acc=size_data.baseline_acc.mean() * 100,
                    baseline_acc_a=group_a_data.baseline_acc.mean() * 100,
                    baseline_acc_b=group_b_data.baseline_acc.mean() * 100,
                    train_acc=size_data.train_acc.mean() * 100,
                    test_acc=size_data.test_acc.mean() * 100,
                    test_acc_a=group_a_data.test_acc.mean() * 100,
                    test_acc_b=group_b_data.test_acc.mean() * 100,
                ), ignore_index=True)

    return results

# %% [markdown]
# ## Sample size experiments

# %%
model_names = ["threshold", "lr", "nn_8", "nn_32"] + list(dp_model_zoo.keys())
model_names

# %%
treatment_group_sizes = [100, 500, 1000, 2500]

# %%
sample_size_results = pd.DataFrame()

# %%
if not RESTORE_SAVED_DATA:
    num_models = NUM_SHADOWS + NUM_EVAL_MODELS
    num_targets = NUM_EVAL_MODELS
    method = "shadow_attack_loss"
    global_setup = create_multinomial_setup(**DATA_SETUP)

    is_fair_model = lambda model_name: any(keyword in model for keyword in ["eo", "erp", "dm"])

    for model in model_names:
        print(model)
        model_fn = model_zoo[model]
        sim_res = run_sim_all_sizes(
            model_fn,
            treatment_group_sizes,
            dataset_gen=SimDatasetGen(global_setup),
            seed=BASE_SEED,
            method=method,
            num_models=num_models,
            num_targets=num_targets,
            parallel=True,
        )
        display(sim_res.groupby("treatment_size").agg({
            "vuln_a": ["mean", "std"],
            "vuln_b": ["mean", "std"],
            "baseline_acc": "mean",
            "test_acc": "mean",
        }))

        sample_size_results = sample_size_results.append(
            sim_res.assign(model=model, method=method), ignore_index=True)

# %%
if RESTORE_SAVED_DATA:
    sample_size_results = pd.read_csv(DATA_FILE_FORMAT.format("sample_size"))
else:
    sample_size_results.to_csv(DATA_FILE_FORMAT.format("sample_size"))

# %%
# Compute p-values.
methods = sample_size_results.method.unique()
for method, model, size in itertools.product(methods, model_names, treatment_group_sizes):
    df = sample_size_results.query(f"model == '{model}' and treatment_size == {size} and method == '{method}'")
    t, p = stats.ttest_rel(df.vuln_a, df.vuln_b)
    sample_size_results.loc[df.index, "p"] = p
    sample_size_results.loc[df.index, "t"] = t

sample_size_results.groupby(["method", "model", "treatment_size"]).p.mean().reset_index()

# %%
unrolled_list = []
for method, model, size in itertools.product(methods, model_names, treatment_group_sizes):
    df = sample_size_results.query(f"model == '{model}' and treatment_size == {size} and method == '{method}'").copy()
    dataset_size = size + treatment_group_sizes[-1]
    unrolled = pd.concat([
        pd.DataFrame(dict(vuln=df.vuln_a, target_test_acc=df.test_acc_a)).assign(subgroup="a"),
        pd.DataFrame(dict(vuln=df.vuln_b, target_test_acc=df.test_acc_b)).assign(subgroup="b"),
        pd.DataFrame(dict(vuln=df.vuln_a * size/dataset_size + df.vuln_b * treatment_group_sizes[-1]/dataset_size,
                          target_test_acc=df.test_acc_b)).assign(subgroup="All"),
    ], ignore_index=False).assign(method=method, model=model, treatment_group_size=size)
    unrolled_list.append(unrolled)

summary_data = pd.concat(unrolled_list, ignore_index=True)

# %%
summary_data.replace(renaming_dict).query("subgroup == 'All'"
).groupby(["model"]).agg({
    "target_test_acc": ["mean", "std"],
    "vuln": [normalize_vuln_mean, normalize_vuln_std],
}).reindex([renaming_dict[model_name] for model_name in model_names]).rename(columns={
    "model": "Model",
    "target_test_acc": "Test acc.",
    "vuln": "Overall vuln.",
})

# %%
hue_order = [renaming_dict[k] for k in dp_model_zoo.keys()] + [renaming_dict["lr"]]
plot_method = "shadow_attack_loss"

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_models = list(dp_model_zoo.keys()) + ["lr"]
plot_df = sample_size_results.query(f"model in {plot_models} and method == '{plot_method}'")
plot_df["model"] = plot_df["model"].replace(renaming_dict)

plot_df = plot_df.rename(columns={
    "model": "Model",
    "test_acc": "Test acc.",
    "treatment_size": "Subgroup size",
})

sns.barplot(data=plot_df,
            x="Subgroup size",
            y="Test acc.",
            hue="Model",
            palette="rocket",
            hue_order=hue_order,
            ax=ax)

# plt.legend().remove()
fig.set_tight_layout(tight=True)

# plt.savefig("images/plot_accuracy_dp.pdf")

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_models = list(dp_model_zoo.keys()) + ["lr"]
plot_df = sample_size_results.query(f"model in {plot_models} and method == '{plot_method}'")
plot_df["vuln_diff"] = 2 * (plot_df.vuln_b - plot_df.vuln_a) * 100
plot_df["model"] = plot_df["model"].replace(renaming_dict)

plot_df = plot_df.rename(columns={
    "model": "Model",
    "treatment_size": "Subgroup size",
    "vuln_a": "Majority vuln.",
    "vuln_b": "Subgroup vuln.",
    "vuln_diff": "Disparity",
})

g = sns.barplot(
    data=plot_df,
    x="Subgroup size", y="Disparity",
    hue="Model",
    ax=ax,
    palette="rocket",
    hue_order=hue_order,
    alpha=1,
)

add_significance_markers(ax, sample_size_results.query(f"method == '{plot_method}'"),
                         hue_values=plot_models, x_values=treatment_group_sizes)

# ax.set_xscale("log")
ax.set_xticklabels(treatment_group_sizes)
plt.legend().remove()
fig.set_tight_layout(tight=True)

# plt.savefig("images/plot_sample_size_dp.pdf")

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_models = list(dp_model_zoo.keys()) + ["lr"]
plot_df = sample_size_results.query(f"model in {plot_models} and method == '{plot_method}'")
plot_df["model"] = plot_df["model"].replace(renaming_dict)
plot_df["acc_diff"] = plot_df["test_acc_b"] - plot_df["test_acc_a"]

plot_df = plot_df.rename(columns={
    "model": "Model",
    "treatment_size": "Subgroup size",
    "test_acc": "Test acc." ,
    "test_acc_a": "Majority group acc." ,
    "test_acc_b": "Minority group acc." ,
    "acc_diff": "Acc. discrepancy",
})

sns.barplot(data=plot_df,
            x="Subgroup size", y="Minority group acc.",
            hue="Model",
            palette="rocket",
            hue_order=hue_order,
            ax=ax)

# ax.set_xscale("log")
# ax.get_legend().remove()
ax.set_xticklabels(treatment_group_sizes)
fig.set_tight_layout(tight=True)

# plt.savefig("images/plot_accuracy_dp.pdf")

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_models = ["lr", "nn_8", "nn_32"]
plot_df = sample_size_results.query(f"model in {plot_models}")
plot_df["gen_gap"] = plot_df.train_acc - plot_df.test_acc
plot_df.groupby(["model", "treatment_size"]).gen_gap.mean().reset_index()
plot_df = plot_df.replace(renaming_dict)
plot_df = plot_df.rename(columns={
    "model": "Model",
    "treatment_size": "Subgroup size",
    "gen_gap": "Overfitting",
})

sns.barplot(data=plot_df, hue="Model", x="Subgroup size", y="Overfitting",
            ax=ax)
ax.set_xticklabels(treatment_group_sizes)

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_models = ["lr", "nn_8", "nn_32"]
plot_df = sample_size_results.query(f"model in {plot_models} and method == '{plot_method}'")
plot_df["vuln_diff"] = (plot_df.vuln_b - plot_df.vuln_a) * 100
plot_df["model"] = plot_df["model"].replace(renaming_dict)

plot_df = plot_df.rename(columns={
    "model": "Model",
    "treatment_size": "Subgroup size",
    "vuln_b": "Subgroup vuln.",
    "vuln_diff": "Disparity",
})

sns.barplot(data=plot_df,
            x="Subgroup size", y="Disparity",
            hue="Model",
            ax=ax)
# ax.set_xscale("log")

add_significance_markers(ax, sample_size_results.query(f"method == '{plot_method}'"),
                         hue_values=plot_models, x_values=treatment_group_sizes)

ax.set_xticklabels(treatment_group_sizes)
ax.legend()
fig.set_tight_layout(tight=True)

# plt.savefig("images/plot_sample_size_normal.pdf")

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_df = sample_size_results.query(f"model in ['lr', 'nn_8', 'nn_32'] and method == '{plot_method}'")
plot_df["model"] = plot_df["model"].replace(renaming_dict)
plot_df["acc_diff"] = plot_df["test_acc_b"] - plot_df["test_acc_a"]

plot_df = plot_df.rename(columns={
    "model": "Model",
    "treatment_size": "Subgroup size",
    "test_acc": "Test acc." ,
    "test_acc_a": "Majority group acc." ,
    "test_acc_b": "Minority group acc." ,
    "acc_diff": "Acc. discrepancy",
})

sns.barplot(data=plot_df,
             x="Subgroup size", y="Test acc.",
             hue="Model", ax=ax)
# ax.set_xscale("log")
fig.set_tight_layout(tight=True)
ax.set_xticklabels(treatment_group_sizes)
ax.get_legend().remove()

# plt.savefig("images/plot_accuracy_normal.pdf")


# %%
model_names = ["lr", "lr_dm_threshold", "lr_eo_threshold"]
treatment_group_sizes = [100, 500, 1000, 2500]

# %%
fair_results = pd.DataFrame()

# %%
if not RESTORE_SAVED_DATA:
    num_models = NUM_SHADOWS + NUM_EVAL_MODELS
    num_targets = NUM_EVAL_MODELS
    methods = ["shadow_attack_preds", "shadow_attack_loss"]
    global_setup = create_multinomial_setup(num_features=10, seed=BASE_SEED)

    for model, method in itertools.product(model_names, methods):
        print(model, method)
        model_fn = model_zoo[model]
        sim_res = run_sim_all_sizes(
            model_fn,
            treatment_group_sizes,
            dataset_gen=SimDatasetGen(global_setup),
            seed=BASE_SEED,
            method=method,
            num_models=num_models,
            num_targets=num_targets,
            sensitive_fit=any(keyword in model for keyword in ["eo", "erp", "dm"]),
            sensitive_pred=("threshold" in model),
            parallel=True,
        )
        display(sim_res.groupby("treatment_size").agg({
            "vuln_a": ["mean", "std"],
            "vuln_b": ["mean", "std"],
            "baseline_acc": "mean",
            "test_acc": "mean",
        }))

        fair_results = fair_results.append(
            sim_res.assign(model=model, method=method), ignore_index=True)

# %%
if RESTORE_SAVED_DATA:
    fair_results = pd.read_csv(DATA_FILE_FORMAT.format("fair"))
else:
    fair_results.to_csv(DATA_FILE_FORMAT.format("fair"))

# %%
# Compute p-values.
methods = fair_results.method.unique()
for method, model, size in itertools.product(methods, model_names, treatment_group_sizes):
    df = fair_results.query(f"model == '{model}' and treatment_size == {size} and method == '{method}'")
    t, p = stats.ttest_rel(df.vuln_a, df.vuln_b)
    fair_results.loc[df.index, "p"] = p
    fair_results.loc[df.index, "t"] = t

fair_results.groupby(["method", "model", "treatment_size"]).p.mean().reset_index()

# %%
unrolled_list = []
for method, model, size in itertools.product(methods, model_names, treatment_group_sizes):
    df = fair_results.query(f"model == '{model}' and treatment_size == {size} and method == '{method}'").copy()
    dataset_size = size + treatment_group_sizes[-1]
    unrolled = pd.concat([
        pd.DataFrame(dict(vuln=df.vuln_a, target_test_acc=df.test_acc_a)).assign(subgroup="a"),
        pd.DataFrame(dict(vuln=df.vuln_b, target_test_acc=df.test_acc_b)).assign(subgroup="b"),
        pd.DataFrame(dict(vuln=df.vuln_a * size/dataset_size + df.vuln_b * treatment_group_sizes[-1]/dataset_size,
                          target_test_acc=df.test_acc_b)).assign(subgroup="All"),
    ], ignore_index=False).assign(method=method, model=model, treatment_group_size=size)
    unrolled_list.append(unrolled)

summary_data = pd.concat(unrolled_list, ignore_index=True)

# %%
summary_data.replace(renaming_dict).query("subgroup == 'All'"
).groupby(["model"]).agg({
    "target_test_acc": ["mean", "std"],
    "vuln": [normalize_vuln_mean, normalize_vuln_std],
}).reindex([renaming_dict[model_name] for model_name in model_names]).rename(columns={
    "model": "Model",
    "target_test_acc": "Test acc.",
    "vuln": "Overall vuln.",
})

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_models = ["lr", "lr_dm_threshold", "lr_eo_threshold"]
plot_df = fair_results.query(f"model in {plot_models}")
plot_df["model"] = plot_df["model"].replace(renaming_dict)
plot_df["acc_diff"] = plot_df["test_acc_b"] - plot_df["test_acc_a"]

plot_df = plot_df.rename(columns={
    "model": "Model",
    "treatment_size": "Subgroup size",
    "test_acc": "Test acc." ,
    "test_acc_a": "Majority group acc." ,
    "test_acc_b": "Minority group acc." ,
    "acc_diff": "Acc. discrepancy",
})

sns.barplot(data=plot_df,
             x="Subgroup size", y="Minority group acc.",
             hue="Model", ax=ax)
# ax.set_xscale("log")
ax.set_xticklabels(treatment_group_sizes)
fig.set_tight_layout(tight=True)
# ax.get_legend().remove()


# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_models = ["lr", "lr_dm_threshold", "lr_eo_threshold"]
plot_method = "shadow_attack_preds"
plot_df = fair_results.query(f"model in {plot_models} and method == '{plot_method}'")
plot_df["vuln_diff"] = 2 * (plot_df.vuln_b - plot_df.vuln_a) * 100
plot_df["model"] = plot_df["model"].replace(renaming_dict)

plot_df = plot_df.rename(columns={
    "model": "Model",
    "treatment_size": "Subgroup size",
    "vuln_b": "Subgroup vuln.",
    "vuln_diff": "Disparity",
})

sns.barplot(
    data=plot_df,
    x="Subgroup size", y="Disparity",
    hue="Model",
    ax=ax,
    hue_order=[renaming_dict[model] for model in plot_models]
)
ax.set_ylim(-0.5, 2.2)
# ax.set_xscale("log")

add_significance_markers(ax, fair_results.query(f"method == '{plot_method}'"),
                         hue_values=plot_models, x_values=treatment_group_sizes)

ax.set_xticklabels(treatment_group_sizes)
fig.set_tight_layout(tight=True)

# plt.savefig("images/plot_sample_size_fair_preds.pdf")

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_models = ["lr", "lr_dm_threshold", "lr_eo_threshold"]
plot_method = "shadow_attack_loss"
plot_df = fair_results.query(f"model in {plot_models} and method == '{plot_method}'")
plot_df["vuln_diff"] = 2 * (plot_df.vuln_b - plot_df.vuln_a) * 100
plot_df["model"] = plot_df["model"].replace(renaming_dict)

plot_df = plot_df.rename(columns={
    "model": "Model",
    "treatment_size": "Subgroup size",
    "vuln_b": "Subgroup vuln.",
    "vuln_diff": "Disparity",
})

sns.barplot(
    data=plot_df,
    x="Subgroup size", y="Disparity",
    hue="Model",
    ax=ax,
    hue_order=[renaming_dict[model] for model in plot_models]
)
ax.set_ylim(-0.5, 2.2)
# ax.set_xscale("log")

add_significance_markers(ax, fair_results.query(f"method == '{plot_method}'"),
                         hue_values=plot_models, x_values=treatment_group_sizes)

ax.set_xticklabels(treatment_group_sizes)
fig.set_tight_layout(tight=True)
plt.legend().remove()

# plt.savefig("images/plot_sample_size_fair_loss.pdf")

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_models = ["lr", "lr_dm_threshold", "lr_eo_threshold"]
plot_df = fair_results.query(f"model in {plot_models}")
plot_df["model"] = plot_df["model"].replace(renaming_dict)

plot_df = plot_df.rename(columns={
    "model": "Model",
    "treatment_size": "Subgroup size",
    "test_acc": "Test acc." ,
})

sns.barplot(data=plot_df,
            x="Subgroup size", y="Test acc.",
            hue="Model", ax=ax)
# ax.set_xscale("log")
ax.get_legend().remove()
ax.set_xticklabels(treatment_group_sizes)
fig.set_tight_layout(tight=True)

# plt.savefig("images/plot_accuracy_fair.pdf")

# %% [markdown]
# ## Bias study

# %%
bias_treatment_group_sizes = [25, 50, 75, 100, 150, 200, 500, 1000]

# %%
bias_results = pd.DataFrame()

# %%
if not RESTORE_SAVED_DATA:
    num_models = NUM_SHADOWS + NUM_EVAL_MODELS
    num_targets = NUM_EVAL_MODELS
    global_setup = create_multinomial_setup(**DATA_SETUP)

    bias_methods = ["best_loss_threshold", "average_loss_threshold", "shadow_attack_loss"]
    for i, method in enumerate(bias_methods):
        print(method)
        model_fn = model_zoo["threshold"]
        sim_res = run_sim_all_sizes(
            model_fn,
            bias_treatment_group_sizes,
            dataset_gen=SimDatasetGen(global_setup),
            seed=BASE_SEED,  # Use the same datasets.
            method=method,
            num_models=num_models,
            num_targets=num_targets if "shadow" in method else None,
            parallel=True,
        )
        bias_results = bias_results.append(
            sim_res.assign(method=method), ignore_index=True)

# %%
if RESTORE_SAVED_DATA:
    bias_results = pd.read_csv(DATA_FILE_FORMAT.format("bias"))
else:
    bias_results.to_csv(DATA_FILE_FORMAT.format("bias"))

# %%
for method, size in itertools.product(bias_methods, bias_treatment_group_sizes):
    df = bias_results.query(f"treatment_size == {size} and method == '{method}'")
    t, p = stats.ttest_rel(df.vuln_a, df.vuln_b)
    bias_results.loc[df.index, "p"] = p
    bias_results.loc[df.index, "t"] = t

bias_results.groupby(["method", "treatment_size"]).p.mean().reset_index()

# %%
fig, ax = plt.subplots(figsize=(12, 8))

plot_df = bias_results.copy()
plot_df["vuln_diff"] = 2 * (plot_df.vuln_b - plot_df.vuln_a) * 100
plot_df = plot_df.replace({
    "best_loss_threshold": "Opt. loss threshold",
    "average_loss_threshold": "Avg. loss threshold",
    "shadow_attack_loss": "Shadow model attack",
})

plot_df = plot_df.rename(columns={
    "treatment_size": "Subgroup size",
    "vuln_b": "Subgroup vuln.",
    "vuln_diff": "Disparity estimate",
})

sns.barplot(
    data=plot_df,
    x="Subgroup size", y="Disparity estimate",
    hue="method",
    ax=ax,
)
# ax.set_xscale("log")

add_significance_markers(ax, bias_results,
                         x_values=bias_treatment_group_sizes, hue_values=bias_methods,
                         x_label="treatment_size", hue_label="method")

ax.set_xticklabels(bias_treatment_group_sizes)
fig.set_tight_layout(tight=True)

# plt.savefig("images/plot_estimation_bias.pdf")