From e328aeec8e66dfbd16cd7862f27a9ceb5fabc77b Mon Sep 17 00:00:00 2001 From: Qiming Date: Fri, 1 Nov 2024 18:39:38 -0400 Subject: [PATCH 1/4] 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 b4ee25f12..053a4387a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,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 dc76d2d4f..8d93e9e64 100644 --- a/tidy3d/components/data/data_array.py +++ b/tidy3d/components/data/data_array.py @@ -813,6 +813,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 72a5d5287..21501a6c5 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, @@ -187,6 +191,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.""" @@ -326,6 +337,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() @@ -348,6 +362,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.""" From a0896098e5116c6035dfdd7a7b3f911d324b1fa2 Mon Sep 17 00:00:00 2001 From: Qiming Date: Tue, 5 Nov 2024 17:28:13 -0500 Subject: [PATCH 2/4] normalize --- tidy3d/plugins/mode/mode_solver.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tidy3d/plugins/mode/mode_solver.py b/tidy3d/plugins/mode/mode_solver.py index 21501a6c5..7e14908b8 100644 --- a/tidy3d/plugins/mode/mode_solver.py +++ b/tidy3d/plugins/mode/mode_solver.py @@ -389,6 +389,8 @@ def rotated_mode_solver_data(self) -> ModeSolverData: solver_temp_data=solver_temp_data, ) + self._normalize_modes(mode_solver_data=rotated_mode_data) + return rotated_mode_data @cached_property @@ -643,6 +645,8 @@ def _mode_rotation( for name in rotated_field_data } + rotated_data_arrays["n_complex"] = solver_ref_data_cylindrical["n_complex"] + return solver_temp_data.updated_copy(**rotated_data_arrays) @cached_property From bbc09f0f557d204e094c1073ba12d97dc398e760 Mon Sep 17 00:00:00 2001 From: momchil Date: Thu, 7 Nov 2024 12:38:59 +0100 Subject: [PATCH 3/4] PolySlab rotated and translated methods that return PolySlab rather than Tranformed --- tidy3d/components/geometry/polyslab.py | 55 +++++++++++++++++++++++++- 1 file changed, 54 insertions(+), 1 deletion(-) diff --git a/tidy3d/components/geometry/polyslab.py b/tidy3d/components/geometry/polyslab.py index 7a3c7ba98..78d337838 100644 --- a/tidy3d/components/geometry/polyslab.py +++ b/tidy3d/components/geometry/polyslab.py @@ -4,7 +4,7 @@ from copy import copy from math import isclose -from typing import List, Tuple +from typing import List, Tuple, Union import autograd.numpy as np import pydantic.v1 as pydantic @@ -19,6 +19,7 @@ from ..autograd import AutogradFieldMap, TracedVertices, get_static from ..autograd.derivative_utils import DerivativeInfo from ..base import cached_property, skip_if_fields_missing +from ..transformation import RotationAroundAxis from ..types import ( ArrayFloat2D, ArrayLike, @@ -1518,6 +1519,58 @@ def normalize_vect(arr: np.ndarray) -> np.ndarray: norm = np.where(norm == 0, 1, norm) return arr / norm + def translated(self, x: float, y: float, z: float) -> PolySlab: + """Return a translated copy of this geometry. + + Parameters + ---------- + x : float + Translation along x. + y : float + Translation along y. + z : float + Translation along z. + + Returns + ------- + :class:`PolySlab` + Translated copy of this ``PolySlab``. + """ + + t_normal, t_plane = self.pop_axis((x, y, z), axis=self.axis) + translated_vertices = np.array(self.vertices) + np.array(t_plane)[None, :] + translated_slab_bounds = (self.slab_bounds[0] + t_normal, self.slab_bounds[1] + t_normal) + return self.updated_copy(vertices=translated_vertices, slab_bounds=translated_slab_bounds) + + def rotated(self, angle: float, axis: Union[Axis, Coordinate]) -> PolySlab: + """Return a rotated copy of this geometry. + + Parameters + ---------- + angle : float + Rotation angle (in radians). + axis : Union[int, Tuple[float, float, float]] + Axis of rotation: 0, 1, or 2 for x, y, and z, respectively, or a 3D vector. + + Returns + ------- + :class:`PolySlab` + Rotated copy of this ``PolySlab``. + """ + _, plane_axs = self.pop_axis([0, 1, 2], self.axis) + if (isinstance(axis, int) and axis == self.axis) or ( + isinstance(axis, tuple) and all(axis[ax] == 0 for ax in plane_axs) + ): + verts_3d = np.zeros((3, self.vertices.shape[0])) + verts_3d[plane_axs[0], :] = self.vertices[:, 0] + verts_3d[plane_axs[1], :] = self.vertices[:, 1] + rotation = RotationAroundAxis(angle=angle, axis=axis) + rotated_vertices = rotation.rotate_vector(verts_3d) + rotated_vertices = np.take(rotated_vertices, plane_axs, axis=0).T + return self.updated_copy(vertices=rotated_vertices) + + return super().rotated(angle=angle, axis=axis) + class ComplexPolySlabBase(PolySlab): """Interface for dividing a complex polyslab where self-intersecting polygon can From 2517468bcbf121f8fd42250724e8ceef5fcf670b Mon Sep 17 00:00:00 2001 From: Qiming Date: Thu, 7 Nov 2024 17:19:59 -0500 Subject: [PATCH 4/4] Fixed mode rotation for non-colocated field components --- CHANGELOG.md | 3 +- tidy3d/components/mode.py | 23 +- tidy3d/plugins/mode/mode_solver.py | 360 ++++++++++++++++++----------- 3 files changed, 228 insertions(+), 158 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 053a4387a..c7f13fda3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,7 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Differentiable `smooth_min`, `smooth_max`, and `least_squares` functions in `tidy3d.plugins.autograd`. - Differential operators `grad` and `value_and_grad` in `tidy3d.plugins.autograd` that behave similarly to the autograd operators but support auxiliary data via `aux_data=True` as well as differentiation w.r.t. `DataArray`. - `@scalar_objective` decorator in `tidy3d.plugins.autograd` that wraps objective functions to ensure they return a scalar value and performs additional checks to ensure compatibility of objective functions with autograd. Used by default in `tidy3d.plugins.autograd.value_and_grad` as well as `tidy3d.plugins.autograd.grad`. - +- `bend_angle_rotation` in `mode_spec` to improve accuracy in some cases when both `bend_radius` and `angle_theta` are defined." ### Changed - `CustomMedium` design regions require far less data when performing inverse design by reducing adjoint field monitor size for dims with one pixel. @@ -33,7 +33,6 @@ 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/components/mode.py b/tidy3d/components/mode.py index 87f8529a3..c20d3055f 100644 --- a/tidy3d/components/mode.py +++ b/tidy3d/components/mode.py @@ -131,12 +131,12 @@ class ModeSpec(Tidy3dBaseModel): 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.", + "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 correct geometries in the simulation with rotated structures.", ) track_freq: Union[TrackFreq, None] = pd.Field( @@ -172,17 +172,6 @@ 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 7e14908b8..a6e440ff3 100644 --- a/tidy3d/plugins/mode/mode_solver.py +++ b/tidy3d/plugins/mode/mode_solver.py @@ -337,7 +337,11 @@ 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: + if ( + self.mode_spec.bend_angle_rotation + and self.mode_spec.bend_radius is not None + and np.abs(self.mode_spec.angle_theta) > 0 + ): return self.rotated_mode_solver_data # Compute data on the Yee grid @@ -374,22 +378,64 @@ def rotated_mode_solver_data(self) -> ModeSolverData: solver_ref = self.rotated_structures_copy solver_ref_data = solver_ref.data_raw - # Transform the mode solution from Cartesian to cylindrical + # The reference data should always be colocated to convert to cylindrical coordinates + n_complex = solver_ref_data.n_complex + if not solver_ref.colocate: + solver_ref_data = solver_ref._colocate_data(mode_solver_data=solver_ref_data) + + # Transform the colocated mode solution from cartesian to cylindrical coordinates 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 + # processs of mode rotation + try: + solver = self.reduced_simulation_copy + except Exception as e: + solver = self + log.warning( + "Mode solver reduced_simulation_copy failed. " + "Falling back to non-reduced simulation, which may be slower. " + f"Exception: {str(e)}" + ) - # compute mode solution by rotating the reference data to the monitor plane - rotated_mode_data = self._mode_rotation( + # Compute mode solution by rotating the reference data to the monitor plane + rotated_mode_fields = self._mode_rotation( solver_ref_data_cylindrical=solver_ref_data_cylindrical, - solver_temp_data=solver_temp_data, + solver=solver, ) - self._normalize_modes(mode_solver_data=rotated_mode_data) + # finite grid corrections + grid_factors = solver._grid_correction( + simulation=solver.simulation, + plane=solver.plane, + mode_spec=solver.mode_spec, + n_complex=n_complex, + direction=solver.direction, + ) + + # make mode solver data on the Yee grid + mode_solver_monitor = solver.to_mode_solver_monitor(name=MODE_MONITOR_NAME, colocate=False) + grid_expanded = solver.simulation.discretize_monitor(mode_solver_monitor) + rotated_mode_data = ModeSolverData( + monitor=mode_solver_monitor, + symmetry=solver.simulation.symmetry, + symmetry_center=solver.simulation.center, + grid_expanded=grid_expanded, + grid_primal_correction=grid_factors[0], + grid_dual_correction=grid_factors[1], + **rotated_mode_fields, + ) + + # self._normalize_modes(mode_solver_data=rotated_mode_data) + + # filter polarization if requested + if self.mode_spec.filter_pol is not None: + self._filter_polarization(mode_solver_data=rotated_mode_data) + + # sort modes if requested + if self.mode_spec.track_freq and len(self.freqs) > 1: + rotated_mode_data = rotated_mode_data.overlap_sort(self.mode_spec.track_freq) + + self._field_decay_warning(rotated_mode_data.symmetry_expanded) return rotated_mode_data @@ -399,7 +445,7 @@ def rotated_structures_copy(self): to the simulation and updates the ModeSpec to disable bend correction and reset angles to normal.""" - rotated_structures = self._rotate_structure() + rotated_structures = self._rotate_structures() 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 @@ -407,9 +453,9 @@ def rotated_structures_copy(self): 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." + def _rotate_structures(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) @@ -464,25 +510,27 @@ def _car_2_cyn( """Convert Cartesian field data to cylindrical coordinates centered the rotated bend center.""" - # Extract coordinate values from one of the six field components + # Extract coordinate values to evaluate smallest cell size from one of the six field components + # in the bend plane and coords along the axial axis 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 + lateral_axis = 3 - self.normal_axis - self.bend_axis_3d + 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] + pts[lateral_axis] -= self.rotated_bend_center[lateral_axis] + rho = np.sort(np.abs(pts[lateral_axis])) - rho = np.sort(np.sqrt(pts[idx_u] ** 2 + pts[idx_v] ** 2)) - theta = np.unique(np.arctan2(pts[idx_v], pts[idx_u])) + theta = np.atleast_1d( + np.arctan2( + self.plane.center[idx_v] - self.rotated_bend_center[idx_v], + self.plane.center[idx_u] - self.rotated_bend_center[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( @@ -493,22 +541,30 @@ def _car_2_cyn( 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) + axes = ("x", "y", "z") + axial_name, plane_names = self.plane.pop_axis(axes, 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], - } + plane_name_normal = axes[self.normal_axis] + plane_name_lateral = axes[lateral_axis] + + # Determine which coordinate transformation to use + if cmp_1 == plane_name_lateral: + + def compute_sel_coord_value(_rho): + return _rho * cos_theta + self.rotated_bend_center[idx_u] + elif cmp_2 == plane_name_lateral: + + def compute_sel_coord_value(_rho): + return _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 + lateral_coord_value = compute_sel_coord_value(_rho) fields_at_point = { field_name: getattr(mode_solver_data, field_name) - .sel(sel_coords, method="nearest") - .isel({cmp_1: 0, cmp_2: 0}) + .sel({plane_name_normal: self.plane.center[self.normal_axis]}, method="nearest") + .interp({plane_name_lateral: lateral_coord_value}, method="linear") + .isel({plane_name_lateral: 0}) .values for field_name in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz") } @@ -547,123 +603,149 @@ def _mode_rotation( solver_ref_data_cylindrical: Dict[ Union[ScalarModeFieldCylindricalDataArray, ModeIndexDataArray] ], - solver_temp_data: ModeSolverData, + solver: ModeSolver, ) -> 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 - } + """Rotate the mode solver solution from the reference plane in cylindrical coordinates + to the practical plane.""" + rotated_data_arrays = {} + for field_name in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz"): + if self.colocate is True: + # Get colocation coordinates in the solver plane + normal_dim, 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_coords[normal_dim] = np.atleast_1d(self.plane.center[self.normal_axis]) + x = colocate_coords["x"] + y = colocate_coords["y"] + z = colocate_coords["z"] + xyz_coords = [x.copy(), y.copy(), z.copy()] + else: + # Extract coordinate values from one of the six field components + xyz_coords = solver.grid_snapped[field_name].to_list + x, y, z = (coord.copy() for coord in xyz_coords) + + f = list(self.freqs) + mode_index = np.arange(self.mode_spec.num_modes) + + axial = solver_ref_data_cylindrical["Er"].axial + + # Initialize output arrays + shape = (x.size, y.size, z.size, len(f), mode_index.size) + rotated_field_cmp = np.zeros(shape, dtype=complex) + + idx_w, idx_uv = self.plane.pop_axis((0, 1, 2), axis=self.bend_axis_3d) + idx_u, idx_v = idx_uv + + pts = [coord.copy() for coord in xyz_coords] + + 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]) + axial = pts[idx_w] + + 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): + for axial_idx, axial_value in enumerate(axial): + 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, f=f_idx) + .interp(rho=rho, axial=axial_value, 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 + if field_name == f"E{cmp_1}": + rotated_field_cmp[idx_x, idx_y, idx_z, f_idx, mode_idx] = phase * ( + fields["Er"] * cos_theta - fields["Etheta"] * sin_theta + ) + elif field_name == f"E{cmp_2}": + rotated_field_cmp[idx_x, idx_y, idx_z, f_idx, mode_idx] = phase * ( + fields["Er"] * sin_theta + fields["Etheta"] * cos_theta + ) + elif field_name == f"E{cmp_normal}": + rotated_field_cmp[idx_x, idx_y, idx_z, f_idx, mode_idx] = ( + phase * fields["Eaxial"] + ) + elif field_name == f"H{cmp_1}": + rotated_field_cmp[idx_x, idx_y, idx_z, f_idx, mode_idx] = phase * ( + fields["Hr"] * cos_theta - fields["Htheta"] * sin_theta + ) + elif field_name == f"H{cmp_2}": + rotated_field_cmp[idx_x, idx_y, idx_z, f_idx, mode_idx] = phase * ( + fields["Hr"] * sin_theta + fields["Htheta"] * cos_theta + ) + elif field_name == f"H{cmp_normal}": + rotated_field_cmp[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[field_name] = ScalarModeFieldDataArray( + rotated_field_cmp, coords=coords + ) rotated_data_arrays["n_complex"] = solver_ref_data_cylindrical["n_complex"] - return solver_temp_data.updated_copy(**rotated_data_arrays) + return 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 + """Computes the azimutal angle of the reference modal plane.""" + _, local_coords = self.plane.pop_axis((0, 1, 2), axis=self.bend_axis_3d) + local_coord_x, local_coord_y = local_coords 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], + self.plane.center[local_coord_y] - self.bend_center[local_coord_y], + self.plane.center[local_coord_x] - self.bend_center[local_coord_x], ) return theta_ref @cached_property def bend_center(self) -> List: - """Computes the bend center based on plane center, angle_theta and angle_phi""" + """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 @@ -1245,7 +1327,7 @@ def _grid_correction( # 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 + k_vec = cos_theta * 2 * np.pi * n_complex * n_complex.f / C_0 if direction == "-": k_vec *= -1 phase_primal = np.exp(1j * k_vec * (normal_primal - normal_pos))