Source code for torx.measure.base_fields_m

"""Functions for adding normalization factors to the base fields."""
import xarray as xr
from torx.normalization.normalization_m import Normalization
from torx.autodoc_decorators_m import autodoc_function

def _compute_base_field_normalizations(norm: Normalization):
    """Return a dictionary of normalization values for the base fields."""
    return {
        "tau": norm.tau_0,
        "density": norm.n0,
        "electron_temp": norm.Te0,
        "ion_temp": norm.Ti0,
        "velocity": norm.c_s0,
        "current": (
            norm.c_s0 * norm.elementary_charge * norm.n0
        ).to("kiloampere*meter**-2"),
        "potential": (norm.Te0 / norm.elementary_charge).to(
            "kilovolt"
        ),
        "vorticity": (
            norm.Mi * norm.n0 * norm.Te0 / (
                norm.elementary_charge
                * norm.rho_s0**2
                * norm.B0**2
            )
        ).to("coulomb/meter**3"),
        "apar": norm.beta_0 * norm.B0 * norm.rho_s0,
        "neutrals_density": norm.n0,
        "neutrals_momentum": (norm.n0 * norm.c_s0),
        "neutrals_pressure": (norm.n0 * norm.Ti0),
    }



@autodoc_function
def convert_base_fields_to_SI(snaps: xr.Dataset, norm: Normalization):
    """Iterate over a dataset and sets the name and norm for each array."""
    base_field_normalizations = _compute_base_field_normalizations(norm)

    for key, norm in base_field_normalizations.items():

        try:
            snaps[key].attrs["name"] = key
            snaps[key].attrs["norm"] = norm
        except KeyError:
            continue

    return snaps




@autodoc_function
def convert_base_field_to_SI(
    data_array: xr.DataArray, norm: Normalization, name: str = None
):
    """
    Convert the norm and name of a single data array.

    If the data array hasn't set its 'name', you must provide it as an argument.
    """
    if name is None:
        name = data_array.name

    data_array.attrs["name"] = name
    data_array.attrs["norm"] = _compute_base_field_normalizations(norm)[name]

    return data_array