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

trajectree/quimb/quimb/linalg/approx_spectral.py

@@ -0,0 +1,1109 @@
+"""Use stochastic Lanczos quadrature to approximate spectral function sums of
+any operator which has an efficient representation of action on a vector.
+"""
+
+import functools
+import random
+import warnings
+from math import exp, inf, log2, nan, sqrt
+
+import numpy as np
+import scipy.linalg as scla
+from scipy.ndimage import uniform_filter1d
+
+from ..core import divide_update_, dot, njit, prod, ptr, subtract_update_, vdot
+from ..gen.rand import rand_phase, rand_rademacher, randn, seed_rand
+from ..linalg.mpi_launcher import get_mpi_pool
+from ..utils import (
+    default_to_neutral_style,
+    find_library,
+    format_number_with_error,
+    int2tup,
+    raise_cant_find_library_function,
+)
+from ..utils import progbar as Progbar
+
+if find_library("cotengra") and find_library("autoray"):
+    from ..tensor.tensor_1d import MatrixProductOperator
+    from ..tensor.tensor_approx_spectral import construct_lanczos_tridiag_MPO
+    from ..tensor.tensor_core import Tensor
+else:
+    reqs = "[cotengra,autoray]"
+    Tensor = raise_cant_find_library_function(reqs)
+    construct_lanczos_tridiag_MPO = raise_cant_find_library_function(reqs)
+
+
+# --------------------------------------------------------------------------- #
+#                  'Lazy' representation tensor contractions                  #
+# --------------------------------------------------------------------------- #
+
+
+def lazy_ptr_linop(psi_ab, dims, sysa, **linop_opts):
+    r"""A linear operator representing action of partially tracing a bipartite
+    state, then multiplying another 'unipartite' state::
+
+          ( | )
+        +-------+
+        | psi_a |   ______
+        +_______+  /      \
+           a|      |b     |
+        +-------------+   |
+        |  psi_ab.H   |   |
+        +_____________+   |
+                          |
+        +-------------+   |
+        |   psi_ab    |   |
+        +_____________+   |
+           a|      |b     |
+            |      \______/
+
+    Parameters
+    ----------
+    psi_ab : ket
+        State to partially trace and dot with another ket, with
+        size ``prod(dims)``.
+    dims : sequence of int, optional
+        The sub dimensions of ``psi_ab``.
+    sysa : int or sequence of int, optional
+        Index(es) of the 'a' subsystem(s) to keep.
+    """
+    sysa = int2tup(sysa)
+
+    Kab = Tensor(
+        np.asarray(psi_ab).reshape(dims),
+        inds=[
+            ("kA{}" if i in sysa else "xB{}").format(i)
+            for i in range(len(dims))
+        ],
+    )
+
+    Bab = Tensor(
+        Kab.data.conjugate(),
+        inds=[
+            ("bA{}" if i in sysa else "xB{}").format(i)
+            for i in range(len(dims))
+        ],
+    )
+
+    return (Kab & Bab).aslinearoperator(
+        [f"kA{i}" for i in sysa], [f"bA{i}" for i in sysa], **linop_opts
+    )
+
+
+def lazy_ptr_ppt_linop(psi_abc, dims, sysa, sysb, **linop_opts):
+    r"""A linear operator representing action of partially tracing a tripartite
+    state, partially transposing the remaining bipartite state, then
+    multiplying another bipartite state::
+
+             ( | )
+        +--------------+
+        |   psi_ab     |
+        +______________+  _____
+         a|  ____   b|   /     \
+          | /   a\   |   |c    |
+          | | +-------------+  |
+          | | |  psi_abc.H  |  |
+          \ / +-------------+  |
+           X                   |
+          / \ +-------------+  |
+          | | |   psi_abc   |  |
+          | | +-------------+  |
+          | \____/a  |b  |c    |
+         a|          |   \_____/
+
+    Parameters
+    ----------
+    psi_abc : ket
+        State to partially trace, partially transpose, then dot with another
+        ket, with size ``prod(dims)``.
+        ``prod(dims[sysa] + dims[sysb])``.
+    dims : sequence of int
+        The sub dimensions of ``psi_abc``.
+    sysa : int or sequence of int, optional
+        Index(es) of the 'a' subsystem(s) to keep, with respect to all
+        the dimensions, ``dims``, (i.e. pre-partial trace).
+    sysa : int or sequence of int, optional
+        Index(es) of the 'b' subsystem(s) to keep, with respect to all
+        the dimensions, ``dims``, (i.e. pre-partial trace).
+    """
+    sysa, sysb = int2tup(sysa), int2tup(sysb)
+    sys_ab = sorted(sysa + sysb)
+
+    Kabc = Tensor(
+        np.asarray(psi_abc).reshape(dims),
+        inds=[
+            ("kA{}" if i in sysa else "kB{}" if i in sysb else "xC{}").format(
+                i
+            )
+            for i in range(len(dims))
+        ],
+    )
+
+    Babc = Tensor(
+        Kabc.data.conjugate(),
+        inds=[
+            ("bA{}" if i in sysa else "bB{}" if i in sysb else "xC{}").format(
+                i
+            )
+            for i in range(len(dims))
+        ],
+    )
+
+    return (Kabc & Babc).aslinearoperator(
+        [("bA{}" if i in sysa else "kB{}").format(i) for i in sys_ab],
+        [("kA{}" if i in sysa else "bB{}").format(i) for i in sys_ab],
+        **linop_opts,
+    )
+
+
+# --------------------------------------------------------------------------- #
+#                         Lanczos tri-diag technique                          #
+# --------------------------------------------------------------------------- #
+
+
+def inner(a, b):
+    """Inner product between two vectors"""
+    return vdot(a, b).real
+
+
+def norm_fro(a):
+    """'Frobenius' norm of a vector."""
+    return sqrt(inner(a, a))
+
+
+def norm_fro_approx(A, **kwargs):
+    r"""Calculate the approximate frobenius norm of any hermitian linear
+    operator:
+
+    .. math::
+
+        \mathrm{Tr} \left[ A^{\dagger} A \right]
+
+    Parameters
+    ----------
+    A : linear operator like
+        Operator with a dot method, assumed to be hermitian, to estimate the
+        frobenius norm of.
+    kwargs
+        Supplied to :func:`approx_spectral_function`.
+
+    Returns
+    -------
+    float
+    """
+    return approx_spectral_function(A, lambda x: x**2, **kwargs) ** 0.5
+
+
+def random_rect(
+    shape,
+    dist="rademacher",
+    orthog=False,
+    norm=True,
+    seed=False,
+    dtype=complex,
+):
+    """Generate a random array optionally orthogonal.
+
+    Parameters
+    ----------
+    shape : tuple of int
+        The shape of array.
+    dist : {'guassian', 'rademacher'}
+        Distribution of the random variables.
+    orthog : bool or operator.
+        Orthogonalize the columns if more than one.
+    norm : bool
+        Explicitly normalize the frobenius norm to 1.
+    """
+    if seed:
+        # needs to be truly random so e.g. MPI processes don't overlap
+        seed_rand(random.SystemRandom().randint(0, 2**32 - 1))
+
+    if dist == "rademacher":
+        V = rand_rademacher(shape, scale=1 / sqrt(prod(shape)), dtype=dtype)
+        # already normalized
+
+    elif dist == "gaussian":
+        V = randn(shape, scale=1 / (prod(shape) ** 0.5 * 2**0.5), dtype=dtype)
+        if norm:
+            V /= norm_fro(V)
+
+    elif dist == "phase":
+        V = rand_phase(shape, scale=1 / sqrt(prod(shape)), dtype=dtype)
+        # already normalized
+
+    else:
+        raise ValueError(f"`dist={dist}` not understood.")
+
+    if orthog and min(shape) > 1:
+        V = scla.orth(V)
+        V /= sqrt(min(V.shape))
+
+    return V
+
+
+def construct_lanczos_tridiag(
+    A,
+    K,
+    v0=None,
+    bsz=1,
+    k_min=10,
+    orthog=False,
+    beta_tol=1e-6,
+    seed=False,
+    v0_opts=None,
+):
+    """Construct the tridiagonal lanczos matrix using only matvec operators.
+    This is a generator that iteratively yields the alpha and beta digaonals
+    at each step.
+
+    Parameters
+    ----------
+    A : dense array, sparse matrix or linear operator
+        The operator to approximate, must implement ``.dot`` method to compute
+        its action on a vector.
+    K : int, optional
+        The maximum number of iterations and thus rank of the matrix to find.
+    v0 : vector, optional
+        The starting vector to iterate with, default to random.
+    bsz : int, optional
+        The block size (number of columns) of random vectors to iterate with.
+    k_min : int, optional
+        The minimum size of the krylov subspace for form.
+    orthog : bool, optional
+        If True, perform full re-orthogonalization for each new vector.
+    beta_tol : float, optional
+        The 'breakdown' tolerance. If the next beta ceofficient in the lanczos
+        matrix is less that this, implying that the full non-null space has
+        been found, terminate early.
+    seed : bool, optional
+        If True, seed the numpy random generator with a system random int.
+
+    Yields
+    ------
+    alpha : sequence of float of length k
+        The diagonal entries of the lanczos matrix.
+    beta : sequence of float of length k
+        The off-diagonal entries of the lanczos matrix, with the last entry
+        the 'look' forward value.
+    scaling : float
+        How to scale the overall weights.
+    """
+    d = A.shape[0]
+
+    if bsz == 1:
+        v_shp = (d,)
+    else:
+        orthog = False
+        v_shp = (d, bsz)
+
+    alpha = np.zeros(K + 1, dtype=get_equivalent_real_dtype(A.dtype))
+    beta = np.zeros(K + 2, dtype=get_equivalent_real_dtype(A.dtype))
+    beta[1] = sqrt(prod(v_shp))  # by construction
+
+    if v0 is None:
+        if v0_opts is None:
+            v0_opts = {}
+        q = random_rect(v_shp, seed=seed, dtype=A.dtype, **v0_opts)
+    else:
+        q = v0.astype(A.dtype)
+        divide_update_(q, norm_fro(q), q)
+    v = np.zeros_like(q)
+
+    if orthog:
+        Q = np.copy(q).reshape(-1, 1)
+
+    for j in range(1, K + 1):
+        r = dot(A, q)
+        subtract_update_(r, beta[j], v)
+        alpha[j] = inner(q, r)
+        subtract_update_(r, alpha[j], q)
+
+        # perform full orthogonalization
+        if orthog:
+            r -= Q.dot(Q.conj().T.dot(r))
+
+        beta[j + 1] = norm_fro(r)
+
+        # check for convergence
+        if abs(beta[j + 1]) < beta_tol:
+            yield (
+                alpha[1 : j + 1].copy(),
+                beta[2 : j + 2].copy(),
+                beta[1] ** 2 / bsz,
+            )
+            break
+
+        v[()] = q
+        divide_update_(r, beta[j + 1], q)
+
+        # keep all vectors
+        if orthog:
+            Q = np.concatenate((Q, q.reshape(-1, 1)), axis=1)
+
+        if j >= k_min:
+            yield (
+                alpha[1 : j + 1].copy(),
+                beta[2 : j + 2].copy(),
+                beta[1] ** 2 / bsz,
+            )
+
+
+def lanczos_tridiag_eig(alpha, beta, check_finite=True):
+    """Find the eigen-values and -vectors of the Lanczos triadiagonal matrix.
+
+    Parameters
+    ----------
+    alpha : array of float
+        The diagonal.
+    beta : array of float
+        The k={-1, 1} off-diagonal. Only first ``len(alpha) - 1`` entries used.
+    """
+    Tk_banded = np.empty((2, alpha.size), dtype=alpha.dtype)
+    Tk_banded[1, -1] = 0.0  # sometimes can get nan here? -> breaks eig_banded
+    Tk_banded[0, :] = alpha
+    Tk_banded[1, : beta.size] = beta
+
+    try:
+        tl, tv = scla.eig_banded(
+            Tk_banded, lower=True, check_finite=check_finite
+        )
+
+    # sometimes get no convergence -> use dense hermitian method
+    except scla.LinAlgError:  # pragma: no cover
+        tl, tv = np.linalg.eigh(
+            np.diag(alpha) + np.diag(beta[: alpha.size - 1], -1), UPLO="L"
+        )
+
+    return tl, tv
+
+
+def calc_trace_fn_tridiag(tl, tv, f, pos=True):
+    """Spectral ritz function sum, weighted by ritz vectors."""
+    return sum(
+        tv[0, i] ** 2 * f(max(tl[i], 0.0) if pos else tl[i])
+        for i in range(tl.size)
+    )
+
+
+@njit
+def ext_per_trim(x, p=0.6, s=1.0):  # pragma: no cover
+    r"""Extended percentile trimmed-mean. Makes the mean robust to asymmetric
+    outliers, while using all data when it is nicely clustered. This can be
+    visualized roughly as::
+
+            |--------|=========|--------|
+        x     x   xx xx xxxxx xxx   xx     x      x           x
+
+    Where the inner range contains the central ``p`` proportion of the data,
+    and the outer ranges entends this by a factor of ``s`` either side.
+
+    Parameters
+    ----------
+    x : array
+        Data to trim.
+    p : Proportion of data used to define the 'central' percentile.
+        For example, p=0.5 gives the inter-quartile range.
+    s : Include data up to this factor times the central 'percentile' range
+        away from the central percentile itself.
+
+    Returns
+    xt : array
+        Trimmed data.
+    """
+    lb = np.percentile(x, 100 * (1 - p) / 2)
+    ub = np.percentile(x, 100 * (1 + p) / 2)
+    ib = ub - lb
+
+    trimmed_x = x[(lb - s * ib < x) & (x < ub + s * ib)]
+
+    return trimmed_x
+
+
+@njit  # pragma: no cover
+def nbsum(xs):
+    tot = 0
+    for x in xs:
+        tot += x
+    return tot
+
+
+@njit  # pragma: no cover
+def std(xs):
+    """Simple standard deviation - don't invoke numpy for small lists."""
+    N = len(xs)
+    xm = nbsum(xs) / N
+    var = nbsum([(x - xm) ** 2 for x in xs]) / N
+    return var**0.5
+
+
+def calc_est_fit(estimates, conv_n, tau):
+    """Make estimate by fitting exponential convergence to estimates."""
+    n = len(estimates)
+
+    if n < conv_n:
+        return nan, inf
+
+    # iteration number, fit function to inverse this to get k->infinity
+    ks = np.arange(1, len(estimates) + 1)
+
+    # smooth data with a running mean
+    smoothed_estimates = uniform_filter1d(estimates, n // 2)
+
+    # ignore this amount of the initial estimates and fit later part only
+    ni = n // 2
+
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+
+            # fit the inverse data with a line, weighting recent ests more
+            popt, pcov = np.polyfit(
+                x=(1 / ks[ni:]),
+                y=smoothed_estimates[ni:],
+                w=ks[ni:],
+                deg=1,
+                cov=True,
+            )
+
+        # estimate of function at 1 / k = 0 and standard error
+        est, err = popt[-1], abs(pcov[-1, -1]) ** 0.5
+
+    except (ValueError, RuntimeError):
+        est, err = nan, inf
+
+    return est, err
+
+
+def calc_est_window(estimates, conv_n):
+    """Make estimate from mean of last ``m`` samples, following:
+
+    1. Take between ``conv_n`` and 12 estimates.
+    2. Pair the estimates as they are alternate upper/lower bounds
+    3. Compute the standard error on the paired estimates.
+    """
+    m_est = min(max(conv_n, len(estimates) // 8), 12)
+    est = sum(estimates[-m_est:]) / len(estimates[-m_est:])
+
+    if len(estimates) > conv_n:
+        # check for convergence using variance of paired last m estimates
+        #   -> paired because estimates alternate between upper and lower bound
+        paired_ests = tuple(
+            (a + b) / 2
+            for a, b in zip(estimates[-m_est::2], estimates[-m_est + 1 :: 2])
+        )
+        err = std(paired_ests) / (m_est / 2) ** 0.5
+    else:
+        err = inf
+
+    return est, err
+
+
+def single_random_estimate(
+    A,
+    K,
+    bsz,
+    beta_tol,
+    v0,
+    f,
+    pos,
+    tau,
+    tol_scale,
+    k_min=10,
+    verbosity=0,
+    *,
+    seed=None,
+    v0_opts=None,
+    info=None,
+    **lanczos_opts,
+):
+    # choose normal (any LinearOperator) or MPO lanczos tridiag construction
+    if isinstance(A, MatrixProductOperator):
+        lanc_fn = construct_lanczos_tridiag_MPO
+    else:
+        lanc_fn = construct_lanczos_tridiag
+        lanczos_opts["bsz"] = bsz
+
+    estimates_raw = []
+    estimates_window = []
+    estimates_fit = []
+    estimates = []
+
+    # the number of samples to check standard deviation convergence with
+    conv_n = 6  # 3 pairs
+
+    # iteratively build the lanczos matrix, checking for convergence
+    for alpha, beta, scaling in lanc_fn(
+        A,
+        K=K,
+        beta_tol=beta_tol,
+        seed=seed,
+        k_min=k_min - 2 * conv_n,
+        v0=v0() if callable(v0) else v0,
+        v0_opts=v0_opts,
+        **lanczos_opts,
+    ):
+        try:
+            Tl, Tv = lanczos_tridiag_eig(alpha, beta, check_finite=False)
+            Gf = scaling * calc_trace_fn_tridiag(Tl, Tv, f=f, pos=pos)
+        except scla.LinAlgError:  # pragma: no cover
+            warnings.warn("Approx Spectral Gf tri-eig didn't converge.")
+            estimates_raw.append(np.nan)
+            continue
+
+        k = alpha.size
+        estimates_raw.append(Gf)
+
+        # check for break-down convergence (e.g. found entire subspace)
+        #     in which case latest estimate should be accurate
+        if abs(beta[-1]) < beta_tol:
+            if verbosity >= 2:
+                print(f"k={k}: Beta breadown, returning {Gf}.")
+            est = Gf
+            estimates.append(est)
+            break
+
+        # compute an estimate and error using a window of the last few results
+        win_est, win_err = calc_est_window(estimates_raw, conv_n)
+        estimates_window.append(win_est)
+
+        # try and compute an estimate and error using exponential fit
+        fit_est, fit_err = calc_est_fit(estimates_window, conv_n, tau)
+        estimates_fit.append(fit_est)
+
+        # take whichever has lowest error
+        est, err = min(
+            (win_est, win_err),
+            (fit_est, fit_err),
+            key=lambda est_err: est_err[1],
+        )
+        estimates.append(est)
+        converged = err < tau * (abs(win_est) + tol_scale)
+
+        if verbosity >= 2:
+            if verbosity >= 3:
+                print(f"est_win={win_est}, err_win={win_err}")
+                print(f"est_fit={fit_est}, err_fit={fit_err}")
+            print(f"k={k}: Gf={Gf}, Est={est}, Err={err}")
+            if converged:
+                print(f"k={k}: Converged to tau {tau}.")
+
+        if converged:
+            break
+
+    if verbosity >= 1:
+        print(f"k={k}: Returning estimate {est}.")
+
+    if info is not None:
+        if "estimates_raw" in info:
+            info["estimates_raw"].append(estimates_raw)
+        if "estimates_window" in info:
+            info["estimates_window"].append(estimates_window)
+        if "estimates_fit" in info:
+            info["estimates_fit"].append(estimates_fit)
+        if "estimates" in info:
+            info["estimates"].append(estimates)
+
+    return est
+
+
+def calc_stats(samples, mean_p, mean_s, tol, tol_scale):
+    """Get an estimate from samples."""
+    samples = np.array(samples)
+
+    xtrim = ext_per_trim(samples, p=mean_p, s=mean_s)
+
+    # sometimes everything is an outlier...
+    if xtrim.size == 0:  # pragma: no cover
+        estimate, sdev = np.mean(samples), std(samples)
+    else:
+        estimate, sdev = np.mean(xtrim), std(xtrim)
+
+    err = sdev / len(samples) ** 0.5
+
+    converged = err < tol * (abs(estimate) + tol_scale)
+
+    return estimate, err, converged
+
+
+def get_single_precision_dtype(dtype):
+    if np.issubdtype(dtype, np.complexfloating):
+        return np.complex64
+    elif np.issubdtype(dtype, np.floating):
+        return np.float32
+    else:
+        raise ValueError(f"dtype {dtype} not understood.")
+
+
+def get_equivalent_real_dtype(dtype):
+    if dtype in ("float64", "complex128"):
+        return "float64"
+    elif dtype in ("float32", "complex64"):
+        return "float32"
+    else:
+        raise ValueError(f"dtype {dtype} not understood.")
+
+
+@default_to_neutral_style
+def plot_approx_spectral_info(info):
+    from matplotlib import pyplot as plt
+    from matplotlib.ticker import MaxNLocator
+
+    fig, axs = plt.subplots(
+        ncols=2,
+        figsize=(8, 4),
+        sharey=True,
+        gridspec_kw={"width_ratios": [3, 1]},
+    )
+    plt.subplots_adjust(wspace=0.0)
+
+    Z = info["estimate"]
+
+    alpha = len(info["estimates_raw"])**-(1 / 6)
+
+    # plot the raw kyrlov runs
+    for x in info["estimates_raw"]:
+        axs[0].plot(x, ".-", alpha=alpha, lw=1 / 2, zorder=-10, markersize=1)
+    axs[0].axhline(Z - info["error"], color="grey", linestyle="--")
+    axs[0].axhline(Z + info["error"], color="grey", linestyle="--")
+    axs[0].axhline(Z, color="black", linestyle="--")
+    axs[0].set_rasterization_zorder(-5)
+    axs[0].set_xlabel("krylov iteration (offset)")
+    axs[0].xaxis.set_major_locator(MaxNLocator(integer=True))
+    axs[0].set_ylabel("$Tr[f(x)]$ approximation")
+
+    # plot the overall final samples
+    axs[1].hist(
+        info["samples"],
+        bins=round(len(info["samples"])**0.5),
+        orientation="horizontal",
+        color=(0.2, 0.6, 1.0),
+    )
+    axs[1].axhline(Z - info["error"], color="grey", linestyle="--")
+    axs[1].axhline(Z + info["error"], color="grey", linestyle="--")
+    axs[1].axhline(Z, color="black", linestyle="--")
+    axs[1].set_xlabel("sample count")
+    axs[1].set_title(
+        "estimate ≈ " + format_number_with_error(Z, info["error"]),
+        ha="right",
+    )
+
+    # plot the correlation between raw and fitted estimates
+    iax = axs[0].inset_axes((0.03, 0.6, 0.3, 0.3))
+    iax.set_aspect("equal")
+    x = [es[-1] for es in info["estimates"]]
+    y = [es[-1] for es in info["estimates_raw"]]
+    iax.scatter(x, y, marker=".", alpha=alpha, color=(0.3, 0.7, 0.3), s=1)
+
+    return fig, axs
+
+
+def approx_spectral_function(
+    A,
+    f,
+    tol=1e-2,
+    *,
+    bsz=1,
+    R=1024,
+    R_min=3,
+    tol_scale=1,
+    tau=1e-4,
+    k_min=10,
+    k_max=512,
+    beta_tol=1e-6,
+    mpi=False,
+    mean_p=0.7,
+    mean_s=1.0,
+    pos=False,
+    v0=None,
+    verbosity=0,
+    single_precision="AUTO",
+    info=None,
+    progbar=False,
+    plot=False,
+    **lanczos_opts,
+):
+    """Approximate a spectral function, that is, the quantity ``Tr(f(A))``.
+
+    Parameters
+    ----------
+    A : dense array, sparse matrix or LinearOperator
+        Operator to approximate spectral function for. Should implement
+        ``A.dot(vec)``.
+    f : callable
+        Scalar function with which to act on approximate eigenvalues.
+    tol : float, optional
+        Relative convergence tolerance threshold for error on mean of repeats.
+        This can pretty much be relied on as the overall accuracy. See also
+        ``tol_scale`` and ``tau``. Default: 1%.
+    bsz : int, optional
+        Number of simultenous vector columns to use at once, 1 equating to the
+        standard lanczos method. If ``bsz > 1`` then ``A`` must implement
+        matrix-matrix multiplication. This is a more performant way of
+        essentially increasing ``R``, at the cost of more memory. Default: 1.
+    R : int, optional
+        The number of repeats with different initial random vectors to perform.
+        Increasing this should increase accuracy as ``sqrt(R)``. Cost of
+        algorithm thus scales linearly with ``R``. If ``tol`` is non-zero, this
+        is the maximum number of repeats.
+    R_min : int, optional
+        The minimum number of repeats to perform. Default: 3.
+    tau : float, optional
+        The relative tolerance required for a single lanczos run to converge.
+        This needs to be small enough that each estimate with a single random
+        vector produces an unbiased sample of the operators spectrum..
+    k_min : int, optional
+        The minimum size of the krylov subspace to form for each sample.
+    k_max : int, optional
+        The maximum size of the kyrlov space to form. Cost of algorithm scales
+        linearly with ``K``. If ``tau`` is non-zero, this is the maximum size
+        matrix to form.
+    tol_scale : float, optional
+        This sets the overall expected scale of each estimate, so that an
+        absolute tolerance can be used for values near zero. Default: 1.
+    beta_tol : float, optional
+        The 'breakdown' tolerance. If the next beta ceofficient in the lanczos
+        matrix is less that this, implying that the full non-null space has
+        been found, terminate early. Default: 1e-6.
+    mpi : bool, optional
+        Whether to parallelize repeat runs over MPI processes.
+    mean_p : float, optional
+        Factor for robustly finding mean and err of repeat estimates,
+        see :func:`ext_per_trim`.
+    mean_s : float, optional
+        Factor for robustly finding mean and err of repeat estimates,
+        see :func:`ext_per_trim`.
+    v0 : vector, or callable
+        Initial vector to iterate with, sets ``R=1`` if given. If callable, the
+        function to produce a random intial vector (sequence).
+    pos : bool, optional
+        If True, make sure any approximate eigenvalues are positive by
+        clipping below 0.
+    verbosity : {0, 1, 2}, optional
+        How much information to print while computing.
+    single_precision : {'AUTO', False, True}, optional
+        Try and convert the operator to single precision. This can lead to much
+        faster operation, especially if a GPU is available. Additionally,
+        double precision is not really needed given the stochastic nature of
+        the algorithm.
+    lanczos_opts
+        Supplied to
+        :func:`~quimb.linalg.approx_spectral.single_random_estimate` or
+        :func:`~quimb.linalg.approx_spectral.construct_lanczos_tridiag`.
+
+
+    Returns
+    -------
+    scalar
+        The approximate value ``Tr(f(a))``.
+
+    See Also
+    --------
+    construct_lanczos_tridiag
+    """
+    if single_precision == "AUTO":
+        single_precision = hasattr(A, "astype")
+    if single_precision:
+        A = A.astype(get_single_precision_dtype(A.dtype))
+
+    if (v0 is not None) and not callable(v0):
+        # we only have one sample to run
+        R = 1
+    else:
+        R = max(1, int(R / bsz))
+
+    if tau is None:
+        # require better precision for the lanczos procedure, otherwise biased
+        tau = tol / 1000
+
+    if verbosity:
+        print(f"LANCZOS f(A) CALC: tol={tol}, tau={tau}, R={R}, bsz={bsz}")
+
+    if plot:
+        # need to store all the info
+        if info is None:
+            info = {}
+        info.setdefault('estimate', None)
+        info.setdefault('error', None)
+        info.setdefault('samples', None)
+        info.setdefault('estimates_raw', [])
+        info.setdefault('estimates_window', [])
+        info.setdefault('estimates_fit', [])
+        info.setdefault('estimates', [])
+
+    # generate repeat estimates
+    kwargs = {
+        "A": A,
+        "K": k_max,
+        "bsz": bsz,
+        "beta_tol": beta_tol,
+        "v0": v0,
+        "f": f,
+        "pos": pos,
+        "tau": tau,
+        "k_min": k_min,
+        "tol_scale": tol_scale,
+        "verbosity": verbosity,
+        "info": info,
+        **lanczos_opts,
+    }
+
+    if not mpi:
+
+        def gen_results():
+            for _ in range(R):
+                yield single_random_estimate(**kwargs)
+
+    else:
+        pool = get_mpi_pool()
+        kwargs["seed"] = True
+        fs = [pool.submit(single_random_estimate, **kwargs) for _ in range(R)]
+
+        def gen_results():
+            for f in fs:
+                yield f.result()
+
+    if progbar:
+        pbar = Progbar(total=R)
+    else:
+        pbar = None
+
+    # iterate through estimates, waiting for convergence
+    results = gen_results()
+    estimate = None
+    samples = []
+    for _ in range(R):
+        samples.append(next(results))
+
+        if verbosity >= 1:
+            print(f"Repeat {len(samples)}: estimate is {samples[-1]}")
+
+        # wait a few iterations before checking error on mean breakout
+        if len(samples) >= R_min:
+            estimate, err, converged = calc_stats(
+                samples, mean_p, mean_s, tol, tol_scale
+            )
+            if verbosity >= 1:
+                print(f"Total estimate = {estimate} ± {err}")
+            if converged:
+                if verbosity >= 1:
+                    print(f"Repeat {len(samples)}: converged to tol {tol}")
+                break
+
+        if pbar:
+            if len(samples) < R_min:
+                estimate, err, _ = calc_stats(
+                    samples, mean_p, mean_s, tol, tol_scale
+                )
+            pbar.set_description(format_number_with_error(estimate, err))
+
+        if pbar:
+            pbar.update()
+    if pbar:
+        pbar.close()
+
+    if mpi:
+        # deal with remaining futures
+        extra_futures = []
+        for f in fs:
+            if f.done() or f.running():
+                extra_futures.append(f)
+            else:
+                f.cancel()
+
+        if extra_futures:
+            # might as well combine finished samples
+            samples.extend(f.result() for f in extra_futures)
+            estimate, err, converged = calc_stats(
+                samples, mean_p, mean_s, tol, tol_scale
+            )
+
+    if estimate is None:
+        estimate, err, _ = calc_stats(samples, mean_p, mean_s, tol, tol_scale)
+
+    if verbosity >= 1:
+        print(f"ESTIMATE is {estimate} ± {err}")
+
+    if info is not None:
+        if "samples" in info:
+            info["samples"] = samples
+        if "error" in info:
+            info["error"] = err
+        if "estimate" in info:
+            info["estimate"] = estimate
+
+    if plot:
+        info["fig"], info["axs"] = plot_approx_spectral_info(info)
+
+    return estimate
+
+
+@functools.wraps(approx_spectral_function)
+def tr_abs_approx(*args, **kwargs):
+    return approx_spectral_function(*args, f=abs, **kwargs)
+
+
+@functools.wraps(approx_spectral_function)
+def tr_exp_approx(*args, **kwargs):
+    return approx_spectral_function(*args, f=exp, **kwargs)
+
+
+@functools.wraps(approx_spectral_function)
+def tr_sqrt_approx(*args, **kwargs):
+    return approx_spectral_function(*args, f=sqrt, pos=True, **kwargs)
+
+
+def xlogx(x):
+    return x * log2(x) if x > 0 else 0.0
+
+
+@functools.wraps(approx_spectral_function)
+def tr_xlogx_approx(*args, **kwargs):
+    return approx_spectral_function(*args, f=xlogx, **kwargs)
+
+
+# --------------------------------------------------------------------------- #
+#                             Specific quantities                             #
+# --------------------------------------------------------------------------- #
+
+
+def entropy_subsys_approx(psi_ab, dims, sysa, backend=None, **kwargs):
+    """Approximate the (Von Neumann) entropy of a pure state's subsystem.
+
+    Parameters
+    ----------
+    psi_ab : ket
+        Bipartite state to partially trace and find entopy of.
+    dims : sequence of int, optional
+        The sub dimensions of ``psi_ab``.
+    sysa : int or sequence of int, optional
+        Index(es) of the 'a' subsystem(s) to keep.
+    kwargs
+        Supplied to :func:`approx_spectral_function`.
+    """
+    lo = lazy_ptr_linop(psi_ab, dims=dims, sysa=sysa, backend=backend)
+    return -tr_xlogx_approx(lo, **kwargs)
+
+
+def tr_sqrt_subsys_approx(psi_ab, dims, sysa, backend=None, **kwargs):
+    """Approximate the trace sqrt of a pure state's subsystem.
+
+    Parameters
+    ----------
+    psi_ab : ket
+        Bipartite state to partially trace and find trace sqrt of.
+    dims : sequence of int, optional
+        The sub dimensions of ``psi_ab``.
+    sysa : int or sequence of int, optional
+        Index(es) of the 'a' subsystem(s) to keep.
+    kwargs
+        Supplied to :func:`approx_spectral_function`.
+    """
+    lo = lazy_ptr_linop(psi_ab, dims=dims, sysa=sysa, backend=backend)
+    return tr_sqrt_approx(lo, **kwargs)
+
+
+def norm_ppt_subsys_approx(psi_abc, dims, sysa, sysb, backend=None, **kwargs):
+    """Estimate the norm of the partial transpose of a pure state's subsystem."""
+    lo = lazy_ptr_ppt_linop(
+        psi_abc, dims=dims, sysa=sysa, sysb=sysb, backend=backend
+    )
+    return tr_abs_approx(lo, **kwargs)
+
+
+def logneg_subsys_approx(psi_abc, dims, sysa, sysb, **kwargs):
+    """Estimate the logarithmic negativity of a pure state's subsystem.
+
+    Parameters
+    ----------
+    psi_abc : ket
+        Pure tripartite state, for which estimate the entanglement between
+        'a' and 'b'.
+    dims : sequence of int
+        The sub dimensions of ``psi_abc``.
+    sysa : int or sequence of int, optional
+        Index(es) of the 'a' subsystem(s) to keep, with respect to all
+        the dimensions, ``dims``, (i.e. pre-partial trace).
+    sysa : int or sequence of int, optional
+        Index(es) of the 'b' subsystem(s) to keep, with respect to all
+        the dimensions, ``dims``, (i.e. pre-partial trace).
+    kwargs
+        Supplied to :func:`approx_spectral_function`.
+    """
+    nrm = norm_ppt_subsys_approx(psi_abc, dims, sysa, sysb, **kwargs)
+    return max(0.0, log2(nrm))
+
+
+def negativity_subsys_approx(psi_abc, dims, sysa, sysb, **kwargs):
+    """Estimate the negativity of a pure state's subsystem.
+
+    Parameters
+    ----------
+    psi_abc : ket
+        Pure tripartite state, for which estimate the entanglement between
+        'a' and 'b'.
+    dims : sequence of int
+        The sub dimensions of ``psi_abc``.
+    sysa : int or sequence of int, optional
+        Index(es) of the 'a' subsystem(s) to keep, with respect to all
+        the dimensions, ``dims``, (i.e. pre-partial trace).
+    sysa : int or sequence of int, optional
+        Index(es) of the 'b' subsystem(s) to keep, with respect to all
+        the dimensions, ``dims``, (i.e. pre-partial trace).
+    kwargs
+        Supplied to :func:`approx_spectral_function`.
+    """
+    nrm = norm_ppt_subsys_approx(psi_abc, dims, sysa, sysb, **kwargs)
+    return max(0.0, (nrm - 1) / 2)
+
+
+def gen_bipartite_spectral_fn(exact_fn, approx_fn, pure_default):
+    """Generate a function that computes a spectral quantity of the subsystem
+    of a pure state. Automatically computes for the smaller subsystem, or
+    switches to the approximate method for large subsystems.
+
+    Parameters
+    ----------
+    exact_fn : callable
+        The function that computes the quantity on a density matrix, with
+        signature: ``exact_fn(rho_a, rank=...)``.
+    approx_fn : callable
+        The function that approximately computes the quantity using a lazy
+        representation of the whole system. With signature
+        ``approx_fn(psi_ab, dims, sysa, **approx_opts)``.
+    pure_default : float
+        The default value when the whole state is the subsystem.
+
+    Returns
+    -------
+    bipartite_spectral_fn : callable
+        The function, with signature:
+        ``(psi_ab, dims, sysa, approx_thresh=2**13, **approx_opts)``
+    """
+
+    def bipartite_spectral_fn(
+        psi_ab, dims, sysa, approx_thresh=2**13, **approx_opts
+    ):
+        sysa = int2tup(sysa)
+        sz_a = prod(d for i, d in enumerate(dims) if i in sysa)
+        sz_b = prod(dims) // sz_a
+
+        # pure state
+        if sz_b == 1:
+            return pure_default
+
+        # also check if system b is smaller, since spectrum is same for both
+        if sz_b < sz_a:
+            # if so swap things around
+            sz_a = sz_b
+            sysb = [i for i in range(len(dims)) if i not in sysa]
+            sysa = sysb
+
+        # check whether to use approx lanczos method
+        if (approx_thresh is not None) and (sz_a >= approx_thresh):
+            return approx_fn(psi_ab, dims, sysa, **approx_opts)
+
+        rho_a = ptr(psi_ab, dims, sysa)
+        return exact_fn(rho_a)
+
+    return bipartite_spectral_fn