Пример #1
0
def test_meshgrid(as_midp):
    vec1 = (0, 1)
    vec2 = (-1, 0, 1)
    vec3 = (2, 3, 4, 5)

    # Sparse meshgrid
    mgx = np.array(vec1)[:, np.newaxis, np.newaxis]
    mgy = np.array(vec2)[np.newaxis, :, np.newaxis]
    mgz = np.array(vec3)[np.newaxis, np.newaxis, :]

    grid = TensorGrid(vec1, vec2, vec3, as_midp=as_midp)
    xx, yy, zz = grid.meshgrid()
    assert all_equal(mgx, xx)
    assert all_equal(mgy, yy)
    assert all_equal(mgz, zz)

    xx, yy, zz = grid.meshgrid()
    assert all_equal(mgx, xx)
    assert all_equal(mgy, yy)
    assert all_equal(mgz, zz)

    # Fortran ordering
    grid = TensorGrid(vec1, vec2, vec3, order='F', as_midp=as_midp)
    xx, yy, zz = grid.meshgrid()
    assert all(arr.flags.f_contiguous for arr in (xx, yy, zz))