Yee grid gradients (NanoComp#1285)

* add yee grid to array slice

* add yee grid to array slice

* fix gradients

* rebase

* try better gradients

* fix interpolation weights

* cleanup meepgeom

* cleanup from revert

* fingers crossed

* fix multifreq bug

* add support for python2

* found more freq bugs

* mpi memory fixes

* fix memory leaks

* fix dtft term

* minor filter fix

* comment out freq patch and cleanup scalars

* restore mpb.cpp

Co-authored-by: Alec Hammond <ahammond36@gatech.edu>

python/adjoint/filter_source.py

@@ -66,12 +66,12 @@ def nuttall_dtft(self,f,f0):
         a = [0.355768, 0.4873960, 0.144232, 0.012604]
         return self.cos_window_fd(a,f,f0)
     def dtft(self,y,f):
-        return np.matmul(np.exp(1j*2*np.pi*f[:,np.newaxis]*np.arange(y.size)*self.dt), y)
+        return np.matmul(np.exp(1j*2*np.pi*f[:,np.newaxis]*np.arange(y.size)*self.dt), y)*self.dt/np.sqrt(2*np.pi)
 
     def __call__(self,t):
         if t > self.T:
             return 0
-        vec = self.nuttall(t,self.center_frequencies)
+        vec = self.nuttall(t,self.center_frequencies) / (self.dt/np.sqrt(2*np.pi)) # compensate for meep dtft
         return np.inner(vec,self.nodes)
 
     def func(self):

python/adjoint/filters.py

@@ -279,7 +279,7 @@ def gaussian_filter(x,sigma,Lx,Ly,resolution,symmetries=[]):
     kernel = kernel / np.sum(kernel.flatten()) # Normalize the filter
 
     # Filter the response
-    y = mpa.simple_2d_filter(x,kernel,Lx,Ly,resolution,symmetries)
+    y = simple_2d_filter(x,kernel,Lx,Ly,resolution,symmetries)
 
     return y
 

python/adjoint/objective.py

@@ -38,17 +38,20 @@ def __init__(self,sim,volume,mode,forward=True,kpoint_func=None,**kwargs):
 
     def register_monitors(self,frequencies):
         self.frequencies = np.asarray(frequencies)
-        self.monitor = self.sim.add_mode_monitor(frequencies,mp.FluxRegion(center=self.volume.center,size=self.volume.size))
+        self.monitor = self.sim.add_mode_monitor(frequencies,mp.FluxRegion(center=self.volume.center,size=self.volume.size),yee_grid=True)
         self.normal_direction = self.monitor.normal_direction
         return self.monitor
 
-    def place_adjoint_source(self,dJ,dt):
+    def place_adjoint_source(self,dJ):
         '''Places an equivalent eigenmode monitor facing the opposite direction. Calculates the 
         correct scaling/time profile.
 
         dJ ........ the user needs to pass the dJ/dMonitor evaluation
         dt ........ the timestep size from sim.fields.dt of the forward sim
         '''
+        dt = self.sim.fields.dt
+        T = self.sim.meep_time()
+
         dJ = np.atleast_1d(dJ)
         # determine starting kpoint for reverse mode eigenmode source
         direction_scalar = 1 if self.forward else -1
@@ -76,28 +79,39 @@ def place_adjoint_source(self,dJ,dt):
             dV = 1/self.sim.resolution * 1/self.sim.resolution
         else:
             dV = 1/self.sim.resolution * 1/self.sim.resolution * 1/self.sim.resolution
-        da_dE = 0.5*(dV * self.cscale)
-        scale = da_dE * dJ * 1j * 2 * np.pi * self.frequencies / np.array([self.time_src.fourier_transform(f) for f in self.frequencies]) # final scale factor
+
+        da_dE = 0.5 * self.cscale # scalar popping out of derivative
+        iomega = (1.0 - np.exp(-1j * (2 * np.pi * self.frequencies) * dt)) * (1.0 / dt) # scaled frequency factor with discrete time derivative fix
+
+        src = self.time_src
+
+        # 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) # dtft
+
+        # we need to compensate for the phase added by the time envelope at our freq of interest
+        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))
+
         if self.frequencies.size == 1:
             # Single frequency simulations. We need to drive it with a time profile.
-            src = self.time_src
-            amp = scale
+            amp = da_dE * dV * dJ * iomega / fwd_dtft / adj_src_phase # final scale factor
         else:
