Source code for torx.measure.heat_flux_m

"""Functions related to heat fluxes."""
import xarray as xr
import numpy as np
from torx import make_xarray
from torx.grid import Grid2D
from torx.measure import electric_field
from torx.equilibrium.equilibrium_m import EquilibriumBaseClass
from torx.normalization.normalization_m import Normalization
from .velocities_m import ExB_velocity
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def electrostatic_ExB_heat_flux(
    grid: Grid2D,
    equi: EquilibriumBaseClass,
    norm: Normalization,
    es_pot: xr.DataArray,
    E_par: xr.DataArray,
    E_perp: xr.DataArray,
    **kwargs
):
    """Calculate the electrostatic cross field ExB heat flux."""
    assert es_pot.shape == E_par.shape
    assert es_pot.shape == E_perp.shape

    # Total energy
    E = E_par + E_perp
    E.attrs["norm"] = E_par.norm

    E_field = electric_field(grid, es_pot)
    v_exb = ExB_velocity(
        E_field, grid=grid, equi=equi, norm=norm, component="radial", **kwargs
    )

    return make_xarray(
        v_exb * E,
        norm=v_exb.norm * E.norm,
    )




@autodoc_function
def electromagnetic_heat_flux(
    grid: Grid2D,
    equi: EquilibriumBaseClass,
    norm: Normalization,
    A_par: xr.DataArray,
    Q_par: xr.DataArray,
    Q_perp: xr.DataArray,
    **kwargs
):
    """Calculate the electromagnetic cross field heat flux."""
    assert A_par.shape == Q_par.shape
    assert A_par.shape == Q_perp.shape

    # Total energy flux
    Q = Q_par + Q_perp
    Q.attrs["norm"] = Q_par.norm

    E_field = grid.perpendicular_gradient(A_par)
    v_exb = ExB_velocity(
        E_field, grid=grid, equi=equi, norm=norm, component="radial", **kwargs
    )

    return make_xarray(
        v_exb * Q,
        norm=v_exb.norm * Q.norm,
    )