Python numpy 模块，complex_() 实例源码

def besselj(n, z):
    """Bessel function of first kind of order n at kr.
    Wraps scipy.special.jn(n, z).

    Parameters
    ----------
    n : array_like
       Order
    kr: array_like
       Argument

    Returns
    -------
    J : array_like
       Values of Bessel function of order n at position z
    """
    return scy.jv(n, _np.complex_(z))

def sphankel1(n, kr):
    """Spherical Hankel (first kind) of order n at kr

    Parameters
    ----------
    n : array_like
       Order
    kr: array_like
       Argument

    Returns
    -------
    hn1 : complex float
       Spherical Hankel function hn (first kind)
    """
    n, kr = scalar_broadcast_match(n, kr)
    hn1 = _np.full(n.shape, _np.nan, dtype=_np.complex_)
    kr_nonzero = kr != 0
    hn1[kr_nonzero] = _np.sqrt(_np.pi / 2) / _np.lib.scimath.sqrt(kr[kr_nonzero]) * hankel1(n[kr_nonzero] + 0.5, kr[kr_nonzero])
    return hn1

def sphankel2(n, kr):
    """Spherical Hankel (second kind) of order n at kr

    Parameters
    ----------
    n : array_like
       Order
    kr: array_like
       Argument

    Returns
    -------
    hn2 : complex float
       Spherical Hankel function hn (second kind)
    """
    n, kr = scalar_broadcast_match(n, kr)
    hn2 = _np.full(n.shape, _np.nan, dtype=_np.complex_)
    kr_nonzero = kr != 0
    hn2[kr_nonzero] = _np.sqrt(_np.pi / 2) / _np.lib.scimath.sqrt(kr[kr_nonzero]) * hankel2(n[kr_nonzero] + 0.5, kr[kr_nonzero])
    return hn2

def dsphankel2(n, kr):
    """Derivative spherical Hankel (second kind) of order n at kr

    Parameters
    ----------
    n : array_like
       Order
    kr: array_like
       Argument

    Returns
    -------
    dhn2 : complex float
       Derivative of spherical Hankel function hn' (second kind)
    """
    n, kr = scalar_broadcast_match(n, kr)
    dhn2 = _np.full(n.shape, _np.nan, dtype=_np.complex_)
    kr_nonzero = kr != 0
    dhn2[kr_nonzero] = 0.5 * (sphankel2(n[kr_nonzero] - 1, kr[kr_nonzero]) - sphankel2(n[kr_nonzero] + 1, kr[kr_nonzero]) - sphankel2(n[kr_nonzero], kr[kr_nonzero]) / kr[kr_nonzero])
    return dhn2

def _random_vec(sites, ldim, randstate=None, dtype=np.complex_):
    """Returns a random complex vector (normalized to ||x||_2 = 1) of shape
    (ldim,) * sites, i.e. a pure state with local dimension `ldim` living on
    `sites` sites.

    :param sites: Number of local sites
    :param ldim: Local ldimension
    :param randstate: numpy.random.RandomState instance or None
    :returns: numpy.ndarray of shape (ldim,) * sites

    >>> psi = _random_vec(5, 2); psi.shape
    (2, 2, 2, 2, 2)
    >>> np.abs(np.vdot(psi, psi) - 1) < 1e-6
    True
    """
    shape = (ldim, ) * sites
    psi = _randfuncs[dtype](shape, randstate=randstate)
    psi /= np.linalg.norm(psi)
    return psi

def _random_op(sites, ldim, hermitian=False, normalized=False, randstate=None,
               dtype=np.complex_):
    """Returns a random operator  of shape (ldim,ldim) * sites with local
    dimension `ldim` living on `sites` sites in global form.

    :param sites: Number of local sites
    :param ldim: Local ldimension
    :param hermitian: Return only the hermitian part (default False)
    :param normalized: Normalize to Frobenius norm=1 (default False)
    :param randstate: numpy.random.RandomState instance or None
    :returns: numpy.ndarray of shape (ldim,ldim) * sites

    >>> A = _random_op(3, 2); A.shape
    (2, 2, 2, 2, 2, 2)
    """
    op = _randfuncs[dtype]((ldim**sites,) * 2, randstate=randstate)
    if hermitian:
        op += np.transpose(op).conj()
    if normalized:
        op /= np.linalg.norm(op)
    return op.reshape((ldim,) * 2 * sites)

def _standard_normal(shape, randstate=np.random, dtype=np.float_):
    """Generates a standard normal numpy array of given shape and dtype, i.e.
    this function is equivalent to `randstate.randn(*shape)` for real dtype and
    `randstate.randn(*shape) + 1.j * randstate.randn(shape)` for complex dtype.

    :param tuple shape: Shape of array to be returned
    :param randstate: An instance of :class:`numpy.random.RandomState` (default is
        ``np.random``))
    :param dtype: ``np.float_`` (default) or `np.complex_`

    Returns
    -------

    A: An array of given shape and dtype with standard normal entries

    """
    if dtype == np.float_:
        return randstate.randn(*shape)
    elif dtype == np.complex_:
        return randstate.randn(*shape) + 1.j * randstate.randn(*shape)
    else:
        raise ValueError('{} is not a valid dtype.'.format(dtype))

