"""Collections of DataArrays."""
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Union, Dict, Callable, Any, Tuple

import xarray as xr
import numpy as np
import pydantic.v1 as pd
from matplotlib.tri import Triangulation
from matplotlib import pyplot as plt
import numbers

from .data_array import DataArray, DATA_ARRAY_MAP
from .data_array import ScalarFieldDataArray, ScalarFieldTimeDataArray, ScalarModeFieldDataArray
from .data_array import ModeIndexDataArray, GroupIndexDataArray, ModeDispersionDataArray
from .data_array import TriangleMeshDataArray
from .data_array import TimeDataArray
from .data_array import PointDataArray, IndexedDataArray, CellDataArray, SpatialDataArray
from .data_array import EMEScalarFieldDataArray, EMEScalarModeFieldDataArray

from ..viz import equal_aspect, add_ax_if_none, plot_params_grid
from ..base import Tidy3dBaseModel, cached_property
from ..base import skip_if_fields_missing
from ..types import Axis, Bound, ArrayLike, Ax, Coordinate, Literal, annotate_type
from ...packaging import vtk, requires_vtk
from ...exceptions import DataError, ValidationError, Tidy3dNotImplementedError
from ...constants import PICOSECOND_PER_NANOMETER_PER_KILOMETER, inf
from ...log import log


DEFAULT_MAX_SAMPLES_PER_STEP = 10_000
DEFAULT_MAX_CELLS_PER_STEP = 10_000
DEFAULT_TOLERANCE_CELL_FINDING = 1e-6


class Dataset(Tidy3dBaseModel, ABC):
    """Abstract base class for objects that store collections of `:class:`.DataArray`s."""



[docs]
class AbstractFieldDataset(Dataset, ABC):
    """Collection of scalar fields with some symmetry properties."""

    @property
    @abstractmethod
    def field_components(self) -> Dict[str, DataArray]:
        """Maps the field components to their associated data."""


[docs]
    def apply_phase(self, phase: float) -> AbstractFieldDataset:
        """Create a copy where all elements are phase-shifted by a value (in radians)."""
        if phase == 0.0:
            return self
        phasor = np.exp(1j * phase)
        field_components_shifted = {}
        for fld_name, fld_cmp in self.field_components.items():
            fld_cmp_shifted = phasor * fld_cmp
            field_components_shifted[fld_name] = fld_cmp_shifted
        return self.updated_copy(**field_components_shifted)


    @property
    @abstractmethod
    def grid_locations(self) -> Dict[str, str]:
        """Maps field components to the string key of their grid locations on the yee lattice."""

    @property
    @abstractmethod
    def symmetry_eigenvalues(self) -> Dict[str, Callable[[Axis], float]]:
        """Maps field components to their (positive) symmetry eigenvalues."""


[docs]
    def package_colocate_results(self, centered_fields: Dict[str, ScalarFieldDataArray]) -> Any:
        """How to package the dictionary of fields computed via self.colocate()."""
        return xr.Dataset(centered_fields)



[docs]
    def colocate(self, x=None, y=None, z=None) -> xr.Dataset:
        """Colocate all of the data at a set of x, y, z coordinates.

        Parameters
        ----------
        x : Optional[array-like] = None
            x coordinates of locations.
            If not supplied, does not try to colocate on this dimension.
        y : Optional[array-like] = None
            y coordinates of locations.
            If not supplied, does not try to colocate on this dimension.
        z : Optional[array-like] = None
            z coordinates of locations.
            If not supplied, does not try to colocate on this dimension.

        Returns
        -------
        xr.Dataset
            Dataset containing all fields at the same spatial locations.
            For more details refer to `xarray's Documentation <https://tinyurl.com/cyca3krz>`_.

        Note
        ----
        For many operations (such as flux calculations and plotting),
        it is important that the fields are colocated at the same spatial locations.
        Be sure to apply this method to your field data in those cases.
        """

        if hasattr(self, "monitor") and self.monitor.colocate:
            with log as consolidated_logger:
                consolidated_logger.warning(
                    "Colocating data that has already been colocated during the solver "
                    "run. For most accurate results when colocating to custom coordinates set "
                    "'Monitor.colocate' to 'False' to use the raw data on the Yee grid "
                    "and avoid double interpolation. Note: the default value was changed to 'True' "
                    "in Tidy3D version 2.4.0."
                )

        # convert supplied coordinates to array and assign string mapping to them
        supplied_coord_map = {k: np.array(v) for k, v in zip("xyz", (x, y, z)) if v is not None}

        # dict of data arrays to combine in dataset and return
        centered_fields = {}

        # loop through field components
        for field_name, field_data in self.field_components.items():
            # loop through x, y, z dimensions and raise an error if only one element along dim
            for coord_name, coords_supplied in supplied_coord_map.items():
                coord_data = np.array(field_data.coords[coord_name])
                if coord_data.size == 1:
                    raise DataError(
                        f"colocate given {coord_name}={coords_supplied}, but "
                        f"data only has one coordinate at {coord_name}={coord_data[0]}. "
                        "Therefore, can't colocate along this dimension. "
                        f"supply {coord_name}=None to skip it."
                    )

            centered_fields[field_name] = field_data.interp(
                **supplied_coord_map, kwargs={"bounds_error": True}
            )

        # combine all centered fields in a dataset
        return self.package_colocate_results(centered_fields)




EMScalarFieldType = Union[
    ScalarFieldDataArray,
    ScalarFieldTimeDataArray,
    ScalarModeFieldDataArray,
    EMEScalarModeFieldDataArray,
    EMEScalarFieldDataArray,
]


class ElectromagneticFieldDataset(AbstractFieldDataset, ABC):
    """Stores a collection of E and H fields with x, y, z components."""

    Ex: EMScalarFieldType = pd.Field(
        None,
        title="Ex",
        description="Spatial distribution of the x-component of the electric field.",
    )
    Ey: EMScalarFieldType = pd.Field(
        None,
        title="Ey",
        description="Spatial distribution of the y-component of the electric field.",
    )
    Ez: EMScalarFieldType = pd.Field(
        None,
        title="Ez",
        description="Spatial distribution of the z-component of the electric field.",
    )
    Hx: EMScalarFieldType = pd.Field(
        None,
        title="Hx",
        description="Spatial distribution of the x-component of the magnetic field.",
    )
    Hy: EMScalarFieldType = pd.Field(
        None,
        title="Hy",
        description="Spatial distribution of the y-component of the magnetic field.",
    )
    Hz: EMScalarFieldType = pd.Field(
        None,
        title="Hz",
        description="Spatial distribution of the z-component of the magnetic field.",
    )

    @property
    def field_components(self) -> Dict[str, DataArray]:
        """Maps the field components to their associated data."""
        fields = {
            "Ex": self.Ex,
            "Ey": self.Ey,
            "Ez": self.Ez,
            "Hx": self.Hx,
            "Hy": self.Hy,
            "Hz": self.Hz,
        }
        return {field_name: field for field_name, field in fields.items() if field is not None}

    @property
    def grid_locations(self) -> Dict[str, str]:
        """Maps field components to the string key of their grid locations on the yee lattice."""
        return dict(Ex="Ex", Ey="Ey", Ez="Ez", Hx="Hx", Hy="Hy", Hz="Hz")

    @property
    def symmetry_eigenvalues(self) -> Dict[str, Callable[[Axis], float]]:
        """Maps field components to their (positive) symmetry eigenvalues."""

        return dict(
            Ex=lambda dim: -1 if (dim == 0) else +1,
            Ey=lambda dim: -1 if (dim == 1) else +1,
            Ez=lambda dim: -1 if (dim == 2) else +1,
            Hx=lambda dim: +1 if (dim == 0) else -1,
            Hy=lambda dim: +1 if (dim == 1) else -1,
            Hz=lambda dim: +1 if (dim == 2) else -1,
        )



[docs]
class FieldDataset(ElectromagneticFieldDataset):
    """Dataset storing a collection of the scalar components of E and H fields in the freq. domain

    Example
    -------
    >>> x = [-1,1]
    >>> y = [-2,0,2]
    >>> z = [-3,-1,1,3]
    >>> f = [2e14, 3e14]
    >>> coords = dict(x=x, y=y, z=z, f=f)
    >>> scalar_field = ScalarFieldDataArray((1+1j) * np.random.random((2,3,4,2)), coords=coords)
    >>> data = FieldDataset(Ex=scalar_field, Hz=scalar_field)
    """

    Ex: ScalarFieldDataArray = pd.Field(
        None,
        title="Ex",
        description="Spatial distribution of the x-component of the electric field.",
    )
    Ey: ScalarFieldDataArray = pd.Field(
        None,
        title="Ey",
        description="Spatial distribution of the y-component of the electric field.",
    )
    Ez: ScalarFieldDataArray = pd.Field(
        None,
        title="Ez",
        description="Spatial distribution of the z-component of the electric field.",
    )
    Hx: ScalarFieldDataArray = pd.Field(
        None,
        title="Hx",
        description="Spatial distribution of the x-component of the magnetic field.",
    )
    Hy: ScalarFieldDataArray = pd.Field(
        None,
        title="Hy",
        description="Spatial distribution of the y-component of the magnetic field.",
    )
    Hz: ScalarFieldDataArray = pd.Field(
        None,
        title="Hz",
        description="Spatial distribution of the z-component of the magnetic field.",
    )




[docs]
class FieldTimeDataset(ElectromagneticFieldDataset):
    """Dataset storing a collection of the scalar components of E and H fields in the time domain

    Example
    -------
    >>> x = [-1,1]
    >>> y = [-2,0,2]
    >>> z = [-3,-1,1,3]
    >>> t = [0, 1e-12, 2e-12]
    >>> coords = dict(x=x, y=y, z=z, t=t)
    >>> scalar_field = ScalarFieldTimeDataArray(np.random.random((2,3,4,3)), coords=coords)
    >>> data = FieldTimeDataset(Ex=scalar_field, Hz=scalar_field)
    """

    Ex: ScalarFieldTimeDataArray = pd.Field(
        None,
        title="Ex",
        description="Spatial distribution of the x-component of the electric field.",
    )
    Ey: ScalarFieldTimeDataArray = pd.Field(
        None,
        title="Ey",
        description="Spatial distribution of the y-component of the electric field.",
    )
    Ez: ScalarFieldTimeDataArray = pd.Field(
        None,
        title="Ez",
        description="Spatial distribution of the z-component of the electric field.",
    )
    Hx: ScalarFieldTimeDataArray = pd.Field(
        None,
        title="Hx",
        description="Spatial distribution of the x-component of the magnetic field.",
    )
    Hy: ScalarFieldTimeDataArray = pd.Field(
        None,
        title="Hy",
        description="Spatial distribution of the y-component of the magnetic field.",
    )
    Hz: ScalarFieldTimeDataArray = pd.Field(
        None,
        title="Hz",
        description="Spatial distribution of the z-component of the magnetic field.",
    )


[docs]
    def apply_phase(self, phase: float) -> AbstractFieldDataset:
        """Create a copy where all elements are phase-shifted by a value (in radians)."""

        if phase != 0.0:
            raise ValueError("Can't apply phase to time-domain field data, which is real-valued.")

        return self





