def average(a, axis=None, weights=None, returned=False):
axis = _conv_axis_arg(axis)
ta = _makeMaskedArg(a)
maresult = numpy.ma.average(ta, axis, weights, returned)
axes, attributes, id, grid = _extractMetadata(a,
omit=axis,
omitall=(axis is None))
if returned:
if isinstance(maresult, tuple):
maresult, wresult = maresult
else:
# ok it's masked constant need to return both things by hand
wresult = numpy.ma.masked
F = getattr(a, "fill_value", 1.e20)
r1 = TransientVariable(maresult,
axes=axes,
attributes=attributes,
grid=grid,
id=id,
no_update_from=True,
fill_value=F)
if returned:
F = getattr(a, "fill_value", 1.e20)
w1 = TransientVariable(wresult,
axes=axes,
grid=grid,
id=id,
no_update_from=True,
fill_value=F)
return r1, w1
else:
return r1
def mask_OCNFRAC(mv, ocnfrac):
"""Applies a mask computed from the variable ocnfrac (which comes from the CAM variable OCNFRAC)
to the specified variable, and returns the resulting new variable. Both mv and ocnfrac should
have the same domain."""
if mv is None or ocnfrac is None:
return None
ocnmask = ocnfrac < 0.5
currentmask = mv.mask
if currentmask is numpy.ma.nomask:
newmask = ocnmask
else:
newmask = numpy.logical_or(currentmask, ocnmask)
jfpvar = TransientVariable(mv, mask=newmask)
return TransientVariable(mv, mask=newmask)
def masked_values(data,
value,
rtol=1.e-5,
atol=1.e-8,
copy=1,
savespace=0,
axes=None,
attributes=None,
id=None):
"""
Masked Values
Parameters
----------
masked_values(data, value, rtol=1.e-5, atol=1.e-8)
Create a masked array; mask is None if possible.
May share data values with original array, but not recommended.
Masked where abs(data-value)<= atol + rtol * abs(value)
"""
maresult = numpy.ma.masked_values(_makeMaskedArg(data),
value,
rtol=rtol,
atol=atol,
copy=copy)
axes, attributes, id, grid = _extractMetadata(data, axes, attributes, id)
F = getattr(data, "fill_value", 1.e20)
return TransientVariable(maresult,
axes=axes,
attributes=attributes,
id=id,
grid=grid,
fill_value=F)
def squeeze(x):
"call numpy.squeeze on ndarray and rebuild tvariable."
# ta = _makeMaskedArg(x)
maresult = numpy.squeeze(x._data)
axes, attributes, id, grid = _extractMetadata(x)
return TransientVariable(
maresult, axes=axes, attributes=attributes, grid=grid, id=id)
def masked_array(a,
mask=None,
fill_value=None,
axes=None,
attributes=None,
id=None):
"""
Masked Array
Parameters
----------
masked_array(a, mask=None) = array(a, mask=mask, copy=0, fill_value=fill_value)
Use fill_value(a) if None.
"""
maresult = numpy.ma.masked_array(_makeMaskedArg(a),
mask=mask,
fill_value=fill_value)
axes, attributes, id, grid = _extractMetadata(a, axes, attributes, id)
F = getattr(a, "fill_value", 1.e20)
return TransientVariable(maresult,
axes=axes,
attributes=attributes,
id=id,
grid=grid,
fill_value=F)
def resize(a, new_shape, axes=None, attributes=None, id=None, grid=None):
"""resize(a, new_shape)
Resize
Returns
-------
a new array with the specified shape.
The original array's total size can be any size.
"""
ignore, attributes, id, ignore = _extractMetadata(a, axes, attributes, id)
if axes is not None:
axesshape = [len(item) for item in axes]
if axesshape != list(new_shape):
raise CDMSError('axes must be shaped %s' % repr(new_shape))
ta = _makeMaskedArg(a)
maresult = numpy.ma.resize(ta, new_shape)
F = getattr(a, "fill_value", 1.e20)
return TransientVariable(maresult,
axes=axes,
attributes=attributes,
id=id,
grid=grid,
fill_value=F)
def right_shift(a, n):
"Right shift n bits"
ta = _makeMaskedArg(a)
tb = _makeMaskedArg(n)
maresult = numpy.ma.right_shift(ta, numpy.ma.filled(tb))
axes, attributes, id, grid = _extractMetadata(a)
return TransientVariable(
maresult, axes=axes, attributes=attributes, grid=grid, id=id)
def power(a, b, third=None):
"a**b"
ta = _makeMaskedArg(a)
tb = _makeMaskedArg(b)
maresult = numpy.ma.power(ta, tb, third)
axes, attributes, id, grid = _extractMetadata(a)
return TransientVariable(
maresult, axes=axes, attributes=attributes, grid=grid, id=id)
def product(a, axis=0, dtype=None):
"Product of elements along axis."
ta = _makeMaskedArg(a)
maresult = numpy.ma.product(ta, axis, dtype=dtype)
axes, attributes, id, grid = _extractMetadata(a, omit=axis)
F = getattr(a, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, grid=grid, id=id, fill_value=F)
def diagonal(a, offset=0, axis1=0, axis2=1):
"""diagonal(a, offset=0, axis1=0, axis2 = 1) returns the given
diagonals defined by the two dimensions of the array.
"""
F = getattr(a, "fill_value", 1.e20)
return TransientVariable(numpy.ma.diagonal(_makeMaskedArg(a), offset,
axis1, axis2),
fill_value=F)
def __call__(self, a, axis=0, **kwargs):
axis = _conv_axis_arg(axis)
ta = _makeMaskedArg(a)
maresult = self.mafunc(ta, axis=axis, **kwargs)
axes, attributes, id, grid = _extractMetadata(
a, omit=axis, omitall=(axis is None))
return TransientVariable(
maresult, axes=axes, attributes=attributes, id=id, grid=grid)
def ones(myshape, typecode=float, savespace=0, axes=None,
attributes=None, id=None, grid=None, dtype=None):
"""ones(n, typecode=float, savespace=0, axes=None, attributes=None, id=None) =
an array of all ones of the given length or shape."""
dtype = _convdtype(dtype, typecode)
maresult = numpy.ma.ones(myshape, dtype=dtype)
return TransientVariable(
maresult, axes=axes, attributes=attributes, id=id, grid=grid)
def outerproduct(a, b):
"""outerproduct(a,b) = {a[i]*b[j]}, has shape (len(a),len(b))"""
ta = asVariable(a, writeable=1)
tb = asVariable(b, writeable=1)
maresult = numpy.ma.outerproduct(ta, tb)
axes = (ta.getAxis(0), tb.getAxis(0))
F = getattr(a, "fill_value", 1.e20)
return TransientVariable(maresult, axes=axes, fill_value=F)
def __call__(self, a, **kwargs):
axes, attributes, id, grid = _extractMetadata(a)
maresult = self.mafunc(_makeMaskedArg(a), **kwargs)
return TransientVariable(maresult,
axes=axes,
attributes=attributes,
id=id,
grid=grid)
def masked_inside(x, v1, v2):
"x with mask of all values of x that are inside [v1,v2]"
tx = _makeMaskedArg(x)
maresult = numpy.ma.masked_inside(tx, v1, v2)
axes, attributes, id, grid = _extractMetadata(x)
F = getattr(x, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, id=id, grid=grid, fill_value=F)
def masked_not_equal(x, value):
"masked_not_equal(x, value) = x masked where x != value"
tx = _makeMaskedArg(x)
maresult = numpy.ma.masked_not_equal(tx, value)
axes, attributes, id, grid = _extractMetadata(x)
F = getattr(x, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, id=id, grid=grid, fill_value=F)
def masked_object(data, value, copy=1, savespace=0,
axes=None, attributes=None, id=None):
"Create array masked where exactly data equal to value"
maresult = numpy.ma.masked_object(_makeMaskedArg(data), value, copy=copy)
axes, attributes, id, grid = _extractMetadata(data, axes, attributes, id)
F = getattr(data, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, id=id, grid=grid, fill_value=F)
def compress(a, b):
__doc__ = numpy.ma.__doc__ # noqa
import warnings
warnings.warn(
"arguments order for compress function has changed\nit is now: " +
"MV2.copmress(array,condition), if your code seems to not react " +
"or act wrong to a call to compress, please check this", Warning)
return TransientVariable(numpy.ma.compress(a, b), copy=1)
def concatenate(arrays, axis=0, axisid=None, axisattributes=None):
"""Concatenate the arrays along the given axis. Give the extended axis the id and
attributes provided - by default, those of the first array."""
tarrays = [_makeMaskedArg(a) for a in arrays]
maresult = numpy.ma.concatenate(arrays, axis=axis)
if len(arrays) > 1:
varattributes = None
varid = None
axes = commonDomain(tarrays[0], tarrays[1], omit=axis)
grid = commonGrid(tarrays[0], tarrays[1], axes)
for i in range(len(arrays) - 2):
if axes is None:
break
axes = commonAxes(tarrays[i + 2], axes, omit=axis)
grid = commonGrid1(tarrays[i + 2], grid, axes)
else:
axes = tarrays[0].getAxisList()
varattributes = tarrays[0].attributes
varid = tarrays[0].id
if (isinstance(tarrays[0], TransientVariable)):
grid = tarrays[0].getGrid()
else:
grid = None
if axes is not None:
if axisid is None:
axisid = tarrays[0].getAxis(axis).id
allunitsequal = True
try:
allunits = tarrays[0].getAxis(axis).units
except BaseException:
allunits = None
for t in tarrays[1:]:
try:
tunits = t.getAxis(axis).units
except BaseException:
tunits = None
if tunits != allunits:
allunitsequal = False
if allunitsequal:
if axisattributes is None:
axisattributes = tarrays[0].getAxis(axis).attributes
axes[axis] = axisConcatenate([t.getAxis(axis) for t in tarrays],
axisid, axisattributes)
# If the grid doesn't match the axislist (e.g., catenation was on
# latitude) then omit it.
if grid is not None:
for item in grid.getAxisList():
if item not in axes:
grid = None
F = getattr(arrays[0], "fill_value", 1.e20)
return TransientVariable(maresult,
axes=axes,
attributes=varattributes,
id=varid,
grid=grid,
fill_value=F)
def accumulate(self, target, axis=0):
ttarget = _makeMaskedArg(target)
maresult = self.mafunc.accumulate(ttarget, axis=axis)
axes, attributes, id, grid = _extractMetadata(target, omit=axis)
return TransientVariable(maresult,
axes=axes,
attributes=attributes,
id=id,
grid=grid)
def where(condition, x, y):
"where(condition, x, y) is x where condition is true, y otherwise"
# axes = commonDomain(x,y)
# grid = commonGrid(x,y,axes)
maresult = numpy.ma.where(condition, _makeMaskedArg(x), _makeMaskedArg(y))
axes, attributes, id, grid = _extractMetadata(condition)
F = getattr(x, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, grid=grid, id=id, fill_value=F)
def __call__(self, a, b, **kwargs):
id = "variable_%i" % TransientVariable.variable_count
TransientVariable.variable_count += 1
axes = commonDomain(a, b)
grid = commonGrid(a, b, axes)
ta = _makeMaskedArg(a)
tb = _makeMaskedArg(b)
maresult = self.mafunc(ta, tb, **kwargs)
return TransientVariable(
maresult, axes=axes, grid=grid, no_update_from=True, id=id)
def sort(a, axis=-1):
maresult = numpy.ma.sort(a.asma(), axis)
axes, attributes, id, grid = _extractMetadata(a)
sortaxis = axes[axis]
if (grid is not None) and (sortaxis in grid.getAxisList()):
grid = None
axes[axis] = TransientAxis(numpy.arange(len(sortaxis)))
F = getattr(a, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, grid=grid, id=id, fill_value=F)
def sum(a, axis=None, fill_value=0, dtype=None):
"Sum of elements along a certain axis."
axis = _conv_axis_arg(axis)
ta = _makeMaskedArg(a)
maresult = numpy.ma.sum(ta, axis, dtype=dtype)
axes, attributes, id, grid = _extractMetadata(
a, omit=axis, omitall=(axis is None))
F = getattr(a, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, grid=grid, id=id, fill_value=F)
def masked_equal(x, value):
"""masked_equal(x, value) = x masked where x == value
For floating point consider masked_values(x, value) instead.
"""
tx = _makeMaskedArg(x)
maresult = numpy.ma.masked_equal(tx, value)
axes, attributes, id, grid = _extractMetadata(x)
F = getattr(x, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, id=id, grid=grid, fill_value=F)
def reshape(a, newshape, axes=None, attributes=None, id=None, grid=None):
ignore, attributes, id, ignore = _extractMetadata(a, axes, attributes, id)
if axes is not None:
axesshape = [len(item) for item in axes]
if axesshape != list(newshape):
raise CDMSError('axes must be shaped %s' % repr(newshape))
ta = _makeMaskedArg(a)
maresult = numpy.ma.reshape(ta, newshape)
F = getattr(a, "fill_value", 1.e20)
return TransientVariable(maresult, axes=axes, attributes=attributes,
id=id, grid=grid, no_update_from=True, fill_value=F)
def argsort(x, axis=-1, fill_value=None):
"""Treating masked values as if they have the value fill_value,
return sort indices for sorting along given axis.
if fill_value is None, use fill_value(x)
"""
tx = _makeMaskedArg(x)
maresult = numpy.ma.argsort(tx, axis=axis, fill_value=fill_value)
axes, attributes, id, grid = _extractMetadata(x)
F = getattr(x, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, id=id, grid=grid, fill_value=F)
def count(a, axis=None):
"Count of the non-masked elements in a, or along a certain axis."
if axis is None:
return numpy.ma.count(a, axis)
else:
ta = _makeMaskedArg(a)
maresult = numpy.ma.count(ta, axis)
axes, attributes, id, grid = _extractMetadata(a, omit=axis)
F = getattr(a, "fill_value", 1.e20)
return TransientVariable(
maresult, axes=axes, attributes=attributes, grid=grid, id=id, fill_value=F)
def choose(myindices, t):
"""Returns an array shaped like indices containing elements chosen
from t.
If an element of t is the special element masked, any element
of the result that "chooses" that element is masked.
The result has only the default axes.
"""
maresult = numpy.ma.choose(myindices, list(map(_makeMaskedArg, t)))
F = getattr(t, "fill_value", 1.e20)
return TransientVariable(maresult, fill_value=F)
def test_returns_annual_cycle_for_a_dataset_variable(self):
# Mock a Dataset object and get_climo_variable()
dataset_mock = Mock()
dataset_mock.get_climo_variable.return_value = TransientVariable(
data=np.zeros((2, 2)),
attributes={
"id": "PRECNT",
"long_name": "long_name",
"units": "units"
},
axes=[
TransientAxis(np.zeros(2), id="latitude"),
TransientAxis(np.zeros(2), id="longitude"),
],
)
# Generate expected and result
expected = TransientVariable(
data=np.zeros((12, 2, 2)),
attributes={
"id": "PRECNT",
"long_name": "long_name",
"units": "units"
},
axes=[
TransientAxis(
np.arange(1, 13),
id="time",
attributes={
"axis": "T",
"calendar": "gregorian"
},
),
TransientAxis(np.zeros(2), id="latitude"),
TransientAxis(np.zeros(2), id="longitude"),
],
)
result = _create_annual_cycle(dataset_mock, variable="PRECNT")
# Check data are equal
np.array_equal(result.data, expected.data)
# Check attributes are equal. Must delete "name" attribute since they differ.
result.deleteattribute("name")
expected.deleteattribute("name")
self.assertDictEqual(result.attributes, expected.attributes)
# Check time, latitude, and longitude axes are equal
np.array_equal(result.getAxis(0)[:], expected.getAxis(0)[:])
np.array_equal(result.getLatitude()[:], expected.getLatitude()[:])
np.array_equal(result.getLongitude()[:], expected.getLongitude()[:])
