Yee grid gradients

python/adjoint/objective.py

@@ -4,6 +4,19 @@
 import numpy as np
 import meep as mp
 from .filter_source import FilteredSource
+from .optimization_problem import YeeDims, atleast_3d
+
+# ---------------------------------------------- #
+# general-purpose constants and utility routines
+# ---------------------------------------------- #
+ORIGIN           = np.zeros(3)
+XHAT, YHAT, ZHAT = [ mp.Vector3(a) for a in [[1.,0,0], [0,1.,0], [0,0,1.]] ]
+E_CPTS           = [mp.Ex, mp.Ey, mp.Ez]
+H_CPTS           = [mp.Hx, mp.Hy, mp.Hz]
+EH_CPTS          = E_CPTS + H_CPTS
+EH_TRANSVERSE    = [ [mp.Ey, mp.Ez, mp.Hy, mp.Hz],
+                     [mp.Ez, mp.Ex, mp.Hz, mp.Hx],
+                     [mp.Ex, mp.Ey, mp.Hx, mp.Hy] ]
 
 class ObjectiveQuantitiy(ABC):
     @abstractmethod
@@ -30,16 +43,17 @@ def __init__(self,sim,volume,mode,forward=True,kpoint_func=None,**kwargs):
         self.volume=volume
         self.mode=mode
         self.forward = 0 if forward else 1
-        self.normal_direction = None
+        self.normal_direction = self.volume.swigobj.normal_direction()
         self.kpoint_func = kpoint_func
         self.eval = None
         self.EigenMode_kwargs = kwargs
-        return
+        self.components = EH_TRANSVERSE[self.normal_direction]
 
     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.normal_direction = self.monitor.normal_direction
+        self.nf = self.frequencies.size
+        #self.monitor = self.sim.add_mode_monitor(frequencies,mp.FluxRegion(center=self.volume.center,size=self.volume.size))
+        self.monitor = self.sim.add_dft_fields(self.components,frequencies,where=self.volume,yee_grid=True)
         return self.monitor
 
     def place_adjoint_source(self,dJ,dt):
@@ -76,8 +90,9 @@ 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
+
+        iomega = ((1.0 - np.exp(-1j * (2 * np.pi * self.frequencies) * dt)) * (1.0 / dt))
+        scale = self.dx * dV * dJ * iomega / np.array([self.time_src.fourier_transform(f) for f in self.frequencies]) # final scale factor
         if self.frequencies.size == 1:
             # Single frequency simulations. We need to drive it with a time profile.
             src = self.time_src
@@ -88,30 +103,69 @@ def place_adjoint_source(self,dJ,dt):
             # 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
             amp = 1
-        # generate source object
-        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)
+
+        sign  =  -1.0 if self.forward else 1.0
+        signs = [ -1.0, +1.0, -1.0*sign, +1.0*sign ]
+        # -------------------------------------- #
+        # Place sources
+        # -------------------------------------- #
+        self.source = []
+        for ci,c in enumerate(self.components):
+            for ix,rx in enumerate(np.linspace(self.min_corners[ci].x,self.max_corners[ci].x,self.component_metadata[ci][0])):
+                for iy,ry in enumerate(np.linspace(self.min_corners[ci].y,self.max_corners[ci].y,self.component_metadata[ci][1])):
+                    for iz,rz in enumerate(np.linspace(self.min_corners[ci].z,self.max_corners[ci].z,self.component_metadata[ci][2])):
+                        cur_amp = amp * self.EH[ci][ix,iy,iz] * signs[ci]
+                        self.source += [mp.Source(src,component=self.components[3-ci],amplitude=cur_amp,size=mp.Vector3(),center=mp.Vector3(rx,ry,rz))]
 
         return self.source
 
     def __call__(self):
         # We just need a workable time profile, so just grab the first available time profile and use that.
         self.time_src = self.sim.sources[0].src
 