def test_operations_typesafety(nr_sites, local_dim, rank, rgen):
    # create a real MPA
    mpo1 = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              randstate=rgen, dtype=np.float_)
    mpo2 = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              randstate=rgen, dtype=np.complex_)

    assert mpo1.dtype == np.float_
    assert mpo2.dtype == np.complex_

    assert (mpo1 + mpo1).dtype == np.float_
    assert (mpo1 + mpo2).dtype == np.complex_
    assert (mpo2 + mpo1).dtype == np.complex_

    assert mp.sumup((mpo1, mpo1)).dtype == np.float_
    assert mp.sumup((mpo1, mpo2)).dtype == np.complex_
    assert mp.sumup((mpo2, mpo1)).dtype == np.complex_

    assert (mpo1 - mpo1).dtype == np.float_
    assert (mpo1 - mpo2).dtype == np.complex_
    assert (mpo2 - mpo1).dtype == np.complex_

    mpo1 += mpo2
    assert mpo1.dtype == np.complex_

def test_inject_many(local_dim, rank, rgen):
    """Calling mp.inject() repeatedly vs. calling it with sequence arguments"""
    mpa = factory.random_mpa(3, local_dim, rank, rgen, normalized=True,
                             dtype=np.complex_)
    inj_lt = [factory._zrandn(s, rgen) for s in [(2, 3), (1,), (2, 2), (3, 2)]]

    mpa_inj1 = mp.inject(mpa, 1, None, [inj_lt[0]])
    mpa_inj1 = mp.inject(mpa_inj1, 2, 1, inj_lt[0])
    mpa_inj1 = mp.inject(mpa_inj1, 4, None, [inj_lt[2]])
    mpa_inj2 = mp.inject(mpa, [1, 2], [2, None], [inj_lt[0], [inj_lt[2]]])
    mpa_inj3 = mp.inject(mpa, [1, 2], [2, 1], [inj_lt[0], inj_lt[2]])
    assert_mpa_almost_equal(mpa_inj1, mpa_inj2, True)
    assert_mpa_almost_equal(mpa_inj1, mpa_inj3, True)

    inj_lt = [inj_lt[:2], inj_lt[2:]]
    mpa_inj1 = mp.inject(mpa, 1, None, inj_lt[0])
    mpa_inj1 = mp.inject(mpa_inj1, 4, inject_ten=inj_lt[1])
    mpa_inj2 = mp.inject(mpa, [1, 2], None, inj_lt)
    assert_mpa_almost_equal(mpa_inj1, mpa_inj2, True)

