geomstats · ninamiolane · Dec 17, 2021 · Sep 7, 2021 · Sep 7, 2021 · Sep 7, 2021
diff --git a/geomstats/_backend/__init__.py b/geomstats/_backend/__init__.py
@@ -89,6 +89,7 @@
         "polygamma",
         "power",
         "prod",
+        "ravel_tril_indices",
         "real",
         "repeat",
         "reshape",
@@ -117,7 +118,9 @@
         "triu",
         "tril_indices",
         "triu_indices",
+        "tril_to_vec",
         "triu_to_vec",
+        "vec_to_diag",
         "unique",
         "vectorize",
         "vstack",

diff --git a/geomstats/_backend/autograd/__init__.py b/geomstats/_backend/autograd/__init__.py
@@ -92,8 +92,8 @@
     trace,
     transpose,
     tril,
-    triu,
     tril_indices,
+    triu,
     triu_indices,
     uint8,
     unique,
@@ -385,12 +385,26 @@ def array_from_sparse(indices, data, target_shape):
     return array(coo_matrix((data, list(zip(*indices))), target_shape).todense())
 
 
+def tril_to_vec(x, k=0):
+    """ """
+    n = x.shape[-1]
+    rows, cols = tril_indices(n, k=k)
+    return x[..., rows, cols]
+
+
 def triu_to_vec(x, k=0):
+    """ """
     n = x.shape[-1]
     rows, cols = triu_indices(n, k=k)
     return x[..., rows, cols]
 
 
