From b1c38bc1039acd5df9eece0abe98e4876b7f061e Mon Sep 17 00:00:00 2001 From: Qiming Date: Thu, 7 Nov 2024 17:19:59 -0500 Subject: [PATCH] Fixed mode rotation for non-colocated field components --- tidy3d/components/mode.py | 23 +- tidy3d/plugins/mode/mode_solver.py | 360 ++++++++++++++++++----------- 2 files changed, 227 insertions(+), 156 deletions(-) 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..189c6575f 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, y, z] + else: + # Extract coordinate values from one of the six field components + xyz_coords = solver.grid_snapped[field_name].to_list + x, y, z = 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 = 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 + ) + if 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 + ) + if field_name == f"E{cmp_normal}": + rotated_field_cmp[idx_x, idx_y, idx_z, f_idx, mode_idx] = ( + phase * fields["Eaxial"] + ) + if 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 + ) + if 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 + ) + if 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))