KayleeJunyu commited on
Commit
a8456b4
·
verified ·
1 Parent(s): 843d390

Upload 178 files

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. README.md +24 -6
  2. app.py +632 -0
  3. app_utils.py +435 -0
  4. checkpoints/a_r_s_f/mimic_beta9_gelu_dgauss_1_lr3/checkpoint.pt +3 -0
  5. checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt +3 -0
  6. checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt +3 -0
  7. checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint_current.pt +3 -0
  8. checkpoints/m_b_v_s/sup_pgm/checkpoint.pt +3 -0
  9. checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt +3 -0
  10. checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt +3 -0
  11. checkpoints/t_i_d/sup_pgm/checkpoint.pt +3 -0
  12. data/mimic_subset/0.jpg +0 -0
  13. data/mimic_subset/1.jpg +0 -0
  14. data/mimic_subset/10.jpg +0 -0
  15. data/mimic_subset/11.jpg +0 -0
  16. data/mimic_subset/12.jpg +0 -0
  17. data/mimic_subset/13.jpg +0 -0
  18. data/mimic_subset/14.jpg +0 -0
  19. data/mimic_subset/15.jpg +0 -0
  20. data/mimic_subset/16.jpg +0 -0
  21. data/mimic_subset/17.jpg +0 -0
  22. data/mimic_subset/18.jpg +0 -0
  23. data/mimic_subset/19.jpg +0 -0
  24. data/mimic_subset/2.jpg +0 -0
  25. data/mimic_subset/20.jpg +0 -0
  26. data/mimic_subset/21.jpg +0 -0
  27. data/mimic_subset/22.jpg +0 -0
  28. data/mimic_subset/23.jpg +0 -0
  29. data/mimic_subset/24.jpg +0 -0
  30. data/mimic_subset/25.jpg +0 -0
  31. data/mimic_subset/26.jpg +0 -0
  32. data/mimic_subset/27.jpg +0 -0
  33. data/mimic_subset/28.jpg +0 -0
  34. data/mimic_subset/29.jpg +0 -0
  35. data/mimic_subset/3.jpg +0 -0
  36. data/mimic_subset/30.jpg +0 -0
  37. data/mimic_subset/31.jpg +0 -0
  38. data/mimic_subset/32.jpg +0 -0
  39. data/mimic_subset/33.jpg +0 -0
  40. data/mimic_subset/34.jpg +0 -0
  41. data/mimic_subset/35.jpg +0 -0
  42. data/mimic_subset/36.jpg +0 -0
  43. data/mimic_subset/37.jpg +0 -0
  44. data/mimic_subset/38.jpg +0 -0
  45. data/mimic_subset/39.jpg +0 -0
  46. data/mimic_subset/4.jpg +0 -0
  47. data/mimic_subset/40.jpg +0 -0
  48. data/mimic_subset/41.jpg +0 -0
  49. data/mimic_subset/42.jpg +0 -0
  50. data/mimic_subset/43.jpg +0 -0
README.md CHANGED
@@ -1,12 +1,30 @@
1
  ---
2
- title: Counterfactual Demo
3
- emoji: 🐢
4
- colorFrom: red
5
- colorTo: indigo
6
  sdk: gradio
7
- sdk_version: 4.14.0
8
  app_file: app.py
9
  pinned: false
 
 
10
  ---
11
 
12
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ title: Counterfactuals
3
+ emoji: 🌖
4
+ colorFrom: purple
5
+ colorTo: green
6
  sdk: gradio
7
+ sdk_version: 3.27.0
8
  app_file: app.py
9
  pinned: false
10
+ license: mit
11
+ duplicated_from: fabio-deep/counterfactuals
12
  ---
13
 
