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

trajectree/quimb/quimb/gen/rand.py

@@ -0,0 +1,713 @@
+"""Functions for generating random quantum objects and states."""
+
+import math
+import random
+from functools import wraps
+from numbers import Integral
+from itertools import count, chain
+from concurrent.futures import wait
+
+import numpy as np
+import numpy.random
+import scipy.sparse as sp
+
+from ..core import (
+    _NUM_THREAD_WORKERS,
+    complex_array,
+    dag,
+    dot,
+    get_thread_pool,
+    kron,
+    nmlz,
+    prod,
+    ptr,
+    pvectorize,
+    qarray,
+    qu,
+    rdmul,
+    vectorize,
+)
+
+
+class _RGenHandler:
+    """Private object that handles pool of random number generators for
+    parallel number generation - seeding them & changing the underlying bit
+    generators.
+    """
+
+    def __init__(self, initial_seed=None, initial_bitgen=None):
+        self.rgs = []
+        self.set_seed(initial_seed)
+        self.set_bitgen(initial_bitgen)
+
+    def set_bitgen(self, bitgen):
+        """Set the core underlying bit-generator.
+
+        Parameters
+        ----------
+        bitgen : {None, str}
+            Which bit generator to use, either from numpy or `randomgen` -
+            https://bashtage.github.io/randomgen/bit_generators/index.html.
+        """
+
+        if bitgen is None:
+            self.gen_fn = numpy.random.default_rng
+
+        else:
+            try:
+                # try numpy first
+                bg = getattr(numpy.random, bitgen)
+            except AttributeError:
+                # and then (now largely deprecated) randomgen
+                import randomgen
+
+                bg = getattr(randomgen, bitgen)
+
+            def gen(s):
+                return numpy.random.Generator(bg(s))
+
+            self.gen_fn = gen
+
+        # delete any old rgens
+        self.rgs = []
+
+    def set_seed(self, seed=None):
+        """Set the seed for the bit generators.
+
+        Parameters
+        ----------
+        seed : {None, int}, optional
+            Seed supplied to `numpy.random.SeedSequence`. None will randomly
+            the generators (default).
+        """
+        seq = numpy.random.SeedSequence(seed)
+
+        # compute seeds in batches of 4 for perf
+        self.seeds = iter(chain.from_iterable(seq.spawn(4) for _ in count()))
+
+        # delete any old rgens
+        self.rgs = []
+
+    def get_rgens(self, num_threads):
+        """Get a list of the :class:`numpy.random.Generator` instances, having
+        made sure there are at least ``num_threads``.
+
+        Parameters
+        ----------
+        num_threads : int
+            The number of generators to return.
+
+        Returns
+        -------
+        list[numpy.random.Generator]
+        """
+        num_gens = len(self.rgs)
+
+        if num_gens < num_threads:
+            self.rgs.extend(
+                self.gen_fn(next(self.seeds))
+                for _ in range(num_gens, num_threads)
+            )
+
+        return self.rgs[:num_threads]
+
+
+_RG_HANDLER = _RGenHandler()
+
+
+def seed_rand(seed):
+    """See the random number generators, by instantiating a new set of bit
+    generators with a 'seed sequence'.
+    """
+    global _RG_HANDLER
+    return _RG_HANDLER.set_seed(seed)
+
+
+def set_rand_bitgen(bitgen):
+    """Set the core bit generator type to use, from either ``numpy`` or
+    ``randomgen``.
+
+    Parameters
+    ----------
+    bitgen : {'PCG64', 'SFC64', 'MT19937', 'Philox', str}
+        Which bit generator to use.
+    """
+    global _RG_HANDLER
+    return _RG_HANDLER.set_bitgen(bitgen)
+
+
+def _get_rgens(num_threads):
+    global _RG_HANDLER
+    return _RG_HANDLER.get_rgens(num_threads)
+
+
+def randn(
+    shape=(),
+    dtype=float,
+    scale=1.0,
+    loc=0.0,
+    num_threads=None,
+    seed=None,
+    dist="normal",
+):
+    """Fast multithreaded generation of random normally distributed data.
+
+    Parameters
+    ----------
+    shape : tuple[int]
+        The shape of the output random array.
+    dtype : {'complex128', 'float64', 'complex64' 'float32'}, optional
+        The data-type of the output array.
+    scale : float, optional
+        A multiplicative scale for the random numbers.
+    loc : float, optional
+        An additive location for the random numbers.
+    num_threads : int, optional
+        How many threads to use. If ``None``, decide automatically.
+    dist : {'normal', 'uniform', 'rademacher', 'exp'}, optional
+        Type of random number to generate.
+    """
+    if seed is not None:
+        seed_rand(seed)
+
+    if isinstance(shape, Integral):
+        d = shape
+        shape = (shape,)
+    else:
+        d = prod(shape)
+
+    if dist == "rademacher":
+        # not parallelized for now
+        return rand_rademacher(shape, scale=scale, loc=loc, dtype=dtype)
+
+    if num_threads is None:
+        # only multi-thread for big ``d``
+        if d <= 32768:
+            num_threads = 1
+        else:
+            num_threads = _NUM_THREAD_WORKERS
+
+    rgs = _get_rgens(num_threads)
+
+    gen_method = {
+        "uniform": "random",
+        "normal": "standard_normal",
+        "exp": "standard_exponential",
+    }.get(dist, dist)
+
+    # sequential generation
+    if num_threads <= 1:
+
+        def create(d, dtype):
+            out = np.empty(d, dtype)
+            getattr(rgs[0], gen_method)(out=out, dtype=dtype)
+            return out
+
+    # threaded generation
+    else:
+        pool = get_thread_pool()
+        S = math.ceil(d / num_threads)
+
+        def _fill(gen, out, dtype, first, last):
+            getattr(gen, gen_method)(out=out[first:last], dtype=dtype)
+
+        def create(d, dtype):
+            out = np.empty(d, dtype)
+            # submit thread work
+            fs = [
+                pool.submit(_fill, gen, out, dtype, i * S, (i + 1) * S)
+                for i, gen in enumerate(rgs)
+            ]
+            wait(fs)
+            return out
+
+    if np.issubdtype(dtype, np.floating):
+        out = create(d, dtype)
+
+    elif np.issubdtype(dtype, np.complexfloating):
+        # need to sum two real arrays if generating complex numbers
+        if np.issubdtype(dtype, np.complex64):
+            sub_dtype = np.float32
+        else:
+            sub_dtype = np.float64
+
+        out = complex_array(create(d, sub_dtype), create(d, sub_dtype))
+
+    else:
+        raise ValueError(f"dtype {dtype} not understood.")
+
+    if out.dtype != dtype:
+        out = out.astype(dtype)
+
+    if scale != 1.0:
+        out *= scale
+    if loc != 0.0:
+        out += loc
+
+    return out.reshape(shape)
+
+
+@wraps(randn)
+def rand(*args, **kwargs):
+    kwargs.setdefault("dist", "uniform")
+    return randn(*args, **kwargs)
+
+
+def random_seed_fn(fn):
+    """Modify ``fn`` to take a ``seed`` argument (so as to seed the random
+    generators once-only at beginning of function not every ``randn`` call).
+    """
+
+    @wraps(fn)
+    def wrapped_fn(*args, seed=None, **kwargs):
+        if seed is not None:
+            seed_rand(seed)
+        return fn(*args, **kwargs)
+
+    return wrapped_fn
+
+
+def _randint(*args, **kwargs):
+    return _get_rgens(1)[0].integers(*args, **kwargs)
+
+
+def _choice(*args, **kwargs):
+    return _get_rgens(1)[0].choice(*args, **kwargs)
+
+
+choice = random_seed_fn(_choice)
+
+
+@random_seed_fn
+def rand_rademacher(shape, scale=1, loc=0.0, dtype=float):
+    """ """
+    if np.issubdtype(dtype, np.floating):
+        entries = loc + np.array([1.0, -1.0]) * scale
+        need2convert = dtype not in (float, np.float64)
+
+    elif np.issubdtype(dtype, np.complexfloating):
+        entries = loc + np.array([1.0, -1.0, 1.0j, -1.0j]) * scale
+        need2convert = dtype not in (complex, np.complex128)
+
+    else:
+        raise TypeError(
+            f"dtype {dtype} not understood - should be float or complex."
+        )
+
+    x = _choice(entries, shape)
+    if need2convert:
+        x = x.astype(dtype)
+
+    return x
+
+
+def _phase_to_complex_base(x):
+    return 1j * math.sin(x) + math.cos(x)
+
+
+_phase_sigs = ["complex64(float32)", "complex128(float64)"]
+_phase_to_complex_seq = vectorize(_phase_sigs)(_phase_to_complex_base)
+"""Turn array of phases into unit circle complex numbers - sequential.
+"""
+_phase_to_complex_par = pvectorize(_phase_sigs)(_phase_to_complex_base)
+"""Turn array of phases into unit circle complex numbers - parallel.
+"""
+
+
+def phase_to_complex(x):
+    if x.size >= 512:
+        return _phase_to_complex_par(x)
+    # XXX: this is not as fast as numexpr - investigate?
+    return _phase_to_complex_seq(x)
+
+
+@random_seed_fn
+def rand_phase(shape, scale=1, dtype=complex):
+    """Generate random complex numbers distributed on the unit sphere."""
+    if not np.issubdtype(dtype, np.complexfloating):
+        raise ValueError(f"dtype must be complex, got '{dtype}'.")
+
+    if np.issubdtype(dtype, np.complex64):
+        sub_dtype = np.float32
+    else:
+        sub_dtype = np.float64
+
+    phi = randn(shape, dtype=sub_dtype, scale=2 * math.pi, dist="uniform")
+    z = phase_to_complex(phi)
+    if scale != 1:
+        z *= scale
+
+    return z
+
+
+def get_rand_fill_fn(
+    dist="normal",
+    loc=0.0,
+    scale=1.0,
+    seed=None,
+    dtype="float64",
+):
+    """Get a callable with the given random distribution and parameters, that
+    has signature ``fill_fn(shape) -> array``.
+
+    Parameters
+    ----------
+    dist : {'normal', 'uniform', 'rademacher', 'exp'}, optional
+        Type of random number to generate, defaults to 'normal'.
+    loc : float, optional
+        An additive offset to add to the random numbers.
+    scale : float, optional
+        A multiplicative factor to scale the random numbers by.
+    seed : int, optional
+        A random seed.
+    dtype : {'float64', 'complex128', 'float32', 'complex64'}, optional
+        The underlying data type.
+
+    Returns
+    -------
+    callable
+    """
+    if seed is not None:
+        seed_rand(seed)
+
+    def fill_fn(shape=()):
+        return randn(shape, dtype=dtype, dist=dist, loc=loc, scale=scale)
+
+    return fill_fn
+
+
+def rand_matrix(
+    d,
+    scaled=True,
+    sparse=False,
+    stype="csr",
+    density=None,
+    dtype=complex,
+    seed=None,
+):
+    """Generate a random matrix of order `d` with normally distributed
+    entries. If `scaled` is `True`, then in the limit of large `d` the
+    eigenvalues will be distributed on the unit complex disk.
+
+    Parameters
+    ----------
+    d : int
+        Matrix dimension.
+    scaled : bool, optional
+        Whether to scale the matrices values such that its spectrum
+        approximately lies on the unit disk (for dense matrices).
+    sparse : bool, optional
+        Whether to produce a sparse matrix.
+    stype : {'csr', 'csc', 'coo', ...}, optional
+        The type of sparse matrix if ``sparse=True``.
+    density : float, optional
+        Target density of non-zero elements for the sparse matrix. By default
+        aims for about 10 entries per row.
+    dtype : {complex, float}, optional
+        The data type of the matrix elements.
+
+    Returns
+    -------
+    mat : qarray or sparse matrix
+        Random matrix.
+    """
+    if np.issubdtype(dtype, np.floating):
+        iscomplex = False
+    elif np.issubdtype(dtype, np.complexfloating):
+        iscomplex = True
+    else:
+        raise TypeError(
+            f"dtype {dtype} not understood - should be " "float or complex."
+        )
+
+    # handle seed manually since standard python random.seed might be called
+    if seed is not None:
+        seed_rand(seed)
+
+    if sparse:
+        # Aim for 10 non-zero values per row, but betwen 1 and d/2
+        density = min(10, d / 2) / d if density is None else density
+        density = min(
+            max(
+                d**-2,
+                density,
+            ),
+            1.0,
+        )
+        nnz = round(density * d * d)
+
+        if density > 0.1:
+            # take special care to avoid duplicates
+            if seed is not None:
+                random.seed(seed)
+            ijs = random.sample(range(0, d**2), k=nnz)
+        else:
+            ijs = _randint(0, d * d, size=nnz)
+
+        # want to sample nnz unique (d, d) pairs without building list
+        i, j = np.divmod(ijs, d)
+
+        data = randn(nnz, dtype=dtype)
+        mat = sp.coo_matrix((data, (i, j)), shape=(d, d)).asformat(stype)
+    else:
+        density = 1.0
+        mat = qarray(randn((d, d), dtype=dtype))
+
+    if scaled:
+        mat /= ((2 if iscomplex else 1) * d * density) ** 0.5
+
+    return mat
+
+
+@random_seed_fn
+def rand_herm(d, sparse=False, density=None, dtype=complex):
+    """Generate a random hermitian operator of order `d` with normally
+    distributed entries. In the limit of large `d` the spectrum will be a
+    semi-circular distribution between [-1, 1].
+
+    See Also
+    --------
+    rand_matrix, rand_pos, rand_rho, rand_uni
+    """
+    if sparse:
+        density = 10 / d if density is None else density
+        density = min(max(density, d**-2), 1 - d**-2)
+        density /= 2  # to account of herm construction
+
+    herm = rand_matrix(
+        d, scaled=True, sparse=sparse, density=density, dtype=dtype
+    )
+
+    if sparse:
+        herm.data /= 2**1.5
+    else:
+        herm /= 2**1.5
+
+    herm += dag(herm)
+
+    return herm
+
+
+@random_seed_fn
+def rand_pos(d, sparse=False, density=None, dtype=complex):
+    """Generate a random positive operator of size `d`, with normally
+    distributed entries. In the limit of large `d` the spectrum will lie
+    between [0, 1].
+
+    See Also
+    --------
+    rand_matrix, rand_herm, rand_rho, rand_uni
+    """
+    if sparse:
+        density = 10 / d if density is None else density
+        density = min(max(density, d**-2), 1 - d**-2)
+        density = 0.5 * (density / d) ** 0.5  # to account for pos construction
+
+    pos = rand_matrix(
+        d, scaled=True, sparse=sparse, density=density, dtype=dtype
+    )
+
+    return dot(pos, dag(pos))
+
+
+@random_seed_fn
+def rand_rho(d, sparse=False, density=None, dtype=complex):
+    """Generate a random positive operator of size `d` with normally
+    distributed entries and unit trace.
+
+    See Also
+    --------
+    rand_matrix, rand_herm, rand_pos, rand_uni
+    """
+    return nmlz(rand_pos(d, sparse=sparse, density=density, dtype=dtype))
+
+
+@random_seed_fn
+def rand_uni(d, dtype=complex):
+    """Generate a random unitary operator of size `d`, distributed according to
+    the Haar measure.
+
+    See Also
+    --------
+    rand_matrix, rand_herm, rand_pos, rand_rho
+    """
+    q, r = np.linalg.qr(rand_matrix(d, dtype=dtype))
+    r = np.diagonal(r)
+    r = r / np.abs(r)  # read-only so not inplace
+    return rdmul(q, r)
+
+
+@random_seed_fn
+def rand_ket(d, sparse=False, stype="csr", density=0.01, dtype=complex):
+    """Generates a ket of length `d` with normally distributed entries."""
+    if sparse:
+        ket = sp.random(d, 1, format=stype, density=density)
+        ket.data = randn((ket.nnz,), dtype=dtype)
+    else:
+        ket = qarray(randn((d, 1), dtype=dtype))
+    return nmlz(ket)
+
+
+@random_seed_fn
+def rand_haar_state(d, dtype=complex):
+    """Generate a random state of dimension `d` according to the Haar
+    distribution.
+    """
+    u = rand_uni(d, dtype=dtype)
+    return u[:, [0]]
+
+
+@random_seed_fn
+def gen_rand_haar_states(d, reps, dtype=complex):
+    """Generate many random Haar states, recycling a random unitary operator
+    by using all of its columns (not a good idea?).
+    """
+    for rep in range(reps):
+        cyc = rep % d
+        if cyc == 0:
+            u = rand_uni(d, dtype=dtype)
+        yield u[:, [cyc]]
+
+
+@random_seed_fn
+def rand_mix(d, tr_d_min=None, tr_d_max=None, mode="rand", dtype=complex):
+    """Constructs a random mixed state by tracing out a random ket
+    where the composite system varies in size between 2 and d. This produces
+    a spread of states including more purity but has no real meaning.
+    """
+    if tr_d_min is None:
+        tr_d_min = 2
+    if tr_d_max is None:
+        tr_d_max = d
+
+    m = _randint(tr_d_min, tr_d_max)
+    if mode == "rand":
+        psi = rand_ket(d * m, dtype=dtype)
+    elif mode == "haar":
+        psi = rand_haar_state(d * m, dtype=dtype)
+
+    return ptr(psi, [d, m], 0)
+
+
+@random_seed_fn
+def rand_product_state(n, qtype=None, dtype=complex):
+    """Generates a ket of `n` many random pure qubits."""
+
+    def gen_rand_pure_qubits(n):
+        for _ in range(n):
+            (u,) = rand(1)
+            (v,) = rand(1)
+            phi = 2 * np.pi * u
+            theta = np.arccos(2 * v - 1)
+            yield qu(
+                [
+                    [np.cos(theta / 2.0)],
+                    [np.sin(theta / 2.0) * np.exp(1.0j * phi)],
+                ],
+                qtype=qtype,
+                dtype=dtype,
+            )
+
+    return kron(*gen_rand_pure_qubits(n))
+
+
+@random_seed_fn
+def rand_matrix_product_state(
+    n, bond_dim, phys_dim=2, dtype=complex, cyclic=False, trans_invar=False
+):
+    """Generate a random matrix product state (in dense form, see
+    :func:`~quimb.tensor.MPS_rand_state` for tensor network form).
+
+    Parameters
+    ----------
+    n : int
+        Number of sites.
+    bond_dim : int
+        Dimension of the bond (virtual) indices.
+    phys_dim : int, optional
+        Physical dimension of each local site, defaults to 2 (qubits).
+    cyclic : bool (optional)
+        Whether to impose cyclic boundary conditions on the entanglement
+        structure.
+    trans_invar : bool (optional)
+        Whether to generate a translationally invariant state,
+        requires cyclic=True.
+
+    Returns
+    -------
+    ket : qarray
+        The random state, with shape (phys_dim**n, 1)
+
+    """
+    from quimb.tensor import MPS_rand_state
+
+    mps = MPS_rand_state(
+        n,
+        bond_dim,
+        phys_dim=phys_dim,
+        dtype=dtype,
+        cyclic=cyclic,
+        trans_invar=trans_invar,
+    )
+    return mps.to_dense()
+
+
+rand_mps = rand_matrix_product_state
+
+
+@random_seed_fn
+def rand_seperable(dims, num_mix=10, dtype=complex):
+    """Generate a random, mixed, seperable state. E.g rand_seperable([2, 2])
+    for a mixed two qubit state with no entanglement.
+
+    Parameters
+    ----------
+        dims : tuple of int
+            The local dimensions across which to be seperable.
+        num_mix : int, optional
+            How many individual product states to sum together, each with
+            random weight.
+
+    Returns
+    -------
+        qarray
+            Mixed seperable state.
+    """
+
+    def gen_single_sites():
+        for dim in dims:
+            yield rand_rho(dim, dtype=dtype)
+
+    weights = rand(num_mix)
+
+    def gen_single_states():
+        for w in weights:
+            yield w * kron(*gen_single_sites())
+
+    return sum(gen_single_states()) / np.sum(weights)
+
+
+@random_seed_fn
+def rand_iso(n, m, dtype=complex):
+    """Generate a random isometry of shape ``(n, m)``."""
+    data = randn((n, m), dtype=dtype)
+
+    q, _ = np.linalg.qr(data if n >= m else data.T)
+    q = q.astype(dtype)
+
+    return q if (n >= m) else q.T
+
+
+@random_seed_fn
+def rand_mera(n, invariant=False, dtype=complex):
+    """Generate a random mera state of ``n`` qubits, which must be a power
+    of 2. This uses ``quimb.tensor``.
+    """
+    import quimb.tensor as qt
+
+    if invariant:
+        constructor = qt.MERA.rand_invar
+    else:
+        constructor = qt.MERA.rand
+
+    return constructor(n, dtype=dtype).to_qarray()