Source code for torx.measure.temperature_m

"""
Functions used to calculate temperatures from densities, flows and energies.

On the definition of temperature:
We define Tlab = 2/3 * Wlab / n - 1/3 * m * u**2, where Wlab is the lab frame
energy, n the density, m the mass and u the mean flow. Since temperatures in
a kinetic code can be distinguished between parallel and perpendicular to the
magnetic field, and the perpendicular direction has one more degree of freedom,
typically Ttot = 1/3 * Tpar + 2/3 * Tperp. Combining both, we have
Tpar = 2 * Wpar / n - m * u**2 and Tperp = Wperp / n. These formulas are a
matter of choice and make most sense with this definition.
"""
import xarray as xr
import numpy as np
from torx import make_xarray
from torx.normalization.normalization_m import Normalization
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def parallel_temperature(
    params,
    norm: Normalization,
    n: xr.DataArray,
    E_par: xr.DataArray,
    u_par: xr.DataArray,
    spec: str,
):
    """Calculate the parallel (lab frame) temperature."""
    assert n.shape == E_par.shape
    assert n.shape == u_par.shape

    v_par = u_par / n

    # Convert v_par back to normalized units to subtract it from the parallel energy
    idx = params["params_species"]["names"].index(spec)
    v_par = v_par / np.sqrt(2.0 / params["params_species"]["masses"][idx])

    return make_xarray(
            2.0 * E_par / n - v_par**2,
        norm=norm.Ti0.to("eV"),
        name="Parallel temperature",
    )




@autodoc_function
def perpendicular_temperature(
    params, norm: Normalization, n: xr.DataArray, E_perp: xr.DataArray
):
    """Calculate the perpendicular temperature."""
    assert n.shape == E_perp.shape

    return make_xarray(
        E_perp / n,
        norm=norm.Ti0.to("eV"),
        name="Perpendicular temperature",
    )




@autodoc_function
def total_temperature(
    params,
    norm: Normalization,
    n: xr.DataArray,
    E_par: xr.DataArray,
    E_perp: xr.DataArray,
    u_par: xr.DataArray,
    spec,
):
    """Calculate the total (lab frame) temperature."""
    assert n.shape == E_par.shape
    assert n.shape == E_perp.shape
    assert n.shape == u_par.shape

    v_par = u_par / n

    # Convert v_par back to normalized units to subtract it from the parallel energy
    idx = params["params_species"]["names"].index(spec)
    v_par = v_par / np.sqrt(2.0 / params["params_species"]["masses"][idx])

    return make_xarray(
        1.0/3.0 * parallel_temperature(params, norm, n, E_par, u_par, spec)
        + 2.0/3.0 * perpendicular_temperature(params, norm, n, E_perp),
        norm=norm.Ti0.to("eV"),
        name="Total temperature",
    )