mirror of
https://github.com/meteoinfo/MeteoInfo.git
synced 2025-12-08 20:36:05 +00:00
344 lines
12 KiB
Python
344 lines
12 KiB
Python
import math
|
|
|
|
from mipylib.numeric import core as _nx
|
|
from ..core.numerictypes import ScalarType, find_common_type
|
|
|
|
__all__ = ['r_','c_','mgrid']
|
|
|
|
class nd_grid(object):
|
|
"""
|
|
Construct a multi-dimensional "meshgrid".
|
|
|
|
``grid = nd_grid()`` creates an instance which will return a mesh-grid
|
|
when indexed. The dimension and number of the output arrays are equal
|
|
to the number of indexing dimensions. If the step length is not a
|
|
complex number, then the stop is not inclusive.
|
|
|
|
However, if the step length is a **complex number** (e.g. 5j), then the
|
|
integer part of its magnitude is interpreted as specifying the
|
|
number of points to create between the start and stop values, where
|
|
the stop value **is inclusive**.
|
|
|
|
If instantiated with an argument of ``sparse=True``, the mesh-grid is
|
|
open (or not fleshed out) so that only one-dimension of each returned
|
|
argument is greater than 1.
|
|
|
|
Parameters
|
|
----------
|
|
sparse : bool, optional
|
|
Whether the grid is sparse or not. Default is False.
|
|
|
|
Notes
|
|
-----
|
|
The instances of `nd_grid`, `mgrid`, approximately defined as::
|
|
|
|
mgrid = nd_grid(sparse=False)
|
|
|
|
Users should use these pre-defined instances instead of using `nd_grid`
|
|
directly.
|
|
"""
|
|
|
|
def __init__(self, sparse=False):
|
|
self.sparse = sparse
|
|
|
|
def __getitem__(self, key):
|
|
try:
|
|
num_list = []
|
|
for k in range(len(key)):
|
|
step = key[k].step
|
|
start = key[k].start
|
|
stop = key[k].stop
|
|
if start is None:
|
|
start = 0
|
|
if step is None:
|
|
step = 1
|
|
if isinstance(step, complex):
|
|
step = abs(step)
|
|
length = int(step)
|
|
num_list.insert(0, _nx.linspace(start, stop, length))
|
|
else:
|
|
num_list.insert(0, _nx.arange(start, stop, step))
|
|
nn = _nx.meshgrid(*num_list)
|
|
return nn[::-1]
|
|
except (IndexError, TypeError):
|
|
step = key.step
|
|
stop = key.stop
|
|
start = key.start
|
|
if start is None:
|
|
start = 0
|
|
if isinstance(step, complex):
|
|
# Prevent the (potential) creation of integer arrays
|
|
step_float = abs(step)
|
|
length = int(step_float)
|
|
return _nx.linspace(start, stop, length)
|
|
else:
|
|
return _nx.arange(start, stop, step)
|
|
|
|
|
|
class MGridClass(nd_grid):
|
|
"""
|
|
An instance which returns a dense multi-dimensional "meshgrid".
|
|
|
|
An instance which returns a dense (or fleshed out) mesh-grid
|
|
when indexed, so that each returned argument has the same shape.
|
|
The dimensions and number of the output arrays are equal to the
|
|
number of indexing dimensions. If the step length is not a complex
|
|
number, then the stop is not inclusive.
|
|
|
|
However, if the step length is a **complex number** (e.g. 5j), then
|
|
the integer part of its magnitude is interpreted as specifying the
|
|
number of points to create between the start and stop values, where
|
|
the stop value **is inclusive**.
|
|
|
|
Returns
|
|
-------
|
|
mesh-grid `ndarrays` all of the same dimensions
|
|
|
|
See Also
|
|
--------
|
|
ogrid : like `mgrid` but returns open (not fleshed out) mesh grids
|
|
meshgrid: return coordinate matrices from coordinate vectors
|
|
r_ : array concatenator
|
|
:ref:`how-to-partition`
|
|
|
|
Examples
|
|
--------
|
|
>>> np.mgrid[0:5, 0:5]
|
|
array([[[0, 0, 0, 0, 0],
|
|
[1, 1, 1, 1, 1],
|
|
[2, 2, 2, 2, 2],
|
|
[3, 3, 3, 3, 3],
|
|
[4, 4, 4, 4, 4]],
|
|
[[0, 1, 2, 3, 4],
|
|
[0, 1, 2, 3, 4],
|
|
[0, 1, 2, 3, 4],
|
|
[0, 1, 2, 3, 4],
|
|
[0, 1, 2, 3, 4]]])
|
|
>>> np.mgrid[-1:1:5j]
|
|
array([-1. , -0.5, 0. , 0.5, 1. ])
|
|
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(MGridClass, self).__init__(sparse=False)
|
|
|
|
|
|
mgrid = MGridClass()
|
|
|
|
class AxisConcatenator:
|
|
"""
|
|
Translates slice objects to concatenation along an axis.
|
|
For detailed documentation on usage, see `r_`.
|
|
"""
|
|
# allow ma.mr_ to override this
|
|
concatenate = staticmethod(_nx.concatenate)
|
|
|
|
def __init__(self, axis=0, ndmin=1, trans1d=-1):
|
|
self.axis = axis
|
|
self.trans1d = trans1d
|
|
self.ndmin = ndmin
|
|
|
|
def __getitem__(self, key):
|
|
if not isinstance(key, tuple):
|
|
key = (key,)
|
|
|
|
# copy attributes, since they can be overridden in the first argument
|
|
trans1d = self.trans1d
|
|
ndmin = self.ndmin
|
|
axis = self.axis
|
|
|
|
objs = []
|
|
scalars = []
|
|
arraytypes = []
|
|
scalartypes = []
|
|
|
|
for k, item in enumerate(key):
|
|
scalar = False
|
|
if isinstance(item, slice):
|
|
step = item.step
|
|
start = item.start
|
|
stop = item.stop
|
|
if start is None:
|
|
start = 0
|
|
if step is None:
|
|
step = 1
|
|
if isinstance(step, complex):
|
|
size = int(abs(step))
|
|
newobj = _nx.linspace(start, stop, num=size)
|
|
else:
|
|
newobj = _nx.arange(start, stop, step)
|
|
if ndmin > 1:
|
|
newobj = _nx.array(newobj, copy=False, ndmin=ndmin)
|
|
if trans1d != -1:
|
|
newobj = newobj.swapaxes(-1, trans1d)
|
|
elif isinstance(item, str):
|
|
if k != 0:
|
|
raise ValueError("special directives must be the "
|
|
"first entry.")
|
|
if item in ('r', 'c'):
|
|
matrix = True
|
|
col = (item == 'c')
|
|
continue
|
|
if ',' in item:
|
|
vec = item.split(',')
|
|
try:
|
|
axis, ndmin = [int(x) for x in vec[:2]]
|
|
if len(vec) == 3:
|
|
trans1d = int(vec[2])
|
|
continue
|
|
except Exception as e:
|
|
raise ValueError(
|
|
"unknown special directive {!r}".format(item)
|
|
)
|
|
try:
|
|
axis = int(item)
|
|
continue
|
|
except (ValueError, TypeError):
|
|
raise ValueError("unknown special directive")
|
|
elif type(item) in ScalarType:
|
|
newobj = _nx.array(item, ndmin=ndmin)
|
|
scalars.append(len(objs))
|
|
scalar = True
|
|
scalartypes.append(newobj.dtype)
|
|
else:
|
|
item_ndim = _nx.ndim(item)
|
|
newobj = _nx.array(item, copy=False, subok=True, ndmin=ndmin)
|
|
if trans1d != -1 and item_ndim < ndmin:
|
|
k2 = ndmin - item_ndim
|
|
k1 = trans1d
|
|
if k1 < 0:
|
|
k1 += k2 + 1
|
|
defaxes = list(range(ndmin))
|
|
axes = defaxes[:k1] + defaxes[k2:] + defaxes[k1:k2]
|
|
newobj = newobj.transpose(axes)
|
|
objs.append(newobj)
|
|
if not scalar and isinstance(newobj, _nx.NDArray):
|
|
arraytypes.append(newobj.dtype)
|
|
|
|
# Ensure that scalars won't up-cast unless warranted
|
|
final_dtype = find_common_type(arraytypes, scalartypes)
|
|
if final_dtype is not None:
|
|
for k in scalars:
|
|
objs[k] = objs[k].astype(final_dtype)
|
|
|
|
res = self.concatenate(tuple(objs), axis=axis)
|
|
|
|
return res
|
|
|
|
def __len__(self):
|
|
return 0
|
|
|
|
# separate classes are used here instead of just making r_ = concatentor(0),
|
|
# etc. because otherwise we couldn't get the doc string to come out right
|
|
# in help(r_)
|
|
|
|
class RClass(AxisConcatenator):
|
|
"""
|
|
Translates slice objects to concatenation along the first axis.
|
|
This is a simple way to build up arrays quickly. There are two use cases.
|
|
1. If the index expression contains comma separated arrays, then stack
|
|
them along their first axis.
|
|
2. If the index expression contains slice notation or scalars then create
|
|
a 1-D array with a range indicated by the slice notation.
|
|
If slice notation is used, the syntax ``start:stop:step`` is equivalent
|
|
to ``np.arange(start, stop, step)`` inside of the brackets. However, if
|
|
``step`` is an imaginary number (i.e. 100j) then its integer portion is
|
|
interpreted as a number-of-points desired and the start and stop are
|
|
inclusive. In other words ``start:stop:stepj`` is interpreted as
|
|
``np.linspace(start, stop, step, endpoint=1)`` inside of the brackets.
|
|
After expansion of slice notation, all comma separated sequences are
|
|
concatenated together.
|
|
Optional character strings placed as the first element of the index
|
|
expression can be used to change the output. The strings 'r' or 'c' result
|
|
in matrix output. If the result is 1-D and 'r' is specified a 1 x N (row)
|
|
matrix is produced. If the result is 1-D and 'c' is specified, then a N x 1
|
|
(column) matrix is produced. If the result is 2-D then both provide the
|
|
same matrix result.
|
|
A string integer specifies which axis to stack multiple comma separated
|
|
arrays along. A string of two comma-separated integers allows indication
|
|
of the minimum number of dimensions to force each entry into as the
|
|
second integer (the axis to concatenate along is still the first integer).
|
|
A string with three comma-separated integers allows specification of the
|
|
axis to concatenate along, the minimum number of dimensions to force the
|
|
entries to, and which axis should contain the start of the arrays which
|
|
are less than the specified number of dimensions. In other words the third
|
|
integer allows you to specify where the 1's should be placed in the shape
|
|
of the arrays that have their shapes upgraded. By default, they are placed
|
|
in the front of the shape tuple. The third argument allows you to specify
|
|
where the start of the array should be instead. Thus, a third argument of
|
|
'0' would place the 1's at the end of the array shape. Negative integers
|
|
specify where in the new shape tuple the last dimension of upgraded arrays
|
|
should be placed, so the default is '-1'.
|
|
Parameters
|
|
----------
|
|
Not a function, so takes no parameters
|
|
Returns
|
|
-------
|
|
A concatenated ndarray or matrix.
|
|
See Also
|
|
--------
|
|
concatenate : Join a sequence of arrays along an existing axis.
|
|
c_ : Translates slice objects to concatenation along the second axis.
|
|
Examples
|
|
--------
|
|
>>> np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])]
|
|
array([1, 2, 3, ..., 4, 5, 6])
|
|
>>> np.r_[-1:1:6j, [0]*3, 5, 6]
|
|
array([-1. , -0.6, -0.2, 0.2, 0.6, 1. , 0. , 0. , 0. , 5. , 6. ])
|
|
String integers specify the axis to concatenate along or the minimum
|
|
number of dimensions to force entries into.
|
|
>>> a = np.array([[0, 1, 2], [3, 4, 5]])
|
|
>>> np.r_['-1', a, a] # concatenate along last axis
|
|
array([[0, 1, 2, 0, 1, 2],
|
|
[3, 4, 5, 3, 4, 5]])
|
|
>>> np.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, dim>=2
|
|
array([[1, 2, 3],
|
|
[4, 5, 6]])
|
|
>>> np.r_['0,2,0', [1,2,3], [4,5,6]]
|
|
array([[1],
|
|
[2],
|
|
[3],
|
|
[4],
|
|
[5],
|
|
[6]])
|
|
>>> np.r_['1,2,0', [1,2,3], [4,5,6]]
|
|
array([[1, 4],
|
|
[2, 5],
|
|
[3, 6]])
|
|
Using 'r' or 'c' as a first string argument creates a matrix.
|
|
>>> np.r_['r',[1,2,3], [4,5,6]]
|
|
matrix([[1, 2, 3, 4, 5, 6]])
|
|
"""
|
|
|
|
def __init__(self):
|
|
AxisConcatenator.__init__(self, 0)
|
|
|
|
r_ = RClass()
|
|
|
|
class CClass(AxisConcatenator):
|
|
"""
|
|
Translates slice objects to concatenation along the second axis.
|
|
This is short-hand for ``np.r_['-1,2,0', index expression]``, which is
|
|
useful because of its common occurrence. In particular, arrays will be
|
|
stacked along their last axis after being upgraded to at least 2-D with
|
|
1's post-pended to the shape (column vectors made out of 1-D arrays).
|
|
|
|
See Also
|
|
--------
|
|
column_stack : Stack 1-D arrays as columns into a 2-D array.
|
|
r_ : For more detailed documentation.
|
|
Examples
|
|
--------
|
|
>>> np.c_[np.array([1,2,3]), np.array([4,5,6])]
|
|
array([[1, 4],
|
|
[2, 5],
|
|
[3, 6]])
|
|
>>> np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])]
|
|
array([[1, 2, 3, ..., 4, 5, 6]])
|
|
"""
|
|
|
|
def __init__(self):
|
|
AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0)
|
|
|
|
|
|
c_ = CClass() |