-        # 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.cscale = ob.cscale # pull scaling factor
+        # store the array metadata for each field component
+        self.component_metadata = []; self.min_corners = []; self.max_corners = []
+        for c in self.components:
+            s, min_corner, max_corner = self.sim.get_array_slice_dimensions(c,vol=self.volume)
+            self.component_metadata += [s]; self.min_corners += [min_corner]; self.max_corners += [max_corner]
+
+        # pull the dft fields on the yee grid
+        self.EH = [np.zeros((cm[0],cm[1],cm[2],self.nf),dtype=np.complex128) for cm in self.component_metadata]
+        for ci,c in enumerate(self.components):
+            for f in range(self.nf):
+                self.EH[ci][:,:,:,f] = atleast_3d(self.sim.get_dft_array(self.monitor,c,f))
+
+        # compute and store the eigenmode profiles on the yee grid
+        direction = mp.NO_DIRECTION if self.kpoint_func else self.normal_direction
+        eh = [np.zeros((cm[0],cm[1],cm[2],self.nf),dtype=np.complex128) for cm in self.component_metadata]
+        for f in range(self.nf):
+            kpoint = self.kpoint_func(self.frequencies[f],1) if self.kpoint_func else mp.Vector3()
+            eig_data = self.sim.get_eigenmode(frequency=self.frequencies[f],direction=direction,
+                kpoint=kpoint,where=self.volume,band_num=self.mode,resolution=2*self.sim.resolution,
+                **self.EigenMode_kwargs)
+            for ci,c in enumerate(self.components):
+                for ix,rx in enumerate(np.linspace(self.min_corners[ci].x,self.max_corners[ci].x,self.component_metadata[ci][0])):
+                    for iy,ry in enumerate(np.linspace(self.min_corners[ci].y,self.max_corners[ci].y,self.component_metadata[ci][1])):
+                        for iz,rz in enumerate(np.linspace(self.min_corners[ci].z,self.max_corners[ci].z,self.component_metadata[ci][2])):
+                            eh[ci][ix,iy,iz,f] = eig_data.amplitude(mp.Vector3(rx,ry,rz),c) 
+
+        # compute eigenmode coefficient for each freq
+        cp = np.sum(np.conj(eh[2]) * self.EH[1],axis=(0,1,2)) - np.sum(np.conj(eh[3]) * self.EH[0],axis=(0,1,2))
+        cm = np.sum(np.conj(eh[0]) * self.EH[3],axis=(0,1,2)) - np.sum(np.conj(eh[1]) * self.EH[2],axis=(0,1,2))
+        self.eval = -(cp-cm)/4 if self.forward else (cp+cm)/4
+
+        if (((self.volume.size.y == 0) and (self.volume.size.z == 0)) or 
+            ((self.volume.size.x == 0) and (self.volume.size.z == 0)) or
+            ((self.volume.size.x == 0) and (self.volume.size.y == 0))):
+            self.dx = self.sim.resolution ** (-1)
+        elif (((self.volume.size.y != 0) and (self.volume.size.z != 0)) or 
+            ((self.volume.size.x != 0) and (self.volume.size.z != 0)) or
+            ((self.volume.size.x != 0) and (self.volume.size.y != 0))):
+            self.dx = self.sim.resolution ** (-2)
+        else:
+            self.dx = self.sim.resolution ** (-3)
 
-        return self.eval
+        return self.eval / self.dx
     def get_evaluation(self):
         '''Returns the requested eigenmode coefficient.
         '''

python/adjoint/optimization_problem.py

@@ -5,6 +5,7 @@
 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()):
@@ -16,9 +17,11 @@ def __init__(self,design_parameters,volume=None, size=None, center=mp.Vector3())
     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)")
+        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
@@ -176,17 +179,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.forward_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,11 +205,11 @@ 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((c[0],c[1],c[2],self.nf),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)
@@ -222,22 +228,32 @@ def adjoint_run(self,objective_idx=0):
         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
+            self.adjoint_sources += m.place_adjoint_source(dJ,self.dt) # place the appropriate adjoint 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]
 
         # 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))
+        '''from matplotlib import pyplot as plt
+        self.sim.run(until=5)
+        plt.figure(figsize=(10,10))
+        for i,(c,t) in enumerate(zip([mp.Ez,mp.Ex,mp.Hz,mp.Hx],['Ez','Ex','Hz','Hx'])):
+            print(c)
+            plt.subplot(2,2,i+1)
+            self.sim.plot2D(fields=c)
+            plt.title(t)
+        plt.show()
+        quit()'''
+        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((c[0],c[1],c[2],self.nf),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"
@@ -309,7 +325,7 @@ def calculate_fd_gradient(self,num_gradients=1,db=1e-4,design_variables_idx=0,fi
             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 = []
@@ -332,7 +348,7 @@ 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 = []
@@ -380,7 +396,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 +405,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 +428,18 @@ 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 = [[np.zeros(di,dtype=np.complex128) 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][:,:,:] = atleast_3d(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))
+                        print(current_fields[mi][ic].shape)
+                        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()