Source code for torx.specializations.genex.grid_helpers_m

"""Interface GENE-X grid files."""
import xarray as xr
import numpy as np
from pathlib import Path
from torx.fileio import filepath_resolver
from torx.grid.grid_2d_m import Grid2D
from torx.grid.grid_3d_m import Grid3D
from torx.autodoc_decorators_m import autodoc_function
from typing import Union



@autodoc_function
def build_grid_from_filepath_genex(directory_path: Path) -> Union[Grid2D, Grid3D]:
    """Build a GENE-X grid from a a given directory path."""
    grid_dataset = xr.open_dataset(
        filepath_resolver(directory_path, "mesh.nc"),
        group = "RZ_grid"
    )

    return build_grid_genex(grid_dataset)




@autodoc_function
def build_grid_genex(mesh_group: xr.Dataset) -> Union[Grid2D, Grid3D]:
    """Build a GENE-X grid given a netcdf RZ_grid group out of the mesh.nc."""
    x_unstructured, y_unstructured = \
        _get_unstructured_grid_arrays_from_file(mesh_group)

    if isinstance(x_unstructured, xr.DataArray) \
        and isinstance(y_unstructured, xr.DataArray):
        x_unstructured = x_unstructured.rename({"dim_phi": "phi"})
        y_unstructured = y_unstructured.rename({"dim_phi": "phi"})

        grid = Grid3D.from_rz(x_unstructured, y_unstructured)

    else:
        grid = Grid2D(x_unstructured, y_unstructured)

    return grid


def _get_unstructured_grid_arrays_from_file(grid_group: xr.Dataset):
    """Return the unstructured x and y arrays from the grid group."""
    # Fix to support old and new mesh files where names were different
    if("x" in grid_group):
        R_str = "x"
        Z_str = "y"
    else:
        R_str = "R"
        Z_str = "Z"

    if("dim_RZ_grid_phi" in grid_group[R_str].dims):
        x_unstructured = grid_group[R_str].where(grid_group["is_ghost"] == 0, drop=True)
        y_unstructured = grid_group[Z_str].where(grid_group["is_ghost"] == 0, drop=True)

    elif("dim_phi" in grid_group[R_str].dims):
        x_unstructured = grid_group[R_str].where(grid_group["not_ghost"] == 1, drop=True)
        y_unstructured = grid_group[Z_str].where(grid_group["not_ghost"] == 1, drop=True)

    else:
        x_unstructured = np.array(grid_group[R_str])
        y_unstructured = np.array(grid_group[Z_str])

    return x_unstructured, y_unstructured

@autodoc_function
def get_mesh_variable_genex(directory_path: Path, group, variable):
    """Return the specified variable from the given group in the mesh file."""
    grid_file = filepath_resolver(directory_path, "mesh.nc")
    grid_group = xr.open_dataset(grid_file, group=group)

    var = grid_group[variable].rename({"dim_phi": "phi", "dim_RZ": "points"})
    return var



@autodoc_function
def mask_ghost_and_filler(directory_path: Path, ds: xr.DataArray):
    """
    Mask the ghost and filler points if required.

    If there are no ghost or filler points, returns the dataset unchanged.
    """
    grid_file = filepath_resolver(directory_path, "mesh.nc")
    grid_group = xr.open_dataset(grid_file, group="RZ_grid")

    if "not_filler" in grid_group.keys():
        not_filler = grid_group["not_filler"]\
            .rename({"dim_phi": "phi", "dim_RZ": "points"})
    elif "is_filler" in grid_group.keys():
        not_filler = 1 - grid_group["is_filler"]\
            .rename({"dim_phi": "phi", "dim_RZ": "points"})
    else:
        not_filler = None

    if not_filler is not None:
        ds = ds.where(not_filler, drop=True)

    if "not_ghost" in grid_group.keys():
        not_ghost = grid_group["not_ghost"]\
            .rename({"dim_phi": "phi", "dim_RZ": "points"})
    elif "is_ghost" in grid_group.keys():
        not_ghost = 1 - grid_group["is_ghost"]\
            .rename({"dim_phi": "phi", "dim_RZ": "points"})
    else:
        not_ghost = None

    if not_ghost is not None:
        ds = ds.where(not_ghost, drop=True)

    return ds




@autodoc_function
def get_phi_grid_genex(directory_path: Path):
    """Return the GENE-X phi grid."""
    grid_file = filepath_resolver(directory_path, "mesh.nc")
    grid_group = xr.open_dataset(grid_file, group="phi_grid")

    return grid_group["phi"].rename({"dim_phi": "phi"})




@autodoc_function
def get_velocity_space_grid_genex(directory_path: Path):
    """Return the GENE-X velocity space grid."""
    grid_file = filepath_resolver(directory_path, "mesh.nc")

    vp_dataset = xr.open_dataset(grid_file, group="vp_grid")
    mu_dataset = xr.open_dataset(grid_file, group="mu_grid")

    return vp_dataset["vp"], mu_dataset["mu"]




@autodoc_function
def get_magnetic_field_genex(directory_path: Path):
    """Return the GENE-X magnetic field data."""
    try:
        # Backwards compatible version
        mag_file = filepath_resolver(directory_path, "magnetic_field.nc")
        dataset = xr.open_dataset(mag_file, chunks="auto")
        dataset = dataset.rename({"dim_RZ": "points"})
    except:
        mag_file = filepath_resolver(directory_path, "mesh.nc")
        dataset = xr.open_dataset(mag_file, group="magnetic_field", chunks="auto")
        dataset = dataset.rename({"dim_phi":"phi", "dim_RZ": "points"})
        dataset = mask_ghost_and_filler(directory_path, dataset)

    return dataset