Source code for torx.measure.pressure_m

"""Pressure components."""
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 electron_pressure(
    density: xr.DataArray, electron_temp: xr.DataArray, norm: Normalization
):
    """Return the electron thermal pressure."""
    assert density.shape == electron_temp.shape

    return make_xarray(
        density * electron_temp,
        norm=(norm.n0 * norm.Te0).to("Pa"),
        name="Electron thermal pressure",
    )




@autodoc_function
def ion_pressure(density: xr.DataArray, ion_temp: xr.DataArray, norm: Normalization):
    """Return the ion thermal pressure."""
    assert density.shape == ion_temp.shape

    return make_xarray(
        density * ion_temp,
        norm=(norm.n0 * norm.Ti0).to("Pa"),
        name="Ion thermal pressure",
    )




@autodoc_function
def total_pressure(
    density: xr.DataArray,
    electron_temp: xr.DataArray,
    ion_temp: xr.DataArray,
    norm: Normalization,
):
    """Return the total thermal pressure."""
    return (
        electron_pressure(density, electron_temp, norm)
        + ion_pressure(density, ion_temp, norm) / norm.zeta
    )