Source code for torx.equilibrium.initialize_equi_from_params_m

"""Functionality to initialize equilibria from parameter files."""
from pathlib import Path
from torx.equilibrium.check_equilibrium_m import check_equilibrium
from torx.equilibrium.numerical_m import NumericalEquilibrium
from torx.equilibrium.slab_m import SlabEquilibrium
from torx.equilibrium.circular_m import CircularEquilibrium
from torx.equilibrium.circular_toroidal_m import CircularToroidalEquilibrium
from torx.equilibrium.salpha_m import SalphaEquilibrium
from torx.equilibrium.dommaschk_m import DommaschkEquilibrium
from torx.equilibrium.flare_m import FlareEquilibrium
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def initialize_equi_from_params(filepath: Path, params: dict):
    """Initialize the equilibrium from a given filepath and parameter file."""
    equi_type = check_equilibrium(params)

    if equi_type == "SLAB":
        return SlabEquilibrium(**params["equi_slab_params"])

    elif equi_type == "CIRCULAR":
        return CircularEquilibrium(**params["equi_circular_params"])

    elif equi_type == "CIRCULAR_TOROIDAL":
        return CircularToroidalEquilibrium(**params["equi_circular_toroidal_params"])

    elif equi_type == "SALPHA":
        return SalphaEquilibrium(**params["equi_salpha_params"])

    elif equi_type == "DOMMASCHK":
        return DommaschkEquilibrium(
                    filepath=filepath,
                    **params["params_equi_dommaschk"],
                    **params["params_equi_dommaschk_fitting_coef"]
                )

    elif equi_type == "NUMERICAL":
        return NumericalEquilibrium.initialize_from_params(filepath, params)

    elif equi_type =="FLARE":
        flare_params = {}
        for k in params.keys():
            if "flare" in k:
                flare_params = flare_params | params[k]

        flare_params = flare_params | params["params_mesh"]

        return FlareEquilibrium(filepath=filepath, **flare_params)

    else:
        raise NotImplementedError(f"No quick initialization available for {equi_type}")