Source code for torx.grid.mesh_m

"""Low-level interface to the PARALLAX mesh."""
import xarray as xr
import numpy as np
from inspect import getmodule
import torx.grid.multigrid_operations_m as operators
from torx.autodoc_decorators_m import autodoc_class



@autodoc_class
class Mesh:
    """Represents a PARALLAX mesh type (2D)."""



    def __init__(self, dataset: xr.Dataset):
        """Initialize the mesh object."""
        self._dataset = dataset.rename(
            n_points="points",
            n_points_inner="points_inner",
            n_points_boundary="points_boundary",
            n_points_ghost="points_ghost",
        )


    @property
    def r_u(self) -> np.ndarray:
        """Unstructured Cartesian radial coordinate."""
        return operators.calc_r_u(self).values

    @property
    def r_s(self) -> np.ndarray:
        """Structured Cartesian radial coordinate."""
        return operators.calc_r_s(self)

    @property
    def z_u(self) -> np.ndarray:
        """Unstructured Cartesian vertical coordinate."""
        return operators.calc_z_u(self).values

    @property
    def z_s(self) -> np.ndarray:
        """Structured Cartesian vertical coordinate."""
        return operators.calc_z_s(self)



    def vector_to_matrix(self, unstructured_data: xr.DataArray) -> xr.DataArray:
        """Perform vector to matrix operation on the mesh."""
        return operators.vector_to_matrix(self, unstructured_data)




    def matrix_to_vector(self, structured_data: xr.DataArray) -> xr.DataArray:
        """Perform matrix to vector operation on the mesh."""
        return operators.matrix_to_vector(self, structured_data)




    def inner_to_full(self) -> xr.DataArray:
        """Perform inner to full operation on the mesh."""
        return operators.inner_to_full(self)


    @property
    def points(self) -> int:
        """Total number of points."""
        return self._dataset.sizes["points"]

    @property
    def points_inner(self) -> int:
        """Number of inner points."""
        return self._dataset.sizes["points_inner"]

    @property
    def points_boundary(self) -> int:
        """Number of boundary points."""
        return self._dataset.sizes["points_boundary"]

    @property
    def points_ghost(self) -> int:
        """Number of ghost points."""
        return self._dataset.sizes["points_ghost"]

    @property
    def size_neighbor(self) -> int:
        """Number of neighboring points in each direction."""
        # Check that dx and dy is the same so it does not matter which one
        # we return (if they are not, the mesh is most likely corrupt)
        assert self._dataset.index_neighbor.sizes["dx"] \
                == self._dataset.index_neighbor.sizes["dy"]
        return self._dataset.index_neighbor.sizes["dx"]

    @property
    def phi(self) -> float:
        """Toroidal position of this mesh."""
        return self._dataset.phi

    @property
    def lvl(self) -> int:
        """Multigrid level."""
        return self._dataset.lvl

    @property
    def lvst(self) -> int:
        """Discrete grid distance."""
        return self._dataset.lvst

    @property
    def spacing_f(self) -> float:
        """Grid spacing on finest multigrid level."""
        return self._dataset.spacing_f

    @property
    def spacing_c(self) -> float:
        """Grid spacing on this multigrid level."""
        return self._dataset.spacing_c

    @property
    def xmin(self) -> float:
        """Minimum x value (used for grid construction)."""
        return self._dataset.xmin

    @property
    def ymin(self) -> float:
        """Minimum y value (used for grid construction)."""
        return self._dataset.ymin

    @property
    def nx_f(self) -> int:
        """x-dimension of rectangular Cartesian unmasked grid (finest level)."""
        return self._dataset.nx_f

    @property
    def ny_f(self) -> int:
        """y-dimension of rectangular Cartesian unmasked grid (finest level)."""
        return self._dataset.ny_f

    @property
    def size_ghost_layer(self) -> int:
        """Number of perpendicular ghost points."""
        return self._dataset.size_ghost_layer

    @property
    def yperiodic(self) -> int:
        """Return periodicity in SLAB equi, y-direction."""
        return self._dataset.yperiodic

    @property
    def n_points_red(self) -> int:
        """Number of 'red' points for red-black Gauss-Seidel."""
        return self._dataset.n_points_red

    @property
    def n_points_black(self) -> int:
        """Number of 'black' points for red-black Gauss-Seidel."""
        return self._dataset.n_points_black

    @property
    def cart_i(self) -> xr.DataArray:
        """Horizontal integer position of each point."""
        return self._dataset.cart_i

    @property
    def cart_j(self) -> xr.DataArray:
        """Vertical integer position of each point."""
        return self._dataset.cart_j

    @property
    def pinfo(self) -> xr.DataArray:
        """
        Grid descriptor.

        1 for inner, 2 for boundary, 3 for ghost.
        """
        return self._dataset.pinfo

    @property
    def district(self) -> xr.DataArray:
        """District of point."""
        return self._dataset.district

    @property
    def inner_indices(self) -> xr.DataArray:
        """Indices of inner points."""
        return self._dataset.inner_indices

    @property
    def boundary_indices(self) -> xr.DataArray:
        """Indices of boundary points."""
        return self._dataset.boundary_indices

    @property
    def ghost_indices(self) -> xr.DataArray:
        """Indices of ghost points."""
        return self._dataset.ghost_indices

    @property
    def index_neighbor(self) -> xr.DataArray:
        """Indices of perpendicular neighbors."""
        return self._dataset.index_neighbor

    @property
    def redblack_indices(self) -> xr.DataArray:
        """
        Index list of redblack points.

        Eventually needed for red-black Gauss Seidel smoother.
        First number of red entries are indices of red points.
        Following number of black entries are indices of black points.
        Final number of boundary + number of ghost entries are indices of
        boundary and ghost points.
        """
        return self._dataset.redblack_indices

    @property
    def size(self) -> int:
        """
        Return the number of points in the mesh.

        Enables interchanging 'mesh' with 'grid'.
        """
        return self.points