def test_inject_vs_chain(nr_sites, local_dim, rank, rgen):
    """Compare mp.inject() with mp.chain()"""
    if nr_sites == 1:
        return
    mpa = factory.random_mpa(nr_sites // 2, local_dim, rank, rgen,
                             dtype=np.complex_, normalized=True)
    pten = [factory._zrandn((local_dim,) * 2) for _ in range(nr_sites // 2)]
    pten_mpa = mp.MPArray.from_kron(pten)

    outer1 = mp.chain((pten_mpa, mpa))
    outer2 = mp.inject(mpa, 0, inject_ten=pten)
    assert_mpa_almost_equal(outer1, outer2, True)

    outer1 = mp.chain((mpa, pten_mpa))
    outer2 = mp.inject(mpa, [len(mpa)], [None], inject_ten=[pten])
    assert_mpa_almost_equal(outer1, outer2, True)

def test_povm_ic_mpa(nr_sites, local_dim, rank, rgen):
    # Check that the tensor product of the PauliGen POVM is IC.
    paulis = povm.pauli_povm(local_dim)
    inv_map = mp_from_array_repeat(paulis.linear_inversion_map, nr_sites)
    probab_map = mp_from_array_repeat(paulis.probability_map, nr_sites)
    reconstruction_map = mp.dot(inv_map, probab_map)

    eye = factory.eye(nr_sites, local_dim**2)
    assert mp.norm(reconstruction_map - eye) < 1e-5

    # Check linear inversion for a particular example MPA.
    # Linear inversion works for arbitrary matrices, not only for states,
    # so we test it for an arbitrary MPA.
    # Normalize, otherwise the absolute error check below will not work.
    mpa = factory.random_mpa(nr_sites, local_dim**2, rank,
                             dtype=np.complex_, randstate=rgen, normalized=True)
    probabs = mp.dot(probab_map, mpa)
    recons = mp.dot(inv_map, probabs)
    assert mp.norm(recons - mpa) < 1e-6

def test_mppovm_pmf_as_array_pmps(
        nr_sites, local_dim, rank, startsite, width, rgen):
    if hasattr(local_dim, '__len__'):
        pdims = [d for d, _ in local_dim]
        mppaulis = mp.chain(
            povm.MPPovm.from_local_povm(povm.pauli_povm(d), 1)
            for d in pdims[startsite:startsite + width]
        )
    else:
        pdims = local_dim
        local_dim = (local_dim, local_dim)
        mppaulis = povm.MPPovm.from_local_povm(povm.pauli_povm(pdims), width)
    mppaulis = mppaulis.embed(nr_sites, startsite, pdims)
    pmps = factory.random_mpa(nr_sites, local_dim, rank,
                              dtype=np.complex_, randstate=rgen, normalized=True)
    rho = mpsmpo.pmps_to_mpo(pmps)
    expect_rho = mppaulis.pmf_as_array(rho, 'mpdo')

    for impl in ['default', 'pmps-ltr', 'pmps-symm']:
        expect_pmps = mppaulis.pmf_as_array(pmps, 'pmps', impl=impl)
        assert_array_almost_equal(expect_rho, expect_pmps, err_msg=impl)

def test_pmps_to_mpo(nr_sites, local_dim, rank, rgen):
    if (nr_sites % 2) != 0:
        return
    nr_sites = nr_sites // 2
    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              dtype=np.complex_, randstate=rgen)
    rho_mp = mm.pmps_to_mpo(pmps).to_array_global()

    # Local form is what we will use: One system site, one ancilla site, etc
    purification = pmps.to_array()
    # Convert to a density matrix
    purification = np.outer(purification, purification.conj())
    purification = purification.reshape((local_dim,) * (2 * 2 * nr_sites))
    # Trace out the ancilla sites
    traceout = tuple(range(1, 2 * nr_sites, 2))
    rho_np = utils.partial_trace(purification, traceout)

    # Here, we need global form
    assert_array_almost_equal(rho_mp, rho_np)

def test_against_cmath(self):
        import cmath

        points = [-1-1j, -1+1j, +1-1j, +1+1j]
        name_map = {'arcsin': 'asin', 'arccos': 'acos', 'arctan': 'atan',
                    'arcsinh': 'asinh', 'arccosh': 'acosh', 'arctanh': 'atanh'}
        atol = 4*np.finfo(np.complex).eps
        for func in self.funcs:
            fname = func.__name__.split('.')[-1]
            cname = name_map.get(fname, fname)
            try:
                cfunc = getattr(cmath, cname)
            except AttributeError:
                continue
            for p in points:
                a = complex(func(np.complex_(p)))
                b = cfunc(p)
                assert_(abs(a - b) < atol, "%s %s: %s; cmath: %s" % (fname, p, a, b))

def default(self, obj):
        # convert dates and numpy objects in a json serializable format
        if isinstance(obj, datetime):
            return obj.strftime('%Y-%m-%dT%H:%M:%SZ')
        elif isinstance(obj, date):
            return obj.strftime('%Y-%m-%d')
        elif type(obj) in (np.int_, np.intc, np.intp, np.int8, np.int16,
                           np.int32, np.int64, np.uint8, np.uint16,
                           np.uint32, np.uint64):
            return int(obj)
        elif type(obj) in (np.bool_,):
            return bool(obj)
        elif type(obj) in (np.float_, np.float16, np.float32, np.float64,
                           np.complex_, np.complex64, np.complex128):
            return float(obj)

        # Let the base class default method raise the TypeError
        return json.JSONEncoder.default(self, obj)

def test_against_cmath(self):
        import cmath

        points = [-1-1j, -1+1j, +1-1j, +1+1j]
        name_map = {'arcsin': 'asin', 'arccos': 'acos', 'arctan': 'atan',
                    'arcsinh': 'asinh', 'arccosh': 'acosh', 'arctanh': 'atanh'}
        atol = 4*np.finfo(np.complex).eps
        for func in self.funcs:
            fname = func.__name__.split('.')[-1]
            cname = name_map.get(fname, fname)
            try:
                cfunc = getattr(cmath, cname)
            except AttributeError:
                continue
            for p in points:
                a = complex(func(np.complex_(p)))
                b = cfunc(p)
                assert_(abs(a - b) < atol, "%s %s: %s; cmath: %s" % (fname, p, a, b))

def test_against_cmath(self):
        import cmath

        points = [-1-1j, -1+1j, +1-1j, +1+1j]
        name_map = {'arcsin': 'asin', 'arccos': 'acos', 'arctan': 'atan',
                    'arcsinh': 'asinh', 'arccosh': 'acosh', 'arctanh': 'atanh'}
        atol = 4*np.finfo(np.complex).eps
        for func in self.funcs:
            fname = func.__name__.split('.')[-1]
            cname = name_map.get(fname, fname)
            try:
                cfunc = getattr(cmath, cname)
            except AttributeError:
                continue
            for p in points:
                a = complex(func(np.complex_(p)))
                b = cfunc(p)
                assert_(abs(a - b) < atol, "%s %s: %s; cmath: %s" % (fname, p, a, b))

def test_against_cmath(self):
        import cmath

        points = [-1-1j, -1+1j, +1-1j, +1+1j]
        name_map = {'arcsin': 'asin', 'arccos': 'acos', 'arctan': 'atan',
                    'arcsinh': 'asinh', 'arccosh': 'acosh', 'arctanh': 'atanh'}
        atol = 4*np.finfo(np.complex).eps
        for func in self.funcs:
            fname = func.__name__.split('.')[-1]
            cname = name_map.get(fname, fname)
            try:
                cfunc = getattr(cmath, cname)
            except AttributeError:
                continue
            for p in points:
                a = complex(func(np.complex_(p)))
                b = cfunc(p)
                assert_(abs(a - b) < atol, "%s %s: %s; cmath: %s" % (fname, p, a, b))

def test_against_cmath(self):
        import cmath

        points = [-1-1j, -1+1j, +1-1j, +1+1j]
        name_map = {'arcsin': 'asin', 'arccos': 'acos', 'arctan': 'atan',
                    'arcsinh': 'asinh', 'arccosh': 'acosh', 'arctanh': 'atanh'}
        atol = 4*np.finfo(np.complex).eps
        for func in self.funcs:
            fname = func.__name__.split('.')[-1]
            cname = name_map.get(fname, fname)
            try:
                cfunc = getattr(cmath, cname)
            except AttributeError:
                continue
            for p in points:
                a = complex(func(np.complex_(p)))
                b = cfunc(p)
                assert_(abs(a - b) < atol, "%s %s: %s; cmath: %s" % (fname, p, a, b))

def default(self, obj):
        # convert dates and numpy objects in a json serializable format
        if isinstance(obj, datetime):
            return obj.strftime('%Y-%m-%dT%H:%M:%SZ')
        elif isinstance(obj, date):
            return obj.strftime('%Y-%m-%d')
        elif type(obj) in [np.int_, np.intc, np.intp, np.int8, np.int16,
                           np.int32, np.int64, np.uint8, np.uint16,
                           np.uint32, np.uint64]:
            return int(obj)
        elif type(obj) in [np.bool_]:
            return bool(obj)
        elif type(obj) in [np.float_, np.float16, np.float32, np.float64,
                           np.complex_, np.complex64, np.complex128]:
            return float(obj)

        # Let the base class default method raise the TypeError
        return json.JSONEncoder.default(self, obj)

def test_against_cmath(self):
        import cmath

        points = [-1-1j, -1+1j, +1-1j, +1+1j]
        name_map = {'arcsin': 'asin', 'arccos': 'acos', 'arctan': 'atan',
                    'arcsinh': 'asinh', 'arccosh': 'acosh', 'arctanh': 'atanh'}
        atol = 4*np.finfo(np.complex).eps
        for func in self.funcs:
            fname = func.__name__.split('.')[-1]
            cname = name_map.get(fname, fname)
            try:
                cfunc = getattr(cmath, cname)
            except AttributeError:
                continue
            for p in points:
                a = complex(func(np.complex_(p)))
                b = cfunc(p)
                assert_(abs(a - b) < atol, "%s %s: %s; cmath: %s" % (fname, p, a, b))

def test_against_cmath(self):
        import cmath

        points = [-1-1j, -1+1j, +1-1j, +1+1j]
        name_map = {'arcsin': 'asin', 'arccos': 'acos', 'arctan': 'atan',
                    'arcsinh': 'asinh', 'arccosh': 'acosh', 'arctanh': 'atanh'}
        atol = 4*np.finfo(np.complex).eps
        for func in self.funcs:
            fname = func.__name__.split('.')[-1]
            cname = name_map.get(fname, fname)
            try:
                cfunc = getattr(cmath, cname)
            except AttributeError:
                continue
            for p in points:
                a = complex(func(np.complex_(p)))
                b = cfunc(p)
                assert_(abs(a - b) < atol, "%s %s: %s; cmath: %s" % (fname, p, a, b))

def construct_b_matrix(self, PSI, GAMMA):
        """Construct B matrix as in D. F. V. James et al. Phys. Rev. A, 64, 052312 (2001).

        :param array PSI: :math:`\psi_\\nu` vector with :math:`\\nu=1,...,16`, computed in __init__
        :param array GAMMA: :math:`\Gamma` matrices, computed in __init__

        :return: :math:`B_{\\nu,\mu} = \\langle \psi_\\nu \\rvert  \Gamma_\mu  \\lvert \psi_\\nu \\rangle`
        :rtype: numpy array
        """

        B = np.complex_(np.zeros((16,16)))

        for i in range(16):
            for j in range(16):
                B[i,j] = np.dot(np.conj(PSI[i]) , np.dot(GAMMA[j], PSI[i]))
        return B

def rho_phys(self, t):
        """Positive semidefinite matrix based on t values.

        :param numpy_array t: tvalues

        :return: A positive semidefinite matrix which is an estimation of the actual density matrix.
        :rtype: numpy matrix
        """
        T = np.complex_(np.matrix([
                        [t[0],              0,              0,              0],
                        [t[4]+1j*t[5],      t[1],           0,              0],
                        [t[10]+1j*t[11],    t[6]+1j*t[7],   t[2],           0],
                        [t[14]+1j*t[15],    t[12]+1j*t[13], t[8]+1j*t[9],   t[3]]
                        ]))

        TdagT = np.dot(T.conj().T , T)
        norm = np.trace(TdagT)

        return TdagT/norm

def spbessel(n, kr):
    """Spherical Bessel function (first kind) of order n at kr

    Parameters
    ----------
    n : array_like
       Order
    kr: array_like
       Argument

    Returns
    -------
    J : complex float
       Spherical Bessel
    """
    n, kr = scalar_broadcast_match(n, kr)

    if _np.any(n < 0) | _np.any(kr < 0) | _np.any(_np.mod(n, 1) != 0):
        J = _np.zeros(kr.shape, dtype=_np.complex_)

        kr_non_zero = kr != 0
        J[kr_non_zero] = _np.lib.scimath.sqrt(_np.pi / 2 / kr[kr_non_zero]) * besselj(n[kr_non_zero] + 0.5, kr[kr_non_zero])
        J[_np.logical_and(kr == 0, n == 0)] = 1
    else:
        J = scy.spherical_jn(n.astype(_np.int), kr)
    return _np.squeeze(J)

def spneumann(n, kr):
    """Spherical Neumann (Bessel second kind) of order n at kr

    Parameters
    ----------
    n : array_like
       Order
    kr: array_like
       Argument

    Returns
    -------
    Yv : complex float
       Spherical Neumann (Bessel second kind)
    """
    n, kr = scalar_broadcast_match(n, kr)

    if _np.any(n < 0) | _np.any(_np.mod(n, 1) != 0):
        Yv = _np.full(kr.shape, _np.nan, dtype=_np.complex_)

        kr_non_zero = kr != 0
        Yv[kr_non_zero] = _np.lib.scimath.sqrt(_np.pi / 2 / kr[kr_non_zero]) * neumann(n[kr_non_zero] + 0.5, kr[kr_non_zero])
        Yv[kr < 0] = -Yv[kr < 0]
    else:
        Yv = scy.spherical_yn(n.astype(_np.int), kr)
        Yv[_np.isinf(Yv)] = _np.nan  # return possible infs as nan to stay consistent
    return _np.squeeze(Yv)

def random_lowrank(rows, cols, rank, randstate=np.random, dtype=np.float_):
    """Returns a random lowrank matrix of given shape and dtype"""
    if dtype == np.float_:
        A = randstate.randn(rows, rank)
        B = randstate.randn(cols, rank)
    elif dtype == np.complex_:
        A = randstate.randn(rows, rank) + 1.j * randstate.randn(rows, rank)
        B = randstate.randn(cols, rank) + 1.j * randstate.randn(cols, rank)
    else:
        raise ValueError("{} is not a valid dtype".format(dtype))

    C = A.dot(B.conj().T)
    return C / np.linalg.norm(C)

def random_mpo(sites, ldim, rank, randstate=None, hermitian=False,
               normalized=True, force_rank=False):
    """Returns an hermitian MPO with randomly choosen local tensors

    :param sites: Number of sites
    :param ldim: Local dimension
    :param rank: Rank
    :param randstate: numpy.random.RandomState instance or None
    :param hermitian: Is the operator supposed to be hermitian
    :param normalized: Operator should have unit norm
    :param force_rank: If True, the rank is exaclty `rank`.
        Otherwise, it might be reduced if we reach the maximum sensible rank.
    :returns: randomly choosen matrix product operator

    >>> mpo = random_mpo(4, 2, 10, force_rank=True)
    >>> mpo.ranks, mpo.shape
    ((10, 10, 10), ((2, 2), (2, 2), (2, 2), (2, 2)))
    >>> mpo.canonical_form
    (0, 4)

    """
    mpo = random_mpa(sites, (ldim,) * 2, rank, randstate=randstate,
                     force_rank=force_rank, dtype=np.complex_)

    if hermitian:
        # make mpa Herimitan in place, without increasing rank:
        ltens = (l + l.swapaxes(1, 2).conj() for l in mpo.lt)
        mpo = mp.MPArray(ltens)
    if normalized:
        # we do this with a copy to ensure the returned state is not
        # normalized
        mpo /= mp.norm(mpo.copy())

    return mpo

def test_div_mpo_scalar(nr_sites, local_dim, rank, rgen):
    mpo = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                             dtype=np.complex_, randstate=rgen)
    # FIXME Change behavior of to_array
    # For nr_sites == 1, changing `mpo` below will change `op` as
    # well, unless we call .copy().
    op = mpo.to_array_global().copy()
    scalar = rgen.randn() + 1.j * rgen.randn()

    assert_array_almost_equal(op / scalar, (mpo / scalar).to_array_global())

    mpo /= scalar
    assert_array_almost_equal(op / scalar, mpo.to_array_global())

