From 5935ee5b25afabbe37b20c1c8d20bd3e4071690b Mon Sep 17 00:00:00 2001 From: daquintero Date: Tue, 12 Mar 2024 14:36:55 +0100 Subject: [PATCH 1/2] FEAT: Trying to list the goals of the meeting --- docs/notebooks | 2 +- tidy3d/__init__.py | 2 +- tidy3d/components/mode/LICENSE | 25 + tidy3d/components/mode/README.md | 4 + tidy3d/components/mode/__init__.py | 5 + tidy3d/components/mode/derivatives.py | 172 +++ tidy3d/components/mode/mode_solver.py | 1011 +++++++++++++++++ tidy3d/components/mode/solver.py | 745 ++++++++++++ tidy3d/components/mode/transforms.py | 114 ++ tidy3d/components/{mode.py => mode_spec.py} | 0 tidy3d/components/monitor.py | 2 +- tidy3d/components/source.py | 2 +- tidy3d/plugins/mode/mode_solver.py | 2 +- tidy3d/plugins/smatrix/smatrix.py | 2 +- .../waveguide/rectangular_dielectric.py | 2 +- 15 files changed, 2083 insertions(+), 7 deletions(-) create mode 100644 tidy3d/components/mode/LICENSE create mode 100644 tidy3d/components/mode/README.md create mode 100644 tidy3d/components/mode/__init__.py create mode 100644 tidy3d/components/mode/derivatives.py create mode 100644 tidy3d/components/mode/mode_solver.py create mode 100644 tidy3d/components/mode/solver.py create mode 100644 tidy3d/components/mode/transforms.py rename tidy3d/components/{mode.py => mode_spec.py} (100%) diff --git a/docs/notebooks b/docs/notebooks index 4a01deecd..0bf4aa017 160000 --- a/docs/notebooks +++ b/docs/notebooks @@ -1 +1 @@ -Subproject commit 4a01deecdcd137f26c5f2a1ea05052aeaa10eb50 +Subproject commit 0bf4aa017670b4d593d4fb91390058ff025fc72d diff --git a/tidy3d/__init__.py b/tidy3d/__init__.py index 738d98359..c8067c927 100644 --- a/tidy3d/__init__.py +++ b/tidy3d/__init__.py @@ -32,7 +32,7 @@ from .components.structure import Structure, MeshOverrideStructure # modes -from .components.mode import ModeSpec +from .components.mode_spec import ModeSpec # apodization from .components.apodization import ApodizationSpec diff --git a/tidy3d/components/mode/LICENSE b/tidy3d/components/mode/LICENSE new file mode 100644 index 000000000..dc4a960c6 --- /dev/null +++ b/tidy3d/components/mode/LICENSE @@ -0,0 +1,25 @@ +Some functions in the current module are derived from the ceviche package +https://github.com/fancompute/ceviche + +MIT License + +Copyright (c) 2021 Flexcompute +Copyright (c) 2019 Tyler Hughes + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/tidy3d/components/mode/README.md b/tidy3d/components/mode/README.md new file mode 100644 index 000000000..cb6bcbe39 --- /dev/null +++ b/tidy3d/components/mode/README.md @@ -0,0 +1,4 @@ +# Mode Solver for propagating EM modes + + + diff --git a/tidy3d/components/mode/__init__.py b/tidy3d/components/mode/__init__.py new file mode 100644 index 000000000..cd239940d --- /dev/null +++ b/tidy3d/components/mode/__init__.py @@ -0,0 +1,5 @@ +""" Imports from mode solver plugin. """ + +from .mode_solver import ModeSolver, ModeSolverData + +__all__ = ["ModeSolver", "ModeSolverData"] diff --git a/tidy3d/components/mode/derivatives.py b/tidy3d/components/mode/derivatives.py new file mode 100644 index 000000000..85825d461 --- /dev/null +++ b/tidy3d/components/mode/derivatives.py @@ -0,0 +1,172 @@ +""" Finite-difference derivatives and PML absorption operators expressed as sparse matrices. """ + +import numpy as np +import scipy.sparse as sp + +from ...constants import EPSILON_0, ETA_0 + + +def make_dxf(dls, shape, pmc): + """Forward derivative in x.""" + Nx, Ny = shape + if Nx == 1: + return sp.csr_matrix((Ny, Ny)) + dxf = sp.csr_matrix(sp.diags([-1, 1], [0, 1], shape=(Nx, Nx))) + if not pmc: + dxf[0, 0] = 0.0 + dxf = sp.diags(1 / dls).dot(dxf) + dxf = sp.kron(dxf, sp.eye(Ny)) + return dxf + + +def make_dxb(dls, shape, pmc): + """Backward derivative in x.""" + Nx, Ny = shape + if Nx == 1: + return sp.csr_matrix((Ny, Ny)) + dxb = sp.csr_matrix(sp.diags([1, -1], [0, -1], shape=(Nx, Nx))) + if pmc: + dxb[0, 0] = 2.0 + else: + dxb[0, 0] = 0.0 + dxb = sp.diags(1 / dls).dot(dxb) + dxb = sp.kron(dxb, sp.eye(Ny)) + return dxb + + +def make_dyf(dls, shape, pmc): + """Forward derivative in y.""" + Nx, Ny = shape + if Ny == 1: + return sp.csr_matrix((Nx, Nx)) + dyf = sp.csr_matrix(sp.diags([-1, 1], [0, 1], shape=(Ny, Ny))) + if not pmc: + dyf[0, 0] = 0.0 + dyf = sp.diags(1 / dls).dot(dyf) + dyf = sp.kron(sp.eye(Nx), dyf) + return dyf + + +def make_dyb(dls, shape, pmc): + """Backward derivative in y.""" + Nx, Ny = shape + if Ny == 1: + return sp.csr_matrix((Nx, Nx)) + dyb = sp.csr_matrix(sp.diags([1, -1], [0, -1], shape=(Ny, Ny))) + if pmc: + dyb[0, 0] = 2.0 + else: + dyb[0, 0] = 0.0 + dyb = sp.diags(1 / dls).dot(dyb) + dyb = sp.kron(sp.eye(Nx), dyb) + return dyb + + +def create_d_matrices(shape, dlf, dlb, dmin_pmc=(False, False)): + """Make the derivative matrices without PML. If dmin_pmc is True, the + 'backward' derivative in that dimension will be set to implement PMC + boundary, otherwise it will be set to PEC.""" + + dxf = make_dxf(dlf[0], shape, dmin_pmc[0]) + dxb = make_dxb(dlb[0], shape, dmin_pmc[0]) + dyf = make_dyf(dlf[1], shape, dmin_pmc[1]) + dyb = make_dyb(dlb[1], shape, dmin_pmc[1]) + + return (dxf, dxb, dyf, dyb) + + +def create_s_matrices(omega, shape, npml, dlf, dlb, dmin_pml=(True, True)): + """Makes the 'S-matrices'. When dotted with derivative matrices, they add + PML. If dmin_pml is set to False, PML will not be applied on the "bottom" + side of the domain.""" + + # strip out some information needed + Nx, Ny = shape + N = Nx * Ny + nx_pml, ny_pml = npml + + # Create the sfactor in each direction and for 'f' and 'b' + s_vector_x_f = create_sfactor("f", omega, dlf[0], Nx, nx_pml, dmin_pml[0]) + s_vector_x_b = create_sfactor("b", omega, dlb[0], Nx, nx_pml, dmin_pml[0]) + s_vector_y_f = create_sfactor("f", omega, dlf[1], Ny, ny_pml, dmin_pml[1]) + s_vector_y_b = create_sfactor("b", omega, dlb[1], Ny, ny_pml, dmin_pml[1]) + + # Fill the 2d space with layers of appropriate s-factors + sx_f_2d = np.zeros(shape, dtype=np.complex128) + sx_b_2d = np.zeros(shape, dtype=np.complex128) + sy_f_2d = np.zeros(shape, dtype=np.complex128) + sy_b_2d = np.zeros(shape, dtype=np.complex128) + + # Insert the cross sections into the S-grids (could be done more elegantly) + for i in range(Ny): + sx_f_2d[:, i] = 1 / s_vector_x_f + sx_b_2d[:, i] = 1 / s_vector_x_b + for i in range(Nx): + sy_f_2d[i, :] = 1 / s_vector_y_f + sy_b_2d[i, :] = 1 / s_vector_y_b + + # Reshape the 2d s-factors into a 1D s-vecay + sx_f_vec = sx_f_2d.flatten() + sx_b_vec = sx_b_2d.flatten() + sy_f_vec = sy_f_2d.flatten() + sy_b_vec = sy_b_2d.flatten() + + # Construct the 1D total s-vector into a diagonal matrix + sx_f = sp.spdiags(sx_f_vec, 0, N, N) + sx_b = sp.spdiags(sx_b_vec, 0, N, N) + sy_f = sp.spdiags(sy_f_vec, 0, N, N) + sy_b = sp.spdiags(sy_b_vec, 0, N, N) + + return sx_f, sx_b, sy_f, sy_b + + +def create_sfactor(direction, omega, dls, N, n_pml, dmin_pml): + """Creates the S-factor cross section needed in the S-matrices""" + + # For no PNL, this should just be identity matrix. + if n_pml == 0: + return np.ones(N, dtype=np.complex128) + + # Otherwise, get different profiles for forward and reverse derivatives. + if direction == "f": + return create_sfactor_f(omega, dls, N, n_pml, dmin_pml) + if direction == "b": + return create_sfactor_b(omega, dls, N, n_pml, dmin_pml) + + raise ValueError(f"Direction value {direction} not recognized") + + +def create_sfactor_f(omega, dls, N, n_pml, dmin_pml): + """S-factor profile applied after forward derivative matrix, i.e. applied to H-field + locations.""" + sfactor_array = np.ones(N, dtype=np.complex128) + for i in range(N): + if i <= n_pml - 1 and dmin_pml: + sfactor_array[i] = s_value(dls[0], (n_pml - i - 0.5) / n_pml, omega) + elif i >= N - n_pml: + sfactor_array[i] = s_value(dls[-1], (i - (N - n_pml) + 0.5) / n_pml, omega) + return sfactor_array + + +def create_sfactor_b(omega, dls, N, n_pml, dmin_pml): + """S-factor profile applied after backward derivative matrix, i.e. applied to E-field + locations.""" + sfactor_array = np.ones(N, dtype=np.complex128) + for i in range(N): + if i < n_pml and dmin_pml: + sfactor_array[i] = s_value(dls[0], (n_pml - i) / n_pml, omega) + elif i > N - n_pml: + sfactor_array[i] = s_value(dls[-1], (i - (N - n_pml)) / n_pml, omega) + return sfactor_array + + +def sig_w(dl, step, sorder=3): + """Fictional conductivity, note that these values might need tuning""" + sig_max = 0.8 * (sorder + 1) / (ETA_0 * dl) + return sig_max * step**sorder + + +def s_value(dl, step, omega): + """S-value to use in the S-matrices""" + # print(step) + return 1 - 1j * sig_w(dl, step) / (omega * EPSILON_0) diff --git a/tidy3d/components/mode/mode_solver.py b/tidy3d/components/mode/mode_solver.py new file mode 100644 index 000000000..4b25c1b18 --- /dev/null +++ b/tidy3d/components/mode/mode_solver.py @@ -0,0 +1,1011 @@ +"""Solve for modes in a 2D cross-sectional plane in a simulation, assuming translational +invariance along a given propagation axis. +""" + +from __future__ import annotations +from typing import List, Tuple, Dict +from math import isclose + +import numpy as np +import pydantic.v1 as pydantic +import xarray as xr + +from ...log import log +from ...components.base_sim.simulation import AbstractSimulation +from ...components.base import Tidy3dBaseModel, cached_property, skip_if_fields_missing +from ...components.boundary import PECBoundary, BoundarySpec, Boundary, PML, StablePML, Absorber +from ...components.geometry.base import Box +from ...components.simulation import Simulation +from ...components.grid.grid import Grid +from ...components.mode_spec import ModeSpec +from ...components.monitor import ModeSolverMonitor, ModeMonitor +from ...components.medium import FullyAnisotropicMedium +from ...components.source import ModeSource, SourceTime +from ...components.types import Direction, FreqArray, Ax, Literal, Axis, Symmetry, PlotScale +from ...components.types import ArrayComplex3D, ArrayComplex4D, ArrayFloat1D, EpsSpecType +from ...components.data.data_array import ModeIndexDataArray, ScalarModeFieldDataArray +from ...components.data.data_array import FreqModeDataArray +from ...components.data.sim_data import SimulationData +from ...components.data.monitor_data import ModeSolverData + +from ...components.validators import validate_freqs_min, validate_freqs_not_empty +from ...exceptions import ValidationError, SetupError +from ...constants import C_0 + +# Importing the local solver may not work if e.g. scipy is not installed +IMPORT_ERROR_MSG = """Could not import local solver, 'ModeSolver' objects can still be constructed +but will have to be run through the server. +""" +try: + from .solver import compute_modes + + LOCAL_SOLVER_IMPORTED = True +except ImportError: + log.warning(IMPORT_ERROR_MSG) + LOCAL_SOLVER_IMPORTED = False + +FIELD = Tuple[ArrayComplex3D, ArrayComplex3D, ArrayComplex3D] +MODE_MONITOR_NAME = "<<>>" + +# Warning for field intensity at edges over total field intensity larger than this value +FIELD_DECAY_CUTOFF = 1e-2 + +# Maximum allowed size of the field data produced by the mode solver +MAX_MODES_DATA_SIZE_GB = 20 + + +class ModeSolver(AbstractSimulation): + """ + Interface for solving electromagnetic eigenmodes in a 2D plane with translational + invariance in the third dimension. + + See Also + -------- + + :class:`ModeSource`: + Injects current source to excite modal profile on finite extent plane. + + **Notebooks:** + * `Waveguide Y junction <../../notebooks/YJunction.html>`_ + * `Photonic crystal waveguide polarization filter <../../../notebooks/PhotonicCrystalWaveguidePolarizationFilter.html>`_ + + **Lectures:** + * `Prelude to Integrated Photonics Simulation: Mode Injection `_ + """ + + plane: Box = pydantic.Field( + ..., title="Plane", description="Cross-sectional plane in which the mode will be computed." + ) + + mode_spec: ModeSpec = pydantic.Field( + ..., + title="Mode specification", + description="Container with specifications about the modes to be solved for.", + ) + + freqs: FreqArray = pydantic.Field( + ..., title="Frequencies", description="A list of frequencies at which to solve." + ) + + direction: Direction = pydantic.Field( + "+", + title="Propagation direction", + description="Direction of waveguide mode propagation along the axis defined by its normal " + "dimension.", + ) + + colocate: bool = pydantic.Field( + True, + title="Colocate fields", + description="Toggle whether fields should be colocated to grid cell boundaries (i.e. " + "primal grid nodes). Default is ``True``.", + ) + + @pydantic.validator("plane", always=True) + def is_plane(cls, val): + """Raise validation error if not planar.""" + if val.size.count(0.0) != 1: + raise ValidationError(f"ModeSolver plane must be planar, given size={val}") + return val + + _freqs_not_empty = validate_freqs_not_empty() + _freqs_lower_bound = validate_freqs_min() + + @pydantic.validator("plane", always=True) + @skip_if_fields_missing(["simulation"]) + def plane_in_sim_bounds(cls, val, values): + """Check that the plane is at least partially inside the simulation bounds.""" + sim_center = values.get("simulation").center + sim_size = values.get("simulation").size + sim_box = Box(size=sim_size, center=sim_center) + + if not sim_box.intersects(val): + raise SetupError("'ModeSolver.plane' must intersect 'ModeSolver.simulation'.") + return val + + @cached_property + def normal_axis(self) -> Axis: + """Axis normal to the mode plane.""" + return self.plane.size.index(0.0) + + @cached_property + def solver_symmetry(self) -> Tuple[Symmetry, Symmetry]: + """Get symmetry for solver for propagation along self.normal axis.""" + mode_symmetry = list(self.simulation.symmetry) + for dim in range(3): + if self.simulation.center[dim] != self.plane.center[dim]: + mode_symmetry[dim] = 0 + _, solver_sym = self.plane.pop_axis(mode_symmetry, axis=self.normal_axis) + return solver_sym + + def _get_solver_grid( + self, keep_additional_layers: bool = False, truncate_symmetry: bool = True + ) -> Grid: + """Grid for the mode solver, not snapped to plane or simulation zero dims, and optionally + corrected for symmetries. + + Parameters + ---------- + keep_additional_layers : bool = False + Do not discard layers of cells in front and behind the main layer of cells. Together they + represent the region where custom medium data is needed for proper subpixel. + truncate_symmetry : bool = True + Truncate to symmetry quadrant if symmetry present. + + Returns + ------- + :class:.`Grid` + The resulting grid. + """ + + monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME, colocate=False) + + span_inds = self.simulation._discretize_inds_monitor(monitor) + + # Remove extension along monitor normal + if not keep_additional_layers: + span_inds[self.normal_axis, 0] += 1 + span_inds[self.normal_axis, 1] -= 1 + + # Do not extend if simulation has a single pixel along a dimension + for dim, num_cells in enumerate(self.simulation.grid.num_cells): + if num_cells <= 1: + span_inds[dim] = [0, 1] + + # Truncate to symmetry quadrant if symmetry present + if truncate_symmetry: + _, plane_inds = Box.pop_axis([0, 1, 2], self.normal_axis) + for dim, sym in enumerate(self.solver_symmetry): + if sym != 0: + span_inds[plane_inds[dim], 0] += np.diff(span_inds[plane_inds[dim]]) // 2 + + return self.simulation._subgrid(span_inds=span_inds) + + @cached_property + def _solver_grid(self) -> Grid: + """Grid for the mode solver, not snapped to plane or simulation zero dims, and also with + a small correction for symmetries. We don't do the snapping yet because 0-sized cells are + currently confusing to the subpixel averaging. The final data coordinates along the + plane normal dimension and dimensions where the simulation domain is 2D will be correctly + set after the solve.""" + + return self._get_solver_grid(keep_additional_layers=False, truncate_symmetry=True) + + @cached_property + def _num_cells_freqs_modes(self) -> Tuple[int, int, int]: + """Get the number of spatial points, number of freqs, and number of modes requested.""" + num_cells = np.prod(self._solver_grid.num_cells) + num_modes = self.mode_spec.num_modes + num_freqs = len(self.freqs) + return num_cells, num_freqs, num_modes + + def solve(self) -> ModeSolverData: + """:class:`.ModeSolverData` containing the field and effective index data. + + Returns + ------- + ModeSolverData + :class:`.ModeSolverData` object containing the effective index and mode fields. + """ + log.warning( + "Use the remote mode solver with subpixel averaging for better accuracy through " + "'tidy3d.plugins.mode.web.run(...)'.", + log_once=True, + ) + return self.data + + def _freqs_for_group_index(self) -> FreqArray: + """Get frequencies used to compute group index.""" + f_step = self.mode_spec.group_index_step + fractional_steps = (1 - f_step, 1, 1 + f_step) + return np.outer(self.freqs, fractional_steps).flatten() + + def _remove_freqs_for_group_index(self) -> FreqArray: + """Remove frequencies used to compute group index. + + Returns + ------- + FreqArray + Filtered frequency array with only original values. + """ + return np.array(self.freqs[1 : len(self.freqs) : 3]) + + def _get_data_with_group_index(self) -> ModeSolverData: + """:class:`.ModeSolverData` with fields, effective and group indices on unexpanded grid. + + Returns + ------- + ModeSolverData + :class:`.ModeSolverData` object containing the effective and group indices, and mode + fields. + """ + + # create a copy with the required frequencies for numerical differentiation + mode_spec = self.mode_spec.copy(update={"group_index_step": False}) + mode_solver = self.copy( + update={"freqs": self._freqs_for_group_index(), "mode_spec": mode_spec} + ) + + return mode_solver.data_raw._group_index_post_process(self.mode_spec.group_index_step) + + @cached_property + def grid_snapped(self) -> Grid: + """The solver grid snapped to the plane normal and to simulation 0-sized dims if any.""" + grid_snapped = self._solver_grid.snap_to_box_zero_dim(self.plane) + return self.simulation._snap_zero_dim(grid_snapped) + + @cached_property + def data_raw(self) -> ModeSolverData: + """:class:`.ModeSolverData` containing the field and effective index on unexpanded grid. + + Returns + ------- + ModeSolverData + :class:`.ModeSolverData` object containing the effective index and mode fields. + """ + + if self.mode_spec.group_index_step > 0: + return self._get_data_with_group_index() + + # Compute data on the Yee grid + mode_solver_data = self._data_on_yee_grid() + + # Colocate to grid boundaries if requested + if self.colocate: + mode_solver_data = self._colocate_data(mode_solver_data=mode_solver_data) + + # normalize modes + self._normalize_modes(mode_solver_data=mode_solver_data) + + # filter polarization if requested + if self.mode_spec.filter_pol is not None: + self._filter_polarization(mode_solver_data=mode_solver_data) + + # sort modes if requested + if self.mode_spec.track_freq and len(self.freqs) > 1: + mode_solver_data = mode_solver_data.overlap_sort(self.mode_spec.track_freq) + + self._field_decay_warning(mode_solver_data.symmetry_expanded) + + return mode_solver_data + + def _data_on_yee_grid(self) -> ModeSolverData: + """Solve for all modes, and construct data with fields on the Yee grid.""" + _, _solver_coords = self.plane.pop_axis( + self._solver_grid.boundaries.to_list, axis=self.normal_axis + ) + + # Compute and store the modes at all frequencies + n_complex, fields, eps_spec = self._solve_all_freqs( + coords=_solver_coords, symmetry=self.solver_symmetry + ) + + # start a dictionary storing the data arrays for the ModeSolverData + index_data = ModeIndexDataArray( + np.stack(n_complex, axis=0), + coords=dict( + f=list(self.freqs), + mode_index=np.arange(self.mode_spec.num_modes), + ), + ) + data_dict = {"n_complex": index_data} + + # Construct the field data on Yee grid + for field_name in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz"): + xyz_coords = self.grid_snapped[field_name].to_list + scalar_field_data = ScalarModeFieldDataArray( + np.stack([field_freq[field_name] for field_freq in fields], axis=-2), + coords=dict( + x=xyz_coords[0], + y=xyz_coords[1], + z=xyz_coords[2], + f=list(self.freqs), + mode_index=np.arange(self.mode_spec.num_modes), + ), + ) + data_dict[field_name] = scalar_field_data + + # finite grid corrections + grid_factors = self._grid_correction( + simulation=self.simulation, + plane=self.plane, + mode_spec=self.mode_spec, + n_complex=index_data, + direction=self.direction, + ) + + # make mode solver data on the Yee grid + mode_solver_monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME, colocate=False) + grid_expanded = self.simulation.discretize_monitor(mode_solver_monitor) + mode_solver_data = ModeSolverData( + monitor=mode_solver_monitor, + symmetry=self.simulation.symmetry, + symmetry_center=self.simulation.center, + grid_expanded=grid_expanded, + grid_primal_correction=grid_factors[0], + grid_dual_correction=grid_factors[1], + eps_spec=eps_spec, + **data_dict, + ) + + return mode_solver_data + + def _colocate_data(self, mode_solver_data: ModeSolverData) -> ModeSolverData: + """Colocate data to Yee grid boundaries.""" + + # Get colocation coordinates in the solver plane + _, plane_dims = self.plane.pop_axis("xyz", self.normal_axis) + colocate_coords = {} + for dim, sym in zip(plane_dims, self.solver_symmetry): + coords = self.grid_snapped.boundaries.to_dict[dim] + if len(coords) > 2: + if sym == 0: + colocate_coords[dim] = coords[1:-1] + else: + colocate_coords[dim] = coords[:-1] + + # Colocate input data to new coordinates + data_dict_colocated = {} + for key, field in mode_solver_data.symmetry_expanded.field_components.items(): + data_dict_colocated[key] = field.interp(**colocate_coords).astype(field.dtype) + + # Update data + mode_solver_monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME) + grid_expanded = self.simulation.discretize_monitor(mode_solver_monitor) + data_dict_colocated.update({"monitor": mode_solver_monitor, "grid_expanded": grid_expanded}) + mode_solver_data = mode_solver_data._updated(update=data_dict_colocated) + + return mode_solver_data + + def _normalize_modes(self, mode_solver_data: ModeSolverData): + """Normalize modes. Note: this modifies ``mode_solver_data`` in-place.""" + scaling = np.sqrt(np.abs(mode_solver_data.flux)) + for field in mode_solver_data.field_components.values(): + field /= scaling + + def _filter_polarization(self, mode_solver_data: ModeSolverData): + """Filter polarization. Note: this modifies ``mode_solver_data`` in-place.""" + pol_frac = mode_solver_data.pol_fraction + for ifreq in range(len(self.freqs)): + te_frac = pol_frac.te.isel(f=ifreq) + if self.mode_spec.filter_pol == "te": + sort_inds = np.concatenate( + ( + np.where(te_frac >= 0.5)[0], + np.where(te_frac < 0.5)[0], + np.where(np.isnan(te_frac))[0], + ) + ) + elif self.mode_spec.filter_pol == "tm": + sort_inds = np.concatenate( + ( + np.where(te_frac <= 0.5)[0], + np.where(te_frac > 0.5)[0], + np.where(np.isnan(te_frac))[0], + ) + ) + for data in list(mode_solver_data.field_components.values()) + [ + mode_solver_data.n_complex, + mode_solver_data.grid_primal_correction, + mode_solver_data.grid_dual_correction, + ]: + data.values[..., ifreq, :] = data.values[..., ifreq, sort_inds] + + @cached_property + def data(self) -> ModeSolverData: + """:class:`.ModeSolverData` containing the field and effective index data. + + Returns + ------- + ModeSolverData + :class:`.ModeSolverData` object containing the effective index and mode fields. + """ + mode_solver_data = self.data_raw + return mode_solver_data.symmetry_expanded_copy + + @cached_property + def sim_data(self) -> SimulationData: + """:class:`.SimulationData` object containing the :class:`.ModeSolverData` for this object. + + Returns + ------- + SimulationData + :class:`.SimulationData` object containing the effective index and mode fields. + """ + monitor_data = self.data + new_monitors = list(self.simulation.monitors) + [monitor_data.monitor] + new_simulation = self.simulation.copy(update=dict(monitors=new_monitors)) + return SimulationData(simulation=new_simulation, data=(monitor_data,)) + + def _get_epsilon(self, freq: float) -> ArrayComplex4D: + """Compute the epsilon tensor in the plane. Order of components is xx, xy, xz, yx, etc.""" + eps_keys = ["Ex", "Exy", "Exz", "Eyx", "Ey", "Eyz", "Ezx", "Ezy", "Ez"] + eps_tensor = [ + self.simulation.epsilon_on_grid(self._solver_grid, key, freq) for key in eps_keys + ] + return np.stack(eps_tensor, axis=0) + + def _solver_eps(self, freq: float) -> ArrayComplex4D: + """Diagonal permittivity in the shape needed by solver, with normal axis rotated to z.""" + + # Get diagonal epsilon components in the plane + eps_tensor = self._get_epsilon(freq) + + # get rid of normal axis + eps_tensor = np.take(eps_tensor, indices=[0], axis=1 + self.normal_axis) + eps_tensor = np.squeeze(eps_tensor, axis=1 + self.normal_axis) + + # convert to into 3-by-3 representation for easier axis swap + flat_shape = np.shape(eps_tensor) # 9 components flat + tensor_shape = [3, 3] + list(flat_shape[1:]) # 3-by-3 matrix + eps_tensor = eps_tensor.reshape(tensor_shape) + + # swap axes to plane coordinates (normal_axis goes to z) + if self.normal_axis == 0: + # swap x and y + eps_tensor[[0, 1], :, ...] = eps_tensor[[1, 0], :, ...] + eps_tensor[:, [0, 1], ...] = eps_tensor[:, [1, 0], ...] + if self.normal_axis <= 1: + # swap x (normal_axis==0) or y (normal_axis==1) and z + eps_tensor[[1, 2], :, ...] = eps_tensor[[2, 1], :, ...] + eps_tensor[:, [1, 2], ...] = eps_tensor[:, [2, 1], ...] + + # back to "flat" representation + eps_tensor = eps_tensor.reshape(flat_shape) + + # construct eps to feed to mode solver + return eps_tensor + + def _solve_all_freqs( + self, + coords: Tuple[ArrayFloat1D, ArrayFloat1D], + symmetry: Tuple[Symmetry, Symmetry], + ) -> Tuple[List[float], List[Dict[str, ArrayComplex4D]], List[EpsSpecType]]: + """Call the mode solver at all requested frequencies.""" + + fields = [] + n_complex = [] + eps_spec = [] + for freq in self.freqs: + n_freq, fields_freq, eps_spec_freq = self._solve_single_freq( + freq=freq, coords=coords, symmetry=symmetry + ) + fields.append(fields_freq) + n_complex.append(n_freq) + eps_spec.append(eps_spec_freq) + + return n_complex, fields, eps_spec + + def _solve_single_freq( + self, + freq: float, + coords: Tuple[ArrayFloat1D, ArrayFloat1D], + symmetry: Tuple[Symmetry, Symmetry], + ) -> Tuple[float, Dict[str, ArrayComplex4D], EpsSpecType]: + """Call the mode solver at a single frequency. + + The fields are rotated from propagation coordinates back to global coordinates. + """ + + if not LOCAL_SOLVER_IMPORTED: + raise ImportError(IMPORT_ERROR_MSG) + + solver_fields, n_complex, eps_spec = compute_modes( + eps_cross=self._solver_eps(freq), + coords=coords, + freq=freq, + mode_spec=self.mode_spec, + symmetry=symmetry, + direction=self.direction, + ) + + fields = {key: [] for key in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz")} + for mode_index in range(self.mode_spec.num_modes): + # Get E and H fields at the current mode_index + ((Ex, Ey, Ez), (Hx, Hy, Hz)) = self._process_fields(solver_fields, mode_index) + + # Note: back in original coordinates + fields_mode = {"Ex": Ex, "Ey": Ey, "Ez": Ez, "Hx": Hx, "Hy": Hy, "Hz": Hz} + for field_name, field in fields_mode.items(): + fields[field_name].append(field) + + for field_name, field in fields.items(): + fields[field_name] = np.stack(field, axis=-1) + + return n_complex, fields, eps_spec + + def _rotate_field_coords(self, field: FIELD) -> FIELD: + """Move the propagation axis=z to the proper order in the array.""" + f_x, f_y, f_z = np.moveaxis(field, source=3, destination=1 + self.normal_axis) + return np.stack(self.plane.unpop_axis(f_z, (f_x, f_y), axis=self.normal_axis), axis=0) + + def _process_fields( + self, mode_fields: ArrayComplex4D, mode_index: pydantic.NonNegativeInt + ) -> Tuple[FIELD, FIELD]: + """Transform solver fields to simulation axes and set gauge.""" + + # Separate E and H fields (in solver coordinates) + E, H = mode_fields[..., mode_index] + + # Set gauge to highest-amplitude in-plane E being real and positive + ind_max = np.argmax(np.abs(E[:2])) + phi = np.angle(E[:2].ravel()[ind_max]) + E *= np.exp(-1j * phi) + H *= np.exp(-1j * phi) + + # Rotate back to original coordinates + (Ex, Ey, Ez) = self._rotate_field_coords(E) + (Hx, Hy, Hz) = self._rotate_field_coords(H) + + # apply -1 to H fields if a reflection was involved in the rotation + if self.normal_axis == 1: + Hx *= -1 + Hy *= -1 + Hz *= -1 + + return ((Ex, Ey, Ez), (Hx, Hy, Hz)) + + def _field_decay_warning(self, field_data: ModeSolverData): + """Warn if any of the modes do not decay at the edges.""" + _, plane_dims = self.plane.pop_axis(["x", "y", "z"], axis=self.normal_axis) + field_sizes = field_data.Ex.sizes + for freq_index in range(field_sizes["f"]): + for mode_index in range(field_sizes["mode_index"]): + e_edge, e_norm = 0, 0 + # Sum up the total field intensity + for E in (field_data.Ex, field_data.Ey, field_data.Ez): + e_norm += np.sum(np.abs(E[{"f": freq_index, "mode_index": mode_index}]) ** 2) + # Sum up the field intensity at the edges + if field_sizes[plane_dims[0]] > 1: + for E in (field_data.Ex, field_data.Ey, field_data.Ez): + isel = {plane_dims[0]: [0, -1], "f": freq_index, "mode_index": mode_index} + e_edge += np.sum(np.abs(E[isel]) ** 2) + if field_sizes[plane_dims[1]] > 1: + for E in (field_data.Ex, field_data.Ey, field_data.Ez): + isel = {plane_dims[1]: [0, -1], "f": freq_index, "mode_index": mode_index} + e_edge += np.sum(np.abs(E[isel]) ** 2) + # Warn if needed + if e_edge / e_norm > FIELD_DECAY_CUTOFF: + log.warning( + f"Mode field at frequency index {freq_index}, mode index {mode_index} does " + "not decay at the plane boundaries." + ) + + @staticmethod + def _grid_correction( + simulation: Simulation, + plane: Box, + mode_spec: ModeSpec, + n_complex: ModeIndexDataArray, + direction: Direction, + ) -> [FreqModeDataArray, FreqModeDataArray]: + """Correct the fields due to propagation on the grid. + + Return a copy of the :class:`.ModeSolverData` with the fields renormalized to account + for propagation on a finite grid along the propagation direction. The fields are assumed to + have ``E exp(1j k r)`` dependence on the finite grid and are then resampled using linear + interpolation to the exact position of the mode plane. This is needed to correctly compute + overlap with fields that come from a :class:`.FieldMonitor` placed in the same grid. + + Parameters + ---------- + grid : :class:`.Grid` + Numerical grid on which the modes are assumed to propagate. + + Returns + ------- + :class:`.ModeSolverData` + Copy of the data with renormalized fields. + """ + normal_axis = plane.size.index(0.0) + normal_pos = plane.center[normal_axis] + normal_dim = "xyz"[normal_axis] + + # Primal and dual grid along the normal direction, + # i.e. locations of the tangential E-field and H-field components, respectively + grid = simulation.grid + normal_primal = grid.boundaries.to_list[normal_axis] + normal_primal = xr.DataArray(normal_primal, coords={normal_dim: normal_primal}) + normal_dual = grid.centers.to_list[normal_axis] + normal_dual = xr.DataArray(normal_dual, coords={normal_dim: normal_dual}) + + # Propagation phase at the primal and dual locations. The k-vector is along the propagation + # direction, so angle_theta has to be taken into account. The distance along the propagation + # direction is the distance along the normal direction over cosine(theta). + cos_theta = np.cos(mode_spec.angle_theta) + k_vec = 2 * np.pi * n_complex * n_complex.f / C_0 / cos_theta + if direction == "-": + k_vec *= -1 + phase_primal = np.exp(1j * k_vec * (normal_primal - normal_pos)) + phase_dual = np.exp(1j * k_vec * (normal_dual - normal_pos)) + + # Fields are modified by a linear interpolation to the exact monitor position + if normal_primal.size > 1: + phase_primal = phase_primal.interp(**{normal_dim: normal_pos}) + else: + phase_primal = phase_primal.squeeze(dim=normal_dim) + if normal_dual.size > 1: + phase_dual = phase_dual.interp(**{normal_dim: normal_pos}) + else: + phase_dual = phase_dual.squeeze(dim=normal_dim) + + return FreqModeDataArray(phase_primal), FreqModeDataArray(phase_dual) + + @property + def _is_tensorial(self) -> bool: + """Whether the mode computation should be fully tensorial. This is either due to fully + anisotropic media, or due to an angled waveguide, in which case the transformed eps and mu + become tensorial. A separate check is done inside the solver, which looks at the actual + eps and mu and uses a tolerance to determine whether to invoke the tensorial solver, so + the actual behavior may differ from what's predicted by this property.""" + return abs(self.mode_spec.angle_theta) > 0 or self._has_fully_anisotropic_media + + @cached_property + def _intersecting_media(self) -> List: + """List of media (including simulation background) intersecting the mode plane.""" + total_structures = [self.simulation.scene.background_structure] + total_structures += list(self.simulation.structures) + return self.simulation.scene.intersecting_media(self.plane, total_structures) + + @cached_property + def _has_fully_anisotropic_media(self) -> bool: + """Check if there are any fully anisotropic media in the plane of the mode.""" + if np.any( + [isinstance(mat, FullyAnisotropicMedium) for mat in self.simulation.scene.mediums] + ): + for int_mat in self._intersecting_media: + if isinstance(int_mat, FullyAnisotropicMedium): + return True + return False + + @cached_property + def _has_complex_eps(self) -> bool: + """Check if there are media with a complex-valued epsilon in the plane of the mode. + A separate check is done inside the solver, which looks at the actual + eps and mu and uses a tolerance to determine whether to use real or complex fields, so + the actual behavior may differ from what's predicted by this property.""" + check_freqs = np.unique([np.amin(self.freqs), np.amax(self.freqs), np.mean(self.freqs)]) + for int_mat in self._intersecting_media: + for freq in check_freqs: + max_imag_eps = np.amax(np.abs(np.imag(int_mat.eps_model(freq)))) + if not isclose(max_imag_eps, 0): + return False + return True + + def to_source( + self, + source_time: SourceTime, + direction: Direction = None, + mode_index: pydantic.NonNegativeInt = 0, + ) -> ModeSource: + """Creates :class:`.ModeSource` from a :class:`ModeSolver` instance plus additional + specifications. + + Parameters + ---------- + source_time: :class:`.SourceTime` + Specification of the source time-dependence. + direction : Direction = None + Whether source will inject in ``"+"`` or ``"-"`` direction relative to plane normal. + If not specified, uses the direction from the mode solver. + mode_index : int = 0 + Index into the list of modes returned by mode solver to use in source. + + Returns + ------- + :class:`.ModeSource` + Mode source with specifications taken from the ModeSolver instance and the method + inputs. + """ + + if direction is None: + direction = self.direction + + return ModeSource( + center=self.plane.center, + size=self.plane.size, + source_time=source_time, + mode_spec=self.mode_spec, + mode_index=mode_index, + direction=direction, + ) + + def to_monitor(self, freqs: List[float] = None, name: str = None) -> ModeMonitor: + """Creates :class:`ModeMonitor` from a :class:`ModeSolver` instance plus additional + specifications. + + Parameters + ---------- + freqs : List[float] + Frequencies to include in Monitor (Hz). + If not specified, passes ``self.freqs``. + name : str + Required name of monitor. + + Returns + ------- + :class:`.ModeMonitor` + Mode monitor with specifications taken from the ModeSolver instance and the method + inputs. + """ + + if freqs is None: + freqs = self.freqs + + if name is None: + raise ValueError( + "A 'name' must be passed to 'ModeSolver.to_monitor'. " + "The default value of 'None' is for backwards compatibility and is not accepted." + ) + + return ModeMonitor( + center=self.plane.center, + size=self.plane.size, + freqs=freqs, + mode_spec=self.mode_spec, + name=name, + ) + + def to_mode_solver_monitor(self, name: str, colocate: bool = None) -> ModeSolverMonitor: + """Creates :class:`ModeSolverMonitor` from a :class:`ModeSolver` instance. + + Parameters + ---------- + name : str + Name of the monitor. + colocate : bool + Whether to colocate fields or compute on the Yee grid. If not provided, the value + set in the :class:`ModeSolver` instance is used. + + Returns + ------- + :class:`.ModeSolverMonitor` + Mode monitor with specifications taken from the ModeSolver instance and ``name``. + """ + + if colocate is None: + colocate = self.colocate + + return ModeSolverMonitor( + size=self.plane.size, + center=self.plane.center, + mode_spec=self.mode_spec, + freqs=self.freqs, + direction=self.direction, + colocate=colocate, + name=name, + ) + + def sim_with_source( + self, + source_time: SourceTime, + direction: Direction = None, + mode_index: pydantic.NonNegativeInt = 0, + ) -> Simulation: + """Creates :class:`Simulation` from a :class:`ModeSolver`. Creates a copy of + the ModeSolver's original simulation with a ModeSource added corresponding to + the ModeSolver parameters. + + Parameters + ---------- + source_time: :class:`.SourceTime` + Specification of the source time-dependence. + direction : Direction = None + Whether source will inject in ``"+"`` or ``"-"`` direction relative to plane normal. + If not specified, uses the direction from the mode solver. + mode_index : int = 0 + Index into the list of modes returned by mode solver to use in source. + + Returns + ------- + :class:`.Simulation` + Copy of the simulation with a :class:`.ModeSource` with specifications taken + from the ModeSolver instance and the method inputs. + """ + + mode_source = self.to_source( + mode_index=mode_index, direction=direction, source_time=source_time + ) + new_sources = list(self.simulation.sources) + [mode_source] + new_sim = self.simulation.updated_copy(sources=new_sources) + return new_sim + + def sim_with_monitor( + self, + freqs: List[float] = None, + name: str = None, + ) -> Simulation: + """Creates :class:`.Simulation` from a :class:`ModeSolver`. Creates a copy of + the ModeSolver's original simulation with a mode monitor added corresponding to + the ModeSolver parameters. + + Parameters + ---------- + freqs : List[float] = None + Frequencies to include in Monitor (Hz). + If not specified, uses the frequencies from the mode solver. + name : str + Required name of monitor. + + Returns + ------- + :class:`.Simulation` + Copy of the simulation with a :class:`.ModeMonitor` with specifications taken + from the ModeSolver instance and the method inputs. + """ + + mode_monitor = self.to_monitor(freqs=freqs, name=name) + new_monitors = list(self.simulation.monitors) + [mode_monitor] + new_sim = self.simulation.updated_copy(monitors=new_monitors) + return new_sim + + def sim_with_mode_solver_monitor( + self, + name: str, + ) -> Simulation: + """Creates :class:`Simulation` from a :class:`ModeSolver`. Creates a + copy of the ModeSolver's original simulation with a mode solver monitor + added corresponding to the ModeSolver parameters. + + Parameters + ---------- + name : str + Name of the monitor. + + Returns + ------- + :class:`.Simulation` + Copy of the simulation with a :class:`.ModeSolverMonitor` with specifications taken + from the ModeSolver instance and ``name``. + """ + mode_solver_monitor = self.to_mode_solver_monitor(name=name) + new_monitors = list(self.simulation.monitors) + [mode_solver_monitor] + new_sim = self.simulation.updated_copy(monitors=new_monitors) + return new_sim + + def plot_field( + self, + field_name: str, + val: Literal["real", "imag", "abs"] = "real", + scale: PlotScale = "lin", + eps_alpha: float = 0.2, + robust: bool = True, + vmin: float = None, + vmax: float = None, + ax: Ax = None, + **sel_kwargs, + ) -> Ax: + """Plot the field for a :class:`.ModeSolverData` with :class:`.Simulation` plot overlaid. + + Parameters + ---------- + field_name : str + Name of ``field`` component to plot (eg. ``'Ex'``). + Also accepts ``'E'`` and ``'H'`` to plot the vector magnitudes of the electric and + magnetic fields, and ``'S'`` for the Poynting vector. + val : Literal['real', 'imag', 'abs', 'abs^2', 'dB'] = 'real' + Which part of the field to plot. + eps_alpha : float = 0.2 + Opacity of the structure permittivity. + Must be between 0 and 1 (inclusive). + robust : bool = True + If True and vmin or vmax are absent, uses the 2nd and 98th percentiles of the data + to compute the color limits. This helps in visualizing the field patterns especially + in the presence of a source. + 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. + ax : matplotlib.axes._subplots.Axes = None + matplotlib axes to plot on, if not specified, one is created. + sel_kwargs : keyword arguments used to perform ``.sel()`` selection in the monitor data. + These kwargs can select over the spatial dimensions (``x``, ``y``, ``z``), + frequency or time dimensions (``f``, ``t``) or `mode_index`, if applicable. + For the plotting to work appropriately, the resulting data after selection must contain + only two coordinates with len > 1. + Furthermore, these should be spatial coordinates (``x``, ``y``, or ``z``). + + Returns + ------- + matplotlib.axes._subplots.Axes + The supplied or created matplotlib axes. + """ + + sim_data = self.sim_data + sim_data.plot_field( + field_monitor_name=MODE_MONITOR_NAME, + field_name=field_name, + val=val, + scale=scale, + eps_alpha=eps_alpha, + robust=robust, + vmin=vmin, + vmax=vmax, + ax=ax, + **sel_kwargs, + ) + + def _validate_modes_size(self): + """Make sure that the total size of the modes fields is not too large.""" + monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME) + num_cells = self.simulation._monitor_num_cells(monitor) + # size in GB + total_size = monitor._storage_size_solver(num_cells=num_cells, tmesh=[]) / 1e9 + if total_size > MAX_MODES_DATA_SIZE_GB: + raise SetupError( + f"Mode solver has {total_size:.2f}GB of estimated storage, " + f"a maximum of {MAX_MODES_DATA_SIZE_GB:.2f}GB is allowed. Consider making the " + "mode plane smaller, or decreasing the resolution or number of requested " + "frequencies or modes." + ) + + def validate_pre_upload(self, source_required: bool = True): + self._validate_modes_size() + + @cached_property + def reduced_simulation_copy(self): + """Strip objects not used by the mode solver from simulation object. + This might significantly reduce upload time in the presence of custom mediums. + """ + + # we preserve extra cells along the normal direction to ensure there is enough data for + # subpixel + extended_grid = self._get_solver_grid(keep_additional_layers=True, truncate_symmetry=False) + grids_1d = extended_grid.boundaries + new_sim_box = Box.from_bounds( + rmin=(grids_1d.x[0], grids_1d.y[0], grids_1d.z[0]), + rmax=(grids_1d.x[-1], grids_1d.y[-1], grids_1d.z[-1]), + ) + + # remove PML, Absorers, etc, to avoid unnecessary cells + bspec = self.simulation.boundary_spec + + new_bspec_dict = {} + for axis in "xyz": + bcomp = bspec[axis] + for bside, sign in zip([bcomp.plus, bcomp.minus], "+-"): + if isinstance(bside, (PML, StablePML, Absorber)): + new_bspec_dict[axis + sign] = PECBoundary() + else: + new_bspec_dict[axis + sign] = bside + + new_bspec = BoundarySpec( + x=Boundary(plus=new_bspec_dict["x+"], minus=new_bspec_dict["x-"]), + y=Boundary(plus=new_bspec_dict["y+"], minus=new_bspec_dict["y-"]), + z=Boundary(plus=new_bspec_dict["z+"], minus=new_bspec_dict["z-"]), + ) + + # extract sub-simulation removing everything irrelevant + new_sim = self.simulation.subsection( + region=new_sim_box, + monitors=[], + sources=[], + grid_spec="identical", + boundary_spec=new_bspec, + remove_outside_custom_mediums=True, + remove_outside_structures=True, + ) + + return self.updated_copy(simulation=new_sim) diff --git a/tidy3d/components/mode/solver.py b/tidy3d/components/mode/solver.py new file mode 100644 index 000000000..35b03f46a --- /dev/null +++ b/tidy3d/components/mode/solver.py @@ -0,0 +1,745 @@ +"""Mode solver for propagating EM modes.""" +from typing import Tuple + +import numpy as np +import scipy.sparse as sp +import scipy.sparse.linalg as spl + +from ...components.types import Numpy, ModeSolverType, EpsSpecType +from ...components.base import Tidy3dBaseModel +from ...constants import ETA_0, C_0, fp_eps, pec_val +from .derivatives import create_d_matrices as d_mats +from .derivatives import create_s_matrices as s_mats +from .transforms import radial_transform, angled_transform + +# Consider vec to be complex if norm(vec.imag)/norm(vec) > TOL_COMPLEX +TOL_COMPLEX = fp_eps +# Tolerance for eigs +TOL_EIGS = fp_eps +# Tolerance for deciding on the matrix to be diagonal or tensorial +TOL_TENSORIAL = 1e-6 +# shift target neff by this value, both rel and abs, whichever results in larger shift +TARGET_SHIFT = 10 * fp_eps + + +class EigSolver(Tidy3dBaseModel): + """Interface for computing eigenvalues given permittivity and mode spec. + It's a collection of static methods. + """ + + @classmethod + def compute_modes( + cls, + eps_cross, + coords, + freq, + mode_spec, + symmetry=(0, 0), + direction="+", + ) -> Tuple[Numpy, Numpy, EpsSpecType]: + """ + Solve for the modes of a waveguide cross-section. + + Parameters + ---------- + eps_cross : array_like or tuple of array_like + Either a single 2D array defining the relative permittivity in the cross-section, + or nine 2D arrays defining the permittivity at the Ex, Ey, and Ez locations + of the Yee grid in the order xx, xy, xz, yx, yy, yz, zx, zy, zz. + coords : List[Numpy] + Two 1D arrays with each with size one larger than the corresponding axis of + ``eps_cross``. + Defines a (potentially non-uniform) Cartesian grid on which the modes are computed. + freq : float + (Hertz) Frequency at which the eigenmodes are computed. + mode_spec : ModeSpec + ``ModeSpec`` object containing specifications of the mode solver. + direction : Union["+", "-"] + Direction of mode propagation. + + Returns + ------- + Tuple[Numpy, Numpy, str] + The first array gives the E and H fields for all modes, the second one gives the complex + effective index. The last variable describes permittivity characterization on the mode + solver's plane ("diagonal", "tensorial_real", or "tensorial_complex"). + """ + + num_modes = mode_spec.num_modes + bend_radius = mode_spec.bend_radius + bend_axis = mode_spec.bend_axis + angle_theta = mode_spec.angle_theta + angle_phi = mode_spec.angle_phi + omega = 2 * np.pi * freq + k0 = omega / C_0 + + if isinstance(eps_cross, Numpy): + eps_xx, eps_xy, eps_xz, eps_yx, eps_yy, eps_yz, eps_zx, eps_zy, eps_zz = eps_cross + elif len(eps_cross) == 9: + eps_xx, eps_xy, eps_xz, eps_yx, eps_yy, eps_yz, eps_zx, eps_zy, eps_zz = ( + np.copy(e) for e in eps_cross + ) + else: + raise ValueError("Wrong input to mode solver pemittivity!") + + Nx, Ny = eps_xx.shape + N = eps_xx.size + + if len(coords[0]) != Nx + 1 or len(coords[1]) != Ny + 1: + raise ValueError("Mismatch between 'coords' and 'esp_cross' shapes.") + new_coords = [np.copy(c) for c in coords] + + """We work with full tensorial epsilon in mu to handle the most general cases that can + be introduced by coordinate transformations. In the solver, we distinguish the case when + these tensors are still diagonal, in which case the matrix for diagonalization has shape + (2N, 2N), and the full tensorial case, in which case it has shape (4N, 4N).""" + eps_tensor = np.zeros((3, 3, N), dtype=np.complex128) + mu_tensor = np.zeros((3, 3, N), dtype=np.complex128) + for row, eps_row in enumerate( + [[eps_xx, eps_xy, eps_xz], [eps_yx, eps_yy, eps_yz], [eps_zx, eps_zy, eps_zz]] + ): + mu_tensor[row, row, :] = 1.0 + for col, eps in enumerate(eps_row): + eps_tensor[row, col, :] = eps.ravel() + + # Get Jacobian of all coordinate transformations. Initialize as identity (same as mu so far) + jac_e = np.real(np.copy(mu_tensor)) + jac_h = np.real(np.copy(mu_tensor)) + + if bend_radius is not None: + new_coords, jac_e, jac_h = radial_transform(new_coords, bend_radius, bend_axis) + + if np.abs(angle_theta) > 0: + new_coords, jac_e_tmp, jac_h_tmp = angled_transform(new_coords, angle_theta, angle_phi) + jac_e = np.einsum("ij...,jp...->ip...", jac_e_tmp, jac_e) + jac_h = np.einsum("ij...,jp...->ip...", jac_h_tmp, jac_h) + + """We also need to keep track of the transformation of the k-vector. This is + the eigenvalue of the momentum operator assuming some sort of translational invariance and is + different from just the transformation of the derivative operator. For example, in a bent + waveguide, there is strictly speaking no k-vector in the original coordinates as the system + is not translationally invariant there. However, if we define kz = R k_phi, then the + effective index approaches that for a straight-waveguide in the limit of infinite radius. + Since we use w = R phi in the radial_transform, there is nothing else needed in the k transform. + For the angled_transform, the transformation between k-vectors follows from writing the field as + E' exp(i k_p w) in transformed coordinates, and identifying this with + E exp(i k_x x + i k_y y + i k_z z) in the original ones.""" + kxy = np.cos(angle_theta) ** 2 + kz = np.cos(angle_theta) * np.sin(angle_theta) + kp_to_k = np.array([kxy * np.sin(angle_phi), kxy * np.cos(angle_phi), kz]) + + # Transform epsilon and mu + jac_e_det = np.linalg.det(np.moveaxis(jac_e, [0, 1], [-2, -1])) + jac_h_det = np.linalg.det(np.moveaxis(jac_h, [0, 1], [-2, -1])) + eps_tensor = np.einsum("ij...,jp...->ip...", jac_e, eps_tensor) # J.dot(eps) + eps_tensor = np.einsum("ij...,pj...->ip...", eps_tensor, jac_e) # (J.dot(eps)).dot(J.T) + eps_tensor /= jac_e_det + mu_tensor = np.einsum("ij...,jp...->ip...", jac_h, mu_tensor) + mu_tensor = np.einsum("ij...,pj...->ip...", mu_tensor, jac_h) + mu_tensor /= jac_h_det + + """ Boundaries are imposed through the derivative matrices. The forward derivative matrices + always impose PEC boundary at the xmax and ymax interfaces, and on the xmin and ymin + interfaces unless PMC symmetry is present. If so, the PMC boundary is imposed through the + backward derivative matrices.""" + dmin_pmc = [sym == 1 for sym in symmetry] + + # Primal grid steps for E-field derivatives + dl_f = [new_cs[1:] - new_cs[:-1] for new_cs in new_coords] + # Dual grid steps for H-field derivatives + dl_tmp = [(dl[:-1] + dl[1:]) / 2 for dl in dl_f] + dl_b = [np.hstack((d1[0], d2)) for d1, d2 in zip(dl_f, dl_tmp)] + + # Derivative matrices with PEC boundaries by default and optional PMC at the near end + der_mats_tmp = d_mats((Nx, Ny), dl_f, dl_b, dmin_pmc) + + # PML matrices; do not impose PML on the bottom when symmetry present + dmin_pml = np.array(symmetry) == 0 + pml_mats = s_mats(omega, (Nx, Ny), mode_spec.num_pml, dl_f, dl_b, dmin_pml) + + # Add the PML on top of the derivatives; normalize by k0 to match the EM-possible notation + der_mats = [Smat.dot(Dmat) / k0 for Smat, Dmat in zip(pml_mats, der_mats_tmp)] + + # Determine initial guess value for the solver in transformed coordinates + if mode_spec.target_neff is None: + eps_physical = np.array(eps_cross) + eps_physical = eps_physical[np.abs(eps_physical) < np.abs(pec_val)] + n_max = np.sqrt(np.max(np.abs(eps_physical))) + target = n_max + else: + target = mode_spec.target_neff + target_neff_p = target / np.linalg.norm(kp_to_k) + + # shift target_neff slightly to avoid cases where the shiftted matrix is exactly singular + if abs(TARGET_SHIFT) > abs(target_neff_p * TARGET_SHIFT): + target_neff_p += TARGET_SHIFT + else: + target_neff_p *= 1 + TARGET_SHIFT + + # Solve for the modes + E, H, neff, keff, eps_spec = cls.solver_em( + Nx, + Ny, + eps_tensor, + mu_tensor, + der_mats, + num_modes, + target_neff_p, + mode_spec.precision, + direction, + ) + + # Transform back to original axes, E = J^T E' + E = np.sum(jac_e[..., None] * E[:, None, ...], axis=0) + E = E.reshape((3, Nx, Ny, 1, num_modes)) + H = np.sum(jac_h[..., None] * H[:, None, ...], axis=0) + H = H.reshape((3, Nx, Ny, 1, num_modes)) + neff = neff * np.linalg.norm(kp_to_k) + keff = keff * np.linalg.norm(kp_to_k) + + fields = np.stack((E, H), axis=0) + + if mode_spec.precision == "single": + # Recast to single precision which may have changed due to earlier manipulations + fields = fields.astype(np.complex64) + + return fields, neff + 1j * keff, eps_spec + + @classmethod + def solver_em( + cls, + Nx, + Ny, + eps_tensor, + mu_tensor, + der_mats, + num_modes, + neff_guess, + mat_precision, + direction, + ): + """Solve for the electromagnetic modes of a system defined by in-plane permittivity and + permeability and assuming translational invariance in the normal direction. + + Parameters + ---------- + Nx : int + Number of grids along x-direction. + Ny : int + Number of grids along y-direction. + eps_tensor : np.ndarray + Shape (3, 3, N), the permittivity tensor at every point in the plane. + mu_tensor : np.ndarray + Shape (3, 3, N), the permittivity tensor at every point in the plane. + der_mats : List[scipy.sparse.csr_matrix] + The sparce derivative matrices dxf, dxb, dyf, dyb, including the PML. + num_modes : int + Number of modes to solve for. + neff_guess : float + Initial guess for the effective index. + mat_precision : Union['single', 'double'] + Single or double-point precision in eigensolver. + direction : Union["+", "-"] + Direction of mode propagation. + + Returns + ------- + E : np.ndarray + Electric field of the eigenmodes, shape (3, N, num_modes). + H : np.ndarray + Magnetic field of the eigenmodes, shape (3, N, num_modes). + neff : np.ndarray + Real part of the effective index, shape (num_modes, ). + keff : np.ndarray + Imaginary part of the effective index, shape (num_modes, ). + eps_spec : Union["diagonal", "tensorial_real", "tensorial_complex"] + Permittivity characterization on the mode solver's plane. + """ + + # use a high-conductivity model for locations associated with a PEC + def conductivity_model_for_pec(eps, threshold=0.9 * pec_val): + """PEC entries associated with 'eps' are converted to a high-conductivity model.""" + eps = eps.astype(complex) + eps[eps <= threshold] = 1 + 1j * np.abs(pec_val) + return eps + + eps_tensor = conductivity_model_for_pec(eps_tensor) + + # Determine if ``eps`` and ``mu`` are diagonal or tensorial + off_diagonals = (np.ones((3, 3)) - np.eye(3)).astype(bool) + eps_offd = np.abs(eps_tensor[off_diagonals]) + mu_offd = np.abs(mu_tensor[off_diagonals]) + is_tensorial = np.any(eps_offd > TOL_TENSORIAL) or np.any(mu_offd > TOL_TENSORIAL) + + # initial vector for eigensolver in correct data type + vec_init = cls.set_initial_vec(Nx, Ny, is_tensorial=is_tensorial) + + # call solver + kwargs = { + "eps": eps_tensor, + "mu": mu_tensor, + "der_mats": der_mats, + "num_modes": num_modes, + "neff_guess": neff_guess, + "vec_init": vec_init, + "mat_precision": mat_precision, + } + + is_eps_complex = cls.isinstance_complex(eps_tensor) + + if not is_tensorial: + eps_spec = "diagonal" + E, H, neff, keff = cls.solver_diagonal(**kwargs) + if direction == "-": + H[0] *= -1 + H[1] *= -1 + E[2] *= -1 + + elif not is_eps_complex: + eps_spec = "tensorial_real" + E, H, neff, keff = cls.solver_tensorial(**kwargs, direction="+") + if direction == "-": + E = np.conj(E) + H = -np.conj(H) + + else: + eps_spec = "tensorial_complex" + E, H, neff, keff = cls.solver_tensorial(**kwargs, direction=direction) + + return E, H, neff, keff, eps_spec + + @classmethod + def matrix_data_type(cls, eps, mu, der_mats, mat_precision, is_tensorial): + """Determine data type that should be used for the matrix for diagonalization.""" + mat_dtype = np.float32 + # In tensorial case, even though the matrix can be real, the + # expected eigenvalue is purely imaginary. So for now we enforce + # the matrix to be complex type so that it will look for the right eigenvalues. + if is_tensorial: + mat_dtype = np.complex128 if mat_precision == "double" else np.complex64 + else: + # 1) check if complex or not + complex_solver = ( + cls.isinstance_complex(eps) + or cls.isinstance_complex(mu) + or np.any([cls.isinstance_complex(f) for f in der_mats]) + ) + # 2) determine precision + if complex_solver: + mat_dtype = np.complex128 if mat_precision == "double" else np.complex64 + else: + if mat_precision == "double": + mat_dtype = np.float64 + + return mat_dtype + + @classmethod + def trim_small_values(cls, mat: sp.csr_matrix, tol: float) -> sp.csr_matrix: + """Eliminate elements of matrix ``mat`` for which ``abs(element) / abs(max_element) < tol``, + or ``np.abs(mat_data) < tol``. This operates in-place on mat so there is no return. + """ + max_element = np.amax(np.abs(mat)) + mat.data *= np.logical_or(np.abs(mat.data) / max_element > tol, np.abs(mat.data) > tol) + mat.eliminate_zeros() + + @classmethod + def solver_diagonal(cls, eps, mu, der_mats, num_modes, neff_guess, vec_init, mat_precision): + """EM eigenmode solver assuming ``eps`` and ``mu`` are diagonal everywhere.""" + + # code associated with these options is included below in case it's useful in the future + enable_incidence_matrices = False + enable_preconditioner = False + analyze_conditioning = False + + def incidence_matrix_for_pec(eps_vec, threshold=0.9 * np.abs(pec_val)): + """Incidence matrix indicating non-PEC entries associated with 'eps_vec'.""" + nnz = eps_vec[np.abs(eps_vec) < threshold] + eps_nz = eps_vec.copy() + eps_nz[np.abs(eps_vec) >= threshold] = 0 + rows = np.arange(0, len(nnz)) + cols = np.argwhere(eps_nz).flatten() + dnz = sp.csr_matrix(([1] * len(nnz), (rows, cols)), shape=(len(rows), len(eps_vec))) + return dnz + + mode_solver_type = "diagonal" + N = eps.shape[-1] + + # Unpack eps, mu and derivatives + eps_xx = eps[0, 0, :] + eps_yy = eps[1, 1, :] + eps_zz = eps[2, 2, :] + mu_xx = mu[0, 0, :] + mu_yy = mu[1, 1, :] + mu_zz = mu[2, 2, :] + dxf, dxb, dyf, dyb = der_mats + + def any_pec(eps_vec, threshold=0.9 * np.abs(pec_val)): + """Check if there are any PEC values in the given permittivity array.""" + return np.any(np.abs(eps_vec) >= threshold) + + if any(any_pec(i) for i in [eps_xx, eps_yy, eps_zz]): + enable_preconditioner = True + + # Compute the matrix for diagonalization + inv_eps_zz = sp.spdiags(1 / eps_zz, [0], N, N) + inv_mu_zz = sp.spdiags(1 / mu_zz, [0], N, N) + + if enable_incidence_matrices: + dnz_xx, dnz_yy, dnz_zz = (incidence_matrix_for_pec(i) for i in [eps_xx, eps_yy, eps_zz]) + dnz = sp.block_diag((dnz_xx, dnz_yy), format="csr") + inv_eps_zz = (dnz_zz.T) * dnz_zz * inv_eps_zz * (dnz_zz.T) * dnz_zz + + p11 = -dxf.dot(inv_eps_zz).dot(dyb) + p12 = dxf.dot(inv_eps_zz).dot(dxb) + sp.spdiags(mu_yy, [0], N, N) + p21 = -dyf.dot(inv_eps_zz).dot(dyb) - sp.spdiags(mu_xx, [0], N, N) + p22 = dyf.dot(inv_eps_zz).dot(dxb) + q11 = -dxb.dot(inv_mu_zz).dot(dyf) + q12 = dxb.dot(inv_mu_zz).dot(dxf) + sp.spdiags(eps_yy, [0], N, N) + q21 = -dyb.dot(inv_mu_zz).dot(dyf) - sp.spdiags(eps_xx, [0], N, N) + q22 = dyb.dot(inv_mu_zz).dot(dxf) + + pmat = sp.bmat([[p11, p12], [p21, p22]]) + qmat = sp.bmat([[q11, q12], [q21, q22]]) + mat = pmat.dot(qmat) + + # Cast matrix to target data type + mat_dtype = cls.matrix_data_type(eps, mu, der_mats, mat_precision, is_tensorial=False) + mat = cls.type_conversion(mat, mat_dtype) + + # Trim small values in single precision case + if mat_precision == "single": + cls.trim_small_values(mat, tol=fp_eps) + + # Casting starting vector to target data type + vec_init = cls.type_conversion(vec_init, mat_dtype) + + # Starting eigenvalue guess in target data type + eig_guess = cls.type_conversion(np.array([-(neff_guess**2)]), mat_dtype)[0] + + if enable_incidence_matrices: + mat = dnz * mat * dnz.T + vec_init = dnz * vec_init + + if enable_preconditioner: + precon = sp.diags(1 / mat.diagonal()) + mat = mat * precon + else: + precon = None + + if analyze_conditioning: + aca = mat.conjugate().T * mat + aac = mat * mat.conjugate().T + diff = aca - aac + print(spl.norm(diff, ord=np.inf), spl.norm(aca, ord=np.inf), spl.norm(aac, ord=np.inf)) + print(spl.norm(diff, ord="fro"), spl.norm(aca, ord="fro"), spl.norm(aac, ord="fro")) + + # Call the eigensolver. The eigenvalues are -(neff + 1j * keff)**2 + vals, vecs = cls.solver_eigs( + mat, + num_modes, + vec_init, + guess_value=eig_guess, + mode_solver_type=mode_solver_type, + M=precon, + ) + + if enable_preconditioner: + vecs = precon * vecs + + if enable_incidence_matrices: + vecs = dnz.T * vecs + + neff, keff = cls.eigs_to_effective_index(vals, mode_solver_type) + + # Sort by descending neff + sort_inds = np.argsort(neff)[::-1] + neff = neff[sort_inds] + keff = keff[sort_inds] + vecs = vecs[:, sort_inds] + + # Field components from eigenvectors + Ex = vecs[:N, :] + Ey = vecs[N:, :] + + # Get the other field components + h_field = qmat.dot(vecs) + Hx = h_field[:N, :] / (1j * neff - keff) + Hy = h_field[N:, :] / (1j * neff - keff) + Hz = inv_mu_zz.dot(dxf.dot(Ey) - dyf.dot(Ex)) + Ez = inv_eps_zz.dot(dxb.dot(Hy) - dyb.dot(Hx)) + + # Bundle up + E = np.stack((Ex, Ey, Ez), axis=0) + H = np.stack((Hx, Hy, Hz), axis=0) + + # Return to standard H field units (see CEM notes for H normalization used in solver) + H *= -1j / ETA_0 + + return E, H, neff, keff + + @classmethod + def solver_tensorial( + cls, eps, mu, der_mats, num_modes, neff_guess, vec_init, mat_precision, direction + ): + """EM eigenmode solver assuming ``eps`` or ``mu`` have off-diagonal elements.""" + + mode_solver_type = "tensorial" + N = eps.shape[-1] + dxf, dxb, dyf, dyb = der_mats + + # Compute all blocks of the matrix for diagonalization + inv_eps_zz = sp.spdiags(1 / eps[2, 2, :], [0], N, N) + inv_mu_zz = sp.spdiags(1 / mu[2, 2, :], [0], N, N) + axax = -dxf.dot(sp.spdiags(eps[2, 0, :] / eps[2, 2, :], [0], N, N)) - sp.spdiags( + mu[1, 2, :] / mu[2, 2, :], [0], N, N + ).dot(dyf) + axay = -dxf.dot(sp.spdiags(eps[2, 1, :] / eps[2, 2, :], [0], N, N)) + sp.spdiags( + mu[1, 2, :] / mu[2, 2, :], [0], N, N + ).dot(dxf) + axbx = -dxf.dot(inv_eps_zz).dot(dyb) + sp.spdiags( + mu[1, 0, :] - mu[1, 2, :] * mu[2, 0, :] / mu[2, 2, :], [0], N, N + ) + axby = dxf.dot(inv_eps_zz).dot(dxb) + sp.spdiags( + mu[1, 1, :] - mu[1, 2, :] * mu[2, 1, :] / mu[2, 2, :], [0], N, N + ) + ayax = -dyf.dot(sp.spdiags(eps[2, 0, :] / eps[2, 2, :], [0], N, N)) + sp.spdiags( + mu[0, 2, :] / mu[2, 2, :], [0], N, N + ).dot(dyf) + ayay = -dyf.dot(sp.spdiags(eps[2, 1, :] / eps[2, 2, :], [0], N, N)) - sp.spdiags( + mu[0, 2, :] / mu[2, 2, :], [0], N, N + ).dot(dxf) + aybx = -dyf.dot(inv_eps_zz).dot(dyb) + sp.spdiags( + -mu[0, 0, :] + mu[0, 2, :] * mu[2, 0, :] / mu[2, 2, :], [0], N, N + ) + ayby = dyf.dot(inv_eps_zz).dot(dxb) + sp.spdiags( + -mu[0, 1, :] + mu[0, 2, :] * mu[2, 1, :] / mu[2, 2, :], [0], N, N + ) + bxbx = -dxb.dot(sp.spdiags(mu[2, 0, :] / mu[2, 2, :], [0], N, N)) - sp.spdiags( + eps[1, 2, :] / eps[2, 2, :], [0], N, N + ).dot(dyb) + bxby = -dxb.dot(sp.spdiags(mu[2, 1, :] / mu[2, 2, :], [0], N, N)) + sp.spdiags( + eps[1, 2, :] / eps[2, 2, :], [0], N, N + ).dot(dxb) + bxax = -dxb.dot(inv_mu_zz).dot(dyf) + sp.spdiags( + eps[1, 0, :] - eps[1, 2, :] * eps[2, 0, :] / eps[2, 2, :], [0], N, N + ) + bxay = dxb.dot(inv_mu_zz).dot(dxf) + sp.spdiags( + eps[1, 1, :] - eps[1, 2, :] * eps[2, 1, :] / eps[2, 2, :], [0], N, N + ) + bybx = -dyb.dot(sp.spdiags(mu[2, 0, :] / mu[2, 2, :], [0], N, N)) + sp.spdiags( + eps[0, 2, :] / eps[2, 2, :], [0], N, N + ).dot(dyb) + byby = -dyb.dot(sp.spdiags(mu[2, 1, :] / mu[2, 2, :], [0], N, N)) - sp.spdiags( + eps[0, 2, :] / eps[2, 2, :], [0], N, N + ).dot(dxb) + byax = -dyb.dot(inv_mu_zz).dot(dyf) + sp.spdiags( + -eps[0, 0, :] + eps[0, 2, :] * eps[2, 0, :] / eps[2, 2, :], [0], N, N + ) + byay = dyb.dot(inv_mu_zz).dot(dxf) + sp.spdiags( + -eps[0, 1, :] + eps[0, 2, :] * eps[2, 1, :] / eps[2, 2, :], [0], N, N + ) + + mat = sp.bmat( + [ + [axax, axay, axbx, axby], + [ayax, ayay, aybx, ayby], + [bxax, bxay, bxbx, bxby], + [byax, byay, bybx, byby], + ] + ) + + # The eigenvalues for the matrix above are 1j * (neff + 1j * keff) + # Multiply the matrix by -1j, so that eigenvalues are (neff + 1j * keff) + mat *= -1j + + # change matrix sign for backward direction + if direction == "-": + mat *= -1 + + # Cast matrix to target data type + mat_dtype = cls.matrix_data_type(eps, mu, der_mats, mat_precision, is_tensorial=True) + mat = cls.type_conversion(mat, mat_dtype) + + # Trim small values in single precision case + if mat_precision == "single": + cls.trim_small_values(mat, tol=fp_eps) + + # Casting starting vector to target data type + vec_init = cls.type_conversion(vec_init, mat_dtype) + + # Starting eigenvalue guess in target data type + eig_guess = cls.type_conversion(np.array([neff_guess]), mat_dtype)[0] + + # Call the eigensolver. + vals, vecs = cls.solver_eigs( + mat, + num_modes, + vec_init, + guess_value=eig_guess, + mode_solver_type=mode_solver_type, + ) + neff, keff = cls.eigs_to_effective_index(vals, mode_solver_type) + # Sort by descending real part + sort_inds = np.argsort(neff)[::-1] + neff = neff[sort_inds] + keff = keff[sort_inds] + vecs = vecs[:, sort_inds] + + # Field components from eigenvectors + Ex = vecs[:N, :] + Ey = vecs[N : 2 * N, :] + Hx = vecs[2 * N : 3 * N, :] + Hy = vecs[3 * N :, :] + + # Get the other field components + hxy_term = (-mu[2, 0, :] * Hx.T - mu[2, 1, :] * Hy.T).T + Hz = inv_mu_zz.dot(dxf.dot(Ey) - dyf.dot(Ex) + hxy_term) + exy_term = (-eps[2, 0, :] * Ex.T - eps[2, 1, :] * Ey.T).T + Ez = inv_eps_zz.dot(dxb.dot(Hy) - dyb.dot(Hx) + exy_term) + + # Bundle up + E = np.stack((Ex, Ey, Ez), axis=0) + H = np.stack((Hx, Hy, Hz), axis=0) + + # Return to standard H field units (see CEM notes for H normalization used in solver) + # The minus sign here is suspicious, need to check how modes are used in Mode objects + H *= -1j / ETA_0 + + return E, H, neff, keff + + @classmethod + def solver_eigs(cls, mat, num_modes, vec_init, guess_value=1.0, M=None, **kwargs): + """Find ``num_modes`` eigenmodes of ``mat`` cloest to ``guess_value``. + + Parameters + ---------- + mat : scipy.sparse matrix + Square matrix for diagonalization. + num_modes : int + Number of eigenmodes to compute. + guess_value : float, optional + """ + + values, vectors = spl.eigs( + mat, k=num_modes, sigma=guess_value, tol=TOL_EIGS, v0=vec_init, M=M + ) + return values, vectors + + @classmethod + def isinstance_complex(cls, vec_or_mat, tol=TOL_COMPLEX): + """Check if a numpy array or scipy csr_matrix has complex component by looking at + norm(x.imag)/norm(x)>TOL_COMPLEX + + Parameters + ---------- + vec_or_mat : Union[np.ndarray, sp.csr_matrix] + """ + + if isinstance(vec_or_mat, np.ndarray): + return np.linalg.norm(vec_or_mat.imag) / (np.linalg.norm(vec_or_mat) + fp_eps) > tol + if isinstance(vec_or_mat, sp.csr_matrix): + mat_norm = spl.norm(vec_or_mat) + mat_imag_norm = spl.norm(vec_or_mat.imag) + return mat_imag_norm / (mat_norm + fp_eps) > tol + + raise RuntimeError("Variable type should be either numpy array or scipy csr_matrix.") + + @classmethod + def type_conversion(cls, vec_or_mat, new_dtype): + """Convert vec_or_mat to new_type. + + Parameters + ---------- + vec_or_mat : Union[np.ndarray, sp.csr_matrix] + vec or mat to be converted. + new_dtype : Union[np.complex128, np.complex64, np.float64, np.float32] + Final type of vec or mat + + Returns + ------- + converted_vec_or_mat : Union[np.ndarray, sp.csr_matrix] + """ + + if new_dtype in {np.complex128, np.complex64}: + return vec_or_mat.astype(new_dtype) + if new_dtype in {np.float64, np.float32}: + converted_vec_or_mat = vec_or_mat.real + return converted_vec_or_mat.astype(new_dtype) + + raise RuntimeError("Unsupported new_type.") + + @classmethod + def set_initial_vec(cls, Nx, Ny, is_tensorial=False): + """Set initial vector for eigs: + 1) The field at x=0 and y=0 boundaries are set to 0. This should be + the case for PEC boundaries, but wouldn't hurt for non-PEC boundary; + 2) The vector is np.complex128 by default, and will be converted to + appropriate type afterwards. + + Parameters + ---------- + Nx : int + Number of grids along x-direction. + Ny : int + Number of grids along y-direction. + is_tensorial : bool + diagonal or tensorial eigenvalue problem. + """ + + # The size of the vector is len_multiplier * Nx * Ny + len_multiplier = 2 + if is_tensorial: + len_multiplier *= 2 + + # Initialize the vector + size = (Nx, Ny, len_multiplier) + rng = np.random.default_rng(0) + vec_init = rng.random(size) + 1j * rng.random(size) + + # Set values at the boundary to be 0 + if Nx > 1: + vec_init[0, :, :] = 0 + if Ny > 1: + vec_init[:, 0, :] = 0 + + # Concatenate the vector appropriately + vec_init = np.vstack(vec_init) + return vec_init.flatten("F") + + @classmethod + def eigs_to_effective_index(cls, eig_list: Numpy, mode_solver_type: ModeSolverType): + """Convert obtained eigenvalues to n_eff and k_eff. + + Parameters + ---------- + eig_list : Numpy + Array of eigenvalues + mode_solver_type : ModeSolverType + The type of mode solver problems + + Returns + ------- + Tuple[Numpy, Numpy] + n_eff and k_eff + """ + if eig_list.size == 0: + raise RuntimeError("Could not find any eigenmodes for this waveguide.") + + # for tensorial type, it's simply (neff + 1j * keff) + if mode_solver_type == "tensorial": + return np.real(eig_list), np.imag(eig_list) + + # for diagonal type, eigenvalues are -(neff + 1j * keff)**2 + if mode_solver_type == "diagonal": + vre, vim = -np.real(eig_list), -np.imag(eig_list) + # Real and imaginary part of the effective index + neff = np.sqrt(vre / 2 + np.sqrt(vre**2 + vim**2) / 2) + keff = vim / 2 / (neff + 1e-10) + return neff, keff + + raise RuntimeError(f"Unidentified 'mode_solver_type={mode_solver_type}'.") + + +def compute_modes(*args, **kwargs) -> Tuple[Numpy, Numpy, str]: + """A wrapper around ``EigSolver.compute_modes``, which is used in ``ModeSolver``.""" + return EigSolver.compute_modes(*args, **kwargs) diff --git a/tidy3d/components/mode/transforms.py b/tidy3d/components/mode/transforms.py new file mode 100644 index 000000000..44a0b752e --- /dev/null +++ b/tidy3d/components/mode/transforms.py @@ -0,0 +1,114 @@ +""" Coordinate transformations. + +The Jacobian of a transformation from coordinates r = (x, y, z) into coordinates +r' = (u, v, w) is defined as J_ij = dr'_i/dr_j. Here, z and w are the propagation axes in the +original and transformed planes, respectively, and the coords are only provided in (x, y) and +transformed to (u, v). The Yee grid positions also have to be taken into account. The Jacobian +for the transformation of eps and E is evaluated at the r' positions of E-field components. +Similarly, the jacobian for mu and H is evaluated at the r' positions of H-field components. +Currently, the half-step offset in w is ignored, which should be a pretty good approximation.""" + +import numpy as np + + +def radial_transform(coords, radius, bend_axis): + """Compute the new coordinates and the Jacobian of a polar coordinate transformation. After + offsetting the plane such that its center is a distance of ``radius`` away from the center of + curvature, we have, e.g. for ``bend_axis=='y'``: + + u = (x**2 + z**2) + v = y + w = R acos(x / u) + + These are all evaluated at z = 0 below. + + Parameters + ---------- + coords : tuple + A tuple of two arrays of size Nx + 1, Ny + 1, respectively. + radius : float + Radius of the bend. + bend_axis : 0 or 1 + Axis normal to the bend plane. + + Returns + ------- + new_coords: tuple + Transformed coordinates, same shape as ``coords``. + jac_e: np.ndarrray + Jacobian of the transformation at the E-field positions, shape ``(3, 3, Nx * Ny)``. + jac_h: np.ndarrray + Jacobian of the transformation at the H-field positions, shape ``(3, 3, Nx * Ny)``. + k_to_kp: np.ndarray + A matrix of shape (3, 3) that transforms the k-vector from the original coordinates to the + transformed ones. + """ + + Nx, Ny = coords[0].size - 1, coords[1].size - 1 + norm_axis = 0 if bend_axis == 1 else 1 + + # Center the new coordinates such that the radius is at the center of the plane + u = coords[0] + (norm_axis == 0) * (radius - coords[0][Nx // 2]) + v = coords[1] + (norm_axis == 1) * (radius - coords[1][Ny // 2]) + new_coords = (u, v) + + """The only nontrivial derivative is dwdz and it only depends on the coordinate in the + norm_axis direction (orthogonal to both bend_axis and z). We need to compute that derivative + at the En and Hn positions. + """ + dwdz_e = radius / new_coords[norm_axis][:-1] + dwdz_h = radius / (new_coords[norm_axis][:-1] + new_coords[norm_axis][1:]) * 2 + + jac_e = np.zeros((3, 3, Nx, Ny)) + jac_e[0, 0, :, :] = 1 + jac_e[1, 1, :, :] = 1 + jac_e[2, 2, :, :] = np.expand_dims(dwdz_e, axis=bend_axis) + jac_e = jac_e.reshape((3, 3, -1)) + + jac_h = np.zeros((3, 3, Nx, Ny)) + jac_h[0, 0, :, :] = 1 + jac_h[1, 1, :, :] = 1 + jac_h[2, 2, :, :] = np.expand_dims(dwdz_h, axis=bend_axis) + jac_h = jac_h.reshape((3, 3, -1)) + + return new_coords, jac_e, jac_h + + +def angled_transform(coords, angle_theta, angle_phi): + """Compute the new coordinates and the Jacobian for a transformation that "straightens" + an angled waveguide such that it is translationally invariant in w. The transformation is + u = x - tan(angle) * z + + Parameters + ---------- + coords : tuple + A tuple of two arrays of size Nx + 1, Ny + 1, respectively. + angle_theta : float, optional + (radian) Polar angle from the normal axis. + angle_phi : float, optional + (radian) Azimuth angle in the plane orthogonal to the normal axis. + + Returns + ------- + new_coords: tuple + Transformed coordinates, same shape as ``coords``. + jac_e: np.ndarrray + Jacobian of the transformation at the E-field positions, shape ``(3, 3, Nx * Ny)``. + jac_h: np.ndarrray + Jacobian of the transformation at the H-field positions, shape ``(3, 3, Nx * Ny)``. + """ + + Nx, Ny = coords[0].size - 1, coords[1].size - 1 + + # The new coordinates are exactly the same at z = 0 + new_coords = (np.copy(c) for c in coords) + + # The only nontrivial derivatives are dudz, dvdz and they are constant everywhere + jac = np.zeros((3, 3, Nx * Ny)) + jac[0, 0, :] = 1 + jac[1, 1, :] = 1 + jac[2, 2, :] = 1 + jac[0, 2, :] = -np.tan(angle_theta) * np.cos(angle_phi) + jac[1, 2, :] = -np.tan(angle_theta) * np.sin(angle_phi) + + return new_coords, jac, jac diff --git a/tidy3d/components/mode.py b/tidy3d/components/mode_spec.py similarity index 100% rename from tidy3d/components/mode.py rename to tidy3d/components/mode_spec.py diff --git a/tidy3d/components/monitor.py b/tidy3d/components/monitor.py index 2980559bb..6105d5c22 100644 --- a/tidy3d/components/monitor.py +++ b/tidy3d/components/monitor.py @@ -9,7 +9,7 @@ from .types import Literal, Direction, Coordinate, Axis, ObsGridArray, BoxSurface from .validators import assert_plane, validate_freqs_not_empty, validate_freqs_min from .base import cached_property, Tidy3dBaseModel, skip_if_fields_missing -from .mode import ModeSpec +from .mode_spec import ModeSpec from .apodization import ApodizationSpec from .medium import MediumType from .viz import ARROW_COLOR_MONITOR, ARROW_ALPHA diff --git a/tidy3d/components/source.py b/tidy3d/components/source.py index 804f6917f..025f387f7 100644 --- a/tidy3d/components/source.py +++ b/tidy3d/components/source.py @@ -20,7 +20,7 @@ from .data.validators import validate_no_nans from .data.data_array import TimeDataArray from .geometry.base import Box -from .mode import ModeSpec +from .mode_spec import ModeSpec from .viz import add_ax_if_none, PlotParams, plot_params_source from .viz import ARROW_COLOR_SOURCE, ARROW_ALPHA, ARROW_COLOR_POLARIZATION from ..constants import RADIAN, HERTZ, MICROMETER, GLANCING_CUTOFF diff --git a/tidy3d/plugins/mode/mode_solver.py b/tidy3d/plugins/mode/mode_solver.py index eabbd820d..e73b5d9ea 100644 --- a/tidy3d/plugins/mode/mode_solver.py +++ b/tidy3d/plugins/mode/mode_solver.py @@ -16,7 +16,7 @@ from ...components.geometry.base import Box from ...components.simulation import Simulation from ...components.grid.grid import Grid -from ...components.mode import ModeSpec +from ...components.mode_spec import ModeSpec from ...components.monitor import ModeSolverMonitor, ModeMonitor from ...components.medium import FullyAnisotropicMedium from ...components.source import ModeSource, SourceTime diff --git a/tidy3d/plugins/smatrix/smatrix.py b/tidy3d/plugins/smatrix/smatrix.py index b1945f2dc..b17e5c374 100644 --- a/tidy3d/plugins/smatrix/smatrix.py +++ b/tidy3d/plugins/smatrix/smatrix.py @@ -10,7 +10,7 @@ from ...constants import HERTZ from ...components.simulation import Simulation from ...components.geometry.base import Box -from ...components.mode import ModeSpec +from ...components.mode_spec import ModeSpec from ...components.monitor import ModeMonitor from ...components.source import ModeSource, GaussianPulse from ...components.data.sim_data import SimulationData diff --git a/tidy3d/plugins/waveguide/rectangular_dielectric.py b/tidy3d/plugins/waveguide/rectangular_dielectric.py index 40ca51b97..80056d14a 100644 --- a/tidy3d/plugins/waveguide/rectangular_dielectric.py +++ b/tidy3d/plugins/waveguide/rectangular_dielectric.py @@ -12,7 +12,7 @@ from ...components.geometry.polyslab import PolySlab from ...components.grid.grid_spec import GridSpec from ...components.medium import Medium, MediumType -from ...components.mode import ModeSpec +from ...components.mode_spec import ModeSpec from ...components.simulation import Simulation from ...components.source import ModeSource, GaussianPulse From be58f3e1e08a41ef968a0cd494e8c225355dd15b Mon Sep 17 00:00:00 2001 From: Tyler Hughes Date: Wed, 13 Mar 2024 15:28:59 -0400 Subject: [PATCH 2/2] mode solver refactor first pass --- docs/notebooks | 2 +- tests/test_components/test_mode_solver.py | 777 ++++++++++++++++++++++ tidy3d/components/mode/mode_solver.py | 649 ++++++++++++++++-- tidy3d/plugins/mode/web.py | 4 +- tidy3d/web/api/mode_core.py | 539 +++++++++++++++ 5 files changed, 1924 insertions(+), 47 deletions(-) create mode 100644 tests/test_components/test_mode_solver.py create mode 100644 tidy3d/web/api/mode_core.py diff --git a/docs/notebooks b/docs/notebooks index 0bf4aa017..56a67c5db 160000 --- a/docs/notebooks +++ b/docs/notebooks @@ -1 +1 @@ -Subproject commit 0bf4aa017670b4d593d4fb91390058ff025fc72d +Subproject commit 56a67c5db2754d636f71bf6b2c19bba5baf3a15a diff --git a/tests/test_components/test_mode_solver.py b/tests/test_components/test_mode_solver.py new file mode 100644 index 000000000..38b7f9ddc --- /dev/null +++ b/tests/test_components/test_mode_solver.py @@ -0,0 +1,777 @@ +import pytest +import responses +import numpy as np +import matplotlib.pyplot as plt +import pydantic.v1 as pydantic + +import tidy3d as td + +import tidy3d.plugins.mode.web as msweb +from tidy3d.components.mode import ModeSolver +from tidy3d.components.mode.mode_solver import MODE_MONITOR_NAME +from tidy3d.components.mode.derivatives import create_sfactor_b, create_sfactor_f +from tidy3d.components.mode.solver import compute_modes +from tidy3d.exceptions import SetupError +from ..utils import assert_log_level, log_capture # noqa: F401 +from tidy3d import ScalarFieldDataArray +from tidy3d.web.core.environment import Env + + +WG_MEDIUM = td.Medium(permittivity=4.0, conductivity=1e-4) +WAVEGUIDE = td.Structure(geometry=td.Box(size=(1.5, 100, 1)), medium=WG_MEDIUM) +PLANE = td.Box(center=(0, 0, 0), size=(5, 0, 5)) +SIM_SIZE = (4, 3, 3) +SRC = td.PointDipole( + center=(0, 0, 0), source_time=td.GaussianPulse(freq0=2e14, fwidth=1e13), polarization="Ex" +) + +PROJECT_NAME = "Mode Solver" +TASK_NAME = "Untitled" +MODESOLVER_NAME = "mode_solver" +PROJECT_ID = "Project-ID" +TASK_ID = "Task-ID" +SOLVER_ID = "Solver-ID" + + +@pytest.fixture +def mock_remote_api(monkeypatch): + def void(*args, **kwargs): + return None + + def mock_download(resource_id, remote_filename, to_file, *args, **kwargs): + simulation = td.Simulation( + size=SIM_SIZE, + grid_spec=td.GridSpec(wavelength=1.0), + structures=[WAVEGUIDE], + run_time=1e-12, + symmetry=(1, 0, -1), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + mode_spec = td.ModeSpec( + num_modes=3, + target_neff=2.0, + filter_pol="tm", + precision="double", + track_freq="lowest", + ) + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=PLANE, + mode_spec=mode_spec, + freqs=[td.C_0 / 1.0], + ) + ms.data_raw.to_file(to_file) + + from tidy3d.web.core import http_util as httputil + + monkeypatch.setattr(httputil, "api_key", lambda: "api_key") + monkeypatch.setattr(httputil, "get_version", lambda: td.version.__version__) + monkeypatch.setattr("tidy3d.web.api.mode.upload_file", void) + monkeypatch.setattr("tidy3d.web.api.mode.download_gz_file", mock_download) + monkeypatch.setattr("tidy3d.web.api.mode.download_file", mock_download) + + responses.add( + responses.GET, + f"{Env.current.web_api_endpoint}/tidy3d/project", + match=[responses.matchers.query_param_matcher({"projectName": PROJECT_NAME})], + json={"data": {"projectId": PROJECT_ID, "projectName": PROJECT_NAME}}, + status=200, + ) + + responses.add( + responses.POST, + f"{Env.current.web_api_endpoint}/tidy3d/modesolver/py", + match=[ + responses.matchers.json_params_matcher( + { + "projectId": PROJECT_ID, + "taskName": TASK_NAME, + "modeSolverName": MODESOLVER_NAME, + "fileType": "Gz", + "source": "Python", + "protocolVersion": td.version.__version__, + } + ) + ], + json={ + "data": { + "refId": TASK_ID, + "id": SOLVER_ID, + "status": "draft", + "createdAt": "2023-05-19T16:47:57.190Z", + "charge": 0, + "fileType": "Gz", + } + }, + status=200, + ) + + responses.add( + responses.GET, + f"{Env.current.web_api_endpoint}/tidy3d/modesolver/py/{TASK_ID}/{SOLVER_ID}", + json={ + "data": { + "refId": TASK_ID, + "id": SOLVER_ID, + "status": "success", + "createdAt": "2023-05-19T16:47:57.190Z", + "charge": 0, + "fileType": "Json", + } + }, + status=200, + ) + + responses.add( + responses.POST, + f"{Env.current.web_api_endpoint}/tidy3d/modesolver/py/{TASK_ID}/{SOLVER_ID}/run", + json={ + "data": { + "refId": TASK_ID, + "id": SOLVER_ID, + "status": "queued", + "createdAt": "2023-05-19T16:47:57.190Z", + "charge": 0, + "fileType": "Gz", + } + }, + status=200, + ) + + +def test_compute_modes(): + """Test direct call to `compute_modes`.""" + eps_cross = np.random.rand(10, 10) + coords = np.arange(11) + mode_spec = td.ModeSpec(num_modes=3, target_neff=2.0) + _ = compute_modes( + eps_cross=[eps_cross] * 9, + coords=[coords, coords], + freq=td.C_0 / 1.0, + mode_spec=mode_spec, + direction="-", + ) + + +def compare_colocation(ms): + """Compare mode-solver fields with colocation applied during run or post-run.""" + data_col = ms.solve() + ms_nocol = ms.updated_copy(colocate=False) + data = ms_nocol.solve() + data_at_boundaries = ms_nocol.sim_data.at_boundaries(MODE_MONITOR_NAME) + + for key, field in data_col.field_components.items(): + # Check the colocated data is the same + assert np.allclose(data_at_boundaries[key], field, atol=1e-7) + + # Also check coordinates + for dim, coords1 in field.coords.items(): + # Check that noncolocated data has one extra coordinate in the plane dimensions + if coords1.size > 1 and dim in "xyz": + coords2 = data.field_components[key].coords[dim] + assert coords1.size == coords2.size - 1 + + # Check that colocated coords are the same + assert np.allclose(coords1, data_at_boundaries[key].coords[dim]) + + +def verify_pol_fraction(ms): + """Verify that polarization fraction was successfully filtered.""" + pol_frac = ms.data.pol_fraction + pol_frac_wg = ms.data.pol_fraction_waveguide + filter_pol = ms.mode_spec.filter_pol + + # print(pol_frac.isel(mode_index=0)) + # print(pol_frac_wg.isel(mode_index=0)) + # import matplotlib.pyplot as plt + + # f, ax = plt.subplots(3, 3, tight_layout=True, figsize=(10, 6)) + # for mode_index in range(3): + # ms.plot_field("Ex", "abs", mode_index=mode_index, f=ms.freqs[0], ax=ax[mode_index, 0]) + # ms.plot_field("Ey", "abs", mode_index=mode_index, f=ms.freqs[0], ax=ax[mode_index, 1]) + # ms.plot_field("Ez", "abs", mode_index=mode_index, f=ms.freqs[0], ax=ax[mode_index, 2]) + # plt.show() + + if filter_pol is not None: + assert np.all(pol_frac[filter_pol].isel(mode_index=0) > 0.5) + other_pol = "te" if filter_pol == "tm" else "tm" + # There is no guarantee that the waveguide polarization fraction is also predominantly + # the same as the standard definition, but it is true in the cases we test here + assert np.all( + pol_frac_wg[filter_pol].isel(mode_index=0).values + > pol_frac_wg[other_pol].isel(mode_index=0).values + ) + + +def verify_dtype(ms): + """Verify that the returned fields have the correct dtype w.r.t. the specified precision.""" + + dtype = np.complex64 if ms.mode_spec.precision == "single" else np.complex128 + for field in ms.data.field_components.values(): + print(dtype, field.dtype, type(field.dtype)) + assert dtype == field.dtype + + +def check_ms_reduction(ms): + ms_red = ms.reduced_simulation_copy + grids_1d = ms._solver_grid.boundaries + grids_1d_red = ms_red._solver_grid.boundaries + assert np.allclose(grids_1d.x, grids_1d_red.x) + assert np.allclose(grids_1d.y, grids_1d_red.y) + assert np.allclose(grids_1d.z, grids_1d_red.z) + modes_red = ms.solve() + assert np.allclose(ms.data.n_eff.values, modes_red.n_eff.values) + + +def test_mode_solver_validation(): + """Test invalidate mode solver setups.""" + + simulation = td.Simulation( + size=SIM_SIZE, + grid_spec=td.GridSpec(wavelength=1.0), + run_time=1e-12, + ) + mode_spec = td.ModeSpec( + num_modes=1, + ) + + # frequency is too low + with pytest.raises(pydantic.ValidationError): + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=PLANE, + mode_spec=mode_spec, + freqs=[1.1], + direction="+", + ) + + # frequency not too low + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=PLANE, + mode_spec=mode_spec, + freqs=[1e12], + direction="+", + ) + + # mode data too large + simulation = td.Simulation( + size=SIM_SIZE, + grid_spec=td.GridSpec.uniform(dl=0.001), + run_time=1e-12, + ) + ms = ms.updated_copy(simulation=simulation, freqs=np.linspace(1e12, 2e12, 50)) + + with pytest.raises(SetupError): + ms.validate_pre_upload() + + +@pytest.mark.parametrize("group_index_step, log_level", ((1e-7, "WARNING"), (1e-5, "INFO"))) +def test_mode_solver_group_index_warning(group_index_step, log_level, log_capture): # noqa: F811 + """Test mode solver setups issuing warnings.""" + + simulation = td.Simulation( + size=SIM_SIZE, + grid_spec=td.GridSpec(wavelength=1.0), + run_time=1e-12, + ) + mode_spec = td.ModeSpec( + num_modes=1, + group_index_step=group_index_step, + ) + + _ = ModeSolver.from_simulation( + simulation=simulation, + plane=PLANE, + mode_spec=mode_spec, + freqs=[1e12], + direction="+", + ) + assert_log_level(log_capture, log_level) + + +@pytest.mark.parametrize("local", [True, False]) +@responses.activate +def test_mode_solver_simple(mock_remote_api, local): + """Simple mode solver run (with symmetry)""" + + simulation = td.Simulation( + size=SIM_SIZE, + grid_spec=td.GridSpec(wavelength=1.0), + structures=[WAVEGUIDE], + run_time=1e-12, + symmetry=(0, 0, 1), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + mode_spec = td.ModeSpec( + num_modes=3, + target_neff=2.0, + filter_pol="tm", + precision="double" if local else "single", + track_freq="lowest", + ) + if local: + freqs = [td.C_0 / 0.9, td.C_0 / 1.0, td.C_0 / 1.1] + else: + freqs = [td.C_0 / 1.0] + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=PLANE, + mode_spec=mode_spec, + freqs=freqs, + direction="-", + ) + + if local: + compare_colocation(ms) + verify_pol_fraction(ms) + verify_dtype(ms) + _ = ms.data.to_dataframe() + check_ms_reduction(ms) + + else: + _ = msweb.run(ms) + + # Testing issue 807 functions + freq0 = td.C_0 / 1.55 + source_time = td.GaussianPulse(freq0=freq0, fwidth=freq0 / 10) + nS_add_source = ms.sim_with_source(mode_index=0, direction="+", source_time=source_time) + nS_add_monitor = ms.sim_with_monitor(freqs=freqs, name="mode monitor") + nS_add_mode_solver_monitor = ms.sim_with_mode_solver_monitor(name="mode solver monitor") + assert len(nS_add_source.sources) == len(simulation.sources) + 1 + assert len(nS_add_monitor.monitors) == len(simulation.monitors) + 1 + assert len(nS_add_mode_solver_monitor.monitors) == len(simulation.monitors) + 1 + + +@pytest.mark.parametrize("local", [True, False]) +@responses.activate +def test_mode_solver_custom_medium(mock_remote_api, local, tmp_path): + """Test mode solver can work with custom medium. Consider a waveguide with varying + permittivity along x-direction. The value of n_eff at different x position should be + different. + """ + + # waveguide made of custom medium + x_custom = np.linspace(-0.6, 0.6, 2) + y_custom = [0] + z_custom = [0] + freq0 = td.C_0 / 1.0 + n = np.array([1.5, 5]) + n = n[:, None, None, None] + n_data = ScalarFieldDataArray(n, coords=dict(x=x_custom, y=y_custom, z=z_custom, f=[freq0])) + mat_custom = td.CustomMedium.from_nk(n_data, interp_method="nearest") + + waveguide = td.Structure(geometry=td.Box(size=(100, 0.5, 0.5)), medium=mat_custom) + simulation = td.Simulation( + size=(2, 2, 2), + grid_spec=td.GridSpec(wavelength=1.0), + structures=[waveguide], + run_time=1e-12, + ) + mode_spec = td.ModeSpec( + num_modes=1, + precision="double" if local else "single", + ) + + plane_left = td.Box(center=(-0.5, 0, 0), size=(0.9, 0, 0.9)) + plane_right = td.Box(center=(0.5, 0, 0), size=(0.9, 0, 0.9)) + + n_eff = [] + for plane in [plane_left, plane_right]: + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=plane, + mode_spec=mode_spec, + freqs=[freq0], + direction="+", + ) + modes = ms.solve() if local else msweb.run(ms) + n_eff.append(modes.n_eff.values) + + if local: + check_ms_reduction(ms) + + fname = str(tmp_path / "ms_custom_medium.hdf5") + ms.to_file(fname) + m2 = ModeSolver.from_file(fname) + assert m2 == ms + + if local: + assert n_eff[0] < 1.5 + assert n_eff[1] > 4 + assert n_eff[1] < 5 + + +def test_mode_solver_straight_vs_angled(): + """Compare results for a straight and angled nominally identical waveguides. + Note: results do not match perfectly because of the numerical grid. + """ + simulation = td.Simulation( + size=SIM_SIZE, + grid_spec=td.GridSpec.auto(wavelength=1.0, min_steps_per_wvl=16), + structures=[WAVEGUIDE], + run_time=1e-12, + symmetry=(0, 0, 1), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + mode_spec = td.ModeSpec(num_modes=5, group_index_step=True) + freqs = [td.C_0 / 0.9, td.C_0 / 1.0, td.C_0 / 1.1] + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=PLANE, + mode_spec=mode_spec, + freqs=freqs, + direction="-", + ) + + angle = np.pi / 6 + width, height = WAVEGUIDE.geometry.size[0], WAVEGUIDE.geometry.size[2] + vertices = np.array( + [[-width / 2, -100, 0], [width / 2, -100, 0], [width / 2, 100, 0], [-width / 2, 100, 0]] + ) + vertices = PLANE.rotate_points(vertices.T, axis=[0, 0, 1], angle=-angle).T + vertices = [verts[:2] for verts in vertices] + wg_angled = td.Structure( + geometry=td.PolySlab(vertices=vertices, slab_bounds=(-height / 2, height / 2)), + medium=WG_MEDIUM, + ) + mode_spec_angled = mode_spec.updated_copy(angle_theta=angle) + src_angled = td.ModeSource( + source_time=td.GaussianPulse(freq0=2e14, fwidth=1e13), + center=PLANE.center, + size=PLANE.size, + mode_spec=mode_spec_angled, + direction="-", + mode_index=0, + ) + sim_angled = simulation.updated_copy(structures=[wg_angled], sources=[src_angled]) + # sim_angled.plot(z=0) + # plt.show() + + ms_angled = ModeSolver.from_simulation( + simulation=sim_angled, + plane=PLANE, + mode_spec=mode_spec_angled, + freqs=freqs, + direction="-", + ) + + check_ms_reduction(ms) + check_ms_reduction(ms_angled) + + for key, val in ms.data.modes_info.items(): + tol = 1e-2 + if key == "TE (Ex) fraction": + tol = 0.1 + elif key == "wg TE fraction": + tol = 1.3e-2 + elif key == "mode area": + tol = 2.1e-2 + elif key == "dispersion (ps/(nm km))": + tol = 0.7 + # print( + # key, + # (np.abs(val - ms_angled.data.modes_info[key]) / np.abs(val)).values.max(), + # (np.abs(val - ms_angled.data.modes_info[key]) / np.abs(ms_angled.data.modes_info[key])).values.max(), + # ) + assert np.allclose(val, ms_angled.data.modes_info[key], rtol=tol) + + +def test_mode_solver_angle_bend(): + """Run mode solver with angle and bend and symmetry""" + simulation = td.Simulation( + size=SIM_SIZE, + grid_spec=td.GridSpec(wavelength=1.0), + structures=[WAVEGUIDE], + run_time=1e-12, + symmetry=(-1, 0, 1), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + mode_spec = td.ModeSpec( + num_modes=3, + target_neff=2.0, + bend_radius=3, + bend_axis=0, + angle_theta=np.pi / 3, + angle_phi=np.pi, + track_freq="highest", + ) + # put plane entirely in the symmetry quadrant rather than sitting on its center + plane = td.Box(center=(0, 0.5, 0), size=(1, 0, 1)) + ms = ModeSolver.from_simulation( + simulation=simulation, plane=plane, mode_spec=mode_spec, freqs=[td.C_0 / 1.0], direction="-" + ) + compare_colocation(ms) + verify_pol_fraction(ms) + verify_dtype(ms) + _ = ms.data.to_dataframe() + check_ms_reduction(ms) + + # Plot field + _, ax = plt.subplots(1) + ms.plot_field("Ex", ax=ax, mode_index=1) + plt.close() + + # Create source and monitor + st = td.GaussianPulse(freq0=1.0e12, fwidth=1.0e12) + _ = ms.to_source(source_time=st, direction="-") + _ = ms.to_monitor(freqs=np.array([1.0, 2.0]) * 1e12, name="mode_mnt") + + +def test_mode_solver_2D(): + """Run mode solver in 2D simulations.""" + mode_spec = td.ModeSpec( + num_modes=3, + filter_pol="te", + precision="double", + num_pml=(0, 10), + track_freq="central", + ) + simulation = td.Simulation( + size=(0, SIM_SIZE[1], SIM_SIZE[2]), + grid_spec=td.GridSpec(wavelength=1.0), + structures=[WAVEGUIDE], + run_time=1e-12, + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + ms = ModeSolver.from_simulation( + simulation=simulation, plane=PLANE, mode_spec=mode_spec, freqs=[td.C_0 / 1.0], direction="-" + ) + compare_colocation(ms) + verify_pol_fraction(ms) + verify_dtype(ms) + _ = ms.data.to_dataframe() + check_ms_reduction(ms) + + mode_spec = td.ModeSpec( + num_modes=3, + filter_pol="te", + precision="double", + num_pml=(10, 0), + ) + simulation = td.Simulation( + size=(SIM_SIZE[0], SIM_SIZE[1], 0), + grid_spec=td.GridSpec(wavelength=1.0), + structures=[WAVEGUIDE], + run_time=1e-12, + boundary_spec=td.BoundarySpec.pml(z=False), + sources=[SRC], + ) + ms = ModeSolver.from_simulation( + simulation=simulation, plane=PLANE, mode_spec=mode_spec, freqs=[td.C_0 / 1.0], direction="+" + ) + compare_colocation(ms) + # verify_pol_fraction(ms) + _ = ms.data.to_dataframe() + check_ms_reduction(ms) + + # The simulation and the mode plane are both 0D along the same dimension + simulation = td.Simulation( + size=PLANE.size, + grid_spec=td.GridSpec(wavelength=1.0), + run_time=1e-12, + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + ms = ModeSolver.from_simulation(simulation=simulation, plane=PLANE, mode_spec=mode_spec, freqs=[td.C_0 / 1.0]) + compare_colocation(ms) + verify_pol_fraction(ms) + check_ms_reduction(ms) + + +@pytest.mark.parametrize("local", [True, False]) +@responses.activate +def test_group_index(mock_remote_api, log_capture, local): # noqa: F811 + """Test group index and dispersion calculation""" + + simulation = td.Simulation( + size=(5, 5, 1), + grid_spec=td.GridSpec(wavelength=1.55), + structures=[ + td.Structure( + geometry=td.Box(size=(0.5, 0.22, td.inf)), medium=td.Medium(permittivity=3.48**2) + ) + ], + medium=td.Medium(permittivity=1.44**2), + run_time=1e-12, + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + mode_spec = td.ModeSpec( + num_modes=2, + target_neff=3.0, + precision="double" if local else "single", + track_freq="central", + ) + + if local: + freqs = [td.C_0 / 1.54, td.C_0 / 1.55, td.C_0 / 1.56] + else: + freqs = [td.C_0 / 1.0] + + # No group index calculation by default + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=td.Box(size=(td.inf, td.inf, 0)), + mode_spec=mode_spec, + freqs=freqs, + direction="-", + ) + modes = ms.solve() if local else msweb.run(ms) + if local: + assert modes.n_group is None + assert len(log_capture) == 1 + assert modes.dispersion is None + assert len(log_capture) == 2 + for log_msg in log_capture: + assert log_msg[0] == 30 + assert "ModeSpec" in log_msg[1] + _ = modes.n_group + assert len(log_capture) == 2 + _ = modes.dispersion + assert len(log_capture) == 2 + check_ms_reduction(ms) + + # Group index calculated + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=td.Box(size=(td.inf, td.inf, 0)), + mode_spec=mode_spec.copy(update={"group_index_step": True}), + freqs=freqs, + ) + modes = ms.solve() if local else msweb.run(ms) + if local: + assert (modes.n_group.sel(mode_index=0).values > 3.9).all() + assert (modes.n_group.sel(mode_index=0).values < 4.2).all() + assert (modes.n_group.sel(mode_index=1).values > 3.7).all() + assert (modes.n_group.sel(mode_index=1).values < 4.0).all() + assert (modes.dispersion.sel(mode_index=0).values > 1400).all() + assert (modes.dispersion.sel(mode_index=0).values < 1500).all() + assert (modes.dispersion.sel(mode_index=1).values > -16500).all() + assert (modes.dispersion.sel(mode_index=1).values < -15000).all() + check_ms_reduction(ms) + + +def test_pml_params(): + """Test that mode solver pml parameters are computed correctly. + Profiles start with H-field locations on both sides. On the max side, they also terminate with + an H-field location, i.e. the last E-field parameter is missing. + """ + omega = 1 + N = 100 + dls = np.ones((N,)) + n_pml = 12 + + # Normalized target is just third power scaling with position + # E-field locations for backward derivatives + target_profile = (np.arange(1, n_pml + 1) / n_pml) ** 3 + target_profile = target_profile / target_profile[0] + sf_b = create_sfactor_b(omega, dls, N, n_pml, dmin_pml=True) + assert np.allclose(sf_b[:n_pml] / sf_b[n_pml - 1], target_profile[::-1]) + assert np.allclose(sf_b[N - n_pml + 1 :] / sf_b[N - n_pml + 1], target_profile[:-1]) + + # H-field locations for backward derivatives + target_profile = (np.arange(0.5, n_pml + 0.5, 1) / n_pml) ** 3 + target_profile = target_profile / target_profile[0] + sf_f = create_sfactor_f(omega, dls, N, n_pml, dmin_pml=True) + assert np.allclose(sf_f[:n_pml] / sf_f[n_pml - 1], target_profile[::-1]) + assert np.allclose(sf_f[N - n_pml :] / sf_f[N - n_pml], target_profile) + + +def test_mode_solver_nan_pol_fraction(): + """Test mode solver when eigensolver returns 0 for some modes.""" + wg = td.Structure(geometry=td.Box(size=(0.5, 100, 0.22)), medium=td.Medium(permittivity=12)) + + simulation = td.Simulation( + medium=td.Medium(permittivity=2), + size=SIM_SIZE, + grid_spec=td.GridSpec.auto(wavelength=1.55, min_steps_per_wvl=15), + structures=[wg], + run_time=1e-12, + symmetry=(0, 0, 1), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + + mode_spec = td.ModeSpec( + num_modes=10, + target_neff=3.48, + filter_pol="tm", + precision="single", + track_freq="central", + ) + + freqs = [td.C_0 / 1.55] + + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=td.Box(center=(0, 0, 0), size=(2, 0, 1.1)), + mode_spec=mode_spec, + freqs=freqs, + direction="-", + ) + + md = ms.solve() + check_ms_reduction(ms) + + assert list(np.where(np.isnan(md.pol_fraction.te))[1]) == [8, 9] + + +def test_mode_solver_method_defaults(): + """Test that changes to mode solver default values in methods work.""" + + simulation = td.Simulation( + medium=td.Medium(permittivity=2), + size=SIM_SIZE, + grid_spec=td.GridSpec.auto(wavelength=1.55, min_steps_per_wvl=15), + run_time=1e-12, + symmetry=(0, 0, 1), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + + mode_spec = td.ModeSpec( + num_modes=10, + target_neff=3.48, + filter_pol="tm", + precision="single", + track_freq="central", + ) + + freqs = [td.C_0 / 1.55] + + ms = ModeSolver.from_simulation( + simulation=simulation, + plane=td.Box(center=(0, 0, 0), size=(2, 0, 1.1)), + mode_spec=mode_spec, + freqs=freqs, + direction="-", + ) + + # test defaults + st = td.GaussianPulse(freq0=1.0e12, fwidth=1.0e12) + + src = ms.to_source(source_time=st) + assert src.direction == ms.direction + + src = ms.to_source(source_time=st, direction="+") + assert src.direction != ms.direction + + mnt = ms.to_monitor(name="mode_mnt") + assert np.allclose(mnt.freqs, ms.freqs) + + mnt = ms.to_monitor(name="mode_mnt", freqs=[2e14]) + assert not np.allclose(mnt.freqs, ms.freqs) + + # TODO: fix these + + # sim = ms.sim_with_source(source_time=st) + # assert sim.sources[-1].direction == ms.direction + + # sim = ms.sim_with_monitor(name="test") + # assert np.allclose(sim.monitors[-1].freqs, ms.freqs) diff --git a/tidy3d/components/mode/mode_solver.py b/tidy3d/components/mode/mode_solver.py index 4b25c1b18..e733c3a27 100644 --- a/tidy3d/components/mode/mode_solver.py +++ b/tidy3d/components/mode/mode_solver.py @@ -15,8 +15,12 @@ from ...components.base import Tidy3dBaseModel, cached_property, skip_if_fields_missing from ...components.boundary import PECBoundary, BoundarySpec, Boundary, PML, StablePML, Absorber from ...components.geometry.base import Box +from ...components.structure import Structure from ...components.simulation import Simulation -from ...components.grid.grid import Grid +from ...components.grid.grid import Coords1D, Grid, Coords +from ...components.boundary import BoundarySpec, BlochBoundary, PECBoundary, PMCBoundary, Periodic, Boundary +from ...components.medium import AbstractCustomMedium, Medium2D +from ...components.grid.grid_spec import GridSpec from ...components.mode_spec import ModeSpec from ...components.monitor import ModeSolverMonitor, ModeMonitor from ...components.medium import FullyAnisotropicMedium @@ -27,10 +31,79 @@ from ...components.data.data_array import FreqModeDataArray from ...components.data.sim_data import SimulationData from ...components.data.monitor_data import ModeSolverData +from ...components.viz import add_ax_if_none, equal_aspect + +from ...components.simulation import Simulation + +from ...components.base import cached_property +from ...components.base import skip_if_fields_missing +from ...components.validators import assert_objects_in_sim_bounds +from ...components.validators import validate_mode_objects_symmetry +from ...components.geometry.base import Geometry, Box +from ...components.geometry.mesh import TriangleMesh +from ...components.geometry.utils import flatten_groups, traverse_geometries +from ...components.geometry.utils_2d import get_bounds, set_bounds, get_thickened_geom, subdivide +from ...components.types import Ax, FreqBound, Axis, annotate_type, InterpMethod, Symmetry +from ...components.types import Literal, TYPE_TAG_STR +from ...components.grid.grid import Coords1D, Grid, Coords +from ...components.grid.grid_spec import GridSpec, UniformGrid, AutoGrid, CustomGrid +from ...components.medium import MediumType, AbstractMedium +from ...components.medium import AbstractCustomMedium, Medium2D +from ...components.medium import AnisotropicMedium, FullyAnisotropicMedium, AbstractPerturbationMedium +from ...components.boundary import BoundarySpec, BlochBoundary, PECBoundary, PMCBoundary, Periodic, Boundary +from ...components.boundary import PML, StablePML, Absorber, AbsorberSpec +from ...components.structure import Structure +from ...components.source import SourceType, PlaneWave, GaussianBeam, AstigmaticGaussianBeam, CustomFieldSource +from ...components.source import CustomCurrentSource, CustomSourceTime, ContinuousWave +from ...components.source import TFSF, Source, ModeSource +from ...components.medium import Medium, MediumType3D +from ...components.monitor import MonitorType, Monitor, FreqMonitor, SurfaceIntegrationMonitor +from ...components.monitor import AbstractModeMonitor, FieldMonitor, TimeMonitor +from ...components.monitor import PermittivityMonitor, DiffractionMonitor, AbstractFieldProjectionMonitor +from ...components.monitor import FieldProjectionAngleMonitor, FieldProjectionKSpaceMonitor +from ...components.data.dataset import Dataset +from ...components.data.data_array import SpatialDataArray +from ...components.viz import add_ax_if_none, equal_aspect +from ...components.scene import Scene + +from ...components.viz import PlotParams +from ...components.viz import plot_params_pml, plot_params_override_structures +from ...components.viz import plot_params_pec, plot_params_pmc, plot_params_bloch, plot_sim_3d + from ...components.validators import validate_freqs_min, validate_freqs_not_empty from ...exceptions import ValidationError, SetupError -from ...constants import C_0 +from ...constants import C_0, fp_eps + + +# TODO: dont need all of these + +# minimum number of grid points allowed per central wavelength in a medium +MIN_GRIDS_PER_WVL = 6.0 + +# maximum number of sources +MAX_NUM_SOURCES = 1000 + +# maximum numbers of simulation parameters +MAX_TIME_STEPS = 1e7 +WARN_TIME_STEPS = 1e6 +MAX_GRID_CELLS = 20e9 +MAX_CELLS_TIMES_STEPS = 1e16 +WARN_MONITOR_DATA_SIZE_GB = 10 +MAX_MONITOR_INTERNAL_DATA_SIZE_GB = 50 +MAX_SIMULATION_DATA_SIZE_GB = 50 +WARN_MODE_NUM_CELLS = 1e5 + +# number of grid cells at which we warn about slow Simulation.epsilon() +NUM_CELLS_WARN_EPSILON = 100_000_000 +# number of structures at which we warn about slow Simulation.epsilon() +NUM_STRUCTURES_WARN_EPSILON = 10_000 + +# height of the PML plotting boxes along any dimensions where sim.size[dim] == 0 +PML_HEIGHT_FOR_0_DIMS = 0.02 + +# END todo + # Importing the local solver may not work if e.g. scipy is not installed IMPORT_ERROR_MSG = """Could not import local solver, 'ModeSolver' objects can still be constructed @@ -54,7 +127,7 @@ MAX_MODES_DATA_SIZE_GB = 20 -class ModeSolver(AbstractSimulation): +class ModeSolver(Simulation): """ Interface for solving electromagnetic eigenmodes in a 2D plane with translational invariance in the third dimension. @@ -73,6 +146,21 @@ class ModeSolver(AbstractSimulation): * `Prelude to Integrated Photonics Simulation: Mode Injection `_ """ + # TODO: copied from simulation + grid_spec: GridSpec = pydantic.Field( + GridSpec(), + title="Grid Specification", + description="Specifications for the simulation grid along each of the three directions.", + ) + + # TODO: copied from simulation + boundary_spec: BoundarySpec = pydantic.Field( + BoundarySpec(), + title="Boundaries", + description="Specification of boundary conditions along each dimension. If ``None``, " + "PML boundary conditions are applied on all sides.", + ) + plane: Box = pydantic.Field( ..., title="Plane", description="Cross-sectional plane in which the mode will be computed." ) @@ -111,17 +199,36 @@ def is_plane(cls, val): _freqs_not_empty = validate_freqs_not_empty() _freqs_lower_bound = validate_freqs_min() - @pydantic.validator("plane", always=True) - @skip_if_fields_missing(["simulation"]) - def plane_in_sim_bounds(cls, val, values): - """Check that the plane is at least partially inside the simulation bounds.""" - sim_center = values.get("simulation").center - sim_size = values.get("simulation").size - sim_box = Box(size=sim_size, center=sim_center) - - if not sim_box.intersects(val): - raise SetupError("'ModeSolver.plane' must intersect 'ModeSolver.simulation'.") - return val + # TODO: add validator back + + # @pydantic.validator("plane", always=True) + # @skip_if_fields_missing(["simulation"]) + # def plane_in_sim_bounds(cls, val, values): + # """Check that the plane is at least partially inside the simulation bounds.""" + # sim_center = values.get("simulation").center + # sim_size = values.get("simulation").size + # sim_box = Box(size=sim_size, center=sim_center) + + # if not sim_box.intersects(val): + # raise SetupError("'ModeSolver.plane' must intersect 'ModeSolver.simulation'.") + # return val + + # TODO: fix this up + @property + def simulation(self) -> Simulation: + kwargs = self.dict(exclude={"type"}) + return Simulation(**kwargs) + + + # TODO: fix this up + @classmethod + def from_simulation(cls, simulation: Simulation, **kwargs): + kwargs.update(simulation.dict()) + kwargs = {key: val for key, val in kwargs.items() if key in cls.__fields__} + for key in ('type', 'sources'): + kwargs.pop(key) + + return cls(**kwargs) @cached_property def normal_axis(self) -> Axis: @@ -131,13 +238,338 @@ def normal_axis(self) -> Axis: @cached_property def solver_symmetry(self) -> Tuple[Symmetry, Symmetry]: """Get symmetry for solver for propagation along self.normal axis.""" - mode_symmetry = list(self.simulation.symmetry) + mode_symmetry = list(self.symmetry) for dim in range(3): - if self.simulation.center[dim] != self.plane.center[dim]: + if self.center[dim] != self.plane.center[dim]: mode_symmetry[dim] = 0 _, solver_sym = self.plane.pop_axis(mode_symmetry, axis=self.normal_axis) return solver_sym + + # TODO: copied from Simulation, fix + @cached_property + def _periodic(self) -> Tuple[bool, bool, bool]: + """For each dimension, ``True`` if periodic/Bloch boundaries and ``False`` otherwise. + We check on both sides but in practice there should be no cases in which a periodic/Bloch + BC is on one side only. This is explicitly validated for Bloch, and implicitly done for + periodic, in which case we allow PEC/PMC on the other side, but we replace the periodic + boundary with another PEC/PMC plane upon initialization.""" + periodic = [] + for bcs_1d in self.boundary_spec.to_list: + periodic.append(all(isinstance(bcs, (Periodic, BlochBoundary)) for bcs in bcs_1d)) + return periodic + + # TODO: copied from Simulation, fix + def _discretize_inds_monitor(self, monitor: Monitor): + """Start and stopping indexes for the cells where data needs to be recorded to fully cover + a ``monitor``. This is used during the solver run. The final grid on which a monitor data + lives is computed in ``discretize_monitor``, with the difference being that 0-sized + dimensions of the monitor or the simulation are snapped in post-processing.""" + + # Expand monitor size slightly to break numerical precision in favor of always having + # enough data to span the full monitor. + expand_size = [size + fp_eps if size > fp_eps else size for size in monitor.size] + box_expanded = Box(center=monitor.center, size=expand_size) + # Discretize without extension for now + span_inds = np.array(self.grid.discretize_inds(box_expanded, extend=False)) + + if any(ind[0] >= ind[1] for ind in span_inds): + # At least one dimension has no indexes inside the grid, e.g. monitor is entirely + # outside of the grid + return span_inds + + # Now add extensions, which are specific for monitors and are determined such that data + # colocated to grid boundaries can be interpolated anywhere inside the monitor. + # We always need to expand on the right. + span_inds[:, 1] += 1 + # Non-colocating monitors also need to expand on the left. + if not monitor.colocate: + span_inds[:, 0] -= 1 + return span_inds + + # TODO: copied from simulation + def _subgrid(self, span_inds: np.ndarray, grid: Grid = None): + """Take a subgrid of the simulation grid with cell span defined by ``span_inds`` along the + three dimensions. Optionally, a grid different from the simulation grid can be provided. + The ``span_inds`` can also extend beyond the grid, in which case the grid is padded based + on the boundary conditions of the simulation along the different dimensions.""" + + if not grid: + grid = self.grid + + boundary_dict = {} + for idim, (dim, periodic) in enumerate(zip("xyz", self._periodic)): + ind_beg, ind_end = span_inds[idim] + # ind_end + 1 because we are selecting cell boundaries not cells + boundary_dict[dim] = grid.extended_subspace(idim, ind_beg, ind_end + 1, periodic) + return Grid(boundaries=Coords(**boundary_dict)) + + + # TODO: copied from simulation + @cached_property + def num_pml_layers(self) -> List[Tuple[float, float]]: + """Number of absorbing layers in all three axes and directions (-, +). + + Returns + ------- + List[Tuple[float, float]] + List containing the number of absorber layers in - and + boundaries. + """ + num_layers = [[0, 0], [0, 0], [0, 0]] + + for idx_i, boundary1d in enumerate(self.boundary_spec.to_list): + for idx_j, boundary in enumerate(boundary1d): + if isinstance(boundary, (PML, StablePML, Absorber)): + num_layers[idx_i][idx_j] = boundary.num_layers + + return num_layers + + # TODO: copied from simulation + def _volumetric_structures_grid(self, grid: Grid) -> Tuple[Structure]: + """Generate a tuple of structures wherein any 2D materials are converted to 3D + volumetric equivalents, using ``grid`` as the simulation grid.""" + + if not any(isinstance(medium, Medium2D) for medium in self.scene.mediums): + return self.structures + + def get_dls(geom: Geometry, axis: Axis, num_dls: int) -> List[float]: + """Get grid size around the 2D material.""" + dls = self._discretize_grid(Box.from_bounds(*geom.bounds), grid=grid).sizes.to_list[ + axis + ] + # When 1 dl is requested it is assumed that only an approximate value is needed + # before the 2D material has been snapped to the grid + if num_dls == 1: + return [np.mean(dls)] + + # When 2 dls are requested the 2D geometry should have been snapped to grid, + # so this represents the exact adjacent grid spacing + if len(dls) != num_dls: + raise Tidy3dError( + "Failed to detect grid size around the 2D material. " + "Can't generate volumetric equivalent for this simulation. " + "If you received this error, please create an issue in the Tidy3D " + "github repository." + ) + return dls + + def snap_to_grid(geom: Geometry, axis: Axis) -> Geometry: + """Snap a 2D material to the Yee grid.""" + new_centers = self._discretize_grid( + Box.from_bounds(*geom.bounds), grid=grid + ).boundaries.to_list[axis] + new_center = new_centers[np.argmin(abs(new_centers - get_bounds(geom, axis)[0]))] + return set_bounds(geom, (new_center, new_center), axis) + + # Begin volumetric structures grid + # For 1D and 2D simulations, a nonzero size is needed for the polygon operations in subdivide + placeholder_size = tuple(i if i > 0 else inf for i in self.geometry.size) + simulation_placeholder_geometry = self.geometry.updated_copy( + center=self.geometry.center, size=placeholder_size + ) + + simulation_background = Structure( + geometry=simulation_placeholder_geometry, medium=self.medium + ) + background_structures = [simulation_background] + new_structures = [] + for structure in self.structures: + if not isinstance(structure.medium, Medium2D): + # found a 3D material; keep it + background_structures.append(structure) + new_structures.append(structure) + continue + # otherwise, found a 2D material; replace it with volumetric equivalent + axis = structure.geometry._normal_2dmaterial + geometry = structure.geometry + + # subdivide + avg_axis_dl = get_dls(geometry, axis, 1)[0] + subdivided_geometries = subdivide(geometry, axis, avg_axis_dl, background_structures) + # Create and add volumetric equivalents + background_structures_temp = [] + for subdivided_geometry in subdivided_geometries: + # Snap to the grid and create volumetric equivalent + snapped_geometry = snap_to_grid(subdivided_geometry[0], axis) + snapped_center = get_bounds(snapped_geometry, axis)[0] + dls = get_dls(get_thickened_geom(snapped_geometry, axis, avg_axis_dl), axis, 2) + adjacent_media = [subdivided_geometry[1].medium, subdivided_geometry[2].medium] + + # Create the new volumetric medium + new_medium = structure.medium.volumetric_equivalent( + axis=axis, adjacent_media=adjacent_media, adjacent_dls=dls + ) + + new_bounds = (snapped_center - dls[0] / 2, snapped_center + dls[1] / 2) + temp_geometry = set_bounds(snapped_geometry, bounds=new_bounds, axis=axis) + temp_structure = structure.updated_copy(geometry=temp_geometry, medium=new_medium) + + if structure.medium.is_pec: + pec_delta = fp_eps * max(np.abs(snapped_center), 1.0) + new_bounds = (snapped_center - pec_delta, snapped_center + pec_delta) + new_geometry = set_bounds(snapped_geometry, bounds=new_bounds, axis=axis) + new_structure = structure.updated_copy(geometry=new_geometry, medium=new_medium) + + new_structures.append(new_structure) + background_structures_temp.append(temp_structure) + + background_structures += background_structures_temp + + return tuple(new_structures) + + # TODO: copied from simulation + @cached_property + def volumetric_structures(self) -> Tuple[Structure]: + """Generate a tuple of structures wherein any 2D materials are converted to 3D + volumetric equivalents.""" + return self._volumetric_structures_grid(self.grid) + + # TODO: copied from simulation + def epsilon_on_grid( + self, + grid: Grid, + coord_key: str = "centers", + freq: float = None, + ) -> xr.DataArray: + """Get array of permittivity at a given freq on a given grid. + + Parameters + ---------- + grid : :class:`.Grid` + Grid specifying where to measure the permittivity. + coord_key : str = 'centers' + Specifies at what part of the grid to return the permittivity at. + Accepted values are ``{'centers', 'boundaries', 'Ex', 'Ey', 'Ez', 'Exy', 'Exz', 'Eyx', + 'Eyz', 'Ezx', Ezy'}``. The field values (eg. ``'Ex'``) correspond to the corresponding field + locations on the yee lattice. If field values are selected, the corresponding diagonal + (eg. ``eps_xx`` in case of ``'Ex'``) or off-diagonal (eg. ``eps_xy`` in case of ``'Exy'``) epsilon + component from the epsilon tensor is returned. Otherwise, the average of the main + values is returned. + freq : float = None + The frequency to evaluate the mediums at. + If not specified, evaluates at infinite frequency. + Returns + ------- + xarray.DataArray + Datastructure containing the relative permittivity values and location coordinates. + For details on xarray DataArray objects, + refer to `xarray's Documentation `_. + """ + + grid_cells = np.prod(grid.num_cells) + num_structures = len(self.structures) + if grid_cells > NUM_CELLS_WARN_EPSILON: + log.warning( + f"Requested grid contains {int(grid_cells):.2e} grid cells. " + "Epsilon calculation may be slow." + ) + if num_structures > NUM_STRUCTURES_WARN_EPSILON: + log.warning( + f"Simulation contains {num_structures:.2e} structures. " + "Epsilon calculation may be slow." + ) + + def get_eps(structure: Structure, frequency: float, coords: Coords): + """Select the correct epsilon component if field locations are requested.""" + if coord_key[0] != "E": + return np.mean(structure.eps_diagonal(frequency, coords), axis=0) + row = ["x", "y", "z"].index(coord_key[1]) + if len(coord_key) == 2: # diagonal component in case of Ex, Ey, and Ez + col = row + else: # off-diagonal component in case of Exy, Exz, Eyx, etc + col = ["x", "y", "z"].index(coord_key[2]) + return structure.eps_comp(row, col, frequency, coords) + + def make_eps_data(coords: Coords): + """returns epsilon data on grid of points defined by coords""" + arrays = (np.array(coords.x), np.array(coords.y), np.array(coords.z)) + eps_background = get_eps( + structure=self.scene.background_structure, frequency=freq, coords=coords + ) + shape = tuple(len(array) for array in arrays) + eps_array = eps_background * np.ones(shape, dtype=complex) + # replace 2d materials with volumetric equivalents + with log as consolidated_logger: + for structure in self.volumetric_structures: + # Indexing subset within the bounds of the structure + + inds = structure.geometry._inds_inside_bounds(*arrays) + + # Get permittivity on meshgrid over the reduced coordinates + coords_reduced = tuple(arr[ind] for arr, ind in zip(arrays, inds)) + if any(coords.size == 0 for coords in coords_reduced): + continue + + red_coords = Coords(**dict(zip("xyz", coords_reduced))) + eps_structure = get_eps(structure=structure, frequency=freq, coords=red_coords) + + if structure.medium.nonlinear_spec is not None: + consolidated_logger.warning( + "Evaluating permittivity of a nonlinear " + "medium ignores the nonlinearity." + ) + + if isinstance(structure.geometry, TriangleMesh): + consolidated_logger.warning( + "Client-side permittivity of a 'TriangleMesh' may be " + "inaccurate if the mesh is not unionized. We recommend unionizing " + "all meshes before import. A 'PermittivityMonitor' can be used to " + "obtain the true permittivity and check that the surface mesh is " + "loaded correctly." + ) + + # Update permittivity array at selected indexes within the geometry + is_inside = structure.geometry.inside_meshgrid(*coords_reduced) + eps_array[inds][is_inside] = (eps_structure * is_inside)[is_inside] + + coords = dict(zip("xyz", arrays)) + return xr.DataArray(eps_array, coords=coords, dims=("x", "y", "z")) + + # combine all data into dictionary + if coord_key[0] == "E": + # off-diagonal components are sampled at respective locations (eg. `eps_xy` at `Ex`) + coords = grid[coord_key[0:2]] + else: + coords = grid[coord_key] + return make_eps_data(coords) + + # TODO: copied from simulation + def _snap_zero_dim(self, grid: Grid): + """Snap a grid to the simulation center along any dimension along which simulation is + effectively 0D, defined as having a single pixel. This is more general than just checking + size = 0.""" + size_snapped = [ + size if num_cells > 1 else 0 for num_cells, size in zip(self.grid.num_cells, self.size) + ] + return grid.snap_to_box_zero_dim(Box(center=self.center, size=size_snapped)) + + # TODO: copied from simulation + @cached_property + def grid(self) -> Grid: + """FDTD grid spatial locations and information. + + Returns + ------- + :class:`.Grid` + :class:`.Grid` storing the spatial locations relevant to the simulation. + """ + + # Add a simulation Box as the first structure + structures = [Structure(geometry=self.geometry, medium=self.medium)] + structures += self.structures + + grid = self.grid_spec.make_grid( + structures=structures, + symmetry=self.symmetry, + periodic=self._periodic, + sources=self.sources, + num_pml_layers=self.num_pml_layers, + ) + + # This would AutoGrid the in-plane directions of the 2D materials + # return self._grid_corrections_2dmaterials(grid) + return grid + def _get_solver_grid( self, keep_additional_layers: bool = False, truncate_symmetry: bool = True ) -> Grid: @@ -160,7 +592,7 @@ def _get_solver_grid( monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME, colocate=False) - span_inds = self.simulation._discretize_inds_monitor(monitor) + span_inds = self._discretize_inds_monitor(monitor) # Remove extension along monitor normal if not keep_additional_layers: @@ -168,7 +600,7 @@ def _get_solver_grid( span_inds[self.normal_axis, 1] -= 1 # Do not extend if simulation has a single pixel along a dimension - for dim, num_cells in enumerate(self.simulation.grid.num_cells): + for dim, num_cells in enumerate(self.grid.num_cells): if num_cells <= 1: span_inds[dim] = [0, 1] @@ -179,7 +611,7 @@ def _get_solver_grid( if sym != 0: span_inds[plane_inds[dim], 0] += np.diff(span_inds[plane_inds[dim]]) // 2 - return self.simulation._subgrid(span_inds=span_inds) + return self._subgrid(span_inds=span_inds) @cached_property def _solver_grid(self) -> Grid: @@ -252,7 +684,7 @@ def _get_data_with_group_index(self) -> ModeSolverData: def grid_snapped(self) -> Grid: """The solver grid snapped to the plane normal and to simulation 0-sized dims if any.""" grid_snapped = self._solver_grid.snap_to_box_zero_dim(self.plane) - return self.simulation._snap_zero_dim(grid_snapped) + return self._snap_zero_dim(grid_snapped) @cached_property def data_raw(self) -> ModeSolverData: @@ -327,7 +759,6 @@ def _data_on_yee_grid(self) -> ModeSolverData: # finite grid corrections grid_factors = self._grid_correction( - simulation=self.simulation, plane=self.plane, mode_spec=self.mode_spec, n_complex=index_data, @@ -336,11 +767,11 @@ def _data_on_yee_grid(self) -> ModeSolverData: # make mode solver data on the Yee grid mode_solver_monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME, colocate=False) - grid_expanded = self.simulation.discretize_monitor(mode_solver_monitor) + grid_expanded = self.discretize_monitor(mode_solver_monitor) mode_solver_data = ModeSolverData( monitor=mode_solver_monitor, - symmetry=self.simulation.symmetry, - symmetry_center=self.simulation.center, + symmetry=self.symmetry, + symmetry_center=self.center, grid_expanded=grid_expanded, grid_primal_correction=grid_factors[0], grid_dual_correction=grid_factors[1], @@ -371,7 +802,7 @@ def _colocate_data(self, mode_solver_data: ModeSolverData) -> ModeSolverData: # Update data mode_solver_monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME) - grid_expanded = self.simulation.discretize_monitor(mode_solver_monitor) + grid_expanded = self.discretize_monitor(mode_solver_monitor) data_dict_colocated.update({"monitor": mode_solver_monitor, "grid_expanded": grid_expanded}) mode_solver_data = mode_solver_data._updated(update=data_dict_colocated) @@ -433,15 +864,15 @@ def sim_data(self) -> SimulationData: :class:`.SimulationData` object containing the effective index and mode fields. """ monitor_data = self.data - new_monitors = list(self.simulation.monitors) + [monitor_data.monitor] - new_simulation = self.simulation.copy(update=dict(monitors=new_monitors)) + new_monitors = list(self.monitors) + [monitor_data.monitor] + new_simulation = self.copy(update=dict(monitors=new_monitors)) return SimulationData(simulation=new_simulation, data=(monitor_data,)) def _get_epsilon(self, freq: float) -> ArrayComplex4D: """Compute the epsilon tensor in the plane. Order of components is xx, xy, xz, yx, etc.""" eps_keys = ["Ex", "Exy", "Exz", "Eyx", "Ey", "Eyz", "Ezx", "Ezy", "Ez"] eps_tensor = [ - self.simulation.epsilon_on_grid(self._solver_grid, key, freq) for key in eps_keys + self.epsilon_on_grid(self._solver_grid, key, freq) for key in eps_keys ] return np.stack(eps_tensor, axis=0) @@ -591,9 +1022,8 @@ def _field_decay_warning(self, field_data: ModeSolverData): "not decay at the plane boundaries." ) - @staticmethod def _grid_correction( - simulation: Simulation, + self, plane: Box, mode_spec: ModeSpec, n_complex: ModeIndexDataArray, @@ -623,7 +1053,7 @@ def _grid_correction( # Primal and dual grid along the normal direction, # i.e. locations of the tangential E-field and H-field components, respectively - grid = simulation.grid + grid = self.grid normal_primal = grid.boundaries.to_list[normal_axis] normal_primal = xr.DataArray(normal_primal, coords={normal_dim: normal_primal}) normal_dual = grid.centers.to_list[normal_axis] @@ -663,15 +1093,15 @@ def _is_tensorial(self) -> bool: @cached_property def _intersecting_media(self) -> List: """List of media (including simulation background) intersecting the mode plane.""" - total_structures = [self.simulation.scene.background_structure] - total_structures += list(self.simulation.structures) - return self.simulation.scene.intersecting_media(self.plane, total_structures) + total_structures = [self.scene.background_structure] + total_structures += list(self.structures) + return self.scene.intersecting_media(self.plane, total_structures) @cached_property def _has_fully_anisotropic_media(self) -> bool: """Check if there are any fully anisotropic media in the plane of the mode.""" if np.any( - [isinstance(mat, FullyAnisotropicMedium) for mat in self.simulation.scene.mediums] + [isinstance(mat, FullyAnisotropicMedium) for mat in self.scene.mediums] ): for int_mat in self._intersecting_media: if isinstance(int_mat, FullyAnisotropicMedium): @@ -796,6 +1226,7 @@ def to_mode_solver_monitor(self, name: str, colocate: bool = None) -> ModeSolver name=name, ) + # TODO: fix this method def sim_with_source( self, source_time: SourceTime, @@ -826,8 +1257,9 @@ def sim_with_source( mode_source = self.to_source( mode_index=mode_index, direction=direction, source_time=source_time ) - new_sources = list(self.simulation.sources) + [mode_source] - new_sim = self.simulation.updated_copy(sources=new_sources) + new_sources = list(self.sources) + [mode_source] + # new_sim = self.updated_copy(sources=new_sources) + new_sim = self.copy() return new_sim def sim_with_monitor( @@ -855,8 +1287,8 @@ def sim_with_monitor( """ mode_monitor = self.to_monitor(freqs=freqs, name=name) - new_monitors = list(self.simulation.monitors) + [mode_monitor] - new_sim = self.simulation.updated_copy(monitors=new_monitors) + new_monitors = list(self.monitors) + [mode_monitor] + new_sim = self.updated_copy(monitors=new_monitors) return new_sim def sim_with_mode_solver_monitor( @@ -879,10 +1311,137 @@ def sim_with_mode_solver_monitor( from the ModeSolver instance and ``name``. """ mode_solver_monitor = self.to_mode_solver_monitor(name=name) - new_monitors = list(self.simulation.monitors) + [mode_solver_monitor] - new_sim = self.simulation.updated_copy(monitors=new_monitors) + new_monitors = list(self.monitors) + [mode_solver_monitor] + new_sim = self.updated_copy(monitors=new_monitors) return new_sim + # TODO: this is just copied from Simulation + def discretize_monitor(self, monitor: Monitor) -> Grid: + """Grid on which monitor data corresponding to a given monitor will be computed.""" + span_inds = self._discretize_inds_monitor(monitor) + grid_snapped = self._subgrid(span_inds=span_inds).snap_to_box_zero_dim(monitor) + grid_snapped = self._snap_zero_dim(grid=grid_snapped) + return grid_snapped + + # TODO: this is just copied from Simulation. Where should it go? + + @equal_aspect + @add_ax_if_none + def plot_boundaries( + self, + x: float = None, + y: float = None, + z: float = None, + ax: Ax = None, + **kwargs, + ) -> Ax: + """Plot the simulation boundary conditions as lines on a plane + defined by one nonzero x,y,z coordinate. + + Parameters + ---------- + x : float = None + position of plane in x direction, only one of x, y, z must be specified to define plane. + y : float = None + position of plane in y direction, only one of x, y, z must be specified to define plane. + z : float = None + position of plane in z direction, only one of x, y, z must be specified to define plane. + ax : matplotlib.axes._subplots.Axes = None + Matplotlib axes to plot on, if not specified, one is created. + **kwargs + Optional keyword arguments passed to the matplotlib ``LineCollection``. + For details on accepted values, refer to + `Matplotlib's documentation `_. + + Returns + ------- + matplotlib.axes._subplots.Axes + The supplied or created matplotlib axes. + """ + + def set_plot_params(boundary_edge, lim, side, thickness): + """Return the line plot properties such as color and opacity based on the boundary""" + if isinstance(boundary_edge, PECBoundary): + plot_params = plot_params_pec.copy(deep=True) + elif isinstance(boundary_edge, PMCBoundary): + plot_params = plot_params_pmc.copy(deep=True) + elif isinstance(boundary_edge, BlochBoundary): + plot_params = plot_params_bloch.copy(deep=True) + else: + plot_params = PlotParams(alpha=0) + + # expand axis limit so that the axis ticks and labels aren't covered + new_lim = lim + if plot_params.alpha != 0: + if side == -1: + new_lim = lim - thickness + elif side == 1: + new_lim = lim + thickness + + return plot_params, new_lim + + boundaries = self.boundary_spec.to_list + + normal_axis, _ = self.parse_xyz_kwargs(x=x, y=y, z=z) + _, (dim_u, dim_v) = self.pop_axis([0, 1, 2], axis=normal_axis) + + umin, umax = ax.get_xlim() + vmin, vmax = ax.get_ylim() + + size_factor = 1.0 / 35.0 + thickness_u = (umax - umin) * size_factor + thickness_v = (vmax - vmin) * size_factor + + # boundary along the u axis, minus side + plot_params, ulim_minus = set_plot_params(boundaries[dim_u][0], umin, -1, thickness_u) + rect = mpl.patches.Rectangle( + xy=(umin - thickness_u, vmin), + width=thickness_u, + height=(vmax - vmin), + **plot_params.to_kwargs(), + **kwargs, + ) + ax.add_patch(rect) + + # boundary along the u axis, plus side + plot_params, ulim_plus = set_plot_params(boundaries[dim_u][1], umax, 1, thickness_u) + rect = mpl.patches.Rectangle( + xy=(umax, vmin), + width=thickness_u, + height=(vmax - vmin), + **plot_params.to_kwargs(), + **kwargs, + ) + ax.add_patch(rect) + + # boundary along the v axis, minus side + plot_params, vlim_minus = set_plot_params(boundaries[dim_v][0], vmin, -1, thickness_v) + rect = mpl.patches.Rectangle( + xy=(umin, vmin - thickness_v), + width=(umax - umin), + height=thickness_v, + **plot_params.to_kwargs(), + **kwargs, + ) + ax.add_patch(rect) + + # boundary along the v axis, plus side + plot_params, vlim_plus = set_plot_params(boundaries[dim_v][1], vmax, 1, thickness_v) + rect = mpl.patches.Rectangle( + xy=(umin, vmax), + width=(umax - umin), + height=thickness_v, + **plot_params.to_kwargs(), + **kwargs, + ) + ax.add_patch(rect) + + # ax = self._set_plot_bounds(ax=ax, x=x, y=y, z=z) + ax.set_xlim([ulim_minus, ulim_plus]) + ax.set_ylim([vlim_minus, vlim_plus]) + + return ax + def plot_field( self, field_name: str, @@ -950,7 +1509,7 @@ def plot_field( def _validate_modes_size(self): """Make sure that the total size of the modes fields is not too large.""" monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME) - num_cells = self.simulation._monitor_num_cells(monitor) + num_cells = self._monitor_num_cells(monitor) # size in GB total_size = monitor._storage_size_solver(num_cells=num_cells, tmesh=[]) / 1e9 if total_size > MAX_MODES_DATA_SIZE_GB: @@ -980,7 +1539,7 @@ def reduced_simulation_copy(self): ) # remove PML, Absorers, etc, to avoid unnecessary cells - bspec = self.simulation.boundary_spec + bspec = self.boundary_spec new_bspec_dict = {} for axis in "xyz": @@ -998,7 +1557,7 @@ def reduced_simulation_copy(self): ) # extract sub-simulation removing everything irrelevant - new_sim = self.simulation.subsection( + new_sim = self.subsection( region=new_sim_box, monitors=[], sources=[], @@ -1008,4 +1567,4 @@ def reduced_simulation_copy(self): remove_outside_structures=True, ) - return self.updated_copy(simulation=new_sim) + return new_sim#self.updated_copy(simulation=new_sim) diff --git a/tidy3d/plugins/mode/web.py b/tidy3d/plugins/mode/web.py index cc132832c..f74b4a5fa 100644 --- a/tidy3d/plugins/mode/web.py +++ b/tidy3d/plugins/mode/web.py @@ -1,4 +1,6 @@ """Web API for mode solver""" -from ...web.api.mode import run + +# TODO: reset back to regular mode.run +from ...web.api.mode_core import run as run __all__ = ["run"] diff --git a/tidy3d/web/api/mode_core.py b/tidy3d/web/api/mode_core.py new file mode 100644 index 000000000..ea615fcc8 --- /dev/null +++ b/tidy3d/web/api/mode_core.py @@ -0,0 +1,539 @@ +"""Web API for mode solver""" + +from __future__ import annotations +from typing import Optional, Callable + +from datetime import datetime +import os +import pathlib +import tempfile +import time + +import pydantic.v1 as pydantic +from botocore.exceptions import ClientError + +from ..core.environment import Env +from ...components.simulation import Simulation +from ...components.data.monitor_data import ModeSolverData +from ...components.medium import AbstractCustomMedium +from ...components.types import Literal +from ...exceptions import WebError +from ...log import log, get_logging_console +from ..core.core_config import get_logger_console +from ..core.http_util import http +from ..core.s3utils import download_file, download_gz_file, upload_file +from ..core.task_core import Folder +from ..core.types import ResourceLifecycle, Submittable + +from ...plugins.mode.mode_solver import ModeSolver, MODE_MONITOR_NAME +from ...version import __version__ + +SIMULATION_JSON = "simulation.json" +SIM_FILE_HDF5_GZ = "simulation.hdf5.gz" +MODESOLVER_API = "tidy3d/modesolver/py" +MODESOLVER_JSON = "mode_solver.json" +MODESOLVER_HDF5 = "mode_solver.hdf5" +MODESOLVER_GZ = "mode_solver.hdf5.gz" + +MODESOLVER_LOG = "output/result.log" +MODESOLVER_RESULT = "output/result.hdf5" +MODESOLVER_RESULT_GZ = "output/mode_solver_data.hdf5.gz" + + +def run( + mode_solver: ModeSolver, + task_name: str = "Untitled", + mode_solver_name: str = "mode_solver", + folder_name: str = "Mode Solver", + results_file: str = "mode_solver.hdf5", + verbose: bool = True, + progress_callback_upload: Callable[[float], None] = None, + progress_callback_download: Callable[[float], None] = None, + reduce_simulation: Literal["auto", True, False] = "auto", +) -> ModeSolverData: + """Submits a :class:`.ModeSolver` to server, starts running, monitors progress, downloads, + and loads results as a :class:`.ModeSolverData` object. + + Parameters + ---------- + mode_solver : :class:`.ModeSolver` + Mode solver to upload to server. + task_name : str = "Untitled" + Name of task. + mode_solver_name: str = "mode_solver" + The name of the mode solver to create the in task. + folder_name : str = "Mode Solver" + Name of folder to store task on web UI. + results_file : str = "mode_solver.hdf5" + Path to download results file (.hdf5). + verbose : bool = True + If ``True``, will print status, otherwise, will run silently. + progress_callback_upload : Callable[[float], None] = None + Optional callback function called when uploading file with ``bytes_in_chunk`` as argument. + progress_callback_download : Callable[[float], None] = None + Optional callback function called when downloading file with ``bytes_in_chunk`` as argument. + reduce_simulation : Literal["auto", True, False] = "auto" + Restrict simulation to mode solver region. If "auto", then simulation is automatically + restricted if it contains custom mediums. + Returns + ------- + :class:`.ModeSolverData` + Mode solver data with the calculated results. + """ + + log_level = "DEBUG" if verbose else "INFO" + if verbose: + console = get_logging_console() + + if reduce_simulation == "auto": + sim_mediums = mode_solver.scene.mediums + contains_custom = any(isinstance(med, AbstractCustomMedium) for med in sim_mediums) + reduce_simulation = contains_custom + + if reduce_simulation: + log.warning( + "The associated 'Simulation' object contains custom mediums. It will be " + "automatically restricted to the mode solver plane to reduce data for uploading. " + "To force uploading the original 'Simulation' object use 'reduce_simulation=False'." + " Setting 'reduce_simulation=True' will force simulation reduction in all cases and" + " silence this warning." + ) + + if reduce_simulation: + mode_solver = mode_solver.reduced_simulation_copy + + task = ModeSolverTask.create(mode_solver, task_name, mode_solver_name, folder_name) + if verbose: + console.log( + f"Mode solver created with task_id='{task.task_id}', solver_id='{task.solver_id}'." + ) + task.upload(verbose=verbose, progress_callback=progress_callback_upload) + task.submit() + + # Wait for task to finish + prev_status = "draft" + status = task.status + while status not in ("success", "error", "diverged", "deleted"): + if status != prev_status: + log.log(log_level, f"Mode solver status: {status}") + if verbose: + console.log(f"Mode solver status: {status}") + prev_status = status + time.sleep(0.5) + status = task.get_info().status + + if status == "error": + raise WebError("Error running mode solver.") + + log.log(log_level, f"Mode solver status: {status}") + if verbose: + console.log(f"Mode solver status: {status}") + + if status != "success": + # Our cache discards None, so the user is able to re-run + return None + + return task.get_result( + to_file=results_file, verbose=verbose, progress_callback=progress_callback_download + ) + + +class ModeSolverTask(ResourceLifecycle, Submittable, extra=pydantic.Extra.allow): + """Interface for managing the running of a :class:`.ModeSolver` task on server.""" + + task_id: str = pydantic.Field( + None, + title="task_id", + description="Task ID number, set when the task is created, leave as None.", + alias="refId", + ) + + solver_id: str = pydantic.Field( + None, + title="solver", + description="Solver ID number, set when the task is created, leave as None.", + alias="id", + ) + + real_flex_unit: float = pydantic.Field( + None, title="real FlexCredits", description="Billed FlexCredits.", alias="charge" + ) + + created_at: Optional[datetime] = pydantic.Field( + title="created_at", description="Time at which this task was created.", alias="createdAt" + ) + + status: str = pydantic.Field( + None, + title="status", + description="Mode solver task status.", + ) + + file_type: str = pydantic.Field( + None, + title="file_type", + description="File type used to upload the mode solver.", + alias="fileType", + ) + + mode_solver: ModeSolver = pydantic.Field( + None, + title="mode_solver", + description="Mode solver being run by this task.", + ) + + @classmethod + def create( + cls, + mode_solver: ModeSolver, + task_name: str = "Untitled", + mode_solver_name: str = "mode_solver", + folder_name: str = "Mode Solver", + ) -> ModeSolverTask: + """Create a new mode solver task on the server. + + Parameters + ---------- + mode_solver: :class".ModeSolver" + The object that will be uploaded to server in the submitting phase. + task_name: str = "Untitled" + The name of the task. + mode_solver_name: str = "mode_solver" + The name of the mode solver to create the in task. + folder_name: str = "Mode Solver" + The name of the folder to store the task. + + Returns + ------- + :class:`ModeSolverTask` + :class:`ModeSolverTask` object containing information about the created task. + """ + folder = Folder.get(folder_name, create=True) + + mode_solver.validate_pre_upload() + + # TODO: add this back + # mode_solver.simulation.validate_pre_upload(source_required=False) + + response_body = { + "projectId": folder.folder_id, + "taskName": task_name, + "protocolVersion": __version__, + "modeSolverName": mode_solver_name, + "fileType": "Gz", + "source": "Python", + } + + resp = http.post(MODESOLVER_API, response_body) + + # TODO: actually fix the root cause later. + # sometimes POST for mode returns None and crashes the log.info call. + # For now this catches it and raises a better error that we can debug. + if resp is None: + raise WebError( + "'ModeSolver' POST request returned 'None'. If you received this error, please " + "raise an issue on the Tidy3D front end GitHub repository referencing the following" + f"response body '{response_body}'." + ) + + log.info( + "Mode solver created with task_id='%s', solver_id='%s'.", resp["refId"], resp["id"] + ) + return ModeSolverTask(**resp, mode_solver=mode_solver) + + @classmethod + def get( + cls, + task_id: str, + solver_id: str, + to_file: str = "mode_solver.hdf5", + sim_file: str = "simulation.hdf5", + verbose: bool = True, + progress_callback: Callable[[float], None] = None, + ) -> ModeSolverTask: + """Get mode solver task from the server by id. + + Parameters + ---------- + task_id: str + Unique identifier of the task on server. + solver_id: str + Unique identifier of the mode solver in the task. + to_file: str = "mode_solver.hdf5" + File to store the mode solver downloaded from the task. + sim_file: str = "simulation.hdf5" + File to store the simulation downloaded from the task. + verbose: bool = True + Whether to display progress bars. + progress_callback : Callable[[float], None] = None + Optional callback function called while downloading the data. + + Returns + ------- + :class:`ModeSolverTask` + :class:`ModeSolverTask` object containing information about the task. + """ + resp = http.get(f"{MODESOLVER_API}/{task_id}/{solver_id}") + task = ModeSolverTask(**resp) + mode_solver = task.get_modesolver(to_file, sim_file, verbose, progress_callback) + return task.copy(update={"mode_solver": mode_solver}) + + def get_info(self) -> ModeSolverTask: + """Get the current state of this task on the server. + + Returns + ------- + :class:`ModeSolverTask` + :class:`ModeSolverTask` object containing information about the task, without the mode + solver. + """ + resp = http.get(f"{MODESOLVER_API}/{self.task_id}/{self.solver_id}") + return ModeSolverTask(**resp, mode_solver=self.mode_solver) + + def upload( + self, verbose: bool = True, progress_callback: Callable[[float], None] = None + ) -> None: + """Upload this task's 'mode_solver' to the server. + + Parameters + ---------- + verbose: bool = True + Whether to display progress bars. + progress_callback : Callable[[float], None] = None + Optional callback function called while uploading the data. + """ + mode_solver = self.mode_solver.copy() + + # TODO: something here + sim = mode_solver#.simulation + + # Upload simulation.hdf5.gz for GUI display + file, file_name = tempfile.mkstemp(".hdf5.gz") + os.close(file) + try: + sim.to_hdf5_gz(file_name) + upload_file( + self.task_id, + file_name, + SIM_FILE_HDF5_GZ, + verbose=verbose, + progress_callback=progress_callback, + ) + finally: + os.unlink(file_name) + + # Upload a single HDF5 file with the full data + file, file_name = tempfile.mkstemp(".hdf5.gz") + os.close(file) + try: + mode_solver.to_hdf5_gz(file_name) + upload_file( + self.solver_id, + file_name, + MODESOLVER_GZ, + verbose=verbose, + progress_callback=progress_callback, + ) + finally: + os.unlink(file_name) + + def submit(self): + """Start the execution of this task. + + The mode solver must be uploaded to the server with the :meth:`ModeSolverTask.upload` method + before this step. + """ + http.post( + f"{MODESOLVER_API}/{self.task_id}/{self.solver_id}/run", + {"enableCaching": Env.current.enable_caching}, + ) + + def delete(self): + """Delete the mode solver and its corresponding task from the server.""" + # Delete mode solver + http.delete(f"{MODESOLVER_API}/{self.task_id}/{self.solver_id}") + # Delete parent task + http.delete(f"tidy3d/tasks/{self.task_id}") + + def abort(self): + """Abort the mode solver and its corresponding task from the server.""" + return http.put( + "tidy3d/tasks/abort", json={"taskType": "MODE_SOLVER", "taskId": self.solver_id} + ) + + def get_modesolver( + self, + to_file: str = "mode_solver.hdf5", + sim_file: str = "simulation.hdf5", + verbose: bool = True, + progress_callback: Callable[[float], None] = None, + ) -> ModeSolver: + """Get mode solver associated with this task from the server. + + Parameters + ---------- + to_file: str = "mode_solver.hdf5" + File to store the mode solver downloaded from the task. + sim_file: str = "simulation.hdf5" + File to store the simulation downloaded from the task, if any. + verbose: bool = True + Whether to display progress bars. + progress_callback : Callable[[float], None] = None + Optional callback function called while downloading the data. + + Returns + ------- + :class:`ModeSolver` + :class:`ModeSolver` object associated with this task. + """ + if self.file_type == "Gz": + file, file_path = tempfile.mkstemp(".hdf5.gz") + os.close(file) + try: + download_file( + self.solver_id, + MODESOLVER_GZ, + to_file=file_path, + verbose=verbose, + progress_callback=progress_callback, + ) + mode_solver = ModeSolver.from_hdf5_gz(file_path) + finally: + os.unlink(file_path) + + elif self.file_type == "Hdf5": + file, file_path = tempfile.mkstemp(".hdf5") + os.close(file) + try: + download_file( + self.solver_id, + MODESOLVER_HDF5, + to_file=file_path, + verbose=verbose, + progress_callback=progress_callback, + ) + mode_solver = ModeSolver.from_hdf5(file_path) + finally: + os.unlink(file_path) + + else: + file, file_path = tempfile.mkstemp(".json") + os.close(file) + try: + download_file( + self.solver_id, + MODESOLVER_JSON, + to_file=file_path, + verbose=verbose, + progress_callback=progress_callback, + ) + mode_solver_dict = ModeSolver.dict_from_json(file_path) + finally: + os.unlink(file_path) + + download_file( + self.task_id, + SIMULATION_JSON, + to_file=sim_file, + verbose=verbose, + progress_callback=progress_callback, + ) + mode_solver_dict["simulation"] = Simulation.from_json(sim_file) + mode_solver = ModeSolver.parse_obj(mode_solver_dict) + + # Store requested mode solver file + mode_solver.to_file(to_file) + + return mode_solver + + def get_result( + self, + to_file: str = "mode_solver_data.hdf5", + verbose: bool = True, + progress_callback: Callable[[float], None] = None, + ) -> ModeSolverData: + """Get mode solver results for this task from the server. + + Parameters + ---------- + to_file: str = "mode_solver_data.hdf5" + File to store the mode solver downloaded from the task. + verbose: bool = True + Whether to display progress bars. + progress_callback : Callable[[float], None] = None + Optional callback function called while downloading the data. + + Returns + ------- + :class:`.ModeSolverData` + Mode solver data with the calculated results. + """ + + file = None + try: + file = download_gz_file( + resource_id=self.solver_id, + remote_filename=MODESOLVER_RESULT_GZ, + to_file=to_file, + verbose=verbose, + progress_callback=progress_callback, + ) + except ClientError: + if verbose: + console = get_logger_console() + console.log(f"Unable to download '{MODESOLVER_RESULT_GZ}'.") + + if not file: + try: + file = download_file( + resource_id=self.solver_id, + remote_filename=MODESOLVER_RESULT, + to_file=to_file, + verbose=verbose, + progress_callback=progress_callback, + ) + except Exception as e: + raise WebError( + "Failed to download the simulation data file from the server. " + "Please confirm that the task was successfully run." + ) from e + + data = ModeSolverData.from_hdf5(to_file) + data = data.copy( + update={"monitor": self.mode_solver.to_mode_solver_monitor(name=MODE_MONITOR_NAME)} + ) + + self.mode_solver._cached_properties["data_raw"] = data + + # Perform symmetry expansion + return self.mode_solver.data + + def get_log( + self, + to_file: str = "mode_solver.log", + verbose: bool = True, + progress_callback: Callable[[float], None] = None, + ) -> pathlib.Path: + """Get execution log for this task from the server. + + Parameters + ---------- + to_file: str = "mode_solver.log" + File to store the mode solver downloaded from the task. + verbose: bool = True + Whether to display progress bars. + progress_callback : Callable[[float], None] = None + Optional callback function called while downloading the data. + + Returns + ------- + path: pathlib.Path + Path to saved file. + """ + return download_file( + self.solver_id, + MODESOLVER_LOG, + to_file=to_file, + verbose=verbose, + progress_callback=progress_callback, + )