"""Functions to build vectors."""
import xarray as xr
import numpy as np
from typing import Union
from torx import Quantity
from torx.autogrid_decorators_m import autogrid_function
from torx.autodoc_decorators_m import autodoc_function
[docs]
@autodoc_function
def toroidal_vector(
input_phi: Union[np.ndarray, xr.DataArray],
dims: list=None,
coords: dict={},
attrs: dict={},
):
"""
Build a cylindrical vector-field array from a phi component array.
Zeros the (R, Z) components.
"""
# Handle data format, dims and coords
input_phi = _conv_to_xarray(input_phi, dims=dims)
dims, coords = _add_vector_dim(input_phi, dims, coords)
input_phi = input_phi.expand_dims("vector").transpose(*dims)
input_r = xr.zeros_like(input_phi)
input_z = xr.zeros_like(input_phi)
vector_array = xr.concat([input_r, input_phi, input_z], dim="vector")
vector_array = vector_array.chunk(chunks={"vector": 3})
return xr.DataArray(vector_array, coords=coords, attrs=attrs)
[docs]
@autodoc_function
def poloidal_vector(
input_r: Union[np.ndarray, xr.DataArray],
input_z: Union[np.ndarray, xr.DataArray],
dims: list=None,
coords: dict={},
attrs: dict={},
):
"""
Build a cylindrical vector-field array from (R, Z) component arrays.
Zeros the phi component.
dims should be the names of the dimensions (i.e. from input_array.dims)
Optional: coords and attrs are used to set the coords and attributes of
the output array.
"""
assert (np.shape(input_r) == np.shape(input_z)), \
"Shapes of r and z do not match."
# Handle data format, dims and coords
input_r = _conv_to_xarray(input_r, dims=dims)
input_z = _conv_to_xarray(input_z, dims=dims)
dims, coords = _add_vector_dim(input_r, dims, coords)
input_r = input_r.expand_dims("vector").transpose(*dims)
input_z = input_z.expand_dims("vector").transpose(*dims)
input_phi = xr.zeros_like(input_r)
vector_array = xr.concat([input_r, input_phi, input_z], dim="vector")
vector_array = vector_array.chunk(chunks={"vector": 3})
return xr.DataArray(vector_array, coords=coords, attrs=attrs)
[docs]
@autodoc_function
def cylindrical_vector(
input_r: Union[np.ndarray, xr.DataArray],
input_phi: Union[np.ndarray, xr.DataArray],
input_z: Union[np.ndarray, xr.DataArray],
dims: list=None,
coords: dict={},
attrs: dict={},
):
"""
Build a cylindrical vector-field array from (R, phi, Z) component arrays.
dims should be the names of the dimensions (i.e. from input_array.dims)
Optional: coords and attrs are used to set the coords and attributes of
the output array.
"""
assert (np.shape(input_r) == np.shape(input_z)), \
"Shapes of r and z do not match."
# Handle data format, dims and coords
input_r = _conv_to_xarray(input_r, dims=dims)
input_phi = _conv_to_xarray(input_phi, dims=dims)
input_z = _conv_to_xarray(input_z, dims=dims)
dims, coords = _add_vector_dim(input_r, dims, coords)
input_r = input_r.expand_dims("vector").transpose(*dims)
input_phi = input_phi.expand_dims("vector").transpose(*dims)
input_z = input_z.expand_dims("vector").transpose(*dims)
vector_array = xr.concat([input_r, input_phi, input_z], dim="vector")
vector_array = vector_array.chunk(chunks={"vector": 3})
return xr.DataArray(vector_array, coords=coords, attrs=attrs)
[docs]
@autodoc_function
@autogrid_function
def eR_unit_vector(r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return a unit vector pointing in the radial (R) direction."""
if kwargs["is_structured"]:
size = (np.size(z_norm), np.size(r_norm))
r_comp = xr.DataArray(np.ones(size), dims=["Z", "R"])
z_comp = xr.DataArray(np.zeros(size), dims=["Z", "R"])
else:
r_comp = xr.DataArray(np.ones_like(r_norm), dims="points")
z_comp = xr.DataArray(np.zeros_like(r_norm), dims="points")
return poloidal_vector(r_comp, z_comp, attrs={"norm":Quantity(1.0)},)
[docs]
@autodoc_function
@autogrid_function
def ePhi_unit_vector(r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return a unit vector pointing in the toroidal (phi) direction."""
if kwargs["is_structured"]:
size = (np.size(z_norm), np.size(r_norm))
phi_comp = xr.DataArray(np.ones(size), dims=["Z", "R"])
else:
phi_comp = xr.DataArray(np.ones_like(r_norm), dims="points")
return toroidal_vector(phi_comp, attrs={"norm":Quantity(1.0)})
[docs]
@autodoc_function
@autogrid_function
def eZ_unit_vector(r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return a unit vector pointing in the vertical (Z) direction."""
if kwargs["is_structured"]:
size = (np.size(z_norm), np.size(r_norm))
r_comp = xr.DataArray(np.zeros(size), dims=["Z", "R"])
z_comp = xr.DataArray(np.ones(size), dims=["Z", "R"])
else:
r_comp = xr.DataArray(np.zeros_like(r_norm), dims="points")
z_comp = xr.DataArray(np.ones_like(r_norm), dims="points")
return poloidal_vector(r_comp, z_comp, attrs={"norm":Quantity(1.0)},)
def _add_vector_dim(
input_array: xr.DataArray,
dims: list=None,
coords: dict={},
):
"""
Add a 'vector' dimension with coords (eR, ePhi, eZ).
By default adds to the existing 'dims' and 'coords' of the input array.
"""
if (dims is None and isinstance(input_array, xr.DataArray)):
dims = input_array.dims
if (coords is None and isinstance(input_array, xr.DataArray)):
coords = input_array.coords
if isinstance(dims, str):
dims = [dims]
dims = list(dims)
dims.append("vector")
coords = dict(coords)
coords["vector"] = ["eR", "ePhi", "eZ"]
return dims, coords
def _conv_to_xarray(input_array, dims: list=None):
"""Convert input xr.DataArray if not already provided as such."""
if not isinstance(input_array, xr.DataArray):
return xr.DataArray(input_array, dims=dims)
else:
return input_array
