Source code for torx.specializations.grillix.initialize_from_filepath_m

"""Allows initializing GRILLIX data from filepath."""
from pathlib import Path
from warnings import warn

import torx
from torx.fileio import filepath_resolver, read_fortran_namelist
from torx.equilibrium import NumericalEquilibrium, initialize_equi_from_params
from torx.normalization import Normalization
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def initialize_grillix_from_filepath(filepath: Path, load_grid_3D = False):
    """Initialize GRILLIX data from filepath."""
    filepath = Path(filepath)

    try:
        params_filepath = filepath_resolver(filepath, "params_static_data.nml")
    except FileNotFoundError:
        line1 = f"'params_static_data.nml' not found in {filepath}"
        line2 = "  -> falling back to legacy 'params.in'"
        print(f"{line1}\n{line2}")
        params_filepath = filepath_resolver(filepath, "params.in")

    params = read_fortran_namelist(params_filepath)
    params_braginskii = read_fortran_namelist(
        filepath_resolver(filepath, "params_braginskii.in")
    )
    # Combine the two parameter dictionaries
    params = {**params, **params_braginskii}

    equi = initialize_equi_from_params(filepath, params)

    try:
        norm = Normalization.initialize_from_normalization_file(
            filepath_resolver(filepath, "physical_parameters.nml")
        )
        equi.B0 = norm.B0
        R0 = norm.R0
    except FileNotFoundError:
        warn(
            "physical_parameters.nml not found. Normalization will not be useable."\
            f"Using 'axis_Btor' as B0 from {equi.filepath}."
        )
        norm = Normalization(dict())
        equi.B0 = equi.axis_Btor
        R0 = torx.Quantity(1, "m")

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

    if load_grid_3D:
        from torx.grid import Grid3D
        multigrid_file = filepath_resolver(filepath, "multigrids_plane000.nc")
        multigrid_dir = multigrid_file.parent
        grid_cano = Grid3D.from_multigrid_files(multigrid_dir)
        grid_stag = Grid3D.from_multigrid_files(multigrid_dir, staggered = True)

        return grid_cano, grid_stag, equi, params, norm
    else:
        multigrid_file = filepath_resolver(filepath, "multigrid.nc")
        multigrid = torx.grid.Multigrid2D(
            multigrid_file,
            multigrid_group = "",
            load_all_levels = False,
            R0 = R0
        )
        grid = multigrid.get_grid(1)

        return grid, equi, params, norm