"""Contains slab equilibrium."""
import numpy as np
import xarray as xr
from torx import Quantity
from torx.units_handling_m import check_units
from torx.equilibrium.equilibrium_m import EquilibriumBaseClass
from torx.geometry.polygon_2d_m import Polygon2D
from torx.autogrid_decorators_m import autogrid_method
from torx.autodoc_decorators_m import autodoc_class
[docs]
@autodoc_class
class SlabEquilibrium(EquilibriumBaseClass):
"""Represents a Slab equilibrium as defined in PARALLAX."""
[docs]
def __init__(
self,
boxsize: float=64.0,
sol: bool=True,
yperiodic: bool=True,
B0: Quantity=Quantity("1T"),
L_ref: float=1.0,
):
"""
Initialize the SlabEquilibrium.
Arguments to the initializer match the parameters supplied in PARALLAX.
You can either manually specify the arguments,
or pass a **kwargs_dictionary generated from a parameter file.
The L_ref parameter sets the length normalization in meters.
"""
self.B0 = B0
self._norm_length = Quantity(L_ref, "m")
self.x0 = 0.0
self.y0 = 0.0
self.xmin = 0.0
self.xmax = boxsize
self.ymin = 0.0
self.ymax = boxsize
self.rhomin = self.xmin
self.rhomax = self.xmax
self.sol = sol
self.yperiodic = yperiodic
@property
def axis_r_norm(self):
"""Magnetic axis radial position in normalized units."""
return xr.DataArray(self.x0, attrs={"norm":self.R0})
@property
def axis_z_norm(self):
"""Magnetic axis vertical position in normalized units."""
return xr.DataArray(self.y0, attrs={"norm":self.R0})
@property
def B0(self) -> Quantity:
"""Magnetic field normalization."""
return self._norm_magfield
@B0.setter
def B0(self, value: Quantity):
"""Set the magnetic field normalization."""
check_units(value, {"[mass]":1, "[time]":-2, "[current]":-1}, "B0")
self._norm_magfield = value
@property
def R0(self) -> Quantity:
"""Length normalization."""
return self._norm_length
@R0.setter
def R0(self, value: Quantity):
"""Set the length normalization."""
check_units(value, {"[length]":1}, "R0")
self._norm_length = value
[docs]
@autogrid_method
def normalized_flux_surface_label(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""Return the normalized flux surface label (rho_pol)."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
if kwargs["is_structured"]:
r_mesh, z_mesh = np.meshgrid(r_norm, z_norm)
else:
r_mesh = r_norm
z_mesh = z_norm
return xr.DataArray(r_mesh,
attrs={"norm":Quantity(1.0)},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_r(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return the radial magnetic field component."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
rho = self.normalized_flux_surface_label(r_norm, z_norm)
return xr.DataArray(np.zeros_like(rho),
attrs={"norm":self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_z(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return the vertical magnetic field component."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
rho = self.normalized_flux_surface_label(r_norm, z_norm)
return xr.DataArray(np.zeros_like(rho),
attrs={"norm":self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def magfield_component_toroidal(self, r_norm, z_norm, **kwargs) \
-> xr.DataArray:
"""Return the toroidal magnetic field component."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
rho = self.normalized_flux_surface_label(r_norm, z_norm)
return xr.DataArray(np.ones_like(rho),
attrs={"norm":self.B0},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
@autogrid_method
def jacobian(self, r_norm, z_norm, **kwargs) -> xr.DataArray:
"""Return the Jacobian of the circular geometry."""
r_norm = self.convert_length_to_normalized(r_norm)
z_norm = self.convert_length_to_normalized(z_norm)
rho = self.normalized_flux_surface_label(r_norm, z_norm)
return xr.DataArray(np.ones_like(rho),
attrs={"norm":Quantity(1.0)},
dims=kwargs["dims"],
coords=kwargs["coords"])
[docs]
def get_boundary_polygon(self):
"""Return a rectangle representing the boundary of the slab."""
# Define the corners of the slab
x_points = np.array([self.xmin, self.xmax, self.xmax, self.xmin])
y_points = np.array([self.ymin, self.ymin, self.ymax, self.ymax])
return Polygon2D(x_points=x_points, y_points=y_points)
