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

trajectree/quimb/quimb/calc.py

@@ -1,1491 +0,0 @@
-"""Functions for more advanced calculations of quantities and properties of
-quantum objects.
-"""
-import numbers
-import itertools
-import functools
-import collections
-from math import sin, cos, pi, log, log2, sqrt
-
-import numpy as np
-import numpy.linalg as nla
-from scipy.optimize import minimize
-
-from .core import (
-    dag,
-    dop,
-    dot,
-    ensure_qarray,
-    expec,
-    eye,
-    ikron,
-    infer_size,
-    isop,
-    issparse,
-    isvec,
-    kron,
-    njit,
-    prod,
-    ptr,
-    qu,
-    realify,
-    tr,
-    zeroify,
-)
-from .linalg.base_linalg import eigh, eigvalsh, norm, sqrtm, norm_trace_dense
-from .linalg.approx_spectral import (
-    entropy_subsys_approx,
-    gen_bipartite_spectral_fn,
-    logneg_subsys_approx,
-    tr_sqrt_subsys_approx,
-)
-from .gen.operators import pauli
-from .gen.states import basis_vec, bell_state, bloch_state
-from .utils import (
-    frequencies,
-    int2tup,
-    keymap,
-)
-
-from .tensor.contraction import array_contract
-
-
-def fidelity(p1, p2, squared=False):
-    """Fidelity between two quantum states. By default, the unsquared fidelity
-    is used, equivalent in the pure state case to ``|<psi|phi>|``.
-
-    Parameters
-    ----------
-    p1 : vector or operator
-        First state.
-    p2 : vector or operator
-        Second state.
-    squared : bool, optional
-        Whether to use the squared convention or not, by default False.
-
-    Returns
-    -------
-    float
-    """
-    if isvec(p1) or isvec(p2):
-        F = expec(p1, p2).real
-        if not squared:
-            F = max(0.0, F) ** 0.5
-        return F
-
-    sqrho = sqrtm(p1)
-    F = tr(sqrtm(dot(sqrho, dot(p2, sqrho)))).real
-
-    if squared:
-        F = F**2
-
-    return F
-
-
-def purify(rho):
-    """Take state rho and purify it into a wavefunction of squared
-    dimension.
-
-    Parameters
-    ----------
-    rho : operator
-        Density operator to purify.
-
-    Returns
-    -------
-    vector :
-        The purified ket.
-    """
-    d = rho.shape[0]
-    evals, vs = eigh(rho)
-    evals = np.sqrt(np.clip(evals, 0, 1))
-    psi = np.zeros(shape=(d**2, 1), dtype=complex)
-    for i, evals in enumerate(evals.flat):
-        psi += evals * kron(vs[:, [i]], basis_vec(i, d))
-    return qu(psi)
-
-
-@ensure_qarray
-def kraus_op(rho, Ek, dims=None, where=None, check=False):
-    r"""Operate on a mixed state with a set of kraus operators:
-
-    .. math::
-
-        \sigma = \sum_k E_k \rho E_k^{\dagger}
-
-    Parameters
-    ----------
-    rho : (d, d) density matrix
-        The density operator to perform operation on.
-    Ek : (K, d, d) array or sequence of K (d, d) arrays
-        The Kraus operator(s).
-    dims : sequence of int, optional
-        The subdimensions of ``rho``.
-    where : int or sequence of int, optional
-        Which of the subsystems to apply the operation on.
-    check : bool, optional
-        Where to check ``sum_k(Ek.H @ Ek) == 1``.
-
-    Returns
-    -------
-    sigma : density matrix
-        The state after the kraus operation.
-
-    Examples
-    --------
-    The depolarising channel:
-
-    .. code-block:: python3
-
-        In [1]: import quimb as qu
-
-        In [2]: rho = qu.rand_rho(2)
-
-        In [3]: I, X, Y, Z = (qu.pauli(s) for s in 'IXYZ')
-
-        In [4]: [qu.expec(rho, A) for A in (X, Y, Z)]
-        Out[4]: [-0.652176185230622, -0.1301762132792484, 0.22362918368272583]
-
-        In [5]: p = 0.1
-
-        In [6]: Ek = [
-           ...:     (1 - p)**0.5 * I,
-           ...:     (p / 3)**0.5 * X,
-           ...:     (p / 3)**0.5 * Y,
-           ...:     (p / 3)**0.5 * Z
-           ...: ]
-
-        In [7]: sigma = qu.kraus_op(rho, Ek)
-
-        In [8]: [qu.expec(sigma, A) for A in (X, Y, Z)]
-        Out[8]: [-0.5652193605332058, -0.11281938484201527, 0.1938119591916957]
-    """
-    if not isinstance(Ek, np.ndarray):
-        Ek = np.stack(Ek, axis=0)
-
-    if check:
-        SEk = np.einsum("kij,kil", Ek.conj(), Ek)
-        if norm(SEk - eye(Ek.shape[-1]), "fro") > 1e-12:
-            raise ValueError("Did not find ``sum(E_k.H @ Ek) == 1``.")
-
-    if int(dims is None) + int(where is None) == 1:
-        raise ValueError("If `dims` is specified so should `where`.")
-
-    if isinstance(where, numbers.Integral):
-        where = (where,)
-
-    if dims:
-        # the full subdimensions of ``rho``
-        dims = tuple(dims)
-        rho = rho.reshape(dims + dims)
-        N = len(dims)
-
-        # the subdimensions of where the kraus operator should act
-        kdims = tuple(dims[i] for i in where)
-        Ek = Ek.reshape((-1,) + kdims + kdims)
-
-        rho_inds = (
-            *(f"i*{q}" if q in where else f"i{q}" for q in range(N)),
-            *(f"j*{q}" if q in where else f"j{q}" for q in range(N)),
-        )
-        Ei_inds = ("K", *(f"i{q}" for q in where), *(f"i*{q}" for q in where))
-        Ej_inds = ("K", *(f"j{q}" for q in where), *(f"j*{q}" for q in where))
-        out = (*(f"i{q}" for q in range(N)), *(f"j{q}" for q in range(N)))
-    else:
-        Ei_inds = ("K", "i", "i*")
-        rho_inds = ("i*", "j*")
-        Ej_inds = ("K", "j", "j*")
-        out = ("i", "j")
-
-    sigma = array_contract(
-        (Ek, rho, Ek.conj()),
-        (Ei_inds, rho_inds, Ej_inds),
-        out,
-    )
-
-    if dims:
-        sigma = sigma.reshape(prod(dims), prod(dims))
-
-    return sigma
-
-
-@ensure_qarray
-def projector(A, eigenvalue=1.0, tol=1e-12, autoblock=False):
-    """Compute the projector for the target ``eigenvalue`` of operator
-    ``A``.
-
-    Parameters
-    ----------
-    A : operator or tuple[array, operator].
-        The hermitian observable. If a tuple is supplied, assume this is the
-        eigendecomposition of the observable.
-    eigenvalue : float, optional
-        The target eigenvalue to construct the projector for, default: 1.
-    tol : float, optional
-        The tolerance with which to select all eigenvalues.
-    autoblock : bool, optional
-        Whether to use automatic symmetry exploitation.
-
-    Returns
-    -------
-    P : dense matrix
-        The projector onto the target eigenspace.
-    """
-    if isinstance(A, (tuple, list)):
-        el, ev = A
-    else:
-        el, ev = eigh(A, autoblock=autoblock)
-    which = np.argwhere(abs(el - eigenvalue) < tol)
-    P = np.zeros_like(ev)
-    for i in which:
-        vi = ev[:, i]
-        P += vi @ vi.H
-    return P
-
-
-def measure(p, A, eigenvalue=None, tol=1e-12):
-    r"""Measure state ``p`` with observable ``A``, collapsing the state in the
-    process and returning the relevant eigenvalue.
-
-    .. math::
-
-        \left| \psi \right\rangle \rightarrow
-        \frac{
-            P_i \left| \psi \right\rangle
-        }{
-            \sqrt{\left\langle \psi \right| P_i \left| \psi \right\rangle}
-        }
-
-    and
-
-    .. math::
-
-        \rho \rightarrow
-        \frac{
-            P_i \rho P_i^{\dagger}
-        }{
-            \text{Tr} \left[ P_i \rho \right]
-        }
-
-    along with the corresponding eigenvalue :math:`\lambda_i`.
-
-    Parameters
-    ----------
-    p : vector or matrix
-        The quantum state to measure.
-    A : matrix or tuple[array, matrix]
-        The hermitian operator corresponding to an observable. You can supply
-        also a pre-diagonalized operator here as a tuple of eigenvalues and
-        eigenvectors.
-    tol : float, optional
-        The tolerance within which to group eigenspaces.
-    eigenvalue : float, optional
-        If specified, deterministically collapse to this result. Otherwise
-        randomly choose a result as in 'real-life'.
-
-    Returns
-    -------
-    result : float
-        The result of the measurement.
-    p_after : vector or matrix
-        The quantum state post measurement.
-    """
-    if isinstance(A, (tuple, list)):
-        el, ev = A
-    else:
-        el, ev = eigh(A)
-
-    js = np.arange(el.size)
-
-    # compute prob of each eigenvector
-    if isvec(p):
-        pj = (abs(ev.H @ p) ** 2).flatten()
-    else:
-        pj = array_contract(
-            (ev.H, p, ev),
-            ("jk", "kl", "lj"),
-            ("j",),
-        ).real
-
-    # then choose one
-    if eigenvalue is None:
-        j = np.random.choice(js, p=pj)
-        eigenvalue = el[j]
-
-    # now combine whole eigenspace
-    P = projector((el, ev), eigenvalue=eigenvalue, tol=tol)
-    total_prob = np.sum(pj[abs(el - eigenvalue) < tol])
-
-    # now collapse the state
-    if isvec(p):
-        p_after = P @ (p / total_prob**0.5)
-    else:
-        p_after = (P @ p @ P.H) / total_prob
-
-    return eigenvalue, p_after
-
-
-def simulate_counts(p, C, phys_dim=2, seed=None):
-    """Simulate measuring each qubit of ``p`` in the computational basis,
-    producing output like that of ``qiskit``.
-
-    Parameters
-    ----------
-    p : vector or operator
-        The quantum state, assumed to be normalized, as either a ket or density
-        operator.
-    C : int
-        The number of counts to perform.
-    phys_dim : int, optional
-        The assumed size of the subsystems of ``p``, defaults to 2 for qubits.
-
-    Returns
-    -------
-    results : dict[str, int]
-        The counts for each bit string measured.
-
-    Examples
-    --------
-
-    Simulate measuring the state of each qubit in a GHZ-state:
-
-    .. code:: python3
-
-        >>> import quimb as qu
-        >>> psi = qu.ghz_state(3)
-        >>> qu.simulate_counts(psi, 1024)
-        {'000': 514, '111': 510}
-
-    """
-    rng = np.random.default_rng(seed)
-
-    n = infer_size(p, phys_dim)
-    d = phys_dim**n
-
-    if isop(p):
-        pi = np.diag(p).real
-    else:
-        pi = np.multiply(np.conj(p), p).real
-
-    # probability of each basis state
-    pi = pi.reshape(-1)
-
-    # raw counts in terms of integers
-    raw_counts = rng.choice(d, size=C, p=pi)
-
-    # convert to frequencies of binary
-    bin_str = "{:0>" + str(n) + "b}"
-    results = keymap(bin_str.format, frequencies(raw_counts))
-
-    return results
-
-
-def dephase(rho, p, rand_rank=None):
-    """Dephase ``rho`` by amount ``p``, that is, mix it
-    with the maximally mixed state:
-
-        rho -> (1 - p) * rho + p * I / d
-
-    Parameters
-    ----------
-    rho : operator
-        The state.
-    p : float
-        The final proportion of identity.
-    rand_rank : int or float, optional
-        If given, dephase with a random diagonal operator with this many
-        non-zero entries. If float, proportion of full size.
-
-    Returns
-    -------
-    rho_dephase : operator
-        The dephased density operator.
-    """
-    d = rho.shape[0]
-
-    if (rand_rank is None) or (rand_rank == d) or (rand_rank == 1.0):
-        dephaser = eye(d) / d
-
-    else:
-        if not isinstance(rand_rank, numbers.Integral):
-            rand_rank = int(rand_rank * d)
-        rand_rank = min(max(1, rand_rank), d)
-
-        dephaser = np.zeros((d, d))
-        dephaser_diag = np.einsum("aa->a", dephaser)
-        nnz = np.random.choice(np.arange(d), size=rand_rank, replace=False)
-        dephaser_diag[nnz] = 1 / rand_rank
-
-    return (1 - p) * rho + p * dephaser
-
-
-@zeroify
-def entropy(a, rank=None):
-    """Compute the (von Neumann) entropy.
-
-    Parameters
-    ----------
-    a : operator or 1d array
-        Positive operator or list of positive eigenvalues.
-    rank : int (optional)
-        If operator has known rank, then a partial decomposition can be
-        used to accelerate the calculation.
-
-    Returns
-    -------
-    float
-        The von Neumann entropy.
-
-    See Also
-    --------
-    mutinf, entropy_subsys, entropy_subsys_approx
-    """
-    a = np.asarray(a)
-    if np.ndim(a) == 1:
-        evals = a
-    else:
-        if rank is None:
-            evals = eigvalsh(a)
-        else:  # know that not all eigenvalues needed
-            evals = eigvalsh(a, k=rank, which="LM", backend="AUTO")
-
-    evals = evals[evals > 0.0]
-    return np.sum(-evals * np.log2(evals))
-
-
-entropy_subsys = gen_bipartite_spectral_fn(entropy, entropy_subsys_approx, 0.0)
-"""Calculate the entropy of a pure states' subsystem, optionally switching
-to an approximate lanczos method when the subsystem is very large.
-
-Parameters
-----------
-psi_ab : vector
-    Bipartite state.
-dims : sequence of int
-    The sub-dimensions of the state.
-sysa :  sequence of int
-    The indices of which dimensions to calculate the entropy for.
-approx_thresh : int, optional
-    The size of sysa at which to switch to the approx method. Set to
-    ``None`` to never use the approximation.
-**approx_opts
-    Supplied to :func:`entropy_subsys_approx`, if used.
-
-Returns
--------
-float
-    The subsytem entropy.
-
-See Also
---------
-entropy, entropy_subsys_approx, mutinf_subsys
-"""
-
-
-@zeroify
-def mutinf(p, dims=(2, 2), sysa=0, rank=None):
-    """Find the mutual information for a bipartition of a state.
-
-    That is, ``H(A) + H(B) - H(AB)``, for von Neumann entropy ``H``, and two
-    subsystems A and B.
-
-    Parameters
-    ----------
-    p : vector or operator
-        State, can be vector or operator.
-    dims : tuple(int), optional
-        Internal dimensions of state.
-    sysa : int, optional
-        Index of first subsystem, A.
-    sysb : int, optional
-        Index of second subsystem, B.
-    rank : int, optional
-        If known, the rank of rho_ab, to speed calculation of ``H(AB)`` up.
-        For example, if ``p`` comes from tracing out three qubits from a
-        system, then its rank is 2^3 = 8 etc.
-
-    Returns
-    -------
-    float
-
-    See Also
-    --------
-    entropy, mutinf_subsys, entropy_subsys_approx
-    """
-    sysa = int2tup(sysa)
-
-    # mixed combined system
-    if isop(p):
-        # total
-        hab = entropy(p, rank=rank)
-
-        # subsystem a
-        rhoa = ptr(p, dims, sysa)
-        ha = entropy(rhoa)
-
-        # need subsystem b as well
-        sysb = tuple(i for i in range(len(dims)) if i not in sysa)
-        rhob = ptr(p, dims, sysb)
-        hb = entropy(rhob)
-
-    # pure combined system
-    else:
-        hab = 0.0
-        ha = hb = entropy_subsys(p, dims, sysa)
-
-    return ha + hb - hab
-
-
-mutual_information = mutinf
-
-
-def check_dims_and_indices(dims, *syss):
-    """Make sure all indices found in the tuples ``syss`` are in
-    ``range(len(dims))``.
-    """
-    nsys = len(dims)
-    all_sys = sum(syss, ())
-
-    if not all(0 <= i < nsys for i in all_sys):
-        raise ValueError(
-            "Indices specified in `sysa` and `sysb` must be "
-            f"in range({nsys}) for dims {dims}."
-        )
-
-
-def mutinf_subsys(
-    psi_abc, dims, sysa, sysb, approx_thresh=2**13, **approx_opts
-):
-    """Calculate the mutual information of two subsystems of a pure state,
-    possibly using an approximate lanczos method for large subsytems.
-
-    Parameters
-    ----------
-    psi_abc : vector
-        Tri-partite pure state.
-    dims : sequence of int
-        The sub dimensions of the state.
-    sysa : sequence of int
-        The index(es) of the subsystem(s) to consider part of 'A'.
-    sysb : sequence of int
-        The index(es) of the subsystem(s) to consider part of 'B'.
-    approx_thresh : int, optional
-        The size of subsystem at which to switch to the approximate lanczos
-        method. Set to ``None`` to never use the approximation.
-    approx_opts
-        Supplied to :func:`entropy_subsys_approx`, if used.
-
-    Returns
-    -------
-    float
-        The mutual information.
-
-    See Also
-    --------
-    mutinf, entropy_subsys, entropy_subsys_approx, logneg_subsys
-    """
-    sysa, sysb = int2tup(sysa), int2tup(sysb)
-
-    check_dims_and_indices(dims, sysa, sysb)
-
-    sz_a = prod(d for i, d in enumerate(dims) if i in sysa)
-    sz_b = prod(d for i, d in enumerate(dims) if i in sysb)
-    sz_c = prod(dims) // (sz_a * sz_b)
-
-    kws = {"approx_thresh": approx_thresh, **approx_opts}
-
-    if sz_c == 1:
-        hab = 0.0
-        ha = hb = entropy_subsys(psi_abc, dims, sysa, **kws)
-    else:
-        hab = entropy_subsys(psi_abc, dims, sysa + sysb, **kws)
-        ha = entropy_subsys(psi_abc, dims, sysa, **kws)
-        hb = entropy_subsys(psi_abc, dims, sysb, **kws)
-
-    return hb + ha - hab
-
-
-def schmidt_gap(psi_ab, dims, sysa):
-    """Find the schmidt gap of the bipartition of ``psi_ab``. That is, the
-    difference between the two largest eigenvalues of the reduced density
-    operator.
-
-    Parameters
-    ----------
-    psi_ab : vector
-        Bipartite state.
-    dims : sequence of int
-        The sub-dimensions of the state.
-    sysa :  sequence of int
-        The indices of which dimensions to calculate the entropy for.
-
-    Returns
-    -------
-    float
-    """
-    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 1.0
-
-    # also check if system b is smaller, since spectrum is same for both
-    if sz_b < sz_a:
-        # if so swap things around
-        sysb = [i for i in range(len(dims)) if i not in sysa]
-        sysa = sysb
-
-    rho_a = ptr(psi_ab, dims, sysa)
-    el = eigvalsh(rho_a, k=2, which="LM")
-    return abs(el[0] - el[1])
-
-
-def tr_sqrt(A, rank=None):
-    """Return the trace of the sqrt of a positive semidefinite operator."""
-    if rank is None:
-        el = eigvalsh(A, sort=False)
-    else:
-        el = eigvalsh(A, k=rank, which="LM", backend="AUTO")
-    return np.sum(np.sqrt(el[el > 0.0]))
-
-
-tr_sqrt_subsys = gen_bipartite_spectral_fn(tr_sqrt, tr_sqrt_subsys_approx, 1.0)
-"""Compute the trace sqrt of a subsystem, possibly using an approximate
-lanczos method when the subsytem is big.
-
-Parameters
-----------
-psi_ab : vector
-    Bipartite state.
-dims : sequence of int
-    The sub-dimensions of the state.
-sysa :  sequence of int
-    The indices of which dimensions to calculate the trace sqrt for.
-approx_thresh : int, optional
-    The size of sysa at which to switch to the approx method. Set to
-    ``None`` to never use the approximation.
-**approx_opts
-    Supplied to :func:`tr_sqrt_subsys_approx`, if used.
-
-Returns
--------
-float
-    The subsytem entropy.
-
-See Also
---------
-tr_sqrt, tr_sqrt_subsys_approx, partial_transpose_norm
-"""
-
-
-@ensure_qarray
-def partial_transpose(p, dims=(2, 2), sysa=0):
-    """Partial transpose of a density operator.
-
-    Parameters
-    ----------
-    p : operator or vector
-        The state to partially transpose.
-    dims : tuple(int), optional
-        The internal dimensions of the state.
-    sysa : sequence of int
-        The indices of 'system A', everything else assumed to be 'system B'.
-
-    Returns
-    -------
-    operator
-
-    See Also
-    --------
-    logneg, negativity
-    """
-    sysa = int2tup(sysa)
-
-    ndims = len(dims)
-    perm_ket_inds = []
-    perm_bra_inds = []
-
-    for i in range(ndims):
-        if i in sysa:
-            perm_ket_inds.append(i + ndims)
-            perm_bra_inds.append(i)
-        else:
-            perm_ket_inds.append(i)
-            perm_bra_inds.append(i + ndims)
-
-    return (
-        np.asarray(qu(p, "dop"))
-        .reshape((*dims, *dims))
-        .transpose((*perm_ket_inds, *perm_bra_inds))
-        .reshape((prod(dims), prod(dims)))
-    )
-
-
-def partial_transpose_norm(p, dims, sysa):
-    """Compute the norm of the partial transpose for (log)-negativity,
-    taking a shortcut (trace sqrt of reduced subsytem), when system is a
-    vector.
-    """
-    sysa = int2tup(sysa)
-
-    # check for pure bipartition -> easier to calc
-    if isvec(p):
-        sz_a = prod(d for i, d in enumerate(dims) if i in sysa)
-        sz_b = prod(dims) // sz_a
-
-        # check if system b is smaller, since entropy is same for both a & b.
-        if sz_b < sz_a:
-            # if so swap things around
-            sysb = [i for i in range(len(dims)) if i not in sysa]
-            sysa = sysb
-
-        rhoa = ptr(p, dims, sysa)
-        return tr_sqrt(rhoa) ** 2
-
-    return norm_trace_dense(partial_transpose(p, dims, sysa), isherm=True)
-
-
-@zeroify
-def logneg(p, dims=(2, 2), sysa=0):
-    """Compute logarithmic negativity between two subsytems.
-    This is defined as  log_2( | rho_{AB}^{T_B} | ). This only handles
-    bipartitions (including pure states efficiently), and will not trace
-    anything out.
-
-    Parameters
-    ----------
-    p : ket vector or density operator
-        State to compute logarithmic negativity for.
-    dims : tuple(int), optional
-        The internal dimensions of ``p``.
-    sysa : int, optional
-        Index of the first subsystem, A, relative to ``dims``.
-
-    Returns
-    -------
-    float
-
-    See Also
-    --------
-    negativity, partial_transpose, logneg_subsys_approx
-    """
-    return max(0.0, log2(partial_transpose_norm(p, dims, sysa)))
-
-
-logarithmic_negativity = logneg
-
-
-def logneg_subsys(
-    psi_abc, dims, sysa, sysb, approx_thresh=2**13, **approx_opts
-):
-    """Compute the logarithmic negativity between two subsystems of a pure
-    state, possibly using an approximate lanczos for large subsystems. Uses
-    a special method if the two subsystems form a bipartition of the state.
-
-    Parameters
-    ----------
-    psi_abc : vector
-        Tri-partite pure state.
-    dims : sequence of int
-        The sub dimensions of the state.
-    sysa : sequence of int
-        The index(es) of the subsystem(s) to consider part of 'A'.
-    sysb : sequence of int
-        The index(es) of the subsystem(s) to consider part of 'B'.
-    approx_thresh : int, optional
-        The size of subsystem at which to switch to the approximate lanczos
-        method. Set to ``None`` to never use the approximation.
-    approx_opts
-        Supplied to :func:`~quimb.logneg_subsys_approx`, if used.
-
-    Returns
-    -------
-    float
-        The logarithmic negativity.
-
-    See Also
-    --------
-    logneg, mutinf_subsys, logneg_subsys_approx
-    """
-    sysa, sysb = int2tup(sysa), int2tup(sysb)
-
-    check_dims_and_indices(dims, sysa, sysb)
-
-    sz_a = prod(d for i, d in enumerate(dims) if i in sysa)
-    sz_b = prod(d for i, d in enumerate(dims) if i in sysb)
-    sz_ab = sz_a * sz_b
-    sz_c = prod(dims) // sz_ab
-
-    # check for pure bipartition
-    if sz_c == 1:
-        psi_ab_ppt_norm = (
-            tr_sqrt_subsys(
-                psi_abc, dims, sysa, approx_thresh=approx_thresh, **approx_opts
-            )
-            ** 2
-        )
-        return max(log2(psi_ab_ppt_norm), 0.0)
-
-    # check whether to use approx lanczos method
-    if (approx_thresh is not None) and (sz_ab >= approx_thresh):
-        return logneg_subsys_approx(psi_abc, dims, sysa, sysb, **approx_opts)
-
-    rho_ab = ptr(psi_abc, dims, sysa + sysb)
-
-    # need to adjust for new dimensions and indices
-    new_dims, new_sysa = [], []
-    new_inds = iter(range(len(dims)))
-
-    for i, d in enumerate(dims):
-        if i in sysa:
-            new_dims.append(d)
-            new_sysa.append(next(new_inds))
-        elif i in sysb:
-            new_dims.append(d)
-            next(new_inds)  # don't need sysb
-
-    return logneg(rho_ab, new_dims, new_sysa)
-
-
-def negativity(p, dims=(2, 2), sysa=0):
-    """Compute negativity between two subsytems.
-
-    This is defined as  (| rho_{AB}^{T_B} | - 1) / 2. If ``len(dims) > 2``,
-    then the non-target dimensions will be traced out first.
-
-    Parameters
-    ----------
-    p : ket vector or density operator
-        State to compute logarithmic negativity for.
-    dims : tuple(int), optional
-        The internal dimensions of ``p``.
-    sysa : int, optional
-        Index of the first subsystem, A, relative to ``dims``.
-
-    Returns
-    -------
-    float
-
-    See Also
-    --------
-    logneg, partial_transpose, negativity_subsys_approx
-    """
-    return max(0.0, (partial_transpose_norm(p, dims, sysa) - 1) / 2)
-
-
-@zeroify
-def concurrence(p, dims=(2, 2), sysa=0, sysb=1):
-    """Concurrence of two-qubit state.
-
-    If ``len(dims) > 2``, then the non-target dimensions will be traced out
-    first.
-
-    Parameters
-    ----------
-    p : ket vector or density operator
-        State to compute concurrence for.
-    dims : tuple(int), optional
-        The internal dimensions of ``p``.
-    sysa : int, optional
-        Index of the first subsystem, A, relative to ``dims``.
-    sysb : int, optional
-        Index of the first subsystem, B, relative to ``dims``.
-
-    Returns
-    -------
-    float
-    """
-    if len(dims) > 2:
-        p = ptr(p, dims, (sysa, sysb))
-
-    Y = pauli("Y")
-
-    if isop(p):
-        pt = dot(kron(Y, Y), dot(p.conj(), kron(Y, Y)))
-        evals = (nla.eigvals(dot(p, pt)).real ** 2) ** 0.25
-        return max(0, 2 * np.max(evals) - np.sum(evals))
-    else:
-        pt = dot(kron(Y, Y), p.conj())
-        c = np.real(abs(dot(dag(p), pt))).item(0)
-        return max(0, c)
-
-
-def one_way_classical_information(p_ab, prjs, precomp_func=False):
-    """One way classical information for two qubit density operator.
-
-    Parameters
-    ----------
-    p_ab : operator
-        State of two qubits
-    prjs : sequence of matrices
-       The POVMs.
-    precomp_func : bool, optional
-        Whether to return a pre-computed function, closed over the actual
-        state.
-
-    Returns
-    -------
-    float or callable
-        The one-way classical information or the function to compute it for
-        the given state which takes a set of POVMs as its single argument.
-    """
-    p_a = ptr(p_ab, (2, 2), 0)
-    s_a = entropy(p_a)
-
-    def owci(prjs):
-        def gen_paj():
-            for prj in prjs:
-                p_ab_j = dot((eye(2) & prj), p_ab)
-                prob = tr(p_ab_j)
-                p_a_j = ptr(p_ab_j, (2, 2), 0) / prob
-                yield prob, p_a_j
-
-        return s_a - sum(p * entropy(rho) for p, rho in gen_paj())
-
-    return owci if precomp_func else owci(prjs)
-
-
-@zeroify
-def quantum_discord(
-    p,
-    dims=(2, 2),
-    sysa=0,
-    sysb=1,
-    method="COBYLA",
-    tol=1e-12,
-    maxiter=2**14,
-):
-    """Quantum Discord for two qubit density operator.
-
-    If ``len(dims) > 2``, then the non-target dimensions will be traced out
-    first.
-
-    Parameters
-    ----------
-    p : ket vector or density operator
-        State to compute quantum discord for.
-    dims : tuple(int), optional
-        The internal dimensions of ``p``.
-    sysa : int, optional
-        Index of the first subsystem, A, relative to ``dims``.
-    sysb : int, optional
-        Index of the first subsystem, B, relative to ``dims``.
-
-    Returns
-    -------
-    float
-    """
-    if len(dims) > 2:
-        p = ptr(p, dims, (sysa, sysb))
-    else:
-        p = qu(p, "dop")
-    iab = mutual_information(p)
-    owci = one_way_classical_information(p, None, precomp_func=True)
-
-    def trial_qd(a):
-        ax, ay, az = sin(a[0]) * cos(a[1]), sin(a[0]) * sin(a[1]), cos(a[0])
-        prja = bloch_state(ax, ay, az)
-        prjb = eye(2) - prja
-        return iab - owci((prja, prjb))
-
-    opt = minimize(
-        trial_qd,
-        (pi / 2, pi),
-        method=method,
-        bounds=((0, pi), (0, 2 * pi)),
-        tol=tol,
-        options=dict(
-            maxiter=maxiter,
-        ),
-    )
-    if opt.success:
-        return opt.fun
-    else:  # pragma: no cover
-        raise ValueError(opt.message)
-
-
-@zeroify
-def trace_distance(p1, p2):
-    r"""Trace distance between two states:
-
-    .. math::
-
-        \delta(\rho, \sigma)
-        =
-        \frac{1}{2} \left| \rho - \sigma \right|_\mathrm{tr}
-
-    If two wavefunctions are supplied the trace distance will be computed via
-    the more efficient expression:
-
-    .. math::
-
-        \delta(|\psi\rangle\langle\psi|, |\phi\rangle\langle\phi|)
-        =
-        \sqrt{1 - \langle \psi | \phi \rangle^2}
-
-    Parameters
-    ----------
-    p1 : ket or density operator
-        The first state.
-    p2 : ket or density operator
-        The second state.
-
-    Returns
-    -------
-    float
-    """
-    p1_is_op, p2_is_op = isop(p1), isop(p2)
-
-    # If both are pure kets then special case
-    if (not p1_is_op) and (not p2_is_op):
-        return sqrt(1 - expec(p1, p2))
-
-    # Otherwise do full calculation
-    return 0.5 * norm(
-        (p1 if p1_is_op else dop(p1)) - (p2 if p2_is_op else dop(p2)), "tr"
-    )
-
-
-def cprint(psi, prec=6):
-    """Print a state in the computational basis.
-
-    Parameters
-    ----------
-    psi : vector
-        The pure state.
-    prec : int, optional
-        How many significant figures to print.
-
-    Examples
-    --------
-
-        >>> cprint(rand_ket(2**2))
-        (-0.0751069-0.192635j) |00>
-          (0.156837-0.562213j) |01>
-        (-0.307381+0.0291168j) |10>
-          (-0.14302+0.707661j) |11>
-
-        >>> cprint(w_state(4))
-        (0.5+0j) |0001>
-        (0.5+0j) |0010>
-        (0.5+0j) |0100>
-        (0.5+0j) |1000>
-    """
-    d = psi.shape[0]
-    n = int(log2(d))
-    if 2**n != d:
-        raise ValueError("State is not factorizable into computational basis.")
-
-    coeff_str = "{:." + str(prec) + "}"
-
-    coeffs, cs = [], []
-    for c in itertools.product("01", repeat=n):
-        c = "".join(c)
-        ix = int(c, 2)
-        coeff = psi.item(ix)
-        if coeff != 0.0:
-            cs.append(c)
-            coeffs.append(coeff_str.format(coeff))
-
-    max_str_len = max(map(len, coeffs))
-    full_str = "{:>" + str(max_str_len) + "} |{}>"
-    for coeff, c in zip(coeffs, cs):
-        print(full_str.format(coeff, c))
-
-
-def decomp(a, fn, fn_args, fn_d, nmlz_func, mode="p", tol=1e-3):
-    """Decomposes an operator via the Hilbert-schmidt inner product.
-
-    Can both print the decomposition or return it.
-
-    Parameters
-    ----------
-    a : ket or density operator
-        Operator to decompose.
-    fn : callable
-        Function to generate operator/state to decompose with.
-    fn_args :
-        Sequence of args whose permutations will be supplied to ``fn``.
-    fn_d : int
-        The dimension of the operators that `fn` produces.
-    nmlz_func : callable
-        Function to produce a normlization coefficient given the ``n``
-        permutations of operators that will be produced.
-    mode :
-        String, include ``'p'`` to print the decomp and/or ``'c'`` to
-        return OrderedDict, sorted by size of contribution.
-    tol :
-        Print operators with contirbution above ``tol`` only.
-
-    Returns
-    -------
-    None or OrderedDict:
-        Pauli operator name and expec with ``a``.
-
-    See Also
-    --------
-    pauli_decomp, bell_decomp
-    """
-    if isvec(a):
-        a = qu(a, "dop")  # make sure operator
-    n = infer_size(a, base=fn_d)
-
-    # define generator for inner product to iterate over efficiently
-    def calc_name_and_overlap():
-        for perm in itertools.product(fn_args, repeat=n):
-            op = kron(*(fn(x, sparse=True) for x in perm)) * nmlz_func(n)
-            cff = expec(a, op)
-            yield "".join(str(x) for x in perm), cff
-
-    names_cffs = list(calc_name_and_overlap())
-    # sort by descending expec and turn into OrderedDict
-    names_cffs.sort(key=lambda pair: -abs(pair[1]))
-    names_cffs = collections.OrderedDict(names_cffs)
-    # Print decomposition
-    if "p" in mode:
-        for x, cff in names_cffs.items():
-            if abs(cff) < 0.01:
-                break
-            dps = int(round(0.5 - np.log10(1.001 * tol)))  # decimal places
-            print(x, "{: .{prec}f}".format(cff, prec=dps))
-    # Return full calculation
-    if "c" in mode:
-        return names_cffs
-
-
-pauli_decomp = functools.partial(
-    decomp, fn=pauli, fn_args="IXYZ", fn_d=2, nmlz_func=lambda n: 2**-n
-)
-"""Decompose an operator into Paulis."""
-
-bell_decomp = functools.partial(
-    decomp, fn=bell_state, fn_args=(0, 1, 2, 3), fn_d=4, nmlz_func=lambda x: 1
-)
-"""Decompose an operator into bell-states."""
-
-
-def correlation(
-    p, A, B, sysa, sysb, dims=None, sparse=None, precomp_func=False
-):
-    """Calculate the correlation between two sites given two operators.
-
-    Parameters
-    ----------
-    p : ket or density operator
-        State to compute correlations for, ignored if ``precomp_func=True``.
-    A : operator
-        Operator to act on first subsystem.
-    B : operator
-        Operator to act on second subsystem.
-    sysa : int
-        Index of first subsystem.
-    sysb : int
-        Index of second subsystem.
-    dims : tuple of int, optional
-        Internal dimensions of ``p``, will be assumed to be qubits if not
-        given.
-    sparse : bool, optional
-        Whether to compute with sparse operators.
-    precomp_func : bool, optional
-        Whether to return result or single arg function closed
-        over precomputed operator.
-
-    Returns
-    -------
-    float or callable
-        The correlation, <ab> - <a><b>, or a function to compute for an
-        arbitrary state.
-    """
-    if dims is None:
-        sz_p = infer_size(p)
-        dims = (2,) * sz_p
-    if sparse is None:
-        sparse = issparse(A) or issparse(B)
-
-    opts = {
-        "sparse": sparse,
-        "coo_build": sparse,
-        "stype": "csr" if sparse else None,
-    }
-    opab = ikron((A, B), dims, (sysa, sysb), **opts)
-    A = ikron((A,), dims, sysa, **opts)
-    B = ikron((B,), dims, sysb, **opts)
-
-    @realify
-    def corr(state):
-        return expec(opab, state) - expec(A, state) * expec(B, state)
-
-    return corr if precomp_func else corr(p)
-
-
-def pauli_correlations(
-    p, ss=("xx", "yy", "zz"), sysa=0, sysb=1, sum_abs=False, precomp_func=False
-):
-    """Calculate the correlation between sites for a list of operator pairs
-    choisen from the pauli matrices.
-
-    Parameters
-    ----------
-    p : ket or density operator
-        State to compute correlations for. Ignored if ``precomp_func=True``.
-    ss : tuple or str
-        List of pairs specifiying pauli matrices.
-    sysa : int, optional
-        Index of first site.
-    sysb : int, optional
-        Index of second site.
-    sum_abs : bool, optional
-        Whether to sum over the absolute values of each correlation
-    precomp_func : bool, optional
-        whether to return the values or a single argument
-        function closed over precomputed operators etc.
-
-    Returns
-    -------
-    list of float, list of callable, float or callable
-        Either the value(s) of each correlation or the function(s) to compute
-        the correlations for an arbitrary state, depending on ``sum_abs`` and
-        ``precomp_func``.
-    """
-
-    def gen_corr_list():
-        for s1, s2 in ss:
-            yield correlation(
-                p, pauli(s1), pauli(s2), sysa, sysb, precomp_func=precomp_func
-            )
-
-    if sum_abs:
-        if precomp_func:
-            return lambda p: sum((abs(corr(p)) for corr in gen_corr_list()))
-
-        return sum((abs(corr) for corr in gen_corr_list()))
-
-    return tuple(gen_corr_list())
-
-
-def ent_cross_matrix(
-    p, sz_blc=1, ent_fn=logneg, calc_self_ent=True, upscale=False
-):
-    """Calculate the pair-wise function ent_fn  between all sites or blocks
-    of a state.
-
-    Parameters
-    ----------
-    p : ket or density operator
-        State.
-    sz_blc : int
-        Size of the blocks to partition the state into. If the number of
-        individual sites is not a multiple of this then the final (smaller)
-        block will be ignored.
-    ent_fn : callable
-        Bipartite function, notionally entanglement
-    calc_self_ent : bool
-        Whether to calculate the function for each site
-        alone, purified. If the whole state is pure then this is the
-        entanglement with the whole remaining system.
-    upscale : bool, optional
-        Whether, if sz_blc != 1, to upscale the results so that the output
-        array is the same size as if it was.
-
-    Returns
-    -------
-    2D-array
-        array of pairwise ent_fn results.
-    """
-
-    sz_p = infer_size(p)
-    dims = (2,) * sz_p
-    n = sz_p // sz_blc
-    ents = np.empty((n, n))
-
-    ispure = isvec(p)
-    if ispure and sz_blc * 2 == sz_p:  # pure bipartition
-        ent = ent_fn(p, dims=(2**sz_blc, 2**sz_blc)) / sz_blc
-        ents[:, :] = ent
-        if not calc_self_ent:
-            for i in range(n):
-                ents[i, i] = np.nan
-
-    else:
-        # Range over pairwise blocks
-        for i in range(0, sz_p - sz_blc + 1, sz_blc):
-            for j in range(i, sz_p - sz_blc + 1, sz_blc):
-                if i == j:
-                    if calc_self_ent:
-                        rhoa = ptr(p, dims, [i + b for b in range(sz_blc)])
-                        psiap = purify(rhoa)
-                        ent = (
-                            ent_fn(psiap, dims=(2**sz_blc, 2**sz_blc)) / sz_blc
-                        )
-                    else:
-                        ent = np.nan
-                else:
-                    rhoab = ptr(
-                        p,
-                        dims,
-                        [i + b for b in range(sz_blc)]
-                        + [j + b for b in range(sz_blc)],
-                    )
-                    ent = ent_fn(rhoab, dims=(2**sz_blc, 2**sz_blc)) / sz_blc
-                ents[i // sz_blc, j // sz_blc] = ent
-                ents[j // sz_blc, i // sz_blc] = ent
-
-    if upscale:
-        up_ents = np.tile(np.nan, (sz_p, sz_p))
-
-        for i in range(n):
-            for j in range(i, n):
-                up_ents[
-                    i * sz_blc : (i + 1) * sz_blc,
-                    j * sz_blc : (j + 1) * sz_blc,
-                ] = ents[i, j]
-                up_ents[
-                    j * sz_blc : (j + 1) * sz_blc,
-                    i * sz_blc : (i + 1) * sz_blc,
-                ] = ents[j, i]
-
-        ents = up_ents
-
-    return ents
-
-
-def qid(
-    p,
-    dims,
-    inds,
-    precomp_func=False,
-    sparse_comp=True,
-    norm_func=norm,
-    power=2,
-    coeff=1,
-):
-    # Check inputs
-    inds = (inds,) if isinstance(inds, numbers.Number) else inds
-
-    # Construct operators
-    ops_i = tuple(
-        tuple(ikron(pauli(s), dims, ind, sparse=sparse_comp) for s in "xyz")
-        for ind in inds
-    )
-
-    # Define function closed over precomputed operators
-    def qid_func(x):
-        if isvec(x):
-            x = dop(x)
-        return tuple(
-            sum(
-                coeff * norm_func(dot(x, op) - dot(op, x)) ** power
-                for op in ops
-            )
-            for ops in ops_i
-        )
-
-    return qid_func if precomp_func else qid_func(p)
-
-
-def is_degenerate(op, tol=1e-12):
-    """Check if operator has any degenerate eigenvalues, determined relative
-    to mean spacing of all eigenvalues.
-
-    Parameters
-    ----------
-    op : operator or 1d-array
-        Operator or assumed eigenvalues to check degeneracy for.
-    tol : float
-        How much closer than evenly spaced the eigenvalue gap has to be
-        to count as degenerate.
-
-    Returns
-    -------
-    n_dgen : int
-        Number of degenerate eigenvalues.
-    """
-    op = np.asarray(op)
-    if op.ndim != 1:
-        evals = eigvalsh(op)
-    else:
-        evals = op
-    l_gaps = evals[1:] - evals[:-1]
-    l_tol = tol * (evals[-1] - evals[0]) / op.shape[0]
-    return np.count_nonzero(abs(l_gaps) < l_tol)
-
-
-def is_eigenvector(x, A, tol=1e-14):
-    """Determines whether a vector is an eigenvector of an operator.
-
-    Parameters
-    ----------
-    x : vector
-        Vector to check.
-    A : operator
-        Matrix to check.
-    tol : float, optional
-        The variance must be smaller than this value.
-
-    Returns
-    -------
-    bool
-        Whether ``A @ x = l * x`` for some scalar ``l``.
-    """
-    mat_vec = dot(A, x)
-    E = expec(x, mat_vec)
-    E2 = expec(x, dot(A, mat_vec))
-    return abs(E**2 - E2) < tol
-
-
-@njit
-def page_entropy(sz_subsys, sz_total):  # pragma: no cover
-    """Calculate the page entropy, i.e. expected entropy for a subsytem
-    of a random state in Hilbert space.
-
-    Parameters
-    ----------
-    sz_subsys : int
-        Dimension of subsystem.
-    sz_total : int
-        Dimension of total system.
-
-    Returns
-    -------
-    s : float
-        Entropy in bits.
-    """
-    if sz_subsys > sz_total**0.5:
-        sz_subsys = sz_total // sz_subsys
-
-    n = sz_total // sz_subsys
-
-    s = 0
-    for k in range(n + 1, sz_total + 1):
-        s += 1 / k
-    s -= (sz_subsys - 1) / (2 * n)
-
-    # Normalize into bits of entropy
-    return s / log(2)
-
-
-def heisenberg_energy(L):
-    """Get the analytic isotropic heisenberg chain ground energy for length L.
-    Useful for testing. Assumes the heisenberg model is defined with spin
-    operators not pauli matrices (overall factor of 2 smaller). Taken from [1].
-
-    [1] Nickel, Bernie. "Scaling corrections to the ground state energy
-    of the spin-½ isotropic anti-ferromagnetic Heisenberg chain." Journal of
-    Physics Communications 1.5 (2017): 055021
-
-    Parameters
-    ----------
-    L : int
-        The length of the chain.
-
-    Returns
-    -------
-    energy : float
-        The ground state energy.
-    """
-    Einf = (0.5 - 2 * log(2)) * L
-    Efinite = pi**2 / (6 * L)
-    correction = 1 + 0.375 / log(L) ** 3
-    return (Einf - Efinite * correction) / 2