[docs]
class ModeSolverDataset(ElectromagneticFieldDataset):
    """Dataset storing scalar components of E and H fields as a function of freq. and mode_index.

    Example
    -------
    >>> from tidy3d import ModeSpec
    >>> x = [-1,1]
    >>> y = [0]
    >>> z = [-3,-1,1,3]
    >>> f = [2e14, 3e14]
    >>> mode_index = np.arange(5)
    >>> field_coords = dict(x=x, y=y, z=z, f=f, mode_index=mode_index)
    >>> field = ScalarModeFieldDataArray((1+1j)*np.random.random((2,1,4,2,5)), coords=field_coords)
    >>> index_coords = dict(f=f, mode_index=mode_index)
    >>> index_data = ModeIndexDataArray((1+1j) * np.random.random((2,5)), coords=index_coords)
    >>> data = ModeSolverDataset(
    ...     Ex=field,
    ...     Ey=field,
    ...     Ez=field,
    ...     Hx=field,
    ...     Hy=field,
    ...     Hz=field,
    ...     n_complex=index_data
    ... )
    """

    Ex: ScalarModeFieldDataArray = pd.Field(
        None,
        title="Ex",
        description="Spatial distribution of the x-component of the electric field of the mode.",
    )
    Ey: ScalarModeFieldDataArray = pd.Field(
        None,
        title="Ey",
        description="Spatial distribution of the y-component of the electric field of the mode.",
    )
    Ez: ScalarModeFieldDataArray = pd.Field(
        None,
        title="Ez",
        description="Spatial distribution of the z-component of the electric field of the mode.",
    )
    Hx: ScalarModeFieldDataArray = pd.Field(
        None,
        title="Hx",
        description="Spatial distribution of the x-component of the magnetic field of the mode.",
    )
    Hy: ScalarModeFieldDataArray = pd.Field(
        None,
        title="Hy",
        description="Spatial distribution of the y-component of the magnetic field of the mode.",
    )
    Hz: ScalarModeFieldDataArray = pd.Field(
        None,
        title="Hz",
        description="Spatial distribution of the z-component of the magnetic field of the mode.",
    )

    n_complex: ModeIndexDataArray = pd.Field(
        ...,
        title="Propagation Index",
        description="Complex-valued effective propagation constants associated with the mode.",
    )

    n_group_raw: GroupIndexDataArray = pd.Field(
        None,
        alias="n_group",  # This is for backwards compatibility only when loading old data
        title="Group Index",
        description="Index associated with group velocity of the mode.",
    )

    dispersion_raw: ModeDispersionDataArray = pd.Field(
        None,
        title="Dispersion",
        description="Dispersion parameter for the mode.",
        units=PICOSECOND_PER_NANOMETER_PER_KILOMETER,
    )

    @property
    def field_components(self) -> Dict[str, DataArray]:
        """Maps the field components to their associated data."""
        fields = {
            "Ex": self.Ex,
            "Ey": self.Ey,
            "Ez": self.Ez,
            "Hx": self.Hx,
            "Hy": self.Hy,
            "Hz": self.Hz,
        }
        return {field_name: field for field_name, field in fields.items() if field is not None}

    @property
    def n_eff(self) -> ModeIndexDataArray:
        """Real part of the propagation index."""
        return self.n_complex.real

    @property
    def k_eff(self) -> ModeIndexDataArray:
        """Imaginary part of the propagation index."""
        return self.n_complex.imag

    @property
    def n_group(self) -> GroupIndexDataArray:
        """Group index."""
        if self.n_group_raw is None:
            log.warning(
                "The group index was not computed. To calculate group index, pass "
                "'group_index_step = True' in the 'ModeSpec'.",
                log_once=True,
            )
        return self.n_group_raw

    @property
    def dispersion(self) -> ModeDispersionDataArray:
        r"""Dispersion parameter.

        .. math::

           D = -\frac{\lambda}{c_0} \frac{{\rm d}^2 n_{\text{eff}}}{{\rm d}\lambda^2}
        """
        if self.dispersion_raw is None:
            log.warning(
                "The dispersion was not computed. To calculate dispersion, pass "
                "'group_index_step = True' in the 'ModeSpec'.",
                log_once=True,
            )
        return self.dispersion_raw


[docs]
    def plot_field(self, *args, **kwargs):
        """Warn user to use the :class:`.ModeSolver` ``plot_field`` function now."""
        raise DeprecationWarning(
            "The 'plot_field()' method was moved to the 'ModeSolver' object."
            "Once the 'ModeSolver' is constructed, one may call '.plot_field()' on the object and "
            "the modes will be computed and displayed with 'Simulation' overlay."
        )





[docs]
class PermittivityDataset(AbstractFieldDataset):
    """Dataset storing the diagonal components of the permittivity tensor.

    Example
    -------
    >>> x = [-1,1]
    >>> y = [-2,0,2]
    >>> z = [-3,-1,1,3]
    >>> f = [2e14, 3e14]
    >>> coords = dict(x=x, y=y, z=z, f=f)
    >>> sclr_fld = ScalarFieldDataArray((1+1j) * np.random.random((2,3,4,2)), coords=coords)
    >>> data = PermittivityDataset(eps_xx=sclr_fld, eps_yy=sclr_fld, eps_zz=sclr_fld)
    """

    @property
    def field_components(self) -> Dict[str, ScalarFieldDataArray]:
        """Maps the field components to their associated data."""
        return dict(eps_xx=self.eps_xx, eps_yy=self.eps_yy, eps_zz=self.eps_zz)

    @property
    def grid_locations(self) -> Dict[str, str]:
        """Maps field components to the string key of their grid locations on the yee lattice."""
        return dict(eps_xx="Ex", eps_yy="Ey", eps_zz="Ez")

    @property
    def symmetry_eigenvalues(self) -> Dict[str, Callable[[Axis], float]]:
        """Maps field components to their (positive) symmetry eigenvalues."""
        return dict(eps_xx=None, eps_yy=None, eps_zz=None)

    eps_xx: ScalarFieldDataArray = pd.Field(
        ...,
        title="Epsilon xx",
        description="Spatial distribution of the xx-component of the relative permittivity.",
    )
    eps_yy: ScalarFieldDataArray = pd.Field(
        ...,
        title="Epsilon yy",
        description="Spatial distribution of the yy-component of the relative permittivity.",
    )
    eps_zz: ScalarFieldDataArray = pd.Field(
        ...,
        title="Epsilon zz",
        description="Spatial distribution of the zz-component of the relative permittivity.",
    )



class TriangleMeshDataset(Dataset):
    """Dataset for storing triangular surface data."""

    surface_mesh: TriangleMeshDataArray = pd.Field(
        ...,
        title="Surface mesh data",
        description="Dataset containing the surface triangles and corresponding face indices "
        "for a surface mesh.",
    )


class TimeDataset(Dataset):
    """Dataset for storing a function of time."""

    values: TimeDataArray = pd.Field(
        ..., title="Values", description="Values as a function of time."
    )


