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

"""Contains functionality to interface map."""
import xarray as xr
import numpy as np
from pathlib import Path
from inspect import getmodule
from torx.specializations.grillix.operators.parallel_operators_m import (
    parallel_gradient,
    parallel_divergence,
    map_to_staggered,
    map_to_canonical,
)
from torx.autodoc_decorators_m import autodoc_class



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



    @classmethod
    def from_filepath(cls, filepath: Path):
        """Initialize from top-level filepath."""
        return cls(filepath / "trunk" / "map.nc")




    def __init__(self, map_file: Path):
        """Initialize the map."""
        self._map_file = map_file
        self._dataset = xr.open_dataset(self._map_file).rename(n_points_cano="points")


    @property
    def nplanes(self) -> int:
        """Number of poloidal planes."""
        return self._dataset.nplanes

    @property
    def dphi(self) -> float:
        """Angular distance between planes."""
        return self._dataset.dphi

    @property
    def intorder(self) -> int:
        """Interpolation order of the map matrices."""
        return self._dataset.intorder

    @property
    def xorder(self) -> int:
        """Subcell splitting for flux-box integration 2^(xorder)."""
        return self._dataset.xorder

    @property
    def dpar_fwd_half(self) -> xr.DataArray:
        """Parallel distance to the half forward toroidal plane."""
        return self._dataset["dpar_cano_stag_fwd"]

    @property
    def dpar_bwd_half(self) -> xr.DataArray:
        """Parallel distance to the half reverse toroidal plane."""
        return self._dataset["dpar_cano_stag_bwd"]

    @property
    def dpar_fwd_full(self) -> xr.DataArray:
        """Parallel distance to the forward toroidal plane."""
        return self._dataset["dpar_cano_cano_fwd"]

    @property
    def dpar_bwd_full(self) -> xr.DataArray:
        """Parallel distance to the reverse toroidal plane."""
        return self._dataset["dpar_cano_cano_bwd"]

    @property
    def fluxbox_vol(self) -> xr.DataArray:
        """Volume of the flux box around the point (+/- 1/2 phi)."""
        return self._dataset["fluxbox_vol_cano"]

    @property
    def dpar_fwd_half_stag(self) -> xr.DataArray:
        """Parallel distance to the half forward staggered toroidal plane."""
        return self._get_stag_var_from_ds("dpar_stag_cano_fwd")

    @property
    def dpar_bwd_half_stag(self) -> xr.DataArray:
        """Parallel distance to the half reverse staggered toroidal plane."""
        return self._get_stag_var_from_ds("dpar_stag_cano_bwd")

    @property
    def dpar_fwd_full_stag(self) -> xr.DataArray:
        """Parallel distance to the forward staggered toroidal plane."""
        return self._get_stag_var_from_ds("dpar_stag_stag_fwd")

    @property
    def dpar_bwd_full_stag(self) -> xr.DataArray:
        """Parallel distance to the reverse staggered toroidal plane."""
        return self._get_stag_var_from_ds("dpar_stag_stag_bwd")

    @property
    def fluxbox_vol_stag(self) -> xr.DataArray:
        """Volume of the flux box around the staggered point (+/- 1/2 phi)."""
        return self._get_stag_var_from_ds("fluxbox_vol_stag")

    @property
    def mfwd(self) -> xr.Dataset:
        """Map matrix in forward(phi) direction (trace of +dphi/2)."""
        return xr.open_dataset(self._map_file, group="map_stag_cano_fwd")

    @property
    def mbwd(self) -> xr.Dataset:
        """Map matrix in backward(phi) direction (trace of -dphi/2)."""
        return xr.open_dataset(self._map_file, group="map_stag_cano_bwd")

    @property
    def mfwd_stag(self) -> xr.Dataset:
        """Map matrix in forward(phi) direction (trace of +dphi/2)."""
        return xr.open_dataset(self._map_file, group="map_cano_stag_fwd")

    @property
    def mbwd_stag(self) -> xr.Dataset:
        """Map matrix in backward(phi) direction (trace of -dphi/2)."""
        return xr.open_dataset(self._map_file, group="map_cano_stag_bwd")

    @property
    def qfwd(self) -> xr.Dataset:
        """Parallel gradient matrix in forward direction (trace of +dphi/2)."""
        return xr.open_dataset(self._map_file, group="grad_stag_cano_fwd")

    @property
    def qbwd(self) -> xr.Dataset:
        """Parallel gradient matrix in backward direction (trace of -dphi/2)."""
        return xr.open_dataset(self._map_file, group="grad_stag_cano_bwd")

    @property
    def pfwd(self) -> xr.Dataset:
        """
        Parallel divergence matrix in forward(phi/2) direction.

        pfwd = -V^(-1)*qbwd^T*V.
        """
        return xr.open_dataset(self._map_file, group="pdiv_cano_stag_fwd")

    @property
    def pbwd(self) -> xr.Dataset:
        """
        Parallel divergence matrix in backward(-phi/2) direction.

        pbwd = -V^(-1)*qfwd^T*V.
        """
        return xr.open_dataset(self._map_file, group="pdiv_cano_stag_bwd")

    @property
    def mfwd_full(self) -> xr.Dataset:
        """
        Map matrix in forward(phi) direction (trace of dphi).

        Needed for penalization.
        """
        return xr.open_dataset(self._map_file, group="map_cano_cano_fwd")

    @property
    def mbwd_full(self) -> xr.Dataset:
        """
        Map matrix in backward(phi) direction (trace of -dphi).

        Needed for penalization.
        """
        return xr.open_dataset(self._map_file, group="map_cano_cano_bwd")

    @property
    def mfwd_full_stag(self) -> xr.Dataset:
        """
        Map matrix in forward(phi) direction (trace of dphi).

        Needed for penalization.
        """
        return xr.open_dataset(self._map_file, group="map_stag_stag_fwd")

    @property
    def mbwd_full_stag(self) -> xr.Dataset:
        """
        Map matrix in backward(phi) direction (trace of -dphi).

        Needed for penalization.
        """
        return xr.open_dataset(self._map_file, group="map_stag_stag_bwd")



    def parallel_gradient(self, input_array: xr.DataArray) -> xr.DataArray:
        """Parallel gradient."""
        return parallel_gradient(self, input_array)




    def parallel_divergence(self, input_array: xr.DataArray) -> xr.DataArray:
        """Parallel divergence."""
        return parallel_divergence(self, input_array)




    def map_to_staggered(self, input_array: xr.DataArray) -> xr.DataArray:
        """Map to staggered plane."""
        return map_to_staggered(self, input_array)




    def map_to_canonical(self, input_array: xr.DataArray) -> xr.DataArray:
        """Map to canonical plane."""
        return map_to_canonical(self, input_array)


    def _get_stag_var_from_ds(self, staggered_var_name) -> xr.DataArray:
        """
        Check if the equilibrium is axisymmetric.

        Gives an error if it is (for 2D canonical and staggered grids are the same),
        otherwise tries to get and return staggered_var_name from the dataset
        """
        # 1. Check for the condition that invalidates the staggered grid calculation
        # If 'axisymmetric' is 1, the user shouldn't be calling this property.
        is_axisymmetric = self._dataset.attrs.get("axisymmetric")

        if is_axisymmetric:
            raise ValueError(
                f"Cannot retrieve staggered grid data ({staggered_var_name}) "
                f"because the dataset is **Axisymmetric** (2D). "
                f"Use the canonical property (e.g., 'dpar_fwd_half') instead."
            )

        # 2. If the dataset is 3D attempt to retrieve the staggered variable.
        # Use a try...except block to catch the case where the variable is missing.
        try:
            return self._dataset[staggered_var_name]
        except KeyError:
            raise KeyError(
                f"Dataset is non-axisymmetric (3D), but the required staggered "
                f"variable '{staggered_var_name}' was not found in the dataset."
            )