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

"""Contains functionality to interface equi on mesh."""
import xarray as xr
import numpy as np
from pathlib import Path
from inspect import getmodule
import torx
from torx.autodoc_decorators_m import autodoc_class



@autodoc_class
class EquiOnMesh:
    """Low-level interface to the grillix/trunk/equi_on_mesh.nc file."""



    @classmethod
    def from_filepath(cls, filepath: Path):
        """Initialize from top-level filepath."""
        target_file = filepath / "trunk" / "equi_on_mesh.nc"
        if target_file.is_file():
            return cls(target_file)
        else:
            return cls(filepath / "trunk" / "equi_storage.nc")




    def __init__(self, equi_on_mesh_file: Path):
        """Initialize the equi on mesh."""
        self._dataset = xr.open_dataset(equi_on_mesh_file)


    @property
    def rho(self) -> xr.DataArray:
        """Flux surface label on grid points."""
        return self._dataset["rho"].rename(npoints="points")

    @property
    def bx(self) -> xr.DataArray:
        """Magfield bx on grid points."""
        return self._dataset["bx"].rename(npoints="points")

    @property
    def by(self) -> xr.DataArray:
        """Magfield by on grid points."""
        return self._dataset["by"].rename(npoints="points")

    @property
    def btor(self) -> xr.DataArray:
        """Magfield btor on grid points."""
        return self._dataset["btor"].rename(npoints="points")

    @property
    def epol(self) -> xr.DataArray:
        """Magfield unit vector epol on grid points."""
        epol = torx.vector.poloidal_vector(
            input_r=self._dataset["epol"].isel(two=0).rename(npoints="points"),
            input_z=self._dataset["epol"].isel(two=1).rename(npoints="points"),
            dims=["points"],
        )
        return epol

    @property
    def erad(self) -> xr.DataArray:
        """Magfield unit vector erad on grid points."""
        erad = torx.vector.poloidal_vector(
            input_r=self._dataset["erad"].isel(two=0).rename(npoints="points"),
            input_z=self._dataset["erad"].isel(two=1).rename(npoints="points"),
            dims=["points"],
        )
        return erad

    @property
    def district(self) -> xr.DataArray:
        """District on grid points."""
        return self._dataset["district"].rename(npoints="points")