Python numpy.complex_() Examples

def _zeropi_operator_in_product_basis(self, zeropi_operator, zeropi_evecs=None):
        """Helper method that converts a zeropi operator into one in the product basis.

        Returns
        -------
        scipy.sparse.csc_matrix
            operator written in the product basis
        """
        zeropi_dim = self.zeropi_cutoff
        zeta_dim = self.zeta_cutoff

        if zeropi_evecs is None:
            _, zeropi_evecs = self._zeropi.eigensys(evals_count=zeropi_dim)

        op_eigen_basis = sparse.dia_matrix((zeropi_dim, zeropi_dim),
                                           dtype=np.complex_)  # is this guaranteed to be zero?

        op_zeropi = spec_utils.get_matrixelement_table(zeropi_operator, zeropi_evecs)
        for n in range(zeropi_dim):
            for m in range(zeropi_dim):
                op_eigen_basis += op_zeropi[n, m] * op.hubbard_sparse(n, m, zeropi_dim)

        return sparse.kron(op_eigen_basis, sparse.identity(zeta_dim, format='csc', dtype=np.complex_), format='csc') 

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 _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 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_mult_mpo_scalar_normalization(nr_sites, local_dim, rank, rgen):
    if nr_sites < 2:
        # Re-normalization has no effect for nr_sites == 1. There is
        # nothing more to test than :func:`test_mult_mpo_scalar`.
        return

    mpo = factory.random_mpa(nr_sites, (local_dim, local_dim), rank,
                             dtype=np.complex_, randstate=rgen)
    op = mpo.to_array_global()
    scalar = rgen.randn() + 1.j * rgen.randn()

    center = nr_sites // 2
    mpo.canonicalize(left=center - 1, right=center)
    mpo_times_two = scalar * mpo

    assert_array_almost_equal(scalar * op, mpo_times_two.to_array_global())
    assert_correct_normalization(mpo_times_two, center - 1, center)

    mpo *= scalar
    assert_array_almost_equal(scalar * op, mpo.to_array_global())
    assert_correct_normalization(mpo, center - 1, center) 

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 _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_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(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 __init__(self, EJ1, EJ2, EJ3, ECJ1, ECJ2, ECJ3, ECg1, ECg2, ng1, ng2, flux, ncut,
                 truncated_dim=None):
        self.EJ1 = EJ1
        self.EJ2 = EJ2
        self.EJ3 = EJ3
        self.ECJ1 = ECJ1
        self.ECJ2 = ECJ2
        self.ECJ3 = ECJ3
        self.ECg1 = ECg1
        self.ECg2 = ECg2
        self.ng1 = ng1
        self.ng2 = ng2
        self.flux = flux
        self.ncut = ncut
        self.truncated_dim = truncated_dim
        self._sys_type = type(self).__name__
        self._evec_dtype = np.complex_
        self._default_grid = discretization.Grid1d(-np.pi / 2, 3 * np.pi / 2, 100)    # for plotting in phi_j basis
        self._image_filename = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'qubit_pngs/fluxqubit.png') 

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 sparse_kinetic_mat(self):
        """
        Kinetic energy portion of the Hamiltonian.
        TODO: update this method to use single-variable operator methods

        Returns
        -------
        scipy.sparse.csc_matrix
            matrix representing the kinetic energy operator
        """
        pt_count = self.grid.pt_count
        dim_theta = 2 * self.ncut + 1
        identity_phi = sparse.identity(pt_count, format='csc', dtype=np.complex_)
        identity_theta = sparse.identity(dim_theta, format='csc', dtype=np.complex_)

        kinetic_matrix_phi = self.grid.second_derivative_matrix(prefactor=-2.0 * self.ECJ)

        diag_elements = 2.0 * self.ECS * np.square(np.arange(-self.ncut + self.ng, self.ncut + 1 + self.ng))
        kinetic_matrix_theta = sparse.dia_matrix((diag_elements, [0]), shape=(dim_theta, dim_theta)).tocsc()

        kinetic_matrix = (sparse.kron(kinetic_matrix_phi, identity_theta, format='csc')
                          + sparse.kron(identity_phi, kinetic_matrix_theta, format='csc'))

        kinetic_matrix -= 2.0 * self.ECS * self.dCJ * self.i_d_dphi_operator() * self.n_theta_operator()
        return kinetic_matrix 

