"""Grid that provides common functionality for different 2D grid types."""
from abc import ABC
from functools import cached_property
import numpy as np
import xarray as xr
import pandas as pd
from scipy.interpolate import griddata
from scipy.spatial import KDTree
from torx.units import Quantity
from torx.units import check_units
from torx.vector import poloidal_vector
from torx.decorators import AutoDocClass
[docs]
class Grid2D(ABC, AutoDocClass):
"""2D regular grid on a poloidal cross section."""
# Plane of the simulation this grid describes, set by the 3D grid it is
# part of. Both stay None for a grid which is not part of a 3D grid.
phi_index = None
phi = None
# Arrays holding the coordinates of the points of the grid. Only these
# can hold the NaN padding used to align the planes of a 3D grid.
COORDINATE_ARRAYS = ("r_u", "z_u", "r_s", "z_s")
[docs]
def __init__(self,
r_unstructured,
z_unstructured,
norm: Quantity=Quantity("1m"),
init_matrices=False):
"""
Create a grid with the unstructured (R,Z) coordinates provided.
Accepts optionally a filepath to read in GRILLIX trunk data.
Accepts optionally a length quantity to normalize the grid.
"""
self.r_u, self.r_s, r_spacing, r_size = \
self.basis_from_unstructured_array(r_unstructured)
self.z_u, self.z_s, z_spacing, z_size = \
self.basis_from_unstructured_array(z_unstructured)
self.r_u = xr.DataArray(self.r_u, dims="points")
self.r_s = xr.DataArray(self.r_s, dims="R")
self.z_u = xr.DataArray(self.z_u, dims="points")
self.z_s = xr.DataArray(self.z_s, dims="Z")
assert(r_size == z_size), \
"Expected equal sized coordinate arrays but found " \
f"{r_size} != {z_size}"
self.size = r_size
self.shape = (self.r_s.size, self.z_s.size)
self.r_spacing, self.z_spacing = r_spacing, z_spacing
if np.isclose(r_spacing, z_spacing):
self.spacing = r_spacing
self._norm = norm
self._parallel_length_scale = None
self._perpendicular_length_scale = None
self._flipped_z = False
# Built on the first filled vector_to_matrix only, since its nearest
# neighbor search over the whole matrix is expensive for large grids
self.filled_point_index = None
if init_matrices:
self.setup_vector_to_matrix()
self.setup_filled_vector_to_matrix()
self.setup_perpendicular_gradient()
self._vector_to_matrix_initialized = True
self._perpendicular_gradient_initialized = True
else:
self._vector_to_matrix_initialized = False
self._perpendicular_gradient_initialized = False
@property
def R0(self):
"""Normalized radius."""
return self._norm
@R0.setter
def R0(self, value: Quantity):
"""Set the normalized radius."""
check_units(value, {"[length]":1}, "R0")
self._norm = value
def has_padding(self) -> bool:
"""
Check whether the grid contains NaN padding.
Returns
-------
bool
True if any coordinate array of the grid contains a NaN.
"""
return any(np.isnan(getattr(self, name)).any()
for name in self.COORDINATE_ARRAYS)
[docs]
@staticmethod
def basis_from_unstructured_array(x_unstructured):
"""Find the spacing and unique values of unstructured coordinates."""
x_unstructured = np.array(x_unstructured)
assert x_unstructured.ndim == 1
x_structured = np.unique(x_unstructured)
x_structured = x_structured[~np.isnan(x_structured)]
spacing = np.mean(np.diff(x_structured))
assert np.allclose(
np.diff(x_structured), spacing
), "Error: basis vector is not equally spaced."
return x_unstructured, x_structured, spacing, x_unstructured.size
def setup_vector_to_matrix(self):
"""Calculate arrays to improve the performance of vector to matrix."""
tricolumn_data = np.column_stack(
(self.r_u, self.z_u, np.arange(start=0, stop=self.size, dtype=int))
)
pd_dataframe = pd.DataFrame(tricolumn_data, columns=["x", "y", "z"])
# Makes a 2D array of indices
shaped_data = pd_dataframe.pivot_table(
values="z", index="y", columns="x", dropna=False
)
# Flatten array to give a vector
flattened_data = np.array(shaped_data).flatten()
# Replaces NaN with zero (will return the z[0] element in this place)
self.sort_indices = np.nan_to_num(flattened_data).astype("int")
# Add NaN vector to z[self.sort_indices] to replace z[0] with NaN
self.nan_vector = np.zeros_like(flattened_data)
self.nan_vector[np.isnan(flattened_data)] = np.nan
forward_sort = flattened_data[
np.logical_not(np.isnan(self.nan_vector))
].astype(int)
assert forward_sort.size == self.size
self.reverse_sort_indices = np.zeros_like(forward_sort)
self.reverse_sort_indices[forward_sort] = np.arange(forward_sort.size)
def setup_filled_vector_to_matrix(self):
"""
Calculate the array to perform a filled vector to matrix.
Calculates the 'filled point index' array which can be used to
index an unstructured array and return a filled 2D array, using
nearest neighbor interpolation for off-grid points.
"""
point_index = xr.DataArray(np.arange(self.size), dims="points")
r_mesh, z_mesh = np.meshgrid(self.r_s, self.z_s)
self.filled_point_index = griddata(
((self.r_u, self.z_u)),
point_index,
(r_mesh, z_mesh),
method="nearest"
).astype(int)
@cached_property
def _tree_point_index(self) -> np.ndarray:
"""
Return the grid index of every point of the k-d tree.
Grid points with non-finite coordinates are left out of the tree.
The grid size is appended, so that the index the tree returns for a
missing neighbor maps to the grid size.
"""
finite = np.isfinite(self.r_u.values) & np.isfinite(self.z_u.values)
return np.append(np.flatnonzero(finite), self.size)
@cached_property
def _kd_tree(self) -> KDTree:
"""
Return the k-d tree of the grid points.
Built on first use and shared by all nearest neighbor searches on
the grid. It is private, so that displaying the grid with rich does
not build it.
"""
points = self._tree_point_index[:-1]
return KDTree(np.column_stack((self.r_u.values[points],
self.z_u.values[points])))
def nearest_neighbors(
self,
r_query: np.ndarray,
z_query: np.ndarray,
n_neighbors: int,
**query_kwargs,
) -> tuple[np.ndarray, np.ndarray]:
"""
Return the nearest grid points of query points.
The grid points are searched with a k-d tree, which is built on the
first search and then kept, so that the memory needed grows with the
number of query points only, not with the size of the grid. Grid
points with non-finite coordinates are ignored. Where fewer than
n_neighbors grid points are found, for example for query points
with non-finite coordinates, the missing neighbors are returned at
infinite distance with the index self.size, as in KDTree.query.
Parameters
----------
r_query : np.ndarray
Radial coordinates of the query points.
z_query : np.ndarray
Vertical coordinates of the query points, same shape as r_query.
n_neighbors : int
Number of nearest grid points to return.
**query_kwargs
Passed on to scipy.spatial.KDTree.query, for example the norm p
or distance_upper_bound.
Returns
-------
distance : np.ndarray
Distances of the nearest grid points in ascending order, shape
r_query.shape + (n_neighbors,).
index : np.ndarray
Grid indices of the nearest grid points, of the same shape.
"""
r_query = np.asarray(r_query, dtype=float)
z_query = np.asarray(z_query, dtype=float)
distance = np.full(r_query.shape + (n_neighbors,), np.inf)
index = np.full(distance.shape, self.size)
# The k-d tree rejects non-finite query points
finite = np.isfinite(r_query) & np.isfinite(z_query)
finite_distance, tree_index = self._kd_tree.query(
np.column_stack((r_query[finite], z_query[finite])),
k=n_neighbors,
**query_kwargs,
)
distance[finite] = np.reshape(finite_distance, (-1, n_neighbors))
index[finite] = np.reshape(self._tree_point_index[tree_index],
(-1, n_neighbors))
return distance, index
def points_within_grid(
self,
r_test: np.ndarray,
z_test: np.ndarray,
expansion: float = 1.01,
) -> np.ndarray:
"""
Return whether sample points lie within the grid.
A sample point lies within the grid if at least four grid points are
closer to it than spacing * expansion in both R and Z. Sample points
with non-finite coordinates do not lie within the grid, and grid
points with non-finite coordinates are ignored. This tests whether
there is grid around each point, not whether it lies exactly within
the boundary of the grid, unlike Polygon2D.points_inside.
Parameters
----------
r_test : np.ndarray
Radial coordinates of the sample points.
z_test : np.ndarray
Vertical coordinates of the sample points, same shape as r_test.
expansion : float, optional
Factor widening the search distance beyond one grid spacing.
Returns
-------
np.ndarray
Boolean array of the shape of r_test, True for points within the
grid.
"""
# At least four grid points are close enough in both R and Z if the
# fourth nearest one is, in the maximum norm. Neighbors beyond the
# search distance are returned at infinite distance.
search_distance = self.spacing * expansion
distance, _ = self.nearest_neighbors(
r_test,
z_test,
4,
p=np.inf,
distance_upper_bound=search_distance,
)
return distance[..., -1] < search_distance
def nearest_non_collinear_neighbors(
self,
r_query: np.ndarray,
z_query: np.ndarray,
n_candidates: int = 16,
) -> np.ndarray:
"""
Return the nearest grid points of query points for a gradient stencil.
For every query point, returns the nearest and second nearest grid
point and the nearest grid point which is not collinear with these
two. Grid points at equal distance are taken in the order the k-d
tree returns them, since any of them gives a valid stencil. Three
points count as collinear if the sine of the angle between the
edges from the nearest one to the other two is below 1e-8, so that
the test does not depend on the grid spacing.
Parameters
----------
r_query : np.ndarray
Radial coordinates of the query points, shape (n_query,).
z_query : np.ndarray
Vertical coordinates of the query points, shape (n_query,).
n_candidates : int, optional
Number of nearest grid points searched first for a non-collinear
one. The search is widened up to the whole grid where none is
found.
Returns
-------
np.ndarray
The grid indices of the nearest, the second nearest and the
nearest non-collinear grid point, shape (3, n_query). Where all
grid points are collinear, the third nearest one is returned
instead. Query points with non-finite coordinates get the first
three grid points, since nothing is nearest to them.
"""
r_query = np.asarray(r_query, dtype=float)
z_query = np.asarray(z_query, dtype=float)
finite = np.isfinite(r_query) & np.isfinite(z_query)
grid_points = np.column_stack((self.r_u.values, self.z_u.values))
n_points = self._kd_tree.n
k = min(max(n_candidates, 3), n_points)
while True:
_, index = self.nearest_neighbors(
r_query[finite], z_query[finite], k)
# Offsets of the second nearest and the candidate points from
# the nearest one, which span a triangle of zero area if
# collinear
offset = grid_points[index] - grid_points[index[:, :1]]
det = offset[:, 1:2, 0] * offset[:, 2:, 1] \
- offset[:, 1:2, 1] * offset[:, 2:, 0]
# Dividing the area by the edge lengths gives the sine of the
# angle between the edges, which does not scale with the grid
length = np.hypot(offset[..., 0], offset[..., 1])
non_collinear = np.abs(det) > 1e-8 * length[:, 1:2] * length[:, 2:]
if k == n_points or np.all(np.any(non_collinear, axis=-1)):
break
k = min(2 * k, n_points)
first_non_collinear = 2 + np.argmax(non_collinear, axis=-1)
neighbors = np.tile(np.arange(3)[:, None], r_query.size)
neighbors[:, finite] = (
index[:, 0],
index[:, 1],
index[np.arange(index.shape[0]), first_non_collinear],
)
return neighbors
def setup_perpendicular_gradient(self):
"""
Calculate array to create a perpendicular gradient operator on the grid.
Create neighbor index arrays for left, right, up and down neighbors
that can be used to evaluate the perpendicular gradient
"""
point_index = xr.DataArray(np.arange(self.size), dims="points")
# The idea is to simply store the left, right, up, down neighbors for
# the structured (R, Z) data and then convert it back to the
# unstructured form.
point_matrix = self.vector_to_matrix(point_index)
points = point_matrix.values
right_matrix = point_matrix.copy()
left_matrix = point_matrix.copy()
up_matrix = point_matrix.copy()
down_matrix = point_matrix.copy()
# We need to take care of the boundaries which are given by NaN
# neighbors. Here we set the index to the point itself which can be
# used straightforwardly in a forward stencil. The matrix has the
# dims (Z, R).
right_matrix.values[:, :-1] = np.where(
np.isnan(points[:, 1:]), points[:, :-1], points[:, 1:])
left_matrix.values[:, 1:] = np.where(
np.isnan(points[:, :-1]), points[:, 1:], points[:, :-1])
up_matrix.values[:-1, :] = np.where(
np.isnan(points[1:, :]), points[:-1, :], points[1:, :])
down_matrix.values[1:, :] = np.where(
np.isnan(points[:-1, :]), points[1:, :], points[:-1, :])
self.right_index = self.matrix_to_vector(right_matrix).astype(int)
self.left_index = self.matrix_to_vector(left_matrix).astype(int)
self.up_index = self.matrix_to_vector(up_matrix).astype(int)
self.down_index = self.matrix_to_vector(down_matrix).astype(int)
def vector_to_matrix(
self, unstructured_data, filled: bool=False, allow_rechunk: bool=True
):
"""
Convert an array from the unstructured to structured representation.
Converts tricolumn data 'z(..., points, ....)'
to matrix data 'z(..., R, Z, ...)'
If unstructured_data is supplied as an xarray, the 'points' dimension
is shaped. If unstructured_data is supplied as a numpy array, it can
only be 1D
"""
if not isinstance(unstructured_data, xr.DataArray):
assert unstructured_data.ndim == 1
unstructured_data = xr.DataArray(unstructured_data, dims="points")
assert unstructured_data.sizes["points"] == self.size, \
("Unstructured data not the same size as input data. "
f"Unstructured data size: {unstructured_data.sizes['points']}, "
f"Input data size: {self.size}.")
if not self._vector_to_matrix_initialized:
self.setup_vector_to_matrix()
self._vector_to_matrix_initialized = True
if filled and self.filled_point_index is None:
self.setup_filled_vector_to_matrix()
structured_data = xr.apply_ufunc(
_stack,
unstructured_data,
kwargs=dict(
filled=filled,
sort_indices=self.sort_indices,
nan_vector=self.nan_vector,
filled_point_index=self.filled_point_index,
shape=(self.z_s.size, self.r_s.size),
),
input_core_dims=[
["points"],
],
output_core_dims=[["Z", "R"]],
keep_attrs=True,
vectorize=True,
dask="parallelized",
dask_gufunc_kwargs=dict(
output_sizes={"Z": self.z_s.size, "R": self.r_s.size},
allow_rechunk=allow_rechunk,
),
output_dtypes=[np.float64],
).assign_coords(R=self.r_s, Z=self.z_s)
return structured_data
def matrix_to_vector(self, structured_data, allow_rechunk: bool = True):
"""
Convert an array from the structured to unstructured representation.
Converts matrix data 'z(..., R, Z, ...)'
to tricolumn data 'z(..., points, ....)'
"""
if not isinstance(structured_data, xr.DataArray):
assert structured_data.ndim == 2
structured_data = xr.DataArray(structured_data, dims=("Z", "R"))
assert structured_data.sizes["R"] == self.r_s.size
assert structured_data.sizes["Z"] == self.z_s.size
if not self._vector_to_matrix_initialized:
self.setup_vector_to_matrix()
self._vector_to_matrix_initialized = True
unstructured_data = xr.apply_ufunc(
_flatten,
structured_data,
kwargs=dict(
nan_vector=self.nan_vector,
reverse_sort_indices=self.reverse_sort_indices,
),
input_core_dims=[
["Z", "R"],
],
output_core_dims=[["points"]],
keep_attrs=True,
vectorize=True,
dask="parallelized",
dask_gufunc_kwargs=dict(
output_sizes={"points": self.size}, allow_rechunk=allow_rechunk
),
output_dtypes=[np.float64],
)
return unstructured_data
def perpendicular_gradient(self, array: xr.DataArray):
"""Return the perpendicular gradient of the given array as a vector."""
grad_r = self._gradient_R(array)
grad_z = self._gradient_Z(array)
return poloidal_vector(grad_r, grad_z,
dims=array.dims,
attrs=grad_r.attrs)
def perpendicular_gradient_component(self, array: xr.DataArray, \
component="R"):
"""
Return a specific component of the perpendicular gradient.
Component can either be a string specifying "R" or "Z", or
alternatively an xarray DataArray representing a vector where to
project the gradient to.
"""
if isinstance(component, str):
if component == "R":
return self._gradient_R(array)
elif component == "Z":
return self._gradient_Z(array)
else:
raise RuntimeError(f"Component {component} is not available " \
"for perpendicular gradient.")
elif isinstance(component, xr.DataArray):
assert "vector" in component.dims, \
"Component should be a vector with dimension vector " \
"in the DataArray."
grad_r = self._gradient_R(array)
grad_z = self._gradient_Z(array)
grad_proj = component.sel(vector="eR") * grad_r \
+ component.sel(vector="eZ") * grad_z
if ("norm" in array.attrs and "norm" in component.attrs):
grad_proj.attrs["norm"] = component.norm * grad_r.norm
return grad_proj
def radial_perpendicular_gradient(self, array: xr.DataArray, equi):
"""Return the radial component of the perpendicular gradient."""
e_rho = equi.magfield_vector_radial(self.r_u, self.z_u, normalize=True)
return self.perpendicular_gradient_component(array, e_rho)
def _gradient_R(self, array: xr.DataArray):
"""
Return the R component of the perpendicular gradient.
Base method to calculate the R component of the perpendicular gradient,
utilizing the neighbor index arrays created in the setup routine. Uses
second order central differences in the interior and first order
forward/backward differences at the boundaries. Works on structured or
unstructured data. The latter is preferred for speed.
"""
if not self._vector_to_matrix_initialized:
self.setup_vector_to_matrix()
self._vector_to_matrix_initialized = True
if not self._perpendicular_gradient_initialized:
self.setup_perpendicular_gradient()
self._perpendicular_gradient_initialized = True
if not "points" in array.dims:
assert ("R" in array.dims and "Z" in array.dims), \
"Neither dimension points, R or Z were found in the array."
# If array given in structured R, Z, we convert to unstructured
# points, perform the calculation and convert back in the end
array = self.matrix_to_vector(array)
convert_to_RZ = True
else:
convert_to_RZ = False
grad_r = (array.isel({"points":self.right_index}) \
- array.isel({"points":self.left_index})) \
* 0.5 / self.r_spacing
# Apply correction at edge point to achieve first order forward stencil
edge = (self.right_index == self.right_index.points) \
| (self.left_index == self.left_index.points)
edge = xr.DataArray(edge, dims="points")
grad_r *= 1 + edge
# If no norm attribute given, assume normalized units
if "norm" in array.attrs:
grad_r.attrs["norm"] = array.norm / self.R0
if convert_to_RZ:
return self.vector_to_matrix(grad_r)
else:
return grad_r
def _gradient_Z(self, array: xr.DataArray):
"""
Return the Z component of the perpendicular gradient.
Base method to calculate the Z component of the perpendicular gradient,
utilizing the neighbor index arrays created in the setup routine. Uses
second order central differences in the interior and first order
forward/backward differences at the boundaries. Works on structured or
unstructured data. The latter is preferred for speed.
"""
if not self._vector_to_matrix_initialized:
self.setup_vector_to_matrix()
self._vector_to_matrix_initialized = True
if not self._perpendicular_gradient_initialized:
self.setup_perpendicular_gradient()
self._perpendicular_gradient_initialized = True
if not "points" in array.dims:
assert ("R" in array.dims and "Z" in array.dims), \
"Neither dimension points, R or Z were found in the array."
array = self.matrix_to_vector(array)
convert_to_RZ = True
else:
convert_to_RZ = False
grad_z = (array.isel({"points":self.up_index}) \
- array.isel({"points":self.down_index})) \
* 0.5 / self.z_spacing
# Apply correction at edge point to achieve first order forward stencil
edge = (self.up_index == self.up_index.points) \
| (self.down_index == self.down_index.points)
edge = xr.DataArray(edge, dims="points")
grad_z *= 1 + edge
# If no norm attribute given, assume normalized units
if "norm" in array.attrs:
grad_z.attrs["norm"] = array.norm / self.R0
if convert_to_RZ:
return self.vector_to_matrix(grad_z)
else:
return grad_z
def get_unstructured_index(self, rval, zval):
"""Return index of r_u and z_u that correspond to rval and zval."""
dist2 = (self.r_u - rval)**2 + (self.z_u - zval)**2
ui = np.argmin(dist2.data).item()
return ui
def coords_like(self, array: xr.DataArray):
"""Return (R,Z) coordinates that fit the dims of the given array."""
assert(isinstance(array, xr.DataArray)), \
"Array must be of type xarray DataArray."
if("R" in array.dims and "Z" in array.dims):
r_norm = self.r_s
z_norm = self.z_s
elif("points" in array.dims):
r_norm = self.r_u
z_norm = self.z_u
else:
r_norm = None
z_norm = None
return r_norm, z_norm
def _stack(unstructured_data, filled, sort_indices, nan_vector,
filled_point_index, shape):
"""
Return the inner shaping loop for vector_to_matrix.
The matrix has the given shape (n_Z, n_R). filled_point_index is only
used if filled is True.
"""
if filled:
return unstructured_data[filled_point_index]
return (unstructured_data[sort_indices] + nan_vector).reshape(shape)
def _flatten(input_array, nan_vector, reverse_sort_indices):
"""Return the inner flatten loop for matrix_to_vector."""
return input_array.flatten()[np.logical_not(np.isnan(nan_vector))][
reverse_sort_indices
]
