Source code for torx.specializations.genex.initialize_from_filepath_m

"""Allows initializing all GENE-X data from filepath."""
from pathlib import Path
from numbers import Integral
from typing import Sequence, Union
from warnings import warn

from .grid_helpers_m import build_grid_from_filepath_genex
from .normalization_helpers_m import (
    create_physical_parameters_dict_from_filepath_genex,
)
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.units import Normalization
from torx.decorators import autodoc_function



@autodoc_function
def initialize_genex_from_filepath(
    filepath: Path,
    planes: Union[int, str, Sequence[int]] = 0
):
    """
    Initialize GENE-X data from filepath.

    Parameters
    ----------
    filepath : Path
        Path of the GENE-X output directory.
    planes : int, str or Sequence[int]
        Plane of a 3D run to return the 2D grid of. Give a list of indices
        for a 3D grid holding those planes, or 'all' for a 3D grid holding
        every plane. A 2D run has a single plane, so its grid is returned
        whatever is given here.

    Returns
    -------
    tuple
        The grid, the equilibrium, the parameters and the normalization.
        The grid is a Grid2D for a single plane and a Grid3D otherwise.

    Raises
    ------
    ValueError
        If planes is a string other than 'all'.
    """
    filepath = Path(filepath)

    if isinstance(planes, str):
        if planes.lower() != "all":
            raise ValueError(
                f"Got planes='{planes}', expected the index of a plane, "
                "a list of indices or 'all' for the whole 3D grid."
            )
        wanted_planes = None
    elif isinstance(planes, Integral):
        wanted_planes = int(planes)
    else:
        wanted_planes = [int(index) for index in planes]

    #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"))
    params_in = read_fortran_namelist(
        filepath_resolver(filepath, "params_in.txt"))

    equi_type = int(params_out["params_mesh"]["equilibrium_type"])
    if EquiType(equi_type) == EquiType.DOMMASCHK:
        params_out["params_equi_dommaschk"] = params_in["params_equi_dommaschk"]
        params = params_out
    elif EquiType(equi_type) == EquiType.NUMERICAL:
        params_out["equi_numerical_params"] = params_in["equi_numerical_params"]
        params = params_out
    else:
        params = params_in

    norm = Normalization(
        create_physical_parameters_dict_from_filepath_genex(filepath))

    grid = build_grid_from_filepath_genex(filepath)
    if isinstance(grid, Grid3D):
        grid.set_R0(norm.R0)
        if isinstance(wanted_planes, int):
            grid = grid.isel_phi(phi_index=wanted_planes)
        elif wanted_planes is not None:
            grid = grid.isel_planes(wanted_planes)
    elif isinstance(grid, Grid2D):
        if wanted_planes != 0 and wanted_planes is not None:
            warn(
                f"Plane {planes} was asked for, but this run has a single "
                "plane. Its 2D grid is returned."
            )
        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