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

trajectree/quimb/quimb/linalg/approx_spectral.py

@@ -1,1109 +0,0 @@
-"""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