+def vec_to_diag(vec):
+    """Convert vector to diagonal matrix."""
+    d = vec.shape[-1]
+    return np.squeeze(vec[..., None, :] * np.eye(d)[None, :, :])
+
+
 def mat_from_diag_triu_tril(diag, tri_upp, tri_low):
     """Build matrix from given components.
 
@@ -415,3 +429,12 @@ def mat_from_diag_triu_tril(diag, tri_upp, tri_low):
     mat[..., j, k] = tri_upp
     mat[..., k, j] = tri_low
     return mat
+
+
+def ravel_tril_indices(n, k=0, m=None):
+    if m is None:
+        size = (n, n)
+    else:
+        size = (n, m)
+    idxs = np.tril_indices(n, k, m)
+    return np.ravel_multi_index(idxs, size)
diff --git a/geomstats/_backend/numpy/__init__.py b/geomstats/_backend/numpy/__init__.py
@@ -92,8 +92,8 @@
     trace,
     transpose,
     tril,
-    triu,
     tril_indices,
+    triu,
     triu_indices,
     uint8,
     unique,
@@ -385,6 +385,18 @@ def array_from_sparse(indices, data, target_shape):
     return array(coo_matrix((data, list(zip(*indices))), target_shape).todense())
 
 
+def vec_to_diag(vec):
+    """Convert vector to diagonal matrix."""
+    d = vec.shape[-1]
+    return np.squeeze(vec[..., None, :] * np.eye(d)[None, :, :])
+
+
+def tril_to_vec(x, k=0):
+    n = x.shape[-1]
+    rows, cols = tril_indices(n, k=k)
+    return x[..., rows, cols]
+
+
 def triu_to_vec(x, k=0):
     n = x.shape[-1]
     rows, cols = triu_indices(n, k=k)
@@ -415,3 +427,12 @@ def mat_from_diag_triu_tril(diag, tri_upp, tri_low):
     mat[..., j, k] = tri_upp
     mat[..., k, j] = tri_low
     return mat
+
+
+def ravel_tril_indices(n, k=0, m=None):
+    if m is None:
+        size = (n, n)
+    else:
+        size = (n, m)
+    idxs = np.tril_indices(n, k, m)
+    return np.ravel_multi_index(idxs, size)
diff --git a/geomstats/_backend/pytorch/__init__.py b/geomstats/_backend/pytorch/__init__.py
@@ -12,6 +12,9 @@
 from torch import (
     ceil,
     clip,
+    complex32,
+    complex64,
+    complex128,
     cos,
     cosh,
     cross,
@@ -24,23 +27,10 @@
     flatten,
     float32,
     float64,
-    complex32,
-    complex64,
-    complex128,
     floor,
 )
 from torch import fmod as mod
-from torch import (
-    greater,
-    hstack,
-    imag,
-    int32,
-    int64,
-    isnan,
-    less,
-    log,
-    logical_or,
-)
+from torch import greater, hstack, imag, int32, int64, isnan, less, log, logical_or
 from torch import max as amax
 from torch import mean, meshgrid
 from torch import min as amin
@@ -57,8 +47,6 @@
     std,
     tan,
     tanh,
-    tril,
-    triu,
     uint8,
     unique,
     vstack,
@@ -514,6 +502,14 @@ def diag_indices(*args, **kwargs):
     return tuple(map(torch.from_numpy, _np.diag_indices(*args, **kwargs)))
 
 
+def tril(mat, k=0):
+    return torch.tril(mat, diagonal=k)
+
+
+def triu(mat, k=0):
+    return torch.triu(mat, diagonal=k)
+
+
 def tril_indices(n, k=0, m=None):
     if m is None:
         m = n
@@ -780,6 +776,16 @@ def vectorize(x, pyfunc, multiple_args=False, **kwargs):
     return stack(list(map(pyfunc, x)))
 
 
+def vec_to_diag(vec):
+    return torch.diag_embed(vec, offset=0)
+
+
+def tril_to_vec(x, k=0):
+    n = x.shape[-1]
+    rows, cols = tril_indices(n, k=k)
+    return x[..., rows, cols]
+
+
 def triu_to_vec(x, k=0):
     n = x.shape[-1]
     rows, cols = triu_indices(n, k=k)
@@ -812,6 +818,15 @@ def mat_from_diag_triu_tril(diag, tri_upp, tri_low):
     return mat
 
 
-def sort(a, axis=- 1):
+def ravel_tril_indices(n, k=0, m=None):
+    if m is None:
+        size = (n, n)
+    else:
+        size = (n, m)
+    idxs = _np.tril_indices(n, k, m)
+    return torch.from_numpy(_np.ravel_multi_index(idxs, size))
+
+
+def sort(a, axis=-1):
     sorted_a, _ = torch.sort(a, dim=axis)
     return sorted_a
diff --git a/geomstats/_backend/pytorch/linalg.py b/geomstats/_backend/pytorch/linalg.py
@@ -42,6 +42,7 @@ def backward(ctx, grad):
         return Logm._logm(backward_tensor).to(tensor.dtype)[..., :n, n:]
 
 
+cholesky = torch.linalg.cholesky
 eig = torch.linalg.eig
 eigh = torch.linalg.eigh
 eigvalsh = torch.linalg.eigvalsh
@@ -53,10 +54,6 @@ def backward(ctx, grad):
 logm = Logm.apply
 
 
-def cholesky(a):
-    return torch.cholesky(a, upper=False)
-
-
 def sqrtm(x):
     np_sqrtm = np.vectorize(scipy.linalg.sqrtm, signature="(n,m)->(n,m)")(x)
     return torch.as_tensor(np_sqrtm, dtype=x.dtype)

diff --git a/geomstats/_backend/tensorflow/__init__.py b/geomstats/_backend/tensorflow/__init__.py
@@ -763,16 +763,26 @@ def cumprod(a, axis=None):
     return tf.math.cumprod(a, axis=axis)
 
 
-def tril(m, k=0):
-    if k != 0:
-        raise NotImplementedError("Only k=0 supported so far")
-    return tf.linalg.band_part(m, -1, 0)
-
-
-def triu(m, k=0):
-    if k != 0:
-        raise NotImplementedError("Only k=0 supported so far")
-    return tf.linalg.band_part(m, 0, -1)
+# (sait) there is tf.experimental.tril (we can use it once it moves to stable)
+def tril(mat, k=0):
+    if k not in (0, -1):
+        raise NotImplementedError("Only k=0 and k=-1 supported so far")
+    tril = tf.linalg.band_part(mat, -1, 0)
+    if k == 0:
+        return tril
+    zero_diag = tf.zeros(mat.shape[:-1])
+    return tf.linalg.set_diag(tril, zero_diag)
+
+
+# TODO(sait) use tf.experimental.triu once it becomes stable.
+def triu(mat, k=0):
+    if k not in (0, 1):
+        raise NotImplementedError("Only k=0 and k=1 supported so far")
+    triu = tf.linalg.band_part(mat, 0, -1)
+    if k == 0:
+        return triu
+    zero_diag = tf.zeros(mat.shape[:-1])
+    return tf.linalg.set_diag(triu, zero_diag)
 
 
 def diag_indices(*args, **kwargs):
@@ -802,6 +812,19 @@ def where(condition, x=None, y=None):
     return tf.where(condition, x, y)
 
 
+def tril_to_vec(x, k=0):
+    n = x.shape[-1]
+    axis = 1 if x.ndim == 3 else 0
+    mask = tf.ones((n, n))
+    mask_a = tf.linalg.band_part(mask, -1, 0)
+    if k < 0:
+        mask_b = tf.linalg.band_part(mask, -k - 1, 0)
+    else:
+        mask_b = tf.zeros_like(mask_a)
+    mask = tf.cast(mask_a - mask_b, dtype=tf.bool)
+    return tf.boolean_mask(x, mask, axis=axis)
+
+
 def triu_to_vec(x, k=0):
     n = x.shape[-1]
     axis = 1 if x.ndim == 3 else 0
@@ -823,8 +846,13 @@ def tile(x, multiples):
     return tf.tile(x_reshape, multiples)
 
 
+def vec_to_diag(vec):
  
  
+    """Convert vec to diagonal matrix"""
+    return tf.linalg.diag(vec)
+
+
 def vec_to_triu(vec):
-    """Take vec and forms strictly upper traingular matrix.
+    """Take vec and forms strictly upper triangular matrix.
 
     Parameters
     ---------
@@ -896,3 +924,17 @@ def mat_from_diag_triu_tril(diag, tri_upp, tri_low):
     triu_tril_mat = triu_mat + tril_mat
     mat = tf.linalg.set_diag(triu_tril_mat, diag)
     return mat
+
+
+def _ravel_multi_index(multi_index, shape):
+    strides = tf.math.cumprod(shape, exclusive=True, reverse=True)
+    return tf.reduce_sum(multi_index * tf.expand_dims(strides, 1), axis=0)
+
+
+def ravel_tril_indices(n, k=0, m=None):
+    if m is None:
+        size = (n, n)
+    else:
+        size = (n, m)
+    idxs = tril_indices(n, k, m)
+    return _ravel_multi_index(idxs, size)