Merge branch 'main' of github.com:mackelab/labproject into main

configs/conf_dims.yaml

@@ -1,10 +1,11 @@
 exp_log_name: "scaling_dims" # optional but recommended
 data: ["multivariate_normal"]
-experiments: ["ScaleDimSW", "ScaleDimMMD"]
-n: [3000]
+experiments: ["ScaleDimSW", "ScaleDimC2ST", "ScaleDimMMD"]
+n: [1000]
 d: [100]
 distort: ["shift_one", "shift_all"]
 dim_sizes: [5, 10, 15, 20, 25, 30, 40, 60, 80, 100]
+mmd_bandwidth: [10., 10.]
 
 seed: 1404
 runs: 5

labproject/data.py

@@ -288,7 +288,6 @@ def multivariate_normal(n=3000, dims=100, means=None, vars=None, distort=None):
         idx = 0
         shift = torch.zeros(n) + 1
         samples[:, idx] = samples[:, idx] + shift
-    print(f"First 5 rows of dataset distorted: {samples[:5, :5]}")
     return samples
 
 
@@ -527,7 +526,9 @@ def imagenet_validation_embedding(n, d=2048, device="cpu", save_path="data"):
 
 
 @register_dataset("imagenet_conditional_model")
-def imagenet_conditional_model(n, d=2048, label:Optional[int]=None, device="cpu", permute_if_no_label=True, save_path="data"):
+def imagenet_conditional_model(
+    n, d=2048, label: Optional[int] = None, device="cpu", permute_if_no_label=True, save_path="data"
+):
     r"""Get the conditional model embeddings for ImageNet
 
     Args:
@@ -549,9 +550,11 @@ def imagenet_conditional_model(n, d=2048, label:Optional[int]=None, device="cpu"
     if label is not None:
         conditional_embeddings = conditional_embeddings[label]
     else:
-        conditional_embeddings = conditional_embeddings.flatten(0,1)
+        conditional_embeddings = conditional_embeddings.flatten(0, 1)
         if permute_if_no_label:
-            conditional_embeddings = conditional_embeddings[torch.randperm(conditional_embeddings.shape[0])]
+            conditional_embeddings = conditional_embeddings[
+                torch.randperm(conditional_embeddings.shape[0])
+            ]
 
     max_n = conditional_embeddings.shape[0]
 

labproject/experiments.py

@@ -8,6 +8,7 @@
 from labproject.plotting import plot_scaling_metric_dimensionality, plot_scaling_metric_sample_size
 from labproject.metrics.gaussian_squared_wasserstein import gaussian_squared_w2_distance
 import pickle
+import math
 
 
 class Experiment:
@@ -55,7 +56,6 @@ def run_experiment(self, dataset1, dataset2, dataset_size, nb_runs=5, dim_sizes=
             else torch.zeros_like(torch.tensor(dim_sizes))
         )
         final_distances = torch.tensor([torch.mean(d) for d in final_distances])
-        print(f"Final errors: {final_errors}")
         return dim_sizes, final_distances, final_errors
 
     def plot_experiment(