mparray_test.py 文件源码

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)