Trajectree 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

trajectree/quimb/quimb/linalg/rand_linalg.py

@@ -0,0 +1,514 @@
+"""Randomized iterative methods for decompositions.
+"""
+from numbers import Integral
+
+import numpy as np
+import scipy.linalg as sla
+
+from ..gen.rand import randn
+from ..core import dag, dot, njit
+from ..utils import identity
+
+
+def lu_orthog(X):
+    return sla.lu(X, permute_l=True, overwrite_a=True, check_finite=False)[0]
+
+
+def qr_orthog(X):
+    return sla.qr(X, mode="economic", overwrite_a=True, check_finite=False)[0]
+
+
+def orthog(X, lu=False):
+    if lu:
+        return lu_orthog(X)
+    return qr_orthog(X)
+
+
+def QB_to_svd(Q, B, compute_uv=True):
+    UsV = sla.svd(
+        B,
+        full_matrices=False,
+        compute_uv=compute_uv,
+        overwrite_a=True,
+        check_finite=False,
+    )
+
+    if not compute_uv:
+        return UsV
+
+    U, s, V = UsV
+    return dot(Q, U), s, V
+
+
+def trim(arrays, k):
+    if isinstance(arrays, tuple) and len(arrays) == 3:
+        U, s, VH = arrays
+        U, s, VH = U[:, :k], s[:k], VH[:k, :]
+        return U, s, VH
+    if isinstance(arrays, tuple) and len(arrays) == 2:
+        # Q, B factors
+        Q, B = arrays
+        return Q[:, :k], B[:k, :]
+    else:
+        # just singular values
+        return arrays[:k]
+
+
+def possibly_extend_randn(G, k, p, A):
+    # make sure we are using block of the right size by removing or adding
+    kG = G.shape[1]
+    if kG > k + p:
+        # have too many columns
+        G = G[:, : k + p]
+    elif kG < k + p:
+        # have too few columns
+        G_extra = randn((A.shape[1], k + p - kG), dtype=A.dtype)
+        G = np.concatenate((G, G_extra), axis=1)
+    return G
+
+
+def isstring(x, s):
+    if not isinstance(x, str):
+        return False
+    return x == s
+
+
+def rsvd_qb(A, k, q, p, state, AH=None):
+    if AH is None:
+        AH = dag(A)
+
+    # generate first block
+    if isstring(state, "begin-qb"):
+        G = randn((A.shape[1], k + p), dtype=A.dtype)
+    # block already supplied
+    elif len(state) == 1:
+        (G,) = state
+    # mid-way through adaptive algorithm in QB mode
+    if len(state) == 3:
+        Q, B, G = state
+    else:
+        Q = np.empty((A.shape[0], 0), dtype=A.dtype)
+        B = np.empty((0, A.shape[1]), dtype=A.dtype)
+
+    QH, BH = dag(Q), dag(B)
+    G = possibly_extend_randn(G, k, p, A)
+
+    Qi = orthog(dot(A, G) - dot(Q, dot(B, G)), lu=q > 0)
+
+    for i in range(1, q + 1):
+        Qi = orthog(dot(AH, Qi) - dot(BH, dot(QH, Qi)), lu=True)
+        Qi = orthog(dot(A, Qi) - dot(Q, dot(B, Qi)), lu=i != q)
+
+    Qi = orthog(Qi - dot(Q, dot(QH, Qi)))
+    Bi = dag(dot(AH, Qi)) - dot(dot(dag(Qi), Q), B)
+
+    if p > 0:
+        Qi, Bi = trim((Qi, Bi), k)
+
+    Q = np.concatenate((Q, Qi), axis=1)
+    B = np.concatenate((B, Bi), axis=0)
+
+    return Q, B, G
+
+
+def rsvd_core(A, k, compute_uv=True, q=2, p=0, state=None, AH=None):
+    """Core R3SVD algorithm.
+
+    Parameters
+    ----------
+    A : linear operator, shape (m, n)
+        Operator to decompose, assumed m >= n.
+    k : int
+        Number of singular values to find.
+    compute_uv : bool, optional
+        Return the left and right singular vectors.
+    q : int, optional
+        Number of power iterations.
+    p : int, optional
+        Over sampling factor.
+    state : {None, array_like, (), (G0,), (U0, s0, VH0, G0)}, optional
+        Iterate based on these previous results:
+
+            - None: basic mode.
+            - array_like: use this as the initial subspace.
+            - 'begin-svd': begin block iterations, return U, s, VH, G
+            - (G0,) : begin block iterations with this subspace
+            - (U0, s0, VH0, G0): continue block iterations, return G
+
+    """
+    iterating = isinstance(state, (tuple, str))
+    maybe_project_left = maybe_project_right = identity
+
+    if AH is None:
+        AH = dag(A)
+
+    # generate first block
+    if state is None or isstring(state, "begin-svd"):
+        G = randn((A.shape[1], k + p), dtype=A.dtype)
+    # initial block supplied
+    elif hasattr(state, "shape"):
+        G = state
+    elif len(state) == 1:
+        (G,) = state
+    # mid-way through adaptive algorithm in SVD mode
+    elif len(state) == 4:
+        U0, s0, VH0, G = state
+        UH0, V0 = dag(U0), dag(VH0)
+
+        def maybe_project_left(X):
+            X -= dot(U0, dot(UH0, X))
+            return X
+
+        def maybe_project_right(X):
+            X -= dot(V0, dot(VH0, X))
+            return X
+
+    G = possibly_extend_randn(G, k, p, A)
+    G = maybe_project_right(G)
+
+    Q = dot(A, G)
+    Q = maybe_project_left(Q)
+    Q = orthog(Q, lu=q > 0)
+
+    # power iterations with stabilization
+    for i in range(1, q + 1):
+        Q = dot(AH, Q)
+        Q = maybe_project_right(Q)
+        Q = orthog(Q, lu=True)
+
+        Q = dot(A, Q)
+        Q = maybe_project_left(Q)
+        Q = orthog(Q, lu=i < q)
+
+    B = dag(dot(AH, Q))
+    UsVH = QB_to_svd(Q, B, compute_uv=compute_uv or iterating)
+    if p > 0:
+        UsVH = trim(UsVH, k)
+
+    if not iterating:
+        return UsVH
+
+    U, s, VH = UsVH
+
+    if isstring(state, "begin-svd") or len(state) == 1:
+        # first run -> don't need to project or concatenate anything
+        return U, s, VH, G
+
+    U = orthog(maybe_project_left(U))
+    VH = dag(orthog(maybe_project_right(dag(VH))))
+
+    U = np.concatenate((U0, U), axis=1)
+    s = np.concatenate((s0, s))
+    VH = np.concatenate((VH0, VH), axis=0)
+
+    return U, s, VH, G
+
+
+@njit
+def is_sorted(x):  # pragma: no cover
+    for i in range(x.size - 1):
+        if x[i + 1] < x[i]:
+            return False
+    return True
+
+
+def gen_k_steps(start, incr=1.4):
+    yield start
+    step = start
+    while True:
+        yield step
+        step = round(incr * step)
+
+
+def rsvd_iterate(
+    A,
+    eps,
+    compute_uv=True,
+    q=2,
+    p=0,
+    G0=None,
+    k_max=None,
+    k_start=2,
+    k_incr=1.4,
+    AH=None,
+    use_qb=20,
+):
+    """Handle rank-adaptively calling ``rsvd_core``."""
+
+    if AH is None:
+        AH = dag(A)
+
+    # perform first iteration and set initial rank
+    k_steps = gen_k_steps(k_start, k_incr)
+    rank = next(k_steps)
+
+    if use_qb:
+        Q, B, G = rsvd_qb(
+            A, rank, q=q, p=p, AH=AH, state="begin-qb" if G0 is None else (G0,)
+        )
+        U, s, VH = QB_to_svd(Q, B)
+        G -= dot(dag(VH), dot(VH, G))
+    else:
+        U, s, VH, G = rsvd_core(
+            A,
+            rank,
+            q=q,
+            p=p,
+            AH=AH,
+            state="begin-svd" if G0 is None else (G0,),
+        )
+
+    # perform randomized SVD in small blocks
+    while (s[-1] > eps * s[0]) and (rank < k_max):
+        # only step k as far as k_max
+        new_k = min(next(k_steps), k_max - rank)
+        rank += new_k
+
+        if (rank < use_qb) or (use_qb is True):
+            Q, B, G = rsvd_qb(A, new_k, q=q, p=p, state=(Q, B, G), AH=AH)
+            U, s, VH = QB_to_svd(Q, B)
+            G -= dot(dag(VH), dot(VH, G))
+        else:
+            # concatenate new U, s, VH orthogonal to current U, s, VH
+            U, s, VH, G = rsvd_core(
+                A, new_k, q=q, p=p, state=(U, s, VH, G), AH=AH
+            )
+
+    # make sure singular values always sorted in decreasing order
+    if not is_sorted(s):
+        so = np.argsort(s)[::-1]
+        U, s, VH = U[:, so], s[so], VH[so, :]
+
+    return U, s, VH if compute_uv else s
+
+
+@njit
+def count_svdvals_needed(s, eps):  # pragma: no cover
+    n = s.size
+    thresh = eps * s[0]
+    for i in range(n - 1, 0, -1):
+        if s[i - 1] < thresh:
+            n -= 1
+        else:
+            break
+    return n
+
+
+def isdouble(dtype):
+    """Check if ``dtype`` is double precision."""
+    return dtype in ("float64", "complex128")
+
+
+def estimate_rank(
+    A,
+    eps,
+    k_max=None,
+    use_sli=True,
+    k_start=2,
+    k_incr=1.4,
+    q=0,
+    p=0,
+    get_vectors=False,
+    G0=None,
+    AH=None,
+    use_qb=20,
+):
+    """Estimate the rank of an linear operator. Uses a low quality random
+    SVD with a resolution of ~ 10.
+
+    Parameters
+    ----------
+    A : linear operator
+        The operator to find rank of.
+    eps : float
+        Find rank to this relative (compared to largest singular value)
+        precision.
+    k_max : int, optional
+        The maximum rank to find.
+    use_sli : bool, optional
+        Whether to use :func:`scipy.linalg.interpolative.estimate_rank` if
+        possible (double precision and no ``k_max`` set).
+    k_start : int, optional
+        Begin the adaptive SVD with a block of this size.
+    k_incr : float, optional
+        Adaptive rank increment factor. Increase the k-step (from k_start) by
+        this factor each time. Set to 1 to use a constant step.
+    q : int, optional
+        Number of power iterations.
+    get_vectors : bool, optional
+        Return the right singular vectors found in the pass.
+    G0 : , optional
+
+    Returns
+    -------
+    rank : int
+        The rank.
+    VH : array
+        The (adjoint) right singular vectors if ``get_vectors=True``.
+    """
+    if k_max is None:
+        k_max = min(A.shape)
+    if eps <= 0.0:
+        return k_max
+
+    use_sli = (
+        use_sli
+        and (k_max == min(A.shape))
+        and isdouble(A.dtype)
+        and not get_vectors
+    )
+    if use_sli:
+        return sla.interpolative.estimate_rank(A, eps)
+
+    if A.shape[0] < A.shape[1]:
+        A = A.T
+        if get_vectors:
+            raise ValueError
+    if AH is None:
+        AH = dag(A)
+
+    _, s, VH = rsvd_iterate(
+        A,
+        eps,
+        q=q,
+        p=p,
+        G0=G0,
+        AH=AH,
+        use_qb=use_qb,
+        k_start=k_start,
+        k_max=k_max,
+        k_incr=k_incr,
+    )
+
+    rank = count_svdvals_needed(s, eps)
+
+    if get_vectors:
+        return rank, VH[:rank, :]
+    return rank
+
+
+def maybe_flip(UsV, flipped):
+    # if only singular values or only tranposing do nothing
+    if not (isinstance(UsV, tuple) and flipped):
+        return UsV
+    U, s, V = UsV
+    return V.T, s, U.T
+
+
+def rsvd(
+    A,
+    eps_or_k,
+    compute_uv=True,
+    mode="adapt+block",
+    use_qb=20,
+    q=2,
+    p=0,
+    k_max=None,
+    k_start=2,
+    k_incr=1.4,
+    G0=None,
+    AH=None,
+):
+    """Fast, randomized, iterative SVD. Adaptive variant of method due
+    originally to Halko. This scales as ``log(k)`` rather than ``k`` so can be
+    more efficient.
+
+    Parameters
+    ----------
+    A : operator, shape (m, n)
+        The operator to decompose.
+    eps_or_k : float or int
+        Either the relative precision or the number of singular values to
+        target. If precision, this is relative to the largest singular value.
+    compute_uv : bool, optional
+        Whether to return the left and right singular vectors.
+    mode : {'adapt+block', 'adapt', 'block'}, optional
+        How to perform the randomized SVD. If ``eps_or_k`` is an integer then
+        this is implicitly 'block' and ignored. Else:
+
+            - 'adapt+block', perform an initial low quality pass to estimate
+              the rank of ``A``, then use the subspace and rank from that to
+              perform an accurate fully blocked RSVD.
+            - 'adapt', just perform the adaptive randomized SVD.
+
+    q : int, optional
+        The number of power iterations, increase for accuracy at the expense
+        of runtime.
+    p : int, optional
+        Oversampling factor. Perform projections with this many extra columns
+        and then throw then away.
+    k_max : int, optional
+        Maximum adaptive rank. Default: ``min(A.shape)``.
+    k_start : int, optional
+        Initial k when increasing rank adaptively.
+    k_incr : float, optional
+        Adaptive rank increment factor. Increase the k-step (from k_start) by
+        this factor each time. Set to 1 to use a constant step.
+    G0 : array_like, shape (n, k), optional
+        Initial subspace to start iterating on. If not given a random one will
+        be generated.
+
+    Returns
+    -------
+    U, array, shape (m, k)
+        Left singular vectors, if ``compute_uv=True``.
+    s, array, shape (k,)
+        Singular values.
+    V, array, shape (k, n)
+        Right singular vectors, if ``compute_uv=True``.
+    """
+
+    flipped = A.shape[0] < A.shape[1]
+    if flipped:
+        A = A.T
+
+    # 'block' mode -> just perform single pass random SVD
+    if isinstance(eps_or_k, Integral):
+        UsV = rsvd_core(A, eps_or_k, q=q, p=p, state=G0, compute_uv=compute_uv)
+        return maybe_flip(UsV, flipped)
+
+    if k_max is None:
+        k_max = min(A.shape)
+    k_max = min(max(1, k_max), min(A.shape))
+
+    if AH is None:
+        AH = dag(A)
+
+    adaptive_opts = {
+        "k_start": k_start,
+        "k_max": k_max,
+        "k_incr": k_incr,
+        "use_qb": use_qb,
+        "AH": AH,
+        "G0": G0,
+    }
+
+    # 'adapt' mode -> rank adaptively perform SVD to low accuracy
+    if mode == "adapt":
+        UsV = rsvd_iterate(
+            A, eps_or_k, q=q, p=p, compute_uv=compute_uv, **adaptive_opts
+        )
+
+    # 'adapt+block' mode -> use first pass to find rank, then use blocking mode
+    elif mode == "adapt+block":
+        # estimate both rank and get approximate spanning vectors
+        k, VH = estimate_rank(A, eps_or_k, get_vectors=True, **adaptive_opts)
+
+        # reuse vectors to effectively boost number of power iterations by one
+        UsV = rsvd_core(
+            A,
+            k,
+            q=max(q - 1, 0),
+            p=p,
+            AH=AH,
+            state=dag(VH),
+            compute_uv=compute_uv,
+        )
+
+    else:
+        raise ValueError(
+            "``mode`` must  be one of {'adapt+block', 'adapt'} or"
+            " ``k`` should be a integer to use 'block' mode."
+        )
+
+    return maybe_flip(UsV, flipped)