"""
Decorators for automatic grid selection.
Can be added to functions or class methods.
Automatically applies the function to structured or unstructured data.
"""
import functools
import numpy as np
import xarray as xr
[docs]
def autogrid_method(func):
"""
Make class methods on grids work for a variety of input data formats.
Methods work automatically for structured and unstructured data.
"""
# An additional wrapper layer is required to pass additional kwargs
@functools.wraps(func)
def wrapper_func(self, r_norm, z_norm, *args, **kwargs):
kwargs = {**dict(self=self, func=func, r_norm=r_norm, z_norm=z_norm),
**kwargs}
# The coordinate phi can optionally be passed as a third argument
if len(args) == 1:
kwargs["phi"] = args[0]
elif len(args) != 0:
raise Exception("Autogrid methods take at most 3 arguments" \
+ f"({2 + len(args)} given)!")
# NOTE: This is just a wrapper, the magic happens inside the following
# function.
return _wrapped_grid_function(**kwargs)
return wrapper_func
[docs]
def autogrid_function(func):
"""
Make functions on grids work for a variety of input data formats.
Same use as the method version.
"""
@functools.wraps(func)
def wrapper_func(r_norm, z_norm, **kwargs):
kwargs = {**dict(func=func, r_norm=r_norm, z_norm=z_norm), **kwargs}
return _wrapped_grid_function(**kwargs)
return wrapper_func
[docs]
def autocoord_method(func):
"""
Make class methods on arrays work for a variety of input data formats.
Methods work automatically for structured and unstructured data.
"""
@functools.wraps(func)
def wrapper_func(self, *args, r_norm=None, z_norm=None, grid=None,
**kwargs):
kwargs = {**dict(self=self, func=func, args=args, r_norm=r_norm,
z_norm=z_norm, grid=grid), **kwargs}
return _wrapped_coord_function(**kwargs)
return wrapper_func
[docs]
def autocoord_function(func):
"""
Make functions on arrays work for a variety of input data formats.
Same use as the method version.
"""
@functools.wraps(func)
def wrapper_func(*args, r_norm=None, z_norm=None, grid=None, **kwargs):
kwargs = {**dict(func=func, args=args, r_norm=r_norm, z_norm=z_norm,
grid=grid), **kwargs}
return _wrapped_coord_function(**kwargs)
return wrapper_func
def _wrapped_grid_function(func, **kwargs):
"""
Handle different data formats on a grid.
Determines for a function that operators on a grid,
which format of the data is used. This means that the function may be
called with (R, Z) coordinates that are floats or numpy arrays, or xarray
DataArrays that have dimension points or R and Z. Chooses attributes
grid, dims and coords of the resulting xarray automatically.
"""
assert("r_norm" in kwargs)
assert("z_norm" in kwargs)
r_norm = kwargs["r_norm"]
z_norm = kwargs["z_norm"]
if (isinstance(r_norm, xr.DataArray) \
and isinstance(z_norm, xr.DataArray) \
and (np.size(r_norm) > 1 or np.size(z_norm) > 1)):
# For the xarray case we distinct between structured (R, Z) and
# unstructured (points) data.
if ("R" in r_norm.dims and "Z" in z_norm.dims):
dims = ["Z", "R"]
coords = {"R": r_norm.values, "Z": z_norm.values}
is_structured = True
else:
assert("points" in r_norm.dims)
assert("points" in z_norm.dims)
assert(r_norm.shape == z_norm.shape)
dims = ["points"]
coords = {}
is_structured = False
else:
# For the non xarray case we require that r, z are equal in size. One
# further distinction must be made for numpy arrays and floats. Lists
# of floats or similar are not supported.
assert(np.size(r_norm) == np.size(z_norm))
is_structured = False
if isinstance(r_norm, np.ndarray):
dims = ["points"]
else:
assert not isinstance(r_norm, list), \
"List of float not supported for autogrid functions!"
dims = []
coords = {}
kwargs["dims"] = dims
kwargs["coords"] = coords
kwargs["is_structured"] = is_structured
return func(**kwargs)
def _wrapped_coord_function(func, *args, **kwargs):
"""
Handle different data formats on arrays.
Determines for a function that operates on multiple
arrays, which coordinates should be used based on the call. This means that
one can either specify the coordinates manually, or implicit with a given
grid or they may be included within the xarray as coordinates.
"""
args = kwargs["args"]
del kwargs["args"]
# NOTE: For this to work we require one positional argument to be used
# to check and determine dimensions if coordinates are not explicitly
# given
assert(len(args) > 0)
main_arg = args[0]
assert(isinstance(main_arg, xr.DataArray))
assert(("R" in main_arg.dims and "Z" in main_arg.dims) \
or "points" in main_arg.dims)
if ("points" in main_arg.dims):
kwargs["is_structured"] = False
else:
kwargs["is_structured"] = True
# Check that the number of positional arguments is not greater than
# expected. It can be less since keyword version may be chosen. This check
# is useful to catch error cases with too many positional arguments, in
# which case the default error message is not really readable.
n_args_expected = func.__code__.co_argcount \
- func.__code__.co_kwonlyargcount
assert(len(args) <= n_args_expected)
if(kwargs["r_norm"] is not None):
assert(kwargs["z_norm"] is not None)
if not kwargs["is_structured"]:
assert np.shape(kwargs["r_norm"]) == np.shape(kwargs["z_norm"])
elif(kwargs["grid"] is not None):
assert(kwargs["z_norm"] is None)
kwargs["r_norm"], kwargs["z_norm"] = \
kwargs["grid"].coords_like(main_arg)
else:
assert("Z" in main_arg.coords and "R" in main_arg.coords)
kwargs["r_norm"] = main_arg.coords["R"]
kwargs["z_norm"] = main_arg.coords["Z"]
if("self" in kwargs):
self = kwargs["self"]
del kwargs["self"]
return func(self, *args, **kwargs)
else:
return func(*args, **kwargs)
