mirror of
https://github.com/meteoinfo/MeteoInfo.git
synced 2025-12-08 20:36:05 +00:00
251 lines
6.4 KiB
Python
251 lines
6.4 KiB
Python
# coding=utf-8
|
|
|
|
from .numeric import asarray, array, isscalar
|
|
from ._ndarray import NDArray
|
|
from .stride_tricks import broadcast_arrays
|
|
from org.meteoinfo.ndarray.math import ArrayUtil, ArrayMath
|
|
|
|
|
|
__all__ = ['cumprod', 'cumsum', 'ndim', 'nonzero', 'prod', 'ravel', 'searchsorted', 'sum',
|
|
'where']
|
|
|
|
|
|
def ndim(a):
|
|
"""
|
|
Return the number of dimensions of an array.
|
|
|
|
Parameters
|
|
----------
|
|
a : array_like
|
|
Input array. If it is not already a ndarray, a conversion is
|
|
attempted.
|
|
|
|
Returns
|
|
-------
|
|
number_of_dimensions : int
|
|
The number of dimensions in `a`. Scalars are zero-dimensional.
|
|
|
|
See Also
|
|
--------
|
|
ndarray.ndim : equivalent method
|
|
shape : dimensions of array
|
|
ndarray.shape : dimensions of array
|
|
|
|
Examples
|
|
--------
|
|
>>> np.ndim([[1,2,3],[4,5,6]])
|
|
2
|
|
>>> np.ndim(np.array([[1,2,3],[4,5,6]]))
|
|
2
|
|
>>> np.ndim(1)
|
|
0
|
|
"""
|
|
if isscalar(a):
|
|
return 0
|
|
|
|
try:
|
|
return a.ndim
|
|
except AttributeError:
|
|
return asarray(a).ndim
|
|
|
|
|
|
def ravel(a):
|
|
"""
|
|
Return a contiguous flattened array.
|
|
|
|
:param a: (*array*) Input array.
|
|
:return: A contiguous flattened array.
|
|
"""
|
|
if isinstance(a, (list, tuple)):
|
|
a = array(a)
|
|
|
|
return a.ravel()
|
|
|
|
|
|
def nonzero(a):
|
|
"""
|
|
Return the indices of the elements that are non-zero.
|
|
|
|
Returns a tuple of arrays, one for each dimension of a, containing the indices of the
|
|
non-zero elements in that dimension.
|
|
|
|
:param a: (*array_like*) Input array.
|
|
|
|
:returns: (*tuple*) Indices of elements that are non-zero.
|
|
"""
|
|
if isinstance(a, list):
|
|
a = array(a)
|
|
ra = ArrayUtil.nonzero(a.asarray())
|
|
if ra is None:
|
|
return tuple([NDArray([]).astype('int')])
|
|
|
|
r = []
|
|
for aa in ra:
|
|
r.append(NDArray(aa))
|
|
return tuple(r)
|
|
|
|
|
|
def where(condition, *args):
|
|
"""
|
|
Return elements, either from x or y, depending on condition.
|
|
|
|
If only condition is given, return condition.nonzero().
|
|
|
|
Parameters
|
|
----------
|
|
condition : `array_like, bool`
|
|
Where True, yield x, otherwise yield y.
|
|
|
|
x, y : `array_like`
|
|
Values from which to choose. x, y and condition need to be broadcastable to some shape.
|
|
|
|
Returns
|
|
-------
|
|
`array`
|
|
An array with elements from x where condition is True, and elements from y elsewhere.
|
|
"""
|
|
if len(args) == 0:
|
|
return nonzero(condition)
|
|
|
|
x = args[0]
|
|
y = args[1]
|
|
if isinstance(condition, bool):
|
|
return x if condition else y
|
|
else:
|
|
condition = asarray(condition)
|
|
x = asarray(x)
|
|
y = asarray(y)
|
|
x, y = broadcast_arrays(x, y)
|
|
r = ArrayUtil.where(condition._array, x._array, y._array)
|
|
return condition.array_wrap(r)
|
|
|
|
|
|
def searchsorted(a, v, side='left', sorter=None):
|
|
"""
|
|
Find indices where elements should be inserted to maintain order.
|
|
:param a: (*array_like*) Input 1-D array. If sorter is None, then it must be sorted in ascending order,
|
|
otherwise sorter must be an array of indices that sort it.
|
|
:param v: (*array_like*) Values to insert into a.
|
|
:param side: (*str*) [left | right], default is `left`. If `left`, the index of the first suitable location found is given.
|
|
If `right`, return the last such index. If there is no suitable index, return either 0 or N (where N
|
|
is the length of a).
|
|
:param sorter: (*array_like*) Optional array of integer indices that sort array a into ascending order.
|
|
They are typically the result of argsort.
|
|
:return: (*array_like*) Array of insertion points with the same shape as v.
|
|
"""
|
|
if isinstance(a, (list, tuple)):
|
|
a = array(a).asarray()
|
|
elif isinstance(a, NDArray):
|
|
a = a.asarray()
|
|
|
|
if isinstance(v, (list, tuple)):
|
|
v = array(v).asarray()
|
|
elif isinstance(v, NDArray):
|
|
v = v.asarray()
|
|
|
|
left = True if side == 'left' else False
|
|
r = ArrayUtil.searchSorted(a, v, left)
|
|
if isinstance(r, int):
|
|
return r
|
|
else:
|
|
return NDArray(r)
|
|
|
|
|
|
def sum(x, axis=None):
|
|
"""
|
|
Sum of array elements over a given axis.
|
|
|
|
:param x: (*array_like or list*) Input values.
|
|
:param axis: (*int*) Axis along which the standard deviation is computed.
|
|
The default is to compute the standard deviation of the flattened array.
|
|
|
|
:returns: (*array_like*) Sum result
|
|
"""
|
|
if isinstance(x, (list, tuple)):
|
|
if isinstance(x[0], NDArray):
|
|
a = []
|
|
for xx in x:
|
|
a.append(xx.asarray())
|
|
r = ArrayMath.sum(a)
|
|
return x[0].array_wrap(r)
|
|
else:
|
|
x = array(x)
|
|
|
|
if axis is None:
|
|
r = ArrayMath.sum(x.asarray())
|
|
return r
|
|
else:
|
|
r = x.sum(axis)
|
|
return x.array_wrap(r, axis)
|
|
|
|
|
|
def prod(x, axis=None):
|
|
"""
|
|
Product of array elements over a given axis.
|
|
|
|
:param x: (*array_like or list*) Input values.
|
|
:param axis: (*int*) Axis along which the standard deviation is computed.
|
|
The default is to compute the standard deviation of the flattened array.
|
|
|
|
:returns: (*array_like*) Product result
|
|
"""
|
|
if isinstance(x, (list, tuple)):
|
|
x = array(x)
|
|
|
|
if axis is None:
|
|
r = ArrayMath.prod(x.asarray())
|
|
return r
|
|
else:
|
|
r = ArrayMath.prod(x.asarray(), axis)
|
|
x.array_wrap(r, axis)
|
|
|
|
|
|
def cumsum(a, axis=None):
|
|
"""
|
|
Return the cumulative summary of elements along a given axis.
|
|
|
|
Parameters
|
|
----------
|
|
a : array_like
|
|
Input array.
|
|
axis : int, optional
|
|
Axis along which the cumulative summary is computed. By default,
|
|
the input is flattened.
|
|
|
|
Returns
|
|
-------
|
|
cumsum : ndarray
|
|
A new array holding the result is returned.
|
|
"""
|
|
if axis is None:
|
|
r = ArrayMath.cumsum(a._array)
|
|
else:
|
|
r = ArrayMath.cumsum(a._array, axis)
|
|
|
|
return a.array_wrap(r, axis)
|
|
|
|
|
|
def cumprod(a, axis=None):
|
|
"""
|
|
Return the cumulative product of elements along a given axis.
|
|
|
|
Parameters
|
|
----------
|
|
a : array_like
|
|
Input array.
|
|
axis : int, optional
|
|
Axis along which the cumulative product is computed. By default,
|
|
the input is flattened.
|
|
|
|
Returns
|
|
-------
|
|
cumprod : ndarray
|
|
A new array holding the result is returned.
|
|
"""
|
|
if axis is None:
|
|
r = ArrayMath.cumprod(a._array)
|
|
else:
|
|
r = ArrayMath.cumprod(a._array, axis)
|
|
|
|
return a.array_wrap(r, axis)
|