Source code for torx.equilibrium.check_equilibrium_m

"""Checks for what equilibrium type is used for a given run."""
from torx.decorators import autodoc_function, autodoc_class
from enum import Enum



@autodoc_class
class EquiType(Enum):
    """Equilibrium type integer code."""

    CARTHY            = 0
    CERFONS           = 1
    CIRCULAR          = 2
    NUMERICAL         = 3
    SLAB              = 4
    SALPHA            = 5
    DOMMASCHK         = 6
    CIRCULAR_TOROIDAL = 7
    FLARE             = 8




@autodoc_class
class FlareEquiType(Enum):
    """FLARE Equilibrium type string."""

    equi3d_mgrid = "equi3d_mgrid"
    equi3d_coilset = "equi3d_coilset"
    equi3d_hint = "equi3d_hint"


_geometry_dict = {
    "equi_slab_params": "SLAB",
    "equi_salpha_params": "SALPHA",
    "equi_circular_params": "CIRCULAR",
    "equi_carthy_params": "CARTHY",
    "equi_cerfons_params": "CERFONS",
    "equi_numerical_params": "NUMERICAL",
    "params_equi_dommaschk": "DOMMASCHK",
    "equi_circular_toroidal_params": "CIRCULAR_TOROIDAL",
    "equi_params_flare_type": "FLARE",
}



@autodoc_function
def check_equilibrium(params):
    """Return an equilibrium type based on the parameter file."""
    if "equilibrium_type" in params:
        assert "equi_type" in params["equilibrium_type"]
        return params["equilibrium_type"]["equi_type"]
    else:
        for key, value in _geometry_dict.items():
            if key in params:
                return value

    raise RuntimeError(f"Error: no valid equilibrium type found in params")