def put(arr, ids, values): """ The opposite of `take`. The values of `arr` at the locations specified by `ids` are set to the corresponding value of `values`. The following is to test improvments to puts with masked arrays. Places in the code were assuming incorrect put behavior. >>> maskValue = 999999 >>> arr = zeros(3, 'l') >>> ids = MA.masked_values((2, maskValue), maskValue) >>> values = MA.masked_values((4, maskValue), maskValue) >>> put(arr, ids, values) ## this should work >>> print arr [0 0 4] >>> arr = MA.masked_values((maskValue, 5, 10), maskValue) >>> ids = MA.masked_values((2, maskValue), maskValue) >>> values = MA.masked_values((4, maskValue), maskValue) >>> put(arr, ids, values) >>> print arr ## works as expected [-- 5 4] >>> arr = MA.masked_values((maskValue, 5, 10), maskValue) >>> ids = MA.masked_values((maskValue, 2), maskValue) >>> values = MA.masked_values((4, maskValue), maskValue) >>> put(arr, ids, values) >>> print arr ## should be [-- 5 --] maybe?? [-- 5 999999] """ if _isPhysical(arr): arr.put(ids, values) elif MA.isMaskedArray(arr): if NUMERIX.sometrue(MA.getmaskarray(ids)): if numpy_version == 'old': pvalues = MA.array(values, mask=MA.getmaskarray(ids)) else: pvalues = MA.array(values.filled(), mask=MA.getmaskarray(ids)) MA.put(arr, ids.compressed(), pvalues.compressed()) else: MA.put(arr, ids, values) elif MA.isMaskedArray(ids): NUMERIX.put(arr, ids.compressed(), MA.array(values, mask=MA.getmaskarray(ids)).compressed()) else: NUMERIX.put(arr, ids, values)
def take(a, indices, axis=0, fill_value=None): """ Selects the elements of `a` corresponding to `indices`. """ if _isPhysical(a): taken = a.take(indices, axis=axis) elif type(indices) is type(MA.array((0))): ## Replaces `MA.take`. `MA.take` does not always work when ## `indices` is a masked array. ## taken = MA.take(a, MA.filled(indices, 0), axis=axis) mask = MA.getmask(indices) if mask is not MA.nomask: mask = MA.getmaskarray(indices) if taken.shape != mask.shape: mask = MA.repeat(mask[NewAxis, ...], taken.shape[0], axis=0) mask = MA.mask_or(MA.getmask(taken), mask) if mask is not MA.nomask: taken = MA.array(data=taken, mask=mask) else: if MA.getmask(taken) is MA.nomask and numpy_version == 'old': # numpy 1.1 returns normal array when masked array is filled taken = taken.filled() elif type(a) in (type(array((0))), type(()), type([])): taken = NUMERIX.take(a, indices, axis=axis) elif type(a) is type(MA.array((0))): taken = MA.take(a, indices, axis=axis) else: raise TypeError, 'cannot take from %s object: %s' % (type(a), `a`) if fill_value is not None and type(taken) is type(MA.array((0))): taken = taken.filled(fill_value=fill_value) return taken
def take(a, indices, axis=0, fill_value=None): """ Selects the elements of `a` corresponding to `indices`. """ if _isPhysical(a): taken = a.take(indices, axis=axis) elif type(indices) is type(MA.array((0))): ## Replaces `MA.take`. `MA.take` does not always work when ## `indices` is a masked array. ## taken = MA.take(a, MA.filled(indices, 0), axis=axis) mask = MA.getmask(indices) if mask is not MA.nomask: mask = MA.getmaskarray(indices) if taken.shape != mask.shape: mask = MA.repeat(mask[NewAxis, ...], taken.shape[0], axis=0) mask = MA.mask_or(MA.getmask(taken), mask) if mask is not MA.nomask: taken = MA.array(data=taken, mask=mask) else: if MA.getmask(taken) is MA.nomask and numpy_version == 'old': # numpy 1.1 returns normal array when masked array is filled taken = taken.filled() elif type(a) in (type(array((0))), type(()), type([])): taken = NUMERIX.take(a, indices, axis=axis) elif type(a) is type(MA.array((0))): taken = MA.take(a, indices, axis=axis) else: raise TypeError, 'cannot take from %s object: %s' % (type(a), ` a `) if fill_value is not None and type(taken) is type(MA.array((0))): taken = taken.filled(fill_value=fill_value) return taken