mirror of
https://github.com/meteoinfo/MeteoInfo.git
synced 2025-12-08 20:36:05 +00:00
294 lines
9.6 KiB
Python
294 lines
9.6 KiB
Python
|
|
from org.meteoinfo.ndarray.math import ArrayUtil
|
|
from .. import core as np
|
|
|
|
|
|
__all__ = ['unique', 'ediff1d', 'intersect1d']
|
|
|
|
|
|
def _unique1d(ar, return_index=False, return_inverse=False,
|
|
return_counts=False):
|
|
"""
|
|
Find the unique elements of an array, ignoring shape.
|
|
"""
|
|
ar = np.asanyarray(ar).flatten()
|
|
|
|
optional_indices = return_index or return_inverse
|
|
|
|
if optional_indices:
|
|
perm = ar.argsort()
|
|
aux = ar[perm]
|
|
else:
|
|
ar.sort()
|
|
aux = ar
|
|
mask = np.empty(aux.shape, dtype=np.bool_)
|
|
mask[:1] = True
|
|
if aux.shape[0] > 0 and aux.dtype.kind in "cfmM" and np.isnan(aux[-1]):
|
|
if aux.dtype.kind == "c": # for complex all NaNs are considered equivalent
|
|
aux_firstnan = np.searchsorted(np.isnan(aux), True, side='left')
|
|
else:
|
|
aux_firstnan = np.searchsorted(aux, aux[-1], side='left')
|
|
if aux_firstnan > 0:
|
|
mask[1:aux_firstnan] = (
|
|
aux[1:aux_firstnan] != aux[:aux_firstnan - 1])
|
|
mask[aux_firstnan] = True
|
|
mask[aux_firstnan + 1:] = False
|
|
else:
|
|
mask[1:] = aux[1:] != aux[:-1]
|
|
|
|
ret = (aux[mask],)
|
|
if return_index:
|
|
ret += (perm[mask],)
|
|
if return_inverse:
|
|
imask = np.cumsum(mask) - 1
|
|
inv_idx = np.empty(mask.shape, dtype=np.intp)
|
|
inv_idx[perm] = imask
|
|
ret += (inv_idx,)
|
|
if return_counts:
|
|
idx = np.concatenate(np.nonzero(mask) + ([mask.size],))
|
|
ret += (np.diff(idx),)
|
|
return ret
|
|
|
|
|
|
def unique(a, return_index=False, return_inverse=False, return_counts=False, axis=None):
|
|
"""
|
|
Find the unique elements of an array.
|
|
|
|
Returns the sorted unique elements of an array. There are three optional
|
|
outputs in addition to the unique elements:
|
|
|
|
* the indices of the input array that give the unique values
|
|
* the indices of the unique array that reconstruct the input array
|
|
* the number of times each unique value comes up in the input array
|
|
|
|
Parameters
|
|
----------
|
|
ar : array_like
|
|
Input array. Unless `axis` is specified, this will be flattened if it
|
|
is not already 1-D.
|
|
return_index : bool, optional
|
|
If True, also return the indices of `ar` (along the specified axis,
|
|
if provided, or in the flattened array) that result in the unique array.
|
|
return_inverse : bool, optional
|
|
If True, also return the indices of the unique array (for the specified
|
|
axis, if provided) that can be used to reconstruct `ar`.
|
|
return_counts : bool, optional
|
|
If True, also return the number of times each unique item appears
|
|
in `ar`.
|
|
axis : int or None, optional
|
|
The axis to operate on. If None, `ar` will be flattened. If an integer,
|
|
the subarrays indexed by the given axis will be flattened and treated
|
|
as the elements of a 1-D array with the dimension of the given axis,
|
|
see the notes for more details. Object arrays or structured arrays
|
|
that contain objects are not supported if the `axis` kwarg is used. The
|
|
default is None.
|
|
|
|
Returns
|
|
-------
|
|
unique : ndarray
|
|
The sorted unique values.
|
|
unique_indices : ndarray, optional
|
|
The indices of the first occurrences of the unique values in the
|
|
original array. Only provided if `return_index` is True.
|
|
unique_inverse : ndarray, optional
|
|
The indices to reconstruct the original array from the
|
|
unique array. Only provided if `return_inverse` is True.
|
|
unique_counts : ndarray, optional
|
|
The number of times each of the unique values comes up in the
|
|
original array. Only provided if `return_counts` is True.
|
|
|
|
See Also
|
|
--------
|
|
repeat : Repeat elements of an array.
|
|
sort : Return a sorted copy of an array.
|
|
|
|
Notes
|
|
-----
|
|
When an axis is specified the subarrays indexed by the axis are sorted.
|
|
This is done by making the specified axis the first dimension of the array
|
|
(move the axis to the first dimension to keep the order of the other axes)
|
|
and then flattening the subarrays in C order. The flattened subarrays are
|
|
then viewed as a structured type with each element given a label, with the
|
|
effect that we end up with a 1-D array of structured types that can be
|
|
treated in the same way as any other 1-D array. The result is that the
|
|
flattened subarrays are sorted in lexicographic order starting with the
|
|
first element.
|
|
"""
|
|
if isinstance(a, (tuple, list)):
|
|
a = np.array(a)
|
|
|
|
r = ArrayUtil.unique(a._array, axis, return_index, return_inverse, return_counts)
|
|
|
|
return np.NDArray(r[0]) if len(r) == 1 else [np.NDArray(rr) for rr in r]
|
|
|
|
|
|
def intersect1d(ar1, ar2, assume_unique=False, return_indices=False):
|
|
"""
|
|
Find the intersection of two arrays.
|
|
Return the sorted, unique values that are in both of the input arrays.
|
|
|
|
Parameters
|
|
----------
|
|
ar1, ar2 : array_like
|
|
Input arrays. Will be flattened if not already 1D.
|
|
assume_unique : bool
|
|
If True, the input arrays are both assumed to be unique, which
|
|
can speed up the calculation. If True but ``ar1`` or ``ar2`` are not
|
|
unique, incorrect results and out-of-bounds indices could result.
|
|
Default is False.
|
|
return_indices : bool
|
|
If True, the indices which correspond to the intersection of the two
|
|
arrays are returned. The first instance of a value is used if there are
|
|
multiple. Default is False.
|
|
|
|
Returns
|
|
-------
|
|
intersect1d : ndarray
|
|
Sorted 1D array of common and unique elements.
|
|
comm1 : ndarray
|
|
The indices of the first occurrences of the common values in `ar1`.
|
|
Only provided if `return_indices` is True.
|
|
comm2 : ndarray
|
|
The indices of the first occurrences of the common values in `ar2`.
|
|
Only provided if `return_indices` is True.
|
|
|
|
See Also
|
|
--------
|
|
numeric.lib.arraysetops : Module with a number of other functions for
|
|
performing set operations on arrays.
|
|
|
|
Examples
|
|
--------
|
|
>>> np.intersect1d([1, 3, 4, 3], [3, 1, 2, 1])
|
|
array([1, 3])
|
|
To intersect more than two arrays, use functools.reduce:
|
|
>>> from functools import reduce
|
|
>>> reduce(np.intersect1d, ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2]))
|
|
array([3])
|
|
To return the indices of the values common to the input arrays
|
|
along with the intersected values:
|
|
>>> x = np.array([1, 1, 2, 3, 4])
|
|
>>> y = np.array([2, 1, 4, 6])
|
|
>>> xy, x_ind, y_ind = np.intersect1d(x, y, return_indices=True)
|
|
>>> x_ind, y_ind
|
|
(array([0, 2, 4]), array([1, 0, 2]))
|
|
>>> xy, x[x_ind], y[y_ind]
|
|
(array([1, 2, 4]), array([1, 2, 4]), array([1, 2, 4]))
|
|
"""
|
|
ar1 = np.asanyarray(ar1)
|
|
ar2 = np.asanyarray(ar2)
|
|
|
|
if not assume_unique:
|
|
if return_indices:
|
|
ar1, ind1 = unique(ar1, return_index=True)
|
|
ar2, ind2 = unique(ar2, return_index=True)
|
|
else:
|
|
ar1 = unique(ar1)
|
|
ar2 = unique(ar2)
|
|
else:
|
|
ar1 = ar1.ravel()
|
|
ar2 = ar2.ravel()
|
|
|
|
aux = np.concatenate((ar1, ar2))
|
|
if return_indices:
|
|
aux_sort_indices = np.argsort(aux)
|
|
aux = aux[aux_sort_indices]
|
|
else:
|
|
aux.sort()
|
|
|
|
mask = aux[1:] == aux[:-1]
|
|
int1d = aux[:-1][mask]
|
|
|
|
if return_indices:
|
|
ar1_indices = aux_sort_indices[:-1][mask]
|
|
ar2_indices = aux_sort_indices[1:][mask] - ar1.size
|
|
if not assume_unique:
|
|
ar1_indices = ind1[ar1_indices]
|
|
ar2_indices = ind2[ar2_indices]
|
|
|
|
return int1d, ar1_indices, ar2_indices
|
|
else:
|
|
return int1d
|
|
|
|
|
|
def ediff1d(ary, to_end=None, to_begin=None):
|
|
"""
|
|
The differences between consecutive elements of an array.
|
|
|
|
Parameters
|
|
----------
|
|
ary : array_like
|
|
If necessary, will be flattened before the differences are taken.
|
|
to_end : array_like, optional
|
|
Number(s) to append at the end of the returned differences.
|
|
to_begin : array_like, optional
|
|
Number(s) to prepend at the beginning of the returned differences.
|
|
|
|
Returns
|
|
-------
|
|
ediff1d : ndarray
|
|
The differences. Loosely, this is ``ary.flat[1:] - ary.flat[:-1]``.
|
|
|
|
See Also
|
|
--------
|
|
diff, gradient
|
|
|
|
Notes
|
|
-----
|
|
When applied to masked arrays, this function drops the mask information
|
|
if the `to_begin` and/or `to_end` parameters are used.
|
|
|
|
Examples
|
|
--------
|
|
>>> x = np.array([1, 2, 4, 7, 0])
|
|
>>> np.ediff1d(x)
|
|
array([ 1, 2, 3, -7])
|
|
|
|
>>> np.ediff1d(x, to_begin=-99, to_end=np.array([88, 99]))
|
|
array([-99, 1, 2, ..., -7, 88, 99])
|
|
|
|
The returned array is always 1D.
|
|
|
|
>>> y = [[1, 2, 4], [1, 6, 24]]
|
|
>>> np.ediff1d(y)
|
|
array([ 1, 2, -3, 5, 18])
|
|
|
|
"""
|
|
conv = np.asanyarray(ary)
|
|
# Convert to (any) array and ravel:
|
|
ary = conv.ravel()
|
|
|
|
# enforce that the dtype of `ary` is used for the output
|
|
dtype_req = ary.dtype
|
|
|
|
# fast track default case
|
|
if to_begin is None and to_end is None:
|
|
return ary[1:] - ary[:-1]
|
|
|
|
if to_begin is None:
|
|
l_begin = 0
|
|
else:
|
|
to_begin = np.asanyarray(to_begin)
|
|
to_begin = to_begin.ravel()
|
|
l_begin = len(to_begin)
|
|
|
|
if to_end is None:
|
|
l_end = 0
|
|
else:
|
|
to_end = np.asanyarray(to_end)
|
|
to_end = to_end.ravel()
|
|
l_end = len(to_end)
|
|
|
|
# do the calculation in place and copy to_begin and to_end
|
|
l_diff = max(len(ary) - 1, 0)
|
|
result = np.empty_like(ary, shape=l_diff + l_begin + l_end)
|
|
|
|
if l_begin > 0:
|
|
result[:l_begin] = to_begin
|
|
if l_end > 0:
|
|
result[l_begin + l_diff:] = to_end
|
|
|
|
result[l_begin:l_begin+l_diff] = ary[1:] - ary[:-1]
|
|
|
|
return result
|