"""Differential operators working on grids."""
import numpy as np
import xarray as xr
from typing import Union
from torx.grid import Grid2D
from torx.vector import cylindrical_vector, poloidal_vector
from torx.normalization import Normalization
from torx.autodoc_decorators_m import autodoc_function
from torx.autogrid_decorators_m import autocoord_function
[docs]
@autodoc_function
@autocoord_function
def grad(
data: Union[xr.DataArray, xr.Dataset],
norm: Normalization=None,
**kwargs
):
"""
Return the gradient of a given scalar field.
Uses second order central finite differences. There are two intended ways
of usage:
- 1) Input data on structured (R,Z) grid and containing the basis vectors
as coordinates. A grid is constructed internally.
- 2) Input data on unstructured (points), where a corresponding grid must
be provided via the keyword argument 'grid'.
"""
assert "vector" not in data.dims, \
"Gradient operator is only implemented for scalar-fields"
grid = _check_grid(**kwargs)
dims, coords = _check_dims_coords(data, **kwargs)
grad_r = grid._gradient_R(data)
grad_z = grid._gradient_Z(data)
if "phi" in data.dims and isinstance(grid, Grid2D):
grad_phi = grad_phi_2d(data, **kwargs)
grad = cylindrical_vector(grad_r, grad_phi, grad_z,
dims=dims, coords=coords)
else:
grad = poloidal_vector(grad_r, grad_z, dims=dims, coords=coords)
grad_norm = _derivative_normalization(grad, norm=norm)
return grad_norm
@autodoc_function
@autocoord_function
def grad_phi_2d(
data: Union[xr.DataArray, xr.Dataset],
norm: Normalization=None,
**kwargs
):
"""
Gradient in toroidal direction.
Uses naive second order central finite difference in toroidal direction for
axisymmetric data.
"""
grid = _check_grid(**kwargs)
r_loc, _ = grid.coords_like(data)
dphi = 2 * np.pi / data.phi.size
diff_phi = data.roll(phi=-1) - data.roll(phi=1)
grad_phi = diff_phi / ( 2 * dphi * r_loc)
grad_phi_norm = _derivative_normalization(grad_phi, norm=norm)
return grad_phi_norm
[docs]
@autodoc_function
@autocoord_function
def div(
data: Union[xr.DataArray, xr.Dataset],
norm: Normalization=None,
**kwargs
):
"""
Return the divergence of a given vector field.
Uses second order central finite differences. There are two intended ways
of usage:
- 1) Input data on structured (R,Z) grid and containing the basis vectors
as coordinates. A grid is constructed internally.
- 2) Input data on unstructured (points), where a corresponding grid must
be provided via the keyword argument 'grid'.
"""
assert "vector" in data.dims, "Must provide a vector-field"
assert ["eR", "ePhi", "eZ"] in data.vector
grid = _check_grid(**kwargs)
r_loc, _ = grid.coords_like(data)
data_r = data.sel(vector="eR").drop_vars("vector")
data_phi = data.sel(vector="ePhi").drop_vars("vector")
data_z = data.sel(vector="eZ").drop_vars("vector")
dims, _ = _check_dims_coords(data_r, **kwargs)
comp_r = grid._gradient_R(r_loc * data_r) / r_loc
comp_z = grid._gradient_Z(data_z)
if ("phi" in data.dims and isinstance(grid, Grid2D)):
comp_phi = grad_phi_2d(data_phi, **kwargs)
else:
comp_phi = xr.zeros_like(data_r)
div = (comp_r + comp_phi + comp_z).transpose(*dims)
div_norm = _derivative_normalization(div, norm=norm)
return div_norm
[docs]
@autodoc_function
@autocoord_function
def rot(
data: Union[xr.DataArray, xr.Dataset],
norm: Normalization=None,
**kwargs
):
"""
Return the rotation of a given vector field.
Uses second order central finite differences. There are two intended ways
of usage:
- 1) Input data on structured (R,Z) grid and containing the basis vectors
as coordinates. A grid is constructed internally.
- 2) Input data on unstructured (points), where a corresponding grid must
be provided via the keyword argument 'grid'.
"""
assert "vector" in data.dims, "Must provide a vector-field"
assert ["eR", "ePhi", "eZ"] in data.vector
grid = _check_grid(**kwargs)
r_loc, _ = grid.coords_like(data)
dims, coords = _check_dims_coords(data, **kwargs)
if "vector" in dims:
dims.remove("vector")
data_r = data.sel(vector="eR").drop_vars("vector")
data_phi = data.sel(vector="ePhi").drop_vars("vector")
data_z = data.sel(vector="eZ").drop_vars("vector")
# Compute the components that are guaranteed to exist
comp_r = - grid._gradient_Z(data_phi)
comp_phi = grid._gradient_Z(data_r) - grid._gradient_R(data_z)
comp_z = grid._gradient_R(data_phi * r_loc) / r_loc
# Correct by gradient components in toroidal direction if existent
if ("phi" in data.dims and isinstance(grid, Grid2D)):
comp_r += grad_phi_2d(data_z, **kwargs)
comp_z -= grad_phi_2d(data_r, **kwargs)
rot = cylindrical_vector(comp_r, comp_phi, comp_z, dims=dims, coords=coords)
rot_norm = _derivative_normalization(rot, norm=norm)
return rot_norm
def _check_grid(**kwargs):
"""Check for a grid in the kwargs and builds one if required."""
if kwargs["grid"] is None:
try:
r_norm, z_norm = kwargs["r_norm"], kwargs["z_norm"]
except KeyError:
raise RuntimeError("No coordinates found in data arg!")
RR, ZZ = np.meshgrid(r_norm, z_norm)
r_unstructured = np.reshape(RR, RR.size, order="F")
z_unstructured = np.reshape(ZZ, ZZ.size, order="F")
grid = Grid2D(r_unstructured, z_unstructured)
else:
grid = kwargs["grid"]
return grid
def _check_dims_coords(data, **kwargs):
"""Find dims and coords either from the data or the kwargs."""
if "dims" not in kwargs:
dims = list(data.dims)
else:
dims = kwargs["dims"]
if "coords" not in kwargs:
coords = data.coords.variables
else:
coords = kwargs["coords"]
return dims, coords
def _derivative_normalization(data: xr.DataArray,
norm: Normalization=None,
order: int=1):
"""Handle norm attribute of data for spatial derivative of given order."""
if norm is not None:
if "norm" in data.attrs:
data.attrs["norm"] /= norm.R0**order
else:
data.attrs["norm"] = 1 / norm.R0**order
return data
