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

"""Contains functionality to interface polars."""
import xarray as xr
import numpy as np
from pathlib import Path
from inspect import getmodule
from ..operators.csr_operators_m import csr_operator
import typing
from torx.autodoc_decorators_m import autodoc_class



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



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




    def __init__(self, polars_file):
        """Initialize the polars."""
        self._dataset = xr.open_dataset(polars_file)
        self.grid = xr.open_dataset(polars_file, group="grid")
        self._cart_to_polar = xr.open_dataset(polars_file, group="cart_to_polar")
        self._surface_average = xr.open_dataset(polars_file, group="surface_average")
        self._flux_surface_average = xr.open_dataset(
            polars_file, group="flux_surface_average"
        )


    @property
    def intorder(self) -> int:
        """Polar integration order."""
        return self._dataset.intorder

    @property
    def phi(self) -> int:
        """Toroidal angle phi of this polar grid."""
        return self.grid.phi

    @property
    def rhopol_min(self) -> float:
        """Minimum rho value of polar grid."""
        return self.grid.rhopol_min

    @property
    def rhopol_max(self) -> float:
        """Maximum rho value of polar grid."""
        return self.grid.rhopol_max

    @property
    def drho(self) -> float:
        """Radial spacing of polar grid (in units of rho)."""
        return self.grid.drho

    @property
    def dtheta(self) -> float:
        """Angular spacing of polar grid."""
        return self.grid.dtheta

    @property
    def rho(self) -> xr.DataArray:
        """Rho values of polar grid."""
        return self.grid["rho"]

    @property
    def theta(self) -> xr.DataArray:
        """Theta values of polar grid."""
        return self.grid["theta"]

    @property
    def xpol(self) -> xr.DataArray:
        """R values of polar grid points."""
        return self.grid["xpol"]

    @property
    def ypol(self) -> xr.DataArray:
        """Z values of polar grid points."""
        return self.grid["ypol"]



    def cart_to_polar(
        self, input_array: xr.DataArray = None
    ) -> typing.Union[xr.Dataset, xr.DataArray]:
        """Map the input array onto the polar grid."""
        if input_array is None:
            return self._cart_to_polar

        return self.points_to_polars(csr_operator(self._cart_to_polar, input_array))




    def surface_average(
        self, input_array: xr.DataArray = None
    ) -> typing.Union[xr.Dataset, xr.DataArray]:
        """Perform a surface integral."""
        if input_array is None:
            return self._surface_average

        output_array = csr_operator(self._surface_average, input_array)
        return output_array.rename(points="rho").assign_coords(rho=self.rho.values)




    def flux_surface_average(
        self, input_array: xr.DataArray = None
    ) -> typing.Union[xr.Dataset, xr.DataArray]:
        """
        Perform a volume integral.

        For an infinitesimal volume centered on flux surface.
        This is useful is you want to look at fluxes through a surface
        """
        if input_array is None:
            return self._flux_surface_average

        output_array = csr_operator(self._flux_surface_average, input_array)
        return output_array.rename(points="rho").assign_coords(rho=self.rho.values)




    def points_to_polars(self, input_array) -> xr.DataArray:
        """
        Reshapes points to polars.

        Output the polar operators is a 1D array of rho * theta.

        To visualize the result, we need to reshape our array.
        """
        def reshape(array, rho_size, theta_size):
            return array.reshape(rho_size, theta_size)

        output_array = xr.apply_ufunc(
            reshape,
            input_array,
            kwargs=dict(rho_size=self.rho.size, theta_size=self.theta.size),
            keep_attrs=True,
            input_core_dims=[
                ["points"],
            ],
            output_core_dims=[
                ["rho", "theta"],
            ],
            vectorize=True,
            dask="parallelized",
            dask_gufunc_kwargs={
                "output_sizes": dict(rho=self.rho.size, theta=self.theta.size)
            },
            output_dtypes=[np.float64],
        ).assign_coords(rho=self.rho.values, theta=self.theta.values)

        return output_array