Adjoint near2far with multiple far pts (#1417)

* multiple far points * tutorials and minor fixes * tutorials and minor fixes * multiple far pts * multiple far pts * multiple far pts * minor fix * minor fix * update tutorial and minor fix * update tutorial and minor fix * update tutorial and minor fix * update tutorial and minor fix Co-authored-by: Mo Chen <mochen@MacBook-Pro-2.local> Co-authored-by: Mo Chen <mochen@dhcp-10-29-254-74.dyn.MIT.EDU>
NanoComp · Dec 3, 2020 · da4b7db · da4b7db
1 parent fb0be23
commit da4b7db
Show file tree

Hide file tree

Showing 8 changed files with 4,391 additions and 670 deletions.
diff --git a/python/adjoint/objective.py b/python/adjoint/objective.py
@@ -4,7 +4,8 @@
 import numpy as np
 import meep as mp
 from .filter_source import FilteredSource
-from .optimization_problem import atleast_3d, Grid
+from .optimization_problem import Grid
+from meep.simulation import py_v3_to_vec
 
 class ObjectiveQuantitiy(ABC):
     @abstractmethod
@@ -87,7 +88,7 @@ def place_adjoint_source(self,dJ):
                     size=self.volume.size,
                     center=self.volume.center,
                     **self.EigenMode_kwargs)]
-        
+
         return self.source
 
     def __call__(self):
@@ -129,24 +130,26 @@ def place_adjoint_source(self,dJ):
         self.sources = []
         scale = adj_src_scale(self, dt)
 
+        x_dim, y_dim, z_dim = len(self.dg.x), len(self.dg.y), len(self.dg.z)
+
         if self.num_freq == 1:
-            amp = -atleast_3d(dJ[0]) * scale
+            amp = -dJ[0].copy().reshape(x_dim, y_dim, z_dim) * scale
             src = self.time_src
             if self.component in [mp.Hx, mp.Hy, mp.Hz]:
                 amp = -amp
-            for zi in range(len(self.dg.z)):
-                for yi in range(len(self.dg.y)):
-                    for xi in range(len(self.dg.x)):
+            for zi in range(z_dim):
+                for yi in range(y_dim):
+                    for xi in range(x_dim):
                         if amp[xi, yi, zi] != 0:
                             self.sources += [mp.Source(src, component=self.component, amplitude=amp[xi, yi, zi],
                             center=mp.Vector3(self.dg.x[xi], self.dg.y[yi], self.dg.z[zi]))]
         else:
-            dJ_4d = np.array([atleast_3d(dJ[f]) for f in range(self.num_freq)])
+            dJ_4d = np.array([dJ[f].copy().reshape(x_dim, y_dim, z_dim) for f in range(self.num_freq)])
             if self.component in [mp.Hx, mp.Hy, mp.Hz]:
                 dJ_4d = -dJ_4d
