"""Defines 3D grids made up of collections of 2D grids."""
from abc import ABC
import numpy as np
import xarray as xr
from pathlib import Path
from typing import Optional, Sequence, Union
from .grid_2d_m import Grid2D
from torx.units import Quantity
from torx.arrays import combine_obj_list
from torx.units import check_units
from torx.decorators import autodoc_class
from functools import cached_property
[docs]
@autodoc_class
class Grid3D(ABC):
"""Represents a 3D grid, composed of an array of 2D grids at each phi."""
[docs]
def __init__(self,
grids: xr.DataArray,
plane_indices: Optional[Sequence[int]] = None):
"""
Initialize the 3D grid.
Parameters
----------
grids : xr.DataArray
Array holding the Grid2D of every plane.
plane_indices : Sequence[int], optional
Index of every plane in the simulation it belongs to. Only
needed when the array holds a part of the planes, since the
position of a plane is used otherwise.
"""
self._assert_no_padding(grids)
self._set_plane(grids, plane_indices)
self.grid_array = grids
self.name = 'Grid3D'
@staticmethod
def _assert_no_padding(grid_array: xr.DataArray) -> None:
"""
Check that no plane of a 3D grid holds padding.
Parameters
----------
grid_array : xr.DataArray
Array of Grid2D objects, one for each plane.
Raises
------
ValueError
If the 2D grid of any plane holds padding.
"""
for index in np.ndindex(grid_array.shape):
if grid_array.values[index].has_padding():
raise ValueError(
f"The 2D grid of plane {index[0]} holds NaN padding. "
"Planes of a 3D grid must be built from the points of "
"their own plane only, see Grid3D.from_rz."
)
@staticmethod
def _set_plane(grid_array: xr.DataArray,
plane_indices: Optional[Sequence[int]]) -> None:
"""
Tell the 2D grid of every plane which plane of a simulation it is.
Parameters
----------
grid_array : xr.DataArray
Array of Grid2D objects, one for each plane.
plane_indices : Sequence[int] or None
Index of every plane in the simulation it belongs to. The position
of a plane in the array is used when None is given, which is only
correct if the array holds all planes of the simulation.
Raises
------
ValueError
If an index is given for a different number of planes.
"""
if plane_indices is None:
plane_indices = range(grid_array.size)
elif len(plane_indices) != grid_array.size:
raise ValueError(
f"Got {len(plane_indices)} plane indices "
f"for {grid_array.size} planes."
)
phi_values = grid_array.coords["phi"].values
for position, grid in enumerate(grid_array.values.flat):
grid.phi_index = int(plane_indices[position])
grid.phi = phi_values[position]
[docs]
@classmethod
def from_rz(cls,
r_unstructured,
z_unstructured):
"""Create the 3D grid from xarrays of unstructured R and Z grids."""
assert isinstance(r_unstructured, xr.DataArray) \
and isinstance(z_unstructured, xr.DataArray), \
"Unstructured grids must be Xarray DataArrays!"
# Helper function to remove NaNs and build a 2D grid for each plane.
def _build_2d_grid_no_nans(r_plane, z_plane):
valid_mask = ~np.isnan(r_plane) & ~np.isnan(z_plane)
r_clean = r_plane[valid_mask]
z_clean = z_plane[valid_mask]
return Grid2D(r_clean, z_clean)
grids = xr.apply_ufunc(_build_2d_grid_no_nans,
r_unstructured, z_unstructured,
input_core_dims=[["dim_RZ"], ["dim_RZ"]],
output_core_dims=[[]],
vectorize=True,
output_dtypes=[object])
return cls(grids)
[docs]
@classmethod
def from_multigrid_files(
cls,
dirpath: Path,
planes: Optional[Union[int, Sequence[int]]] = None,
staggered: bool = False
):
"""
Create a 3D grid object from a directory of multigrid files.
Parameters
----------
dirpath : str
Path to the directory containing multigrid files.
planes : list or None
List of plane indices to load. If None, all planes are loaded.
"""
from torx.specializations.grillix import grid_2d_from_multigrid_file
# Find all files matching the pattern
path_list = sorted(list(dirpath.glob("multigrids_plane*.nc")))
if not path_list:
raise FileNotFoundError(
"Error: No files matching 'multigrids_plane*.nc' "
"were found in the directory."
)
all_filepaths = np.atleast_1d(path_list)
all_phi = np.linspace(0, 2*np.pi, len(all_filepaths), endpoint=False)
# Filter files based on planes
if planes is None:
phi = all_phi
plane_indices = None
filepaths = xr.DataArray(
all_filepaths,
dims="phi"
)
else:
planes = np.atleast_1d(planes)
phi = all_phi[planes]
filepaths = xr.DataArray(
all_filepaths[planes],
dims="phi"
)
# The position of a plane is not its index in the simulation
# when only a part of the planes is loaded.
plane_indices = planes
grids = xr.apply_ufunc(grid_2d_from_multigrid_file,
filepaths, staggered,
input_core_dims=[[], []],
output_core_dims=[[]],
vectorize=True,
output_dtypes=[object]).assign_coords(phi=phi)
return cls(grids, plane_indices=plane_indices)
@property
def grid_array(self):
"""Return the arrays of values in the grid."""
return self._grid_array
@grid_array.setter
def grid_array(self, value):
"""
Set the arrays of values in the grid.
Everything derived from the 2D grids is dropped, since it describes
the planes which were held before.
"""
self._grid_array = value
self.clear_cached_properties()
def _map_over_planes(self, function, dtype) -> xr.DataArray:
"""
Apply a function to the 2D grid of every plane.
Parameters
----------
function : callable
Function taking the Grid2D of a plane and returning a value.
dtype : type
Data type of the returned values.
Returns
-------
xr.DataArray
Result for every plane, along the dimension of the planes.
"""
values = np.empty(self._grid_array.shape, dtype=dtype)
for index in np.ndindex(self._grid_array.shape):
values[index] = function(self._grid_array.values[index])
return self._grid_array.copy(data=values)
@cached_property
def npoints(self) -> xr.DataArray:
"""Number of grid points in each plane."""
return self._map_over_planes(lambda grid: grid.size, int)
@cached_property
def shape(self) -> xr.DataArray:
"""Shape of the structured grid of each plane."""
return self._map_over_planes(lambda grid: grid.shape, object)
@cached_property
def r_u(self) -> xr.DataArray:
"""Unstructured R values."""
return combine_obj_list(self.grid_array.values, "r_u", "phi",
self.grid_array.coords["phi"].values)
@cached_property
def z_u(self) -> xr.DataArray:
"""Unstructured Z values."""
return combine_obj_list(self.grid_array.values, "z_u", "phi",
self.grid_array.coords["phi"].values)
@cached_property
def r_s(self) -> xr.DataArray:
"""Structured R values."""
return combine_obj_list(self.grid_array.values, "r_s", "phi",
self.grid_array.coords["phi"].values)
@cached_property
def z_s(self) -> xr.DataArray:
"""Structured Z values."""
return combine_obj_list(self.grid_array.values, "z_s", "phi",
self.grid_array.coords["phi"].values)
[docs]
def clear_cached_properties(self):
"""Clear all attributes decorated with @cached_property."""
# We iterate over the class members to find cached_property descriptors
cls = self.__class__
for name, value in cls.__dict__.items():
if isinstance(value, cached_property):
# If the property has been computed,
# it exists in the instance __dict__
if name in self.__dict__:
delattr(self, name)
def __getattr__(self, name):
"""
Give access to the attributes of the array of 2D grids.
The name of a dimension or coordinate gives the coordinate itself, so
that grid.phi gives the phi value of every plane. Any other name is
looked up on the array of 2D grids, which is what makes grid.sel and
grid.isel work.
"""
# Private names are never passed on.
if name.startswith("_"):
raise AttributeError(f"'Grid3D' object has no attribute '{name}'")
# This is only called when the normal lookup fails, so the array of
# 2D grids has to be read from the instance itself.
grid_array = self.__dict__.get("_grid_array")
if grid_array is None:
raise AttributeError(f"'Grid3D' object has no attribute '{name}'")
if name in grid_array.dims or name in grid_array.coords:
return grid_array.coords[name]
try:
return getattr(grid_array, name)
except AttributeError:
raise AttributeError(
f"'Grid3D' object has no attribute '{name}'"
) from None
[docs]
def __getitem__(self, key):
"""
Return the coordinate of a dimension of the 3D grid.
Parameters
----------
key : str
Name of a dimension or coordinate, such as 'phi'.
Returns
-------
xr.DataArray
Value of the coordinate for every plane.
"""
grid_array = self._grid_array
if isinstance(key, str) and \
(key in grid_array.dims or key in grid_array.coords):
return grid_array.coords[key]
raise KeyError(key)
[docs]
def set_R0(self, value: Quantity):
"""Set the normalized radius for all of the 2D grids."""
check_units(value, {"[length]":1}, "R0")
for grid in self._grid_array.values.flat:
grid.R0 = value
[docs]
def sel_phi(self, phi_value=0.0) -> Grid2D:
"""
Return the 2D grid of the plane at the given phi value.
The grid is the one held by the 3D grid, so the index arrays it has
built already are kept when the same plane is selected again.
Parameters
----------
phi_value : float
Value of the phi coordinate of the wanted plane.
Returns
-------
Grid2D
Grid of the selected plane.
"""
return self.sel(phi=phi_value).item()
[docs]
def isel_phi(self, phi_index=0) -> Grid2D:
"""
Return the 2D grid of the plane at the given phi index.
The grid is the one held by the 3D grid, so the index arrays it has
built already are kept when the same plane is selected again.
Parameters
----------
phi_index : int
Index of the wanted plane along phi.
Returns
-------
Grid2D
Grid of the selected plane.
"""
return self.isel(phi=phi_index).item()
[docs]
def isel_planes(self, phi_indices: Sequence[int]) -> "Grid3D":
"""
Return a 3D grid holding the planes at the given phi indices.
The 2D grids are the ones held by this grid, so the index arrays
they have built already are kept.
Parameters
----------
phi_indices : Sequence[int]
Positions along phi of the wanted planes.
Returns
-------
Grid3D
Grid of the selected planes, which keep the index of the plane
of the simulation they describe.
"""
planes = self.isel(phi=list(phi_indices))
return type(self)(
planes,
plane_indices=[grid.phi_index for grid in planes.values.flat]
)
[docs]
def perpendicular_gradient(self, array: xr.DataArray) -> xr.DataArray:
"""Return the perpendicular gradient of the given array as a vector."""
# Rename dimensions if necessary.
if 'planes' in array.dims:
array = array.rename({'planes': 'phi'})
if 'npoints_max' in array.dims:
array = array.rename({'npoints_max': 'points'})
# Check that the array has the expected dimensions.
assert ("points" in array.dims and "phi" in array.dims), \
"The array must have dimensions 'points' and 'phi'."
# Make sure to keep the staggered attribute if it exists.
has_staggered = "staggered" in array.attrs
if has_staggered:
staggered_val = array.attrs["staggered"]
# Helper function that applies the perp grad to a single 2D grid.
# It pads the resulting grad with NaNs for vectorized computation.
def _plane_perp_grad(grid_2d_obj, data_slice, n_points):
if hasattr(grid_2d_obj, 'item'):
grid_2d_obj = grid_2d_obj.item()
# Must set points dim or perpendicular_gradient fails.
array_2d_valid = xr.DataArray(data_slice, dims=['points'])
grad_2d = grid_2d_obj.perpendicular_gradient(array_2d_valid)
n = grad_2d.sizes["points"]
if n < n_points:
grad_2d = grad_2d.pad(
points=(0, n_points - n), constant_values=np.nan
)
return grad_2d
# Chunk beforehand to avoid dask issues with objects.
max_points = int(self.npoints.max())
grid_array = self.grid_array.chunk({"phi": 1})
# Use xr.apply_ufunc to calculate the gradient plane by plane.
gradient = xr.apply_ufunc(
_plane_perp_grad,
grid_array,
array,
max_points,
input_core_dims=[[], ['points'], []],
output_core_dims=[['points', 'vector']],
exclude_dims=set(['points']),
vectorize=True,
dask="parallelized",
output_dtypes=[np.float64],
keep_attrs=True,
dask_gufunc_kwargs={
"output_sizes": {
"points": max_points,
"vector": 3
}
}
)
# If the input array had a "staggered" attribute, preserve it.
if has_staggered:
gradient.attrs["staggered"] = staggered_val
return gradient
