Skip to content

Commit

Permalink
vectorization
Browse files Browse the repository at this point in the history
  • Loading branch information
@QimingFlex
QimingFlex committed Nov 15, 2024
1 parent 8236774 commit b3921cd
Showing 1 changed file with 82 additions and 92 deletions.
174 changes: 82 additions & 92 deletions tidy3d/plugins/mode/mode_solver.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -403,6 +403,9 @@ def rotated_mode_solver_data(self) -> ModeSolverData:
solver=solver,
)

# To do: `eps_spec` is now enforced to be "tensorial_complex" for negative direction mode monitor
# usage.
eps_spec = ["tensorial_complex"]
# finite grid corrections
grid_factors = solver._grid_correction(
simulation=solver.simulation,
Expand All @@ -422,7 +425,7 @@ def rotated_mode_solver_data(self) -> ModeSolverData:
grid_expanded=grid_expanded,
grid_primal_correction=grid_factors[0],
grid_dual_correction=grid_factors[1],
eps_spec=solver_ref_data.eps_spec,
eps_spec=eps_spec,
**rotated_mode_fields,
)

Expand Down Expand Up @@ -550,38 +553,31 @@ def _car_2_cyn(

# 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]
lateral_coord_value = 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]
lateral_coord_value = rho * sin_theta + self.rotated_bend_center[idx_v]

# Transform fields from Cartesian to cylindrical coordinates
for i, _rho in enumerate(rho):
lateral_coord_value = compute_sel_coord_value(_rho)
fields_at_point = {
field_name: getattr(mode_solver_data, field_name)
.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")
}

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
]
fields_at_point = {
field_name: getattr(mode_solver_data, field_name)
.sel({plane_name_normal: self.plane.center[self.normal_axis]}, method="nearest")
.interp({plane_name_lateral: lateral_coord_value}, method="linear")
.values
for field_name in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz")
}

for field_type in ["E", "H"]:
field_data_cylindrical[field_type + "r"][:, 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"][:, 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"][:, 0] = fields_at_point[
field_type + axial_name
]

coords = {
"rho": rho,
Expand Down Expand Up @@ -632,7 +628,7 @@ def _mode_rotation(
xyz_coords = solver.grid_snapped[field_name].to_list
x, y, z = (coord.copy() for coord in xyz_coords)

f = list(self.freqs)
f = self.freqs
mode_index = np.arange(self.mode_spec.num_modes)

axial = solver_ref_data_cylindrical["Er"].axial
Expand Down Expand Up @@ -661,67 +657,61 @@ def _mode_rotation(
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"]
)
k0 = 2 * np.pi * f / C_0
n_eff = solver_ref_data_cylindrical["n_complex"].values
beta = k0[:, None] * n_eff * np.abs(self.mode_spec.bend_radius)

for axial_idx, axial_value in enumerate(axial):
# Interpolate field components
fields = {
f"{field}{comp}": solver_ref_data_cylindrical[f"{field}{comp}"]
.isel(theta=0)
.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[:, None, None] * beta)
else:
phase = np.exp(sign * -1j * theta_rel[:, None, None] * 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, :, :] = phase * (
fields["Er"] * cos_theta[:, None, None]
- fields["Etheta"] * sin_theta[:, None, None]
)
elif field_name == f"E{cmp_2}":
rotated_field_cmp[idx_x, idx_y, idx_z, :, :] = phase * (
fields["Er"] * sin_theta[:, None, None]
+ fields["Etheta"] * cos_theta[:, None, None]
)
elif field_name == f"E{cmp_normal}":
rotated_field_cmp[idx_x, idx_y, idx_z, :, :] = phase * fields["Eaxial"]
elif field_name == f"H{cmp_1}":
rotated_field_cmp[idx_x, idx_y, idx_z, :, :] = phase * (
fields["Hr"] * cos_theta[:, None, None]
- fields["Htheta"] * sin_theta[:, None, None]
)
elif field_name == f"H{cmp_2}":
rotated_field_cmp[idx_x, idx_y, idx_z, :, :] = phase * (
fields["Hr"] * sin_theta[:, None, None]
+ fields["Htheta"] * cos_theta[:, None, None]
)
elif field_name == f"H{cmp_normal}":
rotated_field_cmp[idx_x, idx_y, idx_z, :, :] = 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(
Expand Down

0 comments on commit b3921cd

Please sign in to comment.