def __new__(cls, input_array, *args, grid=None, **kwargs):
# ToDo: sort out kwargs
dtype = kwargs.pop('dtype', None)
order = kwargs.pop('order', 'K')
subok = kwargs.pop('subok', False)
ndmin = kwargs.pop('ndmin', 0)
super().__new__(cls,
input_array,
dtype=dtype,
order=order,
subok=subok,
ndmin=ndmin)
obj = np.ma.asarray(input_array).view(cls)
if grid is not None:
obj.grid = grid
elif not len(args) == 0 or not len(kwargs) == 0:
obj.grid = dm.Grid(*args, **kwargs)
else:
# make a default grid:
if input_array.ndim <= 3:
axes_names = ['x', 'y', 'z']
else:
axes_names = ['x%i' for i in range(input_array.ndim)]
axes = {}
for d, n in zip(axes_names, input_array.shape):
axes[d] = np.arange(n)
obj.grid = dm.Grid(**axes)
return obj
def __getitem__(self, item):
'''
item : int, str, slice, ierable
'''
if isinstance(item, str):
# by name
if not item in self.vars:
# if it does not exist, add empty axis: ToDo: really?
print('needs to be checked, is weird behaviour')
self.add_axis(item)
idx = self.vars.index(item)
return self.axes[idx]
elif isinstance(item, Number):
return self[(item, )]
elif isinstance(item, slice):
return self[(item, )]
elif isinstance(item, list):
if all([isinstance(i, str) for i in item]):
new_obj = dm.Grid()
for var in item:
new_obj.axes.append(self[var])
return new_obj
elif all([isinstance(i, (np.integer, int)) for i in item]):
return self[(item, )]
else:
raise IndexError('Cannot process list of indices %s' % item)
elif isinstance(item, tuple):
if all([isinstance(i, str) for i in item]):
return self[list(item)]
item = self.convert_slice(item)
new_obj = dm.Grid()
for i in range(len(self)):
if i < len(item):
assert item[i] is not Ellipsis
if isinstance(item[i], (np.integer, int)):
# we can skip this axisesnion, as it is one element
continue
new_obj.axes.append(self.axes[i][item[i]])
else:
new_obj.axes.append(self.axes[i])
return new_obj
elif isinstance(item, Iterable):
new_axes = []
for it in item:
new_axes.append(self[it])
return dm.Grid(*new_axes)
else:
raise KeyError('Cannot get key from %s' % type(item))
def __init__(self, *args, **kwargs):
'''
Set the grid
'''
self._data = {}
self._grid = dm.Grid()
if len(args) == 0 and len(kwargs) > 0 and all(
[isinstance(v, dm.GridArray) for v in kwargs.values()]):
for n, d in kwargs.items():
self.add_data(n, d)
elif len(args) == 1 and len(kwargs) == 0 and isinstance(
args[0], dm.Grid):
self._grid = args[0]
else:
self._grid = dm.Grid(*args, **kwargs)
def add_data(self, var, data):
'''Add data
Parameters
----------
var : str
name of data
data : GridArray, GridData, Array
'''
if callable(data):
self._data[var] = data
return
if isinstance(data, (dm.GridArray, GridData)):
if self._grid is None or not self._grid.initialized:
self._grid = data.grid
self._setup_plotting_methods()
else:
assert self._grid == data.grid
if var in self._grid.vars:
raise ValueError('Variable `%s` is already a grid dimension!' %
var)
if isinstance(data, dm.GridData):
assert len(data.data_vars) == 1
data = data[0]
if self.ndim == 0:
print('adding default grid')
# make a default grid:
if data.ndim <= 3 and var not in ['x', 'y', 'z']:
axes_names = ['x', 'y', 'z']
else:
axes_names = ['x%i' for i in range(data.ndim + 1)]
axes_names.delete(var)
axes = {}
for d, n in zip(axes_names, data.shape):
axes[d] = np.arange(n)
self._grid = dm.Grid(**axes)
if data.ndim < self.ndim and self.shape[-1] == 1:
# add new axis
data = data[..., np.newaxis]
data = np.ma.asarray(data)
if not data.shape[:self.ndim] == self.shape:
raise ValueError(
'Incompatible data of shape %s for grid of shape %s' %
(data.shape, self.shape))
self._data[var] = data
def __init__(self, grid=None, data=None, *args, **kwargs):
'''
grid : dm.Grid
data : dm.PointData, dm.GridData or dm.GridArray
-> if PointArray, sample is needed
'''
# set up grid
if grid is None:
if len(args) == 1 and len(kwargs) == 0 and isinstance(
args[0], dm.Grid):
self.grid = args[0]
else:
self.grid = dm.Grid(*args, **kwargs)
#assert set(self.grid.vars) <= set(data.vars)
if not self.grid.initialized:
self.grid.initialize(data)
# create sample (grid variables) and remove from data
self.sample = [
data.get_array(var, flat=True) for var in self.grid.vars
]
# The following is a bit wonky, ToDo
if isinstance(data, dm.GridArray):
self.single = True
self.data = dm.PointData(test=data.flat())
else:
self.single = False
data_vars = data.vars
for var in self.grid.vars:
data_vars.remove(var)
if isinstance(data, dm.PointData):
self.data = data[data_vars]
elif isinstance(data, dm.GridData):
self.data = dm.PointData()
for var in data_vars:
self.data[var] = data.get_array(var)
else:
raise ValueError()
self.indices = None
self.group = None
def generate_destination(self, *args, **kwargs):
'''Correctly set up a destination data format
depending on the supplied input
Parameters
----------
args, kwargs
'''
if len(args) == 1 and len(kwargs) == 0:
dest = args[0]
if isinstance(dest, dm.GridData):
grid = dest.grid
grid.initialize(self.source)
return dm.GridData(grid)
if isinstance(dest, dm.Grid):
grid = dest
grid.initialize(self.source)
return dm.GridData(grid)
if isinstance(dest, dm.GridArray):
grid = dest.grid
grid.initialize(self.source)
return dm.GridArray(np.empty(grid.shape), grid=grid)
if isinstance(dest, dm.PointData):
# check which vars we need:
if self.source_has_grid:
return dest[self.source.grid.vars]
else:
return dest
# check if source has a grid and if any args are in there
if isinstance(self.source, (dm.GridData, dm.GridArray)):
dims = []
for arg in args:
# in this case the source may have a grid, get those edges
if isinstance(arg, str):
if arg in self.source.grid.vars:
dims.append(self.source.grid[arg])
continue
dims.append(arg)
args = dims
# instantiate
grid = dm.Grid(*args, **kwargs)
grid.initialize(self.source)
if isinstance(self.source, dm.GridArray):
return dm.GridArray(np.empty(grid.shape), grid=grid)
return dm.GridData(grid)
def T(self):
'''transpose'''
return dm.Grid(*list(self)[::-1])
def __deepcopy__(self, memo=None):
return dm.Grid(*copy.deepcopy(self.axes, memo=memo))