Source code for torx.specializations.genex.snaps_access_neut_m

"""Access data stored in GENE-X neutrals diagnostic files."""
import xarray as xr
import numpy as np
import os
import re
from pathlib import Path
from glob import glob

from torx.fileio import filepath_resolver
from .grid_helpers_m import mask_ghost_and_filler
from torx.autodoc_decorators_m import autodoc_function
from torx.performance import auto_chunk
from torx.specializations.genex.snaps_access_m import has_parts, find_all

def load_neut_group(path: Path, group: str):
    """Load the given group from the neutrals.nc file."""
    # NOTE: Time is only present in the file if no group is specified, thus
    #       fetch it first before data load
    dataset = xr.open_dataset(path, chunks="auto")
    time = dataset["time"].values

    dataset = xr.open_dataset(path, group=group, chunks="auto")
    dataset = dataset.rename(
        {"dim_time": "tau", "dim_phi": "phi", "dim_RZ": "points"}
    )

    dataset = dataset.assign_coords({"tau": time})
    dataset = dataset.assign_coords({"phi": dataset["phi"]})
    dataset.attrs["name"] = "n"
    return dataset

def load_from_neut_file(directory_path: Path, species: str):
    """Load a given variable from the 'neutrals.nc' file."""
    if not has_parts(directory_path):
        path = filepath_resolver(directory_path, "neutrals.nc")
        if not path.exists():
            raise FileNotFoundError("There is no neutrals.nc in the "\
                                    "given directory!")

        dataset = load_neut_group(path, species)
    else:
        paths = find_all(directory_path, "neutrals.nc")
        # Sort paths naturally by the appending number
        paths = sorted(paths, key=lambda x: int(x.parent.stem.split("_")[-1]))
        dataset = xr.combine_by_coords(
            [load_neut_group(path, species) for path in paths]
        )

    return dataset["n"]



@autodoc_function
def load_snaps_neut_genex(
    directory_path: Path, species: str, mask_ghost: bool=True,
    **loader_kwargs
):
    """
    Load a single 2d variable from the GENE-X neutrals diagnostic files.

    Uses the directory path of the GENE-X output directory, the neutrals
    species name and the variable name.
    """
    ds = load_from_neut_file(directory_path, species)

    ds = mask_ghost_and_filler(directory_path, ds)

    return auto_chunk(ds)