-            # TODO: In theory we should be able drive the source without normalizing out the time profile.
-            # But for some reason, there is a frequency dependent scaling discrepency. It works now for 
-            # multiple monitors and multiple sources, but we should figure out why this is.
-            src = FilteredSource(self.time_src.frequency,self.frequencies,scale,dt,self.time_src) # generate source from broadband response
+            # multi frequency simulations
+            scale = da_dE * dV * dJ * iomega / adj_src_phase
+            src = FilteredSource(self.time_src.frequency,self.frequencies,scale,dt) # generate source from broadband response
             amp = 1
+
         # generate source object
-        self.source = mp.EigenModeSource(src,
+        self.source = [mp.EigenModeSource(src,
                     eig_band=self.mode,
                     direction=mp.NO_DIRECTION,
                     eig_kpoint=k0,
                     amplitude=amp,
                     eig_match_freq=True,
                     size=self.volume.size,
                     center=self.volume.center,
-                    **self.EigenMode_kwargs)
+                    **self.EigenMode_kwargs)]
 
         return self.source
 
@@ -108,7 +122,7 @@ def __call__(self):
         # Eigenmode data
         direction = mp.NO_DIRECTION if self.kpoint_func else mp.AUTOMATIC
         ob = self.sim.get_eigenmode_coefficients(self.monitor,[self.mode],direction=direction,kpoint_func=self.kpoint_func,**self.EigenMode_kwargs)
-        self.eval = np.squeeze(ob.alpha[:,:,self.forward]) # record eigenmode coefficients for scaling   
+        self.eval = np.squeeze(ob.alpha[:,:,self.forward]) # record eigenmode coefficients for scaling 
         self.cscale = ob.cscale # pull scaling factor
 
         return self.eval

python/adjoint/optimization_problem.py

@@ -5,28 +5,35 @@
 from collections import namedtuple
 
 Grid = namedtuple('Grid', ['x', 'y', 'z', 'w'])
+YeeDims = namedtuple('YeeDims', ['Ex','Ey','Ez'])
 
 class DesignRegion(object):
-    def __init__(self,design_parameters,volume=None, size=None, center=mp.Vector3()):
+    def __init__(self,design_parameters,volume=None, size=None, center=mp.Vector3(), MaterialGrid=None):
         self.volume = volume if volume else 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
+        self.MaterialGrid=MaterialGrid
     def update_design_parameters(self,design_parameters):
         self.design_parameters.update_parameters(design_parameters)
-    def get_gradient(self,fields_a,fields_f,frequencies,geom_list,f):
-        # sanitize the input
-        if (fields_a.ndim < 5) or (fields_f.ndim < 5):
-            raise ValueError("Fields arrays must have 5 dimensions (x,y,z,frequency,component)")
+    def get_gradient(self,sim,fields_a,fields_f,frequencies):
+        for c in range(3):
+            fields_a[c] = fields_a[c].flatten(order='C')
+            fields_f[c] = fields_f[c].flatten(order='C')
+        fields_a = np.concatenate(fields_a)
+        fields_f = np.concatenate(fields_f)
         num_freqs = np.array(frequencies).size
 
-        grad = np.zeros((self.num_design_params*num_freqs,)) # preallocate
+        grad = np.zeros((num_freqs,self.num_design_params)) # preallocate
 
+        geom_list = sim.geometry
+        f = sim.fields
+        vol = sim._fit_volume_to_simulation(self.volume)
         # compute the gradient
-        mp._get_gradient(grad,fields_a,fields_f,self.volume,np.array(frequencies),geom_list,f)
+        mp._get_gradient(grad,fields_a,fields_f,vol,np.array(frequencies),geom_list,f) 
 