14
+ Code for the **ICML 2023** paper:
15
+
16
+ [**High Fidelity Image Counterfactuals with Probabilistic Causal Models**](https://arxiv.org/abs/2306.15764)
17
+
18
+ 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>\
19
+ <sup>1</sup>Imperial College London, <sup>2</sup>Microsoft Research Cambridge, UK
20
+
21
+ ```
22
+ @misc{ribeiro2023high,
23
+ title={High Fidelity Image Counterfactuals with Probabilistic Causal Models},
24
+ author={Fabio De Sousa Ribeiro and Tian Xia and Miguel Monteiro and Nick Pawlowski and Ben Glocker},
25
+ year={2023},
26
+ eprint={2306.15764},
27
+ archivePrefix={arXiv},
28
+ primaryClass={cs.LG}
29
+ }
30
+ ```
app.py ADDED
@@ -0,0 +1,632 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import gradio as gr
4
+ import matplotlib.pylab as plt
5
+ import torch.nn.functional as F
6
+
7
+ from vae import HVAE
8
+ from datasets import morphomnist, ukbb, mimic, get_attr_max_min
9
+ from pgm.flow_pgm import MorphoMNISTPGM, FlowPGM, ChestPGM
10
+ from app_utils import (
11
+ mnist_graph,
12
+ brain_graph,
13
+ chest_graph,
14
+ vae_preprocess,
15
+ normalize,
16
+ preprocess_brain,
17
+ get_fig_arr,
18
+ postprocess,
19
+ MidpointNormalize,
20
+ )
21
+
22
+ DATA, MODELS = {}, {}
23
+ for k in ["Morpho-MNIST", "Brain MRI", "Chest X-ray"]:
24
+ DATA[k], MODELS[k] = {}, {}
25
+
26
+ # mnist
27
+ DIGITS = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
28
+ # brain
29
+ MRISEQ_CAT = ["T1", "T2-FLAIR"] # 0,1
30
+ SEX_CAT = ["female", "male"] # 0,1
31
+ HEIGHT, WIDTH = 270, 270
32
+ # chest
33
+ SEX_CAT_CHEST = ["male", "female"] # 0,1
34
+ RACE_CAT = ["white", "asian", "black"] # 0,1,2
35
+ FIND_CAT = ["no disease", "pleural effusion"]
36
+ DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
37
+
38
+
39
+ class Hparams:
40
+ def update(self, dict):
41
+ for k, v in dict.items():
42
+ setattr(self, k, v)
43
+
44
+
45
+ def get_paths(dataset_id):
46
+ if "MNIST" in dataset_id:
47
+ data_path = "./data/morphomnist"
48
+ pgm_path = "./checkpoints/t_i_d/sup_pgm/checkpoint.pt"
49
+ vae_path = "./checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt"
50
+ elif "Brain" in dataset_id:
51
+ data_path = "./data/ukbb_subset"
52
+ pgm_path = "./checkpoints/m_b_v_s/sup_pgm/checkpoint.pt"
53
+ vae_path = "./checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt"
54
+ elif "Chest" in dataset_id:
55
+ data_path = "./data/mimic_subset"
56
+ pgm_path = "./checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt"
57
+ vae_path = [
58
+ "./checkpoints/a_r_s_f/mimic_beta9_gelu_dgauss_1_lr3/checkpoint.pt", # base vae
59
+ "./checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt", # cf trained DSCM
60
+ ]
61
+ return data_path, vae_path, pgm_path
62
+
63
+
64
+ def load_pgm(dataset_id, pgm_path):
65
+ checkpoint = torch.load(pgm_path, map_location=DEVICE)
66
+ args = Hparams()
67
+ args.update(checkpoint["hparams"])
68
+ args.device = DEVICE
69
+ if "MNIST" in dataset_id:
70
+ pgm = MorphoMNISTPGM(args).to(args.device)
71
+ elif "Brain" in dataset_id:
72
+ pgm = FlowPGM(args).to(args.device)
73
+ elif "Chest" in dataset_id:
74
+ pgm = ChestPGM(args).to(args.device)
75
+ pgm.load_state_dict(checkpoint["ema_model_state_dict"])
76
+ MODELS[dataset_id]["pgm"] = pgm
77
+ MODELS[dataset_id]["pgm_args"] = args
78
+
79
+
80
+ def load_vae(dataset_id, vae_path):
81
+ if "Chest" in dataset_id:
82
+ vae_path, dscm_path = vae_path[0], vae_path[1]
83
+ checkpoint = torch.load(vae_path, map_location=DEVICE)
84
+ args = Hparams()
85
+ args.update(checkpoint["hparams"])
86
+ # backwards compatibility hack
87
+ if not hasattr(args, "vae"):
88
+ args.vae = "hierarchical"
89
+ if not hasattr(args, "cond_prior"):
90
+ args.cond_prior = False
91
+ if hasattr(args, "free_bits"):
92
+ args.kl_free_bits = args.free_bits
93
+ args.device = DEVICE
94
+ vae = HVAE(args).to(args.device)
95
+
96
+ if "Chest" in dataset_id:
97
+ dscm_ckpt = torch.load(dscm_path, map_location=DEVICE)
98
+ vae.load_state_dict(
99
+ {
100
+ k[4:]: v
101
+ for k, v in dscm_ckpt["ema_model_state_dict"].items()
102
+ if "vae." in k
103
+ }
104
+ )
105
+ else:
106
+ vae.load_state_dict(checkpoint["ema_model_state_dict"])
107
+ MODELS[dataset_id]["vae"] = vae
108
+ MODELS[dataset_id]["vae_args"] = args
109
+
110
+
111
+ def get_dataloader(dataset_id, data_path):
112
+ MODELS[dataset_id]["pgm_args"].data_dir = data_path
113
+ args = MODELS[dataset_id]["pgm_args"]
114
+ if "MNIST" in dataset_id:
115
+ datasets = morphomnist(args)
116
+ elif "Brain" in dataset_id:
117
+ datasets = ukbb(args)
118
+ elif "Chest" in dataset_id:
119
+ datasets = mimic(args)
120
+ DATA[dataset_id]["test"] = torch.utils.data.DataLoader(
121
+ datasets["test"], shuffle=False, batch_size=args.bs, num_workers=4
122
+ )
123
+
124
+
125
+ def load_dataset(dataset_id):
126
+ data_path, _, pgm_path = get_paths(dataset_id)
127
+ checkpoint = torch.load(pgm_path, map_location=DEVICE)
128
+ args = Hparams()
129
+ args.update(checkpoint["hparams"])
130
+ args.device = DEVICE
131
+ MODELS[dataset_id]["pgm_args"] = args
132
+ get_dataloader(dataset_id, data_path)
133
+
134
+
135
+ def load_model(dataset_id):
136
+ _, vae_path, pgm_path = get_paths(dataset_id)
137
+ load_pgm(dataset_id, pgm_path)
138
+ load_vae(dataset_id, vae_path)
139
+
140
+
141
+ @torch.no_grad()
142
+ def counterfactual_inference(dataset_id, obs, do_pa):
143
+ pa = {k: v.clone() for k, v in obs.items() if k != "x"}
144
+ cf_pa = MODELS[dataset_id]["pgm"].counterfactual(
145
+ obs=pa, intervention=do_pa, num_particles=1
146
+ )
147
+ args, vae = MODELS[dataset_id]["vae_args"], MODELS[dataset_id]["vae"]
148
+ _pa = vae_preprocess(args, {k: v.clone() for k, v in pa.items()})
149
+ _cf_pa = vae_preprocess(args, {k: v.clone() for k, v in cf_pa.items()})
150
+ z_t = 0.1 if "mnist" in args.hps else 1.0
151
+ z = vae.abduct(x=obs["x"], parents=_pa, t=z_t)
152
+ if vae.cond_prior:
153
+ z = [z[j]["z"] for j in range(len(z))]
154
+ px_loc, px_scale = vae.forward_latents(latents=z, parents=_pa)
155
+ cf_loc, cf_scale = vae.forward_latents(latents=z, parents=_cf_pa)
156
+ u = (obs["x"] - px_loc) / px_scale.clamp(min=1e-12)
157
+ u_t = 0.1 if "mnist" in args.hps else 1.0 # cf sampling temp
158
+ cf_scale = cf_scale * u_t
159
+ cf_x = torch.clamp(cf_loc + cf_scale * u, min=-1, max=1)
160
+ return {"cf_x": cf_x, "rec_x": px_loc, "cf_pa": cf_pa}
161
+
162
+
163
+ def get_obs_item(dataset_id, idx=None):
164
+ if idx is None:
165
+ n_test = len(DATA[dataset_id]["test"].dataset)
166
+ idx = torch.randperm(n_test)[0]
167
+ idx = int(idx)
168
+ return idx, DATA[dataset_id]["test"].dataset.__getitem__(idx)
169
+
170
+
171
+ def get_mnist_obs(idx=None):
172
+ dataset_id = "Morpho-MNIST"
173
+ if not DATA[dataset_id]:
174
+ load_dataset(dataset_id)
175
+ idx, obs = get_obs_item(dataset_id, idx)
176
+ x = get_fig_arr(obs["x"].clone().squeeze().numpy())
177
+ t = (obs["thickness"].clone() + 1) / 2 * (6.255515 - 0.87598526) + 0.87598526
178
+ i = (obs["intensity"].clone() + 1) / 2 * (254.90317 - 66.601204) + 66.601204
179
+ y = DIGITS[obs["digit"].clone().argmax(-1)]
180
+ return (idx, x, float(np.round(t, 2)), float(np.round(i, 2)), y)
181
+
182
+
183
+ def get_brain_obs(idx=None):
184
+ dataset_id = "Brain MRI"
185
+ if not DATA[dataset_id]:
186
+ load_dataset(dataset_id)
187
+ idx, obs = get_obs_item(dataset_id, idx)
188
+ x = get_fig_arr(obs["x"].clone().squeeze().numpy())
189
+ m = MRISEQ_CAT[int(obs["mri_seq"].clone().item())]
190
+ s = SEX_CAT[int(obs["sex"].clone().item())]
191
+ a = obs["age"].clone().item()
192
+ b = obs["brain_volume"].clone().item() / 1000 # in ml
193
+ v = obs["ventricle_volume"].clone().item() / 1000 # in ml
194
+ return (idx, x, m, s, a, float(np.round(b, 2)), float(np.round(v, 2)))
195
+
196
+
197
+ def get_chest_obs(idx=None):
198
+ dataset_id = "Chest X-ray"
199
+ if not DATA[dataset_id]:
200
+ load_dataset(dataset_id)
201
+ idx, obs = get_obs_item(dataset_id, idx)
202
+ x = get_fig_arr(postprocess(obs["x"].clone()))
203
+ s = SEX_CAT_CHEST[int(obs["sex"].clone().squeeze().numpy())]
204
+ f = FIND_CAT[int(obs["finding"].clone().squeeze().numpy())]
205
+ r = RACE_CAT[obs["race"].clone().squeeze().numpy().argmax(-1)]
206
+ a = (obs["age"].clone().squeeze().numpy() + 1) * 50
207
+ return (idx, x, r, s, f, float(np.round(a, 1)))
208
+
209
+
210
+ def infer_mnist_cf(*args):
211
+ dataset_id = "Morpho-MNIST"
212
+ idx, _, t, i, y, do_t, do_i, do_y = args
213
+ n_particles = 32
214
+ # preprocess
215
+ obs = DATA[dataset_id]["test"].dataset.__getitem__(int(idx))
216
+ obs["x"] = (obs["x"] - 127.5) / 127.5
217
+ for k, v in obs.items():
218
+ obs[k] = v.view(1, 1) if len(v.shape) < 1 else v.unsqueeze(0)
219
+ obs[k] = obs[k].to(MODELS[dataset_id]["vae_args"].device).float()
220
+ if n_particles > 1:
221
+ ndims = (1,) * 3 if k == "x" else (1,)
222
+ obs[k] = obs[k].repeat(n_particles, *ndims)
223
+ # intervention(s)
224
+ do_pa = {}
225
+ if do_t:
226
+ do_pa["thickness"] = torch.tensor(
227
+ normalize(t, x_max=6.255515, x_min=0.87598526)
228
+ ).view(1, 1)
229
+ if do_i:
230
+ do_pa["intensity"] = torch.tensor(
231
+ normalize(i, x_max=254.90317, x_min=66.601204)
232
+ ).view(1, 1)
233
+ if do_y:
234
+ do_pa["digit"] = F.one_hot(torch.tensor(DIGITS.index(y)), num_classes=10).view(
235
+ 1, 10
236
+ )
237
+
238
+ for k, v in do_pa.items():
239
+ do_pa[k] = (
240
+ v.to(MODELS[dataset_id]["vae_args"].device).float().repeat(n_particles, 1)
241
+ )
242
+ # infer counterfactual
243
+ out = counterfactual_inference(dataset_id, obs, do_pa)
244
+ # avg cf particles
245
+ cf_x = out["cf_x"].mean(0)
246
+ cf_x_std = out["cf_x"].std(0)
247
+ rec_x = out["rec_x"].mean(0)
248
+ cf_t = out["cf_pa"]["thickness"].mean(0)
249
+ cf_i = out["cf_pa"]["intensity"].mean(0)
250
+ cf_y = out["cf_pa"]["digit"].mean(0)
251
+ # post process
252
+ cf_x = postprocess(cf_x)
253
+ cf_x_std = cf_x_std.squeeze().detach().cpu().numpy()
254
+ rec_x = postprocess(rec_x)
255
+ cf_t = np.round((cf_t.item() + 1) / 2 * (6.255515 - 0.87598526) + 0.87598526, 2)
256
+ cf_i = np.round((cf_i.item() + 1) / 2 * (254.90317 - 66.601204) + 66.601204, 2)
257
+ cf_y = DIGITS[cf_y.argmax(-1)]
258
+ # plots
259
+ # plt.close('all')
260
+ effect = cf_x - rec_x
261
+ effect = get_fig_arr(
262
+ effect, cmap="RdBu_r", norm=MidpointNormalize(vmin=-255, midpoint=0, vmax=255)
263
+ )
264
+ cf_x = get_fig_arr(cf_x)
265
+ cf_x_std = get_fig_arr(cf_x_std, cmap="jet")
266
+ return (cf_x, cf_x_std, effect, cf_t, cf_i, cf_y)
267
+
268
+
269
+ def infer_brain_cf(*args):
270
+ dataset_id = "Brain MRI"
271
+ idx, _, m, s, a, b, v = args[:7]
272
+ do_m, do_s, do_a, do_b, do_v = args[7:]
273
+ n_particles = 16
274
+ # preprocessing
275
+ obs = DATA[dataset_id]["test"].dataset.__getitem__(int(idx))
276
+ obs = preprocess_brain(MODELS[dataset_id]["vae_args"], obs)
277
+ for k, _v in obs.items():
278
+ if n_particles > 1:
279
+ ndims = (1,) * 3 if k == "x" else (1,)
280
+ obs[k] = _v.repeat(n_particles, *ndims)
281
+ # interventions(s)
282
+ do_pa = {}
283
+ if do_m:
284
+ do_pa["mri_seq"] = torch.tensor(MRISEQ_CAT.index(m)).view(1, 1)
285
+ if do_s:
286
+ do_pa["sex"] = torch.tensor(SEX_CAT.index(s)).view(1, 1)
287
+ if do_a:
288
+ do_pa["age"] = torch.tensor(a).view(1, 1)
289
+ if do_b:
290
+ do_pa["brain_volume"] = torch.tensor(b * 1000).view(1, 1)
291
+ if do_v:
292
+ do_pa["ventricle_volume"] = torch.tensor(v * 1000).view(1, 1)
293
+ # normalize continuous attributes
294
+ for k in ["age", "brain_volume", "ventricle_volume"]:
295
+ if k in do_pa.keys():
296
+ k_max, k_min = get_attr_max_min(k)
297
+ do_pa[k] = (do_pa[k] - k_min) / (k_max - k_min) # [0,1]
298
+ do_pa[k] = 2 * do_pa[k] - 1 # [-1,1]
299
+
300
+ for k, _v in do_pa.items():
301
+ do_pa[k] = (
302
+ _v.to(MODELS[dataset_id]["vae_args"].device).float().repeat(n_particles, 1)
303
+ )
304
+ # infer counterfactual
305
+ out = counterfactual_inference(dataset_id, obs, do_pa)
306
+ # avg cf particles
307
+ cf_x = out["cf_x"].mean(0)
308
+ cf_x_std = out["cf_x"].std(0)
309
+ rec_x = out["rec_x"].mean(0)
310
+ cf_m = out["cf_pa"]["mri_seq"].mean(0)
311
+ cf_s = out["cf_pa"]["sex"].mean(0)
312
+ # post process
313
+ cf_x = postprocess(cf_x)
314
+ cf_x_std = cf_x_std.squeeze().detach().cpu().numpy()
315
+ rec_x = postprocess(rec_x)
316
+ cf_m = MRISEQ_CAT[int(cf_m.item())]
317
+ cf_s = SEX_CAT[int(cf_s.item())]
318
+ cf_ = {}
319
+ for k in ["age", "brain_volume", "ventricle_volume"]: # unnormalize
320
+ k_max, k_min = get_attr_max_min(k)
321
+ cf_[k] = (out["cf_pa"][k].mean(0).item() + 1) / 2 * (k_max - k_min) + k_min
322
+ # plots
323
+ # plt.close('all')
324
+ effect = cf_x - rec_x
325
+ effect = get_fig_arr(
326
+ effect,
327
+ cmap="RdBu_r",
328
+ norm=MidpointNormalize(vmin=effect.min(), midpoint=0, vmax=effect.max()),
329
+ )
330
+ cf_x = get_fig_arr(cf_x)
331
+ cf_x_std = get_fig_arr(cf_x_std, cmap="jet")
332
+ return (
333
+ cf_x,
334
+ cf_x_std,
335
+ effect,
336
+ cf_m,
337
+ cf_s,
338
+ np.round(cf_["age"], 1),
339
+ np.round(cf_["brain_volume"] / 1000, 2),
340
+ np.round(cf_["ventricle_volume"] / 1000, 2),
341
+ )
342
+
343
+
344
+ def infer_chest_cf(*args):
345
+ dataset_id = "Chest X-ray"
346
+ idx, _, r, s, f, a = args[:6]
347
+ do_r, do_s, do_f, do_a = args[6:]
348
+ n_particles = 16
349
+ # preprocessing
350
+ obs = DATA[dataset_id]["test"].dataset.__getitem__(int(idx))
351
+ for k, v in obs.items():
352
+ obs[k] = v.to(MODELS[dataset_id]["vae_args"].device).float()
353
+ if n_particles > 1:
354
+ ndims = (1,) * 3 if k == "x" else (1,)
355
+ obs[k] = obs[k].repeat(n_particles, *ndims)
356
+ # intervention(s)
357
+ do_pa = {}
358
+ with torch.no_grad():
359
+ if do_s:
360
+ do_pa["sex"] = torch.tensor(SEX_CAT_CHEST.index(s)).view(1, 1)
361
+ if do_f:
362
+ do_pa["finding"] = torch.tensor(FIND_CAT.index(f)).view(1, 1)
363
+ if do_r:
364
+ do_pa["race"] = F.one_hot(
365
+ torch.tensor(RACE_CAT.index(r)), num_classes=3
366
+ ).view(1, 3)
367
+ if do_a:
368
+ do_pa["age"] = torch.tensor(a / 100 * 2 - 1).view(1, 1)
369
+ for k, v in do_pa.items():
370
+ do_pa[k] = (
371
+ v.to(MODELS[dataset_id]["vae_args"].device).float().repeat(n_particles, 1)
372
+ )
373
+ # infer counterfactual
374
+ out = counterfactual_inference(dataset_id, obs, do_pa)
375
+ # avg cf particles
376
+ cf_x = out["cf_x"].mean(0)
377
+ cf_x_std = out["cf_x"].std(0)
378
+ rec_x = out["rec_x"].mean(0)
379
+ cf_r = out["cf_pa"]["race"].mean(0)
380
+ cf_s = out["cf_pa"]["sex"].mean(0)
381
+ cf_f = out["cf_pa"]["finding"].mean(0)
382
+ cf_a = out["cf_pa"]["age"].mean(0)
383
+ # post process
384
+ cf_x = postprocess(cf_x)
385
+ cf_x_std = cf_x_std.squeeze().detach().cpu().numpy()
386
+ rec_x = postprocess(rec_x)
387
+ cf_r = RACE_CAT[cf_r.argmax(-1)]
388
+ cf_s = SEX_CAT_CHEST[int(cf_s.item())]
389
+ cf_f = FIND_CAT[int(cf_f.item())]
390
+ cf_a = (cf_a.item() + 1) * 50
391
+ # plots
392
+ # plt.close('all')
393
+ effect = cf_x - rec_x
394
+ effect = get_fig_arr(
395
+ effect,
396
+ cmap="RdBu_r",
397
+ norm=MidpointNormalize(vmin=effect.min(), midpoint=0, vmax=effect.max()),
398
+ )
399
+ cf_x = get_fig_arr(cf_x)
400
+ cf_x_std = get_fig_arr(cf_x_std, cmap="jet")
401
+ return (cf_x, cf_x_std, effect, cf_r, cf_s, cf_f, np.round(cf_a, 1))
402
+
403
+
404
+ with gr.Blocks(theme=gr.themes.Default()) as demo:
405
+ with gr.Tabs():
406
+
407
+ with gr.TabItem("Brain MRI") as brain_tab:
408
+ brain_id = gr.Textbox(value=brain_tab.label, visible=False)
409
+
410
+ with gr.Row().style(equal_height=True):
411
+ idx_brain = gr.Number(value=0, visible=False)
412
+ with gr.Column(scale=1, min_width=200):
413
+ x_brain = gr.Image(label="Observation", interactive=False).style(
414
+ height=HEIGHT
415
+ )
416
+ with gr.Column(scale=1, min_width=200):
417
+ cf_x_brain = gr.Image(
418
+ label="Counterfactual", interactive=False
419
+ ).style(height=HEIGHT)
420
+ with gr.Column(scale=1, min_width=200):
421
+ cf_x_std_brain = gr.Image(
422
+ label="Counterfactual Uncertainty", interactive=False
423
+ ).style(height=HEIGHT)
424
+ with gr.Column(scale=1, min_width=200):
425
+ effect_brain = gr.Image(
426
+ label="Direct Causal Effect", interactive=False
427
+ ).style(height=HEIGHT)
428
+ with gr.Row():
429
+ with gr.Column(scale=2.55):
430
+ gr.Markdown(
431
+ "**Intervention**"
432
+ # + 20 * "&ensp;"
433
+ # + "[arXiv paper](https://arxiv.org/abs/2306.15764) &ensp; | &ensp; [GitHub code](https://github.com/biomedia-mira/causal-gen)"
434
+ # + "&ensp; | &ensp; Hint: try 90% zoom"
435
+ )
436
+ with gr.Row():
437
+ with gr.Column(min_width=200):
438
+ do_a = gr.Checkbox(label="do(age)", value=False)
439
+ a = gr.Slider(
440
+ label="\u00A0",
441
+ value=50,
442
+ minimum=44,
443
+ maximum=73,
444
+ step=1,
445
+ interactive=False,
446
+ )
447
+ with gr.Column(min_width=200):
448
+ do_s = gr.Checkbox(label="do(sex)", value=False)
449
+ s = gr.Radio(
450
+ ["female", "male"], label="", interactive=False
451
+ )
452
+ with gr.Row():
453
+ with gr.Column(min_width=200):
454
+ do_b = gr.Checkbox(label="do(brain volume)", value=False)
455
+ b = gr.Slider(
456
+ label="\u00A0",
457
+ value=1000,
458
+ minimum=850,
459
+ maximum=1550,
460
+ step=20,
461
+ interactive=False,
462
+ )
463
+ with gr.Column(min_width=200):
464
+ do_v = gr.Checkbox(
465
+ label="do(ventricle volume)", value=False
466
+ )
467
+ v = gr.Slider(
468
+ label="\u00A0",
469
+ value=40,
470
+ minimum=10,
471
+ maximum=125,
472
+ step=2,
473
+ interactive=False,
474
+ )
475
+ with gr.Row():
476
+ new_brain = gr.Button("New Observation")
477
+ reset_brain = gr.Button("Reset", variant="stop")
478
+ submit_brain = gr.Button("Submit", variant="primary")
479
+ with gr.Column(scale=1):
480
+ # gr.Markdown("### &nbsp;")
481
+ causal_graph_brain = gr.Image(
482
+ label="Causal Graph", interactive=False
483
+ ).style(height=340)
484
+
485
+ with gr.TabItem("Chest X-ray") as chest_tab:
486
+ chest_id = gr.Textbox(value=chest_tab.label, visible=False)
487
+
488
+ with gr.Row().style(equal_height=True):
489
+ idx_chest = gr.Number(value=0, visible=False)
490
+ with gr.Column(scale=1, min_width=200):
491
+ x_chest = gr.Image(label="Observation", interactive=False).style(
492
+ height=HEIGHT
493
+ )
494
+ with gr.Column(scale=1, min_width=200):
495
+ cf_x_chest = gr.Image(
496
+ label="Counterfactual", interactive=False
497
+ ).style(height=HEIGHT)
498
+ with gr.Column(scale=1, min_width=200):
499
+ cf_x_std_chest = gr.Image(
500
+ label="Counterfactual Uncertainty", interactive=False
501
+ ).style(height=HEIGHT)
502
+ with gr.Column(scale=1, min_width=200):
503
+ effect_chest = gr.Image(
504
+ label="Direct Causal Effect", interactive=False
505
+ ).style(height=HEIGHT)
506
+
507
+ with gr.Row():
508
+ with gr.Column(scale=2.55):
509
+ gr.Markdown(
510
+ "**Intervention**"
511
+ # + 20 * "&ensp;"
512
+ # + "[arXiv paper](https://arxiv.org/abs/2306.15764) &ensp; | &ensp; [GitHub code](https://github.com/biomedia-mira/causal-gen)"
513
+ # + "&ensp; | &ensp; Hint: try 90% zoom"
514
+ )
515
+ with gr.Row().style(equal_height=True):
516
+ with gr.Column(min_width=200):
517
+ do_a_chest = gr.Checkbox(label="do(age)", value=False)
518
+ a_chest = gr.Slider(
519
+ label="\u00A0", minimum=18, maximum=98, step=1
520
+ )
521
+ with gr.Column(min_width=200):
522
+ do_s_chest = gr.Checkbox(label="do(sex)", value=False)
523
+ s_chest = gr.Radio(
524
+ SEX_CAT_CHEST, label="", interactive=False
525
+ )
526
+
527
+ with gr.Row():
528
+ with gr.Column(min_width=200):
529
+ do_r_chest = gr.Checkbox(label="do(race)", value=False)
530
+ r_chest = gr.Radio(RACE_CAT, label="", interactive=False)
531
+ with gr.Column(min_width=200):
532
+ do_f_chest = gr.Checkbox(label="do(disease)", value=False)
533
+ f_chest = gr.Radio(FIND_CAT, label="", interactive=False)
534
+
535
+ with gr.Row():
536
+ new_chest = gr.Button("New Observation")
537
+ reset_chest = gr.Button("Reset", variant="stop")
538
+ submit_chest = gr.Button("Submit", variant="primary")
539
+ with gr.Column(scale=1):
540
+ # gr.Markdown("### &nbsp;")
541
+ causal_graph_chest = gr.Image(
542
+ label="Causal Graph", interactive=False
543
+ ).style(height=345)
544
+
545
+ # morphomnist
546
+ # do = [do_t, do_i, do_y]
547
+ # obs = [idx, x, t, i, y]
548
+ # cf_out = [cf_x, cf_x_std, effect]
549
+
550
+ # brain
551
+ do_brain = [do_s, do_a, do_b, do_v] # intervention checkboxes
552
+ obs_brain = [idx_brain, x_brain, s, a, b, v] # observed image/attributes
553
+ cf_out_brain = [cf_x_brain, cf_x_std_brain, effect_brain] # counterfactual outputs
554
+
555
+ # chest
556
+ do_chest = [do_r_chest, do_s_chest, do_f_chest, do_a_chest]
557
+ obs_chest = [idx_chest, x_chest, r_chest, s_chest, f_chest, a_chest]
558
+ cf_out_chest = [cf_x_chest, cf_x_std_chest, effect_chest]
559
+
560
+ # on start: load new observations & causal graph
561
+ demo.load(fn=get_brain_obs, inputs=None, outputs=obs_brain)
562
+ demo.load(fn=get_chest_obs, inputs=None, outputs=obs_chest)
563
+
564
+ demo.load(fn=brain_graph, inputs=do_brain, outputs=causal_graph_brain)
565
+ demo.load(fn=chest_graph, inputs=do_chest, outputs=causal_graph_chest)
566
+
567
+ # on tab select: load models
568
+ brain_tab.select(fn=load_model, inputs=brain_id, outputs=None)
569
+ chest_tab.select(fn=load_model, inputs=chest_id, outputs=None)
570
+
571
+ # "new" button: load new observations
572
+ new_chest.click(fn=get_chest_obs, inputs=None, outputs=obs_chest)
573
+ new_brain.click(fn=get_brain_obs, inputs=None, outputs=obs_brain)
574
+
575
+ # "new" button: reset causal graphs
576
+ new_brain.click(fn=brain_graph, inputs=do_brain, outputs=causal_graph_brain)
577
+ new_chest.click(fn=chest_graph, inputs=do_chest, outputs=causal_graph_chest)
578
+
579
+ # "new" button: reset cf output panels
580
+ for _k, _v in zip(
581
+ [new_brain, new_chest], [cf_out_brain, cf_out_chest]
582
+ ):
583
+ _k.click(fn=lambda: (gr.update(value=None),) * 3, inputs=None, outputs=_v)
584
+
585
+ # "reset" button: reload current observations
586
+ reset_brain.click(fn=get_brain_obs, inputs=idx_brain, outputs=obs_brain)
587
+ reset_chest.click(fn=get_chest_obs, inputs=idx_chest, outputs=obs_chest)
588
+
589
+ # "reset" button: deselect intervention checkboxes
590
+ reset_brain.click(
591
+ fn=lambda: (gr.update(value=False),) * len(do_brain),
592
+ inputs=None,
593
+ outputs=do_brain,
594
+ )
595
+ reset_chest.click(
596
+ fn=lambda: (gr.update(value=False),) * len(do_chest),
597
+ inputs=None,
598
+ outputs=do_chest,
599
+ )
600
+
601
+ # "reset" button: reset cf output panels
602
+ for _k, _v in zip(
603
+ [reset_brain, reset_chest], [cf_out_brain, cf_out_chest]
604
+ ):
605
+ _k.click(fn=lambda: plt.close("all"), inputs=None, outputs=None)
606
+ _k.click(fn=lambda: (gr.update(value=None),) * 3, inputs=None, outputs=_v)
607
+
608
+ # enable brain interventions when checkbox is selected & update graph
609
+ for _k, _v in zip(do_brain, [s, a, b, v]):
610
+ _k.change(fn=lambda x: gr.update(interactive=x), inputs=_k, outputs=_v)
611
+ _k.change(brain_graph, inputs=do_brain, outputs=causal_graph_brain)
612
+
613
+ # enable chest interventions when checkbox is selected & update graph
614
+ for _k, _v in zip(do_chest, [r_chest, s_chest, f_chest, a_chest]):
615
+ _k.change(fn=lambda x: gr.update(interactive=x), inputs=_k, outputs=_v)
616
+ _k.change(chest_graph, inputs=do_chest, outputs=causal_graph_chest)
617
+
618
+ # "submit" button: infer countefactuals
619
+ submit_brain.click(
620
+ fn=infer_brain_cf,
621
+ inputs=obs_brain + do_brain,
622
+ outputs=cf_out_brain + [s, a, b, v],
623
+ )
624
+ submit_chest.click(
625
+ fn=infer_chest_cf,
626
+ inputs=obs_chest + do_chest,
627
+ outputs=cf_out_chest + [r_chest, s_chest, f_chest, a_chest],
628
+ )
629
+
630
+ if __name__ == "__main__":
631
+ demo.queue()
632
+ demo.launch()
app_utils.py ADDED
@@ -0,0 +1,435 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import networkx as nx
4
+ import matplotlib.pyplot as plt
5
+
6
+ from PIL import Image
7
+
8
+ from matplotlib import rc, patches, colors
9
+
10
+ rc("font", **{"family": "serif", "serif": ["Roman"]})
11
+ rc("text", usetex=True)
12
+ rc("image", interpolation="none")
13
+ rc("text.latex", preamble=r"\usepackage{amsmath} \usepackage{amssymb}")
14
+
15
+ from datasets import get_attr_max_min
16
+
17
+ HAMMER = np.array(Image.open("./hammer.png").resize((35, 35))) / 255
18
+
19
+
20
+ class MidpointNormalize(colors.Normalize):
21
+ def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):
22
+ self.midpoint = midpoint
23
+ colors.Normalize.__init__(self, vmin, vmax, clip)
24
+
25
+ def __call__(self, value, clip=None):
26
+ v_ext = np.max([np.abs(self.vmin), np.abs(self.vmax)])
27
+ x, y = [-v_ext, self.midpoint, v_ext], [0, 0.5, 1]
28
+ return np.ma.masked_array(np.interp(value, x, y))
29
+
30
+
31
+ def postprocess(x):
32
+ return ((x + 1.0) * 127.5).squeeze().detach().cpu().numpy()
33
+
34
+
35
+ def mnist_graph(*args):
36
+ x, t, i, y = r"$\mathbf{x}$", r"$t$", r"$i$", r"$y$"
37
+ ut, ui, uy = r"$\mathbf{U}_t$", r"$\mathbf{U}_i$", r"$\mathbf{U}_y$"
38
+ zx, ex = r"$\mathbf{z}_{1:L}$", r"$\boldsymbol{\epsilon}$"
39
+
40
+ G = nx.DiGraph()
41
+ G.add_edge(t, x)
42
+ G.add_edge(i, x)
43
+ G.add_edge(y, x)
44
+ G.add_edge(t, i)
45
+ G.add_edge(ut, t)
46
+ G.add_edge(ui, i)
47
+ G.add_edge(uy, y)
48
+ G.add_edge(zx, x)
49
+ G.add_edge(ex, x)
50
+
51
+ pos = {
52
+ y: (0, 0),
53
+ uy: (-1, 0),
54
+ t: (0, 0.5),
55
+ ut: (0, 1),
56
+ x: (1, 0),
57
+ zx: (2, 0.375),
58
+ ex: (2, 0),
59
+ i: (1, 0.5),
60
+ ui: (1, 1),
61
+ }
62
+
63
+ node_c = {}
64
+ for node in G:
65
+ node_c[node] = "lightgrey" if node in [x, t, i, y] else "white"
66
+ node_line_c = {k: "black" for k, _ in node_c.items()}
67
+ edge_c = {e: "black" for e in G.edges}
68
+
69
+ if args[0]: # do_t
70
+ edge_c[(ut, t)] = "lightgrey"
71
+ # G.remove_edge(ut, t)
72
+ node_line_c[t] = "red"
73
+ if args[1]: # do_i
74
+ edge_c[(ui, i)] = "lightgrey"
75
+ edge_c[(t, i)] = "lightgrey"
76
+ # G.remove_edges_from([(ui, i), (t, i)])
77
+ node_line_c[i] = "red"
78
+ if args[2]: # do_y
79
+ edge_c[(uy, y)] = "lightgrey"
80
+ # G.remove_edge(uy, y)
81
+ node_line_c[y] = "red"
82
+
83
+ fs = 30
84
+ options = {
85
+ "font_size": fs,
86
+ "node_size": 3000,
87
+ "node_color": list(node_c.values()),
88
+ "edgecolors": list(node_line_c.values()),
89
+ "edge_color": list(edge_c.values()),
90
+ "linewidths": 2,
91
+ "width": 2,
92
+ }
93
+ plt.close("all")
94
+ fig, ax = plt.subplots(1, 1, figsize=(6, 4.1)) # , constrained_layout=True)
95
+ # fig.patch.set_visible(False)
96
+ ax.margins(x=0.06, y=0.15, tight=False)
97
+ ax.axis("off")
98
+ nx.draw_networkx(G, pos, **options, arrowsize=25, arrowstyle="-|>", ax=ax)
99
+ # need to reuse x, y limits so that the graphs plot the same way before and after removing edges
100
+ x_lim = (-1.348, 2.348)
101
+ y_lim = (-0.215, 1.215)
102
+ ax.set_xlim(x_lim)
103
+ ax.set_ylim(y_lim)
104
+ rect = patches.FancyBboxPatch(
105
+ (1.75, -0.16),
106
+ 0.5,
107
+ 0.7,
108
+ boxstyle="round, pad=0.05, rounding_size=0",
109
+ linewidth=2,
110
+ edgecolor="black",
111
+ facecolor="none",
112
+ linestyle="-",
113
+ )
114
+ ax.add_patch(rect)
115
+ ax.text(1.85, 0.65, r"$\mathbf{U}_{\mathbf{x}}$", fontsize=fs)
116
+
117
+ if args[0]: # do_t
118
+ fig.figimage(HAMMER, 0.26 * fig.bbox.xmax, 0.525 * fig.bbox.ymax, zorder=10)
119
+ if args[1]: # do_i
120
+ fig.figimage(HAMMER, 0.5175 * fig.bbox.xmax, 0.525 * fig.bbox.ymax, zorder=11)
121
+ if args[2]: # do_y
122
+ fig.figimage(HAMMER, 0.26 * fig.bbox.xmax, 0.2 * fig.bbox.ymax, zorder=12)
123
+
124
+ fig.tight_layout()
125
+ fig.canvas.draw()
126
+ return np.array(fig.canvas.renderer.buffer_rgba())
127
+
128
+
129
+ def brain_graph(*args):
130
+ x, m, s, a, b, v = r"$\mathbf{x}$", r"$m$", r"$s$", r"$a$", r"$b$", r"$v$"
131
+ um, us, ua, ub, uv = (
132
+ r"$\mathbf{U}_m$",
133
+ r"$\mathbf{U}_s$",
134
+ r"$\mathbf{U}_a$",
135
+ r"$\mathbf{U}_b$",
136
+ r"$\mathbf{U}_v$",
137
+ )
138
+ zx, ex = r"$\mathbf{z}_{1:L}$", r"$\boldsymbol{\epsilon}$"
139
+
140
+ G = nx.DiGraph()
141
+ G.add_edge(m, x)
142
+ G.add_edge(s, x)
143
+ G.add_edge(b, x)
144
+ G.add_edge(v, x)
145
+ G.add_edge(zx, x)
146
+ G.add_edge(ex, x)
147
+ G.add_edge(a, b)
148
+ G.add_edge(a, v)
149
+ G.add_edge(s, b)
150
+ G.add_edge(um, m)
151
+ G.add_edge(us, s)
152
+ G.add_edge(ua, a)
153
+ G.add_edge(ub, b)
154
+ G.add_edge(uv, v)
155
+
156
+ pos = {
157
+ x: (0, 0),
158
+ zx: (-0.25, -1),
159
+ ex: (0.25, -1),
160
+ a: (0, 1),
161
+ ua: (0, 2),
162
+ s: (1, 0),
163
+ us: (1, -1),
164
+ b: (1, 1),
165
+ ub: (1, 2),
166
+ m: (-1, 0),
167
+ um: (-1, -1),
168
+ v: (-1, 1),
169
+ uv: (-1, 2),
170
+ }
171
+
172
+ node_c = {}
173
+ for node in G:
174
+ node_c[node] = "lightgrey" if node in [x, m, s, a, b, v] else "white"
175
+ node_line_c = {k: "black" for k, _ in node_c.items()}
176
+ edge_c = {e: "black" for e in G.edges}
177
+
178
+ if args[0]: # do_m
179
+ # G.remove_edge(um, m)
180
+ edge_c[(um, m)] = "lightgrey"
181
+ node_line_c[m] = "red"
182
+ if args[1]: # do_s
183
+ # G.remove_edge(us, s)
184
+ edge_c[(us, s)] = "lightgrey"
185
+ node_line_c[s] = "red"
186
+ if args[2]: # do_a
187
+ # G.remove_edge(ua, a)
188
+ edge_c[(ua, a)] = "lightgrey"
189
+ node_line_c[a] = "red"
190
+ if args[3]: # do_b
191
+ # G.remove_edges_from([(ub, b), (s, b), (a, b)])
192
+ edge_c[(ub, b)] = "lightgrey"
193
+ edge_c[(s, b)] = "lightgrey"
194
+ edge_c[(a, b)] = "lightgrey"
195
+ node_line_c[b] = "red"
196
+ if args[4]: # do_v
197
+ # G.remove_edges_from([(uv, v), (a, v), (b, v)])
198
+ edge_c[(uv, v)] = "lightgrey"
199
+ edge_c[(a, v)] = "lightgrey"
200
+ edge_c[(b, v)] = "lightgrey"
201
+ node_line_c[v] = "red"
202
+
203
+ fs = 30
204
+ options = {
205
+ "font_size": fs,
206
+ "node_size": 3000,
207
+ "node_color": list(node_c.values()),
208
+ "edgecolors": list(node_line_c.values()),
209
+ "edge_color": list(edge_c.values()),
210
+ "linewidths": 2,
211
+ "width": 2,
212
+ }
213
+
214
+ plt.close("all")
215
+ fig, ax = plt.subplots(1, 1, figsize=(5, 5)) # , constrained_layout=True)
216
+ # fig.patch.set_visible(False)
217
+ ax.margins(x=0.1, y=0.08, tight=False)
218
+ ax.axis("off")
219
+ nx.draw_networkx(G, pos, **options, arrowsize=25, arrowstyle="-|>", ax=ax)
220
+ # need to reuse x, y limits so that the graphs plot the same way before and after removing edges
221
+ x_lim = (-1.32, 1.32)
222
+ y_lim = (-1.414, 2.414)
223
+ ax.set_xlim(x_lim)
224
+ ax.set_ylim(y_lim)
225
+ rect = patches.FancyBboxPatch(
226
+ (-0.5, -1.325),
227
+ 1,
228
+ 0.65,
229
+ boxstyle="round, pad=0.05, rounding_size=0",
230
+ linewidth=2,
231
+ edgecolor="black",
232
+ facecolor="none",
233
+ linestyle="-",
234
+ )
235
+ ax.add_patch(rect)
236
+ # ax.text(1.85, 0.65, r"$\mathbf{U}_{\mathbf{x}}$", fontsize=fs)
237
+
238
+ if args[0]: # do_m
239
+ fig.figimage(HAMMER, 0.0075 * fig.bbox.xmax, 0.395 * fig.bbox.ymax, zorder=10)
240
+ if args[1]: # do_s
241
+ fig.figimage(HAMMER, 0.72 * fig.bbox.xmax, 0.395 * fig.bbox.ymax, zorder=11)
242
+ if args[2]: # do_a
243
+ fig.figimage(HAMMER, 0.363 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=12)
244
+ if args[3]: # do_b
245
+ fig.figimage(HAMMER, 0.72 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=13)
246
+ if args[4]: # do_v
247
+ fig.figimage(HAMMER, 0.0075 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=14)
248
+ else: # b -> v
249
+ a3 = patches.FancyArrowPatch(
250
+ (0.86, 1.21),
251
+ (-0.86, 1.21),
252
+ connectionstyle="arc3,rad=.3",
253
+ linewidth=2,
254
+ arrowstyle="simple, head_width=10, head_length=10",
255
+ color="k",
256
+ )
257
+ ax.add_patch(a3)
258
+ # print(ax.get_xlim())
259
+ # print(ax.get_ylim())
260
+ fig.tight_layout()
261
+ fig.canvas.draw()
262
+ return np.array(fig.canvas.renderer.buffer_rgba())
263
+
264
+
265
+ def chest_graph(*args):
266
+ x, a, d, r, s = r"$\mathbf{x}$", r"$a$", r"$d$", r"$r$", r"$s$"
267
+ ua, ud, ur, us = (
268
+ r"$\mathbf{U}_a$",
269
+ r"$\mathbf{U}_d$",
270
+ r"$\mathbf{U}_r$",
271
+ r"$\mathbf{U}_s$",
272
+ )
273
+ zx, ex = r"$\mathbf{z}_{1:L}$", r"$\boldsymbol{\epsilon}$"
274
+
275
+ G = nx.DiGraph()
276
+ G.add_edge(ua, a)
277
+ G.add_edge(ud, d)
278
+ G.add_edge(ur, r)
279
+ G.add_edge(us, s)
280
+ G.add_edge(a, d)
281
+ G.add_edge(d, x)
282
+ G.add_edge(r, x)
283
+ G.add_edge(s, x)
284
+ G.add_edge(ex, x)
285
+ G.add_edge(zx, x)
286
+ G.add_edge(a, x)
287
+
288
+ pos = {
289
+ x: (0, 0),
290
+ a: (-1, 1),
291
+ d: (0, 1),
292
+ r: (1, 1),
293
+ s: (1, 0),
294
+ ua: (-1, 2),
295
+ ud: (0, 2),
296
+ ur: (1, 2),
297
+ us: (1, -1),
298
+ zx: (-0.25, -1),
299
+ ex: (0.25, -1),
300
+ }
301
+
302
+ node_c = {}
303
+ for node in G:
304
+ node_c[node] = "lightgrey" if node in [x, a, d, r, s] else "white"
305
+
306
+ edge_c = {e: "black" for e in G.edges}
307
+ node_line_c = {k: "black" for k, _ in node_c.items()}
308
+
309
+ if args[0]: # do_r
310
+ # G.remove_edge(ur, r)
311
+ edge_c[(ur, r)] = "lightgrey"
312
+ node_line_c[r] = "red"
313
+ if args[1]: # do_s
314
+ # G.remove_edges_from([(us, s)])
315
+ edge_c[(us, s)] = "lightgrey"
316
+ node_line_c[s] = "red"
317
+ if args[2]: # do_f (do_d)
318
+ # G.remove_edges_from([(ud, d), (a, d)])
319
+ edge_c[(ud, d)] = "lightgrey"
320
+ edge_c[(a, d)] = "lightgrey"
321
+ node_line_c[d] = "red"
322
+ if args[3]: # do_a
323
+ # G.remove_edge(ua, a)
324
+ edge_c[(ua, a)] = "lightgrey"
325
+ node_line_c[a] = "red"
326
+
327
+ fs = 30
328
+ options = {
329
+ "font_size": fs,
330
+ "node_size": 3000,
331
+ "node_color": list(node_c.values()),
332
+ "edgecolors": list(node_line_c.values()),
333
+ "edge_color": list(edge_c.values()),
334
+ "linewidths": 2,
335
+ "width": 2,
336
+ }
337
+ plt.close("all")
338
+ fig, ax = plt.subplots(1, 1, figsize=(5, 5)) # , constrained_layout=True)
339
+ # fig.patch.set_visible(False)
340
+ ax.margins(x=0.1, y=0.08, tight=False)
341
+ ax.axis("off")
342
+ nx.draw_networkx(G, pos, **options, arrowsize=25, arrowstyle="-|>", ax=ax)
343
+ # need to reuse x, y limits so that the graphs plot the same way before and after removing edges
344
+ x_lim = (-1.32, 1.32)
345
+ y_lim = (-1.414, 2.414)
346
+ ax.set_xlim(x_lim)
347
+ ax.set_ylim(y_lim)
348
+ rect = patches.FancyBboxPatch(
349
+ (-0.5, -1.325),
350
+ 1,
351
+ 0.65,
352
+ boxstyle="round, pad=0.05, rounding_size=0",
353
+ linewidth=2,
354
+ edgecolor="black",
355
+ facecolor="none",
356
+ linestyle="-",
357
+ )
358
+ ax.add_patch(rect)
359
+ ax.text(-0.9, -1.075, r"$\mathbf{U}_{\mathbf{x}}$", fontsize=fs)
360
+
361
+ if args[0]: # do_r
362
+ fig.figimage(HAMMER, 0.72 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=10)
363
+ if args[1]: # do_s
364
+ fig.figimage(HAMMER, 0.72 * fig.bbox.xmax, 0.395 * fig.bbox.ymax, zorder=11)
365
+ if args[2]: # do_f
366
+ fig.figimage(HAMMER, 0.363 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=12)
367
+ if args[3]: # do_a
368
+ fig.figimage(HAMMER, 0.0075 * fig.bbox.xmax, 0.64 * fig.bbox.ymax, zorder=13)
369
+
370
+ fig.tight_layout()
371
+ fig.canvas.draw()
372
+ return np.array(fig.canvas.renderer.buffer_rgba())
373
+
374
+
375
+ def vae_preprocess(args, pa):
376
+ if "ukbb" in args.hps:
377
+ # preprocessing ukbb parents for the vae which was originally trained using
378
+ # log standardized parents. The pgm was trained using [-1,1] normalization
379
+ # first undo [-1,1] parent preprocessing back to original range
380
+ for k, v in pa.items():
381
+ if k != "mri_seq" and k != "sex":
382
+ pa[k] = (v + 1) / 2 # [-1,1] -> [0,1]
383
+ _max, _min = get_attr_max_min(k)
384
+ pa[k] = pa[k] * (_max - _min) + _min
385
+ # log_standardize parents for vae input
386
+ for k, v in pa.items():
387
+ logpa_k = torch.log(v.clamp(min=1e-12))
388
+ if k == "age":
389
+ pa[k] = (logpa_k - 4.112339973449707) / 0.11769197136163712
390
+ elif k == "brain_volume":
391
+ pa[k] = (logpa_k - 13.965583801269531) / 0.09537758678197861
392
+ elif k == "ventricle_volume":
393
+ pa[k] = (logpa_k - 10.345998764038086) / 0.43127763271331787
394
+ # concatenate parents expand to input res for conditioning the vae
395
+ pa = torch.cat(
396
+ [pa[k] if len(pa[k].shape) > 1 else pa[k][..., None] for k in args.parents_x],
397
+ dim=1,
398
+ )
399
+ pa = (
400
+ pa[..., None, None].repeat(1, 1, *(args.input_res,) * 2).to(args.device).float()
401
+ )
402
+ return pa
403
+
404
+
405
+ def preprocess_brain(args, obs):
406
+ obs["x"] = (obs["x"][None, ...].float().to(args.device) - 127.5) / 127.5 # [-1,1]
407
+ # for all other variables except x
408
+ for k in [k for k in obs.keys() if k != "x"]:
409
+ obs[k] = obs[k].float().to(args.device).view(1, 1)
410
+ if k in ["age", "brain_volume", "ventricle_volume"]:
411
+ k_max, k_min = get_attr_max_min(k)
412
+ obs[k] = (obs[k] - k_min) / (k_max - k_min) # [0,1]
413
+ obs[k] = 2 * obs[k] - 1 # [-1,1]
414
+ return obs
415
+
416
+
417
+ def get_fig_arr(x, width=4, height=4, dpi=144, cmap="Greys_r", norm=None):
418
+ fig = plt.figure(figsize=(width, height), dpi=dpi)
419
+ ax = plt.axes([0, 0, 1, 1], frameon=False)
420
+ if cmap == "Greys_r":
421
+ ax.imshow(x, cmap=cmap, vmin=0, vmax=255)
422
+ else:
423
+ ax.imshow(x, cmap=cmap, norm=norm)
424
+ ax.axis("off")
425
+ fig.canvas.draw()
426
+ return np.array(fig.canvas.renderer.buffer_rgba())
427
+
428
+
429
+ def normalize(x, x_min=None, x_max=None, zero_one=False):
430
+ if x_min is None:
431
+ x_min = x.min()
432
+ if x_max is None:
433
+ x_max = x.max()
434
+ x = (x - x_min) / (x_max - x_min) # [0,1]
435
+ 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 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:eaab2e017c753c90183239bf6e149a0cebe83155a1749eab95305f068bdb552f
3
+ size 134
checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:114b8f0f1677605417c99dc7f091ece828e89bc75ee9550e9d0f8048aa43d587
3
+ size 134
checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:109dda6b72d38390aabd6e590a27bd6eeb2eb6c92c3ecd753a2fe7c014ac319d
3
+ size 132
checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint_current.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e686cff0a6a6656c248acdc0873f3e5711665d1420217d89ae550b3bba9f8d59
3
+ size 132
checkpoints/m_b_v_s/sup_pgm/checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8f1a873caf6a96f607034e7cbc8d8abaab5971d1297e51160b07b4ffe39b2c66
3
+ size 132
checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0eb4df03d7573187ee962297735e626af6104769d50f2d5211656cc81276a79b
3
+ size 134
checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f3122e11cce58d2e2cafcdb952bb8dbca57e65298aaced86eb700131e2cff632
3
+ size 133
checkpoints/t_i_d/sup_pgm/checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8cb53124e1a91a8f4254d4cb3eaa1a0d2623b79eb8ec453207d2e6bc65ebcfb6
3
+ size 130
data/mimic_subset/0.jpg ADDED
data/mimic_subset/1.jpg ADDED
data/mimic_subset/10.jpg ADDED
data/mimic_subset/11.jpg ADDED
data/mimic_subset/12.jpg ADDED
data/mimic_subset/13.jpg ADDED
data/mimic_subset/14.jpg ADDED
data/mimic_subset/15.jpg ADDED
data/mimic_subset/16.jpg ADDED
data/mimic_subset/17.jpg ADDED
data/mimic_subset/18.jpg ADDED
data/mimic_subset/19.jpg ADDED
data/mimic_subset/2.jpg ADDED
data/mimic_subset/20.jpg ADDED
data/mimic_subset/21.jpg ADDED
data/mimic_subset/22.jpg ADDED
data/mimic_subset/23.jpg ADDED
data/mimic_subset/24.jpg ADDED
data/mimic_subset/25.jpg ADDED
data/mimic_subset/26.jpg ADDED
data/mimic_subset/27.jpg ADDED
data/mimic_subset/28.jpg ADDED
data/mimic_subset/29.jpg ADDED
data/mimic_subset/3.jpg ADDED
data/mimic_subset/30.jpg ADDED
data/mimic_subset/31.jpg ADDED
data/mimic_subset/32.jpg ADDED
data/mimic_subset/33.jpg ADDED
data/mimic_subset/34.jpg ADDED
data/mimic_subset/35.jpg ADDED
data/mimic_subset/36.jpg ADDED
data/mimic_subset/37.jpg ADDED
data/mimic_subset/38.jpg ADDED
data/mimic_subset/39.jpg ADDED
data/mimic_subset/4.jpg ADDED
data/mimic_subset/40.jpg ADDED
data/mimic_subset/41.jpg ADDED
data/mimic_subset/42.jpg ADDED
data/mimic_subset/43.jpg ADDED