Based backend on latest geomstats + fixed errors arising

Author: FlorianPfaff

pyrecest/_backend/README.md

@@ -1,4 +1,4 @@
-# _Backend Folder
+# Backend Interface
 
 This folder contains code from the Geomstats project, adjusted for pyRecEst by Florian Pfaff. The original version of Geomstats is authored by Nina Miolane et al., and is a Python package geared towards Riemannian Geometry in Machine Learning.
 

pyrecest/_backend/__init__.py

@@ -47,6 +47,7 @@ def get_backend_name():
         "argmin",
         "array",
         "array_from_sparse",
+        "asarray",
         "as_dtype",
         "assignment",
         "assignment_by_sum",
@@ -87,6 +88,7 @@ def get_backend_name():
         "get_default_cdtype",
         "get_slice",
         "greater",
+        "has_autodiff",
         "hsplit",
         "hstack",
         "imag",
@@ -205,6 +207,7 @@ def get_backend_name():
         "jacobian_vec",
         "jacobian_and_hessian",
         "value_and_grad",
+        "value_and_jacobian",
         "value_jacobian_and_hessian",
     ],
     "linalg": [
@@ -218,9 +221,11 @@ def get_backend_name():
         "inv",
         "is_single_matrix_pd",
         "logm",
+        "matrix_power",
         "norm",
         "qr",
         "quadratic_assignment",
+        "polar",
         "solve",
         "solve_sylvester",
         "sqrtm",

pyrecest/_backend/_backend_config.py

@@ -1,8 +1,8 @@
 pytorch_atol = 1e-6
 pytorch_rtol = 1e-5
 
-np_atol = 1e-12
-np_rtol = 1e-6
+np_atol = 1e-8
+np_rtol = 1e-5
 
 jax_atol = 1e-6
 jax_rtol = 1e-5

pyrecest/_backend/_dtype_utils.py

@@ -321,7 +321,7 @@ def _wrapped(x, *args, **kwargs):
                         cmp_dtype = x.dtype
                     else:
                         float_name = dtype_as_str(x.dtype)
-                        cmp_dtype = as_dtype(f"complex{int((float_name[-2:]))*2}")
+                        cmp_dtype = as_dtype(f"complex{int((float_name[-2:])) * 2}")
 
                     if out.dtype != cmp_dtype:
                         return cast(out, cmp_dtype)
@@ -383,7 +383,7 @@ def _np_box_unary_scalar(target=None):
     def _decorator(func):
         @functools.wraps(func)
         def _wrapped(x, *args, **kwargs):
-            if type(x) is float:
+            if isinstance(x, float):
                 return func(x, *args, dtype=_config.DEFAULT_DTYPE, **kwargs)
 
             return func(x, *args, **kwargs)
@@ -407,7 +407,7 @@ def _np_box_binary_scalar(target=None):
     def _decorator(func):
         @functools.wraps(func)
         def _wrapped(x1, x2, *args, **kwargs):
-            if type(x1) is float:
+            if isinstance(x1, float):
                 return func(x1, x2, *args, dtype=_config.DEFAULT_DTYPE, **kwargs)
 
             return func(x1, x2, *args, **kwargs)

pyrecest/_backend/_shared_numpy/__init__.py

@@ -39,7 +39,7 @@
 
 def angle(z, deg=False):
     out = _np.angle(z, deg=deg)
-    if type(z) is float:
+    if isinstance(z, float):
         return cast(out, get_default_dtype())
 
     return out
@@ -63,7 +63,9 @@ def real(x):
 
 def arange(start_or_stop, /, stop=None, step=1, dtype=None, **kwargs):
     if dtype is None and (
-        type(stop) is float or type(step) is float or type(start_or_stop) is float
+        isinstance(stop, float)
+        or isinstance(step, float)
+        or isinstance(start_or_stop, float)
     ):
         dtype = get_default_dtype()
 

pyrecest/_backend/_shared_numpy/_common.py

@@ -70,13 +70,15 @@ def is_array(x):
     return type(x) is _np.ndarray
 
 
-def to_ndarray(x, to_ndim, axis=0):
-    x = _np.array(x)
-    if x.ndim == to_ndim - 1:
+def to_ndarray(x, to_ndim, axis=0, dtype=None):
+    x = _np.asarray(x, dtype=dtype)
+
+    if x.ndim > to_ndim:
+        raise ValueError("The ndim cannot be adapted properly.")
+
+    while x.ndim < to_ndim:
         x = _np.expand_dims(x, axis=axis)
 
-    if x.ndim != 0 and x.ndim < to_ndim:
-        raise ValueError("The ndim was not adapted properly.")
     return x
 
 

pyrecest/_backend/_shared_numpy/linalg.py

@@ -8,14 +8,18 @@
 atol = _common.atol
 
 
+def _transpose(array):
+    axes = list(range(0, array.ndim))
+    axes[-2], axes[-1] = axes[-1], axes[-2]
+    return _np.transpose(array, axes=axes)
+
+
 def _is_symmetric(x, tol=atol):
-    new_x = _to_ndarray(x, to_ndim=3)
-    return (_np.abs(new_x - _np.transpose(new_x, axes=(0, 2, 1))) < tol).all()
+    return (_np.abs(x - _transpose(x)) < tol).all()
 
 
 def _is_hermitian(x, tol=atol):
-    new_x = _to_ndarray(x, to_ndim=3)
-    return (_np.abs(new_x - _np.conj(_np.transpose(new_x, axes=(0, 2, 1)))) < tol).all()
+    return (_np.abs(x - _np.conj(_transpose(x))) < tol).all()
 
 
 _diag_vec = _np.vectorize(_np.diag, signature="(n)->(n,n)")
@@ -26,38 +30,30 @@ def _is_hermitian(x, tol=atol):
 
 
 def logm(x):
-    ndim = x.ndim
-    new_x = _to_ndarray(x, to_ndim=3)
-
-    if _is_symmetric(new_x) and new_x.dtype not in [_np.complex64, _np.complex128]:
-        eigvals, eigvecs = _np.linalg.eigh(new_x)
+    if _is_symmetric(x) and x.dtype not in [_np.complex64, _np.complex128]:
+        eigvals, eigvecs = _np.linalg.eigh(x)
         if (eigvals > 0).all():
             eigvals = _np.log(eigvals)
             eigvals = _diag_vec(eigvals)
-            transp_eigvecs = _np.transpose(eigvecs, axes=(0, 2, 1))
+            transp_eigvecs = _transpose(eigvecs)
             result = _np.matmul(eigvecs, eigvals)
             result = _np.matmul(result, transp_eigvecs)
         else:
-            result = _logm_vec(new_x)
+            result = _logm_vec(x)
     else:
-        result = _logm_vec(new_x)
+        result = _logm_vec(x)
 
-    if ndim == 2:
-        return result[0]
     return result
 
 
 def solve_sylvester(a, b, q, tol=atol):
     if a.shape == b.shape:
-        axes = (0, 2, 1) if a.ndim == 3 else (1, 0)
-        if _np.all(_np.isclose(a, b)) and _np.all(
-            _np.abs(a - _np.transpose(a, axes)) < tol
-        ):
+        if _np.all(_np.isclose(a, b)) and _np.all(_np.abs(a - _transpose(a)) < tol):
             eigvals, eigvecs = _np.linalg.eigh(a)
             if _np.all(eigvals >= tol):
-                tilde_q = _np.transpose(eigvecs, axes) @ q @ eigvecs
+                tilde_q = _transpose(eigvecs) @ q @ eigvecs
                 tilde_x = tilde_q / (eigvals[..., :, None] + eigvals[..., None, :])
-                return eigvecs @ tilde_x @ _np.transpose(eigvecs, axes)
+                return eigvecs @ tilde_x @ _transpose(eigvecs)
 
     return _np.vectorize(
         _scipy.linalg.solve_sylvester, signature="(m,m),(n,n),(m,n)->(m,n)"
@@ -102,4 +98,41 @@ def fractional_matrix_power(A, t):
     if A.ndim == 2:
         return _scipy.linalg.fractional_matrix_power(A, t)
 
-    return _np.stack([_scipy.linalg.fractional_matrix_power(A_, t) for A_ in A])
+    return _np.stack([_scipy.linalg.fractional_matrix_power(A_, t) for A_ in A])
+
+
+def polar(*args, **kwargs):
+    """Polar decomposition of a matrix."""
+    return _np.vectorize(
+        _scipy.linalg.polar, signature="(n,n)->(n,n),(n,n)", excluded=["side"]
+    )(*args, **kwargs)
+
+
+def solve(a, b):
+    """
+    Solve a linear matrix equation, or system of linear scalar equations.
+
+    Computes the "exact" solution, `x`, of the well-determined, i.e., full
+    rank, linear matrix equation `ax = b`.
+
+    Parameters
+    ----------
+    a : array-like, shape=[..., M, M]
+        Coefficient matrix.
+    b : array-like, shape=[..., M]
+        Ordinate or "dependent variable" values".
+
+    Returns
+    -------
+    x : array-like, shape=[..., M]
+        Solution to the system a x = b.
+    """
+    batch_shape = a.shape[:-2]
+    if batch_shape:
+        b = _np.expand_dims(b, axis=-1)
+
+    res = _np.linalg.solve(a, b)
+    if batch_shape:
+        return res[..., 0]
+
+    return res

pyrecest/_backend/autograd/__init__.py

@@ -9,6 +9,7 @@
     any,
     argmax,
     argmin,
+    asarray,
     broadcast_arrays,
     broadcast_to,
     clip,
@@ -61,7 +62,6 @@
     take,
     tile,
     transpose,
-    trapz,
     tril,
     tril_indices,
     triu,
@@ -72,7 +72,12 @@
     where,
     zeros_like,
 )
-from autograd.numpy import trapz as trapezoid
+
+try:
+    from autograd.numpy import trapezoid
+except ImportError:
+    from autograd.numpy import trapz as trapezoid
+
 from autograd.scipy.special import erf, gamma, polygamma  # NOQA
 
 from .._shared_numpy import (
@@ -123,9 +128,11 @@
     vec_to_diag,
     vectorize,
 )
-from . import autodiff  # NOQA
-from . import linalg  # NOQA
-from . import random  # NOQA
+from . import (
+    autodiff,  # NOQA
+    linalg,  # NOQA
+    random,  # NOQA
+)
 from ._common import (
     _box_binary_scalar,
     _box_unary_scalar,
@@ -153,6 +160,16 @@
 empty = _dyn_update_dtype(target=_np.empty)
 
 
+def has_autodiff():
+    """If allows for automatic differentiation.
+
+    Returns
+    -------
+    has_autodiff : bool
+    """
+    return True
+
+
 def imag(x):
     out = _np.imag(x)
     if is_array(x):

pyrecest/_backend/autograd/autodiff.py

@@ -115,7 +115,7 @@ def wrapped_grad_func(i, ans, *args, **kwargs):
             )
         else:
             raise NotImplementedError(
-                "custom_gradient is not yet implemented " "for more than 3 gradients."
+                "custom_gradient is not yet implemented for more than 3 gradients."
             )
 
         return wrapped_function
@@ -200,7 +200,7 @@ def _value_and_jacobian_op(fun, x):
     return ans, _np.reshape(_np.stack(grads), jacobian_shape)
 
 
-def _value_and_jacobian(fun, point_ndim=1):
+def value_and_jacobian(fun, point_ndim=1):
     def _value_and_jacobian_vec(x):
         if x.ndim == point_ndim:
             return _value_and_jacobian_op(fun)(x)
@@ -335,7 +335,7 @@ def jacobian_and_hessian(func, func_out_ndim=None):
         Function that returns func's jacobian and
         func's hessian values at its inputs args.
     """
-    return _value_and_jacobian(jacobian_vec(func))
+    return value_and_jacobian(jacobian_vec(func))
 
 
 def value_jacobian_and_hessian(func, func_out_ndim=None):
@@ -355,7 +355,7 @@ def value_jacobian_and_hessian(func, func_out_ndim=None):
 
     def _cached_value_and_jacobian(fun, return_cached=False):
         def _jac(x):
-            ans, jac = _value_and_jacobian(fun)(x)
+            ans, jac = value_and_jacobian(fun)(x)
             if not return_cached:
                 cache.append(ans)
                 return jac

pyrecest/_backend/autograd/linalg.py

@@ -12,17 +12,24 @@
     eigh,
     eigvalsh,
     inv,
+    matrix_power,
     matrix_rank,
     norm,
-    solve,
     svd,
 )
 from autograd.scipy.linalg import expm
 from scipy.optimize import quadratic_assignment as _quadratic_assignment
 
-from .._shared_numpy.linalg import fractional_matrix_power, is_single_matrix_pd
+from .._shared_numpy.linalg import (
+    fractional_matrix_power,
+    is_single_matrix_pd,
+    polar,
+    qr,
+    solve,
+    solve_sylvester,
+    sqrtm,
+)
 from .._shared_numpy.linalg import logm as _logm
-from .._shared_numpy.linalg import qr, solve_sylvester, sqrtm
 
 
 def _adjoint(_ans, x, fn):