"""Calculate particle and heat drift fluxes across a surface defined by a LineoutSet."""
import numpy as np
import xarray as xr
from torx.arrays import make_xarray
from torx.grid import Grid3D
from torx.decorators import autodoc_function
from torx.units import Normalization
from torx.equilibrium import EquilibriumBaseClass
from torx.analysis.lineouts import LineoutSet
from torx.analysis import drift_flux
[docs]
@autodoc_function
def rad_particleflux_ExB(
grid: Grid3D,
equi: EquilibriumBaseClass,
lineout_set: LineoutSet,
norm: Normalization,
density: xr.DataArray,
pot: xr.DataArray,
n_samples: int=500,
total: bool=True,
):
"""Calculate the electrostatic ExB particle drift flux."""
return drift_flux(grid, equi, lineout_set, norm, pot,
coeff=density, n_samples=n_samples, total=total)
[docs]
@autodoc_function
def rad_particleflux_dia(
grid: Grid3D,
equi: EquilibriumBaseClass,
lineout_set: LineoutSet,
norm: Normalization,
density: xr.DataArray,
temp: xr.DataArray,
charge: float=1.,
n_samples: int=500,
total: bool=True,
):
"""Calculate the diamagnetic particle drift flux."""
coeff_n = make_xarray(
density,
norm=density.norm/norm.elementary_charge
)
coeff_T = make_xarray(
temp,
norm=temp.norm/norm.elementary_charge
)
prefac = -1/charge
return prefac * (
drift_flux(
grid, equi, lineout_set, norm, temp,
coeff=coeff_n, n_samples=n_samples, total=total,
)
+ drift_flux(
grid, equi, lineout_set, norm, density,
coeff=coeff_T, n_samples=n_samples, total=total,
)
)
[docs]
@autodoc_function
def rad_particleflux_em_e(
grid: Grid3D,
equi: EquilibriumBaseClass,
lineout_set: LineoutSet,
norm: Normalization,
density: xr.DataArray,
apar_fluct: xr.DataArray,
upar: xr.DataArray,
jpar: xr.DataArray,
n_samples: int=500,
total: bool=True,
):
"""Calculate the electromagnetic electron drift flux."""
R = grid.r_u
Z = grid.z_u
phi = grid.coords["phi"]
btor = xr.apply_ufunc(
equi.magfield_component_toroidal,
R,
Z,
phi,
input_core_dims=[["points"], ["points"], []],
output_core_dims=[["points"]],
vectorize=True,
dask='parallelized',
output_dtypes=[np.float64]
)
apar_fluct_over_btor = make_xarray(
apar_fluct/btor,
norm=apar_fluct.norm/norm.B0
)
nvpar_btor = make_xarray(
(jpar - density*upar)*btor,
norm=density.norm*upar.norm*norm.B0
)
return drift_flux(grid, equi, lineout_set, norm, apar_fluct_over_btor,
coeff=nvpar_btor, n_samples=n_samples, total=total)
[docs]
@autodoc_function
def rad_particleflux_em_i(
grid: Grid3D,
equi: EquilibriumBaseClass,
lineout_set: LineoutSet,
norm: Normalization,
density: xr.DataArray,
apar_fluct: xr.DataArray,
upar: xr.DataArray,
n_samples: int=500,
total: bool=True,
):
"""Calculate the electromagnetic ion drift flux."""
R = grid.r_u
Z = grid.z_u
phi = grid.coords["phi"]
btor = xr.apply_ufunc(
equi.magfield_component_toroidal,
R,
Z,
phi,
input_core_dims=[["points"], ["points"], []],
output_core_dims=[["points"]],
vectorize=True,
dask='parallelized',
output_dtypes=[np.float64]
)
apar_fluct_over_btor = make_xarray(
apar_fluct/btor,
norm=apar_fluct.norm/norm.B0
)
nupar_btor = make_xarray(
density*upar*btor,
norm=density.norm*upar.norm*norm.B0
)
return drift_flux(grid, equi, lineout_set, norm, apar_fluct_over_btor,
coeff=nupar_btor, n_samples=n_samples, total=total)
[docs]
@autodoc_function
def rad_heatflux_ExB(
grid: Grid3D,
equi: EquilibriumBaseClass,
lineout_set: LineoutSet,
norm: Normalization,
density: xr.DataArray,
pot: xr.DataArray,
temp: xr.DataArray,
n_samples: int=500,
total: bool=True,
):
"""Calculate the electrostatic ExB electron heat drift flux."""
coeff = make_xarray(
1.5*density*temp,
norm=density.norm*temp.norm
)
return drift_flux(grid, equi, lineout_set, norm, pot,
coeff=coeff, n_samples=n_samples, total=total)
[docs]
@autodoc_function
def rad_heatflux_dia(
grid: Grid3D,
equi: EquilibriumBaseClass,
lineout_set: LineoutSet,
norm: Normalization,
density: xr.DataArray,
temp: xr.DataArray,
charge: float=1.,
n_samples: int=500,
total: bool=True,
):
"""Calculate the diamagnetic electron heat drift flux."""
coeff_n = make_xarray(
2.5*temp*density,
norm=density.norm*temp.norm/norm.elementary_charge
)
coeff_T = make_xarray(
2.5*temp*temp,
norm=temp.norm*temp.norm/norm.elementary_charge
)
prefac = -1/charge
return prefac * (
drift_flux(
grid, equi, lineout_set, norm, temp,
coeff=coeff_n, n_samples=n_samples, total=total
) \
+ drift_flux(
grid, equi, lineout_set, norm, density,
coeff=coeff_T, n_samples=n_samples, total=total
)
)
[docs]
@autodoc_function
def rad_heatflux_em_e(
grid: Grid3D,
equi: EquilibriumBaseClass,
lineout_set: LineoutSet,
norm: Normalization,
density: xr.DataArray,
electron_temp: xr.DataArray,
apar_fluct: xr.DataArray,
upar: xr.DataArray,
jpar: xr.DataArray,
n_samples: int=500,
total: bool=True,
):
"""
Calculate the electromagnetic electron heat drift flux.
Note: This does not take into account parallel heat conduction (q_par).
"""
R = grid.r_u
Z = grid.z_u
phi = grid.coords["phi"]
btor = xr.apply_ufunc(
equi.magfield_component_toroidal,
R,
Z,
phi,
input_core_dims=[["points"], ["points"], []],
output_core_dims=[["points"]],
vectorize=True,
dask='parallelized',
output_dtypes=[np.float64]
)
apar_fluct_over_btor = make_xarray(
apar_fluct/btor,
norm = apar_fluct.norm/norm.B0
)
nvpar_Te_btor = make_xarray(
2.5*electron_temp*(jpar - density*upar)*btor,
norm = electron_temp.norm*density.norm*upar.norm*norm.B0
)
return drift_flux(grid, equi, lineout_set, norm, apar_fluct_over_btor,
coeff=nvpar_Te_btor, n_samples=n_samples, total=total)
[docs]
@autodoc_function
def rad_heatflux_em_i(
grid: Grid3D,
equi: EquilibriumBaseClass,
lineout_set: LineoutSet,
norm: Normalization,
density: xr.DataArray,
ion_temp: xr.DataArray,
apar_fluct: xr.DataArray,
upar: xr.DataArray,
n_samples: int=500,
total: bool=True,
):
"""
Calculate the electromagnetic ion heat drift flux.
Note: This does not take into account parallel heat conduction (q_par).
"""
R = grid.r_u
Z = grid.z_u
phi = grid.coords["phi"]
btor = xr.apply_ufunc(
equi.magfield_component_toroidal,
R,
Z,
phi,
input_core_dims=[["points"], ["points"], []],
output_core_dims=[["points"]],
vectorize=True,
dask='parallelized',
output_dtypes=[np.float64]
)
apar_fluct_over_btor = make_xarray(
apar_fluct/btor,
norm = apar_fluct.norm/norm.B0
)
nupar_Ti_btor = make_xarray(
2.5*ion_temp*density*upar*btor,
norm = ion_temp.norm*density.norm*upar.norm*norm.B0
)
return drift_flux(grid, equi, lineout_set, norm, apar_fluct_over_btor,
coeff=nupar_Ti_btor, n_samples=n_samples, total=total)
