def _median_with_nan(values, *args, **kwargs):
""" replace "median" if skipna is False
numpy's median ignore NaNs as long as less than 50%
modify this behaviour and return NaN just as any other operation would
"""
if _hasbottleneck:
result = bottleneck.median(values, *args, **kwargs)
else:
result = np.median(values, *args, **kwargs)
if anynan(values):
if np.size(result) == 1:
result = np.nan
else:
axis = kwargs.pop('axis', None)
nans = anynan(values, axis=axis) # determine where the nans should be
result[nans] = np.nan
return result
def __call__(self, values, *args, **kwargs):
# transform numpy array to masked array if needed
if anynan(values):
values = np.ma.array(values, mask=np.isnan(values))
func = getattr(np.ma, self.__name__)
else:
func = getattr(np, self.__name__)
result = func(values, *args, **kwargs)
# transform back to numpy array
if np.ma.isMaskedArray(result):
result = result.filled(np.nan)
return result
def _plot2D(self, funcname, *args, **kwargs):
""" generic plotting function for 2-D plots
"""
if len(self.dims) != 2:
raise NotImplementedError(funcname+" can only be called on two-dimensional dimarrays.")
import matplotlib.pyplot as plt
#ax = plt.gca()
if 'ax' in kwargs:
ax = kwargs.pop('ax')
else:
ax = plt.gca()
colorbar = kwargs.pop('colorbar', False)
# get the actual plotting function
function = getattr(ax, funcname)
# extract information about axis value and labels
# e.g. transform non-numeric data to float and
# set appropriate labels afterwards.
xval, xticks, xticklab, xlab = _get_axis_labels(self.axes[1])
yval, yticks, yticklab, ylab = _get_axis_labels(self.axes[0])
values = self.values
# pcolor does not work with nans
if (funcname == 'pcolormesh' or funcname == 'pcolor') and anynan(values):
values = np.ma.array(values, mask=np.isnan(values))
# make the plot
pc = function(xval, yval, values, **kwargs)
# add labels
if xlab is not None: ax.set_xlabel(xlab)
if ylab is not None: ax.set_ylabel(ylab)
if xticklab is not None: ax.set_xticklabels(xticklab)
if xticks is not None: ax.set_xticks(xticks)
if yticklab is not None: ax.set_yticklabels(yticklab)
if yticks is not None: ax.set_yticks(yticks)
# add colorbar?
if colorbar:
plt.colorbar(pc, ax=ax)
return pc
def assert_equal_dimarrays(actual, expected, metadata=True, approx=False):
assert isinstance(expected, da.DimArray)
assert isinstance(actual, da.DimArray)
assert actual.shape == expected.shape
assert actual.dims == expected.dims
# check the values
if approx or actual.dtype.kind == 'f' and anynan(actual.values):
assert_equal_values = assert_almost_equal
else:
assert_equal_values = assert_equal
assert_equal_values(actual.values, expected.values)
# check the axes
for actual_axis, expected_axis in zip(actual.axes, expected.axes):
assert_equal_axes(actual_axis, expected_axis, metadata=metadata)
# check the metadata
if metadata:
assert_equal_metadata(actual.attrs, expected.attrs)
def test_median(v, axis, use_bottleneck):
if not anynan(v):
assert_allclose(v.median(axis=axis), np.median(v.values, axis=axis))
else:
assert_allclose(v.median(axis=axis, skipna=True), np.ma.median(v.to_MaskedArray(), axis=axis))
def _get_weights(self, axis=None, mirror_nans=True, weights=None):
""" Build array of weights based on individual axes' weights attribute
Parameters
----------
axis : int or str, optional
if None a N-D weight is created, otherwise 1-D only
mirror_nans : bool, optional
mirror `values`'s NaNs in created weights (True by default)
{weights}
Returns
-------
full_weights : DimArray of weights, or None
See Also
--------
DimArray.mean, DimArray.var, DimArray.std
Examples
--------
Simpler examples are found Dimarray.mean's doc.
These are mostly meant as doctest.
>>> from dimarray import DimArray
>>> np.random.seed(0) # to make results reproducible
>>> v = DimArray(np.random.rand(3,2), axes=[[-80, 0, 80], [-180, 180]], dims=['lat','lon'])
>>> v._get_weights() # None
>>> w = lambda x : np.cos(np.radians(v.lat))
>>> v.axes['lat'].weights = w
>>> v._get_weights()
dimarray: 6 non-null elements (0 null)
0 / lat (3): -80 to 80
1 / lon (2): -180 to 180
array([[ 0.17364818, 0.17364818],
[ 1. , 1. ],
[ 0.17364818, 0.17364818]])
>>> v.axes['lat'].weights = None
>>> v._get_weights() # None
Can pass the weigthts as a dictionary
>>> v._get_weights(weights={{'lat':w}})
dimarray: 6 non-null elements (0 null)
0 / lat (3): -80 to 80
1 / lon (2): -180 to 180
array([[ 0.17364818, 0.17364818],
[ 1. , 1. ],
[ 0.17364818, 0.17364818]])
Or by indicating the axis
>>> v._get_weights(weights=w, axis='lat')
dimarray: 6 non-null elements (0 null)
0 / lat (3): -80 to 80
1 / lon (2): -180 to 180
array([[ 0.17364818, 0.17364818],
[ 1. , 1. ],
[ 0.17364818, 0.17364818]])
It possible to mix up axis and dict-input weights
>>> v.axes['lat'].weights = w
>>> v._get_weights(weights={{'lon':lambda x:((x+180)/100)**2}})
dimarray: 6 non-null elements (0 null)
0 / lat (3): -80 to 80
1 / lon (2): -180 to 180
array([[ 0. , 1.5628336],
[ 0. , 9. ],
[ 0. , 1.5628336]])
But if the transformation is required on one axis only,
weights on other axes are not accounted for since they
would not affect the result of the transformation.
>>> v._get_weights(weights={{'lon':lambda x:((x+180)/100)**2}}, axis='lat')
dimarray: 6 non-null elements (0 null)
0 / lat (3): -80 to 80
1 / lon (2): -180 to 180
array([[ 0.17364818, 0.17364818],
[ 1. , 1. ],
[ 0.17364818, 0.17364818]])
What is provided as argument overrides axis attribute.
Here when the weights are provided for one axis only:
>>> v._get_weights(weights=lambda x:((x+180)/100)**2, axis='lat')
dimarray: 6 non-null elements (0 null)
0 / lat (3): -80 to 80
1 / lon (2): -180 to 180
array([[ 1., 1.],
[ 1., 1.],
[ 4., 4.]])
Or when the weight is provided as a numpy array
>>> w = np.array([[ 1., 1.],
... [ 1., 1.],
... [ 4., 4.]])
>>> v._get_weights(weights=w)
dimarray: 6 non-null elements (0 null)
0 / lat (3): -80 to 80
1 / lon (2): -180 to 180
array([[ 1., 1.],
[ 1., 1.],
[ 4., 4.]])
Or even as DimArray (which will be broadcast with the appropriate shape)
>>> wa = DimArray(w[:,0], axes=[v.axes['lat']])
>>> v._get_weights(weights=wa)
dimarray: 6 non-null elements (0 null)
0 / lat (3): -80 to 80
1 / lon (2): -180 to 180
array([[ 1., 1.],
[ 1., 1.],
[ 4., 4.]])
Any nans in the data will induce 0 in the weights,
since the element will then be discarded or the result be NaN anyway.
>>> v.values[1, 0] = np.nan
>>> v._get_weights(weights=w)
dimarray: 6 non-null elements (0 null)
0 / lat (3): -80 to 80
1 / lon (2): -180 to 180
array([[ 1., 1.],
[ 0., 1.],
[ 4., 4.]])
"""
# avoid circular import
from dimarray import DimArray
# dims: axes from which the weights should be computed
if axis is None:
dims = self.dims
elif type(axis) is tuple:
dims = axis
else:
dims = (axis,)
# common case: no weight is provided and no Axis has weights
# if is enough to check on all axes to be reduced, weights on other axes
# make no difference to the reduction transformation
if weights is None and np.all([self.axes[dim].weights is None for dim in dims]):
return None
# for 1-D array, same as axis is None is like axis=0
if axis is None and self.ndim == 1:
axis = 0
# check `weights` type
# special case: if weights is DimArray, just use it as full weight
if isinstance(weights, DimArray):
full_weights = weights.broadcast(self.axes) # broadcast weights on appropriate dimension
# special case II: numpy array of appropriate dimension ==> just use it as full weight
elif isinstance(weights, np.ndarray) and weights.ndim > 1:
assert weights.shape == self.shape, "`weights` must be either 1-D or have same shape as DimArray"
full_weights = DimArray(weights, self.axes)
# make axis weights based by combining single-axis weights (standard form)
else:
array_type = isinstance(weights, np.ndarray) and weights.dim == 1
dict_type = isinstance(weights, dict)
assert weights is None or callable(weights) or array_type or dict_type, "unexpected type for `weights`"
assert weights is None or not (not dict_type and axis is None), "`weights` type incompatible with flattening operation, provide `axis` or use dict-type weights"
assert weights is None or not (not dict_type and type(axis) is tuple), "`weights` type incompatible with flattening operation, use dict-type weights"
# make weights dict-like (operation on single axis)
if not dict_type:
weights = {axis: weights}
# Create weights from the axes (only compute for the axes to be reduced)
full_weights = 1.
for axis in dims:
ax = self.axes[axis]
# get proper weights array: from `Axis` or user-provided on-the-fly
if weights is not None and axis in weights.keys():
axis_weights = weights[axis]
else:
axis_weights = ax.weights
# if weight is None, just skip it
if axis_weights is None:
continue
# broadcast and multiply weights for all axes involved
full_weights = ax._get_weights(axis_weights).reshape(dims) * full_weights
# SPECIAL CASE: NO weights (e.g. weights parameter provided as {} or {'lat':None} and no axis weights attributes):
if full_weights is 1.:
return None
full_weights = full_weights.broadcast(self.axes)
# set 0 in a where there are NaNs.
if mirror_nans and anynan(self.values):
full_weights.values[np.isnan(self.values)] = 0
assert full_weights is not None
return full_weights