Fix tests

Author: pchlenski

manify/manifolds.py

@@ -234,6 +234,7 @@ def sample(
         z_mean = self.mu0 if z_mean is None else z_mean
         z_mean = torch.Tensor(z_mean).reshape(-1, self.ambient_dim).to(self.device)
         n = z_mean.shape[0]
+
         sigma = torch.stack([torch.eye(self.dim)] * n).to(self.device) if sigma is None else sigma
         sigma = torch.Tensor(sigma).reshape(-1, self.dim, self.dim).to(self.device)
         assert sigma.shape == (
@@ -250,7 +251,7 @@ def sample(
 
         # Sample initial vector from N(0, sigma)
         N = torch.distributions.MultivariateNormal(
-            loc=torch.zeros((n, self.dim), device=self.device), covariance_matrix=sigma
+            loc=torch.zeros((n * n_samples, self.dim), device=self.device), covariance_matrix=sigma
         )
         v = N.sample()
 

tests/test_clustering.py

@@ -9,14 +9,14 @@ def test_riemannian_fuzzy_k_means():
     X, _ = pm.gaussian_mixture(num_points=100)
 
     for optimizer in ["adam", "adan"]:
-        kmeans = RiemannianFuzzyKMeans(manifold=pm, n_clusters=5, random_state=42)
+        kmeans = RiemannianFuzzyKMeans(pm=pm, n_clusters=5, random_state=42)
         kmeans.fit(X)
         preds = kmeans.predict(X)
         assert preds.shape == (100,), f"Predictions should have shape (100,) (optimizer: {optimizer})"
 
         # Also test with X as a numpy array
         X_np = X.numpy()
-        kmeans = RiemannianFuzzyKMeans(manifold=pm, n_clusters=5, random_state=42)
+        kmeans = RiemannianFuzzyKMeans(pm=pm, n_clusters=5, random_state=42)
         kmeans.fit(X_np)
         preds_np = kmeans.predict(X_np)
         assert torch.tensor(preds_np).shape == (100,), f"Predictions should have shape (100,) (optimizer: {optimizer})"
@@ -25,7 +25,7 @@ def test_riemannian_fuzzy_k_means():
         )
 
         # Also do a single manifold
-        kmeans = RiemannianFuzzyKMeans(manifold=pm.P[0], n_clusters=5, optimizer=optimizer, random_state=42)
+        kmeans = RiemannianFuzzyKMeans(pm=pm.P[0], n_clusters=5, optimizer=optimizer, random_state=42)
         X0 = pm.factorize(X)[0]
         kmeans.fit(X0)
         preds = kmeans.predict(X0)

tests/test_manifolds.py

@@ -23,6 +23,17 @@ def _shared_tests(M, X1, X2, is_euclidean):
         assert torch.allclose(ip_11, X1 @ X1.T, atol=1e-5), "Euclidean inner products do not match for X1"
         assert torch.allclose(ip_12, X1 @ X2.T, atol=1e-5), "Euclidean inner products do not match for X1 and X2"
 
+    # Sampling shapes should support a variety of inputs
+    stacked_means = torch.stack([M.mu0] * 5)
+    s1 = M.sample(100)
+    assert s1.shape == (100, M.ambient_dim), "Sampled points should have the correct shape"
+    s2 = M.sample(100, z_mean=M.mu0)
+    assert s2.shape == (100, M.ambient_dim), "Sampled points should have the correct shape"
+    s3 = M.sample(z_mean=stacked_means)
+    assert s3.shape == (5, M.ambient_dim), "Sampled points should have the correct shape"
+    s3 = M.sample(100, z_mean=stacked_means)
+    assert s3.shape == (500, M.ambient_dim), "Sampled points should have the correct shape"
+
     # Dists
     dists_11 = M.dist(X1, X1)
     assert dists_11.shape == (10, 10), "Distance shape mismatch for X1"