class UnstructuredGridDataset(Dataset, np.lib.mixins.NDArrayOperatorsMixin, ABC):
    """Abstract base for datasets that store unstructured grid data."""

    points: PointDataArray = pd.Field(
        ...,
        title="Grid Points",
        description="Coordinates of points composing the unstructured grid.",
    )

    values: IndexedDataArray = pd.Field(
        ...,
        title="Point Values",
        description="Values stored at the grid points.",
    )

    cells: CellDataArray = pd.Field(
        ...,
        title="Grid Cells",
        description="Cells composing the unstructured grid specified as connections between grid "
        "points.",
    )

    @property
    def name(self) -> str:
        """Dataset name."""
        # we redirect name to values.name
        return self.values.name

    @property
    def is_complex(self) -> bool:
        """Data type."""
        return np.iscomplexobj(self.values)

    @property
    def _double_type(self):
        """Corresponding double data type."""
        return np.complex128 if self.is_complex else np.float64

    @pd.validator("points", always=True)
    def points_right_dims(cls, val):
        """Check that point coordinates have the right dimensionality."""
        # currently support only the standard axis ordering, that is 01(2)
        axis_coords_expected = np.arange(cls._point_dims())
        axis_coords_given = val.axis.data
        if np.any(axis_coords_given != axis_coords_expected):
            raise ValidationError(
                f"Points array is expected to have {axis_coords_expected} coord values along 'axis'"
                f" (given: {axis_coords_given})."
            )
        return val

    @property
    def is_uniform(self):
        """Whether each element is of equal value in ``values``."""
        return self.values.is_uniform

    @pd.validator("points", always=True)
    def points_right_indexing(cls, val):
        """Check that points are indexed corrrectly."""
        indices_expected = np.arange(len(val.data))
        indices_given = val.index.data
        if np.any(indices_expected != indices_given):
            raise ValidationError(
                "Coordinate 'index' of array 'points' is expected to have values (0, 1, 2, ...). "
                "This can be easily achieved, for example, by using "
                "PointDataArray(data, dims=['index', 'axis'])."
            )
        return val

    @pd.validator("values", always=True)
    def values_right_indexing(cls, val):
        """Check that data values are indexed correctly."""
        # currently support only simple ordered indexing of points, that is, 0, 1, 2, ...
        indices_expected = np.arange(len(val.data))
        indices_given = val.index.data
        if np.any(indices_expected != indices_given):
            raise ValidationError(
                "Coordinate 'index' of array 'values' is expected to have values (0, 1, 2, ...). "
                "This can be easily achieved, for example, by using "
                "IndexedDataArray(data, dims=['index'])."
            )
        return val

    @pd.validator("values", always=True)
    @skip_if_fields_missing(["points"])
    def number_of_values_matches_points(cls, val, values):
        """Check that the number of data values matches the number of grid points."""
        num_values = len(val)

        points = values.get("points")
        num_points = len(points)

        if num_points != num_values:
            raise ValidationError(
                f"The number of data values ({num_values}) does not match the number of grid "
                f"points ({num_points})."
            )
        return val

    @pd.validator("cells", always=True)
    def match_cells_to_vtk_type(cls, val):
        """Check that cell connections does not have duplicate points."""
        if vtk is None:
            return val

        # using val.astype(np.int32/64) directly causes issues when dataarray are later checked ==
        return CellDataArray(val.data.astype(vtk["id_type"], copy=False), coords=val.coords)

    @pd.validator("cells", always=True)
    def cells_right_type(cls, val):
        """Check that cell are of the right type."""
        # only supporting the standard ordering of cell vertices 012(3)
        vertex_coords_expected = np.arange(cls._cell_num_vertices())
        vertex_coords_given = val.vertex_index.data
        if np.any(vertex_coords_given != vertex_coords_expected):
            raise ValidationError(
                f"Cell connections array is expected to have {vertex_coords_expected} coord values"
                f" along 'vertex_index' (given: {vertex_coords_given})."
            )
        return val

    @pd.validator("cells", always=True)
    @skip_if_fields_missing(["points"])
    def check_cell_vertex_range(cls, val, values):
        """Check that cell connections use only defined points."""
        all_point_indices_used = val.data.ravel()
        # skip validation if zero size data
        if len(all_point_indices_used) > 0:
            min_index_used = np.min(all_point_indices_used)
            max_index_used = np.max(all_point_indices_used)

            points = values.get("points")
            num_points = len(points)

            if max_index_used > num_points - 1 or min_index_used < 0:
                raise ValidationError(
                    "Cell connections array uses undefined point indices in the range "
                    f"[{min_index_used}, {max_index_used}]. The valid range of point indices is "
                    f"[0, {num_points-1}]."
                )
        return val

    @classmethod
    def _find_degenerate_cells(cls, cells: CellDataArray):
        """Find explicitly degenerate cells if any.
        That is, cells that use the same point indices for their different vertices.
        """
        indices = cells.data
        # skip validation if zero size data
        degenerate_cell_inds = set()
        if len(indices) > 0:
            for i in range(cls._cell_num_vertices() - 1):
                for j in range(i + 1, cls._cell_num_vertices()):
                    degenerate_cell_inds = degenerate_cell_inds.union(
                        np.where(indices[:, i] == indices[:, j])[0]
                    )

        return degenerate_cell_inds

    @classmethod
    def _remove_degenerate_cells(cls, cells: CellDataArray):
        """Remove explicitly degenerate cells if any.
        That is, cells that use the same point indices for their different vertices.
        """
        degenerate_cells = cls._find_degenerate_cells(cells=cells)
        if len(degenerate_cells) > 0:
            data = np.delete(cells.values, list(degenerate_cells), axis=0)
            cell_index = np.delete(cells.cell_index.values, list(degenerate_cells))
            return CellDataArray(
                data=data, coords=dict(cell_index=cell_index, vertex_index=cells.vertex_index)
            )
        return cells

    @classmethod
    def _remove_unused_points(
        cls, points: PointDataArray, values: IndexedDataArray, cells: CellDataArray
    ):
        """Remove unused points if any.
        That is, points that are not used in any grid cell.
        """

        used_indices = np.unique(cells.values.ravel())
        num_points = len(points)

        if len(used_indices) != num_points or np.any(np.diff(used_indices) != 1):
            min_index = np.min(used_indices)
            map_len = np.max(used_indices) - min_index + 1
            index_map = np.zeros(map_len)
            index_map[used_indices - min_index] = np.arange(len(used_indices))

            cells = CellDataArray(data=index_map[cells.data - min_index], coords=cells.coords)
            points = PointDataArray(points.data[used_indices, :], dims=["index", "axis"])
            values = IndexedDataArray(values.data[used_indices], dims=["index"])

        return points, values, cells

    def clean(self, remove_degenerate_cells=True, remove_unused_points=True):
        """Remove degenerate cells and/or unused points."""
        if remove_degenerate_cells:
            cells = self._remove_degenerate_cells(cells=self.cells)
        else:
            cells = self.cells

        if remove_unused_points:
            points, values, cells = self._remove_unused_points(self.points, self.values, cells)
        else:
            points = self.points
            values = self.values

        return self.updated_copy(points=points, values=values, cells=cells)

    @pd.validator("cells", always=True)
    def warn_degenerate_cells(cls, val):
        """Check that cell connections does not have duplicate points."""
        degenerate_cells = cls._find_degenerate_cells(val)
        num_degenerate_cells = len(degenerate_cells)
        if num_degenerate_cells > 0:
            log.warning(
                f"Unstructured grid contains {num_degenerate_cells} degenerate cell(s). "
                "Such cells can be removed by using function "
                "'.clean(remove_degenerate_cells: bool = True, remove_unused_points: bool = True)'. "
                "For example, 'dataset = dataset.clean()'."
            )
        return val

    @pd.root_validator(pre=True, allow_reuse=True)
    def _warn_if_none(cls, values):
        """Warn if any of data arrays are not loaded."""

        no_data_fields = []
        for field_name in ["points", "cells", "values"]:
            field = values.get(field_name)
            if isinstance(field, str) and field in DATA_ARRAY_MAP.keys():
                no_data_fields.append(field_name)
        if len(no_data_fields) > 0:
            formatted_names = [f"'{fname}'" for fname in no_data_fields]
            log.warning(
                f"Loading {', '.join(formatted_names)} without data. Constructing an empty dataset."
            )
            values["points"] = PointDataArray(
                np.zeros((0, cls._point_dims())), dims=["index", "axis"]
            )
            values["cells"] = CellDataArray(
                np.zeros((0, cls._cell_num_vertices())), dims=["cell_index", "vertex_index"]
            )
            values["values"] = IndexedDataArray(np.zeros(0), dims=["index"])
        return values

    @pd.root_validator(skip_on_failure=True, allow_reuse=True)
    def _warn_unused_points(cls, values):
        """Warn if some points are unused."""
        point_indices = set(np.arange(len(values["values"].data)))
        used_indices = set(values["cells"].values.ravel())

        if not point_indices.issubset(used_indices):
            log.warning(
                "Unstructured grid dataset contains unused points. "
                "Consider calling 'clean()' to remove them."
            )

        return values

    def rename(self, name: str) -> UnstructuredGridDataset:
        """Return a renamed array."""
        return self.updated_copy(values=self.values.rename(name))

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        """Override of numpy functions."""

        out = kwargs.get("out", ())
        for x in inputs + out:
            # Only support operations with the same class or a scalar
            if not isinstance(x, (numbers.Number, type(self))):
                return Tidy3dNotImplementedError

        # Defer to the implementation of the ufunc on unwrapped values.
        inputs = tuple(x.values if isinstance(x, UnstructuredGridDataset) else x for x in inputs)
        if out:
            kwargs["out"] = tuple(
                x.values if isinstance(x, UnstructuredGridDataset) else x for x in out
            )
        result = getattr(ufunc, method)(*inputs, **kwargs)

        if type(result) is tuple:
            # multiple return values
            return tuple(self.updated_copy(values=x) for x in result)
        elif method == "at":
            # no return value
            return None
        else:
            # one return value
            return self.updated_copy(values=result)

    @property
    def real(self) -> UnstructuredGridDataset:
        """Real part of dataset."""
        return self.updated_copy(values=self.values.real)

    @property
    def imag(self) -> UnstructuredGridDataset:
        """Imaginary part of dataset."""
        return self.updated_copy(values=self.values.imag)

    @property
    def abs(self) -> UnstructuredGridDataset:
        """Absolute value of dataset."""
        return self.updated_copy(values=self.values.abs)

    @cached_property
    def bounds(self) -> Bound:
        """Grid bounds."""
        return tuple(np.min(self.points.data, axis=0)), tuple(np.max(self.points.data, axis=0))

    @classmethod
    @abstractmethod
    def _point_dims(cls) -> pd.PositiveInt:
        """Dimensionality of stored grid point coordinates."""

    @cached_property
    @abstractmethod
    def _points_3d_array(self):
        """3D coordinates of grid points."""

    @classmethod
    @abstractmethod
    def _cell_num_vertices(cls) -> pd.PositiveInt:
        """Number of vertices in a cell."""

    @classmethod
    @abstractmethod
    @requires_vtk
    def _vtk_cell_type(cls):
        """VTK cell type to use in the VTK representation."""

    @cached_property
    def _vtk_offsets(self) -> ArrayLike:
        """Offsets array to use in the VTK representation."""
        offsets = np.arange(len(self.cells) + 1) * self._cell_num_vertices()
        if vtk is None:
            return offsets

        return offsets.astype(vtk["id_type"], copy=False)

    @property
    @requires_vtk
    def _vtk_cells(self):
        """VTK cell array to use in the VTK representation."""
        cells = vtk["mod"].vtkCellArray()
        cells.SetData(
            vtk["numpy_to_vtkIdTypeArray"](self._vtk_offsets),
            vtk["numpy_to_vtkIdTypeArray"](self.cells.data.ravel()),
        )
        return cells

    @property
    @requires_vtk
    def _vtk_points(self):
        """VTK point array to use in the VTK representation."""
        pts = vtk["mod"].vtkPoints()
        pts.SetData(vtk["numpy_to_vtk"](self._points_3d_array))
        return pts

    @property
    @requires_vtk
    def _vtk_obj(self):
        """A VTK representation (vtkUnstructuredGrid) of the grid."""

        grid = vtk["mod"].vtkUnstructuredGrid()

        grid.SetPoints(self._vtk_points)
        grid.SetCells(self._vtk_cell_type(), self._vtk_cells)
        if self.is_complex:
            # vtk doesn't support complex numbers
            # so we will store our complex array as a two-component vtk array
            data_values = self.values.values.view("(2,)float")
        else:
            data_values = self.values.values
        point_data_vtk = vtk["numpy_to_vtk"](data_values)
        point_data_vtk.SetName(self.values.name)
        grid.GetPointData().AddArray(point_data_vtk)

        return grid

    @requires_vtk
    def _plane_slice_raw(self, axis: Axis, pos: float):
        """Slice data with a plane and return the resulting VTK object."""

        if pos > self.bounds[1][axis] or pos < self.bounds[0][axis]:
            raise DataError(
                f"Slicing plane (axis: {axis}, pos: {pos}) does not intersect the unstructured grid "
                f"(extent along axis {axis}: {self.bounds[0][axis]}, {self.bounds[1][axis]})."
            )

        origin = [0, 0, 0]
        origin[axis] = pos

        normal = [0, 0, 0]
        # orientation of normal is important for edge (literally) cases
        normal[axis] = -1
        if pos > (self.bounds[0][axis] + self.bounds[1][axis]) / 2:
            normal[axis] = 1

        # create cutting plane
        plane = vtk["mod"].vtkPlane()
        plane.SetOrigin(origin[0], origin[1], origin[2])
        plane.SetNormal(normal[0], normal[1], normal[2])

        # create cutter
        cutter = vtk["mod"].vtkPlaneCutter()
        cutter.SetPlane(plane)
        cutter.SetInputData(self._vtk_obj)
        cutter.InterpolateAttributesOn()
        cutter.Update()

        # clean up the slice
        cleaner = vtk["mod"].vtkCleanPolyData()
        cleaner.SetInputData(cutter.GetOutput())
        cleaner.Update()

        return cleaner.GetOutput()

    @abstractmethod
    @requires_vtk
    def plane_slice(
        self, axis: Axis, pos: float
    ) -> Union[SpatialDataArray, UnstructuredGridDataset]:
        """Slice data with a plane and return the Tidy3D representation of the result
        (``UnstructuredGridDataset``).

        Parameters
        ----------
        axis : Axis
            The normal direction of the slicing plane.
        pos : float
            Position of the slicing plane along its normal direction.

        Returns
        -------
        Union[SpatialDataArray, UnstructuredGridDataset]
            The resulting slice.
        """

    @staticmethod
    @requires_vtk
    def _read_vtkUnstructuredGrid(fname: str):
        """Load a :class:`vtkUnstructuredGrid` from a file."""
        reader = vtk["mod"].vtkXMLUnstructuredGridReader()
        reader.SetFileName(fname)
        reader.Update()
        grid = reader.GetOutput()

        return grid

    @classmethod
    @abstractmethod
    @requires_vtk
    def _from_vtk_obj(
        cls,
        vtk_obj,
        field: str = None,
        remove_degenerate_cells: bool = False,
        remove_unused_points: bool = False,
    ) -> UnstructuredGridDataset:
        """Initialize from a vtk object."""

    @classmethod
    @requires_vtk
    def from_vtu(
        cls,
        file: str,
        field: str = None,
        remove_degenerate_cells: bool = False,
        remove_unused_points: bool = False,
    ) -> UnstructuredGridDataset:
        """Load unstructured data from a vtu file.

        Parameters
        ----------
        fname : str
            Full path to the .vtu file to load the unstructured data from.
        field : str = None
            Name of the field to load.
        remove_degenerate_cells : bool = False
            Remove explicitly degenerate cells.
        remove_unused_points : bool = False
            Remove unused points.

        Returns
        -------
        UnstructuredGridDataset
            Unstructured data.
        """
        grid = cls._read_vtkUnstructuredGrid(file)
        return cls._from_vtk_obj(
            grid,
            field=field,
            remove_degenerate_cells=remove_degenerate_cells,
            remove_unused_points=remove_unused_points,
        )

    @requires_vtk
    def to_vtu(self, fname: str):
        """Exports unstructured grid data into a .vtu file.

        Parameters
        ----------
        fname : str
            Full path to the .vtu file to save the unstructured data to.
        """

        writer = vtk["mod"].vtkXMLUnstructuredGridWriter()
        writer.SetFileName(fname)
        writer.SetInputData(self._vtk_obj)
        writer.Write()

    @classmethod
    @requires_vtk
    def _get_values_from_vtk(
        cls, vtk_obj, num_points: pd.PositiveInt, field: str = None
    ) -> IndexedDataArray:
        """Get point data values from a VTK object."""

        point_data = vtk_obj.GetPointData()
        num_point_arrays = point_data.GetNumberOfArrays()

        if num_point_arrays == 0:
            log.warning(
                "No point data is found in a VTK object. '.values' will be initialized to zeros."
            )
            values_numpy = np.zeros(num_points)
            values_name = None

        else:
            if field is not None:
                array_vtk = point_data.GetAbstractArray(field)
            else:
                array_vtk = point_data.GetAbstractArray(0)

            # currently we assume there is only one point data array provided in the VTK object
            if num_point_arrays > 1 and field is None:
                array_name = array_vtk.GetName()
                log.warning(
                    f"{num_point_arrays} point data arrays are found in a VTK object. "
                    f"Only the first array (name: {array_name}) will be used to initialize "
                    "'.values' while the rest will be ignored."
                )

            # currently we assume data is real or complex scalar
            num_components = array_vtk.GetNumberOfComponents()
            if num_components > 2:
                raise DataError(
                    "Found point data array in a VTK object is expected to have maximum 2 "
                    "components (1 is for real data, 2 is for complex data). "
                    f"Found {num_components} components."
                )

            # check that number of values matches number of grid points
            num_tuples = array_vtk.GetNumberOfTuples()
            if num_tuples != num_points:
                raise DataError(
                    f"The length of found point data array ({num_tuples}) does not match the number"
                    f" of grid points ({num_points})."
                )

            values_numpy = vtk["vtk_to_numpy"](array_vtk)
            values_name = array_vtk.GetName()

            # vtk doesn't support complex numbers
            # we store our complex array as a two-component vtk array
            # so here we convert that into a single component complex array
            if num_components == 2:
                values_numpy = values_numpy.view("complex")[:, 0]

        values = IndexedDataArray(
            values_numpy, coords=dict(index=np.arange(len(values_numpy))), name=values_name
        )

        return values

    @requires_vtk
    def box_clip(self, bounds: Bound) -> UnstructuredGridDataset:
        """Clip the unstructured grid using a box defined by ``bounds``.

        Parameters
        ----------
        bounds : Tuple[float, float, float], Tuple[float, float float]
            Min and max bounds packaged as ``(minx, miny, minz), (maxx, maxy, maxz)``.

        Returns
        -------
        UnstructuredGridDataset
            Clipped grid.
        """

        # make and run a VTK clipper
        clipper = vtk["mod"].vtkBoxClipDataSet()
        clipper.SetOrientation(0)
        clipper.SetBoxClip(
            bounds[0][0], bounds[1][0], bounds[0][1], bounds[1][1], bounds[0][2], bounds[1][2]
        )
        clipper.SetInputData(self._vtk_obj)
        clipper.GenerateClipScalarsOn()
        clipper.GenerateClippedOutputOff()
        clipper.Update()
        clip = clipper.GetOutput()

        # clean grid from unused points
        grid_cleaner = vtk["mod"].vtkRemoveUnusedPoints()
        grid_cleaner.SetInputData(clip)
        grid_cleaner.GenerateOriginalPointIdsOff()
        grid_cleaner.Update()
        clean_clip = grid_cleaner.GetOutput()

        # no intersection check
        if clean_clip.GetNumberOfPoints() == 0:
            raise DataError("Clipping box does not intersect the unstructured grid.")

        return self._from_vtk_obj(
            clean_clip, remove_degenerate_cells=True, remove_unused_points=True
        )

    def interp(
        self,
        x: Union[float, ArrayLike],
        y: Union[float, ArrayLike],
        z: Union[float, ArrayLike],
        fill_value: Union[float, Literal["extrapolate"]] = None,
        use_vtk: bool = False,
        method: Literal["linear", "nearest"] = "linear",
        max_samples_per_step: int = DEFAULT_MAX_SAMPLES_PER_STEP,
        max_cells_per_step: int = DEFAULT_MAX_CELLS_PER_STEP,
        rel_tol: float = DEFAULT_TOLERANCE_CELL_FINDING,
    ) -> SpatialDataArray:
        """Interpolate data at provided x, y, and z.

        Parameters
        ----------
        x : Union[float, ArrayLike]
            x-coordinates of sampling points.
        y : Union[float, ArrayLike]
            y-coordinates of sampling points.
        z : Union[float, ArrayLike]
            z-coordinates of sampling points.
        fill_value : Union[float, Literal["extrapolate"]] = 0
            Value to use when filling points without interpolated values. If ``"extrapolate"`` then
            nearest values are used. Note: in a future version the default value will be changed
            to ``"extrapolate"``.
        use_vtk : bool = False
            Use vtk's interpolation functionality or Tidy3D's own implementation. Note: this
            option will be removed in a future version.
        method: Literal["linear", "nearest"] = "linear"
            Interpolation method to use.
        max_samples_per_step : int = 1e4
            Max number of points to interpolate at per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        max_cells_per_step : int = 1e4
            Max number of cells to interpolate from per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        rel_tol : float = 1e-6
            Relative tolerance when determining whether a point belongs to a cell.

        Returns
        -------
        SpatialDataArray
            Interpolated data.
        """

        if fill_value is None:
            log.warning(
                "Default parameter setting 'fill_value=0' will be changed to "
                "'fill_value=``extrapolate``' in a future version."
            )
            fill_value = 0

        # calculate the resulting array shape
        x = np.atleast_1d(x)
        y = np.atleast_1d(y)
        z = np.atleast_1d(z)

        if method == "nearest":
            interpolated_values = self._interp_nearest(x=x, y=y, z=z)
        else:
            if fill_value == "extrapolate":
                fill_value_actual = np.nan
            else:
                fill_value_actual = fill_value

            if use_vtk:
                log.warning("Note that option 'use_vtk=True' will be removed in future versions.")
                interpolated_values = self._interp_vtk(x=x, y=y, z=z, fill_value=fill_value_actual)
            else:
                interpolated_values = self._interp_py(
                    x=x,
                    y=y,
                    z=z,
                    fill_value=fill_value_actual,
                    max_samples_per_step=max_samples_per_step,
                    max_cells_per_step=max_cells_per_step,
                    rel_tol=rel_tol,
                )

            if fill_value == "extrapolate" and method != "nearest":
                interpolated_values = self._fill_nans_from_nearests(
                    interpolated_values, x=x, y=y, z=z
                )

        return SpatialDataArray(
            interpolated_values, coords=dict(x=x, y=y, z=z), name=self.values.name
        )

    def _interp_nearest(
        self,
        x: ArrayLike,
        y: ArrayLike,
        z: ArrayLike,
    ) -> ArrayLike:
        """Interpolate data at provided x, y, and z using Scipy's nearest neighbor interpolator.

        Parameters
        ----------
        x : ArrayLike
            x-coordinates of sampling points.
        y : ArrayLike
            y-coordinates of sampling points.
        z : ArrayLike
            z-coordinates of sampling points.

        Returns
        -------
        ArrayLike
            Interpolated data.
        """
        from scipy.interpolate import NearestNDInterpolator

        # use scipy's nearest neighbor interpolator
        X, Y, Z = np.meshgrid(x, y, z, indexing="ij")
        interp = NearestNDInterpolator(self._points_3d_array, self.values.values)
        values = interp(X, Y, Z)

        return values

    def _fill_nans_from_nearests(
        self,
        values: ArrayLike,
        x: ArrayLike,
        y: ArrayLike,
        z: ArrayLike,
    ) -> ArrayLike:
        """Replace nan's in ``values`` with nearest data points.

        Parameters
        ----------
        values : ArrayLike
            3D array containing nan's
        x : ArrayLike
            x-coordinates of sampling points.
        y : ArrayLike
            y-coordinates of sampling points.
        z : ArrayLike
            z-coordinates of sampling points.

        Returns
        -------
        ArrayLike
            Data without nan's.
        """

        # locate all nans
        nans = np.isnan(values)

        if np.sum(nans) > 0:
            from scipy.interpolate import NearestNDInterpolator

            # use scipy's nearest neighbor interpolator
            X, Y, Z = np.meshgrid(x, y, z, indexing="ij")
            interp = NearestNDInterpolator(self._points_3d_array, self.values.values)
            values_to_replace_nans = interp(X[nans], Y[nans], Z[nans])
            values[nans] = values_to_replace_nans

        return values

    @requires_vtk
    def _interp_vtk(
        self,
        x: ArrayLike,
        y: ArrayLike,
        z: ArrayLike,
        fill_value: float,
    ) -> ArrayLike:
        """Interpolate data at provided x, y, and z using vtk package.

        Parameters
        ----------
        x : Union[float, ArrayLike]
            x-coordinates of sampling points.
        y : Union[float, ArrayLike]
            y-coordinates of sampling points.
        z : Union[float, ArrayLike]
            z-coordinates of sampling points.
        fill_value : float = 0
            Value to use when filling points without interpolated values.

        Returns
        -------
        ArrayLike
            Interpolated data.
        """

        shape = (len(x), len(y), len(z))

        # create a VTK rectilinear grid to sample onto
        structured_grid = vtk["mod"].vtkRectilinearGrid()
        structured_grid.SetDimensions(shape)
        structured_grid.SetXCoordinates(vtk["numpy_to_vtk"](x))
        structured_grid.SetYCoordinates(vtk["numpy_to_vtk"](y))
        structured_grid.SetZCoordinates(vtk["numpy_to_vtk"](z))

        # create and execute VTK interpolator
        interpolator = vtk["mod"].vtkResampleWithDataSet()
        interpolator.SetInputData(structured_grid)
        interpolator.SetSourceData(self._vtk_obj)
        interpolator.Update()
        interpolated = interpolator.GetOutput()

        # get results in a numpy representation
        array_id = 0 if self.values.name is None else self.values.name
        values_numpy = vtk["vtk_to_numpy"](interpolated.GetPointData().GetAbstractArray(array_id))

        # fill points without interpolated values
        if fill_value != 0:
            mask = vtk["vtk_to_numpy"](
                interpolated.GetPointData().GetAbstractArray("vtkValidPointMask")
            )
            values_numpy[mask != 1] = fill_value

        # VTK arrays are the z-y-x order, reorder interpolation results to x-y-z order
        values_reordered = np.transpose(np.reshape(values_numpy, shape[::-1]), (2, 1, 0))

        return values_reordered

    @abstractmethod
    def _interp_py(
        self,
        x: ArrayLike,
        y: ArrayLike,
        z: ArrayLike,
        fill_value: float,
        max_samples_per_step: int,
        max_cells_per_step: int,
        rel_tol: float,
    ) -> ArrayLike:
        """Dimensionality-specific function (2D and 3D) to interpolate data at provided x, y, and z
        using vectorized python implementation.

        Parameters
        ----------
        x : Union[float, ArrayLike]
            x-coordinates of sampling points.
        y : Union[float, ArrayLike]
            y-coordinates of sampling points.
        z : Union[float, ArrayLike]
            z-coordinates of sampling points.
        fill_value : float
            Value to use when filling points without interpolated values.
        max_samples_per_step : int
            Max number of points to interpolate at per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        max_cells_per_step : int
            Max number of cells to interpolate from per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        rel_tol : float
            Relative tolerance when determining whether a point belongs to a cell.

        Returns
        -------
        ArrayLike
            Interpolated data.
        """

    def _interp_py_general(
        self,
        x: ArrayLike,
        y: ArrayLike,
        z: ArrayLike,
        fill_value: float,
        max_samples_per_step: int,
        max_cells_per_step: int,
        rel_tol: float,
        axis_ignore: Union[Axis, None],
    ) -> ArrayLike:
        """A general function (2D and 3D) to interpolate data at provided x, y, and z using
        vectorized python implementation.

        Parameters
        ----------
        x : Union[float, ArrayLike]
            x-coordinates of sampling points.
        y : Union[float, ArrayLike]
            y-coordinates of sampling points.
        z : Union[float, ArrayLike]
            z-coordinates of sampling points.
        fill_value : float
            Value to use when filling points without interpolated values.
        max_samples_per_step : int
            Max number of points to interpolate at per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        max_cells_per_step : int
            Max number of cells to interpolate from per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        rel_tol : float
            Relative tolerance when determining whether a point belongs to a cell.
        axis_ignore : Union[Axis, None]
            When interpolating from a 2D dataset, must specify normal axis.

        Returns
        -------
        ArrayLike
            Interpolated data.
        """

        # get dimensionality of data
        num_dims = self._point_dims()

        if num_dims == 2 and axis_ignore is None:
            raise DataError("Must porvide 'axis_ignore' when interpolating from a 2d dataset.")

        xyz_grid = [x, y, z]

        if axis_ignore is not None:
            xyz_grid.pop(axis_ignore)

        # get numpy arrays for points and cells
        cell_connections = (
            self.cells.values
        )  # (num_cells, num_cell_vertices), num_cell_vertices=num_cell_faces
        points = self.points.values  # (num_points, num_dims)

        num_cells = len(cell_connections)
        num_points = len(points)

        # compute tolerances based on total size of unstructured grid
        bounds = self.bounds
        size = np.subtract(bounds[1], bounds[0])
        tol = size * rel_tol
        diag_tol = np.linalg.norm(tol)

        # compute (index) positions of unstructured points w.r.t. target Cartesian grid points
        # (i.e. between which Cartesian grid points a given unstructured grid point is located)
        # we perturb grid values in both directions to make sure we don't miss any points
        # due to numerical precision
        xyz_pos_l = np.zeros((num_dims, num_points), dtype=int)
        xyz_pos_r = np.zeros((num_dims, num_points), dtype=int)
        for dim in range(num_dims):
            xyz_pos_l[dim] = np.searchsorted(xyz_grid[dim] + tol[dim], points[:, dim])
            xyz_pos_r[dim] = np.searchsorted(xyz_grid[dim] - tol[dim], points[:, dim])

        # let's allocate an array for resulting values
        # every time we process a chunk of samples, we will write into this array
        interpolated_values = fill_value + np.zeros(
            [len(xyz_comp) for xyz_comp in xyz_grid], dtype=self.values.dtype
        )

        processed_cells_global = 0

        # to ovoid OOM for large datasets, we process only certain number of cells at a time
        while processed_cells_global < num_cells:
            target_processed_cells_global = min(
                num_cells, processed_cells_global + max_cells_per_step
            )

            connections_to_process = cell_connections[
                processed_cells_global:target_processed_cells_global
            ]

            # now we transfer this information to each cell. That is, each cell knows how its vertices
            # positioned relative to Cartesian grid points.
            # (num_dims, num_cells, num_vertices=num_cell_faces)
            xyz_pos_l_per_cell = xyz_pos_l[:, connections_to_process]
            xyz_pos_r_per_cell = xyz_pos_r[:, connections_to_process]

            # taking min/max among all cell vertices (per each dimension separately)
            # we get min and max indices of Cartesian grid points that may receive their values
            # from a given cell.
            # (num_dims, num_cells)
            cell_ind_min = np.min(xyz_pos_l_per_cell, axis=2)
            cell_ind_max = np.max(xyz_pos_r_per_cell, axis=2)

            # calculate number of Cartesian grid points where we will perform interpolation for a given
            # cell. Note that this number is much larger than actually needed, because essentially for
            # each cell we consider all Cartesian grid points that fall into the cell's bounding box.
            # We use word "sample" to represent such Cartesian grid points.
            # (num_cells,)
            num_samples_per_cell = np.prod(cell_ind_max - cell_ind_min, axis=0)

            # find cells that have non-zero number of samples
            # we use "ne" as a shortcut for "non empty"
            ne_cells = num_samples_per_cell > 0  # (num_cells,)
            num_ne_cells = np.sum(ne_cells)
            # indices of cells with non-zero number of samples in the original list of cells
            # (num_cells,)
            ne_cell_inds = np.arange(processed_cells_global, target_processed_cells_global)[
                ne_cells
            ]

            # restrict to non-empty cells only
            num_samples_per_ne_cell = num_samples_per_cell[ne_cells]
            cum_num_samples_per_ne_cell = np.cumsum(num_samples_per_ne_cell)

            ne_cell_ind_min = cell_ind_min[:, ne_cells]
            ne_cell_ind_max = cell_ind_max[:, ne_cells]

            # Next we need to perform actual interpolation at all sample points
            # this is computationally expensive operation and because we try to do everything
            # in the vectorized form, it can require a lot of memory, sometimes even causing OOM errors.
            # To avoid that, we impose restrictions on how many cells/samples can be processed at a time
            # effectivelly performing these operations in chunks.
            # Note that currently this is done sequentially, but could be relatively easy to parallelize

            # start counters of how many cells/samples have been processed
            processed_samples = 0
            processed_cells = 0

            while processed_cells < num_ne_cells:
                # how many cells we would like to process by the end of this step
                target_processed_cells = min(num_ne_cells, processed_cells + max_cells_per_step)

                # find how many cells we can processed based on number of allowed samples
                target_processed_samples = processed_samples + max_samples_per_step
                target_processed_cells_from_samples = (
                    np.searchsorted(cum_num_samples_per_ne_cell, target_processed_samples) + 1
                )

                # take min between the two
                target_processed_cells = min(
                    target_processed_cells, target_processed_cells_from_samples
                )

                # select cells and corresponding samples to process
                step_ne_cell_ind_min = ne_cell_ind_min[:, processed_cells:target_processed_cells]
                step_ne_cell_ind_max = ne_cell_ind_max[:, processed_cells:target_processed_cells]
                step_ne_cell_inds = ne_cell_inds[processed_cells:target_processed_cells]

                # process selected cells and points
                xyz_inds, interpolated = self._interp_py_chunk(
                    xyz_grid=xyz_grid,
                    cell_inds=step_ne_cell_inds,
                    cell_ind_min=step_ne_cell_ind_min,
                    cell_ind_max=step_ne_cell_ind_max,
                    sdf_tol=diag_tol,
                )

                if num_dims == 3:
                    interpolated_values[xyz_inds[0], xyz_inds[1], xyz_inds[2]] = interpolated
                else:
                    interpolated_values[xyz_inds[0], xyz_inds[1]] = interpolated

                processed_cells = target_processed_cells
                processed_samples = cum_num_samples_per_ne_cell[target_processed_cells - 1]

            processed_cells_global = target_processed_cells_global

        # in case of 2d grid broadcast results along normal direction assuming translational
        # invariance
        if num_dims == 2:
            orig_shape = [len(x), len(y), len(z)]
            flat_shape = orig_shape.copy()
            flat_shape[axis_ignore] = 1
            interpolated_values = np.reshape(interpolated_values, flat_shape)
            interpolated_values = np.broadcast_to(
                interpolated_values, (len(x), len(y), len(z))
            ).copy()

        return interpolated_values

    def _interp_py_chunk(
        self,
        xyz_grid: Tuple[ArrayLike[float], ...],
        cell_inds: ArrayLike[int],
        cell_ind_min: ArrayLike[int],
        cell_ind_max: ArrayLike[int],
        sdf_tol: float,
    ) -> Tuple[Tuple[ArrayLike, ...], ArrayLike]:
        """For each cell listed in ``cell_inds`` perform interpolation at a rectilinear subarray of
        xyz_grid given by a (3D) index span (cell_ind_min, cell_ind_max).

        Parameters
        ----------
        xyz_grid : Tuple[ArrayLike[float], ...]
            x, y, and z coordiantes defining rectilinear grid.
        cell_inds : ArrayLike[int]
            Indices of cells to perfrom interpolation from.
        cell_ind_min : ArrayLike[int]
            Starting x, y, and z indices of points for interpolation for each cell.
        cell_ind_max : ArrayLike[int]
            End x, y, and z indices of points for interpolation for each cell.
        sdf_tol : float
            Effective zero level set value, below which a point is considered to be inside a cell.

        Returns
        -------
        Tuple[Tuple[ArrayLike, ...], ArrayLike]
            x, y, and z indices of interpolated values and values themselves.
        """

        # get dimensionality of data
        num_dims = self._point_dims()
        num_cell_faces = self._cell_num_vertices()

        # get mesh info as numpy arrays
        points = self.points.values  # (num_points, num_dims)
        data_values = self.values.values  # (num_points,)
        cell_connections = self.cells.values[cell_inds]

        # compute number of samples to generate per cell
        num_samples_per_cell = np.prod(cell_ind_max - cell_ind_min, axis=0)

        # at this point we know how many samples we need to perform per each cell and we also
        # know span indices of these samples (in x, y, and z arrays)

        # we would like to perform all interpolations in a vectorized form, however, we have
        # a different number of interpolation samples for different cells. Thus, we need to
        # arange all samples in a linear way (flatten). Basically, we want to have data in this
        # form:
        # cell_ind | x_ind | y_ind | z_ind
        # --------------------------------
        #        0 |    23 |     5 |    11
        #        0 |    23 |     5 |    12
        #        0 |    23 |     6 |    11
        #        0 |    23 |     6 |    12
        #        1 |    41 |    11 |     0
        #        1 |    42 |    11 |     0
        #      ... |   ... |   ... |   ...

        # to do that we start with performing arange for each cell, but in vectorized way
        # this gives us something like this
        # [0, 1, 2, 3,   0, 1,   0, 1, 2, 3, 4, 5, 6,   ...]
        # |<-cell 0->|<-cell 1->|<-     cell 2    ->|<- ...

        num_cells = len(num_samples_per_cell)
        num_samples_cumul = num_samples_per_cell.cumsum()
        num_samples_total = num_samples_cumul[-1]

        # one big arange array
        inds_flat = np.arange(num_samples_total)
        # now subtract previous number of samples
        inds_flat[num_samples_per_cell[0] :] -= np.repeat(
            num_samples_cumul[:-1], num_samples_per_cell[1:]
        )

        # convert flat indices into 3d/2d indices as:
        # x_ind = [23, 23, 23, 23,   41, 41,      ...]
        # y_ind = [ 5,  5,  5,  5,    6,  6,      ...]
        # z_ind = [11, 12, 11, 12,    0,  0,      ...]
        #         |<-  cell 0  ->|<- cell 1 ->|<- ...
        num_samples_y = np.repeat(cell_ind_max[1] - cell_ind_min[1], num_samples_per_cell)

        # note: in 2d x, y correspond to (x, y, z).pop(normal_axis)
        if num_dims == 3:
            num_samples_z = np.repeat(cell_ind_max[2] - cell_ind_min[2], num_samples_per_cell)
            inds_flat, z_inds = np.divmod(inds_flat, num_samples_z)

        x_inds, y_inds = np.divmod(inds_flat, num_samples_y)

        start_inds = np.repeat(cell_ind_min, num_samples_per_cell, axis=1)
        x_inds = x_inds + start_inds[0]
        y_inds = y_inds + start_inds[1]
        if num_dims == 3:
            z_inds = z_inds + start_inds[2]

        # finally, we repeat cell indices corresponding number of times to obtain how
        # (x_ind, y_ind, z_ind) map to cell indices. So, now we have four arras:
        # x_ind    = [23, 23, 23, 23,   41, 41,      ...]
        # y_ind    = [ 5,  5,  5,  5,    6,  6,      ...]
        # z_ind    = [11, 12, 11, 12,    0,  0,      ...]
        # cell_map = [ 0,  0,  0,  0,    1,  1,      ...]
        #            |<-  cell 0  ->|<- cell 1 ->|<- ...
        step_cell_map = np.repeat(np.arange(num_cells), num_samples_per_cell)

        # let's put these arrays aside for a moment and perform the second preparatory step
        # specifically, for each face of each cell we will compute normal vector and distance
        # to the opposing cell vertex. This will allows us quickly calculate SDF of a cell at
        # each sample point as well as perform linear interpolation.

        # first, we collect coordinates of cell vertices into a single array
        # (num_cells, num_cell_vertices, num_dims)
        cell_vertices = np.float64(points[cell_connections, :])

        # array for resulting normals and distances
        normal = np.zeros((num_cell_faces, num_cells, num_dims))
        dist = np.zeros((num_cell_faces, num_cells))

        # loop face by face
        # note that by face_ind we denote both index of face in a cell and index of the opposing vertex
        for face_ind in range(num_cell_faces):
            # select vertices forming the given face
            face_pinds = list(np.arange(num_cell_faces))
            face_pinds.pop(face_ind)

            # calculate normal to the face
            # in 3D: cross product of two vectors lying in the face plane
            # in 2D: (-ty, tx) for a vector (tx, ty) along the face
            p0 = cell_vertices[:, face_pinds[0]]
            p01 = cell_vertices[:, face_pinds[1]] - p0
            p0Opp = cell_vertices[:, face_ind] - p0
            if num_dims == 3:
                p02 = cell_vertices[:, face_pinds[2]] - p0
                n = np.cross(p01, p02)
            else:
                n = np.roll(p01, 1, axis=1)
                n[:, 0] = -n[:, 0]
            n_norm = np.linalg.norm(n, axis=1)
            n = n / n_norm[:, None]

            # compute distance to the opposing vertex by taking a dot product between normal
            # and a vector connecting the opposing vertex and the face
            d = np.einsum("ij,ij->i", n, p0Opp)

            # obtained normal direction is arbitrary here. We will orient it such that it points
            # away from the triangle (and distance to the opposing vertex is negative).
            to_flip = d > 0
            d[to_flip] *= -1
            n[to_flip, :] *= -1

            # set distances in degenerate triangles to something positive to ignore later
            dist_zero = d == 0
            if any(dist_zero):
                d[dist_zero] = 1

            # record obtained info
            normal[face_ind] = n
            dist[face_ind] = d

        # now we all set up to proceed with actual interpolation at each sample point
        # the main idea here is that:
        # - we use `cell_map` to grab normals and distances
        #   of cells in which the given sample point is (potentially) located.
        # - use `x_ind, y_ind, z_ind` to find actual coordinates of a given sample point
        # - combine the above two to calculate cell SDF and interpolated value at a given sample
        #   point
        # - having cell SDF at the sample point actually tells us whether its inside the cell
        #   (keep value) or outside of it (discard interpolated value)

        # to perform SDF calculation and interpolation we will loop face by face and recording
        # their contributions. That is,
        # cell_sdf = max(face0_sdf, face1_sdf, ...)
        # interpolated_value = value0 * face0_sdf / dist0_sdf + ...
        # (because face0_sdf / dist0_sdf is linear shape function for vertex0)
        sdf = -inf * np.ones(num_samples_total)
        interpolated = np.zeros(num_samples_total, dtype=self._double_type)

        # coordinates of each sample point
        sample_xyz = np.zeros((num_samples_total, num_dims))
        sample_xyz[:, 0] = xyz_grid[0][x_inds]
        sample_xyz[:, 1] = xyz_grid[1][y_inds]
        if num_dims == 3:
            sample_xyz[:, 2] = xyz_grid[2][z_inds]

        # loop face by face
        for face_ind in range(num_cell_faces):
            # find a vector connecting sample point and face
            if face_ind == 0:
                vertex_ind = 1  # anythin other than 0
                vec = sample_xyz - cell_vertices[step_cell_map, vertex_ind, :]

            if face_ind == 1:  # since three faces share a point only do this once
                vertex_ind = 0  # it belongs to every face 1, 2, and 3
                vec = sample_xyz - cell_vertices[step_cell_map, 0, :]

            # compute distance from every sample point to the face of corresponding cell
            # using dot product
            tmp = normal[face_ind, step_cell_map, :] * vec
            d = np.sum(tmp, axis=1)

            # take max between distance to obtain the overall SDF of a cell
            sdf = np.maximum(sdf, d)

            # perform linear interpolation. Here we use the fact that when computing face SDF
            # at a given point and dividing it by the distance to the opposing vertex we get
            # a linear shape function for that vertex. So, we just need to multiply that by
            # the data value at that vertex to find its contribution into intepolated value.
            # (decomposed in an attempt to reduce memory consumption)
            tmp = self._double_type(data_values[cell_connections[step_cell_map, face_ind]])
            tmp *= d
            tmp /= dist[face_ind, step_cell_map]

            # ignore degenerate cells
            dist_zero = dist[face_ind, step_cell_map] > 0
            if any(dist_zero):
                sdf[dist_zero] = 10 * sdf_tol

            interpolated += tmp

        # The resulting array of interpolated values contain multiple candidate values for
        # every Cartesian point because bounding boxes of cells overlap.
        # Thus, we need to keep only those that come cell actually containing a given point.
        # This can be easily determined by the sign of the cell SDF sampled at a given point.
        valid_samples = sdf < sdf_tol

        interpolated_valid = interpolated[valid_samples]
        xyz_valid_inds = []
        xyz_valid_inds.append(x_inds[valid_samples])
        xyz_valid_inds.append(y_inds[valid_samples])
        if num_dims == 3:
            xyz_valid_inds.append(z_inds[valid_samples])

        return xyz_valid_inds, interpolated_valid

    @abstractmethod
    @requires_vtk
    def sel(
        self,
        x: Union[float, ArrayLike] = None,
        y: Union[float, ArrayLike] = None,
        z: Union[float, ArrayLike] = None,
    ) -> Union[UnstructuredGridDataset, SpatialDataArray]:
        """Extract/interpolate data along one or more Cartesian directions. At least of x, y, and z
        must be provided.

        Parameters
        ----------
        x : Union[float, ArrayLike] = None
            x-coordinate of the slice.
        y : Union[float, ArrayLike] = None
            y-coordinate of the slice.
        z : Union[float, ArrayLike] = None
            z-coordinate of the slice.

        Returns
        -------
        Union[TriangularGridDataset, SpatialDataArray]
            Extracted data.
        """

    @requires_vtk
    def sel_inside(self, bounds: Bound) -> UnstructuredGridDataset:
        """Return a new UnstructuredGridDataset that contains the minimal amount data necessary to
        cover a spatial region defined by ``bounds``.


        Parameters
        ----------
        bounds : Tuple[float, float, float], Tuple[float, float float]
            Min and max bounds packaged as ``(minx, miny, minz), (maxx, maxy, maxz)``.

        Returns
        -------
        UnstructuredGridDataset
            Extracted spatial data array.
        """
        if any(bmin > bmax for bmin, bmax in zip(*bounds)):
            raise DataError(
                "Min and max bounds must be packaged as ``(minx, miny, minz), (maxx, maxy, maxz)``."
            )

        data_bounds = self.bounds
        tol = 1e-6

        # For extracting cells covering target region we use vtk's filter that extract cells based
        # on provided implicit function. However, when we provide to it the implicit function of
        # the entire box, it has a couple of issues coming from the fact that the algorithm
        # eliminates every cells for which the implicit function has positive sign at all vertices.
        # As result, sometimes there are cells that despite overlaping with the target domain still
        # being eliminated. Two common cases:
        # - near corners of the target domain
        # - target domain is very thin
        # That's why we perform selection by sequentially eliminating cells on the outer side of
        # each of the 6 surfaces of the bounding box separately.
        tmp = self._vtk_obj
        for direction in range(2):
            for dim in range(3):
                sign = -1 + 2 * direction
                plane_pos = bounds[direction][dim]

                # Dealing with situation when target region does intersect with any cell:
                # in this case we shift target region so that it barely touches at least some
                # of cells
                if sign < 0 and plane_pos > data_bounds[1][dim] - tol:
                    plane_pos = data_bounds[1][dim] - tol
                if sign > 0 and plane_pos < data_bounds[0][dim] + tol:
                    plane_pos = data_bounds[0][dim] + tol

                # if all cells are on the inside side of the plane for a given surface
                # we don't need to check for intersection
                if plane_pos <= data_bounds[1][dim] and plane_pos >= data_bounds[0][dim]:
                    plane = vtk["mod"].vtkPlane()
                    center = [0, 0, 0]
                    normal = [0, 0, 0]
                    center[dim] = plane_pos
                    normal[dim] = sign
                    plane.SetOrigin(center)
                    plane.SetNormal(normal)
                    extractor = vtk["mod"].vtkExtractGeometry()
                    extractor.SetImplicitFunction(plane)
                    extractor.ExtractInsideOn()
                    extractor.ExtractBoundaryCellsOn()
                    extractor.SetInputData(tmp)
                    extractor.Update()
                    tmp = extractor.GetOutput()

        return self._from_vtk_obj(tmp, remove_degenerate_cells=True, remove_unused_points=True)

    def does_cover(self, bounds: Bound) -> bool:
        """Check whether data fully covers specified by ``bounds`` spatial region. If data contains
        only one point along a given direction, then it is assumed the data is constant along that
        direction and coverage is not checked.


        Parameters
        ----------
        bounds : Tuple[float, float, float], Tuple[float, float float]
            Min and max bounds packaged as ``(minx, miny, minz), (maxx, maxy, maxz)``.

        Returns
        -------
        bool
            Full cover check outcome.
        """

        return all(
            (dmin <= smin and dmax >= smax)
            for dmin, dmax, smin, smax in zip(self.bounds[0], self.bounds[1], bounds[0], bounds[1])
        )

    @requires_vtk
    def reflect(
        self, axis: Axis, center: float, reflection_only: bool = False
    ) -> UnstructuredGridDataset:
        """Reflect unstructured data across the plane define by parameters ``axis`` and ``center``.
        By default the original data is preserved, setting ``reflection_only`` to ``True`` will
        produce only deflected data.

        Parameters
        ----------
        axis : Literal[0, 1, 2]
            Normal direction of the reflection plane.
        center : float
            Location of the reflection plane along its normal direction.
        reflection_only : bool = False
            Return only reflected data.

        Returns
        -------
        UnstructuredGridDataset
            Data after reflextion is performed.
        """

        reflector = vtk["mod"].vtkReflectionFilter()
        reflector.SetPlane([reflector.USE_X, reflector.USE_Y, reflector.USE_Z][axis])
        reflector.SetCenter(center)
        reflector.SetCopyInput(not reflection_only)
        reflector.SetInputData(self._vtk_obj)
        reflector.Update()

        return self._from_vtk_obj(reflector.GetOutput())



