"""Abstract base class for all equilibria."""
from abc import ABC, abstractmethod
import numpy as np
import xarray as xr
import torx
from torx import make_xarray
from torx.vector import poloidal_vector, cylindrical_vector, vector_dot
from torx.autogrid_decorators_m import autogrid_method, autocoord_method
from torx.geometry.polygon_2d_m import Polygon2D
from torx.autodoc_decorators_m import autodoc_class
[docs]
@autodoc_class
class EquilibriumBaseClass(ABC):
"""
Base class for equilibrium objects.
The equilibrium works with normalized length units internally. This
normalization should be set to the axis location by all inheriting classes.
The functions in this class should accept either dimensionless lengths
which are assumed normalized or lengths with a unit which should be
converted to normalized automatically.
"""
@property
@abstractmethod
def axis_r_norm(self) -> xr.DataArray:
"""Magnetic axis radial position in normalized units."""
return make_xarray(1.0)
@property
@abstractmethod
def axis_z_norm(self) -> xr.DataArray:
"""Magnetic axis vertical position in normalized units."""
return make_xarray(0.0)
@property
@abstractmethod
def B0(self) -> torx.Quantity:
"""Reference magnetic field."""
return torx.Quantity("1T")
@property
@abstractmethod
def R0(self) -> torx.Quantity:
"""Reference length scale."""
return torx.Quantity("1m")
[docs]
@abstractmethod
def normalized_flux_surface_label(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""Return flux surface label, 0 at the axis and 1 at the separatrix."""
[docs]
@abstractmethod
def magfield_component_r(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return normalized magnetic field in the R direction."""
[docs]
@abstractmethod
def magfield_component_z(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return normalized magnetic field in the Z direction."""
[docs]
@abstractmethod
def magfield_component_toroidal(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""Return normalized magnetic field in the toroidal direction."""
[docs]
@abstractmethod
def get_boundary_polygon(self, *args) -> Polygon2D:
"""Return the boundary polygon."""
[docs]
def convert_length_to_normalized(self, length):
"""
Convert a quantity of dimension length to normalized units.
With that it can be used in the equilibrium.
"""
if isinstance(length, torx.Quantity):
if length.units != "":
length = length.m * (length.u / self.R0).to("").m
if isinstance(length, xr.DataArray):
if "norm" in length.attrs:
if length.norm != torx.Quantity("dimensionless"):
length = length * (length.norm / self.R0).to("").m
del length.attrs["norm"]
return length
[docs]
@autogrid_method
def jacobian(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return the Jacobian of the cylindrical coordinate system."""
if kwargs["is_structured"]:
r_mesh, _ = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
r_mesh = self.convert_length_to_normalized(r_mesh)
return xr.DataArray(r_mesh,
attrs={"norm":self.R0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_poloidal(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""Return the absolute value of the poloidal magnetic field."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
B_x = self.magfield_component_r(r_norm, z_norm, **kwargs)
B_y = self.magfield_component_z(r_norm, z_norm, **kwargs)
return np.sqrt(B_x**2 + B_y**2)
[docs]
@autogrid_method
def magfield_absolute(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return the absolute value of the magnetic field."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
B_x = self.magfield_component_r(r_norm, z_norm, **kwargs)
B_y = self.magfield_component_z(r_norm, z_norm, **kwargs)
B_t = self.magfield_component_toroidal(r_norm, z_norm, **kwargs)
return np.sqrt(B_x**2 + B_y**2 + B_t**2)
[docs]
@autogrid_method
def magfield_vector_poloidal(self, r_norm, z_norm, normalize=False,
**kwargs) -> xr.DataArray:
"""
Return a vector in the direction of the poloidal magnetic field.
That is along a flux surface with zero toroidal component.
If 'normalize=True' the magnetic field unit vector is returned.
"""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
B_x = self.magfield_component_r(r_norm, z_norm, **kwargs)
B_y = self.magfield_component_z(r_norm, z_norm, **kwargs)
if normalize:
B_abs = np.sqrt(B_x**2 + B_y**2)
B_x = B_x / B_abs
B_y = B_y / B_abs
norm = B_x.norm / B_abs.norm
else:
norm = self.B0
return poloidal_vector(B_x, B_y, attrs={"norm":norm},
dims=kwargs["dims"], coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_vector_radial(self, r_norm, z_norm, normalize=False,
reverse: bool=None, reverse_eps: float=1e-2,
**kwargs) -> xr.DataArray:
"""
Return a vector in the direction normal to the radial magnetic field.
That is across a flux surface.
If 'normalize=True' the magnetic field unit vector is returned.
Optionally, the argument reverse flips the sign of the vector. Per
default the sign is calculated such that the vector points outwards.
This automatic behavior is controlled with the parameter reverse_eps.
"""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
B_x = self.magfield_component_r(r_norm, z_norm, **kwargs)
B_y = self.magfield_component_z(r_norm, z_norm, **kwargs)
if normalize:
B_abs = np.sqrt(B_x**2 + B_y**2)
B_x = B_x / B_abs
B_y = B_y / B_abs
norm = B_x.norm / B_abs.norm
else:
norm = self.B0
if reverse is None:
reverse = not self.radial_unit_vector_points_outwards(
self.axis_r_norm, self.axis_z_norm, eps=reverse_eps, **kwargs
)
if reverse:
r_x = B_y
r_y = -B_x
else:
r_x = -B_y
r_y = B_x
return poloidal_vector(r_x, r_y, attrs={"norm":norm},
dims=kwargs["dims"], coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_vector(self, r_norm, z_norm, normalize=False,
**kwargs) -> xr.DataArray:
"""
Return the equilibrium magnetic field as a vector array.
If 'normalize=True' the magnetic field unit vector is returned.
"""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
B_x = self.magfield_component_r(r_norm, z_norm, **kwargs)
B_y = self.magfield_component_z(r_norm, z_norm, **kwargs)
B_t = self.magfield_component_toroidal(r_norm, z_norm, **kwargs)
if normalize:
B_abs = np.sqrt(B_x**2 + B_y**2 + B_t**2)
B_x = B_x / B_abs
B_y = B_y / B_abs
B_t = B_t / B_abs
norm = B_t.norm / B_abs.norm
else:
norm = self.B0
return cylindrical_vector(B_x, B_t, B_y, attrs={"norm":norm},
dims=kwargs["dims"], coords=kwargs["coords"])
[docs]
def get_separatrix(self, r_norm, z_norm, level: float = 1.0) -> list:
"""
Return the separatrix contour.
Can plot a specific contour via plt.plot(*separatrix[index].T)
where index is a contour index.
"""
from torx.analysis.contour_finding_m import find_contours
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
assert r_norm.ndim == 1 and z_norm.ndim == 1
return find_contours(
r_norm,
z_norm,
self.normalized_flux_surface_label(r_norm, z_norm),
level=level,
)
[docs]
@autocoord_method
def project_vector_to_magfield(
self,
input_vector: xr.DataArray,
r_norm,
z_norm,
direction: str="parallel",
**kwargs) -> xr.DataArray:
"""
Project an input vector in direction related to the magnetic field.
Can be the direction parallel or perpendicular to the magnetic field.
"""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
assert("vector" in input_vector.dims)
# Find the field vector to project upon
assert direction in ["parallel", "radial", "poloidal"],\
"Choose either 'parallel', 'radial' or 'poloidal' as projection direction"
if direction == "parallel":
field_vector = self.magfield_vector(r_norm, z_norm, normalize=True,
**kwargs)
elif direction == "radial":
field_vector = self.magfield_vector_radial(r_norm, z_norm,
normalize=True, **kwargs)
elif direction == "poloidal":
field_vector = self.magfield_vector_poloidal(r_norm, z_norm,
normalize=True,
**kwargs)
projected_vec = vector_dot(input_vector, field_vector)
# Set attributes of the projected vector
projected_vec.attrs = input_vector.attrs
return projected_vec
[docs]
def radial_unit_vector_points_outwards(self, x0, y0, eps=1e-2, **kwargs):
"""Check if the radial unit vector is pointing outwards."""
x0 = self.convert_length_to_normalized(x0)
y0 = self.convert_length_to_normalized(y0)
def erad_sign_x(r, z, **kwargs):
"""Return the sign of the non-reversed x-component of erad."""
b_y = self.magfield_component_z(r, z, **kwargs).values
return -np.sign(b_y)
def erad_sign_y(r, z, **kwargs):
"""Return the sign of the non-reversed y-component of erad."""
b_x = self.magfield_component_r(r, z, **kwargs).values
return np.sign(b_x)
direction = np.sum(
(
erad_sign_x(x0 - eps, y0, **kwargs) < 0,
erad_sign_x(x0 + eps, y0, **kwargs) > 0,
erad_sign_y(x0, y0 + eps, **kwargs) > 0,
erad_sign_y(x0, y0 - eps, **kwargs) < 0,
)
)
# Make sure nothing weird is going on,
# either all point inwards or outwards.
# TODO: Check needs to be revisited for stellarator equilibria.
if type(self).__name__ not in {"DommaschkEquilibrium",
"FlareEquilibrium"}:
assert direction == 0 or direction == 4
return direction > 0