def convert_esys_to_ndarray(esys_qutip):
    """Takes a qutip eigenstates array, as obtained with .eigenstates(), and converts it into a pure numpy array.

    Parameters
    ----------
    esys_qutip: ndarray of qutip.Qobj
        as obtained from qutip `.eigenstates()`

    Returns
    -------
    ndarray
        converted eigenstate data
    """
    evals_count = len(esys_qutip)
    dimension = esys_qutip[0].shape[0]
    esys_ndarray = np.empty((evals_count, dimension), dtype=np.complex_)
    for index, eigenstate in enumerate(esys_qutip):
        esys_ndarray[index] = eigenstate.full()[:, 0]
    return esys_ndarray 

def test_numpy_scalar(self):
        dtype = self.dtype
        if dtype is numpy.bool_:
            x = dtype(True)
        elif issubclass(dtype, numpy.complex_):
            x = dtype(3.2 - 2.4j)
        elif issubclass(dtype, numpy.integer):
            x = dtype(3)
        elif issubclass(dtype, numpy.floating):
            x = dtype(3.2)
        else:
            assert False

        y = cuda.to_gpu(x)
        assert isinstance(y, cuda.ndarray)
        assert y.shape == ()
        assert y.dtype == dtype
        assert y == x 

def test_numpy_scalar(self):
        dtype = self.dtype
        if dtype is numpy.bool_:
            x = dtype(True)
        elif issubclass(dtype, numpy.complex_):
            x = dtype(3.2 - 2.4j)
        elif issubclass(dtype, numpy.integer):
            x = dtype(3)
        elif issubclass(dtype, numpy.floating):
            x = dtype(3.2)
        else:
            assert False

        y = cuda.to_cpu(x)
        assert isinstance(y, numpy.ndarray)
        assert y.shape == ()
        assert y.dtype == dtype
        assert y == x 