[docs]
class TriangularGridDataset(UnstructuredGridDataset):
    """Dataset for storing triangular grid data. Data values are associated with the nodes of
    the grid.

    Note
    ----
    To use full functionality of unstructured datasets one must install ``vtk`` package (``pip
    install tidy3d[vtk]`` or ``pip install vtk``). Otherwise the functionality of unstructured
    datasets is limited to creation, writing to/loading from a file, and arithmetic manipulations.

    Example
    -------
    >>> tri_grid_points = PointDataArray(
    ...     [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]],
    ...     coords=dict(index=np.arange(4), axis=np.arange(2)),
    ... )
    >>>
    >>> tri_grid_cells = CellDataArray(
    ...     [[0, 1, 2], [1, 2, 3]],
    ...     coords=dict(cell_index=np.arange(2), vertex_index=np.arange(3)),
    ... )
    >>>
    >>> tri_grid_values = IndexedDataArray(
    ...     [1.0, 2.0, 3.0, 4.0], coords=dict(index=np.arange(4)),
    ... )
    >>>
    >>> tri_grid = TriangularGridDataset(
    ...     normal_axis=1,
    ...     normal_pos=0,
    ...     points=tri_grid_points,
    ...     cells=tri_grid_cells,
    ...     values=tri_grid_values,
    ... )
    """

    normal_axis: Axis = pd.Field(
        ...,
        title="Grid Axis",
        description="Orientation of the grid.",
    )

    normal_pos: float = pd.Field(
        ...,
        title="Position",
        description="Coordinate of the grid along the normal direction.",
    )

    @cached_property
    def bounds(self) -> Bound:
        """Grid bounds."""
        bounds_2d = super().bounds
        bounds_3d = self._points_2d_to_3d(bounds_2d)
        return tuple(bounds_3d[0]), tuple(bounds_3d[1])

    @classmethod
    def _point_dims(cls) -> pd.PositiveInt:
        """Dimensionality of stored grid point coordinates."""
        return 2

    def _points_2d_to_3d(self, pts: ArrayLike) -> ArrayLike:
        """Convert 2d points into 3d points."""
        return np.insert(pts, obj=self.normal_axis, values=self.normal_pos, axis=1)

    @cached_property
    def _points_3d_array(self) -> ArrayLike:
        """3D representation of grid points."""
        return self._points_2d_to_3d(self.points.data)

    @classmethod
    def _cell_num_vertices(cls) -> pd.PositiveInt:
        """Number of vertices in a cell."""
        return 3

    @classmethod
    @requires_vtk
    def _vtk_cell_type(cls):
        """VTK cell type to use in the VTK representation."""
        return vtk["mod"].VTK_TRIANGLE

    @classmethod
    @requires_vtk
    def _from_vtk_obj(
        cls,
        vtk_obj,
        field=None,
        remove_degenerate_cells: bool = False,
        remove_unused_points: bool = False,
    ):
        """Initialize from a vtkUnstructuredGrid instance."""

        # get points cells data from vtk object
        if isinstance(vtk_obj, vtk["mod"].vtkPolyData):
            cells_vtk = vtk_obj.GetPolys()
        elif isinstance(vtk_obj, vtk["mod"].vtkUnstructuredGrid):
            cells_vtk = vtk_obj.GetCells()

        cells_numpy = vtk["vtk_to_numpy"](cells_vtk.GetConnectivityArray())

        cell_offsets = vtk["vtk_to_numpy"](cells_vtk.GetOffsetsArray())
        if not np.all(np.diff(cell_offsets) == cls._cell_num_vertices()):
            raise DataError(
                "Only triangular 'vtkUnstructuredGrid' or 'vtkPolyData' can be converted into "
                "'TriangularGridDataset'."
            )

        points_numpy = vtk["vtk_to_numpy"](vtk_obj.GetPoints().GetData())

        # data values are read directly into Tidy3D array
        values = cls._get_values_from_vtk(vtk_obj, len(points_numpy), field)

        # detect zero size dimension
        bounds = np.max(points_numpy, axis=0) - np.min(points_numpy, axis=0)
        zero_dims = np.where(np.isclose(bounds, 0))[0]

        if len(zero_dims) != 1:
            raise DataError(
                f"Provided vtk grid does not represent a two dimensional grid. Found zero size dimensions are {zero_dims}."
            )

        normal_axis = zero_dims[0]
        normal_pos = points_numpy[0][normal_axis]
        tan_dims = [0, 1, 2]
        tan_dims.remove(normal_axis)

        # convert 3d coordinates into 2d
        points_2d_numpy = points_numpy[:, tan_dims]

        # create Tidy3D points and cells arrays
        num_cells = len(cells_numpy) // cls._cell_num_vertices()
        cells_numpy = np.reshape(cells_numpy, (num_cells, cls._cell_num_vertices()))

        cells = CellDataArray(
            cells_numpy,
            coords=dict(
                cell_index=np.arange(num_cells), vertex_index=np.arange(cls._cell_num_vertices())
            ),
        )

        points = PointDataArray(
            points_2d_numpy,
            coords=dict(index=np.arange(len(points_numpy)), axis=np.arange(cls._point_dims())),
        )

        if remove_degenerate_cells:
            cells = cls._remove_degenerate_cells(cells=cells)

        if remove_unused_points:
            points, values, cells = cls._remove_unused_points(
                points=points, values=values, cells=cells
            )

        return cls(
            normal_axis=normal_axis,
            normal_pos=normal_pos,
            points=points,
            cells=cells,
            values=values,
        )


