add test & pass test

Author: guillaumehu

src/ott/geometry/geodesic.py

@@ -44,8 +44,8 @@ def compute_sparse_laplacian(G: Array_g, normalize:bool = False) -> Array_g:
     data = jnp.where(data > 0., 1. / jnp.sqrt(data), 0.)
   degree = jesp.BCOO((data, jnp.c_[ixs, ixs]), shape=(n, n))
   if normalize:
-    inv_sqrt_deg = degree.data # TODO: is MM still faster here ?
-    laplacian = inv_sqrt_deg[:, None] * G * inv_sqrt_deg[None, :]
+    id = jesp.BCOO((jnp.ones(n), jnp.c_[ixs, ixs]), shape=(n, n))
+    laplacian = id - degree @ G @ degree
   else:
     laplacian = degree - G
   return laplacian
@@ -184,10 +184,14 @@ def apply_kernel(
   def kernel_matrix(self) -> jnp.ndarray:  # noqa: D102
     n, _ = self.shape
     kernel = self.apply_kernel(jnp.eye(n))
-    return jax.lax.cond(
-        jnp.allclose(kernel, kernel.T, atol=1e-8, rtol=1e-8), lambda x: x,
-        lambda x: (x + x.T) / 2.0, kernel
-    )
+    if isinstance(kernel, jesp.BCOO):
+      # we symmetrize sparse kernel by default
+      return (kernel + kernel.T) * 0.5
+    else:
+      return jax.lax.cond(
+          jnp.allclose(kernel, kernel.T, atol=1e-8, rtol=1e-8), lambda x: x,
+          lambda x: (x + x.T) / 2.0, kernel
+      )
 
   @property
   def cost_matrix(self) -> jnp.ndarray:  # noqa: D102
@@ -266,31 +270,33 @@ def compute_largest_eigenvalue(
   )
   return eigvals[0]
 
-
 def expm_multiply(
     L: jnp.ndarray, X: jnp.ndarray, coeff: jnp.ndarray, eigval: float
 ) -> jnp.ndarray:
-
   # move to sparse matrix
-  L = jax.experimental.sparse.BCOO.fromdense(L, nse=10)
-  X = jax.experimental.sparse.BCOO.fromdense(X, nse=1)
+  is_sparse = isinstance(L, jesp.BCOO)
+  if is_sparse and not isinstance(X, jesp.BCOO):
+    X = jax.experimental.sparse.BCOO.fromdense(X, nse=X.shape[0])
+  
   def body(carry, c):
     T0, T1, Y = carry
-    T1, Y = T1.sort_indices(), Y.sort_indices()
     T2 = (2.0 / eigval) * L @ T1 - 2.0 * T1 - T0
     Y = Y + c * T2
     return (T1, T2, Y), None
 
   T0 = X
-  T0.unique_indices = False # FIXME: rm change attribute
   Y = 0.5 * coeff[0] * T0
   T1 = (1.0 / eigval) * L @ X - T0
   Y = Y + coeff[1] * T1
 
+
   initial_state = (T0, T1, Y)
-  (_, _, Y), _ = jax.lax.scan(body, initial_state, coeff[2:]) # FIXME scan does not work
-                                                              # because metadata unique_indices is changing
-                                                              # for 1st position of carry.
+  if not is_sparse:
+    (_, _, Y), _ = jax.lax.scan(body, initial_state, coeff[2:]) 
+  else:
+    # NOTE: the scan not working for this type of scan
+    for c in coeff[2:]:
+      (T0, T1, Y), _ = body((T0, T1, Y), c)
   return Y
 
 

tests/geometry/geodesic_test.py

@@ -25,7 +25,6 @@
 from ott.problems.linear import linear_problem
 from ott.solvers.linear import sinkhorn
 
-
 def random_graph(
     n: int,
     p: float = 0.3,
@@ -267,24 +266,21 @@ def test_heat_approx(self, normalize: bool, t: float, order: int):
     approx = geom.apply_kernel(jnp.eye(G.shape[0]))
     np.testing.assert_allclose(exact, approx, rtol=1e-1, atol=1e-1)
 
-  def test_sparse_geodesic(self, rng: jax.Array):
-    G = random_graph(20, p=0.5)
-    # G = sparse.BCOO.fromdense(G)
-    geom = geodesic.Geodesic.from_graph(G, t=1.0)
+  @pytest.mark.parametrize("normalize", [True, False])
+  def test_sparse_geodesic(self, normalize):
+    n = 20
+    G = random_graph(n, p=0.5)
+    G_sparse = sparse.BCOO.fromdense(G)
+    geom = geodesic.Geodesic.from_graph(G_sparse, t=5.0, order=10, normalize=normalize)
     kernel_matrix = geom.kernel_matrix
 
-    v = jax.random.normal(rng, (n,))
-    v = v / jnp.linalg.norm(v)
-
-    v = jax.device_put(v)
-    v = sparse.COO.from_numpy(v)
+    gh_heat_kernel = exact_heat_kernel(G, normalize=normalize, t=5.0)
 
-    w = geom.apply_kernel(v, axis=0)
-    w = w.todense()
-    w = jnp.asarray(w)
+    assert isinstance(kernel_matrix, sparse.BCOO)
 
-    np.testing.assert_allclose(w, geom.kernel_matrix @ v.todense())
+    kernel_matrix = kernel_matrix.todense()
+    np.testing.assert_allclose(kernel_matrix, gh_heat_kernel, rtol=1e-1, atol=1e-1)
 
 
 if __name__ == "__main__":
-  pytest.main([__file__])
+  TestGeodesic().test_sparse_geodesic(jax.random.PRNGKey(0))