def test_mppovm_embed_expectation(
        nr_sites, local_dim, rank, startsite, width, rgen):
    if hasattr(local_dim, '__iter__'):
        local_dim2 = local_dim
    else:
        local_dim2 = [local_dim] * nr_sites
    local_dim2 = list(zip(local_dim2, local_dim2))

    # Create a local POVM `red_povm`, embed it onto a larger chain
    # (`full_povm`), and go back to the reduced POVM.
    red_povm = mp.chain(
        mp.povm.MPPovm.from_local_povm(mp.povm.pauli_povm(d), 1)
        for d, _ in local_dim2[startsite:startsite + width]
    )
    full_povm = red_povm.embed(nr_sites, startsite, local_dim)
    axes = [(1, 2) if i < startsite or i >= startsite + width else None
            for i in range(nr_sites)]
    red_povm2 = mp.partialtrace(full_povm, axes, mp.MPArray)
    red_povm2 = mp.prune(red_povm2, singletons=True)
    red_povm2 /= np.prod([d for i, (d, _) in enumerate(local_dim2)
                          if i < startsite or i >= startsite + width])
    assert_almost_equal(mp.normdist(red_povm, red_povm2), 0.0)

    # Test with an arbitrary random MPO instead of an MPDO
    mpo = mp.factory.random_mpa(nr_sites, local_dim2, rank, rgen,
                                dtype=np.complex_, normalized=True)
    mpo_red = next(mp.reductions_mpo(mpo, width, startsites=[startsite]))
    ept = mp.prune(full_povm.pmf(mpo, 'mpdo'), singletons=True).to_array()
    ept_red = red_povm.pmf(mpo_red, 'mpdo').to_array()
    assert_array_almost_equal(ept, ept_red)

