Source code for dask.array.ma

from functools import wraps

import numpy as np

from ..base import normalize_token
from .core import (
    concatenate_lookup,
    tensordot_lookup,
    map_blocks,
    asanyarray,
    blockwise,
)
from .routines import _average
from ..utils import derived_from


@normalize_token.register(np.ma.masked_array)
def normalize_masked_array(x):
    data = normalize_token(x.data)
    mask = normalize_token(x.mask)
    fill_value = normalize_token(x.fill_value)
    return (data, mask, fill_value)


@concatenate_lookup.register(np.ma.masked_array)
def _concatenate(arrays, axis=0):
    out = np.ma.concatenate(arrays, axis=axis)
    fill_values = [i.fill_value for i in arrays if hasattr(i, "fill_value")]
    if any(isinstance(f, np.ndarray) for f in fill_values):
        raise ValueError(
            "Dask doesn't support masked array's with non-scalar `fill_value`s"
        )
    if fill_values:
        # If all the fill_values are the same copy over the fill value
        fill_values = np.unique(fill_values)
        if len(fill_values) == 1:
            out.fill_value = fill_values[0]
    return out


@tensordot_lookup.register(np.ma.masked_array)
def _tensordot(a, b, axes=2):
    # Much of this is stolen from numpy/core/numeric.py::tensordot
    # Please see license at https://github.com/numpy/numpy/blob/master/LICENSE.txt
    try:
        iter(axes)
    except TypeError:
        axes_a = list(range(-axes, 0))
        axes_b = list(range(0, axes))
    else:
        axes_a, axes_b = axes
    try:
        na = len(axes_a)
        axes_a = list(axes_a)
    except TypeError:
        axes_a = [axes_a]
        na = 1
    try:
        nb = len(axes_b)
        axes_b = list(axes_b)
    except TypeError:
        axes_b = [axes_b]
        nb = 1

    # a, b = asarray(a), asarray(b)  # <--- modified
    as_ = a.shape
    nda = a.ndim
    bs = b.shape
    ndb = b.ndim
    equal = True
    if na != nb:
        equal = False
    else:
        for k in range(na):
            if as_[axes_a[k]] != bs[axes_b[k]]:
                equal = False
                break
            if axes_a[k] < 0:
                axes_a[k] += nda
            if axes_b[k] < 0:
                axes_b[k] += ndb
    if not equal:
        raise ValueError("shape-mismatch for sum")

    # Move the axes to sum over to the end of "a"
    # and to the front of "b"
    notin = [k for k in range(nda) if k not in axes_a]
    newaxes_a = notin + axes_a
    N2 = 1
    for axis in axes_a:
        N2 *= as_[axis]
    newshape_a = (-1, N2)
    olda = [as_[axis] for axis in notin]

    notin = [k for k in range(ndb) if k not in axes_b]
    newaxes_b = axes_b + notin
    N2 = 1
    for axis in axes_b:
        N2 *= bs[axis]
    newshape_b = (N2, -1)
    oldb = [bs[axis] for axis in notin]

    at = a.transpose(newaxes_a).reshape(newshape_a)
    bt = b.transpose(newaxes_b).reshape(newshape_b)
    res = np.ma.dot(at, bt)
    return res.reshape(olda + oldb)


