Add script to evaluate W1 distance

Author: plainerman

environment.yml

@@ -17,6 +17,7 @@ dependencies:
   - jaxlib=0.4.23
   - flax=0.8.3
   - notebook=7.0.8
+  - POT=0.9.4
   - pip:
     - dmff @ git+https://github.com/deepmodeling/DMFF@v1.0.0
     - rdkit==2023.3.3

eval/evaluate_wasserstein.py

@@ -0,0 +1,39 @@
+import ot
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm import trange
+import os
+
+if __name__ == '__main__':
+    two_way_shooting = np.load('./out/baselines/mueller/paths-two-way-shooting.npy', allow_pickle=True)
+    two_way_shooting = two_way_shooting.squeeze(2)
+    ours = np.load('./out/toy/mueller_single_gaussian/stochastic_paths.npy')
+
+    assert two_way_shooting.shape == ours.shape, f'Shapes do not match: {two_way_shooting.shape} vs {ours.shape}'
+
+    savedir = './out/evaluation/mueller/'
+    os.makedirs(savedir, exist_ok=True)
+
+    wasserstein = []
+    for t in trange(ours.shape[1]):
+        cur_ground_truth = np.array(two_way_shooting[:, t, :], dtype=np.float64)
+        cur_ours = np.array(ours[:, t, :], dtype=np.float64)
+
+        M = ot.dist(cur_ground_truth, cur_ours, metric='euclidean')
+        w1 = ot.emd2([], [], M)
+        wasserstein.append(w1)
+
+    wasserstein = np.array(wasserstein)
+    print('Median Wasserstein:', np.median(wasserstein))
+    print('Mean Wasserstein:', np.mean(wasserstein))
+    print('Std Wasserstein:', np.std(wasserstein))
+    print('Max Wasserstein:', np.max(wasserstein))
+    print('Min Wasserstein:', np.min(wasserstein))
+
+    plt.plot(wasserstein)
+    plt.xlabel(r'$t$')
+    plt.ylabel('Wasserstein W1 Distance')
+    plt.savefig(f'{savedir}/wasserstein.pdf', bbox_inches='tight')
+    plt.clf()
+
+    print()