def _default(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 isinstance(obj, (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 isinstance(obj, np.bool_):
            return bool(obj)
        elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64,
                              np.complex_, np.complex64, np.complex128)):
            return float(obj)

        raise TypeError(f"Object of type '{obj.__class__.__name__}' is not JSON serializable") 

def get_const_matrix_from_channel(channel, symbols):
        """
        Extract the numeric constant matrix.

        Args:
            channel (matrix): a 4x4 symbolic matrix.
            symbols (list): The full list [x1, ..., xn] of symbols
                used in the matrix.

        Returns:
            matrix: a 4x4 numeric matrix.

        Additional Information:
            Each entry of the 4x4 symbolic input channel matrix is assumed to
            be a polynomial of the form a1x1 + ... + anxn + c. The corresponding
            entry in the output numeric matrix is c.
        """
        from sympy import Poly
        n = channel.rows
        M = numpy.zeros((n, n), dtype=numpy.complex_)
        for (i, j) in itertools.product(range(n), range(n)):
            M[i, j] = numpy.complex(
                Poly(channel[i, j], symbols).coeff_monomial(1))
        return M 

def get_matrix_from_channel(channel, symbol):
        """
        Extract the numeric parameter matrix.

        Args:
            channel (matrix): a 4x4 symbolic matrix.
            symbol (list): a symbol xi

        Returns:
            matrix: a 4x4 numeric matrix.

        Additional Information:
            Each entry of the 4x4 symbolic input channel matrix is assumed to
            be a polynomial of the form a1x1 + ... + anxn + c. The corresponding
            entry in the output numeric matrix is ai.
        """
        from sympy import Poly
        n = channel.rows
        M = numpy.zeros((n, n), dtype=numpy.complex_)
        for (i, j) in itertools.product(range(n), range(n)):
            M[i, j] = numpy.complex(
                Poly(channel[i, j], symbol).coeff_monomial(symbol))
        return M 

def test_against_cmath(self):
        import cmath, sys

        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 wavefunction(self, esys, which=0, phi_grid=None):
        """Returns a fluxonium wave function in `phi` basis

        Parameters
        ----------
        esys: ndarray, ndarray
            eigenvalues, eigenvectors
        which: int, optional
             index of desired wave function (default value = 0)
        phi_grid: Grid1d, optional
            used for setting a custom grid for phi; if None use self._default_grid

        Returns
        -------
        WaveFunction object
        """
        if esys is None:
            evals_count = max(which + 1, 3)
            evals, evecs = self.eigensys(evals_count)
        else:
            evals, evecs = esys
        dim = self.hilbertdim()

        phi_grid = phi_grid or self._default_grid

        phi_basis_labels = phi_grid.make_linspace()
        wavefunc_osc_basis_amplitudes = evecs[:, which]
        phi_wavefunc_amplitudes = np.zeros(phi_grid.pt_count, dtype=np.complex_)
        phi_osc = self.phi_osc()
        for n in range(dim):
            phi_wavefunc_amplitudes += wavefunc_osc_basis_amplitudes[n] * osc.harm_osc_wavefunction(n, phi_basis_labels,
                                                                                                    phi_osc)
        return storage.WaveFunction(basis_labels=phi_basis_labels, amplitudes=phi_wavefunc_amplitudes,
                                    energy=evals[which]) 

def wavefunction(self, esys=None, which=0, phi_grid=None):
        """Return the transmon wave function in phase basis. The specific index of the wavefunction is `which`.
        `esys` can be provided, but if set to `None` then it is calculated on the fly.

        Parameters
        ----------
        esys: tuple(ndarray, ndarray), optional
            if None, the eigensystem is calculated on the fly; otherwise, the provided eigenvalue, eigenvector arrays
            as obtained from `.eigensystem()` are used
        which: int, optional
            eigenfunction index (default value = 0)
        phi_grid: Grid1d, optional
            used for setting a custom grid for phi; if None use self._default_grid

        Returns
        -------
        WaveFunction object
        """
        if esys is None:
            evals_count = max(which + 1, 3)
            esys = self.eigensys(evals_count)
        evals, _ = esys
        n_wavefunc = self.numberbasis_wavefunction(esys, which=which)

        phi_grid = phi_grid or self._default_grid

        phi_basis_labels = phi_grid.make_linspace()
        phi_wavefunc_amplitudes = np.empty(phi_grid.pt_count, dtype=np.complex_)
        for k in range(phi_grid.pt_count):
            phi_wavefunc_amplitudes[k] = ((1j**which / math.sqrt(2 * np.pi)) *
                                          np.sum(n_wavefunc.amplitudes *
                                                 np.exp(1j * phi_basis_labels[k] * n_wavefunc.basis_labels)))
        return storage.WaveFunction(basis_labels=phi_basis_labels, amplitudes=phi_wavefunc_amplitudes,
                                    energy=evals[which])


# — Flux-tunable Cooper pair box / transmon——————————————————————————————————————————— 

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_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 _n_operator(self):
        diag_elements = np.arange(-self.ncut, self.ncut + 1, dtype=np.complex_)
        return np.diag(diag_elements) 

def _phi_operator(self):
        r"""
        Operator :math:`\varphi`, acting only on the `\varphi` Hilbert subspace.


        Returns
        -------
            scipy.sparse.csc_matrix
        """
        pt_count = self.grid.pt_count

        phi_matrix = sparse.dia_matrix((pt_count, pt_count), dtype=np.complex_)
        diag_elements = self.grid.make_linspace()
        phi_matrix.setdiag(diag_elements)
        return phi_matrix 

def __new__(cls, x=0):
        if isinstance(x, afnumpy.ndarray):
            return x.astype(cls)
        elif isinstance(x, numbers.Number):
            return numpy.complex_(x)
        else:
            return afnumpy.array(x).astype(cls) 

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_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)