-            for zi in range(len(self.dg.z)):
-                for yi in range(len(self.dg.y)):
-                    for xi in range(len(self.dg.x)):
+            for zi in range(z_dim):
+                for yi in range(y_dim):
+                    for xi in range(x_dim):
                         final_scale = -dJ_4d[:,xi,yi,zi] * scale
                         src = FilteredSource(self.time_src.frequency,self.frequencies,final_scale,dt)
                         self.sources += [mp.Source(src, component=self.component, amplitude=1,
@@ -157,7 +160,7 @@ def place_adjoint_source(self,dJ):
     def __call__(self):
         self.time_src = self.sim.sources[0].src
         self.dg = Grid(*self.sim.get_array_metadata(dft_cell=self.monitor))
-        self.eval = np.array([self.sim.get_dft_array(self.monitor, self.component, i) for i in range(self.num_freq)]) #Shape = (num_freq, [pts])
+        self.eval = np.array([self.sim.get_dft_array(self.monitor, self.component, i) for i in range(self.num_freq)])
         return self.eval
 
     def get_evaluation(self):
@@ -168,11 +171,12 @@ def get_evaluation(self):
 
 
 class Near2FarFields(ObjectiveQuantitiy):
-    def __init__(self,sim,Near2FarRegions, far_pt):
+    def __init__(self,sim,Near2FarRegions, far_pts):
         self.sim = sim
         self.Near2FarRegions=Near2FarRegions
         self.eval = None
-        self.far_pt = far_pt
+        self.far_pts = far_pts #list of far pts
+        self.nfar_pts = len(far_pts)
         return
 
     def register_monitors(self,frequencies):
@@ -184,10 +188,14 @@ def register_monitors(self,frequencies):
     def place_adjoint_source(self,dJ):
         dt = self.sim.fields.dt # the timestep size from sim.fields.dt of the forward sim
         self.sources = []
+        if dJ.ndim == 4:
+            dJ = np.sum(dJ,axis=0)
         dJ = dJ.flatten()
+        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)
 
-        #TODO far_pts in 3d or cylindrical, perhaps py_v3_to_vec from simulation.py
-        self.all_nearsrcdata = self.monitor.swigobj.near_sourcedata(mp.vec(self.far_pt.x, self.far_pt.y), dJ)
+        self.all_nearsrcdata = self.monitor.swigobj.near_sourcedata(far_pt_vec, farpt_list, self.nfar_pts, dJ)
         for near_data in self.all_nearsrcdata:
             cur_comp = near_data.near_fd_comp
             amp_arr = np.array(near_data.amp_arr).reshape(-1, self.num_freq)
@@ -205,8 +213,7 @@ def place_adjoint_source(self,dJ):
 
     def __call__(self):
         self.time_src = self.sim.sources[0].src
-        self.eval = np.array(self.sim.get_farfield(self.monitor, self.far_pt))
-        self.eval = self.eval.reshape((self.num_freq, 6))
+        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
 
     def get_evaluation(self):

diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py
@@ -61,7 +61,8 @@ def __init__(self,
                 decay_dt=50,
                 decay_fields=[mp.Ez],
                 decay_by=1e-6,
-                minimum_run_time=0
+                minimum_run_time=0,
+                maximum_run_time=None
                  ):
 
         self.sim = simulation
@@ -106,6 +107,7 @@ def __init__(self,
         self.decay_fields=decay_fields
         self.decay_dt=decay_dt
         self.minimum_run_time=minimum_run_time
+        self.maximum_run_time=maximum_run_time
 
         # store sources for finite difference estimations
         self.forward_sources = self.sim.sources
@@ -137,15 +139,13 @@ def __call__(self, rho_vector=None, need_value=True, need_gradient=True):
                 self.forward_run()
                 print("Starting adjoint run...")
                 self.a_E = []
-                for ar in range(len(self.objective_functions)):
-                    self.adjoint_run(ar)
+                self.adjoint_run()
                 print("Calculating gradient...")
                 self.calculate_gradient()
             elif self.current_state == "FWD":
                 print("Starting adjoint run...")
                 self.a_E = []
-                for ar in range(len(self.objective_functions)):
-                    self.adjoint_run(ar)
+                self.adjoint_run()
                 print("Calculating gradient...")
                 self.calculate_gradient()
             else:
@@ -199,7 +199,7 @@ def forward_run(self):
         self.prepare_forward_run()
 
         # Forward 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, True, 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, self.maximum_run_time))
 
         # record objective quantities from user specified monitors
         self.results_list = []
@@ -223,36 +223,37 @@ def forward_run(self):
 
         # update solver's current state
         self.current_state = "FWD"
-    def prepare_adjoint_run(self,objective_idx):
+    def prepare_adjoint_run(self):
         # 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
-            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)
+        self.adjoint_sources = [[] for i in range(len(self.objective_functions))]
+        for ar in range(len(self.objective_functions)):
+            for mi, m in enumerate(self.objective_arguments):
+                dJ = jacobian(self.objective_functions[ar],mi)(*self.results_list) # get gradient of objective w.r.t. monitor
+                if np.any(dJ):
+                    self.adjoint_sources[ar] += m.place_adjoint_source(dJ) # place the appropriate adjoint sources
+
+    def adjoint_run(self):
+        # set up adjoint sources and monitors
+        self.prepare_adjoint_run()
+        for ar in range(len(self.objective_functions)):
+            # Reset the fields
+            self.sim.reset_meep()
 
-        # register design flux
-        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]
+            # Update the sources
+            self.sim.change_sources(self.adjoint_sources[ar])
 
-    def adjoint_run(self,objective_idx=0):
-        # set up adjoint sources and monitors
-        self.prepare_adjoint_run(objective_idx)
+            # register design flux
+            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, True, self.minimum_run_time))
+            # 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, True, self.minimum_run_time, self.maximum_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((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 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))
+            # Store adjoint fields for each design set of design variables in array (x,y,z,field_components,frequencies)
+            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 ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
+                    for f in range(self.nf):
+                        self.a_E[ar][nb][ic][f,:,:,:] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
 
         # update optimizer's state
         self.current_state = "ADJ"
@@ -324,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, True, 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, self.maximum_run_time))
 
 
             # record final objective function value
@@ -348,7 +349,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, True, 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, self.maximum_run_time))
 
             # record final objective function value
             results_list = []
@@ -397,7 +398,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, yee_grid=False, minimum_run_time=0):
+def stop_when_dft_decayed(simob, mon, dt, c, fcen_idx, decay_by, yee_grid=False, minimum_run_time=0, maximum_run_time=None):
     '''Step function that monitors the relative change in DFT fields for a list of monitors.
 
     mon ............. a list of monitors
@@ -422,6 +423,8 @@ def stop_when_dft_decayed(simob, mon, dt, c, fcen_idx, decay_by, yee_grid=False,
     def _stop(sim):
         if sim.round_time() <= dt + closure['t0']:
             return False
+        elif maximum_run_time and sim.round_time() > maximum_run_time:
+            return True
         else:
             previous_fields = closure['previous_fields']