[docs]
    @requires_vtk
    def plane_slice(self, axis: Axis, pos: float) -> SpatialDataArray:
        """Slice data with a plane and return the resulting line as a SpatialDataArray.

        Parameters
        ----------
        axis : Axis
            The normal direction of the slicing plane.
        pos : float
            Position of the slicing plane along its normal direction.

        Returns
        -------
        SpatialDataArray
            The resulting slice.
        """

        if axis == self.normal_axis:
            raise DataError(
                f"Triangular grid (normal: {self.normal_axis}) cannot be sliced by a parallel "
                "plane."
            )

        # perform slicing in vtk and get unprocessed points and values
        slice_vtk = self._plane_slice_raw(axis=axis, pos=pos)
        points_numpy = vtk["vtk_to_numpy"](slice_vtk.GetPoints().GetData())
        values = self._get_values_from_vtk(slice_vtk, len(points_numpy))

        # axis of the resulting line
        slice_axis = 3 - self.normal_axis - axis

        # sort found intersection in ascending order
        sorting = np.argsort(points_numpy[:, slice_axis], kind="mergesort")

        # assemble coords for SpatialDataArray
        coords = [None, None, None]
        coords[axis] = [pos]
        coords[self.normal_axis] = [self.normal_pos]
        coords[slice_axis] = points_numpy[sorting, slice_axis]
        coords_dict = dict(zip("xyz", coords))

        # reshape values from a 1d array into a 3d array
        new_shape = [1, 1, 1]
        new_shape[slice_axis] = len(values)
        values_reshaped = np.reshape(values.data[sorting], new_shape)

        return SpatialDataArray(values_reshaped, coords=coords_dict, name=self.values.name)


    @property
    def _triangulation_obj(self) -> Triangulation:
        """Matplotlib triangular representation of the grid to use in plotting."""
        return Triangulation(self.points[:, 0], self.points[:, 1], self.cells)


