upload

README.md

@@ -1,13 +1,30 @@
 ---
-title: High Fidelity Image Counterfactuals With Probabilistic Causal Models
-emoji: 🏢
-colorFrom: pink
-colorTo: gray
+title: Counterfactuals
+emoji: 🌖
+colorFrom: purple
+colorTo: green
 sdk: gradio
-sdk_version: 4.41.0
+sdk_version: 3.27.0
 app_file: app.py
 pinned: false
 license: mit
+duplicated_from: fabio-deep/counterfactuals
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Code for the **ICML 2023** paper:
+
+[**High Fidelity Image Counterfactuals with Probabilistic Causal Models**](https://arxiv.org/abs/2306.15764)
+
+Fabio De Sousa Ribeiro<sup>1</sup>, Tian Xia<sup>1</sup>, Miguel Monteiro<sup>1</sup>, Nick Pawlowski<sup>2</sup>, Ben Glocker<sup>1</sup>\
+<sup>1</sup>Imperial College London, <sup>2</sup>Microsoft Research Cambridge, UK
+
+```
+@misc{ribeiro2023high,
+      title={High Fidelity Image Counterfactuals with Probabilistic Causal Models}, 
+      author={Fabio De Sousa Ribeiro and Tian Xia and Miguel Monteiro and Nick Pawlowski and Ben Glocker},
+      year={2023},
+      eprint={2306.15764},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG}
+}
+```

app.py

@@ -0,0 +1,702 @@
+import torch
+import numpy as np
+import gradio as gr
+import matplotlib.pylab as plt
+import torch.nn.functional as F
+
+from vae import HVAE
+from datasets import morphomnist, ukbb, mimic, get_attr_max_min
+from pgm.flow_pgm import MorphoMNISTPGM, FlowPGM, ChestPGM
+from app_utils import (
+    mnist_graph,
+    brain_graph,
+    chest_graph,
+    vae_preprocess,
+    normalize,
+    preprocess_brain,
+    get_fig_arr,
+    postprocess,
+    MidpointNormalize,
+)
+
+DATA, MODELS = {}, {}
+for k in ["Morpho-MNIST", "Brain MRI", "Chest X-ray"]:
+    DATA[k], MODELS[k] = {}, {}
+
+# mnist
+DIGITS = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
+# brain
+MRISEQ_CAT = ["T1", "T2-FLAIR"]  # 0,1
+SEX_CAT = ["female", "male"]  # 0,1
+HEIGHT, WIDTH = 270, 270
+# chest
+SEX_CAT_CHEST = ["male", "female"]  # 0,1
+RACE_CAT = ["white", "asian", "black"]  # 0,1,2
+FIND_CAT = ["no disease", "pleural effusion"]
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+class Hparams:
+    def update(self, dict):
+        for k, v in dict.items():
+            setattr(self, k, v)
+
+
+def get_paths(dataset_id):
+    if "MNIST" in dataset_id:
+        data_path = "./data/morphomnist"
+        pgm_path = "./checkpoints/t_i_d/sup_pgm/checkpoint.pt"
+        vae_path = "./checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt"
+    elif "Brain" in dataset_id:
+        data_path = "./data/ukbb_subset"
+        pgm_path = "./checkpoints/m_b_v_s/sup_pgm/checkpoint.pt"
+        vae_path = "./checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt"
+    elif "Chest" in dataset_id:
+        data_path = "./data/mimic_subset"
+        pgm_path = "./checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt"
+        vae_path = [
+            "./checkpoints/a_r_s_f/mimic_beta9_gelu_dgauss_1_lr3/checkpoint.pt",  # base vae
+            "./checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt",  # cf trained DSCM
+        ]
+    return data_path, vae_path, pgm_path
+
+
+def load_pgm(dataset_id, pgm_path):
+    checkpoint = torch.load(pgm_path, map_location=DEVICE)
+    args = Hparams()
+    args.update(checkpoint["hparams"])
+    args.device = DEVICE
+    if "MNIST" in dataset_id:
+        pgm = MorphoMNISTPGM(args).to(args.device)
+    elif "Brain" in dataset_id:
+        pgm = FlowPGM(args).to(args.device)
+    elif "Chest" in dataset_id:
+        pgm = ChestPGM(args).to(args.device)
+    pgm.load_state_dict(checkpoint["ema_model_state_dict"])
+    MODELS[dataset_id]["pgm"] = pgm
+    MODELS[dataset_id]["pgm_args"] = args
+
+
+def load_vae(dataset_id, vae_path):
+    if "Chest" in dataset_id:
+        vae_path, dscm_path = vae_path[0], vae_path[1]
+    checkpoint = torch.load(vae_path, map_location=DEVICE)
+    args = Hparams()
+    args.update(checkpoint["hparams"])
+    # backwards compatibility hack
+    if not hasattr(args, "vae"):
+        args.vae = "hierarchical"
+    if not hasattr(args, "cond_prior"):
+        args.cond_prior = False
+    if hasattr(args, "free_bits"):
+        args.kl_free_bits = args.free_bits
+    args.device = DEVICE
+    vae = HVAE(args).to(args.device)
+
+    if "Chest" in dataset_id:
+        dscm_ckpt = torch.load(dscm_path, map_location=DEVICE)
+        vae.load_state_dict(
+            {
+                k[4:]: v
+                for k, v in dscm_ckpt["ema_model_state_dict"].items()
+                if "vae." in k
+            }
+        )
+    else:
+        vae.load_state_dict(checkpoint["ema_model_state_dict"])
+    MODELS[dataset_id]["vae"] = vae
+    MODELS[dataset_id]["vae_args"] = args
+
+
+def get_dataloader(dataset_id, data_path):
+    MODELS[dataset_id]["pgm_args"].data_dir = data_path
+    args = MODELS[dataset_id]["pgm_args"]
+    if "MNIST" in dataset_id:
+        datasets = morphomnist(args)
+    elif "Brain" in dataset_id:
+        datasets = ukbb(args)
+    elif "Chest" in dataset_id:
+        datasets = mimic(args)
+    DATA[dataset_id]["test"] = torch.utils.data.DataLoader(
+        datasets["test"], shuffle=False, batch_size=args.bs, num_workers=4
+    )
+
+
+def load_dataset(dataset_id):
+    data_path, _, pgm_path = get_paths(dataset_id)
+    checkpoint = torch.load(pgm_path, map_location=DEVICE)
+    args = Hparams()
+    args.update(checkpoint["hparams"])
+    args.device = DEVICE
+    MODELS[dataset_id]["pgm_args"] = args
+    get_dataloader(dataset_id, data_path)
+
+
+def load_model(dataset_id):
+    _, vae_path, pgm_path = get_paths(dataset_id)
+    load_pgm(dataset_id, pgm_path)
+    load_vae(dataset_id, vae_path)
+
+
+@torch.no_grad()
+def counterfactual_inference(dataset_id, obs, do_pa):
+    pa = {k: v.clone() for k, v in obs.items() if k != "x"}
+    cf_pa = MODELS[dataset_id]["pgm"].counterfactual(
+        obs=pa, intervention=do_pa, num_particles=1
+    )
+    args, vae = MODELS[dataset_id]["vae_args"], MODELS[dataset_id]["vae"]
+    _pa = vae_preprocess(args, {k: v.clone() for k, v in pa.items()})
+    _cf_pa = vae_preprocess(args, {k: v.clone() for k, v in cf_pa.items()})
+    z_t = 0.1 if "mnist" in args.hps else 1.0
+    z = vae.abduct(x=obs["x"], parents=_pa, t=z_t)
+    if vae.cond_prior:
+        z = [z[j]["z"] for j in range(len(z))]
+    px_loc, px_scale = vae.forward_latents(latents=z, parents=_pa)
+    cf_loc, cf_scale = vae.forward_latents(latents=z, parents=_cf_pa)
+    u = (obs["x"] - px_loc) / px_scale.clamp(min=1e-12)
+    u_t = 0.1 if "mnist" in args.hps else 1.0  # cf sampling temp
+    cf_scale = cf_scale * u_t
+    cf_x = torch.clamp(cf_loc + cf_scale * u, min=-1, max=1)
+    return {"cf_x": cf_x, "rec_x": px_loc, "cf_pa": cf_pa}
+
+
+def get_obs_item(dataset_id, idx=None):
+    if idx is None:
+        n_test = len(DATA[dataset_id]["test"].dataset)
+        idx = torch.randperm(n_test)[0]
+    idx = int(idx)
+    return idx, DATA[dataset_id]["test"].dataset.__getitem__(idx)
+
+
+def get_mnist_obs(idx=None):
+    dataset_id = "Morpho-MNIST"
+    if not DATA[dataset_id]:
+        load_dataset(dataset_id)
+    idx, obs = get_obs_item(dataset_id, idx)
+    x = get_fig_arr(obs["x"].clone().squeeze().numpy())
+    t = (obs["thickness"].clone() + 1) / 2 * (6.255515 - 0.87598526) + 0.87598526
+    i = (obs["intensity"].clone() + 1) / 2 * (254.90317 - 66.601204) + 66.601204
+    y = DIGITS[obs["digit"].clone().argmax(-1)]
+    return (idx, x, float(np.round(t, 2)), float(np.round(i, 2)), y)
+
+
+def get_brain_obs(idx=None):
+    dataset_id = "Brain MRI"
+    if not DATA[dataset_id]:
+        load_dataset(dataset_id)
+    idx, obs = get_obs_item(dataset_id, idx)
+    x = get_fig_arr(obs["x"].clone().squeeze().numpy())
+    m = MRISEQ_CAT[int(obs["mri_seq"].clone().item())]
+    s = SEX_CAT[int(obs["sex"].clone().item())]
+    a = obs["age"].clone().item()
+    b = obs["brain_volume"].clone().item() / 1000  # in ml
+    v = obs["ventricle_volume"].clone().item() / 1000  # in ml
+    return (idx, x, m, s, a, float(np.round(b, 2)), float(np.round(v, 2)))
+
+
+def get_chest_obs(idx=None):
+    dataset_id = "Chest X-ray"
+    if not DATA[dataset_id]:
+        load_dataset(dataset_id)
+    idx, obs = get_obs_item(dataset_id, idx)
+    x = get_fig_arr(postprocess(obs["x"].clone()))
+    s = SEX_CAT_CHEST[int(obs["sex"].clone().squeeze().numpy())]
+    f = FIND_CAT[int(obs["finding"].clone().squeeze().numpy())]
+    r = RACE_CAT[obs["race"].clone().squeeze().numpy().argmax(-1)]
+    a = (obs["age"].clone().squeeze().numpy() + 1) * 50
+    return (idx, x, r, s, f, float(np.round(a, 1)))
+
+
+def infer_mnist_cf(*args):
+    dataset_id = "Morpho-MNIST"
+    idx, _, t, i, y, do_t, do_i, do_y = args
+    n_particles = 32
+    # preprocess
+    obs = DATA[dataset_id]["test"].dataset.__getitem__(int(idx))
+    obs["x"] = (obs["x"] - 127.5) / 127.5
+    for k, v in obs.items():
+        obs[k] = v.view(1, 1) if len(v.shape) < 1 else v.unsqueeze(0)
+        obs[k] = obs[k].to(MODELS[dataset_id]["vae_args"].device).float()
+        if n_particles > 1:
+            ndims = (1,) * 3 if k == "x" else (1,)
+            obs[k] = obs[k].repeat(n_particles, *ndims)
+    # intervention(s)
+    do_pa = {}
+    if do_t:
+        do_pa["thickness"] = torch.tensor(
+            normalize(t, x_max=6.255515, x_min=0.87598526)
+        ).view(1, 1)
+    if do_i:
+        do_pa["intensity"] = torch.tensor(
+            normalize(i, x_max=254.90317, x_min=66.601204)
+        ).view(1, 1)
+    if do_y:
+        do_pa["digit"] = F.one_hot(torch.tensor(DIGITS.index(y)), num_classes=10).view(
+            1, 10
+        )
+
+    for k, v in do_pa.items():
+        do_pa[k] = (
+            v.to(MODELS[dataset_id]["vae_args"].device).float().repeat(n_particles, 1)
+        )
+    # infer counterfactual
+    out = counterfactual_inference(dataset_id, obs, do_pa)
+    # avg cf particles
+    cf_x = out["cf_x"].mean(0)
+    cf_x_std = out["cf_x"].std(0)
+    rec_x = out["rec_x"].mean(0)
+    cf_t = out["cf_pa"]["thickness"].mean(0)
+    cf_i = out["cf_pa"]["intensity"].mean(0)
+    cf_y = out["cf_pa"]["digit"].mean(0)
+    # post process
+    cf_x = postprocess(cf_x)
+    cf_x_std = cf_x_std.squeeze().detach().cpu().numpy()
+    rec_x = postprocess(rec_x)
+    cf_t = np.round((cf_t.item() + 1) / 2 * (6.255515 - 0.87598526) + 0.87598526, 2)
+    cf_i = np.round((cf_i.item() + 1) / 2 * (254.90317 - 66.601204) + 66.601204, 2)
+    cf_y = DIGITS[cf_y.argmax(-1)]
+    # plots
+    # plt.close('all')
+    effect = cf_x - rec_x
+    effect = get_fig_arr(
+        effect, cmap="RdBu_r", norm=MidpointNormalize(vmin=-255, midpoint=0, vmax=255)
+    )
+    cf_x = get_fig_arr(cf_x)
+    cf_x_std = get_fig_arr(cf_x_std, cmap="jet")
+    return (cf_x, cf_x_std, effect, cf_t, cf_i, cf_y)
+
+
+def infer_brain_cf(*args):
+    dataset_id = "Brain MRI"
+    idx, _, m, s, a, b, v = args[:7]
+    do_m, do_s, do_a, do_b, do_v = args[7:]
+    n_particles = 16
+    # preprocessing
+    obs = DATA[dataset_id]["test"].dataset.__getitem__(int(idx))
+    obs = preprocess_brain(MODELS[dataset_id]["vae_args"], obs)
+    for k, _v in obs.items():
+        if n_particles > 1:
+            ndims = (1,) * 3 if k == "x" else (1,)
+            obs[k] = _v.repeat(n_particles, *ndims)
+    # interventions(s)
+    do_pa = {}
+    if do_m:
+        do_pa["mri_seq"] = torch.tensor(MRISEQ_CAT.index(m)).view(1, 1)
+    if do_s:
+        do_pa["sex"] = torch.tensor(SEX_CAT.index(s)).view(1, 1)
+    if do_a:
+        do_pa["age"] = torch.tensor(a).view(1, 1)
+    if do_b:
+        do_pa["brain_volume"] = torch.tensor(b * 1000).view(1, 1)
+    if do_v:
+        do_pa["ventricle_volume"] = torch.tensor(v * 1000).view(1, 1)
+    # normalize continuous attributes
+    for k in ["age", "brain_volume", "ventricle_volume"]:
+        if k in do_pa.keys():
+            k_max, k_min = get_attr_max_min(k)
+            do_pa[k] = (do_pa[k] - k_min) / (k_max - k_min)  # [0,1]
+            do_pa[k] = 2 * do_pa[k] - 1  # [-1,1]
+
+    for k, _v in do_pa.items():
+        do_pa[k] = (
+            _v.to(MODELS[dataset_id]["vae_args"].device).float().repeat(n_particles, 1)
+        )
+    # infer counterfactual
+    out = counterfactual_inference(dataset_id, obs, do_pa)
+    # avg cf particles
+    cf_x = out["cf_x"].mean(0)
+    cf_x_std = out["cf_x"].std(0)
+    rec_x = out["rec_x"].mean(0)
+    cf_m = out["cf_pa"]["mri_seq"].mean(0)
+    cf_s = out["cf_pa"]["sex"].mean(0)
+    # post process
+    cf_x = postprocess(cf_x)
+    cf_x_std = cf_x_std.squeeze().detach().cpu().numpy()
+    rec_x = postprocess(rec_x)
+    cf_m = MRISEQ_CAT[int(cf_m.item())]
+    cf_s = SEX_CAT[int(cf_s.item())]
+    cf_ = {}
+    for k in ["age", "brain_volume", "ventricle_volume"]:  # unnormalize
+        k_max, k_min = get_attr_max_min(k)
+        cf_[k] = (out["cf_pa"][k].mean(0).item() + 1) / 2 * (k_max - k_min) + k_min
+    # plots
+    # plt.close('all')
+    effect = cf_x - rec_x
+    effect = get_fig_arr(
+        effect,
+        cmap="RdBu_r",
+        norm=MidpointNormalize(vmin=effect.min(), midpoint=0, vmax=effect.max()),
+    )
+    cf_x = get_fig_arr(cf_x)
+    cf_x_std = get_fig_arr(cf_x_std, cmap="jet")
+    return (
+        cf_x,
+        cf_x_std,
+        effect,
+        cf_m,
+        cf_s,
+        np.round(cf_["age"], 1),
+        np.round(cf_["brain_volume"] / 1000, 2),
+        np.round(cf_["ventricle_volume"] / 1000, 2),
+    )
+
+
+def infer_chest_cf(*args):
+    dataset_id = "Chest X-ray"
+    idx, _, r, s, f, a = args[:6]
+    do_r, do_s, do_f, do_a = args[6:]
+    n_particles = 16
+    # preprocessing
+    obs = DATA[dataset_id]["test"].dataset.__getitem__(int(idx))
+    for k, v in obs.items():
+        obs[k] = v.to(MODELS[dataset_id]["vae_args"].device).float()
+        if n_particles > 1:
+            ndims = (1,) * 3 if k == "x" else (1,)
+            obs[k] = obs[k].repeat(n_particles, *ndims)
+    # intervention(s)
+    do_pa = {}
+    with torch.no_grad():
+        if do_s:
+            do_pa["sex"] = torch.tensor(SEX_CAT_CHEST.index(s)).view(1, 1)
+        if do_f:
+            do_pa["finding"] = torch.tensor(FIND_CAT.index(f)).view(1, 1)
+        if do_r:
+            do_pa["race"] = F.one_hot(
+                torch.tensor(RACE_CAT.index(r)), num_classes=3
+            ).view(1, 3)
+        if do_a:
+            do_pa["age"] = torch.tensor(a / 100 * 2 - 1).view(1, 1)
+    for k, v in do_pa.items():
+        do_pa[k] = (
+            v.to(MODELS[dataset_id]["vae_args"].device).float().repeat(n_particles, 1)
+        )
+    # infer counterfactual
+    out = counterfactual_inference(dataset_id, obs, do_pa)
+    # avg cf particles
+    cf_x = out["cf_x"].mean(0)
+    cf_x_std = out["cf_x"].std(0)
+    rec_x = out["rec_x"].mean(0)
+    cf_r = out["cf_pa"]["race"].mean(0)
+    cf_s = out["cf_pa"]["sex"].mean(0)
+    cf_f = out["cf_pa"]["finding"].mean(0)
+    cf_a = out["cf_pa"]["age"].mean(0)
+    # post process
+    cf_x = postprocess(cf_x)
+    cf_x_std = cf_x_std.squeeze().detach().cpu().numpy()
+    rec_x = postprocess(rec_x)
+    cf_r = RACE_CAT[cf_r.argmax(-1)]
+    cf_s = SEX_CAT_CHEST[int(cf_s.item())]
+    cf_f = FIND_CAT[int(cf_f.item())]
+    cf_a = (cf_a.item() + 1) * 50
+    # plots
+    # plt.close('all')
+    effect = cf_x - rec_x
+    effect = get_fig_arr(
+        effect,
+        cmap="RdBu_r",
+        norm=MidpointNormalize(vmin=effect.min(), midpoint=0, vmax=effect.max()),
+    )
+    cf_x = get_fig_arr(cf_x)
+    cf_x_std = get_fig_arr(cf_x_std, cmap="jet")
+    return (cf_x, cf_x_std, effect, cf_r, cf_s, cf_f, np.round(cf_a, 1))
+
+
+with gr.Blocks(theme=gr.themes.Default()) as demo:
+    gr.Markdown("# VIOS-Group/High Fidelity Image Counterfactuals with Probabilistic Causal Models")
+    with gr.Tabs():
+        with gr.TabItem("Morpho-MNIST") as mnist_tab:
+            mnist_id = gr.Textbox(value=mnist_tab.label, visible=False)
+
+            with gr.Row().style(equal_height=True):
+                idx = gr.Number(value=0, visible=False)
+                with gr.Column(scale=1, min_width=200):
+                    x = gr.Image(label="Observation", interactive=False).style(
+                        height=HEIGHT
+                    )
+                with gr.Column(scale=1, min_width=200):
+                    cf_x = gr.Image(label="Counterfactual", interactive=False).style(
+                        height=HEIGHT
+                    )
+                with gr.Column(scale=1, min_width=200):
+                    cf_x_std = gr.Image(
+                        label="Counterfactual Uncertainty", interactive=False
+                    ).style(height=HEIGHT)
+                with gr.Column(scale=1, min_width=200):
+                    effect = gr.Image(
+                        label="Direct Causal Effect", interactive=False
+                    ).style(height=HEIGHT)
+            with gr.Row().style(equal_height=True):
+                with gr.Column(scale=1.75):
+                    gr.Markdown(
+                        "**Intervention**"
+                    )
+                    with gr.Column():
+                        do_y = gr.Checkbox(label="do(digit)", value=False)
+                        y = gr.Radio(DIGITS, label="", interactive=False)
+                        with gr.Row():
+                            with gr.Column(min_width=100):
+                                do_t = gr.Checkbox(label="do(thickness)", value=False)
+                                t = gr.Slider(
+                                    label="\u00A0",
+                                    minimum=0.9,
+                                    maximum=5.5,
+                                    step=0.01,
+                                    interactive=False,
+                                )
+                            with gr.Column(min_width=100):
+                                do_i = gr.Checkbox(label="do(intensity)", value=False)
+                                i = gr.Slider(
+                                    label="\u00A0",
+                                    minimum=50,
+                                    maximum=255,
+                                    step=0.01,
+                                    interactive=False,
+                                )
+                    with gr.Row():
+                        new = gr.Button("New Observation")
+                        reset = gr.Button("Reset", variant="stop")
+                        submit = gr.Button("Submit", variant="primary")
+                with gr.Column(scale=1):
+                    gr.Markdown("### &nbsp;")
+                    causal_graph = gr.Image(
+                        label="Causal Graph", interactive=False
+                    ).style(height=300)
+
+        with gr.TabItem("Brain MRI") as brain_tab:
+            brain_id = gr.Textbox(value=brain_tab.label, visible=False)
+
+            with gr.Row().style(equal_height=True):
+                idx_brain = gr.Number(value=0, visible=False)
+                with gr.Column(scale=1, min_width=200):
+                    x_brain = gr.Image(label="Observation", interactive=False).style(
+                        height=HEIGHT
+                    )
+                with gr.Column(scale=1, min_width=200):
+                    cf_x_brain = gr.Image(
+                        label="Counterfactual", interactive=False
+                    ).style(height=HEIGHT)
+                with gr.Column(scale=1, min_width=200):
+                    cf_x_std_brain = gr.Image(
+                        label="Counterfactual Uncertainty", interactive=False
+                    ).style(height=HEIGHT)
+                with gr.Column(scale=1, min_width=200):
+                    effect_brain = gr.Image(
+                        label="Direct Causal Effect", interactive=False
+                    ).style(height=HEIGHT)
+            with gr.Row():
+                with gr.Column(scale=2.55):
+                    gr.Markdown(
+                        "**Intervention**"
+                    )
+                    with gr.Row():
+                        with gr.Column(min_width=200):
+                            do_m = gr.Checkbox(label="do(MRI sequence)", value=False)
+                            m = gr.Radio(
+                                ["T1", "T2-FLAIR"], label="", interactive=False
+                            )
+                        with gr.Column(min_width=200):
+                            do_s = gr.Checkbox(label="do(sex)", value=False)
+                            s = gr.Radio(
+                                ["female", "male"], label="", interactive=False
+                            )
+                    with gr.Row():
+                        with gr.Column(min_width=100):
+                            do_a = gr.Checkbox(label="do(age)", value=False)
+                            a = gr.Slider(
+                                label="\u00A0",
+                                value=50,
+                                minimum=44,
+                                maximum=73,
+                                step=1,
+                                interactive=False,
+                            )
+                        with gr.Column(min_width=100):
+                            do_b = gr.Checkbox(label="do(brain volume)", value=False)
+                            b = gr.Slider(
+                                label="\u00A0",
+                                value=1000,
+                                minimum=850,
+                                maximum=1550,
+                                step=20,
+                                interactive=False,
+                            )
+                        with gr.Column(min_width=100):
+                            do_v = gr.Checkbox(
+                                label="do(ventricle volume)", value=False
+                            )
+                            v = gr.Slider(
+                                label="\u00A0",
+                                value=40,
+                                minimum=10,
+                                maximum=125,
+                                step=2,
+                                interactive=False,
+                            )
+                    with gr.Row():
+                        new_brain = gr.Button("New Observation")
+                        reset_brain = gr.Button("Reset", variant="stop")
+                        submit_brain = gr.Button("Submit", variant="primary")
+                with gr.Column(scale=1):
+                    # gr.Markdown("### &nbsp;")
+                    causal_graph_brain = gr.Image(
+                        label="Causal Graph", interactive=False
+                    ).style(height=340)
+
+        with gr.TabItem("Chest X-ray") as chest_tab:
+            chest_id = gr.Textbox(value=chest_tab.label, visible=False)
+
+            with gr.Row().style(equal_height=True):
+                idx_chest = gr.Number(value=0, visible=False)
+                with gr.Column(scale=1, min_width=200):
+                    x_chest = gr.Image(label="Observation", interactive=False).style(
+                        height=HEIGHT
+                    )
+                with gr.Column(scale=1, min_width=200):
+                    cf_x_chest = gr.Image(
+                        label="Counterfactual", interactive=False
+                    ).style(height=HEIGHT)
+                with gr.Column(scale=1, min_width=200):
+                    cf_x_std_chest = gr.Image(
+                        label="Counterfactual Uncertainty", interactive=False
+                    ).style(height=HEIGHT)
+                with gr.Column(scale=1, min_width=200):
+                    effect_chest = gr.Image(
+                        label="Direct Causal Effect", interactive=False
+                    ).style(height=HEIGHT)
+
+            with gr.Row():
+                with gr.Column(scale=2.55):
+                    gr.Markdown(
+                        "**Intervention**"
+                    )
+                    with gr.Row().style(equal_height=True):
+                        with gr.Column(min_width=200):
+                            do_f_chest = gr.Checkbox(label="do(disease)", value=False)
+                            f_chest = gr.Radio(FIND_CAT, label="", interactive=False)
+                        with gr.Column(min_width=200):
+                            do_s_chest = gr.Checkbox(label="do(sex)", value=False)
+                            s_chest = gr.Radio(
+                                SEX_CAT_CHEST, label="", interactive=False
+                            )
+
+                    with gr.Row():
+                        with gr.Column(min_width=200):
+                            do_r_chest = gr.Checkbox(label="do(race)", value=False)
+                            r_chest = gr.Radio(RACE_CAT, label="", interactive=False)
+                        with gr.Column(min_width=200):
+                            do_a_chest = gr.Checkbox(label="do(age)", value=False)
+                            a_chest = gr.Slider(
+                                label="\u00A0", minimum=18, maximum=98, step=1
+                            )
+
+                    with gr.Row():
+                        new_chest = gr.Button("New Observation")
+                        reset_chest = gr.Button("Reset", variant="stop")
+                        submit_chest = gr.Button("Submit", variant="primary")
+                with gr.Column(scale=1):
+                    # gr.Markdown("### &nbsp;")
+                    causal_graph_chest = gr.Image(
+                        label="Causal Graph", interactive=False
+                    ).style(height=345)
+
+    # morphomnist
+    do = [do_t, do_i, do_y]
+    obs = [idx, x, t, i, y]
+    cf_out = [cf_x, cf_x_std, effect]
+
+    # brain
+    do_brain = [do_m, do_s, do_a, do_b, do_v]  # intervention checkboxes
+    obs_brain = [idx_brain, x_brain, m, s, a, b, v]  # observed image/attributes
+    cf_out_brain = [cf_x_brain, cf_x_std_brain, effect_brain]  # counterfactual outputs
+
+    # chest
+    do_chest = [do_r_chest, do_s_chest, do_f_chest, do_a_chest]
+    obs_chest = [idx_chest, x_chest, r_chest, s_chest, f_chest, a_chest]
+    cf_out_chest = [cf_x_chest, cf_x_std_chest, effect_chest]
+
+    # on start: load new observations & causal graph
+    demo.load(fn=get_mnist_obs, inputs=None, outputs=obs)
+    demo.load(fn=mnist_graph, inputs=do, outputs=causal_graph)
+    demo.load(fn=load_model, inputs=mnist_id, outputs=None)
+    demo.load(fn=get_brain_obs, inputs=None, outputs=obs_brain)
+    demo.load(fn=get_chest_obs, inputs=None, outputs=obs_chest)
+
+    demo.load(fn=brain_graph, inputs=do_brain, outputs=causal_graph_brain)
+    demo.load(fn=chest_graph, inputs=do_chest, outputs=causal_graph_chest)
+
+    # on tab select: load models
+    brain_tab.select(fn=load_model, inputs=brain_id, outputs=None)
+    chest_tab.select(fn=load_model, inputs=chest_id, outputs=None)
+
+    # "new" button: load new observations
+    new.click(fn=get_mnist_obs, inputs=None, outputs=obs)
+    new_chest.click(fn=get_chest_obs, inputs=None, outputs=obs_chest)
+    new_brain.click(fn=get_brain_obs, inputs=None, outputs=obs_brain)
+
+    # "new" button: reset causal graphs
+    new.click(fn=mnist_graph, inputs=do, outputs=causal_graph)
+    new_brain.click(fn=brain_graph, inputs=do_brain, outputs=causal_graph_brain)
+    new_chest.click(fn=chest_graph, inputs=do_chest, outputs=causal_graph_chest)
+
+    # "new" button: reset cf output panels
+    for _k, _v in zip(
+        [new, new_brain, new_chest], [cf_out, cf_out_brain, cf_out_chest]
+    ):
+        _k.click(fn=lambda: (gr.update(value=None),) * 3, inputs=None, outputs=_v)
+
+    # "reset" button: reload current observations
+    reset.click(fn=get_mnist_obs, inputs=idx, outputs=obs)
+    reset_brain.click(fn=get_brain_obs, inputs=idx_brain, outputs=obs_brain)
+    reset_chest.click(fn=get_chest_obs, inputs=idx_chest, outputs=obs_chest)
+
+    # "reset" button: deselect intervention checkboxes
+    reset.click(fn=lambda: (gr.update(value=False),) * len(do), inputs=None, outputs=do)
+    reset_brain.click(
+        fn=lambda: (gr.update(value=False),) * len(do_brain),
+        inputs=None,
+        outputs=do_brain,
+    )
+    reset_chest.click(
+        fn=lambda: (gr.update(value=False),) * len(do_chest),
+        inputs=None,
+        outputs=do_chest,
+    )
+
+    # "reset" button: reset cf output panels
+    for _k, _v in zip(
+        [reset, reset_brain, reset_chest], [cf_out, cf_out_brain, cf_out_chest]
+    ):
+        _k.click(fn=lambda: plt.close("all"), inputs=None, outputs=None)
+        _k.click(fn=lambda: (gr.update(value=None),) * 3, inputs=None, outputs=_v)
+
+    # enable mnist interventions when checkbox is selected & update graph
+    for _k, _v in zip(do, [t, i, y]):
+        _k.change(fn=lambda x: gr.update(interactive=x), inputs=_k, outputs=_v)
+        _k.change(mnist_graph, inputs=do, outputs=causal_graph)
+
+    # enable brain interventions when checkbox is selected & update graph
+    for _k, _v in zip(do_brain, [m, s, a, b, v]):
+        _k.change(fn=lambda x: gr.update(interactive=x), inputs=_k, outputs=_v)
+        _k.change(brain_graph, inputs=do_brain, outputs=causal_graph_brain)
+
+    # enable chest interventions when checkbox is selected & update graph
+    for _k, _v in zip(do_chest, [r_chest, s_chest, f_chest, a_chest]):
+        _k.change(fn=lambda x: gr.update(interactive=x), inputs=_k, outputs=_v)
+        _k.change(chest_graph, inputs=do_chest, outputs=causal_graph_chest)
+
+    # "submit" button: infer countefactuals
+    submit.click(fn=infer_mnist_cf, inputs=obs + do, outputs=cf_out + [t, i, y])
+    submit_brain.click(
+        fn=infer_brain_cf,
+        inputs=obs_brain + do_brain,
+        outputs=cf_out_brain + [m, s, a, b, v],
+    )
+    submit_chest.click(
+        fn=infer_chest_cf,
+        inputs=obs_chest + do_chest,
+        outputs=cf_out_chest + [r_chest, s_chest, f_chest, a_chest],
+    )
+
+if __name__ == "__main__":
+    demo.queue()
+    demo.launch()

app_utils.py

@@ -0,0 +1,435 @@
+import torch
+import numpy as np
+import networkx as nx
+import matplotlib.pyplot as plt
+
+from PIL import Image
+
+from matplotlib import rc, patches, colors
+
+rc("font", **{"family": "serif", "serif": ["Roman"]})
+rc("text", usetex=True)
+rc("image", interpolation="none")
+rc("text.latex", preamble=r"\usepackage{amsmath} \usepackage{amssymb}")
+
+from datasets import get_attr_max_min
+
+HAMMER = np.array(Image.open("./hammer.png").resize((35, 35))) / 255
+
+
+class MidpointNormalize(colors.Normalize):
+    def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):
+        self.midpoint = midpoint
+        colors.Normalize.__init__(self, vmin, vmax, clip)
+
+    def __call__(self, value, clip=None):
+        v_ext = np.max([np.abs(self.vmin), np.abs(self.vmax)])
+        x, y = [-v_ext, self.midpoint, v_ext], [0, 0.5, 1]
+        return np.ma.masked_array(np.interp(value, x, y))
+
+
+def postprocess(x):
+    return ((x + 1.0) * 127.5).squeeze().detach().cpu().numpy()
+
+
+def mnist_graph(*args):
+    x, t, i, y = r"$\mathbf{x}$", r"$t$", r"$i$", r"$y$"
+    ut, ui, uy = r"$\mathbf{U}_t$", r"$\mathbf{U}_i$", r"$\mathbf{U}_y$"
+    zx, ex = r"$\mathbf{z}_{1:L}$", r"$\boldsymbol{\epsilon}$"
+
+    G = nx.DiGraph()
+    G.add_edge(t, x)
+    G.add_edge(i, x)
+    G.add_edge(y, x)
+    G.add_edge(t, i)
+    G.add_edge(ut, t)
+    G.add_edge(ui, i)
+    G.add_edge(uy, y)
+    G.add_edge(zx, x)
+    G.add_edge(ex, x)
+
+    pos = {
+        y: (0, 0),
+        uy: (-1, 0),
+        t: (0, 0.5),
+        ut: (0, 1),
+        x: (1, 0),
+        zx: (2, 0.375),
+        ex: (2, 0),
+        i: (1, 0.5),
+        ui: (1, 1),
+    }
+
+    node_c = {}
+    for node in G:
+        node_c[node] = "lightgrey" if node in [x, t, i, y] else "white"
+    node_line_c = {k: "black" for k, _ in node_c.items()}
+    edge_c = {e: "black" for e in G.edges}
+
+    if args[0]:  # do_t
+        edge_c[(ut, t)] = "lightgrey"
+        # G.remove_edge(ut, t)
+        node_line_c[t] = "red"
+    if args[1]:  # do_i
+        edge_c[(ui, i)] = "lightgrey"
+        edge_c[(t, i)] = "lightgrey"
+        # G.remove_edges_from([(ui, i), (t, i)])
+        node_line_c[i] = "red"
+    if args[2]:  # do_y
+        edge_c[(uy, y)] = "lightgrey"
+        # G.remove_edge(uy, y)
+        node_line_c[y] = "red"
+
+    fs = 30
+    options = {
+        "font_size": fs,
+        "node_size": 3000,
+        "node_color": list(node_c.values()),
+        "edgecolors": list(node_line_c.values()),
+        "edge_color": list(edge_c.values()),
+        "linewidths": 2,
+        "width": 2,
+    }
+    plt.close("all")
+    fig, ax = plt.subplots(1, 1, figsize=(6, 4.1))  # , constrained_layout=True)
+    # fig.patch.set_visible(False)
+    ax.margins(x=0.06, y=0.15, tight=False)
+    ax.axis("off")
+    nx.draw_networkx(G, pos, **options, arrowsize=25, arrowstyle="-|>", ax=ax)
+    # need to reuse x, y limits so that the graphs plot the same way before and after removing edges
+    x_lim = (-1.348, 2.348)
+    y_lim = (-0.215, 1.215)
+    ax.set_xlim(x_lim)
+    ax.set_ylim(y_lim)
+    rect = patches.FancyBboxPatch(
+        (1.75, -0.16),
+        0.5,
+        0.7,
+        boxstyle="round, pad=0.05, rounding_size=0",
+        linewidth=2,
+        edgecolor="black",
+        facecolor="none",
+        linestyle="-",
+    )
+    ax.add_patch(rect)
+    ax.text(1.85, 0.65, r"$\mathbf{U}_{\mathbf{x}}$", fontsize=fs)
+
+    if args[0]:  # do_t
+        fig.figimage(HAMMER, 0.26 * fig.bbox.xmax, 0.525 * fig.bbox.ymax, zorder=10)
+    if args[1]:  # do_i
+        fig.figimage(HAMMER, 0.5175 * fig.bbox.xmax, 0.525 * fig.bbox.ymax, zorder=11)
+    if args[2]:  # do_y
+        fig.figimage(HAMMER, 0.26 * fig.bbox.xmax, 0.2 * fig.bbox.ymax, zorder=12)
+
+    fig.tight_layout()
+    fig.canvas.draw()
+    return np.array(fig.canvas.renderer.buffer_rgba())
+
+
+def brain_graph(*args):
+    x, m, s, a, b, v = r"$\mathbf{x}$", r"$m$", r"$s$", r"$a$", r"$b$", r"$v$"
+    um, us, ua, ub, uv = (
+        r"$\mathbf{U}_m$",
+        r"$\mathbf{U}_s$",
+        r"$\mathbf{U}_a$",
+        r"$\mathbf{U}_b$",
+        r"$\mathbf{U}_v$",
+    )
+    zx, ex = r"$\mathbf{z}_{1:L}$", r"$\boldsymbol{\epsilon}$"
+
+    G = nx.DiGraph()
+    G.add_edge(m, x)
+    G.add_edge(s, x)
+    G.add_edge(b, x)
+    G.add_edge(v, x)
+    G.add_edge(zx, x)
+    G.add_edge(ex, x)
+    G.add_edge(a, b)
+    G.add_edge(a, v)
+    G.add_edge(s, b)
+    G.add_edge(um, m)
+    G.add_edge(us, s)
+    G.add_edge(ua, a)
+    G.add_edge(ub, b)
+    G.add_edge(uv, v)
+
+    pos = {
+        x: (0, 0),
+        zx: (-0.25, -1),
+        ex: (0.25, -1),
+        a: (0, 1),
+        ua: (0, 2),
+        s: (1, 0),
+        us: (1, -1),
+        b: (1, 1),
+        ub: (1, 2),
+        m: (-1, 0),
+        um: (-1, -1),
+        v: (-1, 1),
+        uv: (-1, 2),
+    }
+
+    node_c = {}
+    for node in G:
+        node_c[node] = "lightgrey" if node in [x, m, s, a, b, v] else "white"
+    node_line_c = {k: "black" for k, _ in node_c.items()}
+    edge_c = {e: "black" for e in G.edges}
+
+    if args[0]:  # do_m
+        # G.remove_edge(um, m)
+        edge_c[(um, m)] = "lightgrey"
+        node_line_c[m] = "red"
+    if args[1]:  # do_s
+        # G.remove_edge(us, s)
+        edge_c[(us, s)] = "lightgrey"
+        node_line_c[s] = "red"
+    if args[2]:  # do_a
+        # G.remove_edge(ua, a)
+        edge_c[(ua, a)] = "lightgrey"
+        node_line_c[a] = "red"
+    if args[3]:  # do_b
+        # G.remove_edges_from([(ub, b), (s, b), (a, b)])
+        edge_c[(ub, b)] = "lightgrey"
+        edge_c[(s, b)] = "lightgrey"
+        edge_c[(a, b)] = "lightgrey"
+        node_line_c[b] = "red"
+    if args[4]:  # do_v
+        # G.remove_edges_from([(uv, v), (a, v), (b, v)])
+        edge_c[(uv, v)] = "lightgrey"
+        edge_c[(a, v)] = "lightgrey"
+        edge_c[(b, v)] = "lightgrey"
+        node_line_c[v] = "red"
+
+    fs = 30
+    options = {
+        "font_size": fs,
+        "node_size": 3000,
+        "node_color": list(node_c.values()),
+        "edgecolors": list(node_line_c.values()),
+        "edge_color": list(edge_c.values()),
+        "linewidths": 2,
+        "width": 2,
+    }
+
+    plt.close("all")
+    fig, ax = plt.subplots(1, 1, figsize=(5, 5))  # , constrained_layout=True)
+    # fig.patch.set_visible(False)
+    ax.margins(x=0.1, y=0.08, tight=False)
+    ax.axis("off")
+    nx.draw_networkx(G, pos, **options, arrowsize=25, arrowstyle="-|>", ax=ax)
+    # need to reuse x, y limits so that the graphs plot the same way before and after removing edges
+    x_lim = (-1.32, 1.32)
+    y_lim = (-1.414, 2.414)
+    ax.set_xlim(x_lim)
+    ax.set_ylim(y_lim)
+    rect = patches.FancyBboxPatch(
+        (-0.5, -1.325),
+        1,
+        0.65,
+        boxstyle="round, pad=0.05, rounding_size=0",
+        linewidth=2,
+        edgecolor="black",
+        facecolor="none",
+        linestyle="-",
+    )
+    ax.add_patch(rect)
+    # ax.text(1.85, 0.65, r"$\mathbf{U}_{\mathbf{x}}$", fontsize=fs)
+
+    if args[0]:  # do_m
+        fig.figimage(HAMMER, 0.0075 * fig.bbox.xmax, 0.395 * fig.bbox.ymax, zorder=10)
+    if args[1]:  # do_s
+        fig.figimage(HAMMER, 0.72 * fig.bbox.xmax, 0.395 * fig.bbox.ymax, zorder=11)
+    if args[2]:  # do_a
+        fig.figimage(HAMMER, 0.363 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=12)
+    if args[3]:  # do_b
+        fig.figimage(HAMMER, 0.72 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=13)
+    if args[4]:  # do_v
+        fig.figimage(HAMMER, 0.0075 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=14)
+    else:  # b -> v
+        a3 = patches.FancyArrowPatch(
+            (0.86, 1.21),
+            (-0.86, 1.21),
+            connectionstyle="arc3,rad=.3",
+            linewidth=2,
+            arrowstyle="simple, head_width=10, head_length=10",
+            color="k",
+        )
+        ax.add_patch(a3)
+    # print(ax.get_xlim())
+    # print(ax.get_ylim())
+    fig.tight_layout()
+    fig.canvas.draw()
+    return np.array(fig.canvas.renderer.buffer_rgba())
+
+
+def chest_graph(*args):
+    x, a, d, r, s = r"$\mathbf{x}$", r"$a$", r"$d$", r"$r$", r"$s$"
+    ua, ud, ur, us = (
+        r"$\mathbf{U}_a$",
+        r"$\mathbf{U}_d$",
+        r"$\mathbf{U}_r$",
+        r"$\mathbf{U}_s$",
+    )
+    zx, ex = r"$\mathbf{z}_{1:L}$", r"$\boldsymbol{\epsilon}$"
+
+    G = nx.DiGraph()
+    G.add_edge(ua, a)
+    G.add_edge(ud, d)
+    G.add_edge(ur, r)
+    G.add_edge(us, s)
+    G.add_edge(a, d)
+    G.add_edge(d, x)
+    G.add_edge(r, x)
+    G.add_edge(s, x)
+    G.add_edge(ex, x)
+    G.add_edge(zx, x)
+    G.add_edge(a, x)
+
+    pos = {
+        x: (0, 0),
+        a: (-1, 1),
+        d: (0, 1),
+        r: (1, 1),
+        s: (1, 0),
+        ua: (-1, 2),
+        ud: (0, 2),
+        ur: (1, 2),
+        us: (1, -1),
+        zx: (-0.25, -1),
+        ex: (0.25, -1),
+    }
+
+    node_c = {}
+    for node in G:
+        node_c[node] = "lightgrey" if node in [x, a, d, r, s] else "white"
+
+    edge_c = {e: "black" for e in G.edges}
+    node_line_c = {k: "black" for k, _ in node_c.items()}
+
+    if args[0]:  # do_r
+        # G.remove_edge(ur, r)
+        edge_c[(ur, r)] = "lightgrey"
+        node_line_c[r] = "red"
+    if args[1]:  # do_s
+        # G.remove_edges_from([(us, s)])
+        edge_c[(us, s)] = "lightgrey"
+        node_line_c[s] = "red"
+    if args[2]:  # do_f (do_d)
+        # G.remove_edges_from([(ud, d), (a, d)])
+        edge_c[(ud, d)] = "lightgrey"
+        edge_c[(a, d)] = "lightgrey"
+        node_line_c[d] = "red"
+    if args[3]:  # do_a
+        # G.remove_edge(ua, a)
+        edge_c[(ua, a)] = "lightgrey"
+        node_line_c[a] = "red"
+
+    fs = 30
+    options = {
+        "font_size": fs,
+        "node_size": 3000,
+        "node_color": list(node_c.values()),
+        "edgecolors": list(node_line_c.values()),
+        "edge_color": list(edge_c.values()),
+        "linewidths": 2,
+        "width": 2,
+    }
+    plt.close("all")
+    fig, ax = plt.subplots(1, 1, figsize=(5, 5))  # , constrained_layout=True)
+    # fig.patch.set_visible(False)
+    ax.margins(x=0.1, y=0.08, tight=False)
+    ax.axis("off")
+    nx.draw_networkx(G, pos, **options, arrowsize=25, arrowstyle="-|>", ax=ax)
+    # need to reuse x, y limits so that the graphs plot the same way before and after removing edges
+    x_lim = (-1.32, 1.32)
+    y_lim = (-1.414, 2.414)
+    ax.set_xlim(x_lim)
+    ax.set_ylim(y_lim)
+    rect = patches.FancyBboxPatch(
+        (-0.5, -1.325),
+        1,
+        0.65,
+        boxstyle="round, pad=0.05, rounding_size=0",
+        linewidth=2,
+        edgecolor="black",
+        facecolor="none",
+        linestyle="-",
+    )
+    ax.add_patch(rect)
+    ax.text(-0.9, -1.075, r"$\mathbf{U}_{\mathbf{x}}$", fontsize=fs)
+
+    if args[0]:  # do_r
+        fig.figimage(HAMMER, 0.72 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=10)
+    if args[1]:  # do_s
+        fig.figimage(HAMMER, 0.72 * fig.bbox.xmax, 0.395 * fig.bbox.ymax, zorder=11)
+    if args[2]:  # do_f
+        fig.figimage(HAMMER, 0.363 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=12)
+    if args[3]:  # do_a
+        fig.figimage(HAMMER, 0.0075 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=13)
+
+    fig.tight_layout()
+    fig.canvas.draw()
+    return np.array(fig.canvas.renderer.buffer_rgba())
+
+
+def vae_preprocess(args, pa):
+    if "ukbb" in args.hps:
+        # preprocessing ukbb parents for the vae which was originally trained using
+        # log standardized parents. The pgm was trained using [-1,1] normalization
+        # first undo [-1,1] parent preprocessing back to original range
+        for k, v in pa.items():
+            if k != "mri_seq" and k != "sex":
+                pa[k] = (v + 1) / 2  # [-1,1] -> [0,1]
+                _max, _min = get_attr_max_min(k)
+                pa[k] = pa[k] * (_max - _min) + _min
+        # log_standardize parents for vae input
+        for k, v in pa.items():
+            logpa_k = torch.log(v.clamp(min=1e-12))
+            if k == "age":
+                pa[k] = (logpa_k - 4.112339973449707) / 0.11769197136163712
+            elif k == "brain_volume":
+                pa[k] = (logpa_k - 13.965583801269531) / 0.09537758678197861
+            elif k == "ventricle_volume":
+                pa[k] = (logpa_k - 10.345998764038086) / 0.43127763271331787
+    # concatenate parents expand to input res for conditioning the vae
+    pa = torch.cat(
+        [pa[k] if len(pa[k].shape) > 1 else pa[k][..., None] for k in args.parents_x],
+        dim=1,
+    )
+    pa = (
+        pa[..., None, None].repeat(1, 1, *(args.input_res,) * 2).to(args.device).float()
+    )
+    return pa
+
+
+def preprocess_brain(args, obs):
+    obs["x"] = (obs["x"][None, ...].float().to(args.device) - 127.5) / 127.5  # [-1,1]
+    # for all other variables except x
+    for k in [k for k in obs.keys() if k != "x"]:
+        obs[k] = obs[k].float().to(args.device).view(1, 1)
+        if k in ["age", "brain_volume", "ventricle_volume"]:
+            k_max, k_min = get_attr_max_min(k)
+            obs[k] = (obs[k] - k_min) / (k_max - k_min)  # [0,1]
+            obs[k] = 2 * obs[k] - 1  # [-1,1]
+    return obs
+
+
+def get_fig_arr(x, width=4, height=4, dpi=144, cmap="Greys_r", norm=None):
+    fig = plt.figure(figsize=(width, height), dpi=dpi)
+    ax = plt.axes([0, 0, 1, 1], frameon=False)
+    if cmap == "Greys_r":
+        ax.imshow(x, cmap=cmap, vmin=0, vmax=255)
+    else:
+        ax.imshow(x, cmap=cmap, norm=norm)
+    ax.axis("off")
+    fig.canvas.draw()
+    return np.array(fig.canvas.renderer.buffer_rgba())
+
+
+def normalize(x, x_min=None, x_max=None, zero_one=False):
+    if x_min is None:
+        x_min = x.min()
+    if x_max is None:
+        x_max = x.max()
+    x = (x - x_min) / (x_max - x_min)  # [0,1]
+    return x if zero_one else 2 * x - 1  # else [-1,1]

checkpoints/a_r_s_f/mimic_beta9_gelu_dgauss_1_lr3/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5a9b5666f4f221bfd464c10cca5a092fca26abcfdd767cfc81904c45df5b7a9
+size 129987615

checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61b4dd71d4e9d8ff21ead93872620d3522e53402685663e1b1c946879238936b
+size 493467505

checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03dd765e5cdd4d6938e8536e72f597a2a5bc8ed43f4b840d5c289017aa555ace
+size 5328833

checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint_current.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5fcb79767cbad89f1ea481da0af49334834440d0c6b42ec4e8243a602825219d
+size 5328833

checkpoints/m_b_v_s/sup_pgm/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ad6a317c8c05946556b74a1886d4a72328b2ab523aa69f56c7bfe0476d5212f
+size 5397757

checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56920771b4112c430ff42f46b2707a8bb75acbaf2be518c640d6be2157be0c54
+size 279070192

checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:014ccfb191fa8307dc9488379e94d527f2e51159c96417c88b7b17abc6e7c1b9
+size 33772784

checkpoints/t_i_d/sup_pgm/checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42585f99a2f5825645f606b99ccbc7d14a2815eb7723bace03a3d94e3c2c8856
+size 39144