mirror of
https://github.com/meteoinfo/MeteoInfo.git
synced 2025-12-08 20:36:05 +00:00
368 lines
12 KiB
Python
368 lines
12 KiB
Python
from mipylib import numeric as np
|
|
from ..core import numeric as _nx
|
|
from ..core import dtype
|
|
from ..core._ndarray import NDArray
|
|
from ..core.fromnumeric import (ravel, nonzero)
|
|
from org.meteoinfo.ndarray.math import ArrayMath
|
|
|
|
__all__ = ['angle','extract', 'place', 'grid_edge', 'gradient']
|
|
|
|
|
|
def extract(condition, arr):
|
|
"""
|
|
Return the elements of an array that satisfy some condition.
|
|
This is equivalent to ``np.compress(ravel(condition), ravel(arr))``. If
|
|
`condition` is boolean ``np.extract`` is equivalent to ``arr[condition]``.
|
|
Note that `place` does the exact opposite of `extract`.
|
|
Parameters
|
|
----------
|
|
condition : array_like
|
|
An array whose nonzero or True entries indicate the elements of `arr`
|
|
to extract.
|
|
arr : array_like
|
|
Input array of the same size as `condition`.
|
|
Returns
|
|
-------
|
|
extract : ndarray
|
|
Rank 1 array of values from `arr` where `condition` is True.
|
|
See Also
|
|
--------
|
|
take, put, copyto, compress, place
|
|
Examples
|
|
--------
|
|
>>> arr = np.arange(12).reshape((3, 4))
|
|
>>> arr
|
|
array([[ 0, 1, 2, 3],
|
|
[ 4, 5, 6, 7],
|
|
[ 8, 9, 10, 11]])
|
|
>>> condition = np.mod(arr, 3)==0
|
|
>>> condition
|
|
array([[ True, False, False, True],
|
|
[False, False, True, False],
|
|
[False, True, False, False]])
|
|
>>> np.extract(condition, arr)
|
|
array([0, 3, 6, 9])
|
|
If `condition` is boolean:
|
|
>>> arr[condition]
|
|
array([0, 3, 6, 9])
|
|
"""
|
|
return _nx.take(ravel(arr), nonzero(ravel(condition))[0])
|
|
|
|
|
|
def place(arr, mask, vals):
|
|
"""
|
|
Change elements of an array based on conditional and input values.
|
|
Similar to ``np.copyto(arr, vals, where=mask)``, the difference is that
|
|
`place` uses the first N elements of `vals`, where N is the number of
|
|
True values in `mask`, while `copyto` uses the elements where `mask`
|
|
is True.
|
|
Note that `extract` does the exact opposite of `place`.
|
|
|
|
Parameters
|
|
----------
|
|
arr : ndarray
|
|
Array to put data into.
|
|
mask : array_like
|
|
Boolean mask array. Must have the same size as `a`.
|
|
vals : 1-D sequence
|
|
Values to put into `a`. Only the first N elements are used, where
|
|
N is the number of True values in `mask`. If `vals` is smaller
|
|
than N, it will be repeated, and if elements of `a` are to be masked,
|
|
this sequence must be non-empty.
|
|
|
|
See Also
|
|
--------
|
|
copyto, put, take, extract
|
|
|
|
Examples
|
|
--------
|
|
>>> arr = np.arange(6).reshape(2, 3)
|
|
>>> np.place(arr, arr>2, [44, 55])
|
|
>>> arr
|
|
array([[ 0, 1, 2],
|
|
[44, 55, 44]])
|
|
"""
|
|
if not isinstance(arr, NDArray):
|
|
raise TypeError("argument 1 must be numpy.ndarray, "
|
|
"not {name}".format(name=type(arr).__name__))
|
|
|
|
if isinstance(vals, (list, tuple)):
|
|
vals = NDArray(vals)
|
|
|
|
ArrayMath.place(arr.asarray(), mask.asarray(), vals.asarray())
|
|
|
|
|
|
def grid_edge(x, y):
|
|
"""
|
|
Return grid edge coordinate array.
|
|
:param x: (*array*) X coordinate array with one dimension.
|
|
:param y: (*array*) Y coordinate array width one dimension.
|
|
:return: Grid edge coordinate array of x and y with one dimension.
|
|
"""
|
|
yn = y.size
|
|
xn = x.size
|
|
n = (xn + yn) * 2
|
|
xx = _nx.zeros(n)
|
|
yy = _nx.zeros(n)
|
|
xx[:xn] = x
|
|
yy[:xn] = y[0]
|
|
xx[xn:xn + yn] = x[-1]
|
|
yy[xn:xn + yn] = y
|
|
xx[xn + yn:xn + yn + xn] = x[::-1]
|
|
yy[xn + yn:xn + yn + xn] = y[-1]
|
|
xx[xn + yn + xn:] = x[0]
|
|
yy[xn + yn + xn:] = y[::-1]
|
|
|
|
return xx, yy
|
|
|
|
def angle(z, deg=False):
|
|
"""
|
|
Return the angle of the complex argument.
|
|
|
|
Parameters
|
|
----------
|
|
z : array_like
|
|
A complex number or sequence of complex numbers.
|
|
deg : bool, optional
|
|
Return angle in degrees if True, radians if False (default).
|
|
|
|
Returns
|
|
-------
|
|
angle : ndarray or scalar
|
|
The counterclockwise angle from the positive real axis on the complex
|
|
plane in the range ``(-pi, pi]``, with dtype as double.
|
|
|
|
See Also
|
|
--------
|
|
arctan2
|
|
absolute
|
|
|
|
Notes
|
|
-----
|
|
Although the angle of the complex number 0 is undefined, ``angle(0)``
|
|
returns the value 0.
|
|
|
|
Examples
|
|
--------
|
|
>>> np.angle([1.0, 1.0j, 1+1j]) # in radians
|
|
array([ 0. , 1.57079633, 0.78539816]) # may vary
|
|
>>> np.angle(1+1j, deg=True) # in degrees
|
|
45.0
|
|
|
|
"""
|
|
z = _nx.asanyarray(z)
|
|
if z.dtype == dtype.complex:
|
|
zimag = z.imag
|
|
zreal = z.real
|
|
else:
|
|
zimag = 0
|
|
zreal = z
|
|
|
|
a = _nx.arctan2(zimag, zreal)
|
|
if deg:
|
|
a *= 180 / _nx.pi
|
|
return a
|
|
|
|
def gradient(f, *varargs, **kwargs):
|
|
"""
|
|
Return the gradient of an N-dimensional array.
|
|
|
|
The gradient is computed using second order accurate central differences
|
|
in the interior points and either first or second order accurate one-sides
|
|
(forward or backwards) differences at the boundaries.
|
|
The returned gradient hence has the same shape as the input array.
|
|
|
|
Parameters
|
|
----------
|
|
f : array_like
|
|
An N-dimensional array containing samples of a scalar function.
|
|
varargs : list of scalar or array, optional
|
|
Spacing between f values. Default unitary spacing for all dimensions.
|
|
Spacing can be specified using:
|
|
|
|
1. single scalar to specify a sample distance for all dimensions.
|
|
2. N scalars to specify a constant sample distance for each dimension.
|
|
i.e. `dx`, `dy`, `dz`, ...
|
|
3. N arrays to specify the coordinates of the values along each
|
|
dimension of F. The length of the array must match the size of
|
|
the corresponding dimension
|
|
4. Any combination of N scalars/arrays with the meaning of 2. and 3.
|
|
|
|
If `axis` is given, the number of varargs must equal the number of axes.
|
|
Default: 1.
|
|
|
|
edge_order : {1, 2}, optional
|
|
Gradient is calculated using N-th order accurate differences
|
|
at the boundaries. Default: 1.
|
|
|
|
axis : None or int or tuple of ints, optional
|
|
Gradient is calculated only along the given axis or axes
|
|
The default (axis = None) is to calculate the gradient for all the axes
|
|
of the input array. axis may be negative, in which case it counts from
|
|
the last to the first axis.
|
|
|
|
Returns
|
|
-------
|
|
gradient : ndarray or list of ndarray
|
|
A list of ndarrays (or a single ndarray if there is only one dimension)
|
|
corresponding to the derivatives of f with respect to each dimension.
|
|
Each derivative has the same shape as f.
|
|
"""
|
|
f = _nx.asanyarray(f)
|
|
N = f.ndim # number of dimensions
|
|
|
|
axis = kwargs.pop('axis', None)
|
|
edge_order = kwargs.pop('edge_order', 1)
|
|
|
|
if axis is None:
|
|
axes = tuple(range(N))
|
|
else:
|
|
axes = _nx.normalize_axis_tuple(axis, N)
|
|
|
|
len_axes = len(axes)
|
|
n = len(varargs)
|
|
if n == 0:
|
|
# no spacing argument - use 1 in all axes
|
|
dx = [1.0] * len_axes
|
|
elif n == 1 and np.ndim(varargs[0]) == 0:
|
|
# single scalar for all axes
|
|
dx = varargs * len_axes
|
|
elif n == len_axes:
|
|
# scalar or 1d array for each axis
|
|
dx = list(varargs)
|
|
for i, distances in enumerate(dx):
|
|
distances = _nx.asanyarray(distances)
|
|
if distances.ndim == 0:
|
|
continue
|
|
elif distances.ndim != 1:
|
|
raise ValueError("distances must be either scalars or 1d")
|
|
if len(distances) != f.shape[axes[i]]:
|
|
raise ValueError("when 1d, distances must match "
|
|
"the length of the corresponding dimension")
|
|
if distances.dtype == dtype.int:
|
|
# Convert numpy integer types to float to avoid modular
|
|
# arithmetic in np.diff(distances).
|
|
distances = distances.astype(dtype.float)
|
|
diffx = _nx.diff(distances)
|
|
# if distances are constant reduce to the scalar case
|
|
# since it brings a consistent speedup
|
|
if (diffx == diffx[0]).all():
|
|
diffx = diffx[0]
|
|
dx[i] = diffx
|
|
else:
|
|
raise TypeError("invalid number of arguments")
|
|
|
|
if edge_order > 2:
|
|
raise ValueError("'edge_order' greater than 2 not supported")
|
|
|
|
# use central differences on interior and one-sided differences on the
|
|
# endpoints. This preserves second order-accuracy over the full domain.
|
|
|
|
outvals = []
|
|
|
|
# create slice objects --- initially all are [:, :, ..., :]
|
|
slice1 = [slice(None)]*N
|
|
slice2 = [slice(None)]*N
|
|
slice3 = [slice(None)]*N
|
|
slice4 = [slice(None)]*N
|
|
|
|
otype = f.dtype
|
|
if otype == dtype.int:
|
|
otype = dtype.float
|
|
|
|
for axis, ax_dx in zip(axes, dx):
|
|
if f.shape[axis] < edge_order + 1:
|
|
raise ValueError(
|
|
"Shape of array too small to calculate a numerical gradient, "
|
|
"at least (edge_order + 1) elements are required.")
|
|
# result allocation
|
|
out = np.empty_like(f, dtype=otype)
|
|
|
|
# spacing for the current axis
|
|
uniform_spacing = np.ndim(ax_dx) == 0
|
|
|
|
# Numerical differentiation: 2nd order interior
|
|
slice1[axis] = slice(1, -1)
|
|
slice2[axis] = slice(None, -2)
|
|
slice3[axis] = slice(1, -1)
|
|
slice4[axis] = slice(2, None)
|
|
|
|
if uniform_spacing:
|
|
out[tuple(slice1)] = (f[tuple(slice4)] - f[tuple(slice2)]) / (2. * ax_dx)
|
|
else:
|
|
dx1 = ax_dx[0:-1]
|
|
dx2 = ax_dx[1:]
|
|
a = -(dx2)/(dx1 * (dx1 + dx2))
|
|
b = (dx2 - dx1) / (dx1 * dx2)
|
|
c = dx1 / (dx2 * (dx1 + dx2))
|
|
# fix the shape for broadcasting
|
|
shape = np.ones(N, dtype=dtype.int)
|
|
shape[axis] = -1
|
|
a.shape = b.shape = c.shape = shape
|
|
# 1D equivalent -- out[1:-1] = a * f[:-2] + b * f[1:-1] + c * f[2:]
|
|
out[tuple(slice1)] = a * f[tuple(slice2)] + b * f[tuple(slice3)] + c * f[tuple(slice4)]
|
|
|
|
# Numerical differentiation: 1st order edges
|
|
if edge_order == 1:
|
|
slice1[axis] = 0
|
|
slice2[axis] = 1
|
|
slice3[axis] = 0
|
|
dx_0 = ax_dx if uniform_spacing else ax_dx[0]
|
|
# 1D equivalent -- out[0] = (f[1] - f[0]) / (x[1] - x[0])
|
|
out[tuple(slice1)] = (f[tuple(slice2)] - f[tuple(slice3)]) / dx_0
|
|
|
|
slice1[axis] = -1
|
|
slice2[axis] = -1
|
|
slice3[axis] = -2
|
|
dx_n = ax_dx if uniform_spacing else ax_dx[-1]
|
|
# 1D equivalent -- out[-1] = (f[-1] - f[-2]) / (x[-1] - x[-2])
|
|
out[tuple(slice1)] = (f[tuple(slice2)] - f[tuple(slice3)]) / dx_n
|
|
|
|
# Numerical differentiation: 2nd order edges
|
|
else:
|
|
slice1[axis] = 0
|
|
slice2[axis] = 0
|
|
slice3[axis] = 1
|
|
slice4[axis] = 2
|
|
if uniform_spacing:
|
|
a = -1.5 / ax_dx
|
|
b = 2. / ax_dx
|
|
c = -0.5 / ax_dx
|
|
else:
|
|
dx1 = ax_dx[0]
|
|
dx2 = ax_dx[1]
|
|
a = -(2. * dx1 + dx2)/(dx1 * (dx1 + dx2))
|
|
b = (dx1 + dx2) / (dx1 * dx2)
|
|
c = - dx1 / (dx2 * (dx1 + dx2))
|
|
# 1D equivalent -- out[0] = a * f[0] + b * f[1] + c * f[2]
|
|
out[tuple(slice1)] = a * f[tuple(slice2)] + b * f[tuple(slice3)] + c * f[tuple(slice4)]
|
|
|
|
slice1[axis] = -1
|
|
slice2[axis] = -3
|
|
slice3[axis] = -2
|
|
slice4[axis] = -1
|
|
if uniform_spacing:
|
|
a = 0.5 / ax_dx
|
|
b = -2. / ax_dx
|
|
c = 1.5 / ax_dx
|
|
else:
|
|
dx1 = ax_dx[-2]
|
|
dx2 = ax_dx[-1]
|
|
a = (dx2) / (dx1 * (dx1 + dx2))
|
|
b = - (dx2 + dx1) / (dx1 * dx2)
|
|
c = (2. * dx2 + dx1) / (dx2 * (dx1 + dx2))
|
|
# 1D equivalent -- out[-1] = a * f[-3] + b * f[-2] + c * f[-1]
|
|
out[tuple(slice1)] = a * f[tuple(slice2)] + b * f[tuple(slice3)] + c * f[tuple(slice4)]
|
|
|
|
outvals.append(out)
|
|
|
|
# reset the slice object in this dimension to ":"
|
|
slice1[axis] = slice(None)
|
|
slice2[axis] = slice(None)
|
|
slice3[axis] = slice(None)
|
|
slice4[axis] = slice(None)
|
|
|
|
if len_axes == 1:
|
|
return outvals[0]
|
|
else:
|
|
return outvals |