Beispiel #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))
Beispiel #2
0
def test_meshgrid():
    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)
    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(sparse=True)
    assert all_equal(mgx, xx)
    assert all_equal(mgy, yy)
    assert all_equal(mgz, zz)

    # Dense meshgrid
    mgx = np.empty((2, 3, 4))
    for i in range(2):
        mgx[i, :, :] = vec1[i]

    mgy = np.empty((2, 3, 4))
    for i in range(3):
        mgy[:, i, :] = vec2[i]

    mgz = np.empty((2, 3, 4))
    for i in range(4):
        mgz[:, :, i] = vec3[i]

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

    assert all_equal(xx.shape, (2, 3, 4))