""" A collection of objects and helper methods for defining objective functions used in topology optimization. """ import abc from collections import namedtuple from typing import Callable, List, Optional import numpy as np from meep.simulation import py_v3_to_vec, FluxData, NearToFarData import meep as mp from .filter_source import FilteredSource Grid = namedtuple("Grid", ["x", "y", "z", "w"]) class ObjectiveQuantity(abc.ABC): """A differentiable objective quantity. Attributes: sim: the Meep simulation object used to register the objective quantity. frequencies: the frequencies at which the objective quantity is evaluated. num_freq: the number of frequencies at which the objective quantity is evaluated. """ def __init__(self, sim): self.sim = sim self._eval = None self._frequencies = None @property def frequencies(self): return self._frequencies @property def num_freq(self): return len(self.frequencies) @abc.abstractmethod def __call__(self): """Evaluates the objective quantity.""" @abc.abstractmethod def register_monitors(self, frequencies): """Registers monitors in the forward simulation.""" @abc.abstractmethod def place_adjoint_source(self, dJ): """Places appropriate sources for the adjoint simulation.""" def get_evaluation(self): """Evaluates the objective quantity.""" if self._eval is not None: return self._eval else: raise RuntimeError( "You must first run a forward simulation before requesting the" "evaluation of an objective quantity." ) def _adj_src_scale(self, include_resolution=True): """Calculates the scale for the adjoint sources.""" T = self.sim.meep_time() dt = self.sim.fields.dt src = self._create_time_profile() if include_resolution: num_dims = self.sim._infer_dimensions(self.sim.k_point) dV = 1 / self.sim.resolution**num_dims else: dV = 1 iomega = (1.0 - np.exp(-1j * (2 * np.pi * self._frequencies) * dt)) * ( 1.0 / dt ) # scaled frequency factor with discrete time derivative fix # an ugly way to calcuate the scaled dtft of the forward source y = np.array( [src.swigobj.current(t, dt) for t in np.arange(0, T, dt)] ) # time domain signal fwd_dtft = ( np.matmul( np.exp( 1j * 2 * np.pi * self._frequencies[:, np.newaxis] * np.arange(y.size) * dt ), y, ) * dt / np.sqrt(2 * np.pi) ) # Interestingly, the real parts of the DTFT and Fourier transform match, # but the imaginary parts are very different... # fwd_dtft = src.fourier_transform(src.frequency) # # Note: for some reason, there seems to be an additional phase factor at # the center frequency that needs to be applied to *all* frequencies... src_center_dtft = ( np.matmul( np.exp( 1j * 2 * np.pi * np.array([src.frequency])[:, np.newaxis] * np.arange(y.size) * dt ), y, ) * dt / np.sqrt(2 * np.pi) ) adj_src_phase = np.exp(1j * np.angle(src_center_dtft)) * self.fwidth_scale if self._frequencies.size == 1: # Single-frequency simulations. Requires a time profile. scale = dV * iomega / fwd_dtft / adj_src_phase # final scale factor else: # Multi-frequency simulations. scale = dV * iomega / adj_src_phase # Cmpensate for the fact that real fields take the real part of the # current, which halves the Fourier amplitude at the positive frequency # (i.e. Re[J] = (J + J*)/2). if self.sim.using_real_fields(): scale *= 2 return scale def _create_time_profile(self, fwidth_frac=0.1, adj_cutoff=5): """Creates a time-domain waveform for normalizing the adjoint source(s). For single-frequency objective functions, we should generate a Gaussian pulse with a reasonable bandwidth centered at the given frequency. TODO: The user may specify a scalar-valued objective function across multiple frequencies (e.g. MSE) in which case we should check that all the frequencies fit in the specified bandwidth. """ self.fwidth_scale = np.exp(-2j * np.pi * adj_cutoff / fwidth_frac) return mp.GaussianSource( np.mean(self._frequencies), fwidth=fwidth_frac * np.mean(self._frequencies), cutoff=adj_cutoff, ) class EigenmodeCoefficient(ObjectiveQuantity): """A differentiable frequency-dependent eigenmode coefficient.""" def __init__( self, sim: mp.Simulation, volume: mp.Volume, mode: int, forward: Optional[bool] = True, kpoint_func: Optional[Callable] = None, kpoint_func_overlap_idx: Optional[int] = 0, decimation_factor: Optional[int] = 0, subtracted_dft_fields: Optional[FluxData] = None, **kwargs ): """Initialize an instance of a differentiable frequency-dependent eigenmode coefficient. Args: sim: the Meep simulation object of the problem. volume: the volume over which the eigenmode coefficient is calculated. mode: the eigenmode number. forward: whether the forward or backward mode coefficient is returned as the result of the evaluation. Default is True. kpoint_func: an optional k-point function to use when evaluating the eigenmode coefficient. When specified, this overrides the effect of `forward`. kpoint_func_overlap_idx: the index of the mode coefficient to return when specifying `kpoint_func`. When specified, this overrides the effect of `forward` and should have a value of either 0 or 1. decimation_factor: An integer used to specify the number of timesteps between updates of the DFT fields. The default is 0, at which the value is automatically determined from the Nyquist rate of the bandwidth-limited sources and the DFT monitor. It can be turned off by setting it to 1. subtracted_dft_fields: the DFT fields obtained using `get_flux_data` from a previous normalization run. This is subtracted from the DFT fields of this mode monitor in order to improve the accuracy of the reflectance measurement (i.e., the $S_{11}$ scattering parameter). Default is None. eigenmode_kwargs: additional keyword arguments for EigenModeSource. """ super().__init__(sim) if kpoint_func_overlap_idx not in [0, 1]: raise ValueError( "`kpoint_func_overlap_idx` should be either 0 or 1, but got %d" % (kpoint_func_overlap_idx,) ) self.volume = volume self.mode = mode self.forward = forward self.kpoint_func = kpoint_func self.kpoint_func_overlap_idx = kpoint_func_overlap_idx self.eigenmode_kwargs = kwargs self._monitor = None self._cscale = None self.decimation_factor = decimation_factor self.subtracted_dft_fields = subtracted_dft_fields def register_monitors(self, frequencies): self._frequencies = np.asarray(frequencies) self._monitor = self.sim.add_mode_monitor( frequencies, mp.ModeRegion(center=self.volume.center, size=self.volume.size), yee_grid=True, decimation_factor=self.decimation_factor, ) if self.subtracted_dft_fields is not None: self.sim.load_minus_flux_data( self._monitor, self.subtracted_dft_fields, ) return self._monitor def place_adjoint_source(self, dJ): dJ = np.atleast_1d(dJ) if dJ.ndim == 2: dJ = np.sum(dJ, axis=1) time_src = self._create_time_profile() da_dE = 0.5 * self._cscale scale = self._adj_src_scale() if self.kpoint_func: eig_kpoint = -1 * self.kpoint_func(time_src.frequency, self.mode) else: center_frequency = 0.5 * ( np.min(self.frequencies) + np.max(self.frequencies) ) direction = mp.Vector3( *(np.eye(3)[self._monitor.normal_direction] * np.abs(center_frequency)) ) eig_kpoint = -1 * direction if self.forward else direction if self._frequencies.size == 1: amp = da_dE * dJ * scale src = time_src else: scale = da_dE * dJ * scale src = FilteredSource( time_src.frequency, self._frequencies, scale, self.sim.fields.dt, ) amp = 1 source = mp.EigenModeSource( src, eig_band=self.mode, direction=mp.NO_DIRECTION, eig_kpoint=eig_kpoint, amplitude=amp, eig_match_freq=True, size=self.volume.size, center=self.volume.center, **self.eigenmode_kwargs, ) return [source] def __call__(self): """The values of the eigenmode coefficient at each frequency. Returns: 1D array of eigenmode coefficients for each frequency in self.frequencies. """ if self.kpoint_func: kpoint_func = self.kpoint_func overlap_idx = self.kpoint_func_overlap_idx else: center_frequency = 0.5 * ( np.min(self.frequencies) + np.max(self.frequencies) ) kpoint = mp.Vector3( *(np.eye(3)[self._monitor.normal_direction] * np.abs(center_frequency)) ) kpoint_func = lambda *not_used: kpoint if self.forward else -1 * kpoint overlap_idx = 0 ob = self.sim.get_eigenmode_coefficients( self._monitor, [self.mode], direction=mp.NO_DIRECTION, kpoint_func=kpoint_func, **self.eigenmode_kwargs, ) overlaps = ob.alpha.squeeze(axis=0) assert overlaps.ndim == 2 self._eval = overlaps[:, overlap_idx] self._cscale = ob.cscale return self._eval class FourierFields(ObjectiveQuantity): """A differentiable frequency-dependent Fourier fields (dft_fields)""" def __init__( self, sim: mp.Simulation, volume: mp.Volume, component: int, yee_grid: Optional[bool] = False, decimation_factor: Optional[int] = 0, subtracted_dft_fields: Optional[FluxData] = None, ): """Initialize an instance of differentiable Fourier fields instance. Args: sim: the Meep simulation object of the problem. volume: the volume over which the eigenmode coefficient is calculated. Due to an unresolved bug, the size must not be zero in at least one direction. component: field component (e.g. mp.Ex, mp.Hz, etc.) of the Fourier fields yee_grid: if True, the Fourier fields are evaluated at the corresponding Yee grid points; otherwise, they are interpolated fields at the center of each voxel. Default is False decimation_factor: An integer used to specify the number of timesteps between updates of the DFT fields. The default is 0, at which the value is automatically determined from the Nyquist rate of the bandwidth-limited sources and the DFT monitor. It can be turned off by setting it to 1. subtracted_dft_fields: the DFT fields obtained using `get_flux_data` from a previous normalization run. This is subtracted from the DFT fields of this mode monitor in order to improve the accuracy of the reflectance measurement (i.e., the $S_{11}$ scattering parameter). Default is None. """ super().__init__(sim) self.volume = sim._fit_volume_to_simulation(volume) self.component = component self.yee_grid = yee_grid self.decimation_factor = decimation_factor self.subtracted_dft_fields = subtracted_dft_fields def register_monitors(self, frequencies): self._frequencies = np.asarray(frequencies) self._monitor = self.sim.add_dft_fields( [self.component], self._frequencies, where=self.volume, yee_grid=self.yee_grid, decimation_factor=self.decimation_factor, ) if self.subtracted_dft_fields is not None: self.sim.load_minus_flux_data( self._monitor, self.subtracted_dft_fields, ) return self._monitor def place_adjoint_source(self, dJ): time_src = self._create_time_profile() sources = [] mon_size = self.sim.fields.dft_monitor_size( self._monitor.swigobj, self.volume.swigobj, self.component ) dJ = dJ.astype(np.complex128) if ( np.prod(mon_size) * self.num_freq != dJ.size and np.prod(mon_size) * self.num_freq**2 != dJ.size ): raise ValueError("The format of J is incorrect!") # The objective function J is a vector. Each component corresponds to a frequency. if np.prod(mon_size) * self.num_freq**2 == dJ.size and self.num_freq > 1: dJ = np.sum(dJ, axis=1) """The adjoint solver requires the objective function to be scalar valued with regard to objective arguments and position, but the function may be vector valued with regard to frequency. In this case, the Jacobian will be of the form [F,F,...] where F is the number of frequencies. Because of linearity, we can sum across the second frequency dimension to calculate a frequency scale factor for each point (rather than a scale vector). """ all_fouriersrcdata = self._monitor.swigobj.fourier_sourcedata( self.volume.swigobj, self.component, self.sim.fields, dJ ) for fourier_data in all_fouriersrcdata: amp_arr = np.array(fourier_data.amp_arr).reshape(-1, self.num_freq) scale = amp_arr * self._adj_src_scale(include_resolution=False) if self.num_freq == 1: sources += [ mp.IndexedSource( time_src, fourier_data, scale[:, 0], not self.yee_grid ) ] else: src = FilteredSource( time_src.frequency, self._frequencies, scale, self.sim.fields.dt ) (num_basis, num_pts) = src.nodes.shape for basis_i in range(num_basis): sources += [ mp.IndexedSource( src.time_src_bf[basis_i], fourier_data, src.nodes[basis_i], not self.yee_grid, ) ] return sources def __call__(self): """The values of Fourier Fields at each frequency Returns: array of Fourier Fields with dimension k+1 where k is the dimension of self.volume The first axis corresponds to the index of frequency, and the rest k axis are for the spatial indices of points in the monitor """ self._eval = np.array( [ self.sim.get_dft_array(self._monitor, self.component, i) for i in range(self.num_freq) ] ) return self._eval class Near2FarFields(ObjectiveQuantity): """A differentiable near2far field transformation""" def __init__( self, sim: mp.Simulation, Near2FarRegions: List[mp.Near2FarRegion], far_pts: List[mp.Vector3], nperiods: Optional[int] = 1, decimation_factor: Optional[int] = 0, norm_near_fields: Optional[NearToFarData] = None, ): """Initialize an instance of differentiable Fourier fields instance. Args: sim: the Meep simulation object of the problem. Near2FarRegions: List of mp.Near2FarRegion over which the near fields are collected far_pts: list of far points at which fields are computed nperiods: If nperiods > 1, sum of 2*nperiods+1 Bloch-periodic copies of near fields is computed to approximate the lattice sum from Bloch periodic boundary condition. Default is 1 (no sum). decimation_factor: An integer used to specify the number of timesteps between updates of the DFT fields. The default is 0, at which the value is automatically determined from the Nyquist rate of the bandwidth-limited sources and the DFT monitor. It can be turned off by setting it to 1. norm_near_fields: the DFT fields obtained using `get_near2far_data` from a previous normalization run. This is subtracted from the DFT fields of this near2far monitor in order to improve the accuracy of the reflectance measurement (i.e., the $S_{11}$ scattering parameter). Default is None. """ super().__init__(sim) self.Near2FarRegions = Near2FarRegions self.far_pts = far_pts # list of far pts self._nfar_pts = len(far_pts) self.decimation_factor = decimation_factor self.norm_near_fields = norm_near_fields self.nperiods = nperiods def register_monitors(self, frequencies): self._frequencies = np.asarray(frequencies) self._monitor = self.sim.add_near2far( self._frequencies, *self.Near2FarRegions, nperiods=self.nperiods, decimation_factor=self.decimation_factor, ) if self.norm_near_fields is not None: self.sim.load_minus_near2far_data( self._monitor, self.norm_near_fields, ) return self._monitor def place_adjoint_source(self, dJ): time_src = self._create_time_profile() sources = [] if dJ.ndim == 4: dJ = np.sum(dJ, axis=0) farpt_list = np.array([list(pi) for pi in self.far_pts]).flatten() far_pt0 = self.far_pts[0] far_pt_vec = py_v3_to_vec( self.sim.dimensions, far_pt0, self.sim.is_cylindrical, ) all_nearsrcdata = self._monitor.swigobj.near_sourcedata( far_pt_vec, farpt_list, self._nfar_pts, dJ ) for near_data in all_nearsrcdata: cur_comp = near_data.near_fd_comp amp_arr = np.array(near_data.amp_arr).reshape(-1, self.num_freq) scale = amp_arr * self._adj_src_scale(include_resolution=False) if self.num_freq == 1: sources += [mp.IndexedSource(time_src, near_data, scale[:, 0])] else: src = FilteredSource( time_src.frequency, self._frequencies, scale, self.sim.fields.dt, ) (num_basis, num_pts) = src.nodes.shape for basis_i in range(num_basis): sources += [ mp.IndexedSource( src.time_src_bf[basis_i], near_data, src.nodes[basis_i], ) ] return sources def __call__(self): """The values of far fields at each points at each frequency Returns: 3D array of far fields. The first axis is the index of far field points in self.far_pts; the second axis is the index of frequency; and the third is the index of component in [mp.Ex(mp.Er), mp.Ey(mp.Ep), mp.Ez, mp.Hx(mp.Hr), mp.Hy(mp.Hp), mp.Hz] """ self._eval = np.array( [self.sim.get_farfield(self._monitor, far_pt) for far_pt in self.far_pts] ).reshape((self._nfar_pts, self.num_freq, 6)) return self._eval class LDOS(ObjectiveQuantity): """A differentiable LDOS""" def __init__(self, sim: mp.Simulation, **kwargs): """Initialize a differentiable LDOS instance Args: sim: the Meep simulation object of the problem. """ super().__init__(sim) self.srckwarg = kwargs def register_monitors(self, frequencies): self._frequencies = np.asarray(frequencies) self._forward_src = self.sim.sources return def place_adjoint_source(self, dJ): time_src = self._create_time_profile() if dJ.ndim == 2: dJ = np.sum(dJ, axis=1) dJ = dJ.flatten() sources = [] forward_f_scale = np.array( [self._ldos_scale / self._ldos_Jdata[k] for k in range(self.num_freq)] ) if self._frequencies.size == 1: amp = (dJ * self._adj_src_scale(False) * forward_f_scale)[0] src = time_src else: scale = dJ * self._adj_src_scale(False) * forward_f_scale src = FilteredSource( time_src.frequency, self._frequencies, scale, self.sim.fields.dt, ) amp = 1 for forward_src_i in self._forward_src: if isinstance(forward_src_i, mp.EigenModeSource): src_i = mp.EigenModeSource( src, component=forward_src_i.component, eig_kpoint=forward_src_i.eig_kpoint, amplitude=amp, eig_band=forward_src_i.eig_band, size=forward_src_i.size, center=forward_src_i.center, **self.srckwarg, ) else: src_i = mp.Source( src, component=forward_src_i.component, amplitude=amp, size=forward_src_i.size, center=forward_src_i.center, **self.srckwarg, ) if mp.is_electric(src_i.component): src_i.amplitude *= -1 sources += [src_i] return sources def __call__(self): """The values of LDOS at each frequency Returns: 1D array of LDOS corresponding to each of self.frequencies """ self._eval = self.sim.ldos_data self._ldos_scale = self.sim.ldos_scale self._ldos_Jdata = self.sim.ldos_Jdata return np.array(self._eval)