def _check_xyz(*args, **kwargs):
nargs = len(args)
if nargs == 1:
x, y = [None]*2
z = asarray(args[0])
elif nargs == 3:
x, y, z = [asarray(a) for a in args]
else:
raise TypeError("_check_xyz: wrong number of arguments")
try:
nx, ny = shape(z)
except:
raise ValueError("z must be 2D, not %dD" % len(shape(z)))
indexing = kwargs.get('indexing', 'ij')
if x is None and y is None:
if indexing == 'ij':
nx, ny = ny, nx # swap
x, y = meshgrid(list(range(ny)), list(range(nx)), indexing=indexing)
else:
if indexing == 'ij':
assert shape(x) == (nx,ny) or shape(x) == (nx,1) or len(x) == nx, \
"_check_xyz: x has shape %s, expected %s, %s, or %s" % \
(shape(x), (nx,ny), (nx,1), (nx,))
assert shape(y) == (nx,ny) or shape(y) == (1,ny) or len(y) == ny, \
"_check_xyz: y has shape %s, expected %s, %s, or %s" % \
(shape(y), (nx,ny), (1,ny), (ny,))
else:
assert shape(x) == (nx,ny) or shape(x) == (1,ny) or len(x) == ny, \
"_check_xyz: x has shape %s, expected %s, %s, or %s" % \
(shape(x), (nx,ny), (1,ny), (ny,))
assert shape(y) == (nx,ny) or shape(y) == (nx,1) or len(y) == nx, \
"_check_xyz: y has shape %s, expected %s, %s, or %s" % \
(shape(y), (nx,ny), (nx,1), (nx,))
return x, y, z
def _check_xyz(*args, **kwargs):
nargs = len(args)
if nargs == 1:
x, y = [None] * 2
z = asarray(args[0])
elif nargs == 3:
x, y, z = [asarray(a) for a in args]
else:
raise TypeError("_check_xyz: wrong number of arguments")
try:
nx, ny = shape(z)
except:
raise ValueError("z must be 2D, not %dD" % len(shape(z)))
indexing = kwargs.get('indexing', 'ij')
if x is None and y is None:
if indexing == 'ij':
nx, ny = ny, nx # swap
x, y = meshgrid(list(range(ny)), list(range(nx)), indexing=indexing)
else:
if indexing == 'ij':
assert shape(x) == (nx,ny) or shape(x) == (nx,1) or len(x) == nx, \
"_check_xyz: x has shape %s, expected %s, %s, or %s" % \
(shape(x), (nx,ny), (nx,1), (nx,))
assert shape(y) == (nx,ny) or shape(y) == (1,ny) or len(y) == ny, \
"_check_xyz: y has shape %s, expected %s, %s, or %s" % \
(shape(y), (nx,ny), (1,ny), (ny,))
else:
assert shape(x) == (nx,ny) or shape(x) == (1,ny) or len(x) == ny, \
"_check_xyz: x has shape %s, expected %s, %s, or %s" % \
(shape(x), (nx,ny), (1,ny), (ny,))
assert shape(y) == (nx,ny) or shape(y) == (nx,1) or len(y) == nx, \
"_check_xyz: y has shape %s, expected %s, %s, or %s" % \
(shape(y), (nx,ny), (nx,1), (nx,))
return x, y, z
def _check_xyzv(*args, **kwargs):
nargs = len(args)
if nargs == 1:
x, y, z = [None]*3
v = asarray(args[0])
elif nargs == 4:
x, y, z, v = [asarray(a) for a in args]
else:
raise ValueError("_check_xyzv: wrong number of arguments")
try:
nx, ny, nz = shape(v)
except:
raise ValueError('_check_xyzv: v must be 3D, not %dD' % len(shape(v)))
indexing = kwargs.get('indexing', 'ij')
if x is None and y is None and z is None:
if indexing == 'ij':
ny, nx = nx, nz # swap
x, y, z = meshgrid(list(range(ny)),
list(range(nx)),
list(range(nz)), indexing=indexing)
else:
if indexing == 'ij':
assert shape(x)==(nx,ny,nz) or shape(x)==(nx,1,1) or \
shape(x)==(nx,), \
"_check_xyzv: x has shape %s, expected %s, %s, or %s" % \
(shape(x), (nx,ny,nz), (nx,1,1), (nx,))
if shape(x) == (nx,ny,nz):
assert shape(y) == (nx,ny,nz), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (nx,ny,nz), shape(y))
assert shape(z) == (nx,ny,nz), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (nx,ny,nz), shape(z))
elif shape(x) == (nx,1,1):
assert shape(y) == (1,ny,1), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (1,ny,1), shape(y))
assert shape(z) == (1,1,nz), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (1,1,nz), shape(z))
else: # shape(x) == (nx,)
assert shape(y) == (ny,), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (ny,), shape(y))
assert shape(z) == (nz,), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (nz,), shape(z))
else:
assert shape(x)==(nx,ny,nz) or shape(x)==(1,ny,1) or \
shape(x)==(ny,), \
"_check_xyzv: x has shape %s, expected %s, %s, or %s" % \
(shape(x), (nx,ny,nz), (1,ny,1), (ny,))
if shape(x) == (nx,ny,nz):
assert shape(y) == (nx,ny,nz), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (nx,ny,nz), shape(y))
assert shape(z) == (nx,ny,nz), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (nx,ny,nz), shape(z))
elif shape(x) == (1,ny,1):
assert shape(y) == (nx,1,1), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (nx,1,1), shape(y))
assert shape(z) == (1,1,nz), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (1,1,nz), shape(z))
else: # shape(x) == (ny,)
assert shape(y) == (nx,), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (nx,), shape(y))
assert shape(z) == (nz,), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (nz,), shape(z))
return x, y, z, v
def _check_xyzuvw(*args, **kwargs):
nargs = len(args)
if nargs == 4:
x, y = [None]*2
z, u, v, w = [asarray(a) for a in args]
elif nargs == 6:
x, y, z, u, v, w = [asarray(a) for a in args]
else:
raise TypeError("_check_xyzuvw: wrong number of arguments")
indexing = kwargs.get('indexing', 'xy')
us = shape(u)
assert us == shape(v) == shape(w), \
"_check_xyzuvw: u, v, and w must be of same shape"
if len(us) == 1:
if x is None and y is None:
x = seq(us[0]-1)
y = seq(us[0]-1)
else:
assert shape(x) == us, \
"_check_xyuv: x has shape %s, expected %s" % (shape(x), us)
assert shape(y) == us, \
"_check_xyuv: y has shape %s, expected %s" % (shape(y), us)
assert shape(z) == us, \
"_check_xyuv: z has shape %s, expected %s" % (shape(z), us)
elif len(us) == 2:
nx, ny = us
if x is None and y is None:
x, y, z = _check_xyz(z, indexing=indexing)
else:
x, y, z = _check_xyz(x, y, z, indexing=indexing)
assert shape(z) == us, \
"_check_xyzuvw: z, u, v, and w must be of same shape"
elif len(us) == 3:
nx, ny, nz = us
if x is None and y is None:
if indexing == 'ij':
nx, ny = ny, nx # swap
x, y, junk = meshgrid(seq(ny-1), seq(nx-1), seq(nz-1))
else:
if indexing == 'ij':
assert shape(x)==us or shape(x)==(nx,1,1) or shape(x)==(nx,), \
"_check_xyzuvw: x has shape %s, expected %s, %s, or %s"\
% (shape(x), us, (nx,1,1), (nx,))
assert shape(y)==us or shape(y)==(1,ny,1) or shape(y)==(ny,), \
"_check_xyzuvw: y has shape %s, expected %s, %s, or %s"\
% (shape(y), us, (1,ny,1), (ny,))
else:
assert shape(x)==us or shape(x)==(1,ny,1) or shape(x)==(ny,), \
"_check_xyzuvw: x has shape %s, expected %s, %s, or %s"\
% (shape(x), us, (1,ny,1), (ny,))
assert shape(y)==us or shape(y)==(nx,1,1) or shape(y)==(nx,), \
"_check_xyzuvw: y has shape %s, expected %s, %s, or %s"\
% (shape(y), us, (nx,1,1), (nx,))
assert shape(z) == us or shape(z) == (1,1,nz) or shape(z) == (nz,), \
"_check_xyzuvw: z has shape %s, expected %s, %s, or %s" % \
(shape(z), us, (1,1,nz), (nz,))
else:
raise ValueError(
"_check_xyzuvw: u must be 1D, 2D, or 3D, not %dD" % len(us))
return x, y, z, u, v, w
def _check_xyzv(*args, **kwargs):
nargs = len(args)
if nargs == 1:
x, y, z = [None] * 3
v = asarray(args[0])
elif nargs == 4:
x, y, z, v = [asarray(a) for a in args]
else:
raise ValueError("_check_xyzv: wrong number of arguments")
try:
nx, ny, nz = shape(v)
except:
raise ValueError('_check_xyzv: v must be 3D, not %dD' % len(shape(v)))
indexing = kwargs.get('indexing', 'ij')
if x is None and y is None and z is None:
if indexing == 'ij':
ny, nx = nx, nz # swap
x, y, z = meshgrid(list(range(ny)),
list(range(nx)),
list(range(nz)),
indexing=indexing)
else:
if indexing == 'ij':
assert shape(x)==(nx,ny,nz) or shape(x)==(nx,1,1) or \
shape(x)==(nx,), \
"_check_xyzv: x has shape %s, expected %s, %s, or %s" % \
(shape(x), (nx,ny,nz), (nx,1,1), (nx,))
if shape(x) == (nx, ny, nz):
assert shape(y) == (nx,ny,nz), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (nx,ny,nz), shape(y))
assert shape(z) == (nx,ny,nz), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (nx,ny,nz), shape(z))
elif shape(x) == (nx, 1, 1):
assert shape(y) == (1,ny,1), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (1,ny,1), shape(y))
assert shape(z) == (1,1,nz), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (1,1,nz), shape(z))
else: # shape(x) == (nx,)
assert shape(y) == (ny,), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (ny,), shape(y))
assert shape(z) == (nz,), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (nz,), shape(z))
else:
assert shape(x)==(nx,ny,nz) or shape(x)==(1,ny,1) or \
shape(x)==(ny,), \
"_check_xyzv: x has shape %s, expected %s, %s, or %s" % \
(shape(x), (nx,ny,nz), (1,ny,1), (ny,))
if shape(x) == (nx, ny, nz):
assert shape(y) == (nx,ny,nz), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (nx,ny,nz), shape(y))
assert shape(z) == (nx,ny,nz), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (nx,ny,nz), shape(z))
elif shape(x) == (1, ny, 1):
assert shape(y) == (nx,1,1), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (nx,1,1), shape(y))
assert shape(z) == (1,1,nz), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (1,1,nz), shape(z))
else: # shape(x) == (ny,)
assert shape(y) == (nx,), \
"_check_xyzv: x has shape %s, expected y to be %s, " \
"not %s" % (shape(x), (nx,), shape(y))
assert shape(z) == (nz,), \
"_check_xyzv: x has shape %s, expected z to be %s, " \
"not %s" % (shape(x), (nz,), shape(z))
return x, y, z, v
def _check_xyzuvw(*args, **kwargs):
nargs = len(args)
if nargs == 4:
x, y = [None] * 2
z, u, v, w = [asarray(a) for a in args]
elif nargs == 6:
x, y, z, u, v, w = [asarray(a) for a in args]
else:
raise TypeError("_check_xyzuvw: wrong number of arguments")
indexing = kwargs.get('indexing', 'xy')
us = shape(u)
assert us == shape(v) == shape(w), \
"_check_xyzuvw: u, v, and w must be of same shape"
if len(us) == 1:
if x is None and y is None:
x = seq(us[0] - 1)
y = seq(us[0] - 1)
else:
assert shape(x) == us, \
"_check_xyuv: x has shape %s, expected %s" % (shape(x), us)
assert shape(y) == us, \
"_check_xyuv: y has shape %s, expected %s" % (shape(y), us)
assert shape(z) == us, \
"_check_xyuv: z has shape %s, expected %s" % (shape(z), us)
elif len(us) == 2:
nx, ny = us
if x is None and y is None:
x, y, z = _check_xyz(z, indexing=indexing)
else:
x, y, z = _check_xyz(x, y, z, indexing=indexing)
assert shape(z) == us, \
"_check_xyzuvw: z, u, v, and w must be of same shape"
elif len(us) == 3:
nx, ny, nz = us
if x is None and y is None:
if indexing == 'ij':
nx, ny = ny, nx # swap
x, y, junk = meshgrid(seq(ny - 1), seq(nx - 1), seq(nz - 1))
else:
if indexing == 'ij':
assert shape(x)==us or shape(x)==(nx,1,1) or shape(x)==(nx,), \
"_check_xyzuvw: x has shape %s, expected %s, %s, or %s"\
% (shape(x), us, (nx,1,1), (nx,))
assert shape(y)==us or shape(y)==(1,ny,1) or shape(y)==(ny,), \
"_check_xyzuvw: y has shape %s, expected %s, %s, or %s"\
% (shape(y), us, (1,ny,1), (ny,))
else:
assert shape(x)==us or shape(x)==(1,ny,1) or shape(x)==(ny,), \
"_check_xyzuvw: x has shape %s, expected %s, %s, or %s"\
% (shape(x), us, (1,ny,1), (ny,))
assert shape(y)==us or shape(y)==(nx,1,1) or shape(y)==(nx,), \
"_check_xyzuvw: y has shape %s, expected %s, %s, or %s"\
% (shape(y), us, (nx,1,1), (nx,))
assert shape(z) == us or shape(z) == (1,1,nz) or shape(z) == (nz,), \
"_check_xyzuvw: z has shape %s, expected %s, %s, or %s" % \
(shape(z), us, (1,1,nz), (nz,))
else:
raise ValueError("_check_xyzuvw: u must be 1D, 2D, or 3D, not %dD" %
len(us))
return x, y, z, u, v, w