def test_mppovm_expectation_pmps(nr_sites, width, local_dim, rank, rgen):
    paulis = povm.pauli_povm(local_dim)
    mppaulis = povm.MPPovm.from_local_povm(paulis, width)
    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              dtype=np.complex_, randstate=rgen)
    rho = mpsmpo.pmps_to_mpo(pmps)
    expect_psi = list(mppaulis.expectations(pmps, mode='pmps'))
    expect_rho = list(mppaulis.expectations(rho))

    assert len(expect_psi) == len(expect_rho)
    for e_rho, e_psi in zip(expect_rho, expect_psi):
        assert_array_almost_equal(e_rho.to_array(), e_psi.to_array())

def test_mppovm_pmf_as_array_pmps_benchmark(
        nr_sites, local_dim, rank, startsite, width, impl, rgen, benchmark):
    pauli_y = povm.pauli_parts(local_dim)[1]
    mpp_y = povm.MPPovm.from_local_povm(pauli_y, width) \
                       .embed(nr_sites, startsite, local_dim)
    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              dtype=np.complex_, randstate=rgen, normalized=True)
    benchmark(lambda: mpp_y.pmf_as_array(pmps, 'pmps', impl=impl))

def test_eig_benchmark(
        nr_sites, local_dim, rank, ev_rank, rgen, benchmark):
    mpo = factory.random_mpo(nr_sites, local_dim, rank, randstate=rgen,
                             hermitian=True, normalized=True)
    mpo.canonicalize()
    mps = factory.random_mpa(nr_sites, local_dim, rank, randstate=rgen,
                             dtype=np.complex_, normalized=True)
    mpo = mpo + mp.mps_to_mpo(mps)

    benchmark(
        mp.eig,
        mpo, startvec_rank=ev_rank, randstate=rgen,
        var_sites=1, num_sweeps=1,
    )

