"""Perpendicular gradient operators."""
from numba import njit
import xarray as xr
import numpy as np
from torx.autodoc_decorators_m import autodoc_function
@njit
def calculate_gradient(
array: np.ndarray,
down2: np.ndarray,
down1: np.ndarray,
up1: np.ndarray,
up2: np.ndarray,
spacing: float,
):
"""
Calculate a second-order finite difference for equally spaced points.
https://web.media.mit.edu/~crtaylor/calculator.html
"""
result = (
array[down2 - 1] - 8 * array[down1 - 1] + 8 * array[up1 - 1] - array[up2 - 1]
) / (12 * spacing)
result[down2 == 0] = np.nan
result[down1 == 0] = np.nan
result[up1 == 0] = np.nan
result[up2 == 0] = np.nan
return result
@njit
def calculate_2d_gradient(
array: np.ndarray,
down2: np.ndarray,
down1: np.ndarray,
up1: np.ndarray,
up2: np.ndarray,
left2: np.ndarray,
left1: np.ndarray,
right1: np.ndarray,
right2: np.ndarray,
spacing: float,
):
"""Calculate the 2d gradient."""
return np.vstack(
(
calculate_gradient(array, left2, left1, right1, right2, spacing),
np.zeros(array.size),
calculate_gradient(array, down2, down1, up1, up2, spacing),
)
).T
[docs]
@autodoc_function
def perpendicular_gradient(mesh, input_array):
"""Calculate the perpendicular gradient of an array."""
kwargs = dict(
down2=mesh.index_neighbor.sel(dx=+0, dy=-2).values,
down1=mesh.index_neighbor.sel(dx=+0, dy=-1).values,
up1=mesh.index_neighbor.sel(dx=+0, dy=+1).values,
up2=mesh.index_neighbor.sel(dx=+0, dy=+2).values,
left2=mesh.index_neighbor.sel(dx=-2, dy=+0).values,
left1=mesh.index_neighbor.sel(dx=-1, dy=+0).values,
right1=mesh.index_neighbor.sel(dx=+1, dy=+0).values,
right2=mesh.index_neighbor.sel(dx=+2, dy=+0).values,
spacing=mesh.spacing_c,
)
with xr.set_options(keep_attrs=True):
perp_gradient = xr.apply_ufunc(
calculate_2d_gradient,
input_array,
kwargs=kwargs,
input_core_dims=[
["points"],
],
output_core_dims=[["new_points", "vector"]],
vectorize=True,
dask="parallelized",
dask_gufunc_kwargs={
"output_sizes": {"new_points": mesh.points, "vector": 3}
},
output_dtypes=[np.float64],
)
perp_gradient = perp_gradient.rename(new_points="points").assign_coords(
vector=["eR", "ePhi", "eZ"]
)
return perp_gradient
