Source code for torx.specializations.genex.initialize_from_filepath_m

"""Allows initializing all GENE-X data from filepath."""
import xarray as xr
from pathlib import Path
from torx.equilibrium import (
    initialize_equi_from_params, 
    NumericalEquilibrium,
    EquiType
)
from torx.grid import Grid2D, Grid3D

from torx.fileio import filepath_resolver, read_fortran_namelist
from torx.specializations.genex.grid_helpers_m import build_grid_from_filepath_genex
from torx.specializations.genex.normalization_helpers_m import (
    create_physical_parameters_dict_from_filepath_genex,
)
from torx.normalization.normalization_m import Normalization
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def initialize_genex_from_filepath(filepath: Path, load_grid_3D = False):
    """Initialize GENE-X data from filepath."""
    filepath = Path(filepath)

    #TODO: This is a temporary fix only.
    # Need to add possibility to write parameters to a given file in Parallax
    params_out = read_fortran_namelist(filepath_resolver(filepath, "params_out.txt"))
    equi_type = int(params_out["params_mesh"]["equilibrium_type"])
    if EquiType(equi_type) == EquiType.DOMMASCHK:
        params = read_fortran_namelist(filepath_resolver(filepath, "params_in.txt"))
        params_out["params_equi_dommaschk"] = params["params_equi_dommaschk"]
        params = params_out
    elif EquiType(equi_type) == EquiType.NUMERICAL:
        params = read_fortran_namelist(filepath_resolver(filepath, "params_in.txt"))
        params_out["equi_numerical_params"] = params["equi_numerical_params"]
        params = params_out
    else:
        params = read_fortran_namelist(filepath_resolver(filepath, "params_in.txt"))

    norm = Normalization(create_physical_parameters_dict_from_filepath_genex(filepath))

    grid = build_grid_from_filepath_genex(filepath)
    if type(grid) == Grid3D:
        grid.set_R0(norm.R0)
        if not load_grid_3D:
            grid = grid.isel_phi(phi_index=0)
    elif type(grid) == Grid2D:
        grid.R0 = norm.R0
    else:
        raise ValueError("Unsupported grid type")

    equi = initialize_equi_from_params(filepath, params)
    equi.B0 = norm.B0

    if (
        isinstance(equi, NumericalEquilibrium)
        and params["equi_numerical_params"]["flip_z"]
    ):
        equi.flip_Z()

    return grid, equi, params, norm