def test_eig_sum_benchmark(
        nr_sites, local_dim, rank, ev_rank, rgen, benchmark):
    mpo = factory.random_mpo(nr_sites, local_dim, rank, randstate=rgen,
                             hermitian=True, normalized=True)
    mpo.canonicalize()
    mps = factory.random_mpa(nr_sites, local_dim, rank, randstate=rgen,
                             dtype=np.complex_, normalized=True)

    benchmark(
        mp.eig_sum,
        [mpo, mps], startvec_rank=ev_rank, randstate=rgen,
        var_sites=1, num_sweeps=1,
    )

def pytest_namespace():
    return dict(
        # nr_sites, local_dim, rank
        MP_TEST_PARAMETERS=[(1, 7, np.nan), (2, 3, 3), (3, 2, 4), (6, 2, 4),
                            (4, 3, 5), (5, 2, 1)],
        MP_TEST_DTYPES=[np.float_, np.complex_]
    )

def test_pmps_dm_to_array(nr_sites, local_dim, rank, rgen):

    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              randstate=rgen, dtype=np.complex_)
    mpo = mm.pmps_to_mpo(pmps)
    op = mpo.to_array()
    op2 = mm.pmps_dm_to_array(pmps)
    assert_array_almost_equal(op2, op)
    op = mpo.to_array_global()
    op2 = mm.pmps_dm_to_array(pmps, True)
    assert_array_almost_equal(op2, op)

