diff --git a/tidy3d/web/api/autograd/autograd.py b/tidy3d/web/api/autograd/autograd.py index 7958dc57c9..726c12a57e 100644 --- a/tidy3d/web/api/autograd/autograd.py +++ b/tidy3d/web/api/autograd/autograd.py @@ -5,7 +5,6 @@ from os.path import dirname from pathlib import Path -import numpy as np from autograd.builtins import dict as dict_ag from autograd.extend import defvjp, primitive @@ -47,16 +46,6 @@ ) from .forward import postprocess_fwd as _postprocess_fwd_impl from .forward import setup_fwd as _setup_fwd_impl -from .io_utils import ( - get_vjp_traced_fields as _get_vjp_traced_fields_impl, -) -from .io_utils import ( - upload_sim_fields_keys as _upload_sim_fields_keys_impl, -) - -# if True, will plot the adjoint fields on the plane provided. used for debugging only -_INSPECT_ADJOINT_FIELDS = False -_INSPECT_ADJOINT_PLANE = td.Box(center=(0, 0, 0), size=(td.inf, td.inf, 0)) def is_valid_for_autograd(simulation: td.Simulation) -> bool: @@ -623,21 +612,6 @@ def postprocess_fwd( ) -def upload_sim_fields_keys(sim_fields_keys: list[tuple], task_id: str, verbose: bool = False): - """Upload traced simulation field keys for adjoint runs (delegated).""" - return _upload_sim_fields_keys_impl( - sim_fields_keys=sim_fields_keys, task_id=task_id, verbose=verbose - ) - - -""" VJP maker for ADJ pass.""" - - -def get_vjp_traced_fields(task_id_adj: str, verbose: bool) -> AutogradFieldMap: - """Fetch VJP traced fields for a completed adjoint job (delegated).""" - return _get_vjp_traced_fields_impl(task_id_adj=task_id_adj, verbose=verbose) - - def _run_bwd( data_fields_original: AutogradFieldMap, sim_fields_original: AutogradFieldMap, @@ -919,22 +893,6 @@ def setup_adj( ) -def _compute_eps_array(medium, frequencies): - """Deprecated shim, kept for backward compatibility; use ops_backward._compute_eps_array.""" - from .backward import _compute_eps_array as __impl - - return __impl(medium, frequencies) - - -def _slice_field_data( - field_data: dict, freqs: np.ndarray, component_indicator: typing.Optional[str] = None -) -> dict: - """Deprecated shim, kept for backward compatibility; use ops_backward._slice_field_data.""" - from .backward import _slice_field_data as __impl - - return __impl(field_data, freqs, component_indicator) - - def postprocess_adj( sim_data_adj: td.SimulationData, sim_data_orig: td.SimulationData, diff --git a/tidy3d/web/api/autograd/backward.py b/tidy3d/web/api/autograd/backward.py index 799f4dd95f..e4f6f471d0 100644 --- a/tidy3d/web/api/autograd/backward.py +++ b/tidy3d/web/api/autograd/backward.py @@ -10,6 +10,7 @@ from tidy3d.components.autograd.constants import ADJOINT_FREQ_CHUNK_SIZE from tidy3d.components.autograd.derivative_utils import DerivativeInfo from tidy3d.exceptions import AdjointError +from tidy3d.packaging import disable_local_subpixel from .utils import E_to_D, get_derivative_maps @@ -98,6 +99,7 @@ def _slice_field_data( return {k: v.sel(f=freqs) for k, v in field_data.items()} +@disable_local_subpixel def postprocess_adj( sim_data_adj: td.SimulationData, sim_data_orig: td.SimulationData, diff --git a/tidy3d/web/api/autograd/constants_local.py b/tidy3d/web/api/autograd/constants_local.py deleted file mode 100644 index 3429876352..0000000000 --- a/tidy3d/web/api/autograd/constants_local.py +++ /dev/null @@ -1,17 +0,0 @@ -from __future__ import annotations - -# keys for data into auxiliary dictionary (re-exported in autograd.py for tests) -AUX_KEY_SIM_DATA_ORIGINAL = "sim_data" -AUX_KEY_SIM_DATA_FWD = "sim_data_fwd_adjoint" -AUX_KEY_FWD_TASK_ID = "task_id_fwd" -AUX_KEY_SIM_ORIGINAL = "sim_original" - -# server-side auxiliary files to upload/download -SIM_VJP_FILE = "output/autograd_sim_vjp.hdf5" -SIM_FIELDS_KEYS_FILE = "autograd_sim_fields_keys.hdf5" - -# default behaviors -LOCAL_GRADIENT = False - -# directory to store adjoint data for local gradient calculation relative to run path -LOCAL_ADJOINT_DIR = "adjoint_data" diff --git a/tidy3d/web/api/autograd/ops_backward.py b/tidy3d/web/api/autograd/ops_backward.py deleted file mode 100644 index 799f4dd95f..0000000000 --- a/tidy3d/web/api/autograd/ops_backward.py +++ /dev/null @@ -1,320 +0,0 @@ -from __future__ import annotations - -from collections import defaultdict - -import numpy as np -import xarray as xr - -import tidy3d as td -from tidy3d.components.autograd import AutogradFieldMap, get_static -from tidy3d.components.autograd.constants import ADJOINT_FREQ_CHUNK_SIZE -from tidy3d.components.autograd.derivative_utils import DerivativeInfo -from tidy3d.exceptions import AdjointError - -from .utils import E_to_D, get_derivative_maps - - -def setup_adj( - data_fields_vjp: AutogradFieldMap, - sim_data_orig: td.SimulationData, - sim_fields_keys: list[tuple], - max_num_adjoint_per_fwd: int, -) -> list[td.Simulation]: - """Construct an adjoint simulation from a set of data_fields for the VJP.""" - - td.log.info("Running custom vjp (adjoint) pipeline.") - - # filter out any data_fields_vjp with all 0's - data_fields_vjp = { - k: get_static(v) for k, v in data_fields_vjp.items() if not np.allclose(v, 0) - } - - for k, v in data_fields_vjp.items(): - if np.any(np.isnan(v)): - raise AdjointError( - f"NaN values detected for data field {k} in the adjoint pipeline. This may be " - f"due to NaN values in the simulation data or the computed value of your " - f"objective function." - ) - - # if all entries are zero, there is no adjoint sim to run - if not data_fields_vjp: - return [] - - # start with the full simulation data structure and either zero out the fields - # that have no tracer data for them or insert the tracer data - full_sim_data_dict = sim_data_orig._strip_traced_fields( - include_untraced_data_arrays=True, starting_path=("data",) - ) - for path in full_sim_data_dict.keys(): - if path in data_fields_vjp: - full_sim_data_dict[path] = data_fields_vjp[path] - else: - full_sim_data_dict[path] *= 0 - - # insert the raw VJP data into the .data of the original SimulationData - sim_data_vjp = sim_data_orig._insert_traced_fields(field_mapping=full_sim_data_dict) - - # make adjoint simulation from that SimulationData - data_vjp_paths = set(data_fields_vjp.keys()) - - num_monitors = len(sim_data_orig.simulation.monitors) - adjoint_monitors = sim_data_orig.simulation._with_adjoint_monitors(sim_fields_keys).monitors[ - num_monitors: - ] - - sims_adj = sim_data_vjp._make_adjoint_sims( - data_vjp_paths=data_vjp_paths, - adjoint_monitors=adjoint_monitors, - ) - - if len(sims_adj) > max_num_adjoint_per_fwd: - raise AdjointError( - f"Number of adjoint simulations ({len(sims_adj)}) exceeds the maximum allowed " - f"({max_num_adjoint_per_fwd}) per forward simulation. This typically means that " - "there are many frequencies and monitors in the simulation that are being differentiated " - "w.r.t. in the objective function. To proceed, please double-check the simulation " - "setup, increase the 'max_num_adjoint_per_fwd' parameter in the run function, and re-run." - ) - - return sims_adj - - -def _compute_eps_array(medium, frequencies): - """Compute permittivity array for all frequencies.""" - eps_data = [np.mean(medium.eps_model(f)) for f in frequencies] - return td.components.data.data_array.DataArray( - data=np.array(eps_data), dims=("f",), coords={"f": frequencies} - ) - - -def _slice_field_data( - field_data: dict, freqs: np.ndarray, component_indicator: str | None = None -) -> dict: - """Slice field data dictionary along frequency dimension.""" - 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( - sim_data_adj: td.SimulationData, - sim_data_orig: td.SimulationData, - sim_data_fwd: td.SimulationData, - sim_fields_keys: list[tuple], -) -> AutogradFieldMap: - """Postprocess some data from the adjoint simulation into the VJP for the original sim flds.""" - - # map of index into 'structures' to the list of paths we need vjps for - sim_vjp_map = defaultdict(list) - for _, structure_index, *structure_path in sim_fields_keys: - structure_path = tuple(structure_path) - sim_vjp_map[structure_index].append(structure_path) - - # store the derivative values given the forward and adjoint data - sim_fields_vjp = {} - for structure_index, structure_paths in sim_vjp_map.items(): - # grab the forward and adjoint data - 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") - 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 - 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 - - 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=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(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(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) - - # handle background medium if present - if structure.background_medium: - eps_background = _compute_eps_array(structure.background_medium, adjoint_frequencies) - else: - eps_background = None - - # auto permittivity detection for non-box geometries - if not isinstance(structure.geometry, td.Box): - sim_orig = sim_data_orig.simulation - plane_eps = eps_fwd.monitor.geometry - - sim_orig_grid_spec = td.components.grid.grid_spec.GridSpec.from_grid(sim_orig.grid) - - # permittivity without this structure - structs_no_struct = list(sim_orig.structures) - structs_no_struct.pop(structure_index) - sim_no_structure = sim_orig.updated_copy( - structures=structs_no_struct, monitors=[], sources=[], grid_spec=sim_orig_grid_spec - ) - - eps_no_structure_data = [ - sim_no_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 - ) - - 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 - - # compute bounds intersection - struct_bounds = rmin_struct, rmax_struct = structure.geometry.bounds - rmin_sim, rmax_sim = sim_data_orig.simulation.bounds - rmin_intersect = tuple([max(a, b) for a, b in zip(rmin_sim, rmin_struct)]) - rmax_intersect = tuple([min(a, b) for a, b in zip(rmax_sim, rmax_struct)]) - bounds_intersect = (rmin_intersect, rmax_intersect) - - # get chunk size - if None, process all frequencies as one chunk - freq_chunk_size = ADJOINT_FREQ_CHUNK_SIZE - n_freqs = len(adjoint_frequencies) - if freq_chunk_size is None: - freq_chunk_size = n_freqs - - # process in chunks - vjp_value_map = {} - - for chunk_start in range(0, n_freqs, freq_chunk_size): - chunk_end = min(chunk_start + freq_chunk_size, n_freqs) - freq_slice = slice(chunk_start, chunk_end) - - select_adjoint_freqs = adjoint_frequencies[freq_slice] - - # 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( - 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) - eps_background_chunk = ( - eps_background.sel(f=select_adjoint_freqs) if eps_background is not None else None - ) - eps_no_structure_chunk = ( - eps_no_structure.sel(f=select_adjoint_freqs) - if eps_no_structure is not None - else None - ) - eps_inf_structure_chunk = ( - eps_inf_structure.sel(f=select_adjoint_freqs) - if eps_inf_structure is not None - else None - ) - - # create derivative info with sliced data - derivative_info = DerivativeInfo( - 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=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 - vjp_chunk = structure._compute_derivatives(derivative_info) - - # accumulate results - for path, value in vjp_chunk.items(): - if path in vjp_value_map: - val = vjp_value_map[path] - if isinstance(val, (list, tuple)) and isinstance(value, (list, tuple)): - vjp_value_map[path] = type(val)(x + y for x, y in zip(val, value)) - else: - vjp_value_map[path] += value - else: - vjp_value_map[path] = value - - # store vjps in output map - for structure_path, vjp_value in vjp_value_map.items(): - sim_path = ("structures", structure_index, *list(structure_path)) - sim_fields_vjp[sim_path] = vjp_value - - return sim_fields_vjp diff --git a/tidy3d/web/api/autograd/ops_forward.py b/tidy3d/web/api/autograd/ops_forward.py deleted file mode 100644 index 7804f6a8a2..0000000000 --- a/tidy3d/web/api/autograd/ops_forward.py +++ /dev/null @@ -1,41 +0,0 @@ -from __future__ import annotations - -import tidy3d as td -from tidy3d.components.autograd import AutogradFieldMap - - -def setup_fwd( - sim_fields: AutogradFieldMap, - sim_original: td.Simulation, - local_gradient: bool = False, -) -> td.Simulation: - """Return a forward simulation with adjoint monitors attached.""" - - # Always try to build the variant that includes adjoint monitors so that - # errors in monitor placement are caught early. - sim_with_adj_mon = sim_original._with_adjoint_monitors(sim_fields) - return sim_with_adj_mon if local_gradient else sim_original - - -def postprocess_fwd( - sim_data_combined: td.SimulationData, - sim_original: td.Simulation, - aux_data: dict, -) -> AutogradFieldMap: - """Postprocess the combined simulation data into an Autograd field map.""" - - num_mnts_original = len(sim_original.monitors) - sim_data_original, sim_data_fwd = sim_data_combined._split_original_fwd( - num_mnts_original=num_mnts_original - ) - - aux_data["sim_data"] = sim_data_original - aux_data["sim_data_fwd_adjoint"] = sim_data_fwd - - # strip out the tracer AutogradFieldMap for the .data from the original sim - data_traced = sim_data_original._strip_traced_fields( - include_untraced_data_arrays=True, starting_path=("data",) - ) - - # return the AutogradFieldMap that autograd registers as the "output" of the primitive - return data_traced