[docs]
    @equal_aspect
    @add_ax_if_none
    def plot(
        self,
        ax: Ax = None,
        field: bool = True,
        grid: bool = True,
        cbar: bool = True,
        cmap: str = "viridis",
        vmin: float = None,
        vmax: float = None,
        shading: Literal["gourand", "flat"] = "gouraud",
        cbar_kwargs: Dict = None,
        pcolor_kwargs: Dict = None,
    ) -> Ax:
        """Plot the data field and/or the unstructured grid.

        Parameters
        ----------
        ax : matplotlib.axes._subplots.Axes = None
            matplotlib axes to plot on, if not specified, one is created.
        field : bool = True
            Whether to plot the data field.
        grid : bool = True
            Whether to plot the unstructured grid.
        cbar : bool = True
            Display colorbar (only if ``field == True``).
        cmap : str = "viridis"
            Color map to use for plotting.
        vmin : float = None
            The lower bound of data range that the colormap covers. If ``None``, they are
            inferred from the data and other keyword arguments.
        vmax : float = None
            The upper bound of data range that the colormap covers. If ``None``, they are
            inferred from the data and other keyword arguments.
        shading : Literal["gourand", "flat"] = "gourand"
            Type of shading to use when plotting the data field.
        cbar_kwargs : Dict = {}
            Additional parameters passed to colorbar object.

        Returns
        -------
        matplotlib.axes._subplots.Axes
            The supplied or created matplotlib axes.
        """

        if cbar_kwargs is None:
            cbar_kwargs = {}
        if pcolor_kwargs is None:
            pcolor_kwargs = {}
        if not (field or grid):
            raise DataError("Nothing to plot ('field == False', 'grid == False').")

        # plot data field if requested
        if field:
            plot_obj = ax.tripcolor(
                self._triangulation_obj,
                self.values.data,
                shading=shading,
                cmap=cmap,
                vmin=vmin,
                vmax=vmax,
                **pcolor_kwargs,
            )

            if cbar:
                label_kwargs = {}
                if "label" not in cbar_kwargs:
                    label_kwargs["label"] = self.values.name
                plt.colorbar(plot_obj, **cbar_kwargs, **label_kwargs)

        # plot grid if requested
        if grid:
            ax.triplot(
                self._triangulation_obj,
                color=plot_params_grid.edgecolor,
                linewidth=plot_params_grid.linewidth,
            )

        # set labels and titles
        ax_labels = ["x", "y", "z"]
        normal_axis_name = ax_labels.pop(self.normal_axis)
        ax.set_xlabel(ax_labels[0])
        ax.set_ylabel(ax_labels[1])
        ax.set_title(f"{normal_axis_name} = {self.normal_pos}")
        return ax