-        return np.squeeze(grad.reshape(self.num_design_params,num_freqs,order='F'))
+        return np.squeeze(grad).T
 
 class OptimizationProblem(object):
     """Top-level class in the MEEP adjoint module.
@@ -108,6 +115,8 @@ def __init__(self,
         #    ADJ  - The optimizer has already run an adjoint simulation (but not yet calculated the gradient)
         self.current_state = "INIT"
 
+        self.gradient = []
+
     def __call__(self, rho_vector=None, need_value=True, need_gradient=True):
         """Evaluate value and/or gradient of objective function.
         """
@@ -176,17 +185,20 @@ def prepare_forward_run(self):
             self.forward_monitors.append(m.register_monitors(self.frequencies))
 
         # register design region
-        self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=False) for dr in self.design_regions]
+        self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=True) for dr in self.design_regions]
 
         # store design region voxel parameters
-        self.design_grids = [Grid(*self.sim.get_array_metadata(dft_cell=drm)) for drm in self.design_region_monitors]
+        self.design_grids = []
+        for drm in self.design_region_monitors:
+            s = [self.sim.get_array_slice_dimensions(c,vol=drm.where)[0] for c in [mp.Ex,mp.Ey,mp.Ez]]
+            self.design_grids += [YeeDims(*s)]
 
     def forward_run(self):
         # set up monitors
         self.prepare_forward_run()
 
         # Forward run
-        self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, self.minimum_run_time))
+        self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.design_region_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, True, self.minimum_run_time))
 
         # record objective quantities from user specified monitors
         self.results_list = []
@@ -199,53 +211,55 @@ def forward_run(self):
             self.f0 = self.f0[0]
 
         # Store forward fields for each set of design variables in array (x,y,z,field_components,frequencies)
-        self.d_E = [np.zeros((len(dg.x),len(dg.y),len(dg.z),self.nf,3),dtype=np.complex128) for dg in self.design_grids]
+        self.d_E = [[np.zeros((self.nf,c[0],c[1],c[2]),dtype=np.complex128) for c in dg] for dg in self.design_grids]
         for nb, dgm in enumerate(self.design_region_monitors):
-            for f in range(self.nf):
-                for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
-                    self.d_E[nb][:,:,:,f,ic] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
+            for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
+                for f in range(self.nf):
+                    self.d_E[nb][ic][f,:,:,:] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
 
         # store objective function evaluation in memory
         self.f_bank.append(self.f0)
 
         # update solver's current state
         self.current_state = "FWD"
-
-    def adjoint_run(self,objective_idx=0):
-        # Grab the simulation step size from the forward run
-        self.dt = self.sim.fields.dt
-
-        # Prepare adjoint run
-        self.sim.reset_meep()
-
-        # Replace sources with adjoint sources
+    def prepare_adjoint_run(self,objective_idx):
+        # Compute adjoint sources
         self.adjoint_sources = []
         for mi, m in enumerate(self.objective_arguments):
             dJ = jacobian(self.objective_functions[objective_idx],mi)(*self.results_list) # get gradient of objective w.r.t. monitor
-            self.adjoint_sources.append(m.place_adjoint_source(dJ,self.dt)) # place the appropriate adjoint sources
+            if np.any(dJ):
+                self.adjoint_sources += m.place_adjoint_source(dJ) # place the appropriate adjoint sources
+
+        # Reset the fields
+        self.sim.reset_meep()
+
+        # Update the sources
         self.sim.change_sources(self.adjoint_sources)
 
         # register design flux
-        # TODO use yee grid directly
-        self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=False) for dr in self.design_regions]
+        self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=True) for dr in self.design_regions]
+
+    def adjoint_run(self,objective_idx=0):
+        # set up adjoint sources and monitors
+        self.prepare_adjoint_run(objective_idx)
 
         # Adjoint run
-        self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.design_region_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, self.minimum_run_time))
+        self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.design_region_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, True, self.minimum_run_time))
 
         # Store adjoint fields for each design set of design variables in array (x,y,z,field_components,frequencies)
-        self.a_E.append([np.zeros((len(dg.x),len(dg.y),len(dg.z),self.nf,3),dtype=np.complex128) for dg in self.design_grids])
+        self.a_E.append([[np.zeros((self.nf,c[0],c[1],c[2]),dtype=np.complex128) for c in dg] for dg in self.design_grids])
         for nb, dgm in enumerate(self.design_region_monitors):
-            for f in range(self.nf):
-                for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
-                    self.a_E[objective_idx][nb][:,:,:,f,ic] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
+            for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
+                for f in range(self.nf):
+                    self.a_E[objective_idx][nb][ic][f,:,:,:] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
 
         # update optimizer's state
         self.current_state = "ADJ"
 
     def calculate_gradient(self):
         # Iterate through all design regions and calculate gradient
-        self.gradient = [[dr.get_gradient(self.a_E[ar][dri],self.d_E[dri],self.frequencies,self.sim.geometry,self.sim.fields) for dri, dr in enumerate(self.design_regions)] for ar in range(len(self.objective_functions))]
-
+        self.gradient = [[dr.get_gradient(self.sim,self.a_E[ar][dri],self.d_E[dri],self.frequencies) for dri, dr in enumerate(self.design_regions)] for ar in range(len(self.objective_functions))]
+        
         # Cleanup list of lists
         if len(self.gradient) == 1:
             self.gradient = self.gradient[0] # only one objective function
@@ -308,9 +322,9 @@ def calculate_fd_gradient(self,num_gradients=1,db=1e-4,design_variables_idx=0,fi
             self.forward_monitors = []
             for m in self.objective_arguments:
                 self.forward_monitors.append(m.register_monitors(self.frequencies))
-
-            self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, self.minimum_run_time))
-
+            
+            self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, True, self.minimum_run_time))
+            
             # record final objective function value
             results_list = []
             for m in self.objective_arguments:
@@ -332,8 +346,8 @@ def calculate_fd_gradient(self,num_gradients=1,db=1e-4,design_variables_idx=0,fi
                 self.forward_monitors.append(m.register_monitors(self.frequencies))
 
             # add monitor used to track dft convergence
-            self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, self.minimum_run_time))
-
+            self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, True, self.minimum_run_time))
+            
             # record final objective function value
             results_list = []
             for m in self.objective_arguments:
@@ -380,7 +394,7 @@ def plot2D(self,init_opt=False, **kwargs):
 
         self.sim.plot2D(**kwargs)
 
-def stop_when_dft_decayed(simob, mon, dt, c, fcen_idx, decay_by, minimum_run_time=0):
+def stop_when_dft_decayed(simob, mon, dt, c, fcen_idx, decay_by, yee_grid=False, minimum_run_time=0):
     '''Step function that monitors the relative change in DFT fields for a list of monitors.
 
     mon ............. a list of monitors
@@ -389,18 +403,17 @@ def stop_when_dft_decayed(simob, mon, dt, c, fcen_idx, decay_by, minimum_run_tim
     '''
 
     # get monitor dft output array dimensions
-    x = []
-    y = []
-    z = []
+    dims = []
     for m in mon:
-        xi,yi,zi,wi = simob.get_array_metadata(dft_cell=m)
-        x.append(len(xi))
-        y.append(len(yi))
-        z.append(len(zi))
-
+        ci_dims = []
+        for ci in c:
+            comp = ci if yee_grid else mp.Dielectric
+            ci_dims += [simob.get_array_slice_dimensions(comp,vol=m.where)[0]]
+        dims.append(ci_dims)
+
     # Record data in closure so that we can persitently edit
     closure = {
-        'previous_fields': [np.ones((x[mi],y[mi],z[mi],len(c)),dtype=np.complex128) for mi, m in enumerate(mon)],
+        'previous_fields': [[np.ones(di,dtype=np.complex128) for di in d] for d in dims],
         't0': 0
     }
 
@@ -413,17 +426,17 @@ def _stop(sim):
 
             # Pull all the relevant frequency and spatial dft points
             relative_change = []
-            current_fields = [np.zeros((x[mi],y[mi],z[mi],len(c)), dtype=np.complex128) for mi in range(len(mon))]
+            current_fields = [[0 for di in d] for d in dims]
             for mi, m in enumerate(mon):
                 for ic, cc in enumerate(c):
                     if isinstance(m,mp.DftFlux):
-                        current_fields[mi][:,:,:,ic] = atleast_3d(mp.get_fluxes(m)[fcen_idx])
+                        current_fields[mi][ic] = mp.get_fluxes(m)[fcen_idx]
                     elif isinstance(m,mp.DftFields):
-                        current_fields[mi][:,:,:,ic] = atleast_3d(sim.get_dft_array(m, cc, fcen_idx))
+                        current_fields[mi][ic] = atleast_3d(sim.get_dft_array(m, cc, fcen_idx))
                     else:
                         raise TypeError("Monitor of type {} not supported".format(type(m)))
-                relative_change_raw = np.abs(previous_fields[mi] - current_fields[mi]) / np.abs(previous_fields[mi])
-                relative_change.append(np.mean(relative_change_raw.flatten())) # average across space and frequency
+                    relative_change_raw = np.abs(previous_fields[mi][ic] - current_fields[mi][ic]) / np.abs(previous_fields[mi][ic])
+                    relative_change.append(np.mean(relative_change_raw.flatten())) # average across space and frequency
             relative_change = np.mean(relative_change) # average across monitors
             closure['previous_fields'] = current_fields
             closure['t0'] = sim.round_time()