Source code for torx.specializations.grillix.trunk.grid_m

"""Contains functionality to interface grid."""
from torx.fileio import _load_multigrid_dataset
import torx.grid.multigrid_operations_m as operators
from torx.grid.grid_2d_m import Grid2D
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def grid_2d_from_multigrid_file(filepath: str, staggered: bool = False, lvl: int=1) -> "Grid2D":
    """
    Make a Grid2D object from a filepath to a multigrid file.

    This means a grid file from GRILLIX.
    Use staggered = False (True) to get the canonical (staggered) grid.
    """
    if staggered:
        mesh_dataset = _load_multigrid_dataset(filepath,
                                               group=f"staggered/mesh_lvl_{lvl:03d}")
    else:
        mesh_dataset = _load_multigrid_dataset(filepath,
                                               group=f"canonical/mesh_lvl_{lvl:03d}")

    r_u, z_u = operators.calc_r_u(mesh_dataset), operators.calc_z_u(mesh_dataset)

    return Grid2D(r_u, z_u)