Duplicate from fabio-deep/counterfactuals

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+*.pyc

README.md

@@ -0,0 +1,14 @@
+---
+title: Counterfactuals
+emoji: 🌖
+colorFrom: purple
+colorTo: green
+sdk: gradio
+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

app.py

@@ -0,0 +1,556 @@
+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.pop('pa')
+    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:
+    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.launch()

app_utils.py

@@ -0,0 +1,373 @@
+import torch
+import numpy as np
+import networkx as nx
+import matplotlib.pyplot as plt
+
+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
+
+from PIL import Image
+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((.86, 1.21), (-.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