def operation(func, o1, o2, reindex=True, broadcast=True, constructor=None):
""" operation on LaxArray objects
input:
func : operator
o1 : LHS operand: DimArray
o2 : RHS operand: at least: be convertible by np.array())
align, optional: if True, use pandas to align the axes
output:
values: array values
dims : dimension names
"""
if constructor is None:
constructor = o1._constructor
# second operand is not a DimArray: let numpy do the job
if not is_DimArray(o2): # isinstance
if np.ndim(o2) > np.ndim(o1):
raise ValueError(
"bad input: second operand's dimensions not documented")
res = func(o1.values, np.array(o2))
return constructor(res, o1.axes)
# check for first operand (reverse operations)
elif not is_DimArray(o1): # isinstance
if np.ndim(o1) > np.ndim(o2):
raise ValueError(
"bad input: second operand's dimensions not documented")
res = func(np.array(o1), o2.values)
return constructor(res, o2.axes)
# both objects are dimarrays
# Align axes by re-indexing
if reindex:
o1, o2 = align_axes(o1, o2)
# Align dimensions by adding new axes and transposing if necessary
if broadcast:
o1, o2 = align_dims(o1, o2)
# make the new axes
newaxes = o1.axes.copy()
# ...make sure no singleton value is included
for i, ax in enumerate(newaxes):
if ax.values[0] is None:
newaxes[i] = o2.axes[ax.name]
res = func(o1.values, o2.values)
return constructor(res, newaxes)
def operation(func, o1, o2, reindex=True, broadcast=True, constructor=None):
""" operation on LaxArray objects
input:
func : operator
o1 : LHS operand: DimArray
o2 : RHS operand: at least: be convertible by np.array())
align, optional: if True, use pandas to align the axes
output:
values: array values
dims : dimension names
"""
if constructor is None:
constructor = o1._constructor
# second operand is not a DimArray: let numpy do the job
if not is_DimArray(o2): # isinstance
if np.ndim(o2) > np.ndim(o1):
raise ValueError("bad input: second operand's dimensions not documented")
res = func(o1.values, np.array(o2))
return constructor(res, o1.axes)
# check for first operand (reverse operations)
elif not is_DimArray(o1): # isinstance
if np.ndim(o1) > np.ndim(o2):
raise ValueError("bad input: second operand's dimensions not documented")
res = func(np.array(o1), o2.values)
return constructor(res, o2.axes)
# both objects are dimarrays
# Align axes by re-indexing
if reindex:
o1, o2 = align_axes(o1, o2)
# Align dimensions by adding new axes and transposing if necessary
if broadcast:
o1, o2 = align_dims(o1, o2)
# make the new axes
newaxes = o1.axes.copy()
# ...make sure no singleton value is included
for i, ax in enumerate(newaxes):
if ax.values[0] is None:
newaxes[i] = o2.axes[ax.name]
res = func(o1.values, o2.values)
return constructor(res, newaxes)
def stack(arrays, axis, keys=None, align=False):
""" stack arrays along a new dimension (raise error if already existing)
parameters:
----------
arrays: sequence or dict of arrays
axis: str, new dimension along which to stack the array
keys, optional: stack axis values, useful if array is a sequence, or a non-ordered dictionary
align, optional: if True, align axes prior to stacking (Default to False)
returns:
--------
DimArray: joint array
Sea Also:
---------
concatenate: join arrays along an existing dimension
Examples:
---------
>>> a = DimArray([1,2,3])
>>> b = DimArray([11,22,33])
>>> stack([a, b], axis='stackdim', keys=['a','b'])
dimarray: 6 non-null elements (0 null)
dimensions: 'stackdim', 'x0'
0 / stackdim (2): a to b
1 / x0 (3): 0 to 2
array([[ 1, 2, 3],
[11, 22, 33]])
"""
# make a sequence of arrays
arrays, keys = _check_stack_args(arrays, keys)
for a in arrays:
if not is_DimArray(a): raise TypeError('can only stack DimArray instances')
# make sure the stacking dimension is OK (new)
dims = get_dims(*arrays)
axis = _check_stack_axis(axis, dims)
# re-index axes if needed
if align:
arrays = align_axes(*arrays)
# make it a numpy array
data = [a.values for a in arrays]
data = np.array(data)
# new axis
newaxis = Axis(keys, axis)
# find common axes
try:
axes = _get_axes(*arrays)
except ValueError, msg:
if 'axes are not aligned' in repr(msg):
msg = 'axes are not aligned\n ==> Try passing `align=True`'
raise ValueError(msg)
def is_boolean_array(value):
"""
>>> a = DimArray([1,2,3])
>>> is_boolean_array(a)
False
>>> is_boolean_array(a>1)
True
"""
return (isinstance(value, np.ndarray) or is_DimArray(value)) \
and value.dtype is np.dtype('bool')
def is_boolean_index(indices, shape):
""" check if something like a[a>2] is performed for ndim > 1
"""
# indices = np.index_exp[indices]
# if len(shape) > 1 and len(indices) == 1:
if isinstance(indices, np.ndarray) or is_DimArray(indices):
if indices.shape == shape:
if indices.dtype == np.dtype(bool):
return True
return False # otherwise
def is_boolean_array(value):
"""
>>> a = DimArray([1,2,3])
>>> is_boolean_array(a)
False
>>> is_boolean_array(a>1)
True
"""
return (isinstance(value, np.ndarray) or is_DimArray(value)) \
and value.dtype is np.dtype('bool')
def is_boolean_index(indices, shape):
""" check if something like a[a>2] is performed for ndim > 1
"""
#indices = np.index_exp[indices]
#if len(shape) > 1 and len(indices) == 1:
if isinstance(indices, np.ndarray) or is_DimArray(indices):
if indices.shape == shape:
if indices.dtype == np.dtype(bool):
return True
return False # otherwise
def __getitem__(self, ix):
"""
"""
#
# check special cases
#
assert ix is not None, "index is None!"
if self.position_index:
return ix
# boolean indexing ?
if is_DimArray(ix):
ix = ix.values
if type(ix) in (np.ndarray, ) and ix.dtype is np.dtype(bool):
return ix
# make sure (1,) is understood as 1 just as numpy would
elif type(ix) is tuple:
if len(ix) == 1:
ix = ix[0]
# else:
# raise TypeError("index not understood: did you mean a `slice`?")
#
# look up corresponding numpy indices
#
# e.g. 45:56
if type(ix) is slice:
res = self.slice(ix)
elif self._islist(ix):
res = map(self.locate, ix)
else:
res = self.locate(ix)
return res
def __getitem__(self, ix):
"""
"""
#
# check special cases
#
assert ix is not None, "index is None!"
if self.position_index:
return ix
# boolean indexing ?
if is_DimArray(ix):
ix = ix.values
if type(ix) in (np.ndarray,) and ix.dtype is np.dtype(bool):
return ix
# make sure (1,) is understood as 1 just as numpy would
elif type(ix) is tuple:
if len(ix) == 1:
ix = ix[0]
# else:
# raise TypeError("index not understood: did you mean a `slice`?")
#
# look up corresponding numpy indices
#
# e.g. 45:56
if type(ix) is slice:
res = self.slice(ix)
elif self._islist(ix):
res = map(self.locate, ix)
else:
res = self.locate(ix)
return res
def put(
obj,
val,
indices=None,
axis=0,
indexing="values",
tol=TOLERANCE,
convert=False,
inplace=False,
broadcast_arrays=True,
):
""" Put new values into DimArray (inplace)
parameters:
-----------
obj: DimArray (do not provide if method bound to class instance)
val: value to put in: scalar or array-like with appropriate shape or DimArray
indices, optional: see `take` for indexing rules. indices may be omitted if
val is a DimArray (will then deduce `indices` from its axes)
axis: for single index (see help on `take`)
indexing : "position", "values"
convert: convert array to val's type
inplace: True
broadcast_arrays: See documentation on `take`.
returns:
--------
None: (inplace modification)
Examples:
---------
>>> a = DimArray(np.zeros((2,2)), [('d0',['a','b']), ('d1',[10.,20.])])
Index by values
>>> b = a.put(1, indices={'d0': 'b'})
>>> b
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 0.],
[ 1., 1.]])
>>> a['b'] = 2 # slicing equivalent
>>> a
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 0.],
[ 2., 2.]])
Index by position
>>> b = a.put(3, indices=1, axis='d1', indexing="position")
>>> b
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 3.],
[ 2., 3.]])
>>> a.ix[:,1] = 4 # slicing equivalent
>>> a
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 4.],
[ 2., 4.]])
Multi-dimension, multi-index
>>> b = a.put(5, indices={'d0':'b', 'd1':[10.]})
>>> b
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 4.],
[ 5., 4.]])
>>> a["b",[10]] = 6
>>> a
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 4.],
[ 6., 4.]])
Inplace
>>> a.put(6, indices={'d0':'b', 'd1':[10.]}, inplace=True)
Multi-Index tests (not straightforward to get matlab-like behaviour)
>>> big = DimArray(np.zeros((2,3,4,5)))
>>> indices = {'x0':0 ,'x1':[2,1],'x3':[1,4]}
>>> sub = big.take(indices, broadcast_arrays=False)*0
>>> sub.dims == ('x1','x2','x3')
True
>>> sub.shape == (2,4,2)
True
>>> big.put(sub+1, indices, inplace=True, broadcast_arrays=False)
>>> sub2 = big.take(indices, broadcast_arrays=False)
>>> np.all(sub+1 == sub2)
True
"""
assert indexing in ("position", "values"), "invalid mode: " + repr(indexing)
if indices is None:
# DimArray as subarray: infer indices from its axes
if is_DimArray(val):
indices = {ax.name: ax.values for ax in val.axes}
broadcast_arrays = False
elif np.isscalar(val):
raise ValueError("indices must be provided for non-DimArray or non-matching shape. See also `fill` method.")
else:
raise ValueError("indices must be provided for non-DimArray or non-matching shape")
else:
indices = _fill_ellipsis(indices, obj.ndim)
# SPECIAL CASE: full scale boolean array
if is_boolean_index(indices, obj.shape):
return _put(obj, val, np.asarray(indices), inplace=inplace, convert=convert)
indices_numpy = obj.axes.loc(indices, axis=axis, position_index=(indexing == "position"), tol=tol)
# do nothing for full-array, boolean indexing
# if len(indices_numpy) == 1 and isinstance
# Convert to matlab-like indexing
if not broadcast_arrays:
indices_array = array_indices(indices_numpy, obj.shape)
indices_numpy_ = np.ix_(*indices_array)
shp = [len(ix) for ix in indices_array] # get an idea of the shape
## ...first check that val's shape is consistent with originally required indices
# if DimArray, transpose to the right shape
if is_DimArray(val):
newdims = [d for d in obj.dims if d in val.dims] + [d for d in val.dims if d not in obj.dims]
val = val.transpose(newdims)
# only check for n-d array of size and dimensions > 1
if np.size(val) > 1 and np.ndim(val) > 1 and np.any(np.array(shp) > 1):
shp1 = [d for d in shp if d > 1]
shp2 = [d for d in np.shape(val) if d > 1]
if shp1 != shp2:
raise ValueError(
"array is not broadcastable to correct shape (got values: {} but inferred from indices {})".format(
shp2, shp1
)
)
#
# # ...then reshape to new matlab-like form
if np.isscalar(val):
val_ = val
elif np.size(val) == 1:
val_ = np.squeeze(val)
else:
val = np.asarray(val)
val_ = np.reshape(val, shp)
else:
val_ = val
indices_numpy_ = indices_numpy
return _put(obj, val_, indices_numpy_, inplace=inplace, convert=convert)
def put(obj,
val,
indices=None,
axis=0,
indexing="values",
tol=TOLERANCE,
convert=False,
inplace=False,
broadcast_arrays=True):
""" Put new values into DimArray (inplace)
parameters:
-----------
obj: DimArray (do not provide if method bound to class instance)
val: value to put in: scalar or array-like with appropriate shape or DimArray
indices, optional: see `take` for indexing rules. indices may be omitted if
val is a DimArray (will then deduce `indices` from its axes)
axis: for single index (see help on `take`)
indexing : "position", "values"
convert: convert array to val's type
inplace: True
broadcast_arrays: See documentation on `take`.
returns:
--------
None: (inplace modification)
Examples:
---------
>>> a = DimArray(np.zeros((2,2)), [('d0',['a','b']), ('d1',[10.,20.])])
Index by values
>>> b = a.put(1, indices={'d0': 'b'})
>>> b
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 0.],
[ 1., 1.]])
>>> a['b'] = 2 # slicing equivalent
>>> a
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 0.],
[ 2., 2.]])
Index by position
>>> b = a.put(3, indices=1, axis='d1', indexing="position")
>>> b
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 3.],
[ 2., 3.]])
>>> a.ix[:,1] = 4 # slicing equivalent
>>> a
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 4.],
[ 2., 4.]])
Multi-dimension, multi-index
>>> b = a.put(5, indices={'d0':'b', 'd1':[10.]})
>>> b
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 4.],
[ 5., 4.]])
>>> a["b",[10]] = 6
>>> a
dimarray: 4 non-null elements (0 null)
dimensions: 'd0', 'd1'
0 / d0 (2): a to b
1 / d1 (2): 10.0 to 20.0
array([[ 0., 4.],
[ 6., 4.]])
Inplace
>>> a.put(6, indices={'d0':'b', 'd1':[10.]}, inplace=True)
Multi-Index tests (not straightforward to get matlab-like behaviour)
>>> big = DimArray(np.zeros((2,3,4,5)))
>>> indices = {'x0':0 ,'x1':[2,1],'x3':[1,4]}
>>> sub = big.take(indices, broadcast_arrays=False)*0
>>> sub.dims == ('x1','x2','x3')
True
>>> sub.shape == (2,4,2)
True
>>> big.put(sub+1, indices, inplace=True, broadcast_arrays=False)
>>> sub2 = big.take(indices, broadcast_arrays=False)
>>> np.all(sub+1 == sub2)
True
"""
assert indexing in ("position",
"values"), "invalid mode: " + repr(indexing)
if indices is None:
# DimArray as subarray: infer indices from its axes
if is_DimArray(val):
indices = {ax.name: ax.values for ax in val.axes}
broadcast_arrays = False
elif np.isscalar(val):
raise ValueError(
"indices must be provided for non-DimArray or non-matching shape. See also `fill` method."
)
else:
raise ValueError(
"indices must be provided for non-DimArray or non-matching shape"
)
else:
indices = _fill_ellipsis(indices, obj.ndim)
# SPECIAL CASE: full scale boolean array
if is_boolean_index(indices, obj.shape):
return _put(obj,
val,
np.asarray(indices),
inplace=inplace,
convert=convert)
indices_numpy = obj.axes.loc(indices,
axis=axis,
position_index=(indexing == "position"),
tol=tol)
# do nothing for full-array, boolean indexing
#if len(indices_numpy) == 1 and isinstance
# Convert to matlab-like indexing
if not broadcast_arrays:
indices_array = array_indices(indices_numpy, obj.shape)
indices_numpy_ = np.ix_(*indices_array)
shp = [len(ix) for ix in indices_array] # get an idea of the shape
## ...first check that val's shape is consistent with originally required indices
# if DimArray, transpose to the right shape
if is_DimArray(val):
newdims = [d for d in obj.dims if d in val.dims
] + [d for d in val.dims if d not in obj.dims]
val = val.transpose(newdims)
# only check for n-d array of size and dimensions > 1
if np.size(val) > 1 and np.ndim(val) > 1 and np.any(np.array(shp) > 1):
shp1 = [d for d in shp if d > 1]
shp2 = [d for d in np.shape(val) if d > 1]
if shp1 != shp2:
raise ValueError(
'array is not broadcastable to correct shape (got values: {} but inferred from indices {})'
.format(shp2, shp1))
#
# # ...then reshape to new matlab-like form
if np.isscalar(val):
val_ = val
elif np.size(val) == 1:
val_ = np.squeeze(val)
else:
val = np.asarray(val)
val_ = np.reshape(val, shp)
else:
val_ = val
indices_numpy_ = indices_numpy
return _put(obj, val_, indices_numpy_, inplace=inplace, convert=convert)
def broadcast(self, other):
""" broadcast the array along a set of axes by repeating it as necessay
other : DimArray or Axes objects or ordered Dictionary of axis values
Examples:
--------
Create some dummy data:
# ...create some dummy data:
>>> lon = np.linspace(10, 30, 2)
>>> lat = np.linspace(10, 50, 3)
>>> time = np.arange(1950,1955)
>>> ts = da.DimArray.from_kw(np.arange(5), time=time)
>>> cube = da.DimArray.from_kw(np.zeros((3,2,5)), lon=lon, lat=lat, time=time) # lat x lon x time
>>> cube.axes
dimensions: 'lat', 'lon', 'time'
0 / lat (3): 10.0 to 50.0
1 / lon (2): 10.0 to 30.0
2 / time (5): 1950 to 1954
# ...broadcast timeseries to 3D data
>>> ts3D = ts.broadcast(cube) # lat x lon x time
>>> ts3D
dimarray: 30 non-null elements (0 null)
dimensions: 'lat', 'lon', 'time'
0 / lat (3): 10.0 to 50.0
1 / lon (2): 10.0 to 30.0
2 / time (5): 1950 to 1954
array([[[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4]],
<BLANKLINE>
[[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4]],
<BLANKLINE>
[[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4]]])
"""
# Input as axes
if isinstance(other, list):
newaxes = other
# Or as DimArray
elif is_DimArray(other):
newaxes = other.axes
# Or as OrderedDict of axis names, axis values
elif isinstance(other, OrderedDict):
newaxes = [Axis(other[k], k) for k in other]
else:
raise TypeError("should be a DimArray, a list of Axis objects or an OrderedDict of Axis objects")
if len(newaxes) > 0 and not isinstance(newaxes[0], Axis): # just check the first as basic test
raise TypeError("should be a DimArray, a list of Axis objects or an OrderedDict of Axis objects")
newshape = [ax.name for ax in newaxes]
# First give it the right shape
newobj = self.reshape(newshape)
# Then repeat along axes
#for newaxis in newaxes:
for newaxis in reversed(newaxes): # should be faster ( CHECK )
if newobj.axes[newaxis.name].size == 1 and newaxis.size != 1:
newobj = newobj.repeat(newaxis.values, axis=newaxis.name)
return newobj
