from typing import Iterable, List, Tuple import numpy as onp import meep as mp from . import ObjectiveQuantity # Meep field components used to compute adjoint sensitivities _ADJOINT_FIELD_COMPONENTS = [mp.Dx, mp.Dy, mp.Dz] _ADJOINT_FIELD_COMPONENTS_CYL = [mp.Dr, mp.Dp, mp.Dz] # The frequency axis in the array returned by `mp._get_gradient()` _GRADIENT_FREQ_AXIS = 1 # The returned axis order from get_dft_array in cylindrical coordinates _FREQ_AXIS = 0 _RHO_AXIS = 2 _PHI_AXIS = 3 _Z_AXIS = 1 # default finite difference step size when calculating Aᵤ FD_DEFAULT = 1e-3 class DesignRegion: def __init__( self, design_parameters: Iterable[onp.ndarray], volume: mp.Volume = None, size: mp.Vector3 = None, center: mp.Vector3 = mp.Vector3(), ): self.volume = volume or mp.Volume(center=center, size=size) self.size = self.volume.size self.center = self.volume.center self.design_parameters = design_parameters self.num_design_params = design_parameters.num_params def update_design_parameters(self, design_parameters) -> None: self.design_parameters.update_weights(design_parameters) def update_beta(self, beta: float) -> None: self.design_parameters.beta = beta def get_gradient( self, sim: mp.Simulation, fields_a: List[mp.DftFields], fields_f: List[mp.DftFields], frequencies: List[float], finite_difference_step: float, ) -> onp.ndarray: num_freqs = onp.array(frequencies).size """We have the option to linearly scale the gradients up front using the scalegrad parameter (leftover from MPB API). Not currently needed for any existing feature (but available for future use)""" scalegrad = 1 grad = onp.zeros((num_freqs, self.num_design_params)) # preallocate vol = sim._fit_volume_to_simulation(self.volume) # compute the gradient mp._get_gradient( grad, scalegrad, fields_a[0].swigobj, fields_a[1].swigobj, fields_a[2].swigobj, fields_f[0].swigobj, fields_f[1].swigobj, fields_f[2].swigobj, sim.gv, vol.swigobj, onp.array(frequencies), sim.geps, finite_difference_step, ) return onp.squeeze(grad).T def _check_if_cylindrical(sim: mp.Simulation) -> bool: return sim.is_cylindrical or (sim.dimensions == mp.CYLINDRICAL) def _compute_components(sim: mp.Simulation) -> List[int]: return ( _ADJOINT_FIELD_COMPONENTS_CYL if _check_if_cylindrical(sim) else _ADJOINT_FIELD_COMPONENTS ) def _make_at_least_nd(x: onp.ndarray, dims: int = 3) -> onp.ndarray: """Makes an array have at least the specified number of dimensions.""" return onp.reshape(x, x.shape + onp.maximum(dims - x.ndim, 0) * (1,)) def calculate_vjps( simulation: mp.Simulation, design_regions: List[DesignRegion], frequencies: List[float], fwd_fields: List[List[mp.DftFields]], adj_fields: List[List[mp.DftFields]], design_variable_shapes: List[Tuple[int, ...]], sum_freq_partials: bool = True, finite_difference_step: float = FD_DEFAULT, ) -> List[onp.ndarray]: """Calculates the VJP for a given set of forward and adjoint fields.""" vjps = [ design_region.get_gradient( simulation, adj_fields[i], fwd_fields[i], frequencies, finite_difference_step, ) for i, design_region in enumerate(design_regions) ] if sum_freq_partials: vjps = [ onp.sum(vjp, axis=_GRADIENT_FREQ_AXIS).reshape(shape) for vjp, shape in zip(vjps, design_variable_shapes) ] else: vjps = [ vjp.reshape(shape + (-1,)) for vjp, shape in zip(vjps, design_variable_shapes) ] return vjps def register_monitors( monitors: List[ObjectiveQuantity], frequencies: List[float], ) -> None: """Registers a list of monitors.""" for monitor in monitors: monitor.register_monitors(frequencies) def install_design_region_monitors( simulation: mp.Simulation, design_regions: List[DesignRegion], frequencies: List[float], decimation_factor: int = 0, ) -> List[List[mp.DftFields]]: """Installs DFT field monitors at the design regions of the simulation.""" return [ [ simulation.add_dft_fields( [comp], frequencies, where=design_region.volume, yee_grid=True, decimation_factor=decimation_factor, persist=True, ) for comp in _compute_components(simulation) ] for design_region in design_regions ] def gather_monitor_values(monitors: List[ObjectiveQuantity]) -> onp.ndarray: """Gathers the mode monitor overlap values as a rank 2 ndarray. Args: monitors: the mode monitors. Returns: a rank-2 ndarray, where the dimensions are (monitor, frequency), of dtype complex128. Note that these values refer to the mode as oriented (i.e. they are unidirectional). """ monitor_values = [monitor() for monitor in monitors] monitor_values = onp.array(monitor_values) assert monitor_values.ndim in [1, 2] monitor_values = _make_at_least_nd(monitor_values, 2) return monitor_values def validate_and_update_design( design_regions: List[DesignRegion], design_variables: Iterable[onp.ndarray] ) -> None: """Validate the design regions and variables. In particular the design variable should be 1,2,3-D and the design region shape should match the design variable shape after dimension expansion. The arguments are modified in place. Args: design_regions: List of mpa.DesignRegion, design_variables: Iterable with numpy arrays representing design variables. Raises: ValueError if the validation of dimensions fails. """ for i, (design_region, design_variable) in enumerate( zip(design_regions, design_variables) ): if design_variable.ndim not in [1, 2, 3]: raise ValueError( f"Design variables should be 1D, 2D, or 3D, but the design variable at index {i} had a shape of {design_variable.shape}." ) design_region_shape = tuple( int(x) for x in design_region.design_parameters.grid_size ) design_variable_padded_shape = design_variable.shape + (1,) * ( 3 - design_variable.ndim ) if design_variable_padded_shape != design_region_shape: raise ValueError( f"The design variable at index {i} with a shape of {design_variable.shape} is incompatible with the associated design region, which has a shape of {design_region_shape}." ) design_variable = onp.asarray(design_variable, dtype=onp.float64) # Update the design variable in Meep design_region.update_design_parameters(design_variable.flatten()) def create_adjoint_sources( monitors: Iterable[ObjectiveQuantity], monitor_values_grad: onp.ndarray ) -> List[mp.Source]: monitor_values_grad = onp.asarray( monitor_values_grad, dtype=onp.complex64 if mp.is_single_precision() else onp.complex128, ) if not onp.any(monitor_values_grad): raise RuntimeError( "The gradient of all monitor values is zero, which " "means that no adjoint sources can be placed to set " "up an adjoint simulation in Meep. Possible causes " "could be:\n\n" " * the forward simulation was not run for long enough " "to allow the input pulse(s) to reach the monitors" " * the monitor values are disconnected from the " "objective function output." ) adjoint_sources = [] for monitor_idx, monitor in enumerate(monitors): # `dj` for each monitor will have a shape of (num frequencies,) dj = onp.asarray( monitor_values_grad[monitor_idx], dtype=onp.complex64 if mp.is_single_precision() else onp.complex128, ) if onp.any(dj): adjoint_sources += monitor.place_adjoint_source(dj) assert adjoint_sources return adjoint_sources