import numpy as np
from .igrid import CENTERING_CHOICES, CenteringValueError, DimensionError, IField
from .raster import UniformRectilinear
from .rectilinear import Rectilinear
from .structured import Structured
from .unstructured import Unstructured
[docs]def combine_args_to_list(*args, **kwds):
if len(args) == 0:
args = kwds.get("default", [])
args = list(args)
combined_args = []
for arg in args:
if isinstance(arg, str):
combined_args.append(arg)
else:
combined_args.extend(arg)
return combined_args
[docs]class GridField(Unstructured, IField):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._fields = {}
self._field_units = {}
[docs] def add_field(
self,
field_name,
val,
centering="zonal",
units="-",
exist_action="clobber",
time=None,
):
if exist_action not in ("clobber", "append"):
raise ValueError("'exist_action' must be on of {'clobber', 'append'}")
try:
val.shape = val.size
except AttributeError:
val = np.array(val)
if centering not in CENTERING_CHOICES:
raise CenteringValueError(centering)
if centering == "zonal" and val.size != self.get_cell_count():
raise DimensionError(val.size, self.get_cell_count())
elif centering != "zonal" and val.size != self.get_point_count():
raise DimensionError(val.size, self.get_point_count())
self._field_times = {}
if field_name not in self._fields or exist_action == "clobber":
self._fields[field_name] = [val]
self._field_times[field_name] = [time]
else:
self._fields[field_name].append(val)
self._field_times[field_name].append(time)
self._field_units[field_name] = units
# self._fields[field_name].shape = val.size
[docs] def pop_field(self, field_name, *args, **kwds):
return_with_time = kwds.get("return_with_time", False)
field = self._fields[field_name].pop(*args)
if return_with_time:
time = self._field_times[field_name].pop(*args)
return (time, field)
else:
return field
[docs] def get_field(self, field_name):
if len(self._fields[field_name]) == 1:
return self._fields[field_name][0]
else:
return self._fields[field_name]
[docs] def get_field_units(self, field_name):
return self._field_units[field_name]
[docs] def set_field_units(self, field_name, units):
try:
self._field_units[field_name] = units
except KeyError:
pass
[docs] def get_point_fields(self, *args):
names = combine_args_to_list(*args, default=self._fields.keys())
fields = {}
for name in names:
array = self.get_field(name)
if array.size == self.get_point_count():
fields[name] = array
return fields
# fields = {}
# for (field, array) in self._fields.items ():
# if array.size == self.get_point_count ():
# fields[field] = array
# return fields
[docs] def get_cell_fields(self, *args):
names = combine_args_to_list(*args, default=self._fields.keys())
fields = {}
for name in names:
array = self.get_field(name)
if array.size == self.get_cell_count():
fields[name] = array
return fields
[docs] def items(self):
return self._fields.items()
[docs] def keys(self):
return self._fields.keys()
[docs] def values(self):
return self._fields.values()
[docs] def has_field(self, name):
return name in self._fields
[docs]class UnstructuredField(GridField):
pass
[docs]class StructuredField(Structured, GridField):
[docs] def add_field(
self,
field_name,
val,
centering="zonal",
units="-",
exist_action="clobber",
time=None,
):
if not hasattr(val, "ndim"):
val = np.array(val)
# try:
# ndim = val.ndim
# except AttributeError:
# val = np.array (val)
# ndim = val.ndim
if centering == "zonal":
if val.ndim > 1 and np.any(val.shape != self.get_shape() - 1):
raise DimensionError(val.shape, self.get_shape() - 1)
elif centering != "zonal":
if val.ndim > 1 and np.any(val.shape != self.get_shape()):
raise DimensionError(val.shape, self.get_shape())
try:
super().add_field(
field_name,
val,
centering=centering,
units=units,
exist_action=exist_action,
time=time,
)
except (DimensionError, CenteringValueError):
raise
[docs]class RectilinearField(Rectilinear, StructuredField):
pass
[docs]class RasterField(UniformRectilinear, StructuredField):
"""
Create a field that looks like this,
::
(0) --- (1) --- (2)
| | |
| 0 | 1 |
| | |
(3) --- (4) --- (5)
Create the field,
>>> g = RasterField ((2,3), (1,2), (0, 0), indexing='ij')
>>> g.get_cell_count ()
2
>>> g.get_point_count ()
6
Add some data at the points of our grid.
>>> data = np.arange (6)
>>> g.add_field ('var0', data, centering='point')
>>> g.get_field ('var0')
array([0, 1, 2, 3, 4, 5])
The data can be given either as a 1D array or with the same shape
as the point grid. In either case, though, it will be flattened.
>>> data = np.arange (6)
>>> data.shape = (2, 3)
>>> g.add_field ('var0', data, centering='point')
>>> g.get_field ('var0')
array([0, 1, 2, 3, 4, 5])
If the size or shape doesn't match, it's an error.
>>> data = np.arange (2)
>>> g.add_field ('bad var', data, centering='point') # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
...
DimensionError: 2 != 6
>>> data = np.ones ((3, 2))
>>> g.add_field ('bad var', data, centering='point') # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
...
DimensionError: (3, 2) != (2, 3)
Add data to the cells of the grid. Again, the data can be given
as a 1D array or with the same shape as the cell grid and if the
size or shape doesn't match raise an exception.
>>> data = np.arange (2)
>>> g.add_field('var1', data, centering='zonal')
>>> g.get_field('var1')
array([0, 1])
>>> data = np.ones((2, 1))
>>> g.add_field('bad var', data, centering='zonal') # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
...
DimensionError: (2, 1) != (1, 2)
>>> data = np.arange(3)
>>> g.add_field('bad var', data, centering='zonal') # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
...
DimensionError: 3 != 2
>>> data = np.array(3)
>>> g.add_field('bad var', data, centering='zonal') # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
...
DimensionError: 1 != 2
A 1D-Field,
>>> g = RasterField ((6, ), (1.5, ), (0.5, ), indexing='ij')
>>> g.get_point_count ()
6
>>> g.get_cell_count ()
5
>>> point_data = np.arange (6.)
>>> g.add_field ('Point Data', point_data, centering='point')
>>> cell_data = np.arange (5.)*10.
>>> g.add_field ('Cell Data', cell_data, centering='zonal')
>>> g.get_field ('Cell Data')
array([ 0., 10., 20., 30., 40.])
>>> g.get_field ('Point Data')
array([ 0., 1., 2., 3., 4., 5.])
A 3D-Field,
>>> g = RasterField ((4, 3, 2), (1.5, 1., 3), (0.5, 0, -.5 ), indexing='ij')
>>> g.get_point_count ()
24
>>> g.get_cell_count ()
6
>>> g.add_field ('Point Data', g.get_x (), centering='point')
>>> cell_data = np.arange (6.)*10.
>>> g.add_field ('Cell Data', cell_data, centering='zonal')
>>> g.get_field ('Cell Data')
array([ 0., 10., 20., 30., 40., 50.])
>>> g.get_field('Point Data')
array([-0.5, 2.5, -0.5, 2.5, -0.5, 2.5, -0.5, 2.5, -0.5, 2.5, -0.5,
2.5, -0.5, 2.5, -0.5, 2.5, -0.5, 2.5, -0.5, 2.5, -0.5, 2.5,
-0.5, 2.5])
>>> g.get_shape()
array([4, 3, 2])
>>> g.get_x().size == g.get_field('Point Data').size
True
>>> x = g.get_x()
>>> x.shape
(24,)
>>> x.shape = g.get_shape()
>>> x.shape
(4, 3, 2)
"""
pass
if __name__ == "__main__":
import doctest
doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)