[docs]
    def interp(
        self,
        x: Union[float, ArrayLike],
        y: Union[float, ArrayLike],
        z: Union[float, ArrayLike],
        fill_value: Union[float, Literal["extrapolate"]] = None,
        use_vtk: bool = False,
        method: Literal["linear", "nearest"] = "linear",
        ignore_normal_pos: bool = True,
        max_samples_per_step: int = DEFAULT_MAX_SAMPLES_PER_STEP,
        max_cells_per_step: int = DEFAULT_MAX_CELLS_PER_STEP,
        rel_tol: float = DEFAULT_TOLERANCE_CELL_FINDING,
    ) -> SpatialDataArray:
        """Interpolate data at provided x, y, and z. Note that data is assumed to be invariant along
        the dataset's normal direction.

        Parameters
        ----------
        x : Union[float, ArrayLike]
            x-coordinates of sampling points.
        y : Union[float, ArrayLike]
            y-coordinates of sampling points.
        z : Union[float, ArrayLike]
            z-coordinates of sampling points.
        fill_value : Union[float, Literal["extrapolate"]] = 0
            Value to use when filling points without interpolated values. If ``"extrapolate"`` then
            nearest values are used. Note: in a future version the default value will be changed
            to ``"extrapolate"``.
        use_vtk : bool = False
            Use vtk's interpolation functionality or Tidy3D's own implementation. Note: this
            option will be removed in a future version.
        method: Literal["linear", "nearest"] = "linear"
            Interpolation method to use.
        ignore_normal_pos : bool = True
            (Depreciated) Assume data is invariant along the normal direction to the grid plane.
        max_samples_per_step : int = 1e4
            Max number of points to interpolate at per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        max_cells_per_step : int = 1e4
            Max number of cells to interpolate from per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        rel_tol : float = 1e-6
            Relative tolerance when determining whether a point belongs to a cell.

        Returns
        -------
        SpatialDataArray
            Interpolated data.
        """

        if fill_value is None:
            log.warning(
                "Default parameter setting 'fill_value=0' will be changed to "
                "'fill_value=``extrapolate``' in a future version."
            )
            fill_value = 0

        if not ignore_normal_pos:
            log.warning(
                "Parameter 'ignore_normal_pos' is depreciated. It is always assumed that data "
                "contained in 'TriangularGridDataset' is invariant in the normal direction. "
                "That is, 'ignore_normal_pos=True' is used."
            )

        x = np.atleast_1d(x)
        y = np.atleast_1d(y)
        z = np.atleast_1d(z)

        xyz = [x, y, z]
        xyz[self.normal_axis] = [self.normal_pos]
        interp_inplane = super().interp(
            **dict(zip("xyz", xyz)),
            fill_value=fill_value,
            use_vtk=use_vtk,
            method=method,
            max_samples_per_step=max_samples_per_step,
            max_cells_per_step=max_cells_per_step,
        )
        interp_broadcasted = np.broadcast_to(
            interp_inplane, [len(np.atleast_1d(comp)) for comp in [x, y, z]]
        )

        return SpatialDataArray(
            interp_broadcasted, coords=dict(x=x, y=y, z=z), name=self.values.name
        )


    def _interp_py(
        self,
        x: ArrayLike,
        y: ArrayLike,
        z: ArrayLike,
        fill_value: float,
        max_samples_per_step: int,
        max_cells_per_step: int,
        rel_tol: float,
    ) -> ArrayLike:
        """2D-specific function to interpolate data at provided x, y, and z
        using vectorized python implementation.

        Parameters
        ----------
        x : Union[float, ArrayLike]
            x-coordinates of sampling points.
        y : Union[float, ArrayLike]
            y-coordinates of sampling points.
        z : Union[float, ArrayLike]
            z-coordinates of sampling points.
        fill_value : float
            Value to use when filling points without interpolated values.
        max_samples_per_step : int
            Max number of points to interpolate at per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        max_cells_per_step : int
            Max number of cells to interpolate from per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        rel_tol : float
            Relative tolerance when determining whether a point belongs to a cell.

        Returns
        -------
        ArrayLike
            Interpolated data.
        """

        return self._interp_py_general(
            x=x,
            y=y,
            z=z,
            fill_value=fill_value,
            max_samples_per_step=max_samples_per_step,
            max_cells_per_step=max_cells_per_step,
            rel_tol=rel_tol,
            axis_ignore=self.normal_axis,
        )


[docs]
    @requires_vtk
    def sel(
        self,
        x: Union[float, ArrayLike] = None,
        y: Union[float, ArrayLike] = None,
        z: Union[float, ArrayLike] = None,
    ) -> SpatialDataArray:
        """Extract/interpolate data along one or more Cartesian directions. At least of x, y, and z
        must be provided.

        Parameters
        ----------
        x : Union[float, ArrayLike] = None
            x-coordinate of the slice.
        y : Union[float, ArrayLike] = None
            y-coordinate of the slice.
        z : Union[float, ArrayLike] = None
            z-coordinate of the slice.

        Returns
        -------
        SpatialDataArray
            Extracted data.
        """

        xyz = [x, y, z]
        axes = [ind for ind, comp in enumerate(xyz) if comp is not None]
        num_provided = len(axes)

        if self.normal_axis in axes:
            if xyz[self.normal_axis] != self.normal_pos:
                raise DataError(
                    f"No data for {'xyz'[self.normal_axis]} = {xyz[self.normal_axis]} (unstructured"
                    f" grid is defined at {'xyz'[self.normal_axis]} = {self.normal_pos})."
                )

            if num_provided < 3:
                num_provided -= 1
                axes.remove(self.normal_axis)

        if num_provided == 0:
            raise DataError("At least one of 'x', 'y', and 'z' must be specified.")

        if num_provided == 1:
            axis = axes[0]
            return self.plane_slice(axis=axis, pos=xyz[axis])

        if num_provided == 2:
            pos = [x, y, z]
            pos[self.normal_axis] = [self.normal_pos]
            return self.interp(x=pos[0], y=pos[1], z=pos[2])

        if num_provided == 3:
            return self.interp(x=x, y=y, z=z)



[docs]
    @requires_vtk
    def reflect(
        self, axis: Axis, center: float, reflection_only: bool = False
    ) -> UnstructuredGridDataset:
        """Reflect unstructured data across the plane define by parameters ``axis`` and ``center``.
        By default the original data is preserved, setting ``reflection_only`` to ``True`` will
        produce only deflected data.

        Parameters
        ----------
        axis : Literal[0, 1, 2]
            Normal direction of the reflection plane.
        center : float
            Location of the reflection plane along its normal direction.
        reflection_only : bool = False
            Return only reflected data.

        Returns
        -------
        UnstructuredGridDataset
            Data after reflextion is performed.
        """

        # disallow reflecting along normal direction
        if axis == self.normal_axis:
            if reflection_only:
                return self.updated_copy(normal_pos=2 * center - self.normal_pos)
            else:
                raise DataError(
                    "Reflection in the normal direction to the grid is prohibited unless 'reflection_only=True'."
                )

        return super().reflect(axis=axis, center=center, reflection_only=reflection_only)



[docs]
    @requires_vtk
    def sel_inside(self, bounds: Bound) -> TriangularGridDataset:
        """Return a new ``TriangularGridDataset`` that contains the minimal amount data necessary to
        cover a spatial region defined by ``bounds``.


        Parameters
        ----------
        bounds : Tuple[float, float, float], Tuple[float, float float]
            Min and max bounds packaged as ``(minx, miny, minz), (maxx, maxy, maxz)``.

        Returns
        -------
        TriangularGridDataset
            Extracted spatial data array.
        """
        if any(bmin > bmax for bmin, bmax in zip(*bounds)):
            raise DataError(
                "Min and max bounds must be packaged as ``(minx, miny, minz), (maxx, maxy, maxz)``."
            )

        # expand along normal direction
        new_bounds = [list(bounds[0]), list(bounds[1])]

        new_bounds[0][self.normal_axis] = -inf
        new_bounds[1][self.normal_axis] = inf

        return super().sel_inside(new_bounds)



