"""Functions related to velocities."""
import xarray as xr
import numpy as np
from torx.grid import Grid2D
from torx.vector import vector_cross, toroidal_vector, eZ_unit_vector
from torx.autogrid_decorators_m import autocoord_function
from torx.equilibrium.equilibrium_m import EquilibriumBaseClass
from torx.normalization.normalization_m import Normalization
from torx import make_xarray
from torx.autodoc_decorators_m import autodoc_function
[docs]
@autodoc_function
def parallel_ion_velocity_vector(
grid: Grid2D,
equi: EquilibriumBaseClass,
shaped_velocity: xr.DataArray,
**kwargs
) -> xr.DataArray:
"""Convert the parallel ion velocity ion a vector array."""
assert (
"R" in shaped_velocity.dims and "Z" in shaped_velocity.dims
), f"Call vector_to_matrix on velocity before calling \
parallel_ion_velocity_vector"
e_par = equi.magfield_vector(grid.r_s, grid.z_s, normalize=True, **kwargs)
return make_xarray(
e_par * shaped_velocity,
norm=shaped_velocity.norm,
)
[docs]
@autodoc_function
def sound_speed(
electron_temp: xr.DataArray,
ion_temp: xr.DataArray,
norm: Normalization,
adiabatic_index: float = 1.0,
) -> xr.DataArray:
"""
Return the ion sound speed for longitudinal waves.
Use adiabatic_index = 0.0 to exclude ion_temperature from the calculation.
Normally use adiabatic_index = 1.0, but it is possible that 3.0 could also
be valid since some kinetic modeling suggests ions have an additional
degree of freedom.
"""
return make_xarray(
np.sqrt(
electron_temp + adiabatic_index * norm.Z * ion_temp \
/ norm.zeta.magnitude
),
norm=norm.c_s0,
name="Ion sound speed",
)
[docs]
@autodoc_function
@autocoord_function
def ExB_velocity(
electric_field: xr.DataArray,
r_norm,
z_norm,
equi: EquilibriumBaseClass,
norm: Normalization,
component: str=None,
**kwargs
) -> xr.DataArray:
"""
Return the E cross B velocity as a vector array.
Assumes that the magnetic field can be approximated as purely toroidal.
It is possible to calculate the ExB_velocity with the full magnetic field,
but the result is no longer within a poloidal plane.
Normalization:
[v_ExB] = [B]/[B]**2 * [grad_perp] * [phi]
= 1 / B0 * 1/R0 * Te0/e
= Te0/(e * rho_s0 * B0) * rho_s0/R0
We can use rho_s0 = sqrt(Te0 * Mi)/(e * B0), so
[v_ExB] = sqrt(Te0 / Mi) / delta
= c_s0 / delta
Note:
----
In order to get rid of the delta the normalization is multiplied by delta
and the values are divided by delta and we have v_ExB.norm = c_s0.
"""
B_tor = equi.magfield_component_toroidal(r_norm, z_norm, **kwargs)
assert("vector" in electric_field.dims)
v_ExB = vector_cross(electric_field, toroidal_vector(B_tor)) \
/ B_tor**2 / norm.delta.m
v_ExB.attrs["norm"] = electric_field.norm / B_tor.norm * norm.delta
if (component is not None):
v_ExB = equi.project_vector_to_magfield(v_ExB,
r_norm=r_norm,
z_norm=z_norm,
direction=component,
**kwargs)
return v_ExB
[docs]
@autodoc_function
@autocoord_function
def diamagnetic_velocity(
temperature: xr.DataArray,
r_norm,
z_norm,
equi: EquilibriumBaseClass,
norm: Normalization,
spec: str,
component: str=None,
**kwargs
) -> xr.DataArray:
"""
Return the divergent part of the diamagnetic drift velocity for species.
Species can be "ions" or "electrons".
v_dia = 2 T_i / (e B**3) vec(B) cross grad(B)
Units:
[v_dia] = [T_i] / ([e] * [B**3]) [B] * [grad_perp] * [B]
= Ti0 / (e * B0**3) B0 * 1/R0 * B0
= Ti0 / (e * B0 * rho_s0) rho_s0/R0---> sec docstring of ExB-velocity
= sqrt(Ti0 / Mi) / delta
= sqrt(Te0 / Mi) * Ti0 / Te0 / delta
= c_s0 * zeta / delta
Note:
----
In order to get rid of the delta the normalization is multiplied by delta
and the values are divided by delta and we have v_dia.norm = c_s0.
Note:
----
The units are shown for ions, which is why an additional factor of
sqrt(zeta) is needed for correct normalization. For electrons it is not
needed.
"""
assert spec in ["ions", "electrons"], "Invalid species selection!"
eZ = eZ_unit_vector(r_norm, z_norm)
v_dia = 2 * temperature / norm.delta.m * eZ
v_dia = drift_normalization(v_dia, norm, spec)
v_dia.attrs["name"] = f"Diamagnetic drift velocity of {spec}"
if (component is not None):
v_dia = equi.project_vector_to_magfield(v_dia,
r_norm=r_norm,
z_norm=z_norm,
direction=component,
**kwargs)
return v_dia
@autodoc_function
def drift_normalization(
v_drift: xr.DataArray,
norm: Normalization,
spec: str,
) -> xr.DataArray:
"""
Apply correct normalization to drift velocities.
Depends on the species (can be "ions" or "electrons").
"""
# Apply correct normalization depending on the species
if spec == "ions":
v_drift *= (norm.elementary_charge.m / norm.zeta.m)
elif spec == "electrons":
v_drift *= -norm.elementary_charge.m
v_norm = v_drift.attrs.get("norm", None)
if v_norm is not None:
v_norm /= (np.abs(norm.elementary_charge) / norm.zeta)
v_drift.attrs["norm"] = v_norm
else:
v_drift.attrs["norm"] = norm.c_s0
return v_drift
