From 4e6dbc750f3a1ddde0ed14a00bc0830a4a5de106 Mon Sep 17 00:00:00 2001 From: Qiming Date: Fri, 1 Nov 2024 18:39:38 -0400 Subject: [PATCH] Added bend_angle_rotation in mode_spec --- CHANGELOG.md | 1 + tidy3d/__init__.py | 2 + tidy3d/components/data/data_array.py | 18 ++ tidy3d/components/data/dataset.py | 2 + tidy3d/components/mode.py | 23 ++ tidy3d/plugins/mode/mode_solver.py | 347 +++++++++++++++++++++++++++ 6 files changed, 393 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 02390610b..1fd2df141 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Users can manually specify the background medium for a structure to be used for geometry gradient calculations by supplying `Structure.background_permittivity`. This is useful when there are overlapping structures or structures embedded in other mediums. - Autograd functions can now be called directly on `DataArray` (e.g., `np.sum(data_array)`) in objective functions. - Automatic differentiation support for local field projections with `FieldProjectionAngleMonitor` and `FieldProjectionCartesianMonitor` using `FieldProjector.project_fields(far_field_monitor)`. +- `bend_angle_rotation` in `mode_spec` to improve accuracy when both bend and angle are defined." ### Changed - Improved autograd tracer handling in `DataArray`, resulting in significant speedups for differentiation involving large monitors. diff --git a/tidy3d/__init__.py b/tidy3d/__init__.py index 6642b2c95..2496b42cd 100644 --- a/tidy3d/__init__.py +++ b/tidy3d/__init__.py @@ -57,6 +57,7 @@ PointDataArray, ScalarFieldDataArray, ScalarFieldTimeDataArray, + ScalarModeFieldCylindricalDataArray, ScalarModeFieldDataArray, SpatialDataArray, ) @@ -371,6 +372,7 @@ def set_logging_level(level: str) -> None: "FieldProjector", "ScalarFieldDataArray", "ScalarModeFieldDataArray", + "ScalarModeFieldCylindricalDataArray", "ScalarFieldTimeDataArray", "SpatialDataArray", "ModeAmpsDataArray", diff --git a/tidy3d/components/data/data_array.py b/tidy3d/components/data/data_array.py index f37074583..79cd55202 100644 --- a/tidy3d/components/data/data_array.py +++ b/tidy3d/components/data/data_array.py @@ -822,6 +822,24 @@ class ScalarModeFieldDataArray(AbstractSpatialDataArray): _dims = ("x", "y", "z", "f", "mode_index") +class ScalarModeFieldCylindricalDataArray(AbstractSpatialDataArray): + """Spatial distribution of a mode in frequency-domain as a function of mode index. + + Example + ------- + >>> rho = [1,2] + >>> theta = [2,3,4] + >>> axial = [3,4,5,6] + >>> f = [2e14, 3e14] + >>> mode_index = np.arange(5) + >>> coords = dict(rho=rho, theta=theta, axial=axial, f=f, mode_index=mode_index) + >>> fd = ScalarModeFieldCylindricalDataArray((1+1j) * np.random.random((2,3,4,2,5)), coords=coords) + """ + + __slots__ = () + _dims = ("rho", "theta", "axial", "f", "mode_index") + + class FluxDataArray(DataArray): """Flux through a surface in the frequency-domain. diff --git a/tidy3d/components/data/dataset.py b/tidy3d/components/data/dataset.py index dcae12c7a..72c88977f 100644 --- a/tidy3d/components/data/dataset.py +++ b/tidy3d/components/data/dataset.py @@ -32,6 +32,7 @@ PointDataArray, ScalarFieldDataArray, ScalarFieldTimeDataArray, + ScalarModeFieldCylindricalDataArray, ScalarModeFieldDataArray, SpatialDataArray, TimeDataArray, @@ -149,6 +150,7 @@ def colocate(self, x=None, y=None, z=None) -> xr.Dataset: ScalarFieldDataArray, ScalarFieldTimeDataArray, ScalarModeFieldDataArray, + ScalarModeFieldCylindricalDataArray, EMEScalarModeFieldDataArray, EMEScalarFieldDataArray, ] diff --git a/tidy3d/components/mode.py b/tidy3d/components/mode.py index 85d0382f4..87f8529a3 100644 --- a/tidy3d/components/mode.py +++ b/tidy3d/components/mode.py @@ -127,6 +127,18 @@ class ModeSpec(Tidy3dBaseModel): "yz plane, the ``bend_axis`` is always 1 (the global z axis).", ) + bend_angle_rotation: bool = pd.Field( + False, + title="Use fields rotation when both bend and angle are defined", + description="Defines how modes are computed when both a bend and an angle are defined. " + " If `False`, the two coordinate transformations are directly composed. If `True`, the " + "structures in the simulation are first rotated, to compute a mode solution at a reference" + "plane normal to the bend's azimuthal direction. Then, the fields are rotated to align with" + "the mode plane, using the `n_eff` calculated at the reference plane. The second option can" + "produce more accurate results, but more care must be taken for example in ensuring that the" + "original mode plane intersects the right geometries in the simulation with rotated structures.", + ) + track_freq: Union[TrackFreq, None] = pd.Field( "central", title="Mode Tracking Frequency", @@ -160,6 +172,17 @@ def bend_radius_not_zero(cls, val, values): raise SetupError("The magnitude of 'bend_radius' must be larger than 0.") return val + @pd.validator("bend_angle_rotation", always=True) + @skip_if_fields_missing(["bend_radius", "bend_axis"]) + def validate_bend_correction_requirements(cls, val, values): + """Ensure that both ``bend_axis`` and ``bend_radius`` are provided if ``bend_correction`` is enabled.""" + if val is True: + if values.get("bend_axis") is None or values.get("bend_radius") is None: + raise SetupError( + "'bend_correction' can only be enabled when both 'bend_axis' and 'bend_radius' are provided." + ) + return val + @pd.validator("angle_theta", allow_reuse=True, always=True) def glancing_incidence(cls, val): """Warn if close to glancing incidence.""" diff --git a/tidy3d/plugins/mode/mode_solver.py b/tidy3d/plugins/mode/mode_solver.py index e1a9b06e5..3bc5f5ac1 100644 --- a/tidy3d/plugins/mode/mode_solver.py +++ b/tidy3d/plugins/mode/mode_solver.py @@ -19,6 +19,7 @@ from ...components.data.data_array import ( FreqModeDataArray, ModeIndexDataArray, + ScalarModeFieldCylindricalDataArray, ScalarModeFieldDataArray, ) from ...components.data.monitor_data import ModeSolverData @@ -30,8 +31,10 @@ from ...components.medium import FullyAnisotropicMedium from ...components.mode import ModeSpec from ...components.monitor import ModeMonitor, ModeSolverMonitor +from ...components.scene import Scene from ...components.simulation import Simulation from ...components.source import ModeSource, SourceTime +from ...components.structure import Structure from ...components.types import ( TYPE_TAG_STR, ArrayComplex3D, @@ -39,6 +42,7 @@ ArrayFloat1D, Ax, Axis, + Axis2D, Direction, EpsSpecType, FreqArray, @@ -178,6 +182,13 @@ def normal_axis(self) -> Axis: """Axis normal to the mode plane.""" return self.plane.size.index(0.0) + @cached_property + def normal_axis_2d(self) -> Axis2D: + """Axis normal to the mode plane in a 2D plane that is normal to the bend_axis_3d.""" + _, idx_plane = self.plane.pop_axis((0, 1, 2), axis=self.bend_axis_3d) + + return idx_plane.index(self.normal_axis) + @cached_property def solver_symmetry(self) -> Tuple[Symmetry, Symmetry]: """Get symmetry for solver for propagation along self.normal axis.""" @@ -317,6 +328,9 @@ def data_raw(self) -> ModeSolverData: if self.mode_spec.group_index_step > 0: return self._get_data_with_group_index() + if self.mode_spec.bend_angle_rotation: + return self.rotated_mode_solver_data + # Compute data on the Yee grid mode_solver_data = self._data_on_yee_grid() @@ -339,6 +353,339 @@ def data_raw(self) -> ModeSolverData: return mode_solver_data + @cached_property + def bend_axis_3d(self) -> Axis: + """Transform the 2D bend axis to its corresponding 3D axis.""" + _, idx_plane = self.plane.pop_axis((0, 1, 2), axis=self.normal_axis) + return idx_plane[self.mode_spec.bend_axis] + + @cached_property + def rotated_mode_solver_data(self) -> ModeSolverData: + # Create a mode solver with rotated geometries for a 0-degree angle solution + solver_ref = self.rotated_structures_copy + solver_ref_data = solver_ref.data_raw + + # Transform the mode solution from Cartesian to cylindrical + solver_ref_data_cylindrical = self._car_2_cyn(mode_solver_data=solver_ref_data) + + # This solver temp is only used to get the desired modal plane coords for interpolation + # and better be avoid one unneccessary mode solver simulation. + mode_spec = self.mode_spec.updated_copy(bend_angle_rotation=False) + solver_temp = self.updated_copy(mode_spec=mode_spec) + solver_temp_data = solver_temp.data_raw + + # compute mode solution by rotating the reference data to the monitor plane + rotated_mode_data = self._mode_rotation( + solver_ref_data_cylindrical=solver_ref_data_cylindrical, + solver_temp_data=solver_temp_data, + ) + + return rotated_mode_data + + @cached_property + def rotated_structures_copy(self): + """Creates a copy of the original ModeSolver with rotated structures + to the simulation and updates the ModeSpec to disable bend correction + and reset angles to normal.""" + + rotated_structures = self._rotate_structure() + rotated_simulation = self.simulation.updated_copy(structures=rotated_structures) + rotated_mode_spec = self.mode_spec.updated_copy( + bend_angle_rotation=False, angle_theta=0, angle_phi=0 + ) + + return self.updated_copy(simulation=rotated_simulation, mode_spec=rotated_mode_spec) + + def _rotate_structure(self) -> List[Structure]: + "rotate the structures intersecting with modal plane by angle theta if bend_correction" + "is enabeled for bend simulations." + + _, (idx_u, idx_v) = self.plane.pop_axis((0, 1, 2), axis=self.bend_axis_3d) + + mnt_center = self.plane.center + angle_theta = self.mode_spec.angle_theta + angle_phi = self.mode_spec.angle_phi + + theta_map = { + (0, 2): -angle_theta * np.cos(angle_phi), + (0, 1): angle_theta * np.sin(angle_phi), + (1, 2): angle_theta * np.cos(angle_phi), + (1, 0): -angle_theta * np.sin(angle_phi), + (2, 1): -angle_theta * np.cos(angle_phi), + (2, 0): angle_theta * np.sin(angle_phi), + } + theta = theta_map.get((self.normal_axis, self.bend_axis_3d), 0) + + # Get the translation values + translate_coords = [0, 0, 0] + translate_coords[idx_u] = mnt_center[idx_u] + translate_coords[idx_v] = mnt_center[idx_v] + + rotated_structures = [] + for structure in Scene.intersecting_structures(self.plane, self.simulation.structures): + # Rotate and apply translations + geometry = structure.geometry + geometry = ( + geometry.translated( + x=-translate_coords[0], y=-translate_coords[1], z=-translate_coords[2] + ) + .rotated(theta, axis=self.bend_axis_3d) + .translated(x=translate_coords[0], y=translate_coords[1], z=translate_coords[2]) + ) + + rotated_structures.append(structure.updated_copy(geometry=geometry)) + + return rotated_structures + + @cached_property + def rotated_bend_center(self) -> List: + """Calculate the center of the rotated bend such that the modal plane is normal + to the azimuthal direction of the bend.""" + rotated_bend_center = list(self.plane.center) + idx_rotate = 3 - self.normal_axis - self.bend_axis_3d + rotated_bend_center[idx_rotate] -= self.mode_spec.bend_radius + + return rotated_bend_center + + def _car_2_cyn( + self, mode_solver_data: ModeSolverData + ) -> Dict[Union[ScalarModeFieldCylindricalDataArray, ModeIndexDataArray]]: + """Convert Cartesian field data to cylindrical coordinates centered the + rotated bend center.""" + + # Extract coordinate values from one of the six field components + pts = [mode_solver_data.Ex[name].values.copy() for name in ["x", "y", "z"]] + f, mode_index = mode_solver_data.Ex.f, mode_solver_data.Ex.mode_index + + idx_w, idx_uv = self.plane.pop_axis((0, 1, 2), axis=self.bend_axis_3d) + idx_u, idx_v = idx_uv + + # offset coordinates by the rotated bend center in the bend plane + pts[idx_u] -= self.rotated_bend_center[idx_u] + pts[idx_v] -= self.rotated_bend_center[idx_v] + + rho = np.sort(np.sqrt(pts[idx_u] ** 2 + pts[idx_v] ** 2)) + theta = np.unique(np.arctan2(pts[idx_v], pts[idx_u])) + axial = pts[idx_w] + + # Sanity check for theta uniqueness + if len(theta) > 1: + raise ValueError("Theta should have only one unique value.") + + # Initialize output data arrays for cylindrical fields + field_data_cylindrical = { + field_name_cylindrical: np.zeros( + (len(rho), len(theta), len(axial), len(f), len(mode_index)), dtype=complex + ) + for field_name_cylindrical in ("Er", "Etheta", "Eaxial", "Hr", "Htheta", "Haxial") + } + + cos_theta, sin_theta = np.cos(theta), np.sin(theta) + + axial_name, plane_names = self.plane.pop_axis(("x", "y", "z"), axis=self.bend_axis_3d) + cmp_1, cmp_2 = plane_names + sel_coords_func = { + cmp_1: lambda _rho: _rho * cos_theta + self.rotated_bend_center[idx_u], + cmp_2: lambda _rho: _rho * sin_theta + self.rotated_bend_center[idx_v], + } + + # Transform fields from Cartesian to cylindrical coordinates + for i, _rho in enumerate(rho): + sel_coords = {key: func(_rho) for key, func in sel_coords_func.items()} + + # Find nearest values for the original data and populate cylindrical fields + fields_at_point = { + field_name: getattr(mode_solver_data, field_name) + .sel(sel_coords, method="nearest") + .isel({cmp_1: 0, cmp_2: 0}) + .values + for field_name in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz") + } + + for field_type in ["E", "H"]: + field_data_cylindrical[field_type + "r"][i, 0] = ( + fields_at_point[field_type + cmp_1] * cos_theta + + fields_at_point[field_type + cmp_2] * sin_theta + ) + field_data_cylindrical[field_type + "theta"][i, 0] = ( + -fields_at_point[field_type + cmp_1] * sin_theta + + fields_at_point[field_type + cmp_2] * cos_theta + ) + field_data_cylindrical[field_type + "axial"][i, 0] = fields_at_point[ + field_type + axial_name + ] + + coords = { + "rho": rho, + "theta": theta, + "axial": axial, + "f": f, + "mode_index": mode_index, + } + + solver_data_cylindrical = { + name: ScalarModeFieldCylindricalDataArray(field_data_cylindrical[name], coords=coords) + for name in ("Er", "Etheta", "Eaxial", "Hr", "Htheta", "Haxial") + } + solver_data_cylindrical["n_complex"] = mode_solver_data.n_complex + + return solver_data_cylindrical + + def _mode_rotation( + self, + solver_ref_data_cylindrical: Dict[ + Union[ScalarModeFieldCylindricalDataArray, ModeIndexDataArray] + ], + solver_temp_data: ModeSolverData, + ) -> ModeSolverData: + """rotate the mode solver solution from the reference plane in cylindrical coordinates + to the practical plane""" + + # Extract coordinate values from one of the six field components + f, mode_index = solver_temp_data.Ex.f, solver_temp_data.Ex.mode_index + x, y, z = solver_temp_data.Ex.x, solver_temp_data.Ex.y, solver_temp_data.Ex.z + axial = solver_ref_data_cylindrical["Er"].axial + + # Initialize output arrays + shape = (x.size, y.size, z.size, f.size, mode_index.size) + rotated_field_data = { + name: np.zeros(shape, dtype=complex) for name in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz") + } + + # offset coors by bend_center + _, idx_uv = self.plane.pop_axis((0, 1, 2), axis=self.bend_axis_3d) + idx_u, idx_v = idx_uv + + pts = [solver_temp_data.Ex[name].values.copy() for name in ("x", "y", "z")] + pts[idx_u] -= self.bend_center[idx_u] + pts[idx_v] -= self.bend_center[idx_v] + + rho = np.sqrt(pts[idx_u] ** 2 + pts[idx_v] ** 2) + theta = np.arctan2(pts[idx_v], pts[idx_u]) + theta_rel = theta - self.theta_reference + + cos_theta = pts[idx_u] / rho + sin_theta = pts[idx_v] / rho + + cmp_normal, source_names = self.plane.pop_axis(("x", "y", "z"), axis=self.bend_axis_3d) + cmp_1, cmp_2 = source_names + + for mode_idx in range(mode_index.size): + for f_idx, f_val in enumerate(f.values): + for axial_idx in range(axial.size): + k0 = 2 * np.pi * f_val / C_0 + n_eff = ( + solver_ref_data_cylindrical["n_complex"] + .isel(mode_index=mode_idx, f=f_idx) + .values + ) + beta = k0 * n_eff * np.abs(self.mode_spec.bend_radius) + + # Interpolate field components + fields = { + f"{field}{comp}": solver_ref_data_cylindrical[f"{field}{comp}"] + .isel(mode_index=mode_idx, theta=0, axial=axial_idx, f=f_idx) + .interp(rho=rho, method="linear") + .values + for field in ["E", "H"] + for comp in ["r", "theta", "axial"] + } + + # Determine the phase factor based on normal_axis_2d + sign = -1 if self.normal_axis_2d == 0 else 1 + if (self.direction == "+" and self.mode_spec.bend_radius >= 0) or ( + self.direction == "-" and self.mode_spec.bend_radius < 0 + ): + phase = np.exp(sign * 1j * theta_rel * beta) + else: + phase = np.exp(sign * -1j * theta_rel * beta) + + # Set fixed index based on normal_axis + idx = [slice(None)] * 3 + idx[self.normal_axis] = 0 + idx[self.bend_axis_3d] = axial_idx + idx_x, idx_y, idx_z = idx + + # Assign rotated fields + rotated_field_data[f"E{cmp_1}"][idx_x, idx_y, idx_z, f_idx, mode_idx] = ( + phase * (fields["Er"] * cos_theta - fields["Etheta"] * sin_theta) + ) + rotated_field_data[f"E{cmp_2}"][idx_x, idx_y, idx_z, f_idx, mode_idx] = ( + phase * (fields["Er"] * sin_theta + fields["Etheta"] * cos_theta) + ) + rotated_field_data[f"E{cmp_normal}"][idx_x, idx_y, idx_z, f_idx, mode_idx] = ( + phase * fields["Eaxial"] + ) + + rotated_field_data[f"H{cmp_1}"][idx_x, idx_y, idx_z, f_idx, mode_idx] = ( + phase * (fields["Hr"] * cos_theta - fields["Htheta"] * sin_theta) + ) + rotated_field_data[f"H{cmp_2}"][idx_x, idx_y, idx_z, f_idx, mode_idx] = ( + phase * (fields["Hr"] * sin_theta + fields["Htheta"] * cos_theta) + ) + rotated_field_data[f"H{cmp_normal}"][idx_x, idx_y, idx_z, f_idx, mode_idx] = ( + phase * fields["Haxial"] + ) + + # Convert numpy arrays to xarray DataArrays + coords = {"x": x, "y": y, "z": z, "f": f, "mode_index": mode_index} + rotated_data_arrays = { + name: ScalarModeFieldDataArray(rotated_field_data[name], coords=coords) + for name in rotated_field_data + } + + return solver_temp_data.updated_copy(**rotated_data_arrays) + + @cached_property + def theta_reference(self) -> float: + """Computes the azimutal angle of the reference modal plane""" + _, local_coors = self.plane.pop_axis((0, 1, 2), axis=self.bend_axis_3d) + local_coor_x, local_coor_y = local_coors + theta_ref = np.arctan2( + self.plane.center[local_coor_y] - self.bend_center[local_coor_y], + self.plane.center[local_coor_x] - self.bend_center[local_coor_x], + ) + + return theta_ref + + @cached_property + def bend_center(self) -> List: + """Computes the bend center based on plane center, angle_theta and angle_phi""" + _, id_bend_uv = self.plane.pop_axis((0, 1, 2), axis=self.bend_axis_3d) + + id_bend_u, id_bend_v = id_bend_uv + bend_radius = self.mode_spec.bend_radius + theta = self.mode_spec.angle_theta + phi = self.mode_spec.angle_phi + + bend_center = list(self.plane.center) + + angle_map = { + (0, 2): lambda: theta * np.cos(phi), + (0, 1): lambda: theta * np.sin(phi), + (1, 2): lambda: theta * np.cos(phi), + (1, 0): lambda: theta * np.sin(phi), + (2, 0): lambda: theta * np.sin(phi), + (2, 1): lambda: theta * np.cos(phi), + } + + update_map = { + (0, 2): lambda angle: (bend_radius * np.sin(angle), -bend_radius * np.cos(angle)), + (0, 1): lambda angle: (bend_radius * np.sin(angle), -bend_radius * np.cos(angle)), + (1, 2): lambda angle: (-bend_radius * np.cos(angle), bend_radius * np.sin(angle)), + (1, 0): lambda angle: (bend_radius * np.sin(angle), -bend_radius * np.cos(angle)), + (2, 0): lambda angle: (-bend_radius * np.cos(angle), bend_radius * np.sin(angle)), + (2, 1): lambda angle: (-bend_radius * np.cos(angle), bend_radius * np.sin(angle)), + } + + if (self.normal_axis, self.bend_axis_3d) in angle_map: + angle = angle_map[(self.normal_axis, self.bend_axis_3d)]() + delta_u, delta_v = update_map[(self.normal_axis, self.bend_axis_3d)](angle) + bend_center[id_bend_u] = self.plane.center[id_bend_u] + delta_u + bend_center[id_bend_v] = self.plane.center[id_bend_v] + delta_v + + return bend_center + def _data_on_yee_grid(self) -> ModeSolverData: """Solve for all modes, and construct data with fields on the Yee grid."""