diff --git a/tests/test_components/test_autograd_rf_box.py b/tests/test_components/test_autograd_rf_box.py new file mode 100644 index 0000000000..9f88a136ff --- /dev/null +++ b/tests/test_components/test_autograd_rf_box.py @@ -0,0 +1,841 @@ +# test autograd and compares to numerically computed finite difference gradients +from __future__ import annotations + +import operator +import sys + +import autograd as ag +import matplotlib.pylab as plt +import numpy as np +import pytest + +import tidy3d as td +import tidy3d.web as web + +td.config.logging_level = "ERROR" + +PLOT_FD_ADJ_COMPARISON = False +NUM_BOXES_PER_FD_TESTS = 1 +RUN_WITH_FD_CONVERGENCE = True +FD_CONVERGENCE_THRESHOLD = 0.075 +MIN_FD_COMPARISON = 2 +FIXED_MESH_REFINEMENT_FACTOR = 200.0 +SAVE_FD_ADJ_DATA = True +SAVE_FD_LOC = 0 +SAVE_ADJ_LOC = 1 +LOCAL_GRADIENT = True +VERBOSE = False +SHOW_PRINT_STATEMENTS = True +USE_POLYSLAB_FOR_BOX = False +NUMERICAL_RESULTS_DATA_DIR = ( + "./numerical_rf_box_polyslab_test/" if USE_POLYSLAB_FOR_BOX else "./numerical_rf_box_box_test/" +) + +assert NUM_BOXES_PER_FD_TESTS == 1, "Currently only supporting a single box in the numerical test!" + +if PLOT_FD_ADJ_COMPARISON: + pytestmark = pytest.mark.usefixtures("mpl_config_interactive") +else: + pytestmark = pytest.mark.usefixtures("mpl_config_noninteractive") + +if SHOW_PRINT_STATEMENTS: + sys.stdout = sys.stderr + + +def get_sim_geometry(mesh_wvl_um): + return td.Box(size=(4 * mesh_wvl_um, 4 * mesh_wvl_um, 7 * mesh_wvl_um), center=(0, 0, 0)) + + +def make_base_sim( + mesh_wvl_um, + adj_wvl_um, + monitor_size_wvl, + box_for_override, + mesh_refinement_factor, + run_time=1e-8, +): + sim_geometry = get_sim_geometry(mesh_wvl_um) + + sim_size_um = sim_geometry.size + sim_center_um = sim_geometry.center + + boundary_spec = td.BoundarySpec( + x=td.Boundary.pml(), + y=td.Boundary.pml(), + z=td.Boundary.pml(), + ) + + dl_design = mesh_wvl_um / mesh_refinement_factor + + mesh_overrides = [] + mesh_overrides.extend( + [ + td.MeshOverrideStructure( + geometry=box_for_override, + dl=[dl_design, dl_design, dl_design], + ), + ] + ) + + src_size = sim_size_um[0:2] + (0,) + + wl_min_src_um = 0.9 * adj_wvl_um + wl_max_src_um = 1.1 * adj_wvl_um + + fwidth_src = td.C_0 * ((1.0 / wl_min_src_um) - (1.0 / wl_max_src_um)) + freq0 = td.C_0 / adj_wvl_um + + pulse = td.GaussianPulse(freq0=freq0, fwidth=fwidth_src) + src = td.PlaneWave( + center=(0, 0, -2 * mesh_wvl_um), + size=src_size, + source_time=pulse, + direction="+", + pol_angle=np.pi / 4.0, + ) + + field_monitor = td.FieldMonitor( + # shifted center + center=(0, 0, 0.125 * sim_size_um[2]), + # center=(0, 0, -0.75 * mesh_wvl_um), + size=tuple(dim * mesh_wvl_um for dim in monitor_size_wvl), + name="monitor_fields", + freqs=[freq0], + ) + + sim_base = td.Simulation( + center=sim_center_um, + size=sim_size_um, + grid_spec=td.GridSpec.auto( + min_steps_per_wvl=30, + wavelength=mesh_wvl_um, + override_structures=mesh_overrides, + ), + structures=[], + sources=[src], + monitors=[field_monitor], + run_time=run_time, + boundary_spec=boundary_spec, + subpixel=True, + ) + + return sim_base + + +def create_objective_function_2D(create_sim_base, eval_fn, box_z_values, sim_path_dir): + # 2D PEC objective will take in an array of parameters that can be reshaped into + # a Nx4 array where N is the number of boxes and the four parameters are + # center (x,y) and lateral size (x,y) assuming the thickness in z is 0 + # + # box_param_arrays should be a list of such parameters for cases where we + # are running finite difference simulations and want to put together a + # whole batch + def objective(box_param_arrays): + sim_base = create_sim_base() + + simulation_dict = {} + for idx in range(len(box_param_arrays)): + get_boxes_params = box_param_arrays[idx] + num_boxes = len(get_boxes_params) // 2 + reshape_parameters = np.reshape(get_boxes_params, (num_boxes, 2)) + + box_structures = [] + + for box_idx in range(num_boxes): + box_params = reshape_parameters[box_idx] + + if USE_POLYSLAB_FOR_BOX: + min_x = -0.5 * box_params[0] + min_y = -0.5 * box_params[1] + + max_x = 0.5 * box_params[0] + max_y = 0.5 * box_params[1] + + vertices = ((min_x, min_y), (min_x, max_y), (max_x, max_y), (max_x, min_y)) + + polyslab = td.PolySlab( + vertices=vertices, + slab_bounds=(box_z_values[box_idx], box_z_values[box_idx]), + ) + box_structures.append(td.Structure(geometry=polyslab, medium=td.PEC2D)) + else: + box_structures.append( + td.Structure( + geometry=td.Box( + center=(0.0, 0.0, box_z_values[box_idx]), + size=(box_params[0], box_params[1], 0), + ), + medium=td.PEC2D, + ) + ) + + sim_with_block = sim_base.updated_copy( + structures=tuple(list(sim_base.structures) + box_structures) + ) + + simulation_dict[f"numerical_rf_box_2d_testing_{idx}"] = sim_with_block.copy() + + sim_data = web.run_async( + simulation_dict, path_dir=sim_path_dir, local_gradient=LOCAL_GRADIENT, verbose=VERBOSE + ) + + objective_vals = [] + for idx in range(len(box_param_arrays)): + objective_vals.append(eval_fn(sim_data[f"numerical_rf_box_2d_testing_{idx}"])) + + if len(box_param_arrays) == 1: + return objective_vals[0] + + return objective_vals + + return objective + + +def create_objective_function_3D(create_sim_base, eval_fn, box_z_thickneses, sim_path_dir): + # 3D PEC objective will take in an array of parameters that can be reshaped into + # a Nx5 array where N is the number of boxes and the four parameters are + # center (x,y) and lateral size (x,y) and the z position of the box. The box thicknesses + # are set in box_z_thicknesses + # + # box_param_arrays should be a list of such parameters for cases where we + # are running finite difference simulations and want to put together a + # whole batch + def objective(box_param_arrays): + sim_base = create_sim_base() + + simulation_dict = {} + for idx in range(len(box_param_arrays)): + get_boxes_params = box_param_arrays[idx] + num_boxes = len(get_boxes_params) // 3 + reshape_parameters = np.reshape(get_boxes_params, (num_boxes, 3)) + + box_structures = [] + + for box_idx in range(num_boxes): + box_params = reshape_parameters[box_idx] + + if USE_POLYSLAB_FOR_BOX: + min_x = -0.5 * box_params[0] + min_y = -0.5 * box_params[1] + + max_x = 0.5 * box_params[0] + max_y = 0.5 * box_params[1] + + vertices = ((min_x, min_y), (min_x, max_y), (max_x, max_y), (max_x, min_y)) + polyslab = td.PolySlab( + vertices=vertices, + slab_bounds=( + box_params[2] - 0.5 * box_z_thickneses[box_idx], + box_params[2] + 0.5 * box_z_thickneses[box_idx], + ), + ) + box_structures.append(td.Structure(geometry=polyslab, medium=td.PECMedium())) + else: + box_structures.append( + td.Structure( + geometry=td.Box( + center=(0.0, 0.0, box_params[2]), + size=(box_params[0], box_params[1], box_z_thickneses[box_idx]), + ), + medium=td.PECMedium(), + ) + ) + + sim_with_block = sim_base.updated_copy( + structures=tuple(list(sim_base.structures) + box_structures) + ) + + simulation_dict[f"numerical_rf_box_3d_testing_{idx}"] = sim_with_block.copy() + + sim_data = web.run_async( + simulation_dict, path_dir=sim_path_dir, local_gradient=LOCAL_GRADIENT, verbose=VERBOSE + ) + + objective_vals = [] + for idx in range(len(box_param_arrays)): + objective_vals.append(eval_fn(sim_data[f"numerical_rf_box_3d_testing_{idx}"])) + + if len(box_param_arrays) == 1: + return objective_vals[0] + + return objective_vals + + return objective + + +def make_eval_fns(monitor_size_wvl): + num_nonzero_spatial_dims = 3 - np.sum(np.isclose(monitor_size_wvl, 0)) + + def intensity(sim_data): + field_data = sim_data["monitor_fields"] + shape_x, shape_y, shape_z, *_ = field_data.Ex.values.shape + + proj_pol_rotation = field_data.Ex.values - field_data.Ey.values + return np.sum(np.abs(proj_pol_rotation) ** 2) + + eval_fns = [intensity] + eval_fn_names = ["intensity"] + + return eval_fns, eval_fn_names + + +def compare_fd_adj(fd_data, adj_data): + norm_fd = fd_data / np.linalg.norm(fd_data) + norm_adj = adj_data / np.linalg.norm(adj_data) + + dot_fd_adj = np.sum(norm_fd * norm_adj) + + directional_overlap_deg = np.arccos(dot_fd_adj) * 180.0 / np.pi + + def compute_error_statistics(a, b): + relative_error = np.abs(a - b) / np.abs(b) + return np.mean(relative_error), np.std(relative_error) + + error_mean, error_std = compute_error_statistics(fd_data, adj_data) + error_norm_mean, error_norm_std = compute_error_statistics(norm_fd, norm_adj) + + return directional_overlap_deg, error_mean, error_std, error_norm_mean, error_norm_std + + +def generate_boxes(lateral_dim_bounds, mesh_cell_override_size, is_2d, rng): + all_box_parameters = [] + box_z_values = [] + + for _box in range(NUM_BOXES_PER_FD_TESTS): + x_size = rng.uniform(0.6 * mesh_wvl_um, 0.7 * mesh_wvl_um) + y_size = rng.uniform(0.2 * mesh_wvl_um, 0.3 * mesh_wvl_um) + + if is_2d: + all_box_parameters += [x_size, y_size] + box_z_values.append(0.0) + else: + all_box_parameters += [ + x_size, + y_size, + 0.0, + ] + + if is_2d: + return all_box_parameters, box_z_values + else: + return all_box_parameters + + +def run_and_process_fd(all_box_parameters, fd_step, objective): + all_fd_grad_parameters = [] + for fd_step_idx in range(len(fd_step)): + fd_box_parameters = [] + + for param_idx in range(len(all_box_parameters)): + copy_params_up = all_box_parameters.copy() + copy_params_down = all_box_parameters.copy() + + copy_params_up[param_idx] += fd_step[fd_step_idx] + copy_params_down[param_idx] -= fd_step[fd_step_idx] + + fd_box_parameters.append(copy_params_up) + fd_box_parameters.append(copy_params_down) + + all_obj = objective(fd_box_parameters) + + fd_grad_parameters = [] + + for param_idx in range(len(all_box_parameters)): + fd_grad = (all_obj[2 * param_idx] - all_obj[2 * param_idx + 1]) / ( + 2 * fd_step[fd_step_idx] + ) + + fd_grad_parameters.append(fd_grad) + + all_fd_grad_parameters.append(fd_grad_parameters) + + all_fd_grad_parameters = np.array(all_fd_grad_parameters) + + if RUN_WITH_FD_CONVERGENCE: + assert len(fd_step) == 2, "Currently only convergence testing with 2 points" + + argmin_convergence_test = np.argmin(fd_step) + argmax_convergence_test = np.argmax(fd_step) + + relative_diff = ( + all_fd_grad_parameters[argmax_convergence_test] + - all_fd_grad_parameters[argmin_convergence_test] + ) / all_fd_grad_parameters[argmin_convergence_test] + + valid_mask = np.abs(relative_diff) < FD_CONVERGENCE_THRESHOLD + fd_grad = np.squeeze(all_fd_grad_parameters[argmin_convergence_test]) + else: + fd_grad = np.squeeze(all_fd_grad_parameters) + valid_mask = np.ones(fd_grad.shape, dtype=bool) + + return fd_grad, valid_mask + + +mm = 1e3 + +background_indices = [1.0] +mesh_wvls_mm = [15.0, 30.0] +adj_wvls_mm = [15.0, 30.0] + +mesh_refinement_factors = np.linspace(200.0, 300.0, 4) +box_z_thickneses_3d_wvl = np.linspace(0.2, 0.4, 4) + +mesh_wvls_um = [mesh_wvl_mm * mm for mesh_wvl_mm in mesh_wvls_mm] +adj_wvls_um = [adj_wvl_mm * mm for adj_wvl_mm in adj_wvls_mm] + +monitor_size_3d_wvl = (1.0, 1.0, 0) + +rf_2d_test_parameters = [] + +test_number = 0 +for idx in range(len(mesh_wvls_um)): + mesh_wvl_um = mesh_wvls_um[idx] + adj_wvl_um = adj_wvls_um[idx] + + eval_fns, eval_fn_names = make_eval_fns(monitor_size_3d_wvl) + + for mesh_refinement_factor in mesh_refinement_factors: + for eval_fn_idx, eval_fn in enumerate(eval_fns): + rf_2d_test_parameters.append( + { + "mesh_wvl_um": mesh_wvl_um, + "adj_wvl_um": adj_wvl_um, + "monitor_size_wvl": monitor_size_3d_wvl, + "mesh_refinement_factor": mesh_refinement_factor, + "eval_fn": eval_fn, + "eval_fn_name": eval_fn_names[eval_fn_idx], + "test_number": test_number, + } + ) + + test_number += 1 + + +rf_3d_test_parameters = [] + +test_number = 0 +for idx in range(len(mesh_wvls_um)): + mesh_wvl_um = mesh_wvls_um[idx] + adj_wvl_um = adj_wvls_um[idx] + + eval_fns, eval_fn_names = make_eval_fns(monitor_size_3d_wvl) + + for box_z_thickness_wvl in box_z_thickneses_3d_wvl: + for eval_fn_idx, eval_fn in enumerate(eval_fns): + rf_3d_test_parameters.append( + { + "mesh_wvl_um": mesh_wvl_um, + "adj_wvl_um": adj_wvl_um, + "monitor_size_wvl": monitor_size_3d_wvl, + "box_z_thickness_wvl": box_z_thickness_wvl, + "eval_fn": eval_fn, + "eval_fn_name": eval_fn_names[eval_fn_idx], + "test_number": test_number, + } + ) + + test_number += 1 + + +@pytest.mark.numerical +@pytest.mark.parametrize( + "rf_2d_test_parameters, dir_name", + zip( + rf_2d_test_parameters, + ([NUMERICAL_RESULTS_DATA_DIR] if SAVE_FD_ADJ_DATA else [None]) * len(rf_2d_test_parameters), + ), + indirect=["dir_name"], +) +def test_finite_difference_2d_box_pec(rf_2d_test_parameters, rng, tmp_path, create_directory): + """Test a variety of autograd permittivity gradients for 2D PEC boxes by""" + """comparing them to numerical finite difference.""" + + test_number = rf_2d_test_parameters["test_number"] + + ( + mesh_wvl_um, + adj_wvl_um, + monitor_size_wvl, + mesh_refinement_factor, + eval_fn, + eval_fn_name, + test_number, + ) = operator.itemgetter( + "mesh_wvl_um", + "adj_wvl_um", + "monitor_size_wvl", + "mesh_refinement_factor", + "eval_fn", + "eval_fn_name", + "test_number", + )(rf_2d_test_parameters) + + dim_um = 3 * mesh_wvl_um + + thickness_box_placement_um = 0.6 * mesh_wvl_um # 1.2 * mesh_wvl_um + + sim_geometry = get_sim_geometry(mesh_wvl_um) + + box_for_override = td.Box( + center=(0, 0, 0), + size=sim_geometry.size[0:2] + (thickness_box_placement_um + 0.3 * mesh_wvl_um,), + ) + + eval_fns, eval_fn_names = make_eval_fns(monitor_size_wvl) + + sim_path_dir = tmp_path / f"test{test_number}" + sim_path_dir.mkdir() + + mesh_cell_override_size = mesh_wvl_um / mesh_refinement_factor + + if RUN_WITH_FD_CONVERGENCE: + fd_step = np.linspace(2 * mesh_cell_override_size, mesh_cell_override_size, 2) + else: + fd_step = np.array([mesh_cell_override_size]) + + z_planes = np.array([0.0]) + + all_box_parameters, box_z_values = generate_boxes( + lateral_dim_bounds=[-0.5 * dim_um, 0.5 * dim_um], + mesh_cell_override_size=mesh_cell_override_size, + is_2d=True, + rng=rng, + ) + + objective = create_objective_function_2D( + lambda mesh_wvl_um=mesh_wvl_um, + adj_wvl_um=adj_wvl_um, + monitor_size_wvl=monitor_size_wvl, + box_for_override=box_for_override, + mesh_refinement_factor=mesh_refinement_factor: make_base_sim( + mesh_wvl_um=mesh_wvl_um, + adj_wvl_um=adj_wvl_um, + monitor_size_wvl=monitor_size_wvl, + box_for_override=box_for_override, + mesh_refinement_factor=mesh_refinement_factor, + ), + eval_fn, + box_z_values, + sim_path_dir=str(sim_path_dir), + ) + + obj_val_and_grad = ag.value_and_grad(objective) + + obj, adj_grad = obj_val_and_grad([all_box_parameters]) + + fd_grad, valid_mask = run_and_process_fd( + all_box_parameters=all_box_parameters, fd_step=fd_step, objective=objective + ) + + fd_grad = np.array(fd_grad) + adj_grad = np.array(adj_grad) + valid_mask = np.array(valid_mask) + + reshape_fd = np.reshape(fd_grad, (NUM_BOXES_PER_FD_TESTS, 2)) + reshape_adj = np.reshape(adj_grad, (NUM_BOXES_PER_FD_TESTS, 2)) + reshape_valid = np.reshape(valid_mask, (NUM_BOXES_PER_FD_TESTS, 2)) + + all_center_fd = [] + all_center_adj = [] + + all_width_fd = [] + all_width_adj = [] + + for box_idx in range(NUM_BOXES_PER_FD_TESTS): + width_fd = reshape_fd[box_idx] + width_adj = reshape_adj[box_idx] + + width_valid = reshape_valid[box_idx] + + add_widths_fd = [width_fd[idx] for idx in range(len(width_fd)) if width_valid[idx]] + add_widths_adj = [width_adj[idx] for idx in range(len(width_adj)) if width_valid[idx]] + + all_width_fd += add_widths_fd + all_width_adj += add_widths_adj + + all_width_fd = np.array(all_width_fd) + all_width_adj = np.array(all_width_adj) + + assert len(all_width_fd) >= MIN_FD_COMPARISON, "Too many widths were trimmed" + + ( + width_overlap_deg, + width_error_mean, + width_error_std, + width_error_norm_mean, + width_error_norm_std, + ) = compare_fd_adj(all_width_fd, all_width_adj) + + width_data = np.zeros((2, len(all_width_fd))) + + width_data[SAVE_FD_LOC, :] = all_width_fd + width_data[SAVE_ADJ_LOC, :] = all_width_adj + + if SHOW_PRINT_STATEMENTS: + print(f"\n2D PEC Box Test {test_number} Summary:") + print(f"Mesh wavelength (um): {mesh_wvl_um}") + print(f"Adjoint wavelength (um): {adj_wvl_um}") + print(f"Monitor size (wavelengths): {monitor_size_wvl}") + print(f"Mesh refinement factor: {mesh_refinement_factor}") + print(f"Eval function: {eval_fn_name}") + print(f"Width mean (std): {width_error_mean} ({width_error_std})") + print(f"Width norm mean (std): {width_error_norm_mean} ({width_error_norm_std})") + print(f"Width overlap deg: {width_overlap_deg}") + print("\n") + + if SAVE_FD_ADJ_DATA: + np.save( + f"{NUMERICAL_RESULTS_DATA_DIR}/rf_box_2d_width_fd_adj_data_test_{test_number}.npy", + width_data, + ) + + yvals = [width_error_mean, width_error_norm_mean] + yerrs = [width_error_std, width_error_norm_std] + bars = plt.bar( + ["width\nerror", "width\nerror\nnorm"], + yvals, + yerr=yerrs, + color="skyblue", + edgecolor="black", + capsize=5, + ) + + plt.title(f"Vector Alignment (deg):\n{width_overlap_deg:.2f}") + plt.xticks(rotation=45) + + for idx, bar in enumerate(bars): + yval = bar.get_height() + yerr = yerrs[idx] + plt.text( + bar.get_x() + 0.5 * bar.get_width(), + yval + yerr + 0.1, + round(yval, 3), + ha="center", + va="bottom", + ) + + plt.savefig(f"{NUMERICAL_RESULTS_DATA_DIR}/rf_2d_box_summary_plot_test_{test_number}.png") + + +@pytest.mark.numerical +@pytest.mark.parametrize( + "rf_3d_test_parameters, dir_name", + zip( + rf_3d_test_parameters, + ([NUMERICAL_RESULTS_DATA_DIR] if SAVE_FD_ADJ_DATA else [None]) * len(rf_3d_test_parameters), + ), + indirect=["dir_name"], +) +def test_finite_difference_3d_box_pec(rf_3d_test_parameters, rng, tmp_path, create_directory): + """Test a variety of autograd permittivity gradients for 3D PEC boxes by""" + """comparing them to numerical finite difference.""" + + test_number = rf_3d_test_parameters["test_number"] + + ( + mesh_wvl_um, + adj_wvl_um, + monitor_size_wvl, + box_z_thickness_wvl, + eval_fn, + eval_fn_name, + test_number, + ) = operator.itemgetter( + "mesh_wvl_um", + "adj_wvl_um", + "monitor_size_wvl", + "box_z_thickness_wvl", + "eval_fn", + "eval_fn_name", + "test_number", + )(rf_3d_test_parameters) + + dim_um = 3 * mesh_wvl_um + + thickness_box_placement_um = 1.2 * mesh_wvl_um + + sim_geometry = get_sim_geometry(mesh_wvl_um) + + box_for_override = td.Box( + center=(0, 0, 0), + size=sim_geometry.size[0:2] + (thickness_box_placement_um + 0.3 * mesh_wvl_um,), + ) + + eval_fns, eval_fn_names = make_eval_fns(monitor_size_wvl) + + sim_path_dir = tmp_path / f"test{test_number}" + sim_path_dir.mkdir() + + mesh_cell_override_size = mesh_wvl_um / FIXED_MESH_REFINEMENT_FACTOR + + if RUN_WITH_FD_CONVERGENCE: + fd_step = np.linspace(2 * mesh_cell_override_size, mesh_cell_override_size, 2) + else: + fd_step = np.array([mesh_cell_override_size]) + + all_box_parameters = generate_boxes( + lateral_dim_bounds=[-0.5 * dim_um, 0.5 * dim_um], + mesh_cell_override_size=mesh_cell_override_size, + is_2d=False, + rng=rng, + ) + + box_z_thicknesses = box_z_thickness_wvl * adj_wvl_um * np.ones(NUM_BOXES_PER_FD_TESTS) + + objective = create_objective_function_3D( + lambda mesh_wvl_um=mesh_wvl_um, + adj_wvl_um=adj_wvl_um, + monitor_size_wvl=monitor_size_wvl, + box_for_override=box_for_override, + mesh_refinement_factor=FIXED_MESH_REFINEMENT_FACTOR: make_base_sim( + mesh_wvl_um=mesh_wvl_um, + adj_wvl_um=adj_wvl_um, + monitor_size_wvl=monitor_size_wvl, + box_for_override=box_for_override, + mesh_refinement_factor=mesh_refinement_factor, + ), + eval_fn, + box_z_thicknesses, + sim_path_dir=str(sim_path_dir), + ) + + obj_val_and_grad = ag.value_and_grad(objective) + + obj, adj_grad = obj_val_and_grad([all_box_parameters]) + + fd_grad, valid_mask = run_and_process_fd( + all_box_parameters=all_box_parameters, fd_step=fd_step, objective=objective + ) + + fd_grad = np.array(fd_grad) + adj_grad = np.array(adj_grad) + valid_mask = np.array(valid_mask) + + reshape_fd = np.reshape(fd_grad, (NUM_BOXES_PER_FD_TESTS, 3)) + reshape_adj = np.reshape(adj_grad, (NUM_BOXES_PER_FD_TESTS, 3)) + reshape_valid = np.reshape(valid_mask, (NUM_BOXES_PER_FD_TESTS, 3)) + + all_center_fd = [] + all_center_adj = [] + + all_width_fd = [] + all_width_adj = [] + + all_z_coord_fd = [] + all_z_coord_adj = [] + + for box_idx in range(NUM_BOXES_PER_FD_TESTS): + width_fd = reshape_fd[box_idx, 0:2] + width_adj = reshape_adj[box_idx, 0:2] + + width_valid = reshape_valid[box_idx, 0:2] + + add_widths_fd = [width_fd[idx] for idx in range(len(width_fd)) if width_valid[idx]] + add_widths_adj = [width_adj[idx] for idx in range(len(width_adj)) if width_valid[idx]] + + all_width_fd += add_widths_fd + all_width_adj += add_widths_adj + + if reshape_valid[box_idx, 2]: + all_z_coord_fd.append(reshape_fd[box_idx, 2]) + all_z_coord_adj.append(reshape_adj[box_idx, 2]) + + all_width_fd = np.array(all_width_fd) + all_width_adj = np.array(all_width_adj) + + all_z_coord_fd = np.array(all_z_coord_fd) + all_z_coord_adj = np.array(all_z_coord_adj) + + assert len(all_width_fd) >= MIN_FD_COMPARISON, "Too many widths were trimmed" + assert len(all_z_coord_fd) >= 1, "Too many z coordinates were trimmed" + + ( + width_overlap_deg, + width_error_mean, + width_error_std, + width_error_norm_mean, + width_error_norm_std, + ) = compare_fd_adj(all_width_fd, all_width_adj) + ( + z_coord_overlap_deg, + z_coord_error_mean, + z_coord_error_std, + z_coord_error_norm_mean, + z_coord_error_norm_std, + ) = compare_fd_adj(all_z_coord_fd, all_z_coord_adj) + + width_data = np.zeros((2, len(all_width_fd))) + z_coord_data = np.zeros((2, len(all_z_coord_fd))) + + width_data[SAVE_FD_LOC, :] = all_width_fd + width_data[SAVE_ADJ_LOC, :] = all_width_adj + + z_coord_data[SAVE_FD_LOC, :] = all_z_coord_fd + z_coord_data[SAVE_ADJ_LOC, :] = all_z_coord_adj + + if SHOW_PRINT_STATEMENTS: + print(f"\n3D PEC Box Test {test_number} Summary:") + print(f"Mesh wavelength (um): {mesh_wvl_um}") + print(f"Adjoint wavelength (um): {adj_wvl_um}") + print(f"Monitor size (wavelengths): {monitor_size_wvl}") + print(f"Box z thickness (wavelengths): {box_z_thickness_wvl}") + print(f"Eval function: {eval_fn_name}") + print(f"Width mean (std): {width_error_mean} ({width_error_std})") + print(f"Width norm mean (std): {width_error_norm_mean} ({width_error_norm_std})") + print(f"Width overlap deg: {width_overlap_deg}") + print(f"Z mean (std): {z_coord_error_mean} ({z_coord_error_std})") + print(f"Z norm mean (std): {z_coord_error_norm_mean} ({z_coord_error_norm_std})") + print(f"Z overlap deg: {z_coord_overlap_deg}") + print("\n") + + if SAVE_FD_ADJ_DATA: + np.save( + f"{NUMERICAL_RESULTS_DATA_DIR}/rf_box_3d_width_fd_adj_data_test_{test_number}.npy", + width_data, + ) + np.save( + f"{NUMERICAL_RESULTS_DATA_DIR}/rf_box_3d_z_coord_fd_adj_data_test_{test_number}.npy", + z_coord_data, + ) + + yvals = [ + width_error_mean, + width_error_norm_mean, + z_coord_error_mean, + z_coord_error_norm_mean, + ] + yerrs = [ + width_error_std, + width_error_norm_std, + z_coord_error_std, + z_coord_error_norm_std, + ] + bars = plt.bar( + [ + "width\nerror", + "width\nerror\nnorm", + "z\ncoord\nerror", + "z\ncoord\nerror\nnorm", + ], + yvals, + yerr=yerrs, + color="skyblue", + edgecolor="black", + capsize=5, + ) + + plt.title(f"Vector Alignment (deg):\n{width_overlap_deg:.2f}, {z_coord_overlap_deg:.2f}") + plt.ylabel("Relative Error") + + for idx, bar in enumerate(bars): + yval = bar.get_height() + yerr = yerrs[idx] + plt.text( + bar.get_x() + 0.5 * bar.get_width(), + yval + yerr + 0.1, + round(yval, 3), + ha="center", + va="bottom", + ) + + plt.savefig(f"{NUMERICAL_RESULTS_DATA_DIR}/rf_3d_box_summary_plot_test_{test_number}.png") diff --git a/tests/test_components/test_autograd_rf_polyslab.py b/tests/test_components/test_autograd_rf_polyslab.py new file mode 100644 index 0000000000..08d85334e3 --- /dev/null +++ b/tests/test_components/test_autograd_rf_polyslab.py @@ -0,0 +1,706 @@ +# test autograd and compares to numerically computed finite difference gradients +from __future__ import annotations + +import operator +import sys + +import autograd as ag +import matplotlib.pylab as plt +import numpy as np +import pytest + +import tidy3d as td +import tidy3d.web as web + +td.config.logging_level = "ERROR" + +PLOT_FD_ADJ_COMPARISON = False +NUM_VERTICES_PER_FD_TESTS = 10 +RUN_WITH_FD_CONVERGENCE = True +FD_CONVERGENCE_THRESHOLD = 0.05 +MIN_FD_COMPARISON = 2 +FIXED_MESH_REFINEMENT_FACTOR = 200.0 +SAVE_FD_ADJ_DATA = True +SAVE_FD_LOC = 0 +SAVE_ADJ_LOC = 1 +LOCAL_GRADIENT = True +VERBOSE = False +NUMERICAL_RESULTS_DATA_DIR = "./numerical_rf_polyslab_test/" +SHOW_PRINT_STATEMENTS = False + +if PLOT_FD_ADJ_COMPARISON: + pytestmark = pytest.mark.usefixtures("mpl_config_interactive") +else: + pytestmark = pytest.mark.usefixtures("mpl_config_noninteractive") + +if SHOW_PRINT_STATEMENTS: + sys.stdout = sys.stderr + + +def get_sim_geometry(mesh_wvl_um): + return td.Box(size=(5 * mesh_wvl_um, 5 * mesh_wvl_um, 7 * mesh_wvl_um), center=(0, 0, 0)) + + +def make_base_sim( + mesh_wvl_um, + adj_wvl_um, + monitor_size_wvl, + box_for_override, + mesh_refinement_factor, + run_time=1e-8, +): + sim_geometry = get_sim_geometry(mesh_wvl_um) + + sim_size_um = sim_geometry.size + sim_center_um = sim_geometry.center + + boundary_spec = td.BoundarySpec( + x=td.Boundary.pml(), + y=td.Boundary.pml(), + z=td.Boundary.pml(), + ) + + dl_design = mesh_wvl_um / mesh_refinement_factor + + mesh_overrides = [] + mesh_overrides.extend( + [ + td.MeshOverrideStructure( + geometry=box_for_override, + dl=[dl_design, dl_design, dl_design], + ), + ] + ) + + src_size = sim_size_um[0:2] + (0,) + + wl_min_src_um = 0.9 * adj_wvl_um + wl_max_src_um = 1.1 * adj_wvl_um + + fwidth_src = td.C_0 * ((1.0 / wl_min_src_um) - (1.0 / wl_max_src_um)) + freq0 = td.C_0 / adj_wvl_um + + pulse = td.GaussianPulse(freq0=freq0, fwidth=fwidth_src) + src = td.PlaneWave( + center=(0, 0, -2 * mesh_wvl_um), + size=src_size, + source_time=pulse, + direction="+", + pol_angle=np.pi / 4.0, + ) + + field_monitor = td.FieldMonitor( + # shifted center + center=(0, 0, 0.25 * sim_size_um[2]), + size=tuple(dim * mesh_wvl_um for dim in monitor_size_wvl), + name="monitor_fields", + freqs=[freq0], + ) + + sim_base = td.Simulation( + center=sim_center_um, + size=sim_size_um, + grid_spec=td.GridSpec.auto( + min_steps_per_wvl=30, + wavelength=mesh_wvl_um, + override_structures=mesh_overrides, + ), + structures=[], + sources=[src], + monitors=[field_monitor], + run_time=run_time, + boundary_spec=boundary_spec, + subpixel=True, + ) + + return sim_base + + +def create_objective_function_2D(create_sim_base, eval_fn, polyslab_z_value, sim_path_dir): + # 2D PEC objective will take in an array of parameters that can be reshaped into + # an Nx2 array where N is the number of vertices in a polyslab and the parameters are + # center (x,y) and lateral size (x,y) assuming the thickness in z is 0 + # + # box_param_arrays should be a list of such parameters for cases where we + # are running finite difference simulations and want to put together a + # whole batch + def objective(polyslab_param_arrays): + sim_base = create_sim_base() + + simulation_dict = {} + for idx in range(len(polyslab_param_arrays)): + get_polyslab_params = polyslab_param_arrays[idx] + reshape_parameters = np.reshape(get_polyslab_params, (NUM_VERTICES_PER_FD_TESTS, 2)) + + polyslab_structures = [ + td.Structure( + geometry=td.PolySlab( + vertices=tuple(map(tuple, reshape_parameters)), + slab_bounds=(polyslab_z_value, polyslab_z_value), + ), + medium=td.PEC2D, + ) + ] + + sim_with_block = sim_base.updated_copy( + structures=tuple(list(sim_base.structures) + polyslab_structures) + ) + + simulation_dict[f"numerical_rf_polyslab_2d_testing_{idx}"] = sim_with_block.copy() + + sim_data = web.run_async( + simulation_dict, path_dir=sim_path_dir, local_gradient=LOCAL_GRADIENT, verbose=VERBOSE + ) + + objective_vals = [] + for idx in range(len(polyslab_param_arrays)): + objective_vals.append(eval_fn(sim_data[f"numerical_rf_polyslab_2d_testing_{idx}"])) + + if len(polyslab_param_arrays) == 1: + return objective_vals[0] + + return objective_vals + + return objective + + +def create_objective_function_3D( + create_sim_base, eval_fn, polyslab_z_value, polyslab_z_thickness, sim_path_dir +): + # 2D PEC objective will take in an array of parameters that can be reshaped into + # an Nx2 array where N is the number of vertices in a polyslab and the parameters are + # center (x,y) and lateral size (x,y) assuming the thickness in z is 0 + # + # box_param_arrays should be a list of such parameters for cases where we + # are running finite difference simulations and want to put together a + # whole batch + def objective(polyslab_param_arrays): + sim_base = create_sim_base() + + simulation_dict = {} + for idx in range(len(polyslab_param_arrays)): + get_polyslab_params = polyslab_param_arrays[idx] + reshape_parameters = np.reshape(get_polyslab_params, (NUM_VERTICES_PER_FD_TESTS, 2)) + + polyslab_structures = [ + td.Structure( + geometry=td.PolySlab( + vertices=tuple(map(tuple, reshape_parameters)), + slab_bounds=( + polyslab_z_value - 0.5 * polyslab_z_thickness, + polyslab_z_value + 0.5 * polyslab_z_thickness, + ), + ), + medium=td.PECMedium(), + ) + ] + + sim_with_block = sim_base.updated_copy( + structures=tuple(list(sim_base.structures) + polyslab_structures) + ) + + simulation_dict[f"numerical_rf_polyslab_3d_testing_{idx}"] = sim_with_block.copy() + + sim_data = web.run_async( + simulation_dict, path_dir=sim_path_dir, local_gradient=LOCAL_GRADIENT, verbose=VERBOSE + ) + + objective_vals = [] + for idx in range(len(polyslab_param_arrays)): + objective_vals.append(eval_fn(sim_data[f"numerical_rf_polyslab_3d_testing_{idx}"])) + + if len(polyslab_param_arrays) == 1: + return objective_vals[0] + + return objective_vals + + return objective + + +def make_eval_fns(monitor_size_wvl): + num_nonzero_spatial_dims = 3 - np.sum(np.isclose(monitor_size_wvl, 0)) + + def intensity(sim_data): + field_data = sim_data["monitor_fields"] + shape_x, shape_y, shape_z, *_ = field_data.Ex.values.shape + + proj_pol_rotation = field_data.Ex.values - field_data.Ey.values + return np.sum(np.abs(proj_pol_rotation) ** 2) + + eval_fns = [intensity] + eval_fn_names = ["intensity"] + + return eval_fns, eval_fn_names + + +def compare_fd_adj(fd_data, adj_data): + norm_fd = fd_data / np.linalg.norm(fd_data) + norm_adj = adj_data / np.linalg.norm(adj_data) + + dot_fd_adj = np.sum(norm_fd * norm_adj) + + directional_overlap_deg = np.arccos(dot_fd_adj) * 180.0 / np.pi + + def compute_error_statistics(a, b): + relative_error = np.abs(a - b) / np.abs(b) + return np.mean(relative_error), np.std(relative_error) + + error_mean, error_std = compute_error_statistics(fd_data, adj_data) + error_norm_mean, error_norm_std = compute_error_statistics(norm_fd, norm_adj) + + return directional_overlap_deg, error_mean, error_std, error_norm_mean, error_norm_std + + +def generate_polyslab(min_dim_lateral, max_dim_lateral, rng): + x_radius = 0.5 * rng.uniform(min_dim_lateral, max_dim_lateral) + y_radius = 0.5 * rng.uniform(min_dim_lateral, max_dim_lateral) + rotation = rng.uniform(0, 2 * np.pi) + + theta = np.linspace(0, 2 * np.pi, NUM_VERTICES_PER_FD_TESTS, endpoint=False) + + x_values = x_radius * np.cos(theta) + y_values = y_radius * np.sin(theta) + + def rotate(x, y, phi): + x_rotate = x * np.cos(phi) + y * np.sin(phi) + y_rotate = -x * np.sin(phi) + y * np.cos(phi) + + return x_rotate, y_rotate + + x_rotated, y_rotated = rotate(x_values, y_values, rotation) + + return np.array(list(zip(x_rotated, y_rotated))).flatten() + + +def run_and_process_fd(polyslab_parameters, fd_step, objective): + all_fd_grad_parameters = [] + for fd_step_idx in range(len(fd_step)): + fd_polyslab_parameters = [] + + for param_idx in range(len(polyslab_parameters)): + copy_params_up = polyslab_parameters.copy() + copy_params_down = polyslab_parameters.copy() + + copy_params_up[param_idx] += fd_step[fd_step_idx] + copy_params_down[param_idx] -= fd_step[fd_step_idx] + + fd_polyslab_parameters.append(copy_params_up) + fd_polyslab_parameters.append(copy_params_down) + + all_obj = objective(fd_polyslab_parameters) + + fd_grad_parameters = [] + + for param_idx in range(len(polyslab_parameters)): + fd_grad = (all_obj[2 * param_idx] - all_obj[2 * param_idx + 1]) / ( + 2 * fd_step[fd_step_idx] + ) + + fd_grad_parameters.append(fd_grad) + + all_fd_grad_parameters.append(fd_grad_parameters) + + all_fd_grad_parameters = np.array(all_fd_grad_parameters) + + if RUN_WITH_FD_CONVERGENCE: + assert len(fd_step) == 2, "Currently only convergence testing with 2 points" + + argmin_convergence_test = np.argmin(fd_step) + argmax_convergence_test = np.argmax(fd_step) + + relative_diff = ( + all_fd_grad_parameters[argmax_convergence_test] + - all_fd_grad_parameters[argmin_convergence_test] + ) / all_fd_grad_parameters[argmin_convergence_test] + + valid_mask = np.abs(relative_diff) < FD_CONVERGENCE_THRESHOLD + fd_grad = np.squeeze(all_fd_grad_parameters[argmin_convergence_test]) + else: + fd_grad = np.squeeze(all_fd_grad_parameters) + valid_mask = np.ones(fd_grad.shape, dtype=bool) + + return fd_grad, valid_mask + + +mm = 1e3 + +background_indices = [1.0, 1.5] +mesh_wvls_mm = [15.0, 15.0] +adj_wvls_mm = [15.0, 20.0] + +mesh_refinement_factors = np.linspace(200.0, 300.0, 4) + +polyslab_z_thickneses_3d_wvl = np.linspace(0.2, 0.4, 4) + +mesh_wvls_um = [mesh_wvl_mm * mm for mesh_wvl_mm in mesh_wvls_mm] +adj_wvls_um = [adj_wvl_mm * mm for adj_wvl_mm in adj_wvls_mm] + +monitor_size_3d_wvl = (1.0, 1.0, 0) + +rf_2d_test_parameters = [] + +test_number = 0 +for idx in range(len(mesh_wvls_um)): + mesh_wvl_um = mesh_wvls_um[idx] + adj_wvl_um = adj_wvls_um[idx] + + eval_fns, eval_fn_names = make_eval_fns(monitor_size_3d_wvl) + + for mesh_refinement_factor in mesh_refinement_factors: + for eval_fn_idx, eval_fn in enumerate(eval_fns): + rf_2d_test_parameters.append( + { + "mesh_wvl_um": mesh_wvl_um, + "adj_wvl_um": adj_wvl_um, + "monitor_size_wvl": monitor_size_3d_wvl, + "mesh_refinement_factor": mesh_refinement_factor, + "eval_fn": eval_fn, + "eval_fn_name": eval_fn_names[eval_fn_idx], + "test_number": test_number, + } + ) + + test_number += 1 + + +rf_3d_test_parameters = [] + +test_number = 0 +for idx in range(len(mesh_wvls_um)): + mesh_wvl_um = mesh_wvls_um[idx] + adj_wvl_um = adj_wvls_um[idx] + + eval_fns, eval_fn_names = make_eval_fns(monitor_size_3d_wvl) + + for polyslab_z_thickness_wvl in polyslab_z_thickneses_3d_wvl: + for eval_fn_idx, eval_fn in enumerate(eval_fns): + rf_3d_test_parameters.append( + { + "mesh_wvl_um": mesh_wvl_um, + "adj_wvl_um": adj_wvl_um, + "monitor_size_wvl": monitor_size_3d_wvl, + "polyslab_z_thickness_wvl": polyslab_z_thickness_wvl, + "eval_fn": eval_fn, + "eval_fn_name": eval_fn_names[eval_fn_idx], + "test_number": test_number, + } + ) + + test_number += 1 + +import sys + +sys.stdout = sys.stderr + + +@pytest.mark.numerical +@pytest.mark.parametrize( + "rf_2d_test_parameters, dir_name", + zip( + rf_2d_test_parameters, + ([NUMERICAL_RESULTS_DATA_DIR] if SAVE_FD_ADJ_DATA else [None]) * len(rf_2d_test_parameters), + ), + indirect=["dir_name"], +) +def test_finite_difference_2d_polyslab_pec(rf_2d_test_parameters, rng, tmp_path, create_directory): + """Test a variety of autograd permittivity gradients for 2D `PolySlab` PEC by""" + """comparing them to numerical finite difference.""" + + test_number = rf_2d_test_parameters["test_number"] + + ( + mesh_wvl_um, + adj_wvl_um, + monitor_size_wvl, + mesh_refinement_factor, + eval_fn, + eval_fn_name, + test_number, + ) = operator.itemgetter( + "mesh_wvl_um", + "adj_wvl_um", + "monitor_size_wvl", + "mesh_refinement_factor", + "eval_fn", + "eval_fn_name", + "test_number", + )(rf_2d_test_parameters) + + dim_um = 1.5 * mesh_wvl_um + + thickness_box_placement_um = 1.2 * mesh_wvl_um + + sim_geometry = get_sim_geometry(mesh_wvl_um) + + box_for_override = td.Box( + center=(0, 0, 0), + size=sim_geometry.size[0:2] + (thickness_box_placement_um + 0.3 * mesh_wvl_um,), + ) + + eval_fns, eval_fn_names = make_eval_fns(monitor_size_wvl) + + sim_path_dir = tmp_path / f"test{test_number}" + sim_path_dir.mkdir() + + mesh_cell_override_size = mesh_wvl_um / mesh_refinement_factor + + if RUN_WITH_FD_CONVERGENCE: + fd_step = np.linspace(2 * mesh_cell_override_size, mesh_cell_override_size, 2) + else: + fd_step = np.array([mesh_cell_override_size]) + + z_planes = np.linspace(-0.5 * thickness_box_placement_um, 0.5 * thickness_box_placement_um, 3) + + polyslab = generate_polyslab(0.5 * dim_um, dim_um, rng) + + polyslab_z_value = 0 + + objective = create_objective_function_2D( + lambda mesh_wvl_um=mesh_wvl_um, + adj_wvl_um=adj_wvl_um, + monitor_size_wvl=monitor_size_wvl, + box_for_override=box_for_override, + mesh_refinement_factor=mesh_refinement_factor: make_base_sim( + mesh_wvl_um=mesh_wvl_um, + adj_wvl_um=adj_wvl_um, + monitor_size_wvl=monitor_size_wvl, + box_for_override=box_for_override, + mesh_refinement_factor=mesh_refinement_factor, + ), + eval_fn, + polyslab_z_value, + sim_path_dir=str(sim_path_dir), + ) + + obj_val_and_grad = ag.value_and_grad(objective) + + obj, adj_grad = obj_val_and_grad([polyslab]) + + fd_grad, valid_mask = run_and_process_fd( + polyslab_parameters=polyslab, fd_step=fd_step, objective=objective + ) + + fd_grad = np.array(fd_grad) + adj_grad = np.squeeze(np.array(adj_grad)) + valid_mask = np.array(valid_mask) + + vertex_fd = fd_grad[valid_mask] + vertex_adj = adj_grad[valid_mask] + + assert len(vertex_fd) >= MIN_FD_COMPARISON, "Too many vertices were trimmed" + + ( + vertex_overlap_deg, + vertex_error_mean, + vertex_error_std, + vertex_error_norm_mean, + vertex_error_norm_std, + ) = compare_fd_adj(vertex_fd, vertex_adj) + + vertex_data = np.zeros((2, len(vertex_fd))) + + vertex_data[SAVE_FD_LOC, :] = vertex_fd + vertex_data[SAVE_ADJ_LOC, :] = vertex_adj + + if SHOW_PRINT_STATEMENTS: + print(f"\n2D PEC PolySlab Test {test_number} Summary:") + print(f"Mesh wavelength (um): {mesh_wvl_um}") + print(f"Adjoint wavelength (um): {adj_wvl_um}") + print(f"Monitor size (wavelengths): {monitor_size_wvl}") + print(f"Mesh refinement factor: {mesh_refinement_factor}") + print(f"Eval function: {eval_fn_name}") + print(f"Vertex mean (std): {vertex_error_mean} ({vertex_error_std})") + print(f"Vertex norm mean (std): {vertex_error_norm_mean} ({vertex_error_norm_std})") + print(f"Vertex overlap deg: {vertex_overlap_deg}") + print("\n") + + if SAVE_FD_ADJ_DATA: + np.save( + f"{NUMERICAL_RESULTS_DATA_DIR}/rf_polyslab_2d_vertex_fd_adj_data_test_{test_number}.npy", + vertex_data, + ) + + yvals = [vertex_error_mean, vertex_error_norm_mean] + yerrs = [vertex_error_std, vertex_error_norm_std] + bars = plt.bar( + ["vertex\nerror", "vertex\nerror\nnorm"], + yvals, + yerr=yerrs, + color="skyblue", + edgecolor="black", + capsize=5, + ) + + plt.title(f"Vector Alignment (deg):\n{vertex_overlap_deg:.2f}") + plt.xticks(rotation=45) + + for idx, bar in enumerate(bars): + yval = bar.get_height() + yerr = yerrs[idx] + plt.text( + bar.get_x() + 0.5 * bar.get_width(), + yval + yerr + 0.1, + round(yval, 3), + ha="center", + va="bottom", + ) + + plt.savefig( + f"{NUMERICAL_RESULTS_DATA_DIR}/rf_2d_polyslab_summary_plot_test_{test_number}.png" + ) + + +@pytest.mark.numerical +@pytest.mark.parametrize( + "rf_3d_test_parameters, dir_name", + zip( + rf_3d_test_parameters, + ([NUMERICAL_RESULTS_DATA_DIR] if SAVE_FD_ADJ_DATA else [None]) * len(rf_3d_test_parameters), + ), + indirect=["dir_name"], +) +def test_finite_difference_3d_polyslab_pec(rf_3d_test_parameters, rng, tmp_path, create_directory): + """Test a variety of autograd permittivity gradients for 3D PEC `PolySlab` by""" + """comparing them to numerical finite difference.""" + + test_number = rf_3d_test_parameters["test_number"] + + ( + mesh_wvl_um, + adj_wvl_um, + monitor_size_wvl, + polyslab_z_thickness_wvl, + eval_fn, + eval_fn_name, + test_number, + ) = operator.itemgetter( + "mesh_wvl_um", + "adj_wvl_um", + "monitor_size_wvl", + "polyslab_z_thickness_wvl", + "eval_fn", + "eval_fn_name", + "test_number", + )(rf_3d_test_parameters) + + dim_um = 1.5 * mesh_wvl_um + + thickness_box_placement_um = 1.0 * mesh_wvl_um + + sim_geometry = get_sim_geometry(mesh_wvl_um) + + box_for_override = td.Box( + center=(0, 0, 0), + size=sim_geometry.size[0:2] + (thickness_box_placement_um,), + ) + + eval_fns, eval_fn_names = make_eval_fns(monitor_size_wvl) + + sim_path_dir = tmp_path / f"test{test_number}" + sim_path_dir.mkdir() + + mesh_cell_override_size = mesh_wvl_um / FIXED_MESH_REFINEMENT_FACTOR + + if RUN_WITH_FD_CONVERGENCE: + fd_step = np.linspace(2 * mesh_cell_override_size, mesh_cell_override_size, 2) + else: + fd_step = np.array([mesh_cell_override_size]) + + polyslab = generate_polyslab(0.5 * dim_um, dim_um, rng) + polyslab_z_value = 0 + polyslab_z_thickness = polyslab_z_thickness_wvl * adj_wvl_um + + objective = create_objective_function_3D( + lambda mesh_wvl_um=mesh_wvl_um, + adj_wvl_um=adj_wvl_um, + monitor_size_wvl=monitor_size_wvl, + box_for_override=box_for_override, + mesh_refinement_factor=FIXED_MESH_REFINEMENT_FACTOR: make_base_sim( + mesh_wvl_um=mesh_wvl_um, + adj_wvl_um=adj_wvl_um, + monitor_size_wvl=monitor_size_wvl, + box_for_override=box_for_override, + mesh_refinement_factor=mesh_refinement_factor, + ), + eval_fn, + polyslab_z_value, + polyslab_z_thickness, + sim_path_dir=str(sim_path_dir), + ) + + obj_val_and_grad = ag.value_and_grad(objective) + + obj, adj_grad = obj_val_and_grad([polyslab]) + + fd_grad, valid_mask = run_and_process_fd( + polyslab_parameters=polyslab, fd_step=fd_step, objective=objective + ) + + fd_grad = np.array(fd_grad) + adj_grad = np.squeeze(np.array(adj_grad)) + valid_mask = np.array(valid_mask) + + vertex_fd = fd_grad[valid_mask] + vertex_adj = adj_grad[valid_mask] + + assert len(vertex_fd) >= MIN_FD_COMPARISON, "Too many vertices were trimmed" + + ( + vertex_overlap_deg, + vertex_error_mean, + vertex_error_std, + vertex_error_norm_mean, + vertex_error_norm_std, + ) = compare_fd_adj(vertex_fd, vertex_adj) + + vertex_data = np.zeros((2, len(vertex_fd))) + + vertex_data[SAVE_FD_LOC, :] = vertex_fd + vertex_data[SAVE_ADJ_LOC, :] = vertex_adj + + if SHOW_PRINT_STATEMENTS: + print(f"\n3D PEC PolySlab Test {test_number} Summary:") + print(f"Mesh wavelength (um): {mesh_wvl_um}") + print(f"Adjoint wavelength (um): {adj_wvl_um}") + print(f"Monitor size (wavelengths): {monitor_size_wvl}") + print(f"Polyslab z thickness (wavelengths): {polyslab_z_thickness_wvl}") + print(f"Eval function: {eval_fn_name}") + print(f"Vertex mean (std): {vertex_error_mean} ({vertex_error_std})") + print(f"Vertex norm mean (std): {vertex_error_norm_mean} ({vertex_error_norm_std})") + print(f"Vertex overlap deg: {vertex_overlap_deg}") + print("\n") + + if SAVE_FD_ADJ_DATA: + np.save( + f"{NUMERICAL_RESULTS_DATA_DIR}/rf_polyslab_3d_vertex_fd_adj_data_test_{test_number}.npy", + vertex_data, + ) + + yvals = [vertex_error_mean, vertex_error_norm_mean] + yerrs = [vertex_error_std, vertex_error_norm_std] + bars = plt.bar( + ["vertex\nerror", "vertex\nerror\nnorm"], + yvals, + yerr=yerrs, + color="skyblue", + edgecolor="black", + capsize=5, + ) + + plt.title(f"Vector Alignment (deg):\n{vertex_overlap_deg:.2f}") + plt.ylabel("Relative Error") + + for idx, bar in enumerate(bars): + yval = bar.get_height() + yerr = yerrs[idx] + plt.text( + bar.get_x() + 0.5 * bar.get_width(), + yval + yerr + 0.1, + round(yval, 3), + ha="center", + va="bottom", + ) + + plt.savefig( + f"{NUMERICAL_RESULTS_DATA_DIR}/rf_3d_polyslab_summary_plot_test_{test_number}.png" + ) diff --git a/tests/test_components/test_geometry.py b/tests/test_components/test_geometry.py index 0edab2fd35..93f576ae6b 100644 --- a/tests/test_components/test_geometry.py +++ b/tests/test_components/test_geometry.py @@ -1298,3 +1298,136 @@ def test_cleanup_shapely_object(): orig_polygon = shapely.Polygon(exterior_coords) new_polygon = cleanup_shapely_object(orig_polygon, tolerance_ratio=1e-12) assert len(new_polygon.exterior.coords) == 0 # empty / collinear polygons should get deleted + + +def test_snap_coords_outside(): + """Test coordinate snapping for box geometries.""" + num_coords = 101 + coords_to_snap = np.linspace(-1, 1, num_coords) + + coord_face_min = np.mean(coords_to_snap[0:2]) + coord_face_max = np.mean(coords_to_snap[-2:]) + + snap_pt_min = td.Box._snap_coords_outside( + min_max_index=0, snap_coords_values=coords_to_snap, coord_normal_face=coord_face_min + ) + snap_pt_max = td.Box._snap_coords_outside( + min_max_index=1, snap_coords_values=coords_to_snap, coord_normal_face=coord_face_max + ) + + assert snap_pt_min == coords_to_snap[0], "Unexpected snapping point minimum" + assert snap_pt_max == coords_to_snap[-1], "Unexpected snapping point maximum" + + with AssertLogLevel("WARNING", contains_str="Unable to snap coordinates outside"): + snap_pt_min = td.Box._snap_coords_outside( + min_max_index=0, + snap_coords_values=coords_to_snap, + coord_normal_face=coords_to_snap[0] - 0.1, + ) + + assert snap_pt_min == coords_to_snap[0], "Unexpected snapping point minimum" + + with AssertLogLevel("WARNING", contains_str="Unable to snap coordinates outside"): + snap_pt_min = td.Box._snap_coords_outside( + min_max_index=1, + snap_coords_values=coords_to_snap, + coord_normal_face=coords_to_snap[-1] + 0.1, + ) + + assert snap_pt_min == coords_to_snap[-1], "Unexpected snapping point maxium" + + +def test_singularity_correction_pec(): + """Tests singularity correction detection used for PEC gradients of `Box` geometries.""" + + check_3d_size = (1.0, 2.0, 3.0) + + for axis in range(3): + is_2d, zero_dimension, do_singularity_correction = td.Box._check_singularity_correction_pec( + size=check_3d_size, axis_normal=axis + ) + assert not is_2d, "Unexpectedly found 2D geometry" + assert not zero_dimension, "Unexpectedly found zero_dimension" + assert not do_singularity_correction, ( + "Unexpectedly identifying need for singularity correction" + ) + + check_1d_error_size = (1.0, 0.0, 0.0) + + for axis in range(3): + with AssertLogLevel( + "ERROR", contains_str="Derivative of PEC material with less than 2 dimensions" + ): + is_2d, zero_dimension, do_singularity_correction = ( + td.Box._check_singularity_correction_pec(size=check_1d_error_size, axis_normal=axis) + ) + + check_2d_size = 1.0 + + for index_2d in range(3): + check_2d_size = tuple(0.0 if idx == index_2d else 1.0 for idx in range(3)) + + for axis in range(3): + print(f"check 2d size = {check_2d_size}") + + is_2d, zero_dimension, do_singularity_correction = ( + td.Box._check_singularity_correction_pec(size=check_2d_size, axis_normal=axis) + ) + + print(f"is 2d = {is_2d}") + + assert is_2d, "Expected 2D geometry detection" + assert zero_dimension == "xyz"[index_2d], "Wrong zero dimension identified" + + if axis == index_2d: + assert not do_singularity_correction, ( + "Unexpectedly detecting need for singularity correction" + ) + else: + assert do_singularity_correction, ( + "Unexpectedly not detecting need for singularity correction" + ) + + +def test_trim_dims_and_bounds_edge(): + """Tests the dimension and bound trimming for 2D PEC gradient integrations in `Box`.""" + + for zero_dimension_idx in range(3): + box = td.Box( + center=(0.0, 0.0, 0.0), + size=tuple(0.0 if zero_dimension_idx == idx else 1.0 * (idx + 1) for idx in range(3)), + ) + + for axis_normal in range(3): + if axis_normal == zero_dimension_idx: + continue + + dim_normal, dims_perp = box.pop_axis("xyz", axis=axis_normal) + bounds_normal, bounds_perp = box.pop_axis(np.array(box.bounds).T, axis=axis_normal) + bounds_perp = np.array(bounds_perp).T + + trimmed_dims, trimmed_bounds = td.Box._trim_dims_and_bounds_edge( + dims_perp=dims_perp, + bounds_perp=bounds_perp, + zero_dimension="xyz"[zero_dimension_idx], + ) + + expected_trimmed_idx = 0 + for idx in range(3): + if (idx == axis_normal) or (idx == zero_dimension_idx): + continue + + expected_trimmed_idx = idx + + expected_trimmed_dims = ("xyz"[expected_trimmed_idx],) + expected_trimmed_bounds = ( + (-0.5 * (expected_trimmed_idx + 1),), + (0.5 * (expected_trimmed_idx + 1),), + ) + + assert np.all(np.array(expected_trimmed_dims) == np.array(trimmed_dims)), ( + "Unexpected trimmed dims" + ) + assert np.all(np.array(expected_trimmed_bounds) == np.array(trimmed_bounds)), ( + "Unexpected trimmed bounds" + ) diff --git a/tidy3d/components/autograd/derivative_utils.py b/tidy3d/components/autograd/derivative_utils.py index b5a21acde0..ebc1011cd3 100644 --- a/tidy3d/components/autograd/derivative_utils.py +++ b/tidy3d/components/autograd/derivative_utils.py @@ -10,7 +10,7 @@ from tidy3d.components.data.data_array import FreqDataArray, ScalarFieldDataArray from tidy3d.components.types import ArrayLike, Bound, tidycomplex -from tidy3d.constants import C_0, LARGE_NUMBER +from tidy3d.constants import C_0, EPSILON_0, LARGE_NUMBER, MU_0 from tidy3d.log import log from .constants import ( @@ -111,9 +111,32 @@ class DerivativeInfo: Bounds corresponding to the minimum intersection between the structure and the simulation it is contained in.""" + simulation_bounds: Bound + """Simulation bounds. + Bounds corresponding to the simulation domain containing this structure. + Unlike bounds_intersect, this is independent of the structure's bounds and + is purely based on the simulation geometry.""" + frequencies: ArrayLike """Frequencies at which the adjoint gradient should be computed.""" + H_der_map: Optional[FieldData] = None + """Magnetic field gradient map. + Dataset where the field components ("Hx", "Hy", "Hz") store the multiplication + of the forward and adjoint magnetic fields. The tangential component of this + dataset is used when computing adjoint gradients for shifting boundaries of + structures composed of PEC mediums.""" + + H_fwd: Optional[FieldData] = None + """Forward magnetic fields. + Dataset where the field components ("Hx", "Hy", "Hz") represent the forward + magnetic fields used for computing gradients for a given structure.""" + + H_adj: Optional[FieldData] = None + """Adjoint magnetic fields. + Dataset where the field components ("Hx", "Hy", "Hz") represent the adjoint + magnetic fields used for computing gradients for a given structure.""" + # Optional fields with defaults eps_background: Optional[EpsType] = None """Permittivity in background. @@ -139,6 +162,11 @@ class DerivativeInfo: GeometryGroup handling where it is difficult to automatically evaluate the inside and outside relative permittivity for each geometry.""" + is_medium_pec: bool = False + """Indicates if structure material is PEC. + If True, the structure contains a PEC material which changes the gradient + formulation at the boundary compared to the dielectric case.""" + interpolators: Optional[dict] = None """Pre-computed interpolators. Optional pre-computed interpolators for field components and permittivity data. @@ -235,7 +263,6 @@ def creator_func(arr=arr, points=points): data = arr.data.astype( GRADIENT_DTYPE_COMPLEX if np.iscomplexobj(arr.data) else dtype, copy=False ) - # create interpolator with frequency dimension if "f" in arr.dims: freq_coords = arr.coords["f"].data.astype(dtype, copy=False) @@ -251,8 +278,13 @@ def creator_func(arr=arr, points=points): data = data[..., np.newaxis] points_with_freq = (*points, freq_coords) + # If PEC, use nearest interpolation instead of linear to avoid interpolating + # with field values inside the PEC (which are 0). Instead, we make sure to + # choose interplation points such that their nearest location is outside of + # the PEC surface. + method = "nearest" if self.is_medium_pec else "linear" interpolator_obj = RegularGridInterpolator( - points_with_freq, data, method="linear", bounds_error=False, fill_value=None + points_with_freq, data, method=method, bounds_error=False, fill_value=None ) def interpolator(coords): @@ -275,12 +307,16 @@ def interpolator(coords): else: interpolators[component_name] = LazyInterpolator(creator_func) - for group_key, data_dict in [ + # process field interpolators (nested dictionaries) + interpolator_groups = [ ("E_fwd", self.E_fwd), ("E_adj", self.E_adj), ("D_fwd", self.D_fwd), ("D_adj", self.D_adj), - ]: + ] + if self.is_medium_pec: + interpolator_groups += [("H_fwd", self.H_fwd), ("H_adj", self.H_adj)] + for group_key, data_dict in interpolator_groups: _make_lazy_interpolator_group(data_dict, group_key, is_field_group=True) if self.eps_inf_structure is not None: @@ -333,6 +369,39 @@ def evaluate_gradient_at_points( "Please create interpolators using 'create_interpolators()' first." ) + if "eps_no" in interpolators: + eps_out = interpolators["eps_no"](spatial_coords) + else: + # use eps_background if available, otherwise use eps_out + eps_to_prepare = ( + self.eps_background if self.eps_background is not None else self.eps_out + ) + eps_out = self._prepare_epsilon(eps_to_prepare) + + if self.is_medium_pec: + vjps = self._evaluate_pec_gradient_at_points( + spatial_coords, normals, perps1, perps2, interpolators, eps_out + ) + else: + vjps = self._evaluate_dielectric_gradient_at_points( + spatial_coords, normals, perps1, perps2, interpolators, eps_out + ) + + # sum over frequency dimension + vjps = np.sum(vjps, axis=-1) + + return vjps + + def _evaluate_dielectric_gradient_at_points( + self, + spatial_coords: np.ndarray, + normals: np.ndarray, + perps1: np.ndarray, + perps2: np.ndarray, + interpolators: dict, + # todo: type + eps_out, + ) -> np.ndarray: # evaluate all field components at surface points E_fwd_at_coords = { name: interp(spatial_coords) for name, interp in interpolators["E_fwd"].items() @@ -347,6 +416,14 @@ def evaluate_gradient_at_points( name: interp(spatial_coords) for name, interp in interpolators["D_adj"].items() } + if "eps_inf" in interpolators: + eps_in = interpolators["eps_inf"](spatial_coords) + else: + eps_in = self._prepare_epsilon(self.eps_in) + + delta_eps_inv = 1.0 / eps_in - 1.0 / eps_out + delta_eps = eps_in - eps_out + # project fields onto local surface basis (normal + two tangents) D_fwd_norm = self._project_in_basis(D_fwd_at_coords, basis_vector=normals) D_adj_norm = self._project_in_basis(D_adj_at_coords, basis_vector=normals) @@ -361,27 +438,231 @@ def evaluate_gradient_at_points( E_der_perp1 = E_fwd_perp1 * E_adj_perp1 E_der_perp2 = E_fwd_perp2 * E_adj_perp2 - if "eps_inf" in interpolators: - eps_in = interpolators["eps_inf"](spatial_coords) - else: - eps_in = self._prepare_epsilon(self.eps_in) + vjps = -delta_eps_inv * D_der_norm + E_der_perp1 * delta_eps + E_der_perp2 * delta_eps - if "eps_no" in interpolators: - eps_out = interpolators["eps_no"](spatial_coords) - else: - # use eps_background if available, otherwise use eps_out - eps_to_prepare = ( - self.eps_background if self.eps_background is not None else self.eps_out + return vjps + + def _evaluate_pec_gradient_at_points( + self, + spatial_coords: np.ndarray, + normals: np.ndarray, + perps1: np.ndarray, + perps2: np.ndarray, + interpolators: dict, + # todo: type + eps_out, + ) -> np.ndarray: + def _adjust_spatial_coords_pec(grid_centers: dict[str, np.ndarray]): + """Assuming a nearest interpolation, adjust the interpolation points given the grid + defined by `grid_centers` and using `spatial_coords` as a starting point such that we + select a point outside of the PEC boundary. + + *** (nearest point outside boundary) + ^ + | n (normal direction) + | + _.-~'`-._.-~'`-._ (PEC surface) + * (nearest point) + + Parameters + ---------- + grid_centers: dict[str, np.ndarray] + The grid points for a given field component indexed by dimension. These grid points + are used to find the nearest snapping point and adjust the inerpolation coordinates + to ensure we fall outside of the PEC surface. + + Returns + ------- + (np.ndarray, np.ndarray) + (N, 3) array of coordinate centers at which to interpolate such that they line up + with a grid center and are outside the PEC surface + (N,) array of distances from the nearest interpolation points to the desired surface + edge points specified by `spatial_coords` + + """ + grid_ddim = np.zeros_like(normals) + for idx, dim in enumerate("xyz"): + expanded_coords = np.expand_dims(spatial_coords[:, idx], axis=1) + grid_centers_select = grid_centers[dim] + + diff = np.abs(expanded_coords - grid_centers_select) + + nearest_grid = np.argmin(diff, axis=-1) + nearest_grid = np.minimum(np.maximum(nearest_grid, 1), len(grid_centers_select) - 1) + + # compute the local grid spacing near the boundary + grid_ddim[:, idx] = ( + grid_centers_select[nearest_grid] - grid_centers_select[nearest_grid - 1] + ) + + # assuming we move in the normal direction, finds which dimension we need to move the least + # in order to ensure we snap to a point outside the boundary in the worst case (i.e. - the + # nearest point is just inside the surface) + min_movement_index = np.argmin( + np.abs(grid_ddim) / (np.abs(normals) + np.finfo(normals.dtype).min), axis=1 ) - eps_out = self._prepare_epsilon(eps_to_prepare) - delta_eps_inv = 1.0 / eps_in - 1.0 / eps_out - delta_eps = eps_in - eps_out + selection = (np.arange(normals.shape[0]), min_movement_index) + coords_dn = np.expand_dims(np.abs(grid_ddim[selection]), axis=1) + + # adjust coordinates by half a grid point outside boundary such that nearest interpolation + # point snaps to outside the boundary + adjust_spatial_coords = np.squeeze(spatial_coords + normals * 0.5 * coords_dn) + + edge_distance = np.zeros_like(adjust_spatial_coords[:, 0]) + for idx, dim in enumerate("xyz"): + expanded_adjusted_coords = np.expand_dims(adjust_spatial_coords[:, idx], axis=1) + grid_centers_select = grid_centers[dim] + + # find nearest grid point from the adjusted coordinates + diff = np.abs(expanded_adjusted_coords - grid_centers_select) + nearest_grid = np.argmin(diff, axis=-1) + + # compute edge distance from the nearest interpolated point to the boundary edge + edge_distance += ( + np.abs(spatial_coords[:, idx] - grid_centers_select[nearest_grid]) ** 2 + ) + + # this edge distance is useful when correcting for edge singularities from the PEC material + # and is used when the PEC PolySlab structure has zero thickness + edge_distance = np.sqrt(edge_distance) + + return adjust_spatial_coords, edge_distance + + def _snap_coordinate_outside(field_components: FieldData): + """Helper function to perform coordinate adjustment and compute edge distance for each + component in `field_components`. + + Parameters + ---------- + field_components: FieldData + The field components (i.e - Ex, Ey, Ez, Hx, Hy, Hz) that we would like to sample just + outside the PEC surface using nearest interpolation. + + Returns + ------- + dict[str, dict[str, np.ndarray]] + Dictionary mapping each field component name to a dictionary of adjusted coordinates + and edge distances for that component. + """ + adjustment = {} + for name in field_components: + field_component = field_components[name] + field_component_coords = field_component.coords + + adjusted_coords, edge_distance = _adjust_spatial_coords_pec( + { + key: np.array(field_component_coords[key].values) + for key in field_component_coords + } + ) + adjustment[name] = {"coords": adjusted_coords, "edge_distance": edge_distance} + + return adjustment + + def _interpolate_field_components(interp_coords, field_name): + return { + name: interp(interp_coords[name]["coords"]) + for name, interp in interpolators[field_name].items() + } + + # adjust coordinates for PEC to be outside structure bounds and get edge distance for singularity correction. + E_fwd_coords_adjusted = _snap_coordinate_outside(self.E_fwd) + E_adj_coords_adjusted = _snap_coordinate_outside(self.E_adj) + + H_fwd_coords_adjusted = _snap_coordinate_outside(self.H_fwd) + H_adj_coords_adjusted = _snap_coordinate_outside(self.H_adj) + + # using the adjusted coordinates, evaluate all field components at surface points + E_fwd_at_coords = _interpolate_field_components(E_fwd_coords_adjusted, "E_fwd") + E_adj_at_coords = _interpolate_field_components(E_adj_coords_adjusted, "E_adj") + H_fwd_at_coords = _interpolate_field_components(H_fwd_coords_adjusted, "H_fwd") + H_adj_at_coords = _interpolate_field_components(H_adj_coords_adjusted, "H_adj") + + structure_sizes = np.array( + [self.bounds[1][idx] - self.bounds[0][idx] for idx in range(len(self.bounds[0]))] + ) + + is_flat_perp_dim1 = np.isclose(np.abs(np.sum(perps1[0] * structure_sizes)), 0.0) + is_flat_perp_dim2 = np.isclose(np.abs(np.sum(perps2[0] * structure_sizes)), 0.0) + flat_perp_dims = [is_flat_perp_dim1, is_flat_perp_dim2] + + # check if this integration is happening along an edge in which case we will eliminate + # on of the H field integration components and apply singularity correction + pec_line_integration = is_flat_perp_dim1 or is_flat_perp_dim2 + + def _compute_singularity_correction(adjustment_: dict[str, dict[str, np.ndarray]]): + """ + Given the `adjustment_` which contains the distance from the PEC edge each field + component is nearest interpolated at, computes the singularity correction when + working with 2D PEC using the average edge_distance for each component. In the case + of 3D PEC gradients, no singularity correction is applied so an array of ones is returned. + + Parameters + ---------- + adjustment_: dict[str, dict[str, np.ndarray]] + Dictionary that maps field component name to a dictionary containing the coordinate + adjustment and the distance to the PEC edge for those coordinates. The edge distance + is used for 2D PEC singularity correction. + + Returns + ------- + np.ndarray + Returns the singularity correction which has shape (N,) where there are N points in + `spatial_coords` + """ + return ( + ( + 0.5 + * np.pi + * np.mean([adjustment_[name]["edge_distance"] for name in adjustment_], axis=0) + ) + if pec_line_integration + else np.ones_like(spatial_coords, shape=spatial_coords.shape[0]) + ) - vjps = -delta_eps_inv * D_der_norm + E_der_perp1 * delta_eps + E_der_perp2 * delta_eps + E_norm_singularity_correction = np.expand_dims( + _compute_singularity_correction(E_fwd_coords_adjusted), axis=1 + ) + H_perp_singularity_correction = np.expand_dims( + _compute_singularity_correction(H_fwd_coords_adjusted), axis=1 + ) + + E_fwd_norm = self._project_in_basis(E_fwd_at_coords, basis_vector=normals) + E_adj_norm = self._project_in_basis(E_adj_at_coords, basis_vector=normals) + + # compute the normal E contribution to the gradient (the tangential E contribution + # is 0 in the case of PEC since this field component is continuous and thus 0 at + # the boundary) + contrib_E = E_norm_singularity_correction * eps_out * E_fwd_norm * E_adj_norm + vjps = contrib_E + + # compute the tangential H contribution to the gradient (the normal H contribution + # is 0 for PEC) + H_fwd_perp1 = self._project_in_basis(H_fwd_at_coords, basis_vector=perps1) + H_adj_perp1 = self._project_in_basis(H_adj_at_coords, basis_vector=perps1) + + H_fwd_perp2 = self._project_in_basis(H_fwd_at_coords, basis_vector=perps2) + H_adj_perp2 = self._project_in_basis(H_adj_at_coords, basis_vector=perps2) + + H_der_perp1 = H_perp_singularity_correction * H_fwd_perp1 * H_adj_perp1 + H_der_perp2 = H_perp_singularity_correction * H_fwd_perp2 * H_adj_perp2 + + H_integration_components = (H_der_perp1, H_der_perp2) + if pec_line_integration: + # if we are integrating along the line, we choose the H component normal to + # the edge which corresponds to a surface current along the edge whereas the other + # tangential component corresponds to a surface current along the flat dimension. + H_integration_components = tuple( + H_comp for idx, H_comp in enumerate(H_integration_components) if flat_perp_dims[idx] + ) - # sum over frequency dimension - vjps = np.sum(vjps, axis=-1) + # for each of the tangential components we are integrating the H fields over, + # adjust weighting to account for pre-weighting of the source by `EPSILON_0` + # and multiply by appropriate `MU_0` factor + for H_perp in H_integration_components: + contrib_H = MU_0 * H_perp / EPSILON_0 + vjps += contrib_H return vjps diff --git a/tidy3d/components/geometry/base.py b/tidy3d/components/geometry/base.py index b00440b5e7..c764dafb8f 100644 --- a/tidy3d/components/geometry/base.py +++ b/tidy3d/components/geometry/base.py @@ -20,8 +20,14 @@ from tidy3d.compat import _shapely_is_older_than from tidy3d.components.autograd import AutogradFieldMap, TracedCoordinate, TracedSize, get_static from tidy3d.components.autograd.constants import GRADIENT_DTYPE_FLOAT -from tidy3d.components.autograd.derivative_utils import DerivativeInfo, integrate_within_bounds +from tidy3d.components.autograd.derivative_utils import ( + DerivativeInfo, + FieldData, + integrate_within_bounds, +) from tidy3d.components.base import Tidy3dBaseModel, cached_property +from tidy3d.components.data.data_array import ScalarFieldDataArray +from tidy3d.components.geometry.bound_ops import bounds_intersection, bounds_union from tidy3d.components.transformation import ReflectionFromPlane, RotationAroundAxis from tidy3d.components.types import ( ArrayFloat2D, @@ -51,7 +57,7 @@ polygon_patch, set_default_labels_and_title, ) -from tidy3d.constants import LARGE_NUMBER, MICROMETER, RADIAN, fp_eps, inf +from tidy3d.constants import EPSILON_0, LARGE_NUMBER, MICROMETER, MU_0, RADIAN, fp_eps, inf from tidy3d.exceptions import ( SetupError, Tidy3dError, @@ -62,8 +68,6 @@ from tidy3d.log import log from tidy3d.packaging import verify_packages_import -from .bound_ops import bounds_intersection, bounds_union - POLY_GRID_SIZE = 1e-12 POLY_TOLERANCE_RATIO = 1e-12 POLY_DISTANCE_TOLERANCE = 8e-12 @@ -2440,11 +2444,11 @@ def _derivative_face( # normal and tangential dims dim_normal, dims_perp = self.pop_axis("xyz", axis=axis_normal) - fld_normal, flds_perp = self.pop_axis(("Ex", "Ey", "Ez"), axis=axis_normal) + fld_E_normal, flds_E_perp = self.pop_axis(("Ex", "Ey", "Ez"), axis=axis_normal) # fields and bounds - D_normal = derivative_info.D_der_map[fld_normal] - Es_perp = tuple(derivative_info.E_der_map[key] for key in flds_perp) + D_normal = derivative_info.D_der_map[fld_E_normal] + Es_perp = tuple(derivative_info.E_der_map[key] for key in flds_E_perp) bounds_normal, bounds_perp = self.pop_axis( np.array(derivative_info.bounds).T, axis=axis_normal ) @@ -2494,39 +2498,325 @@ def _derivative_face( eps_in_normal, eps_in_perps = self.pop_axis(eps_xyz_inside, axis=axis_normal) eps_out_normal, eps_out_perps = self.pop_axis(eps_xyz_outside, axis=axis_normal) - # compute integration pre-factors - delta_eps_perps = [eps_in - eps_out for eps_in, eps_out in zip(eps_in_perps, eps_out_perps)] - delta_eps_inv_normal = 1.0 / eps_in_normal - 1.0 / eps_out_normal + if derivative_info.is_medium_pec: + return self._derivative_face_pec( + dim_normal=dim_normal, + axis_normal=axis_normal, + min_max_index=min_max_index, + coord_normal_face=coord_normal_face, + dims_perp=dims_perp, + bounds_perp=bounds_perp, + D_normal=D_normal, + H_der_map=derivative_info.H_der_map, + eps_out_normal=eps_out_normal, + ) + else: + return self._derivative_face_dielectric( + dim_normal=dim_normal, + coord_normal_face=coord_normal_face, + dims_perp=dims_perp, + bounds_perp=bounds_perp, + D_normal=D_normal, + Es_perp=Es_perp, + eps_in_normal=eps_in_normal, + eps_out_normal=eps_out_normal, + eps_in_perps=eps_in_perps, + eps_out_perps=eps_out_perps, + ) + + @staticmethod + def _arr_at_face(arr: xr.DataArray, interp_point: float, dim_normal: str) -> xr.DataArray: + """Interpolate the array at the given coordinate unless it only has a single value in the normal + dimnsion already in which case just return the array itself.""" + arr_at_face = ( + arr + if (len(arr.coords[dim_normal]) == 1) + else arr.interp(**{dim_normal: float(interp_point)}, assume_sorted=True) + ) + + return arr_at_face + + @staticmethod + def _integrate_face( + arr_at_face: xr.DataArray, + integration_dims: Union[tuple[str], tuple[str, str]], + integration_bounds: Union[tuple[Coordinate2D], tuple[Coordinate2D, Coordinate2D]], + ) -> complex: + """Perform the integration of the surface and sum over the frequency component of the gradient.""" + + integral_result = integrate_within_bounds( + arr=arr_at_face, + dims=integration_dims, + bounds=integration_bounds, + ) + + return complex(integral_result.sum("f")) + + @staticmethod + def _snap_coords_outside( + min_max_index: int, snap_coords_values: np.ndarray, coord_normal_face: float + ) -> float: + """Snap interpolation coordinate for a PEC face integration to be just outside the surface boundary. + This ensures we don't interpolate with fields that are zero inside of the PEC.""" + + if min_max_index == 0: + min_boundary_mapping = np.where(snap_coords_values > coord_normal_face)[0] + index_face = ( + 0 + if (len(min_boundary_mapping) == 0) + else np.maximum(0, min_boundary_mapping[0] - 1) + ) + + if snap_coords_values[index_face] > coord_normal_face: + log.warning("Unable to snap coordinates outside of min face.") + else: + max_boundary_mapping = np.where(snap_coords_values < coord_normal_face)[0] + index_face = ( + len(snap_coords_values) - 1 + if (len(max_boundary_mapping) == 0) + else np.minimum(len(snap_coords_values) - 1, max_boundary_mapping[-1] + 1) + ) + + if snap_coords_values[index_face] < coord_normal_face: + log.warning("Unable to snap coordinates outside of max face.") - def integrate_face(arr: xr.DataArray) -> complex: - """Interpolate and integrate a scalar field data over the face using bounds.""" + snapped_point = snap_coords_values[index_face] - arr_at_face = arr.interp(**{dim_normal: float(coord_normal_face)}, assume_sorted=True) + return snapped_point - integral_result = integrate_within_bounds( - arr=arr_at_face, - dims=dims_perp, - bounds=bounds_perp, + @staticmethod + def _check_singularity_correction_pec( + size: TracedSize, axis_normal: Axis + ) -> tuple[bool, str, bool]: + """Checks if the box is 2D (i.e. - one of the dimensions is zero) and + identifies the zero dimension if any. Then, checks if we should apply singularity + correction if we are integrating a face with that contains the zero dimension.""" + + # detect whether the box is 2-dimensional + zero_size_map = [s == 0.0 for s in size] + + dimension = 3 - np.sum(zero_size_map) + if dimension < 2: + log.error( + "Derivative of PEC material with less than 2 dimensions is unsupported. " + f"Specified PEC box is {dimension}-dimesional" ) - return complex(integral_result.sum("f")) + do_singularity_correction = False + is_2d = np.any(zero_size_map) + zero_dimension = None + if is_2d: + zero_dim_idx = np.where(zero_size_map)[0][0] + zero_dimension = "xyz"[zero_dim_idx] + + # for singularity correction, need 2-dimensional box and that + # the face we are integrating over is the 1-dimensional (i.e. - + # the normal for the face is not the same as the flat dimension + do_singularity_correction = not (axis_normal == zero_dim_idx) - # compute vjp from field integrals - vjp_value = 0.0 + return is_2d, zero_dimension, do_singularity_correction + + @staticmethod + def _trim_dims_and_bounds_edge( + dims_perp: tuple[str, str], + bounds_perp: tuple[Coordinate2D, Coordinate2D], + zero_dimension: str, + ) -> tuple[tuple[str], tuple[tuple[float], tuple[float]]]: + """Trim the dimensions and bounds for integration along the edge.""" + + # if we are correcting for singularity and integrating over the line, then adjust + # integration dimensions and bounds to exclude the flat dimension + integration_dims = [dim for dim in dims_perp if (not (dim == zero_dimension))] + + integration_bounds = [] + zero_dim_idx = 0 + # find the index into the bounds that corresponds to the flat dimension + for dim in dims_perp: + if dim == zero_dimension: + break + zero_dim_idx += 1 + + # trim the zero dimension from the integration bounds + for bound in bounds_perp: + new_bound = [b for b_idx, b in enumerate(bound) if b_idx != zero_dim_idx] + + integration_bounds.append(new_bound) + + return integration_dims, integration_bounds + + def _derivative_face_dielectric( + self, + dim_normal: str, + coord_normal_face: float, + dims_perp: tuple[str, str], + bounds_perp: tuple[Coordinate2D, Coordinate2D], + D_normal: ScalarFieldDataArray, + Es_perp: tuple[ScalarFieldDataArray, ScalarFieldDataArray], + eps_in_normal: ScalarFieldDataArray, + eps_out_normal: ScalarFieldDataArray, + eps_in_perps: tuple[ScalarFieldDataArray, ScalarFieldDataArray], + eps_out_perps: tuple[ScalarFieldDataArray, ScalarFieldDataArray], + ) -> float: + """Compute derivative with respect to the face using the dielectric form of the gradient. + + Parameters + ---------- + dtype : np.dtype = GRADIENT_DTYPE_FLOAT + Data type for interpolation coordinates and values. + + dim_normal : str + Surface normal of the face + coord_normal_face : float + The coordinate at which to interpolate the surface fields + dims_perp : tuple[str, str] + The perpendicular dimensions of the face + bounds_perp : tuple[Coordinate2D, Coordinate2D] + Bounds of integration along the face + D_normal : ScalarFieldDataArray + D-field component normal to the surface + Es_perp : tuple[ScalarFieldDataArray, ScalarFieldDataArray] + E-field components tangential to the surface + eps_in_normal : ScalarFieldDataArray + Normal component of permittivity inside the surface + eps_out_normal : ScalarFieldDataArray + Normal component of permittivity outside the surface + eps_in_perps : tuple[ScalarFieldDataArray, ScalarFieldDataArray] + Tangential components of permittivity inside the surface + eps_out_perps : tuple[ScalarFieldDataArray, ScalarFieldDataArray] + Tangential components of permittivity outside the surface + + Returns + ------- + float + Surface gradient vjp value + """ + + delta_eps_inv_normal = 1.0 / eps_in_normal - 1.0 / eps_out_normal + + # compute integration pre-factors + delta_eps_perps = [eps_in - eps_out for eps_in, eps_out in zip(eps_in_perps, eps_out_perps)] - # perform D-normal integral integrand_D = -delta_eps_inv_normal * D_normal - integral_D = integrate_face(integrand_D) - vjp_value += integral_D + integral_D = self._integrate_face( + self._arr_at_face(integrand_D, coord_normal_face, dim_normal), + integration_dims=dims_perp, + integration_bounds=bounds_perp, + ) + + vjp_value = integral_D # perform E-perpendicular integrals for E_perp, delta_eps_perp in zip(Es_perp, delta_eps_perps): integrand_E = E_perp * delta_eps_perp - integral_E = integrate_face(integrand_E) + integral_E = self._integrate_face( + self._arr_at_face(integrand_E, coord_normal_face, dim_normal), + integration_dims=dims_perp, + integration_bounds=bounds_perp, + ) + vjp_value += integral_E return np.real(vjp_value) + def _derivative_face_pec( + self, + dim_normal: str, + axis_normal: Axis, + min_max_index: int, + coord_normal_face: float, + dims_perp: tuple[str, str], + bounds_perp: tuple[Coordinate2D, Coordinate2D], + D_normal: ScalarFieldDataArray, + H_der_map: FieldData, + eps_out_normal: ScalarFieldDataArray, + ) -> float: + """Compute derivative with respect to the face using the PEC form of the gradient. + + Parameters + ---------- + dtype : np.dtype = GRADIENT_DTYPE_FLOAT + Data type for interpolation coordinates and values. + + dim_normal : str + Surface normal of the face + axis_normal : Axis + Axis (index) corresponding to the normal of the face + min_max_index : int + Indicator for if the face is on the minimum of the box (0) or the maximum of the Box (1) + coord_normal_face : float + The coordinate at which to interpolate the surface fields + dims_perp : tuple[str, str] + The perpendicular dimensions of the face + bounds_perp : tuple[Coordinate2D, Coordinate2D] + Bounds of integration along the face + D_normal : ScalarFieldDataArray + D-field component normal to the surface + H_der_map : FieldData + Multiplication of H-field components in the Box region + eps_out_normal : ScalarFieldDataArray + Normal component of permittivity outside the surface + + Returns + ------- + float + Surface gradient vjp value + """ + + fld_H_normal, flds_H_perp = self.pop_axis(("Hx", "Hy", "Hz"), axis=axis_normal) + Hs_perp = tuple(H_der_map[key] for key in flds_H_perp) + + is_2d, zero_dimension, do_singularity_correction = self._check_singularity_correction_pec( + self.size, axis_normal + ) + + integration_dims = dims_perp + integration_bounds = bounds_perp + if do_singularity_correction: + integration_dims, integration_bounds = self._trim_dims_and_bounds_edge( + dims_perp, bounds_perp, zero_dimension + ) + + integrand_E = D_normal / np.real(eps_out_normal) + + # apply singularity correction only when integrating along the line + def _apply_singularity_correction(interp_point: float) -> float: + edge_distance = np.abs(interp_point - coord_normal_face) + return (0.5 * np.pi * edge_distance) if do_singularity_correction else 1.0 + + snap_E_coord = self._snap_coords_outside( + min_max_index, integrand_E.coords[dim_normal].values, coord_normal_face + ) + + integral_E = self._integrate_face( + self._arr_at_face(integrand_E, snap_E_coord, dim_normal), + integration_dims=integration_dims, + integration_bounds=integration_bounds, + ) + + vjp_value = _apply_singularity_correction(snap_E_coord) * integral_E + + for H_perp_idx, H_perp in enumerate(Hs_perp): + integrand_H = MU_0 * H_perp / EPSILON_0 + + if (is_2d and not (flds_H_perp[H_perp_idx] == f"H{zero_dimension}")) and ( + dim_normal != zero_dimension + ): + continue + + snap_H_coord = self._snap_coords_outside( + min_max_index, integrand_H.coords[dim_normal].values, coord_normal_face + ) + + integral_H = self._integrate_face( + self._arr_at_face(integrand_H, snap_H_coord, dim_normal), + integration_dims=integration_dims, + integration_bounds=integration_bounds, + ) + + vjp_value += _apply_singularity_correction(snap_H_coord) * integral_H + + return np.real(vjp_value) + """Compound subclasses""" diff --git a/tidy3d/components/geometry/polyslab.py b/tidy3d/components/geometry/polyslab.py index 41a7a2ebf3..d02d16a406 100644 --- a/tidy3d/components/geometry/polyslab.py +++ b/tidy3d/components/geometry/polyslab.py @@ -1447,13 +1447,15 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField """ vjps: AutogradFieldMap = {} - sim_min, sim_max = map(np.asarray, derivative_info.bounds_intersect) - extents = sim_max - sim_min + intersect_min, intersect_max = map(np.asarray, derivative_info.bounds_intersect) + sim_min, sim_max = map(np.asarray, derivative_info.simulation_bounds) + + extents = intersect_max - intersect_min is_2d = np.isclose(extents[self.axis], 0.0) # early return if polyslab is not in simulation domain slab_min, slab_max = self.slab_bounds - if (slab_max <= sim_min[self.axis]) or (slab_min >= sim_max[self.axis]): + if (slab_max < sim_min[self.axis]) or (slab_min > sim_max[self.axis]): log.warning( "'PolySlab' lies completely outside the simulation domain.", log_once=True, @@ -1473,6 +1475,7 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField vjps[path] = self._compute_derivative_vertices( derivative_info, sim_min, sim_max, is_2d, interpolators ) + elif path[0] == "slab_bounds": idx = path[1] face_coord = self.slab_bounds[idx] @@ -1508,6 +1511,12 @@ def _compute_derivative_slab_bounds( is split equally between the two vertices that bound the edge segment. """ # rmin/rmax over the geometry and simulation box + if np.isclose(self.slab_bounds[1] - self.slab_bounds[0], 0.0): + log.warning( + "Computing slab face derivatives for flat structures is not fully supported and " + "may give zero for the derivative. Try using a structure with a small, but nonzero " + "thickness for slab bound derivatives." + ) rmin, rmax = derivative_info.bounds_intersect _, (r1_min, r2_min) = self.pop_axis(rmin, axis=self.axis) _, (r1_max, r2_max) = self.pop_axis(rmax, axis=self.axis) diff --git a/tidy3d/components/structure.py b/tidy3d/components/structure.py index e7168180c5..e89d288b78 100644 --- a/tidy3d/components/structure.py +++ b/tidy3d/components/structure.py @@ -321,11 +321,15 @@ def _make_adjoint_monitors( else: interval_space = AUTOGRAD_MONITOR_INTERVAL_SPACE_POLY + field_components_for_adjoint = [f"E{dim}" for dim in "xyz"] + if self.medium.is_pec: + field_components_for_adjoint += [f"H{dim}" for dim in "xyz"] + mnt_fld = FieldMonitor( size=size, center=center, freqs=freqs, - fields=("Ex", "Ey", "Ez"), + fields=field_components_for_adjoint, name=self._get_monitor_name(index=index, data_type="fld"), interval_space=interval_space, colocate=False, diff --git a/tidy3d/plugins/smatrix/component_modelers/base.py b/tidy3d/plugins/smatrix/component_modelers/base.py index fa503d488a..85fc8ec626 100644 --- a/tidy3d/plugins/smatrix/component_modelers/base.py +++ b/tidy3d/plugins/smatrix/component_modelers/base.py @@ -37,7 +37,7 @@ DEFAULT_DATA_DIR = "." # whether to run gradient calculation for component modeler locally -LOCAL_GRADIENT = False +LOCAL_GRADIENT = True # False LumpedPortType = Union[LumpedPort, CoaxialLumpedPort] TerminalPortType = Union[LumpedPortType, WavePort] diff --git a/tidy3d/web/api/autograd/autograd.py b/tidy3d/web/api/autograd/autograd.py index 02a24d6527..b5b4fccb3f 100644 --- a/tidy3d/web/api/autograd/autograd.py +++ b/tidy3d/web/api/autograd/autograd.py @@ -991,8 +991,7 @@ def _compute_eps_array(medium, frequencies): def _slice_field_data( - field_data: dict, - freqs: np.ndarray, + field_data: dict, freqs: np.ndarray, component_indicator: typing.Optional[str] = None ) -> dict: """Slice field data dictionary along frequency dimension. @@ -1002,13 +1001,18 @@ def _slice_field_data( Dictionary of field components. freqs : np.ndarray Frequencies to select. + component_indicator: str + Component to filter field data selction on. Returns ------- dict Sliced field data dictionary. """ - return {k: v.sel(f=freqs) for k, v in field_data.items()} + if component_indicator: + return {k: v.sel(f=freqs) for k, v in field_data.items() if component_indicator in k} + else: + return {k: v.sel(f=freqs) for k, v in field_data.items()} def postprocess_adj( @@ -1029,33 +1033,39 @@ def postprocess_adj( sim_fields_vjp = {} for structure_index, structure_paths in sim_vjp_map.items(): # grab the forward and adjoint data - E_fwd = sim_data_fwd._get_adjoint_data(structure_index, data_type="fld") + fld_fwd = sim_data_fwd._get_adjoint_data(structure_index, data_type="fld") eps_fwd = sim_data_fwd._get_adjoint_data(structure_index, data_type="eps") - E_adj = sim_data_adj._get_adjoint_data(structure_index, data_type="fld") + fld_adj = sim_data_adj._get_adjoint_data(structure_index, data_type="fld") eps_adj = sim_data_adj._get_adjoint_data(structure_index, data_type="eps") # post normalize the adjoint fields if a single, broadband source - adj_flds_normed = {} - for key, val in E_adj.field_components.items(): - adj_flds_normed[key] = val * sim_data_adj.simulation.post_norm + fwd_flds_adj_normed = {} + for key, val in fld_adj.field_components.items(): + fwd_flds_adj_normed[key] = val * sim_data_adj.simulation.post_norm - E_adj = E_adj.updated_copy(**adj_flds_normed) + fld_adj = fld_adj.updated_copy(**fwd_flds_adj_normed) # maps of the E_fwd * E_adj and D_fwd * D_adj, each as as td.FieldData & 'Ex', 'Ey', 'Ez' der_maps = get_derivative_maps( - fld_fwd=E_fwd, eps_fwd=eps_fwd, fld_adj=E_adj, eps_adj=eps_adj + fld_fwd=fld_fwd, + eps_fwd=eps_fwd, + fld_adj=fld_adj, + eps_adj=eps_adj, ) E_der_map = der_maps["E"] D_der_map = der_maps["D"] + H_der_map = der_maps["H"] + + H_info_exists = H_der_map is not None - D_fwd = E_to_D(E_fwd, eps_fwd) - D_adj = E_to_D(E_adj, eps_fwd) + D_fwd = E_to_D(fld_fwd, eps_fwd) + D_adj = E_to_D(fld_adj, eps_fwd) # compute the derivatives for this structure structure = sim_data_fwd.simulation.structures[structure_index] # compute epsilon arrays for all frequencies - adjoint_frequencies = np.array(E_adj.monitor.freqs) + adjoint_frequencies = np.array(fld_adj.monitor.freqs) eps_in = _compute_eps_array(structure.medium, adjoint_frequencies) eps_out = _compute_eps_array(sim_data_orig.simulation.medium, adjoint_frequencies) @@ -1085,27 +1095,31 @@ def postprocess_adj( for f in adjoint_frequencies ] - # permittivity with infinite structure - structs_inf_struct = list(sim_orig.structures)[structure_index + 1 :] - sim_inf_structure = sim_orig.updated_copy( - structures=structs_inf_struct, - medium=structure.medium, - monitors=[], - sources=[], - grid_spec=sim_orig_grid_spec, - ) - - eps_inf_structure_data = [ - sim_inf_structure.epsilon(box=plane_eps, coord_key="centers", freq=f) - for f in adjoint_frequencies - ] - eps_no_structure = xr.concat(eps_no_structure_data, dim="f").assign_coords( f=adjoint_frequencies ) - eps_inf_structure = xr.concat(eps_inf_structure_data, dim="f").assign_coords( - f=adjoint_frequencies - ) + + if structure.medium.is_pec: + eps_inf_structure = None + else: + # permittivity with infinite structure + structs_inf_struct = list(sim_orig.structures)[structure_index + 1 :] + sim_inf_structure = sim_orig.updated_copy( + structures=structs_inf_struct, + medium=structure.medium, + monitors=[], + sources=[], + grid_spec=sim_orig_grid_spec, + ) + + eps_inf_structure_data = [ + sim_inf_structure.epsilon(box=plane_eps, coord_key="centers", freq=f) + for f in adjoint_frequencies + ] + + eps_inf_structure = xr.concat(eps_inf_structure_data, dim="f").assign_coords( + f=adjoint_frequencies + ) else: eps_no_structure = eps_inf_structure = None @@ -1134,12 +1148,31 @@ def postprocess_adj( # slice field data for current chunk E_der_map_chunk = _slice_field_data(E_der_map.field_components, select_adjoint_freqs) D_der_map_chunk = _slice_field_data(D_der_map.field_components, select_adjoint_freqs) - E_fwd_chunk = _slice_field_data(E_fwd.field_components, select_adjoint_freqs) - E_adj_chunk = _slice_field_data(E_adj.field_components, select_adjoint_freqs) + E_fwd_chunk = _slice_field_data( + fld_fwd.field_components, select_adjoint_freqs, component_indicator="E" + ) + E_adj_chunk = _slice_field_data( + fld_adj.field_components, select_adjoint_freqs, component_indicator="E" + ) D_fwd_chunk = _slice_field_data(D_fwd.field_components, select_adjoint_freqs) D_adj_chunk = _slice_field_data(D_adj.field_components, select_adjoint_freqs) eps_data_chunk = _slice_field_data(eps_fwd.field_components, select_adjoint_freqs) + H_der_map_chunk = None + H_fwd_chunk = None + H_adj_chunk = None + + if H_info_exists: + H_der_map_chunk = _slice_field_data( + H_der_map.field_components, select_adjoint_freqs + ) + H_fwd_chunk = _slice_field_data( + fld_fwd.field_components, select_adjoint_freqs, component_indicator="H" + ) + H_adj_chunk = _slice_field_data( + fld_adj.field_components, select_adjoint_freqs, component_indicator="H" + ) + # slice epsilon arrays eps_in_chunk = eps_in.sel(f=select_adjoint_freqs) eps_out_chunk = eps_out.sel(f=select_adjoint_freqs) @@ -1162,19 +1195,24 @@ def postprocess_adj( paths=structure_paths, E_der_map=E_der_map_chunk, D_der_map=D_der_map_chunk, + H_der_map=H_der_map_chunk, E_fwd=E_fwd_chunk, E_adj=E_adj_chunk, D_fwd=D_fwd_chunk, D_adj=D_adj_chunk, + H_fwd=H_fwd_chunk, + H_adj=H_adj_chunk, eps_data=eps_data_chunk, eps_in=eps_in_chunk, eps_out=eps_out_chunk, eps_background=eps_background_chunk, - frequencies=adjoint_frequencies[freq_slice], # only chunk frequencies + frequencies=select_adjoint_freqs, # only chunk frequencies eps_no_structure=eps_no_structure_chunk, eps_inf_structure=eps_inf_structure_chunk, bounds=struct_bounds, bounds_intersect=bounds_intersect, + simulation_bounds=sim_data_orig.simulation.bounds, + is_medium_pec=structure.medium.is_pec, ) # compute derivatives for chunk diff --git a/tidy3d/web/api/autograd/utils.py b/tidy3d/web/api/autograd/utils.py index 7189096531..5d86eed822 100644 --- a/tidy3d/web/api/autograd/utils.py +++ b/tidy3d/web/api/autograd/utils.py @@ -3,6 +3,7 @@ import typing +import numpy as np import pydantic as pd import tidy3d as td @@ -22,12 +23,23 @@ def get_derivative_maps( """Get electric and displacement field derivative maps.""" der_map_E = derivative_map_E(fld_fwd=fld_fwd, fld_adj=fld_adj) der_map_D = derivative_map_D(fld_fwd=fld_fwd, eps_fwd=eps_fwd, fld_adj=fld_adj, eps_adj=eps_adj) - return {"E": der_map_E, "D": der_map_D} + + make_H_der_map = np.all([f"H{dim}" in fld_fwd.field_components for dim in "xyz"]) + der_map_H = None + if make_H_der_map: + der_map_H = derivative_map_H(fld_fwd=fld_fwd, fld_adj=fld_adj) + + return {"E": der_map_E, "D": der_map_D, "H": der_map_H} def derivative_map_E(fld_fwd: td.FieldData, fld_adj: td.FieldData) -> td.FieldData: """Get td.FieldData where the Ex, Ey, Ez components store the gradients w.r.t. these.""" - return multiply_field_data(fld_fwd, fld_adj) + return multiply_field_data(fld_fwd, fld_adj, fld_key="E") + + +def derivative_map_H(fld_fwd: td.FieldData, fld_adj: td.FieldData) -> td.FieldData: + """Get td.FieldData where the Hx, Hy, Hz components store the gradients w.r.t. these.""" + return multiply_field_data(fld_fwd, fld_adj, fld_key="H") def derivative_map_D( @@ -39,22 +51,24 @@ def derivative_map_D( """Get td.FieldData where the Ex, Ey, Ez components store the gradients w.r.t. D fields.""" fwd_D = E_to_D(fld_data=fld_fwd, eps_data=eps_fwd) adj_D = E_to_D(fld_data=fld_adj, eps_data=eps_adj) - return multiply_field_data(fwd_D, adj_D) + + return multiply_field_data(fwd_D, adj_D, fld_key="E") def E_to_D(fld_data: td.FieldData, eps_data: td.PermittivityData) -> td.FieldData: """Convert electric field to displacement field.""" - return multiply_field_data(fld_data, eps_data) + + return multiply_field_data(fld_data, eps_data, fld_key="E") def multiply_field_data( - fld_1: td.FieldData, fld_2: typing.Union[td.FieldData, td.PermittivityData] + fld_1: td.FieldData, fld_2: typing.Union[td.FieldData, td.PermittivityData], fld_key: str ) -> td.FieldData: """Elementwise multiply two field data objects, writes data into ``fld_1`` copy.""" def get_field_key(dim: str, fld_data: typing.Union[td.FieldData, td.PermittivityData]) -> str: """Get the key corresponding to the scalar field along this dimension.""" - return f"E{dim}" if isinstance(fld_data, td.FieldData) else f"eps_{dim}{dim}" + return f"{fld_key}{dim}" if isinstance(fld_data, td.FieldData) else f"eps_{dim}{dim}" field_components = {} for dim in "xyz":