[docs]
    def does_cover(self, bounds: Bound) -> bool:
        """Check whether data fully covers specified by ``bounds`` spatial region. If data contains
        only one point along a given direction, then it is assumed the data is constant along that
        direction and coverage is not checked.


        Parameters
        ----------
        bounds : Tuple[float, float, float], Tuple[float, float float]
            Min and max bounds packaged as ``(minx, miny, minz), (maxx, maxy, maxz)``.

        Returns
        -------
        bool
            Full cover check outcome.
        """

        # expand along normal direction
        new_bounds = [list(bounds[0]), list(bounds[1])]

        new_bounds[0][self.normal_axis] = self.normal_pos
        new_bounds[1][self.normal_axis] = self.normal_pos

        return super().does_cover(new_bounds)





[docs]
class TetrahedralGridDataset(UnstructuredGridDataset):
    """Dataset for storing tetrahedral grid data. Data values are associated with the nodes of
    the grid.

    Note
    ----
    To use full functionality of unstructured datasets one must install ``vtk`` package (``pip
    install tidy3d[vtk]`` or ``pip install vtk``). Otherwise the functionality of unstructured
    datasets is limited to creation, writing to/loading from a file, and arithmetic manipulations.

    Example
    -------
    >>> tet_grid_points = PointDataArray(
    ...     [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
    ...     coords=dict(index=np.arange(4), axis=np.arange(3)),
    ... )
    >>>
    >>> tet_grid_cells = CellDataArray(
    ...     [[0, 1, 2, 3]],
    ...     coords=dict(cell_index=np.arange(1), vertex_index=np.arange(4)),
    ... )
    >>>
    >>> tet_grid_values = IndexedDataArray(
    ...     [1.0, 2.0, 3.0, 4.0], coords=dict(index=np.arange(4)),
    ... )
    >>>
    >>> tet_grid = TetrahedralGridDataset(
    ...     points=tet_grid_points,
    ...     cells=tet_grid_cells,
    ...     values=tet_grid_values,
    ... )
    """

    @classmethod
    def _point_dims(cls) -> pd.PositiveInt:
        """Dimensionality of stored grid point coordinates."""
        return 3

    @cached_property
    def _points_3d_array(self) -> Bound:
        """3D coordinates of grid points."""
        return self.points.data

    @classmethod
    def _cell_num_vertices(cls) -> pd.PositiveInt:
        """Number of vertices in a cell."""
        return 4

    @classmethod
    @requires_vtk
    def _vtk_cell_type(cls):
        """VTK cell type to use in the VTK representation."""
        return vtk["mod"].VTK_TETRA

    @classmethod
    @requires_vtk
    def _from_vtk_obj(
        cls,
        grid,
        field=None,
        remove_degenerate_cells: bool = False,
        remove_unused_points: bool = False,
    ) -> TetrahedralGridDataset:
        """Initialize from a vtkUnstructuredGrid instance."""

        # read point, cells, and values info from a vtk instance
        cells_numpy = vtk["vtk_to_numpy"](grid.GetCells().GetConnectivityArray())
        points_numpy = vtk["vtk_to_numpy"](grid.GetPoints().GetData())
        values = cls._get_values_from_vtk(grid, len(points_numpy), field)

        # verify cell_types
        cells_types = vtk["vtk_to_numpy"](grid.GetCellTypesArray())
        if not np.all(cells_types == cls._vtk_cell_type()):
            raise DataError("Only tetrahedral 'vtkUnstructuredGrid' is currently supported")

        # pack point and cell information into Tidy3D arrays
        num_cells = len(cells_numpy) // cls._cell_num_vertices()
        cells_numpy = np.reshape(cells_numpy, (num_cells, cls._cell_num_vertices()))

        cells = CellDataArray(
            cells_numpy,
            coords=dict(
                cell_index=np.arange(num_cells), vertex_index=np.arange(cls._cell_num_vertices())
            ),
        )

        points = PointDataArray(
            points_numpy,
            coords=dict(index=np.arange(len(points_numpy)), axis=np.arange(cls._point_dims())),
        )

        if remove_degenerate_cells:
            cells = cls._remove_degenerate_cells(cells=cells)

        if remove_unused_points:
            points, values, cells = cls._remove_unused_points(
                points=points, values=values, cells=cells
            )

        return cls(points=points, cells=cells, values=values)


[docs]
    @requires_vtk
    def plane_slice(self, axis: Axis, pos: float) -> TriangularGridDataset:
        """Slice data with a plane and return the resulting :class:.`TriangularGridDataset`.

        Parameters
        ----------
        axis : Axis
            The normal direction of the slicing plane.
        pos : float
            Position of the slicing plane along its normal direction.

        Returns
        -------
        TriangularGridDataset
            The resulting slice.
        """

        slice_vtk = self._plane_slice_raw(axis=axis, pos=pos)

        return TriangularGridDataset._from_vtk_obj(
            slice_vtk, remove_degenerate_cells=True, remove_unused_points=True
        )



[docs]
    @requires_vtk
    def line_slice(self, axis: Axis, pos: Coordinate) -> SpatialDataArray:
        """Slice data with a line and return the resulting :class:.`SpatialDataArray`.

        Parameters
        ----------
        axis : Axis
            The axis of the slicing line.
        pos : Tuple[float, float, float]
            Position of the slicing line.

        Returns
        -------
        SpatialDataArray
            The resulting slice.
        """

        bounds = self.bounds
        start = list(pos)
        end = list(pos)

        start[axis] = bounds[0][axis]
        end[axis] = bounds[1][axis]

        # create cutting plane
        line = vtk["mod"].vtkLineSource()
        line.SetPoint1(start)
        line.SetPoint2(end)
        line.SetResolution(1)

        # this should be done using vtkProbeLineFilter
        # but for some reason it crashes Python
        # so, we use a workaround:
        # 1) extract cells that are intersected by line (to speed up further slicing)
        # 2) do plane slice along first direction
        # 3) do second plane slice along second direction

        prober = vtk["mod"].vtkExtractCellsAlongPolyLine()
        prober.SetSourceConnection(line.GetOutputPort())
        prober.SetInputData(self._vtk_obj)
        prober.Update()

        extracted_cells_vtk = prober.GetOutput()

        if extracted_cells_vtk.GetNumberOfPoints() == 0:
            raise DataError("Slicing line does not intersect the unstructured grid.")

        extracted_cells = TetrahedralGridDataset._from_vtk_obj(
            extracted_cells_vtk, remove_degenerate_cells=True, remove_unused_points=True
        )

        tan_dims = [0, 1, 2]
        tan_dims.remove(axis)

        # first plane slice
        plane_slice = extracted_cells.plane_slice(axis=tan_dims[0], pos=pos[tan_dims[0]])
        # second plane slice
        line_slice = plane_slice.plane_slice(axis=tan_dims[1], pos=pos[tan_dims[1]])

        return line_slice



[docs]
    @requires_vtk
    def sel(
        self,
        x: Union[float, ArrayLike] = None,
        y: Union[float, ArrayLike] = None,
        z: Union[float, ArrayLike] = None,
    ) -> Union[TriangularGridDataset, SpatialDataArray]:
        """Extract/interpolate data along one or more Cartesian directions. At least of x, y, and z
        must be provided.

        Parameters
        ----------
        x : Union[float, ArrayLike] = None
            x-coordinate of the slice.
        y : Union[float, ArrayLike] = None
            y-coordinate of the slice.
        z : Union[float, ArrayLike] = None
            z-coordinate of the slice.

        Returns
        -------
        Union[TriangularGridDataset, SpatialDataArray]
            Extracted data.
        """

        xyz = [x, y, z]
        axes = [ind for ind, comp in enumerate(xyz) if comp is not None]

        num_provided = len(axes)

        if num_provided < 3 and any(not np.isscalar(comp) for comp in xyz if comp is not None):
            raise DataError(
                "Providing x, y, or z as array is only allowed for interpolation. That is, when all"
                " three x, y, and z are provided or method '.interp()' is used explicitly."
            )

        if num_provided == 0:
            raise DataError("At least one of 'x', 'y', and 'z' must be specified.")

        if num_provided == 1:
            axis = axes[0]
            return self.plane_slice(axis=axis, pos=xyz[axis])

        if num_provided == 2:
            axis = 3 - axes[0] - axes[1]
            xyz[axis] = 0
            return self.line_slice(axis=axis, pos=xyz)

        if num_provided == 3:
            return self.interp(x=x, y=y, z=z)


    def _interp_py(
        self,
        x: ArrayLike,
        y: ArrayLike,
        z: ArrayLike,
        fill_value: float,
        max_samples_per_step: int,
        max_cells_per_step: int,
        rel_tol: float,
    ) -> ArrayLike:
        """3D-specific function to interpolate data at provided x, y, and z
        using vectorized python implementation.

        Parameters
        ----------
        x : Union[float, ArrayLike]
            x-coordinates of sampling points.
        y : Union[float, ArrayLike]
            y-coordinates of sampling points.
        z : Union[float, ArrayLike]
            z-coordinates of sampling points.
        fill_value : float
            Value to use when filling points without interpolated values.
        max_samples_per_step : int
            Max number of points to interpolate at per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        max_cells_per_step : int
            Max number of cells to interpolate from per iteration (used only if `use_vtk=False`).
            Using a higher number may speed up calculations but, at the same time, it increases
            RAM usage.
        rel_tol : float
            Relative tolerance when determining whether a point belongs to a cell.

        Returns
        -------
        ArrayLike
            Interpolated data.
        """

        return self._interp_py_general(
            x=x,
            y=y,
            z=z,
            fill_value=fill_value,
            max_samples_per_step=max_samples_per_step,
            max_cells_per_step=max_cells_per_step,
            rel_tol=rel_tol,
            axis_ignore=None,
        )



UnstructuredGridDatasetType = Union[TriangularGridDataset, TetrahedralGridDataset]
CustomSpatialDataType = Union[SpatialDataArray, UnstructuredGridDatasetType]
CustomSpatialDataTypeAnnotated = Union[SpatialDataArray, annotate_type(UnstructuredGridDatasetType)]


def _get_numpy_array(data_array: Union[ArrayLike, DataArray, UnstructuredGridDataset]) -> ArrayLike:
    """Get numpy representation of dataarray/dataset values."""
    if isinstance(data_array, UnstructuredGridDataset):
        return data_array.values.values
    if isinstance(data_array, xr.DataArray):
        return data_array.values
    return np.array(data_array)


def _zeros_like(
    data_array: Union[ArrayLike, xr.DataArray, UnstructuredGridDataset]
) -> Union[ArrayLike, xr.DataArray, UnstructuredGridDataset]:
    """Get a zeroed replica of dataarray/dataset."""
    if isinstance(data_array, UnstructuredGridDataset):
        return data_array.updated_copy(values=xr.zeros_like(data_array.values))
    if isinstance(data_array, xr.DataArray):
        return xr.zeros_like(data_array)
    return np.zeros_like(data_array)


def _ones_like(
    data_array: Union[ArrayLike, xr.DataArray, UnstructuredGridDataset]
) -> Union[ArrayLike, xr.DataArray, UnstructuredGridDataset]:
    """Get a unity replica of dataarray/dataset."""
    if isinstance(data_array, UnstructuredGridDataset):
        return data_array.updated_copy(values=xr.ones_like(data_array.values))
    if isinstance(data_array, xr.DataArray):
        return xr.ones_like(data_array)
    return np.ones_like(data_array)


def _check_same_coordinates(
    a: Union[ArrayLike, xr.DataArray, UnstructuredGridDataset],
    b: Union[ArrayLike, xr.DataArray, UnstructuredGridDataset],
) -> bool:
    """Check whether two array are defined at the same coordinates."""

    # we can have xarray.DataArray's of different types but still same coordinates
    # we will deal with that case separately
    both_xarrays = isinstance(a, xr.DataArray) and isinstance(b, xr.DataArray)
    if (not both_xarrays) and type(a) != type(b):
        return False

    if isinstance(a, UnstructuredGridDataset):
        if not np.allclose(a.points, b.points) or not np.all(a.cells == b.cells):
            return False

        if isinstance(a, TriangularGridDataset):
            if a.normal_axis != b.normal_axis or a.normal_pos != b.normal_pos:
                return False

    elif isinstance(a, xr.DataArray):
        if a.coords.keys() != b.coords.keys() or a.coords != b.coords:
            return False

    else:
        if np.shape(a) != np.shape(b):
            return False

    return True