Source code for tidy3d.components.grid.grid

"""Defines the FDTD grid."""
from __future__ import annotations
from typing import Tuple, List, Union

import numpy as np
import pydantic.v1 as pd

from ..base import Tidy3dBaseModel
from ..data.data_array import DataArray, SpatialDataArray, ScalarFieldDataArray
from ..data.dataset import UnstructuredGridDatasetType, UnstructuredGridDataset
from ..types import ArrayFloat1D, Axis, InterpMethod, Literal
from ..geometry.base import Box

from ...exceptions import SetupError

# data type of one dimensional coordinate array.
Coords1D = ArrayFloat1D


[docs] class Coords(Tidy3dBaseModel): """Holds data about a set of x,y,z positions on a grid. Example ------- >>> x = np.linspace(-1, 1, 10) >>> y = np.linspace(-1, 1, 11) >>> z = np.linspace(-1, 1, 12) >>> coords = Coords(x=x, y=y, z=z) """ x: Coords1D = pd.Field( ..., title="X Coordinates", description="1-dimensional array of x coordinates." ) y: Coords1D = pd.Field( ..., title="Y Coordinates", description="1-dimensional array of y coordinates." ) z: Coords1D = pd.Field( ..., title="Z Coordinates", description="1-dimensional array of z coordinates." ) @property def to_dict(self): """Return a dict of the three Coord1D objects as numpy arrays.""" return {key: self.dict()[key] for key in "xyz"} @property def to_list(self): """Return a list of the three Coord1D objects as numpy arrays.""" return list(self.to_dict.values()) def _interp_from_xarray( self, array: Union[SpatialDataArray, ScalarFieldDataArray], interp_method: InterpMethod, fill_value: Union[Literal["extrapolate"], float] = "extrapolate", ) -> Union[SpatialDataArray, ScalarFieldDataArray]: """ Similar to ``xarrray.DataArray.interp`` with 2 enhancements: 1) Check if the coordinate of the supplied data are in monotonically increasing order. If they are, apply the faster ``assume_sorted=True``. 2) For axes of single entry, instead of error, apply ``isel()`` along the axis. Parameters ---------- array : Union[:class:`.SpatialDataArray`, :class:`.ScalarFieldDataArray`] Supplied scalar dataset interp_method : :class:`.InterpMethod` Interpolation method. fill_value : Union[Literal['extrapolate'], float] = "extrapolate" Value used to fill in for points outside the data range. If set to 'extrapolate', values will be extrapolated into those regions using the "nearest" method. Returns ------- Union[:class:`.SpatialDataArray`, :class:`.ScalarFieldDataArray`] The interpolated spatial dataset. Note ---- This method is called from a :class:`Coords` instance with the array to be interpolated as an argument, not the other way around. """ # Check which axes need interpolation or selection interp_ax = [] isel_ax = [] for ax in "xyz": if array.sizes[ax] == 1: isel_ax.append(ax) else: interp_ax.append(ax) # apply iselection for the axis containing single entry if len(isel_ax) > 0: array = array.isel({ax: [0] * len(self.to_dict[ax]) for ax in isel_ax}) array = array.assign_coords({ax: self.to_dict[ax] for ax in isel_ax}) if len(interp_ax) == 0: return array # Apply interp for the rest is_sorted = all(np.all(np.diff(array.coords[f]) > 0) for f in interp_ax) interp_param = { "method": interp_method, "assume_sorted": is_sorted, "kwargs": { "bounds_error": False, "fill_value": fill_value, }, } # Mark extrapolated points with nan's to fill in later if fill_value == "extrapolate" and interp_method != "nearest": interp_param["kwargs"]["fill_value"] = np.nan # interpolation interp_array = array.interp({ax: self.to_dict[ax] for ax in interp_ax}, **interp_param) # Fill in nan's with nearest values if fill_value == "extrapolate" and interp_method != "nearest": interp_param["method"] = "nearest" interp_param["kwargs"]["fill_value"] = "extrapolate" nearest_array = array.interp({ax: self.to_dict[ax] for ax in interp_ax}, **interp_param) interp_array.values[:] = np.where( np.isnan(interp_array.values), nearest_array.values, interp_array.values ) return interp_array def _interp_from_unstructured( self, array: UnstructuredGridDatasetType, interp_method: InterpMethod, fill_value: Union[Literal["extrapolate"], float] = "extrapolate", ) -> SpatialDataArray: """ Interpolate from untructured grid onto a Cartesian one. Parameters ---------- array : Union[class:`.TriangularGridDataset`, class:`.TetrahedralGridDataset`] Supplied scalar dataset interp_method : :class:`.InterpMethod` Interpolation method. fill_value : Union[Literal['extrapolate'], float] = "extrapolate" Value used to fill in for points outside the data range. If set to 'extrapolate', values will be extrapolated into those regions using the "nearest" method. Returns ------- :class:`.SpatialDataArray` The interpolated spatial dataset. Note ---- This method is called from a :class:`Coords` instance with the array to be interpolated as an argument, not the other way around. """ interp_array = array.interp( **{ax: self.to_dict[ax] for ax in "xyz"}, method=interp_method, fill_value=fill_value ) return interp_array
[docs] def spatial_interp( self, array: Union[SpatialDataArray, ScalarFieldDataArray, UnstructuredGridDatasetType], interp_method: InterpMethod, fill_value: Union[Literal["extrapolate"], float] = "extrapolate", ) -> Union[SpatialDataArray, ScalarFieldDataArray]: """ Similar to ``xarrray.DataArray.interp`` with 2 enhancements: 1) (if input data is an ``xarrray.DataArray``) Check if the coordinate of the supplied data are in monotonically increasing order. If they are, apply the faster ``assume_sorted=True``. 2) Data is assumed invariant along zero-size dimensions (if any). Parameters ---------- array : Union[ :class:`.SpatialDataArray`, :class:`.ScalarFieldDataArray`, :class:`.TriangularGridDataset`, :class:`.TetrahedralGridDataset`, ] Supplied scalar dataset interp_method : :class:`.InterpMethod` Interpolation method. fill_value : Union[Literal['extrapolate'], float] = "extrapolate" Value used to fill in for points outside the data range. If set to 'extrapolate', values will be extrapolated into those regions using the "nearest" method. Returns ------- Union[:class:`.SpatialDataArray`, :class:`.ScalarFieldDataArray`] The interpolated spatial dataset. Note ---- This method is called from a :class:`Coords` instance with the array to be interpolated as an argument, not the other way around. """ # Check for empty dimensions result_coords = dict(self.to_dict) if any(len(v) == 0 for v in result_coords.values()): if isinstance(array, (SpatialDataArray, ScalarFieldDataArray)): for c in array.coords: if c not in result_coords: result_coords[c] = array.coords[c].values result_shape = tuple(len(v) for v in result_coords.values()) result = DataArray(np.empty(result_shape, dtype=array.dtype), coords=result_coords) return result # interpolation if isinstance(array, UnstructuredGridDataset): return self._interp_from_unstructured( array=array, interp_method=interp_method, fill_value=fill_value ) else: return self._interp_from_xarray( array=array, interp_method=interp_method, fill_value=fill_value )
[docs] class FieldGrid(Tidy3dBaseModel): """Holds the grid data for a single field. Example ------- >>> x = np.linspace(-1, 1, 10) >>> y = np.linspace(-1, 1, 11) >>> z = np.linspace(-1, 1, 12) >>> coords = Coords(x=x, y=y, z=z) >>> field_grid = FieldGrid(x=coords, y=coords, z=coords) """ x: Coords = pd.Field( ..., title="X Positions", description="x,y,z coordinates of the locations of the x-component of a vector field.", ) y: Coords = pd.Field( ..., title="Y Positions", description="x,y,z coordinates of the locations of the y-component of a vector field.", ) z: Coords = pd.Field( ..., title="Z Positions", description="x,y,z coordinates of the locations of the z-component of a vector field.", )
[docs] class YeeGrid(Tidy3dBaseModel): """Holds the yee grid coordinates for each of the E and H positions. Example ------- >>> x = np.linspace(-1, 1, 10) >>> y = np.linspace(-1, 1, 11) >>> z = np.linspace(-1, 1, 12) >>> coords = Coords(x=x, y=y, z=z) >>> field_grid = FieldGrid(x=coords, y=coords, z=coords) >>> yee_grid = YeeGrid(E=field_grid, H=field_grid) >>> Ex_coords = yee_grid.E.x """ E: FieldGrid = pd.Field( ..., title="Electric Field Grid", description="Coordinates of the locations of all three components of the electric field.", ) H: FieldGrid = pd.Field( ..., title="Electric Field Grid", description="Coordinates of the locations of all three components of the magnetic field.", ) @property def grid_dict(self): """The Yee grid coordinates associated to various field components as a dictionary.""" return { "Ex": self.E.x, "Ey": self.E.y, "Ez": self.E.z, "Hx": self.H.x, "Hy": self.H.y, "Hz": self.H.z, }
[docs] class Grid(Tidy3dBaseModel): """Contains all information about the spatial positions of the FDTD grid. Example ------- >>> x = np.linspace(-1, 1, 10) >>> y = np.linspace(-1, 1, 11) >>> z = np.linspace(-1, 1, 12) >>> coords = Coords(x=x, y=y, z=z) >>> grid = Grid(boundaries=coords) >>> centers = grid.centers >>> sizes = grid.sizes >>> yee_grid = grid.yee """ boundaries: Coords = pd.Field( ..., title="Boundary Coordinates", description="x,y,z coordinates of the boundaries between cells, defining the FDTD grid.", ) @staticmethod def _avg(coords1d: Coords1D): """Return average positions of an array of 1D coordinates.""" return (coords1d[1:] + coords1d[:-1]) / 2.0 @staticmethod def _min(coords1d: Coords1D): """Return minus positions of 1D coordinates.""" return coords1d[:-1] @property def centers(self) -> Coords: """Return centers of the cells in the :class:`Grid`. Returns ------- :class:`Coords` centers of the FDTD cells in x,y,z stored as :class:`Coords` object. Example ------- >>> x = np.linspace(-1, 1, 10) >>> y = np.linspace(-1, 1, 11) >>> z = np.linspace(-1, 1, 12) >>> coords = Coords(x=x, y=y, z=z) >>> grid = Grid(boundaries=coords) >>> centers = grid.centers """ return Coords(**{key: self._avg(val) for key, val in self.boundaries.to_dict.items()}) @property def sizes(self) -> Coords: """Return sizes of the cells in the :class:`Grid`. Returns ------- :class:`Coords` Sizes of the FDTD cells in x,y,z stored as :class:`Coords` object. Example ------- >>> x = np.linspace(-1, 1, 10) >>> y = np.linspace(-1, 1, 11) >>> z = np.linspace(-1, 1, 12) >>> coords = Coords(x=x, y=y, z=z) >>> grid = Grid(boundaries=coords) >>> sizes = grid.sizes """ return Coords(**{key: np.diff(val) for key, val in self.boundaries.to_dict.items()}) @property def num_cells(self) -> Tuple[int, int, int]: """Return sizes of the cells in the :class:`Grid`. Returns ------- tuple[int, int, int] Number of cells in the grid in the x, y, z direction. Example ------- >>> x = np.linspace(-1, 1, 10) >>> y = np.linspace(-1, 1, 11) >>> z = np.linspace(-1, 1, 12) >>> coords = Coords(x=x, y=y, z=z) >>> grid = Grid(boundaries=coords) >>> Nx, Ny, Nz = grid.num_cells """ return [len(self.boundaries.dict()[dim]) - 1 for dim in "xyz"] @property def _primal_steps(self) -> Coords: """Return primal steps of the cells in the :class:`Grid`. Returns ------- :class:`Coords` Distances between each of the cell boundaries along each dimension. """ return self.sizes @property def _dual_steps(self) -> Coords: """Return dual steps of the cells in the :class:`Grid`. Returns ------- :class:`Coords` Distances between each of the cell centers along each dimension, with periodicity applied. """ primal_steps = {dim: self._primal_steps.dict()[dim] for dim in "xyz"} dsteps = {key: (psteps + np.roll(psteps, 1)) / 2 for (key, psteps) in primal_steps.items()} return Coords(**dsteps) @property def yee(self) -> YeeGrid: """Return the :class:`YeeGrid` defining the yee cell locations for this :class:`Grid`. Returns ------- :class:`YeeGrid` Stores coordinates of all of the components on the yee lattice. Example ------- >>> x = np.linspace(-1, 1, 10) >>> y = np.linspace(-1, 1, 11) >>> z = np.linspace(-1, 1, 12) >>> coords = Coords(x=x, y=y, z=z) >>> grid = Grid(boundaries=coords) >>> yee_cells = grid.yee >>> Ex_positions = yee_cells.E.x """ yee_e_kwargs = {key: self._yee_e(axis=axis) for axis, key in enumerate("xyz")} yee_h_kwargs = {key: self._yee_h(axis=axis) for axis, key in enumerate("xyz")} yee_e = FieldGrid(**yee_e_kwargs) yee_h = FieldGrid(**yee_h_kwargs) return YeeGrid(E=yee_e, H=yee_h)
[docs] def __getitem__(self, coord_key: str) -> Coords: """quickly get the grid element by grid[key].""" coord_dict = { "centers": self.centers, "sizes": self.sizes, "boundaries": self.boundaries, "Ex": self.yee.E.x, "Ey": self.yee.E.y, "Ez": self.yee.E.z, "Hx": self.yee.H.x, "Hy": self.yee.H.y, "Hz": self.yee.H.z, } if coord_key not in coord_dict: raise SetupError(f"key {coord_key} not found in grid with {list(coord_dict.keys())} ") return coord_dict.get(coord_key)
def _yee_e(self, axis: Axis): """E field yee lattice sites for axis.""" boundary_coords = self.boundaries.to_dict # initially set all to the minus bounds yee_coords = {key: self._min(val) for key, val in boundary_coords.items()} # average the axis index between the cell boundaries key = "xyz"[axis] yee_coords[key] = self._avg(boundary_coords[key]) return Coords(**yee_coords) def _yee_h(self, axis: Axis): """H field yee lattice sites for axis.""" boundary_coords = self.boundaries.to_dict # initially set all to centers yee_coords = {key: self._avg(val) for key, val in boundary_coords.items()} # set the axis index to the minus bounds key = "xyz"[axis] yee_coords[key] = self._min(boundary_coords[key]) return Coords(**yee_coords)
[docs] def discretize_inds(self, box: Box, extend: bool = False) -> List[Tuple[int, int]]: """Start and stopping indexes for the cells that intersect with a :class:`Box`. Parameters ---------- box : :class:`Box` Rectangular geometry within simulation to discretize. extend : bool = False If ``True``, ensure that the returned indexes extend sufficiently in every direction to be able to interpolate any field component at any point within the ``box``, for field components sampled on the Yee grid. Returns ------- List[Tuple[int, int]] The (start, stop) indexes of the cells that intersect with ``box`` in each of the three dimensions. """ pts_min, pts_max = box.bounds boundaries = self.boundaries inds_list = [] # for each dimension for axis, (pt_min, pt_max) in enumerate(zip(pts_min, pts_max)): bound_coords = np.array(boundaries.to_list[axis]) assert pt_min <= pt_max, "min point was greater than max point" # index of smallest coord greater than pt_max inds_gt_pt_max = np.where(bound_coords > pt_max)[0] ind_max = len(bound_coords) - 1 if len(inds_gt_pt_max) == 0 else inds_gt_pt_max[0] # index of largest coord less than or equal to pt_min inds_leq_pt_min = np.where(bound_coords <= pt_min)[0] ind_min = 0 if len(inds_leq_pt_min) == 0 else inds_leq_pt_min[-1] # handle extensions if ind_max > ind_min and extend: # Left side if box.bounds[0][axis] < self.centers.to_list[axis][ind_min]: # Box bounds on the left side are to the left of the closest grid center ind_min -= 1 # We always need an extra pixel on the right for the tangential components ind_max += 1 # store indexes inds_list.append((ind_min, ind_max)) return inds_list
[docs] def extended_subspace( self, axis: Axis, ind_beg: int = 0, ind_end: int = 0, periodic: bool = True, ) -> Coords1D: """Pick a subspace of 1D boundaries within ``range(ind_beg, ind_end)``. If any indexes lie outside of the grid boundaries array, padding is used based on the boundary conditions. For periodic BCs, the zeroth and last element of the grid boundaries are identified. For other BCs, the zeroth and last element of the boundaries are a reflection plane. Parameters ---------- axis : Axis Axis index along which to pick the subspace. ind_beg : int = 0 Starting index for the subspace. ind_end : int = 0 Ending index for the subspace. periodic : bool = True Whether to pad out of bounds indexes with a periodic or reflected pattern. Returns ------- Coords1D The subspace of the grid along ``axis``. """ coords = self.boundaries.to_list[axis] padded_coords = coords num_cells = coords.size - 1 reverse = True while ind_beg < 0: if periodic or not reverse: offset = padded_coords[0] - coords[-1] padded_coords = np.concatenate([coords[:-1] + offset, padded_coords]) reverse = True else: offset = padded_coords[0] + coords[0] padded_coords = np.concatenate([offset - coords[:0:-1], padded_coords]) reverse = False ind_beg += num_cells ind_end += num_cells reverse = True while ind_end >= padded_coords.size: if periodic or not reverse: offset = padded_coords[-1] - coords[0] padded_coords = np.concatenate([padded_coords, coords[1:] + offset]) reverse = True else: offset = padded_coords[-1] + coords[-1] padded_coords = np.concatenate([padded_coords, offset - coords[-2::-1]]) reverse = False return padded_coords[ind_beg:ind_end]
[docs] def snap_to_box_zero_dim(self, box: Box): """Snap a grid to an exact box position for dimensions for which the box is size 0. If the box location is outside of the grid, an error is raised. Parameters ---------- box : :class:`Box` Box to use for the zero dim check. Returns ------- class:`Grid` Snapped copy of the grid. """ boundary_dict = self.boundaries.to_dict.copy() for dim, center, size in zip("xyz", box.center, box.size): # Overwrite grid boundaries with box center if box is size 0 along dimension if size == 0: if boundary_dict[dim][0] > center or boundary_dict[dim][-1] < center: raise ValueError("Cannot snap grid to box center outside of grid domain.") boundary_dict[dim] = np.array([center, center]) return self.updated_copy(boundaries=Coords(**boundary_dict))