Update example to use the right data and to also show-case the discretized manifold

Author: sorenhauberg

examples/local_pca_mnist.py

@@ -9,14 +9,19 @@
 
 def get_subset_mnist(n: int = 1000):
     dataset = MNIST(root="", download=True)
-    N = dataset.data.shape[0]
+    data = dataset.data[dataset.targets==1]
+    N = data.shape[0]
     idx = np.random.choice(np.arange(N), size=n)
-    return dataset.data[idx], dataset.targets[idx]
+    return data[idx]
 
 
 # Read data
-data, targets = get_subset_mnist(n=1000)
-data = data.reshape(data.shape[0], -1)
+data = get_subset_mnist(n=1000)
+data = data.reshape(data.shape[0], -1).to(torch.float)
+cov = torch.cov(data.t())
+values, vectors = torch.linalg.eigh(cov)
+proj = vectors[:, -2:] / values[-2:].sqrt().unsqueeze(0)
+data = data @ proj
 N, D = data.shape
 
 # Parameters for metric
@@ -27,13 +32,14 @@ def get_subset_mnist(n: int = 1000):
 M = stochman.manifold.LocalVarMetric(data=data, sigma=sigma, rho=rho)
 
 # Plot metric and data
-ran = torch.linspace(-2.5, 2.5, 100)
-X, Y = torch.meshgrid([ran, ran])
+plt.figure()
+ran = torch.linspace(-3.0, 3.0, 100)
+X, Y = torch.meshgrid([ran, ran], indexing='ij')
 XY = torch.stack((X.flatten(), Y.flatten()), dim=1)  # 10000x2
 gridM = M.metric(XY)  # 10000x2
-Mim = gridM.sum(dim=1).reshape((100, 100)).detach().numpy().T
+Mim = gridM.sum(dim=1).reshape((100, 100)).detach().t()
 plt.imshow(Mim, extent=(ran[0], ran[-1], ran[0], ran[-1]), origin="lower")
-plt.plot(data[:, 0].numpy(), data[:, 1].numpy(), "w.", markersize=1)
+plt.plot(data[:, 0], data[:, 1], "w.", markersize=1)
 
 # Compute geodesics in parallel
 p0 = data[torch.randint(high=N, size=[10], dtype=torch.long)]  # 10xD
@@ -42,16 +48,30 @@ def get_subset_mnist(n: int = 1000):
 C.plot()
 C.constant_speed(M)
 C.plot()
-plt.show()
 
-# Compute shooting geodesic as a sanity check
-p0 = data[0]  # 1xD
-p1 = data[1]  # 1xD
-C, success = M.connecting_geodesic(p0, p1)
-C.plot()
+# Construct discretized manifold
+DM = stochman.discretized_manifold.DiscretizedManifold()
+DM.fit(M, [ran, ran], batch_size=100)
+
+# Compute discretized geodesics
+plt.figure()
+ran2 = torch.linspace(-3.0, 3.0, 133)
+X2, Y2 = torch.meshgrid([ran2, ran2], indexing='ij')
+XY2 = torch.stack((X2.flatten(), Y2.flatten()), dim=1)  # 10000x2
+DMim = DM.metric(XY2).log().sum(dim=1).view(133, 133).t()
+plt.imshow(DMim, extent=(ran[0], ran[-1], ran[0], ran[-1]), origin="lower")
+plt.plot(data[:, 0], data[:, 1], "w.", markersize=1)
+for k in range(10):
+    p0 = data[torch.randint(high=N, size=[1], dtype=torch.long)]  # 1xD
+    p1 = data[torch.randint(high=N, size=[1], dtype=torch.long)]  # 1xD
+    C = DM.connecting_geodesic(p0, p1)
+    C.plot()
 
 # p = C.begin
 # with torch.no_grad():
 #     v = C.deriv(torch.zeros(1))
 #     c, dc = shooting_geodesic(M, p, v, t=torch.linspace(0, 1, 100))
 # plt.plot(c[:,0,0], c[:,1, 0], 'o')
+
+
+plt.show()