[docs]@derived_from(np.ma) def filled(a, fill_value=None): a = asanyarray(a) return a.map_blocks(np.ma.filled, fill_value=fill_value)
def _wrap_masked(f): @wraps(f) def _(a, value): a = asanyarray(a) value = asanyarray(value) ainds = tuple(range(a.ndim))[::-1] vinds = tuple(range(value.ndim))[::-1] oinds = max(ainds, vinds, key=len) return blockwise(f, oinds, a, ainds, value, vinds, dtype=a.dtype) return _ masked_greater = _wrap_masked(np.ma.masked_greater) masked_greater_equal = _wrap_masked(np.ma.masked_greater_equal) masked_less = _wrap_masked(np.ma.masked_less) masked_less_equal = _wrap_masked(np.ma.masked_less_equal) masked_not_equal = _wrap_masked(np.ma.masked_not_equal)
[docs]@derived_from(np.ma) def masked_equal(a, value): a = asanyarray(a) if getattr(value, "shape", ()): raise ValueError("da.ma.masked_equal doesn't support array `value`s") inds = tuple(range(a.ndim)) return blockwise(np.ma.masked_equal, inds, a, inds, value, (), dtype=a.dtype)
[docs]@derived_from(np.ma) def masked_invalid(a): return asanyarray(a).map_blocks(np.ma.masked_invalid)
[docs]@derived_from(np.ma) def masked_inside(x, v1, v2): x = asanyarray(x) return x.map_blocks(np.ma.masked_inside, v1, v2)
[docs]@derived_from(np.ma) def masked_outside(x, v1, v2): x = asanyarray(x) return x.map_blocks(np.ma.masked_outside, v1, v2)
[docs]@derived_from(np.ma) def masked_where(condition, a): cshape = getattr(condition, "shape", ()) if cshape and cshape != a.shape: raise IndexError( "Inconsistant shape between the condition and the " "input (got %s and %s)" % (cshape, a.shape) ) condition = asanyarray(condition) a = asanyarray(a) ainds = tuple(range(a.ndim)) cinds = tuple(range(condition.ndim)) return blockwise( np.ma.masked_where, ainds, condition, cinds, a, ainds, dtype=a.dtype )
[docs]@derived_from(np.ma) def masked_values(x, value, rtol=1e-05, atol=1e-08, shrink=True): x = asanyarray(x) if getattr(value, "shape", ()): raise ValueError("da.ma.masked_values doesn't support array `value`s") return map_blocks( np.ma.masked_values, x, value, rtol=rtol, atol=atol, shrink=shrink )
[docs]@derived_from(np.ma) def fix_invalid(a, fill_value=None): a = asanyarray(a) return a.map_blocks(np.ma.fix_invalid, fill_value=fill_value)
[docs]@derived_from(np.ma) def getdata(a): a = asanyarray(a) return a.map_blocks(np.ma.getdata)
[docs]@derived_from(np.ma) def getmaskarray(a): a = asanyarray(a) return a.map_blocks(np.ma.getmaskarray)
def _masked_array(data, mask=np.ma.nomask, **kwargs): dtype = kwargs.pop("masked_dtype", None) return np.ma.masked_array(data, mask=mask, dtype=dtype, **kwargs)
[docs]@derived_from(np.ma) def masked_array(data, mask=np.ma.nomask, fill_value=None, **kwargs): data = asanyarray(data) inds = tuple(range(data.ndim)) arginds = [inds, data, inds] if getattr(fill_value, "shape", ()): raise ValueError("non-scalar fill_value not supported") kwargs["fill_value"] = fill_value if mask is not np.ma.nomask: mask = asanyarray(mask) if mask.size == 1: mask = mask.reshape((1,) * data.ndim) elif data.shape != mask.shape: raise np.ma.MaskError( "Mask and data not compatible: data shape " "is %s, and mask shape is " "%s." % (repr(data.shape), repr(mask.shape)) ) arginds.extend([mask, inds]) if "dtype" in kwargs: kwargs["masked_dtype"] = kwargs["dtype"] else: kwargs["dtype"] = data.dtype return blockwise(_masked_array, *arginds, **kwargs)
def _set_fill_value(x, fill_value): if isinstance(x, np.ma.masked_array): x = x.copy() np.ma.set_fill_value(x, fill_value=fill_value) return x
[docs]@derived_from(np.ma) def set_fill_value(a, fill_value): a = asanyarray(a) if getattr(fill_value, "shape", ()): raise ValueError("da.ma.set_fill_value doesn't support array `value`s") fill_value = np.ma.core._check_fill_value(fill_value, a.dtype) res = a.map_blocks(_set_fill_value, fill_value) a.dask = res.dask a.name = res.name
[docs]@derived_from(np.ma) def average(a, axis=None, weights=None, returned=False): return _average(a, axis, weights, returned, is_masked=True)