more hypothesis testing on MMD

flatironinstitute · Sep 6, 2023 · ac9cdf5 · ac9cdf5
1 parent c2a520c
commit ac9cdf5
Show file tree

Hide file tree

Showing 34 changed files with 6,058 additions and 60 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,4 +1,5 @@
 # Project specific files
+ASPIRE-Python/
 messing_around/
 tests/test_simulator.ipynb
 *.estimator

diff --git a/6wxb_files/6wxb_MMD.ipynb b/6wxb_files/6wxb_MMD.ipynb
diff --git a/Lukes_folder/6xwb_images.ipynb → 6wxb_files/6wxb_images.ipynb b/Lukes_folder/6xwb_images.ipynb → 6wxb_files/6wxb_images.ipynb
diff --git a/6wxb_files/MMD.py b/6wxb_files/MMD.py
@@ -0,0 +1,19 @@
+import torch
+
+def MMD_gaussian(X, Y, bandwidth=None):
+
+    # Compute all pairwise distances 
+    Xdists = torch.cdist(X, X)**2
+    Ydists = torch.cdist(Y, Y)**2
+    XYdists = torch.cdist(X, Y)**2
+
+    # Use heuristic for bandwidth if none provided
+    if bandwidth == None:
+        bandwidth = torch.sqrt(torch.median(XYdists)*2)
+
+    # Compute all kernel sums
+    Xterm = torch.exp(-Xdists[None, ...]/bandwidth**2).mean()
+    Yterm = torch.exp(-Ydists[None, ...]/bandwidth**2).mean()
+    XYterm = torch.exp(-XYdists[None, ...]/bandwidth**2).mean()
+    return Xterm + Yterm - 2*XYterm
+
diff --git a/6wxb_files/image_params_mixed_training.json b/6wxb_files/image_params_mixed_training.json
@@ -0,0 +1,11 @@
+{
+    "N_PIXELS": 128,
+    "PIXEL_SIZE": 2.06,
+    "SIGMA": [0.5, 5.0],
+    "MODEL_FILE": "../data/protein_models/6wxb_mixed_models.npy",
+    "SHIFT": 40,
+    "DEFOCUS": [0.5, 5.0],
+    "SNR": [0.01,1.0],
+    "AMP": 0.1,
+    "B_FACTOR": [1.0, 100.0]
+}
diff --git a/6wxb_files/resnet18_encoder.json b/6wxb_files/resnet18_encoder.json
@@ -0,0 +1,12 @@
+{"EMBEDDING": "RESNET18",
+"OUT_DIM": 256,
+"NUM_TRANSFORM": 5,
+"NUM_HIDDEN_FLOW": 10,
+"HIDDEN_DIM_FLOW": 256,
+"MODEL": "NSF",
+"LEARNING_RATE": 0.0003,
+"CLIP_GRADIENT": 5.0,
+"THETA_SHIFT": 50,
+"THETA_SCALE": 50,
+"BATCH_SIZE": 256
+}
diff --git a/6wxb_files/resnet18_fft_encoder.json b/6wxb_files/resnet18_fft_encoder.json
@@ -0,0 +1,12 @@
+{"EMBEDDING": "RESNET18_FFT_FILTER",
+"OUT_DIM": 256,
+"NUM_TRANSFORM": 5,
+"NUM_HIDDEN_FLOW": 10,
+"HIDDEN_DIM_FLOW": 256,
+"MODEL": "NSF",
+"LEARNING_RATE": 0.0003,
+"CLIP_GRADIENT": 5.0,
+"THETA_SHIFT": 50,
+"THETA_SCALE": 50,
+"BATCH_SIZE": 256
+}
diff --git a/Lukes_folder/6wxb/image_params_mixed_training.json b/Lukes_folder/6wxb/image_params_mixed_training.json
@@ -3,10 +3,7 @@
     "PIXEL_SIZE": 2.06,
     "SIGMA": [0.5, 5.0],
     "MODEL_FILE": "../data/protein_models/6wxb_mixed_models.npy",
-    "ROTATIONS": true,
     "SHIFT": true,
-    "CTF": true,
-    "NOISE": true,
     "DEFOCUS": [0.5, 5.0],
     "SNR": [0.001, 0.1],
     "RADIUS_MASK": 100,