def test_pmps_dm_to_array_fast(nr_sites, local_dim, rank, rgen, benchmark):

    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              dtype=np.complex_, normalized=True,
                              randstate=rgen)
    benchmark(mm.pmps_dm_to_array, pmps)

def test_pmps_reduction(nr_sites, local_dim, rank, keep, rgen):

    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              dtype=np.complex_, normalized=True,
                              randstate=rgen)
    rho = mm.pmps_to_mpo(pmps).to_array_global()
    traceout = [pos for pos in range(nr_sites) if pos not in keep]
    red = utils.partial_trace(rho, traceout)
    pmps_red = mm.pmps_reduction(pmps, keep)
    red2 = mm.pmps_to_mpo(pmps_red).to_array_global()
    red2 = red2.reshape([local_dim] * (2 * len(keep)))
    assert_array_almost_equal(red2, red)

def test_pmps_reduction_array_fast(nr_sites, local_dim, rank, keep, rgen,
                                   benchmark):
    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              dtype=np.complex_, normalized=True,
                              randstate=rgen)
    benchmark(lambda: mm.pmps_dm_to_array(mm.pmps_reduction(pmps, keep)))

def test_pmps_reduction_array_slow_noprune(
        nr_sites, local_dim, rank, keep, rgen, benchmark):
    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              dtype=np.complex_, normalized=True,
                              randstate=rgen)
    # NB: The maximal distance between sites of the reduction is
    # limited by the fact that normal numpy builds support arrays with
    # at most 32 indices.
    benchmark(lambda: mm.pmps_to_mpo(mm.pmps_reduction(pmps, keep)).to_array())

def test_pmps_reduction_array_slow_prune(
        nr_sites, local_dim, rank, keep, rgen, benchmark):
    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              dtype=np.complex_, normalized=True,
                              randstate=rgen)
    benchmark(
        lambda: mp.prune(mm.pmps_to_mpo(mm.pmps_reduction(pmps, keep)),
                         singletons=True).to_array()
    )

def test_pmps_reduction_dm_to_array(nr_sites, local_dim, rank, keep, rgen):

    pmps = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                              dtype=np.complex_, randstate=rgen)
    rho = mm.pmps_to_mpo(pmps).to_array_global()
    traceout = [pos for pos in range(nr_sites) if pos not in keep]
    red = utils.partial_trace(rho, traceout)
    pmps_red = mm.pmps_reduction(pmps, keep)
    red2 = mm.pmps_dm_to_array(pmps_red, True)
    assert_array_almost_equal(red2, red)

def test_power_complex(self):
        x = np.array([1+2j, 2+3j, 3+4j])
        assert_equal(x**0, [1., 1., 1.])
        assert_equal(x**1, x)
        assert_almost_equal(x**2, [-3+4j, -5+12j, -7+24j])
        assert_almost_equal(x**3, [(1+2j)**3, (2+3j)**3, (3+4j)**3])
        assert_almost_equal(x**4, [(1+2j)**4, (2+3j)**4, (3+4j)**4])
        assert_almost_equal(x**(-1), [1/(1+2j), 1/(2+3j), 1/(3+4j)])
        assert_almost_equal(x**(-2), [1/(1+2j)**2, 1/(2+3j)**2, 1/(3+4j)**2])
        assert_almost_equal(x**(-3), [(-11+2j)/125, (-46-9j)/2197,
                                      (-117-44j)/15625])
        assert_almost_equal(x**(0.5), [ncu.sqrt(1+2j), ncu.sqrt(2+3j),
                                       ncu.sqrt(3+4j)])
        norm = 1./((x**14)[0])
        assert_almost_equal(x**14 * norm,
                [i * norm for i in [-76443+16124j, 23161315+58317492j,
                                    5583548873 + 2465133864j]])

        # Ticket #836
        def assert_complex_equal(x, y):
            assert_array_equal(x.real, y.real)
            assert_array_equal(x.imag, y.imag)

        for z in [complex(0, np.inf), complex(1, np.inf)]:
            z = np.array([z], dtype=np.complex_)
            with np.errstate(invalid="ignore"):
                assert_complex_equal(z**1, z)
                assert_complex_equal(z**2, z*z)
                assert_complex_equal(z**3, z*z*z)

