# pylint: disable=too-many-lines
""" Container for grid coordinates
Coordinates primarily go into Field objects. The order of coords in
the clist should mirror the data layout, as in if data[ix, iy, iz]
(C-order, ix is fastest varying index) then list should go x, y, z...
this is the default order
types:
"Structured":
"nonuniform_cartesian"
"Cylindrical"
"Spherical"
"Unstructured":
-> Not Implemented <-
FIXME: uniform coordinates are generally unsupported, but they're just
a special case of nonuniform coordinates, so the functionality
is still there... it's just unnatural to use full crd arrays
with uniform coordinates
"""
from __future__ import print_function, division
# from timeit import default_timer as time
import itertools
import sys
import numpy as np
import viscid
from viscid.compat import string_types, izip
from viscid.vutil import subclass_spider
from viscid import sliceutil
__all__ = ['NonuniformFullArrayError', 'arrays2crds', 'wrap_crds', 'extend_arr']
[docs]def arrays2crds(crd_arrs, crd_type='AUTO', crd_names="xyzuvw", **kwargs):
"""make either uniform or nonuniform coordnates given full arrays
Args:
crd_arrs (array-like): for n-dimensional crds, supply a list
of n ndarrays
crd_type (str): passed to wrap_crds, should uniquely specify a
type of coordinate object
crd_names (iterable): names of coordinates in the same order
as crd_arrs. Should always be in xyz order.
"""
clist = []
uniform_clist = []
is_uniform = True
try:
_ = len(crd_arrs[0])
except TypeError:
crd_arrs = [crd_arrs]
for crd_name, arr in zip(crd_names, crd_arrs):
arr = viscid.asarray_datetime64(arr, conservative=True)
clist.append((crd_name, arr))
try:
rtol = 4 * np.finfo(arr.dtype).eps
except ValueError:
rtol = 0
if len(arr) > 1:
diff = arr[1:] - arr[:-1]
else:
diff = [1, 1]
if viscid.is_timedelta_like(diff, conservative=True):
diff = diff / np.ones((1,), dtype=diff.dtype)
if np.allclose(diff[0], diff, rtol=rtol):
uniform_clist.append((crd_name, [arr[0], arr[-1], len(arr)]))
else:
is_uniform = False
# uniform crds don't play nice with datetime axes
if any(viscid.is_time_like(arr, conservative=True) for arr in crd_arrs):
is_uniform = False
if not crd_type:
crd_type = 'AUTO'
if crd_type.startswith('AUTO'):
prefix = 'uniform' if is_uniform else 'nonuniform'
crd_type = crd_type.replace('AUTO', prefix)
cl = uniform_clist if crd_type.startswith('uniform') else clist
crds = wrap_crds(crd_type, cl, **kwargs)
return crds
def lookup_crds_type(type_name):
for cls in subclass_spider(Coordinates):
if cls.istype(type_name):
return cls
return None
[docs]def wrap_crds(crdtype, clist, **kwargs):
""" """
# print(len(clist), clist[0][0], len(clist[0][1]), crytype)
try:
try:
crdtype = lookup_crds_type(crdtype)
except AttributeError:
pass
return crdtype(clist, **kwargs)
except TypeError:
raise NotImplementedError("can not decipher crds: {0}".format(crdtype))
[docs]def extend_arr(x, n=1, cell_fraction=1.0, full_arr=True, default_width=1e-5,
width_arr=None):
"""TODO: docstring"""
x = viscid.asarray_datetime64(x, conservative=True)
try:
n_low, n_high = n
except TypeError:
n_low, n_high = n, n
if width_arr is None:
width_arr = x
if full_arr:
if len(x) > 1:
dxl = cell_fraction * (width_arr[1] - width_arr[0])
dxh = cell_fraction * (width_arr[-1] - width_arr[-2])
else:
dxl = cell_fraction * default_width
dxh = cell_fraction * default_width
arr_low = [x[0] - (i + 1) * dxl for i in range(n_low)][::-1]
arr_high = [x[-1] + (i + 1) * dxh for i in range(n_high)]
ret = np.concatenate([arr_low, x, arr_high])
else:
if cell_fraction != 1.0:
raise ValueError("Linspace arrays require extension with "
"cell_fraction == 1.0")
xl, xh, nx = x
nx = int(nx)
if nx > 1:
dx = (xh - xl) / (nx - 1)
else:
dx = default_width
xl -= (n_low * cell_fraction) * dx
xh += (n_high * cell_fraction) * dx
nx += n_low + n_high
ret = np.array([xl, xh, nx], dtype=x.dtype)
return ret
class Coordinates(object):
"""Base class for all coordinate types
Note:
Subclasses should only call super(...).__init__ after setting
up the axes
"""
_TYPE = "none"
_axes = ["x", "y", "z"]
_Pcrds = None
meta = None
_units_validated = False
_units = None
def __init__(self, units=None, **kwargs):
self.units = units
self.meta = kwargs
@property
def axes(self):
return self._axes
@property
def units(self):
if not self._units_validated:
units = self._units
if not isinstance(units, (list, tuple)):
units = [units] * self.nr_dims
units += [None] * (len(units) - self.nr_dims)
units = ["" if u is None else u.strip() for u in units]
self._units = units
self._units_validated = True
return self._units
@units.setter
def units(self, val):
self._units_validated = False
self._units = val
@property
def crdtype(self):
return self._TYPE
@classmethod
def istype(cls, type_str): # noqa
return cls._TYPE == type_str.lower()
def is_spherical(self):
return "spherical" in self._TYPE
def is_uniform(self):
return self._TYPE.startswith('uniform')
def ind(self, axis):
if isinstance(axis, int):
if axis < len(self.axes):
return axis
else:
raise ValueError()
else:
return self.axes.index(axis)
def _axisvalid(self, ax):
try:
self.ind(ax)
return True
except ValueError:
return False
def get_units(self, axes, allow_invalid=False):
ret = []
for ax in axes:
if self._axisvalid(ax):
ret.append(self.units[self.ind(ax)])
elif allow_invalid:
ret.append('')
else:
raise ValueError("{0} not in axes ({1})".format(ax, self._axes))
return ret
def get_unit(self, axis):
return self.get_units([axis])[0]
def axis_name(self, axis):
if isinstance(axis, string_types):
if axis in self._crds:
return axis
else:
raise KeyError(axis)
else:
return self.axes[axis]
def as_local_coordinates(self):
return self
def as_uv_coordinates(self):
# trying to get self._axes could raise an AttributeError for
# subclasses != StructredCrds
if len(self._axes) == 2:
return self
else:
raise NotImplementedError()
class StructuredCrds(Coordinates):
_TYPE = "structured"
_CENTER = {"none": "", "node": "nc", "cell": "cc",
"face": "fc", "edge": "ec"}
SUFFIXES = list(_CENTER.values())
_axes = ["x", "y", "z"]
_src_crds_nc = None
_src_crds_cc = None
_dtype = None
_reflect_axes = None
has_cc = None
def __init__(self, init_clist, has_cc=True, reflect_axes=None,
dtype=None, full_arrays=True, quiet_init=False, **kwargs):
""" if caled with an init_clist, then the coordinate names
are taken from this list """
self.has_cc = has_cc
self.reflect_axes = reflect_axes
self.dtype = dtype
if init_clist is not None:
self._axes = [d[0].lower() for d in init_clist]
self.clear_crds()
if init_clist is not None:
self._set_crds(init_clist)
super(StructuredCrds, self).__init__(**kwargs)
@property
def xl_nc(self):
return self.get_xl()
@property
def xh_nc(self):
return self.get_xh()
@property
def L_nc(self):
return self.get_L()
@property
def min_dx_nc(self):
return self.get_min_dx()
@property
def xl_cc(self):
return self.get_xl(center='cell')
@property
def xh_cc(self):
return self.get_xh(center='cell')
@property
def L_cc(self):
return self.get_L(center='cell')
@property
def min_dx_cc(self):
return self.get_min_dx(center='cell')
@property
def nr_dims(self):
""" number of spatial dimensions """
return len(self._axes)
@property
def dtype(self):
if self._dtype:
return self._dtype
else:
if self._Pcrds is not None:
return self[self.axes[0]].dtype
elif self._src_crds_nc:
return None
@dtype.setter
def dtype(self, dtype):
self._dtype = dtype
@property
def shape(self):
return self.shape_nc
@property
def shape_nc(self):
return np.asarray([len(self[ax]) for ax in self.axes])
@property
def shape_cc(self):
return np.asarray([len(self[ax + "cc"]) for ax in self.axes])
@property
def size(self):
return self.size_nc
@property
def size_nc(self):
return np.product(self.shape_nc)
@property
def size_cc(self):
return np.product(self.shape_cc)
@property
def reflect_axes(self):
return self._reflect_axes
@reflect_axes.setter
def reflect_axes(self, val):
if not val:
val = []
self._reflect_axes = val
@property
def _crds(self):
"!!! BIG NOTE: THIS PROPERTY IS STILL PRIVATE !!!"
if self._Pcrds is None:
self._fill_crds_dict()
return self._Pcrds
def clear_crds(self):
self._src_crds_nc = {}
self._src_crds_cc = {}
self.clear_cache()
# for d in self.axes:
# for sfx in self.SUFFIXES:
# self._Pcrds[d + sfx] = None
# self._Pcrds[d.upper() + sfx] = None
def clear_cache(self):
self._Pcrds = None
def _set_crds(self, clist):
""" called with a list of lists:
(('x', ndarray), ('y',ndarray), ('z',ndarray))
Note: input crds are assumed to be node centered
"""
for ci in clist:
axis = ci[0]
dat_nc = ci[1]
if axis not in self._axes:
raise KeyError()
if ci[1] is not None and len(ci[1].shape) != 1:
raise ValueError("Coordinates created with mangled nc input")
self._src_crds_nc[axis.lower()] = ci[1]
if len(ci) > 2:
if ci[2] is not None and len(ci[2].shape) != 1:
raise ValueError("Coordinates created with mangled cc input")
self._src_crds_cc[axis.lower()] = ci[2]
def apply_reflections(self):
"""
Returns:
Coordinates with reflections applied
"""
if len(self.reflect_axes) > 0:
return type(self)(self.get_clist(), has_cc=self.has_cc,
dtype=self.dtype)
else:
return self
def _reflect_axis_arr(self, arr): # pylint: disable=no-self-use
return np.array(-1.0 * arr[::-1])
def reflect_fld_arr(self, arr, cc=False, nr_comp=None, nr_comps=None):
fudge = 0 if nr_comp is None else 1
assert len(arr.shape) == len(self._axes) + fudge
rev = [True if ax in self.reflect_axes else False for ax in self._axes]
if cc:
shape = self.shape_cc
else:
shape = self.shape_nc
if nr_comp is not None:
shape = list(shape)
shape.insert(nr_comp, nr_comps)
rev.insert(nr_comp, False)
first, second = sliceutil.make_fwd_slice(shape, [], rev)
return arr[tuple(first)][tuple(second)]
def reflect_slices(self, slices, cc=False, cull_second=True):
"""This is tricky, only use if you know what's going on"""
# assert len(slices) == len(self._axes)
rev = [True if ax in self.reflect_axes else False for ax in self._axes]
if cc:
shape = self.shape_cc
else:
shape = self.shape_nc
first, second = sliceutil.make_fwd_slice(shape, slices, rev,
cull_second=cull_second)
return first, second
def _fill_crds_dict(self):
""" do the math to calc node, cell, face, edge crds from
the src crds """
self._Pcrds = {}
# get node centered crds from the src crds
sfx = self._CENTER["node"]
for axis, arr in self._src_crds_nc.items():
# make axis into string representation
# axis = self.axis_name(axis)
ind = self.ind(axis)
arr = np.array(arr, dtype=arr.dtype.name)
if axis in self.reflect_axes:
arr = self._reflect_axis_arr(arr)
# ====================================================================
# if self.transform_funcs is not None:
# if axis in self.transform_funcs:
# arr = self.transform_funcs[axis](self, arr,
# **self.transform_kwargs)
# elif ind in self.transform_funcs:
# arr = self.transform_funcs[ind](self, arr,
# **self.transform_kwargs)
# ====================================================================
flatarr, openarr = self._ogrid_single(ind, arr)
self._Pcrds[axis.lower()] = flatarr
self._Pcrds[axis.upper()] = openarr
# now with suffix
self._Pcrds[axis.lower() + sfx] = flatarr
self._Pcrds[axis.upper() + sfx] = openarr
# recalculate all cell centers, and refresh face / edges
sfx = self._CENTER["cell"]
for i, a in enumerate(self.axes):
# a = self.axis_name(a) # validate input
if a in self._src_crds_cc:
ccarr = self._src_crds_cc[a]
else:
# doing the cc math this way also works for datetime objects
if len(self._Pcrds[a]) == 1:
ccarr = self._Pcrds[a]
else:
ccarr = self._Pcrds[a][:-1] + 0.5 * (self._Pcrds[a][1:] -
self._Pcrds[a][:-1])
flatarr, openarr = self._ogrid_single(a, ccarr)
self._Pcrds[a + sfx] = flatarr
self._Pcrds[a.upper() + sfx] = openarr
# ok, so this is a little recursive, but it's ok since we set
# _Pcrds above, note however that now we only have nc and cc
# crds in _Pcrds
crds_nc = self.get_crds_nc()
crds_nc_shaped = self.get_crds_nc(shaped=True)
crds_cc = self.get_crds_cc()
crds_cc_shaped = self.get_crds_cc(shaped=True)
# store references to face and edge centers while we're here
sfx = self._CENTER["face"]
for i, a in enumerate(self.axes):
self._Pcrds[a + sfx] = [None] * len(self.axes)
self._Pcrds[a.upper() + sfx] = [None] * len(self.axes)
for j, d in enumerate(self.axes): # pylint: disable=W0612
if i == j:
self._Pcrds[a + sfx][j] = crds_nc[i][:-1]
self._Pcrds[a.upper() + sfx][j] = self._sm1(crds_nc_shaped[i])
else:
self._Pcrds[a + sfx][j] = crds_cc[i]
self._Pcrds[a.upper() + sfx][j] = crds_cc_shaped[i]
# same as face, but swap nc with cc
sfx = self._CENTER["edge"]
for i, a in enumerate(self.axes):
self._Pcrds[a + sfx] = [None] * 3
self._Pcrds[a.upper() + sfx] = [None] * 3
for j in range(3):
if i != j:
self._Pcrds[a + sfx][j] = crds_nc[i][:-1]
self._Pcrds[a.upper() + sfx][j] = self._sm1(crds_nc_shaped[i])
else:
self._Pcrds[a + sfx][j] = crds_cc[i]
self._Pcrds[a.upper() + sfx][j] = crds_cc_shaped[i]
@staticmethod
def _sm1(a):
n = a.shape
slices = [slice(None) if ni <= 1 else slice(None, -1) for ni in n]
return a[tuple(slices)]
def _ogrid_single(self, axis, arr):
""" returns (flat array, open array) """
i = self.ind(axis)
shape = [1] * self.nr_dims
shape[i] = -1
# print(arr.shape)
if len(arr.shape) == 1:
return arr, np.reshape(arr, shape)
elif len(arr.shape) == self.nr_dims:
return np.ravel(arr), arr
else:
raise ValueError()
def get_slice_extent(self, selection, allow_read=False):
"""find value extent of a slice"""
# IMPORTANT: This does not touch the actual crd arrays to maintain
# lazyness
# _, selections, _ = self._parse_slice(selections)
# extent = sliceutil.selections2values(None, selections, self.nr_dims)
sel_list = sliceutil.raw_sel2sel_list(selection)
full_sel_info = sliceutil.fill_nd_sel_list(sel_list, self.axes)
full_sel_list, full_ax_names, full_newdim_flags = full_sel_info
std_sel_list = sliceutil.standardize_sel_list(full_sel_list)
# # FIXME: std_sel_list to extent
# extent = selection2values(None, std_sel_list)
extent = sliceutil.sel_list2values(None, std_sel_list)
if np.any(np.isnan(extent)):
if allow_read:
raise NotImplementedError("Fallback to crd read to find slice "
"extent is not yet implemented.")
raise RuntimeError("Can't infer extent for selection [{0}] "
"without reading crds".format(selection))
# enforce that the values are increasing
for i, ext in enumerate(extent):
if ext[1] < ext[0]:
extent[i] = extent[i][::-1]
extent = np.asarray(extent)
return extent.T
def nc2cc(self, default_width=1e-5):
"""Extend coordinates half a grid cell in all directions
Used for turning node centered fields to cell centered without
changing the data, just the coordinates.
Returns:
New coordinates instance with same type as self
"""
axes = self.axes
crds_nc = self.get_crds_nc()
crds_cc = self.get_crds_cc()
for i, x_cc in enumerate(crds_cc):
crds_nc[i] = extend_arr(x_cc, default_width=default_width,
width_arr=crds_nc[i])
new_clist = [(ax, nc) for ax, nc in zip(axes, crds_nc)]
return type(self)(new_clist)
def _make_slice(self, selection, cc=False):
"""Turns a selection into finalized slices and coordinates
selection is passed through
:py:func:`viscid.sliceutil.raw_sel2sel_list` and
:py:func:`viscid.sliceutil.standardize_sel_list`
Parameters:
selection (str, list): selection
cc (bool): cell centered slice
Returns:
tuple (slices, slcrds, reduced)
* **slices**: list of things that can go straight into an
ndarray's __getitem__; one for each axis in self plus
any new axes
* **slcrds**: a clist for what the coords will be after the
slice
* **reduced**: a list of (axis, location) pairs of which
axes are sliced out
Note: cc is necessary for finding the closest plane, otherwise it
might be off by half a grid cell
"""
sel_list = sliceutil.raw_sel2sel_list(selection)
full_sel_info = sliceutil.fill_nd_sel_list(sel_list, self.axes)
full_sel_list, full_ax_names, full_newdim_flags = full_sel_info
std_sel_list = sliceutil.standardize_sel_list(full_sel_list)
crd_arrs_nc = []
crd_arrs_cc = []
for i, ax, sel in izip(itertools.count(), full_ax_names, std_sel_list):
if ax in self.axes:
assert not full_newdim_flags[i]
crd_arrs_nc.append(self.get_nc(ax))
crd_arrs_cc.append(self.get_cc(ax))
else:
assert sel == np.newaxis and full_newdim_flags[i]
if cc:
crd_arrs_nc.append(np.array([-1e-1, 1e-1], dtype=self.dtype))
crd_arrs_cc.append(np.array([0.0], dtype=self.dtype))
else:
crd_arrs_nc.append(np.array([0.0], dtype=self.dtype))
crd_arrs_cc.append(np.array([0.0], dtype=self.dtype))
crd_arrs = crd_arrs_cc if cc else crd_arrs_nc
idx_sel_list = list(sliceutil.std_sel_list2index(std_sel_list, crd_arrs))
# Figure out what the sel_list is doing. If the slice reduces
# out a dimension, put it in reduced. Also apply the slices to
# the coordinate arrays so they can be used in to create new
# fields later on.
sliced_clist = []
reduced = []
for i, slc in enumerate(idx_sel_list):
axis = full_ax_names[i]
if isinstance(slc, slice):
if not all(s is None or hasattr(s, "__index__")
for s in (slc.start, slc.stop, slc.step)):
raise TypeError("bad sss:", slc)
sliced_clist_item = None
# if using step on a cell centered field, we need to figure
# stuff out, since the data && crd slices can't be the same,
# and it's not as simple as just adding 1 b/c of the stride
if cc and slc.step not in [None, 1, -1]:
crd_cc = crd_arrs_cc[i][slc]
crd_nc = 0.5 * (crd_cc[1:] + crd_cc[:-1])
allslc = slice(slc.start, slc.stop,
slc.step // np.abs(slc.step))
allnc = crd_arrs_nc[i][allslc]
crd_nc = np.concatenate([[allnc[0]], crd_nc, [allnc[-1]]])
sliced_clist_item = [axis, crd_nc, crd_cc]
elif cc and slc.stop is not None:
if slc.stop >= 0:
if slc.step is None or slc.step > 0:
newstop = slc.stop + 1
else:
newstop = slc.stop - 1
else: # slc.stop < 0
# this will slice crds nc, so if we're going backward,
# the extra crd will be included
# print("am i used?", slc.step)
newstop = slc.stop
slc = slice(slc.start, newstop, slc.step)
if sliced_clist_item is None:
cval = crd_arrs_nc[i][slc]
sliced_clist_item = [axis, cval]
sliced_clist.append(sliced_clist_item)
elif isinstance(slc, np.ndarray):
if cc:
cval_cc = crd_arrs_cc[i][slc]
cval_nc = (cval_cc[1:] + cval_cc[:-1]) / 2
cval_nc = extend_arr(cval_nc)
sliced_clist.append([axis, cval_nc, cval_cc])
else:
sliced_clist.append([axis, crd_arrs_nc[i][slc]])
elif hasattr(slc, "__index__"):
cval = crd_arrs_nc[i][slc]
reduced.append([axis, cval])
elif slc == np.newaxis:
cval = crd_arrs_nc[i]
sliced_clist.append([axis, cval])
else:
raise TypeError("{0};; {1}".format(type(slc), slc))
# determine uniform / nonuniform-ness of the resulting crds
if '_' in self._TYPE:
uniform_type, more_type = self._TYPE.split('_', 1)
else:
uniform_type, more_type = self._TYPE, None
any_non_uniform = False
for i, clist_i in enumerate(sliced_clist):
# int crds (including times/dates) are never store at uniform crds
if isinstance(clist_i[1][0], (int, np.integer, np.datetime64)):
any_non_uniform = True
break
diff_nc = np.diff(clist_i[1])
if diff_nc.shape[0] > 0:
# integers / dates / times should have been caught above
rtol = 4 * np.finfo(diff_nc.dtype).eps
if not np.allclose(diff_nc[0], diff_nc, rtol=rtol):
any_non_uniform = True
break
if uniform_type == 'uniform' and any_non_uniform:
uniform_type = 'nonuniform'
if more_type is None:
crds_type_name = uniform_type
else:
crds_type_name = '_'.join([uniform_type, more_type])
crds_type = lookup_crds_type(crds_type_name)
if uniform_type == 'uniform':
for i, slcl in enumerate(sliced_clist):
sliced_clist[i] = slcl[:2]
# print("< slice made", slices)
return idx_sel_list, sliced_clist, reduced, crds_type, full_ax_names
def make_slice(self, selection, cc=False):
slices, sliced_clist, reduced, crds_type, _ = \
self._make_slice(selection, cc=cc)
return slices, sliced_clist, reduced, crds_type
def make_slice_reduce(self, selection, cc=False):
"""make slice, and reduce dims that were not explicitly sliced"""
slices, crdlst, reduced, crds_type = self.make_slice(selection, cc=cc)
# augment slices / reduced
reduced_axes = [t[0] for t in reduced]
for i, axis in enumerate(self.axes):
reduce_axis = False
if axis not in reduced_axes:
if cc and self.shape_cc[i] == 1:
slices[i] = 0
reduced.insert(i, [axis, self.get_cc(axis)[0]])
reduce_axis = True
elif not cc and self.shape_nc[i] == 1:
slices[i] = 0
reduced.insert(i, [axis, self.get_nc(axis)[0]])
reduce_axis = True
# go find which element of crdlst to take out since there
# may already be elements missing
if reduce_axis:
for j, crd in enumerate(crdlst):
if crd is not None and crd[0] == axis:
crdlst[j] = None
break
# remove the newly reduced crds
crdlst = [crd for crd in crdlst if crd is not None]
return slices, crdlst, reduced, crds_type
def make_slice_keep(self, selection, cc=False):
"""make slice, but put back dims that were explicitly reduced"""
slices, crdlst, reduced, crds_type, full_ax_names = \
self._make_slice(selection, cc=cc)
# put reduced dims back, reduced will be in the same order as self.axes
# since make_slice loops over self.axes to do the slices; this enables
# us to call insert in the loop
new_ax_names = [c[0] for c in crdlst]
for axis, loc in reduced: # pylint: disable=W0612
new_axis_ind = full_ax_names.index(axis)
# slices[new_axis_ind] will be an int not a slice since it was
# reduced
loc_ind = slices[new_axis_ind]
crd_nc = self.get_nc(axis)
if cc:
if loc_ind == -1:
crd = crd_nc[-2:]
slc = slice(-1, None)
elif loc_ind == -2:
crd = crd_nc[-2:]
slc = slice(-2, -1)
else:
crd = crd_nc[loc_ind:loc_ind + 2]
slc = slice(loc_ind, loc_ind + 1)
else:
if loc_ind == -1:
crd = crd_nc[-1:]
slc = slice(-1, None)
else:
crd = crd_nc[loc_ind:loc_ind + 1]
slc = slice(loc_ind, loc_ind + 1)
slices[new_axis_ind] = slc
crdlst.insert(new_axis_ind, [axis, crd])
# should be no more reduced crds
reduced = []
# print("MAKE SLICE KEEP : slices", slices, "crdlst", crdlst,
# "reduced", reduced)
return slices, crdlst, reduced, crds_type
def slice(self, selection, cc=False):
"""Get crds that describe a slice (subset) of this grid.
Reduces dims the same way numpy / fields do. Chances are
you want either slice_reduce or slice_keep
"""
slices, crdlst, reduced, crds_type = self.make_slice(selection, cc=cc)
# pass through if nothing happened
if sliceutil.all_slices_none(slices):
return self
cunits = self.get_units((c[0] for c in crdlst), allow_invalid=True)
return wrap_crds(crds_type, crdlst, dtype=self.dtype, units=cunits,
full_arrays=True, quiet_init=True, **self.meta)
def slice_and_reduce(self, selection, cc=False):
"""Get crds that describe a slice (subset) of this grid. Go
through, and if the slice didn't touch a dim with only one crd,
reduce it
"""
slices, crdlst, reduced, crds_type = self.make_slice_reduce(selection,
cc=cc)
# pass through if nothing happened
if sliceutil.all_slices_none(slices):
return self
cunits = self.get_units((c[0] for c in crdlst), allow_invalid=True)
return wrap_crds(crds_type, crdlst, dtype=self.dtype, units=cunits,
full_arrays=True, quiet_init=True, **self.meta)
def slice_and_keep(self, selection, cc=False):
slices, crdlst, reduced, crds_type = self.make_slice_keep(selection,
cc=cc)
# pass through if nothing happened
if sliceutil.all_slices_none(slices):
return self
cunits = self.get_units((c[0] for c in crdlst), allow_invalid=True)
return wrap_crds(crds_type, crdlst, dtype=self.dtype, units=cunits,
full_arrays=True, quiet_init=True, **self.meta)
slice_reduce = slice_and_reduce
slice_keep = slice_and_keep
def slice_interp(self, selection, cc=False):
_, crdlst, _, crds_type = self.make_slice_keep(selection, cc=cc)
# axes, selection, _ = self._parse_slice(selection)
sel_list = sliceutil.raw_sel2sel_list(selection)
full_sel_info = sliceutil.fill_nd_sel_list(sel_list, self.axes)
full_sel_list, full_ax_names, full_newdim_flags = full_sel_info
std_sel_list = sliceutil.standardize_sel_list(full_sel_list)
limits = sliceutil.sel_list2values(None, full_sel_list, len(crdlst))
# for slices that were specified using a float, set that crd
# to the desired float instead of the nearest crd as does slice_keep
for i, slc in enumerate(std_sel_list):
# slc = sliceutil.convert_deprecated_floats(slc, "slc")
_slc_lims = limits[i]
if _slc_lims[0] == _slc_lims[1]:
val = _slc_lims[0]
crdlst[i][1][0] = val
if len(crdlst[i][1]) > 1:
# i'm not sure what this is for...
crdlst[i][1][1] = val
if len(crdlst[i]) > 2:
crdlst[i][2][0] = val
cunits = self.get_units((c[0] for c in crdlst), allow_invalid=True)
return wrap_crds(crds_type, crdlst, dtype=self.dtype, units=cunits,
full_arrays=True, quiet_init=True, **self.meta)
def get_crd(self, axis, shaped=False, center="none"):
"""if axis is not specified, return all coords,
shaped makes axis capitalized and returns ogrid like crds
shaped is only used if axis == None
sfx can be none, node, cell, face, edge
raises KeyError if axis not found
"""
return self.get_crds([axis], shaped=shaped, center=center)[0]
def get_crds(self, axes=None, shaped=False, center="none"):
"""Get coordinate arrays
Parameters:
axes: should be a list of names or integer indices, or None
to return all axes
shaped: boolean if you want a shaped or flat ndarray
(True is the same as capitalizing axis)
center: 'node' 'cell' 'edge' 'face', same as adding a suffix
to axes
Returns:
list of coords as ndarrays
"""
if axes is None:
axes = list(self.axes)
axes = [self.axes[a] if isinstance(a, (int, np.int)) else a for a in axes]
axes = [a.upper() if shaped else a for a in axes]
sfx = self._CENTER[center.lower()]
return [self._crds[self.axis_name(a) + sfx] for a in axes]
# def get_dcrd(self, axis, shaped=False, center="none"):
# """ if axis is not specified, return all coords,
# shaped makes axis capitalized and returns ogrid like crds
# shaped is only used if axis == None
# sfx can be none, node, cell, face, edge
# raises KeyError if axis not found """
# return self.get_dcrds([axis], shaped, center)[0]
# def get_dcrds(self, axes=None, shaped=False, center="none"):
# """ axes: should be a list of names or integer indices, or None
# to return all axes
# shaped: boolean if you want a shaped or flat ndarray
# (True is the same as capitalizing axis)
# center: 'node' 'cell' 'edge' 'face', same as adding a suffix to axes
# returns list of coords as ndarrays """
# if axes == None:
# axes = [a.upper() if shaped else a for a in self.axes]
# if not isinstance(axes, (list, tuple)):
# try:
# axes = [a.upper() if shaped else a for a in axes]
# except TypeError:
# axes = [axes]
# sfx = self._CENTER[center.lower()]
# return [(self._crds[self.axis_name(a) + sfx][1:] -
# self._crds[self.axis_name(a) + sfx][:-1]) for a in axes]
def get_culled_axes(self, ignore=2):
"""Get only good axes
Discard axes whose coords have length <= ignore... useful for
2d fields that only have 1 cell
Returns
list of axes names
"""
return [name for name in self.axes if len(self[name]) > ignore]
def get_clist(self, axes=None, slc=Ellipsis, full_arrays=True, center='node'):
"""??
Returns:
a clist of the coordinates sliced if you wish
Note:
I recommend using ``numpy.s_`` for making the slice
"""
if not full_arrays:
raise NotImplementedError("you need uniform crds for this")
if axes is None:
axes = self.axes
ret = [[axis, self.get_crd(axis, center=center)[slc].copy()] for axis in axes]
# # slc was never None, what was this for?
# if slc is None and full_arrays and center == 'node' and self._src_crds_cc:
# for i, axis in enumerate(axes):
# if axis in self._src_crds_cc:
# ret[i].append(self._src_crds_cc[axis].copy())
return ret
## These methods just return one crd axis
def get_nc(self, axis, shaped=False):
"""returns a flat ndarray of coordinates along a given axis
axis can be crd name as string, or index, as in x==2, y==1, z==2
"""
return self.get_crd(axis, shaped=shaped, center="node")
def get_cc(self, axis, shaped=False):
"""returns a flat ndarray of coordinates along a given axis
axis can be crd name as string, or index, as in x==2, y==1, z==2
"""
return self.get_crd(axis, shaped=shaped, center="cell")
def get_fc(self, axis, shaped=False):
"""returns a flat ndarray of coordinates along a given axis
axis can be crd name as string, or index, as in x==2, y==1, z==2
"""
return self.get_crd(axis, shaped=shaped, center="face")
def get_ec(self, axis, shaped=False):
"""returns a flat ndarray of coordinates along a given axis
axis can be crd name as string, or index, as in x==2, y==1, z==2
"""
return self.get_crd(axis, shaped=shaped, center="edge")
## These methods return all crd axes
def get_crds_nc(self, axes=None, shaped=False):
"""returns all node centered coords as a list of ndarrays, flat if
shaped==False, or shaped if shaped==True
"""
return self.get_crds(axes=axes, center="node", shaped=shaped)
def get_crds_cc(self, axes=None, shaped=False):
"""returns all cell centered coords as a list of ndarrays, flat if
shaped==False, or shaped if shaped==True
"""
return self.get_crds(axes=axes, center="cell", shaped=shaped)
def get_crds_ec(self, axes=None, shaped=False):
"""return all edge centered coords as a list of ndarrays, flat if
shaped==False, or shaped if shaped==True
"""
return self.get_crds(axes=axes, center="edge", shaped=shaped)
def get_crds_fc(self, axes=None, shaped=False):
"""returns all face centered coords as a list of ndarrays, flat if
shaped==False, or shaped if shaped==True
"""
return self.get_crds(axes=axes, center="face", shaped=shaped)
def __array__(self, *args, **kwargs):
return self.get_points()
def points(self, center=None, **kwargs):
return self.get_points(center=center, **kwargs)
def get_points(self, center="none", **kwargs):
"""returns all points in a grid defined by crds as a
nr_dims x nr_points ndarray
"""
crds = self.get_crds(shaped=False, center=center)
shape = [len(c) for c in crds]
arr = np.empty([len(shape)] + [np.prod(shape)])
for i, c in enumerate(crds):
arr[i, :] = np.repeat(np.tile(c, int(np.prod(shape[:i]))),
int(np.prod(shape[i + 1:])))
return arr
def as_mesh(self, center="none"):
crds = self.get_crds(shaped=False, center=center)
shape = [len(c) for c in crds]
pts = self.get_points(center=center)
while len(shape) > 2:
try:
shape.remove(1)
except ValueError:
raise ValueError("Only crds with 2 meaningful dimensions can "
"create a surface mesh")
return pts.reshape([3] + shape)
def get_dx(self, axes=None, center='node'):
"""Get cell widths if center == 'node', or distances between cell
centers if center == 'cell' """
return [x[1:] - x[:-1] if len(x) > 1 else 1.0
for x in self.get_crds(axes, center=center)]
def get_min_dx(self, axes=None, center='node'):
"""Get a minimum cell width for each axis"""
return np.array([np.min(dx) for dx in self.get_dx(axes, center=center)])
def get_xl(self, axes=None, center='node'):
_lst = [x[0] for x in self.get_crds(axes, center=center)]
try:
return np.array(_lst)
except ValueError:
return np.array(_lst, dtype='object')
def get_xh(self, axes=None, center='node'):
_lst = [x[-1] for x in self.get_crds(axes, center=center)]
try:
return np.array(_lst)
except ValueError:
return np.array(_lst, dtype='object')
def get_L(self, axes=None, center='node'):
"""Get lengths"""
return (self.get_xh(axes, center=center) -
self.get_xl(axes, center=center))
@property
def nr_points(self):
return self.get_nr_points()
def get_nr_points(self, center="none", **kwargs):
"""returns the number of points in a grid defined by these crds"""
return np.prod([len(crd) for crd in self.get_crds(center=center)])
def iter_points(self, center="none", **kwargs): # pylint: disable=W0613
"""returns an iterator over all the points in a grid, each
nextitem will be a list of nr_dims numbers
"""
return itertools.product(*self.get_crds(shaped=False, center=center))
def _newax_cval(self, xl, xh, cc):
raise NotImplementedError()
def atleast_3d(self, xl=-1, xh=1, cc=False):
if self.nr_dims >= 3:
return self
else:
current_axes = self._axes
if 'x' in current_axes or 'y' in current_axes or 'z' in current_axes:
axes3 = ['x', 'y', 'z']
target_idx = [0, 1, 2]
for i, ax in reversed(list(enumerate(axes3))):
if ax in current_axes:
axes3.pop(i)
target_idx.pop(i)
else:
axes3 = ['fill_a', 'fill_b', 'fill_c']
target_idx = [3, 3, 3]
new_clist = self.get_clist()
axes3 = axes3[:3 - len(current_axes)]
for idx, ax in zip(target_idx, axes3):
new_clist.insert(idx, (ax, self._newax_cval(xl, xh, cc)))
return type(self)(new_clist)
def print_tree(self):
c = self.get_clist()
for l in c:
print(l[0])
def __len__(self):
return self.nr_dims
def __getitem__(self, axis):
""" returns coord identified by axis (shaped / centering encoded
through capitalization of crd, and suffix), ex: 'Xcc' is shaped cell
centered, 'znc' is flat node centered, 'x' is flat node centered,
2 is self._axes[2]
"""
return self.get_crd(axis)
def __setitem__(self, axis, arr):
"""I recommend against doing this since there may be unintended
side effects
"""
raise RuntimeError("setting crds is deprecated - the constructor "
"does far too much transforming of the input "
"to assume that arr will be in the right form")
# return self._set_crds((axis, arr))
def __delitem__(self, item):
raise ValueError("can not delete crd this way")
def __contains__(self, item):
if item[-2:] in list(self._CENTER.values()):
item = item[:-2].lower()
return item in self._crds
class UniformCrds(StructuredCrds):
_TYPE = "uniform"
_nc_linspace_args = None
_cc_linspace_args = None
xl_nc = None
xh_nc = None
L_nc = None
shape_nc = None
min_dx_nc = None
xl_cc = None
xh_cc = None
shape_cc = None
L_cc = None
min_dx_cc = None
def __init__(self, init_clist, full_arrays=False, dtype='f8',
quiet_init=False, **kwargs):
"""
Args:
init_clist: this should look something like
[('y', [yl, yh, ny]), ('x', [xl, xh, nx])] unless
full_arrays is True.
full_arrays (bool): indicates that clist is given as
full coordinate arrays, like for non-uniform crds
Raises:
NonuniformFullArrayError if full_arrays and crds are not uniform
"""
self.dtype = dtype
if full_arrays:
if not quiet_init:
s = ("DEPRECATION...\n"
"Full arrays for uniform crds shouldn't be used due to \n"
"finite precision errors")
viscid.logger.warning(s)
_nc_linspace_args = [] # pylint: disable=unreachable
for _, arr in init_clist:
if viscid.is_time_like(arr, conservative=True):
raise NotImplementedError("Datetime arrays can't be in "
"uniform crds yet")
arr = np.asarray(arr)
if len(arr) > 1:
diff = arr[1:] - arr[:-1]
else:
diff = np.array([1, 1])
# This allclose is the problem... when slicing, it doesn't
# always pass
rtol = 4 * np.finfo(arr.dtype).eps
if not np.allclose(diff[0], diff, rtol=rtol):
raise NonuniformFullArrayError("Arrays are not uniform, {0}"
"".format(arr))
if len(arr) > 0:
_nc_linspace_args.append([arr[0], arr[-1], len(arr)])
else:
_nc_linspace_args.append([np.nan, np.nan, 0])
else:
for _, arr in init_clist:
if len(arr) != 3:
raise ValueError("is this a full_array?")
if viscid.is_time_like([arr[0], arr[1]], conservative=True):
raise NotImplementedError("Datetime arrays can't be in "
"uniform crds yet")
_nc_linspace_args = [arr for _, arr in init_clist]
self._nc_linspace_args = _nc_linspace_args
self._cc_linspace_args = []
for args in _nc_linspace_args:
half_dx = 0.5 * (args[1] - args[0]) / max(args[2], 1)
cc_args = [args[0] + half_dx, args[1] - half_dx,
max(args[2] - 1, 0)]
self._cc_linspace_args.append(cc_args)
# node centered things
self.xl_nc = np.array([args[0] for args in self._nc_linspace_args],
dtype=self.dtype)
self.xh_nc = np.array([args[1] for args in self._nc_linspace_args],
dtype=self.dtype)
self.shape_nc = np.array([args[2] for args in self._nc_linspace_args],
dtype='int')
self.L_nc = self.xh_nc - self.xl_nc
self.min_dx_nc = self.L_nc / self.shape_nc
self.min_dx_nc = np.ma.masked_values(self.min_dx_nc, 0.0)
# cell centered things
self.xl_cc = np.array([args[0] for args in self._cc_linspace_args],
dtype=self.dtype)
self.xh_cc = np.array([args[1] for args in self._cc_linspace_args],
dtype=self.dtype)
self.shape_cc = np.array([args[2] for args in self._cc_linspace_args],
dtype='int')
self.L_cc = self.xh_cc - self.xl_cc
self.min_dx_cc = self.L_cc / np.clip(self.shape_cc, 1, None)
init_clist = [(axis, None) for axis, _ in init_clist]
super(UniformCrds, self).__init__(init_clist, dtype=self.dtype,
**kwargs)
def _pull_out_axes(self, arrs, axes, center='none'):
center = center.lower()
if axes is None:
axind_list = list(range(len(self._axes)))
else:
axind_list = [] # self._axes.index(ax.lower()) for ax in axes]
for ax in axes:
try:
axind_list.append(self._axes.index(ax.lower()))
except AttributeError:
axind_list.append(ax)
# except ValueError:
# raise
if center == 'none' or center == 'node':
return arrs[0][axind_list]
elif center == 'cell':
return arrs[1][axind_list]
else:
raise ValueError()
def get_min_dx(self, axes=None, center='node'):
"""Get a minimum cell width for each axis"""
ret = self._pull_out_axes([self.min_dx_nc, self.min_dx_cc], axes,
center=center)
return np.ma.masked_values(ret, 0.0)
def get_xl(self, axes=None, center='node'):
return self._pull_out_axes([self.xl_nc, self.xl_cc], axes,
center=center)
def get_xh(self, axes=None, center='node'):
return self._pull_out_axes([self.xh_nc, self.xh_cc], axes,
center=center)
def get_L(self, axes=None, center='node'):
"""Get lengths"""
return self._pull_out_axes([self.L_nc, self.L_cc], axes,
center=center)
@property
def nr_points(self):
return self.get_nr_points()
def get_nr_points(self, center="none", **kwargs):
"""returns the number of points in a grid defined by these crds"""
center = center.lower()
if center == 'none' or center == 'node':
return self.size_nc
else:
return self.size_cc
def nc2cc(self, default_width=1e-5):
"""Extend coordinates half a grid cell in all directions
Used for turning node centered fields to cell centered without
changing the data, just the coordinates.
Returns:
New coordinates instance with same type as self
"""
axes = self.axes
xl, xh = self.get_xl(), self.get_xh()
dx = (xh - xl) / self.shape_nc
# FIXME: I don't think this will extend datetime64/timedelta64 axes
dx = np.choose(np.abs(dx) > 0, [default_width, dx])
xl -= 0.5 * dx
xh += 0.5 * dx
nx = self.shape_nc + 1
new_clist = [(ax, [xl[i], xh[i], nx[i]]) for i, ax in enumerate(axes)]
return type(self)(new_clist)
def _newax_cval(self, xl, xh, cc):
if cc:
return [xl, xh, 2]
else:
x0 = 0.5 * (xl + xh)
return [x0, x0, 1]
def get_clist(self, axes=None, slc=Ellipsis, full_arrays=False, center="node"):
if full_arrays:
return super(UniformCrds, self).get_clist(axes=axes, slc=slc,
center=center)
if slc != Ellipsis:
raise NotImplementedError("use full_arrays=True with slice != Ellipsis"
"for now")
if axes is None:
axes = self.axes
lst = []
for ax in axes:
ax = ax.lower()
if center == "node":
ls_args = self._nc_linspace_args[self.axes.index(ax)]
elif center == "cell":
ls_args = self._cc_linspace_args[self.axes.index(ax)]
else:
raise ValueError("node/cell centering only in here")
if ax in self.reflect_axes:
lst.append([ax, [-ls_args[1], -ls_args[0], ls_args[2]]])
else:
lst.append([ax, list(ls_args)])
return lst
def _fill_crds_dict(self):
assert len(self._nc_linspace_args) == len(self._axes)
self._src_crds_nc = {}
for ax, p in zip(self._axes, self._nc_linspace_args):
self._src_crds_nc[ax] = np.linspace(p[0], p[1], p[2]).astype(self.dtype)
return super(UniformCrds, self)._fill_crds_dict()
class NonuniformCrds(StructuredCrds):
_TYPE = "nonuniform"
def __init__(self, init_clist, full_arrays=True, quiet_init=False, **kwargs):
if not full_arrays:
raise ValueError("did you want Uniform crds?")
super(NonuniformCrds, self).__init__(init_clist, **kwargs)
def _newax_cval(self, xl, xh, cc):
if cc:
return np.array([xl, xh], dtype=self.dtype)
else:
return np.array([0.5 * (xl + xh)], dtype=self.dtype)
def get_clist(self, axes=None, slc=Ellipsis, full_arrays=True, center="node"):
"""??
Returns:
a clist of the coordinates sliced if you wish
Note:
I recommend using ``numpy.s_`` for making the slice
"""
if not full_arrays:
raise ValueError("Unstructured crds only have full arrays")
return super(NonuniformCrds, self).get_clist(axes=axes, slc=slc,
center=center)
class UniformCartesianCrds(UniformCrds):
_TYPE = "uniform_cartesian"
_axes = ["x", "y", "z"]
class NonuniformCartesianCrds(NonuniformCrds):
_TYPE = "nonuniform_cartesian"
_axes = ["x", "y", "z"]
class UniformSphericalCrds(UniformCrds):
_TYPE = "uniform_spherical"
_axes = ["phi", "theta", "r"]
class NonuniformSphericalCrds(NonuniformCrds):
_TYPE = "nonuniform_spherical"
_axes = ["phi", "theta", "r"]
class UnstructuredCrds(Coordinates):
_TYPE = "unstructured"
def __init__(self, **kwargs):
super(UnstructuredCrds, self).__init__(**kwargs)
raise NotImplementedError()
def _main():
print("full array")
x = np.arange(4)
print(viscid.extend_arr(x))
print(viscid.extend_arr(x, n=2))
print(viscid.extend_arr(x, n=1, cell_fraction=0.5))
print(viscid.extend_arr(x, n=2, cell_fraction=0.5))
print(viscid.extend_arr(x, n=1, cell_fraction=0.25))
print(viscid.extend_arr(x, n=2, cell_fraction=0.25))
print(viscid.extend_arr(x, n=(2, 3), cell_fraction=0.25))
print('full array single value')
x = np.array([0.0])
print(viscid.extend_arr(x))
print(viscid.extend_arr(x, n=2))
print(viscid.extend_arr(x, n=1, cell_fraction=0.5))
print(viscid.extend_arr(x, n=(2, 3), cell_fraction=0.5))
print("linspace array")
x = [0.0, 3.0, 4]
print(np.linspace(*viscid.extend_arr(x, full_arr=False)))
print(np.linspace(*viscid.extend_arr(x, full_arr=False, n=2)))
print(np.linspace(*viscid.extend_arr(x, full_arr=False, n=(2, 3))))
print("linspace array, single value")
x = np.array([0.0, 0.0, 1])
print(np.linspace(*viscid.extend_arr(x, full_arr=False)))
print(np.linspace(*viscid.extend_arr(x, full_arr=False, n=2)))
print(np.linspace(*viscid.extend_arr(x, full_arr=False, n=(2, 3))))
if __name__ == "__main__":
_main()
##
## EOF
##