def test_loss_of_precision(self):
        for dtype in [np.complex64, np.complex_]:
            yield self.check_loss_of_precision, dtype

def test_return_dtype(self):
        assert_equal(select(self.conditions, self.choices, 1j).dtype,
                     np.complex_)
        # But the conditions need to be stronger then the scalar default
        # if it is scalar.
        choices = [choice.astype(np.int8) for choice in self.choices]
        assert_equal(select(self.conditions, choices).dtype, np.int8)

        d = np.array([1, 2, 3, np.nan, 5, 7])
        m = np.isnan(d)
        assert_equal(select([m], [d]), [0, 0, 0, np.nan, 0, 0])

def test_power_complex(self):
        x = np.array([1+2j, 2+3j, 3+4j])
        assert_equal(x**0, [1., 1., 1.])
        assert_equal(x**1, x)
        assert_almost_equal(x**2, [-3+4j, -5+12j, -7+24j])
        assert_almost_equal(x**3, [(1+2j)**3, (2+3j)**3, (3+4j)**3])
        assert_almost_equal(x**4, [(1+2j)**4, (2+3j)**4, (3+4j)**4])
        assert_almost_equal(x**(-1), [1/(1+2j), 1/(2+3j), 1/(3+4j)])
        assert_almost_equal(x**(-2), [1/(1+2j)**2, 1/(2+3j)**2, 1/(3+4j)**2])
        assert_almost_equal(x**(-3), [(-11+2j)/125, (-46-9j)/2197,
                                      (-117-44j)/15625])
        assert_almost_equal(x**(0.5), [ncu.sqrt(1+2j), ncu.sqrt(2+3j),
                                       ncu.sqrt(3+4j)])
        norm = 1./((x**14)[0])
        assert_almost_equal(x**14 * norm,
                [i * norm for i in [-76443+16124j, 23161315+58317492j,
                                    5583548873 + 2465133864j]])

        # Ticket #836
        def assert_complex_equal(x, y):
            assert_array_equal(x.real, y.real)
            assert_array_equal(x.imag, y.imag)

        for z in [complex(0, np.inf), complex(1, np.inf)]:
            z = np.array([z], dtype=np.complex_)
            with np.errstate(invalid="ignore"):
                assert_complex_equal(z**1, z)
                assert_complex_equal(z**2, z*z)
                assert_complex_equal(z**3, z*z*z)

def test_loss_of_precision(self):
        for dtype in [np.complex64, np.complex_]:
            yield self.check_loss_of_precision, dtype

def test_return_dtype(self):
        assert_equal(select(self.conditions, self.choices, 1j).dtype,
                     np.complex_)
        # But the conditions need to be stronger then the scalar default
        # if it is scalar.
        choices = [choice.astype(np.int8) for choice in self.choices]
        assert_equal(select(self.conditions, choices).dtype, np.int8)

        d = np.array([1, 2, 3, np.nan, 5, 7])
        m = np.isnan(d)
        assert_equal(select([m], [d]), [0, 0, 0, np.nan, 0, 0])

def test_power_complex(self):
        x = np.array([1+2j, 2+3j, 3+4j])
        assert_equal(x**0, [1., 1., 1.])
        assert_equal(x**1, x)
        assert_almost_equal(x**2, [-3+4j, -5+12j, -7+24j])
        assert_almost_equal(x**3, [(1+2j)**3, (2+3j)**3, (3+4j)**3])
        assert_almost_equal(x**4, [(1+2j)**4, (2+3j)**4, (3+4j)**4])
        assert_almost_equal(x**(-1), [1/(1+2j), 1/(2+3j), 1/(3+4j)])
        assert_almost_equal(x**(-2), [1/(1+2j)**2, 1/(2+3j)**2, 1/(3+4j)**2])
        assert_almost_equal(x**(-3), [(-11+2j)/125, (-46-9j)/2197,
                                      (-117-44j)/15625])
        assert_almost_equal(x**(0.5), [ncu.sqrt(1+2j), ncu.sqrt(2+3j),
                                       ncu.sqrt(3+4j)])
        norm = 1./((x**14)[0])
        assert_almost_equal(x**14 * norm,
                [i * norm for i in [-76443+16124j, 23161315+58317492j,
                                    5583548873 + 2465133864j]])

        # Ticket #836
        def assert_complex_equal(x, y):
            assert_array_equal(x.real, y.real)
            assert_array_equal(x.imag, y.imag)

        for z in [complex(0, np.inf), complex(1, np.inf)]:
            z = np.array([z], dtype=np.complex_)
            with np.errstate(invalid="ignore"):
                assert_complex_equal(z**1, z)
                assert_complex_equal(z**2, z*z)
                assert_complex_equal(z**3, z*z*z)

def test_loss_of_precision(self):
        for dtype in [np.complex64, np.complex_]:
            yield self.check_loss_of_precision, dtype