adjoint for ldos (#2077)

* ldos adjoint

* ldos_scale

* missing ;

* resolve

* test

* ldos adjoint

* ldos_scale

* missing ;

* resolve

* test

Co-authored-by: Mo Chen <mochen@Mos-MacBook-Pro.local>
Co-authored-by: Mo Chen <mochen@dhcp-10-31-50-19.dyn.mit.edu>

python/adjoint/objective.py

@@ -353,3 +353,50 @@ def __call__(self):
             for far_pt in self.far_pts
         ]).reshape((self._nfar_pts, self.num_freq, 6))
         return self._eval
+
+class LDOS(ObjectiveQuantity):
+    def __init__(self, sim, decimation_factor=0, **kwargs):
+        super().__init__(sim)
+        self.decimation_factor = decimation_factor
+        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):
+        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)

python/adjoint/optimization_problem.py

@@ -3,7 +3,7 @@
 from autograd import grad, jacobian
 from collections import namedtuple
 
-from . import utils, DesignRegion
+from . import utils, DesignRegion, LDOS
 
 class OptimizationProblem(object):
     """Top-level class in the MEEP adjoint module.
@@ -175,11 +175,16 @@ def forward_run(self):
         self.prepare_forward_run()
 
         # Forward run
-        self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
-            self.decay_by,
-            self.minimum_run_time,
-            self.maximum_run_time
-        ))
+        if any(isinstance(m, LDOS) for m in self.objective_arguments):
+            self.sim.run(mp.dft_ldos(self.frequencies), until_after_sources=mp.stop_when_dft_decayed(
+                self.decay_by,
+                self.minimum_run_time,
+                self.maximum_run_time))
+        else:
+            self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
+                self.decay_by,
+                self.minimum_run_time,
+                self.maximum_run_time))
 
         # record objective quantities from user specified monitors
         self.results_list = []
@@ -354,11 +359,13 @@ def calculate_fd_gradient(
                 self.forward_monitors.append(
                     m.register_monitors(self.frequencies))
 
-            self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
-                self.decay_by,
-                self.minimum_run_time,
-                self.maximum_run_time
-            ))
+            if any(isinstance(m, LDOS) for m in self.objective_arguments):
+                self.sim.run(mp.dft_ldos(self.frequencies), until_after_sources=mp.stop_when_energy_decayed(dt=1, decay_by=1e-11))
+            else:
+                self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
+                    self.decay_by,
+                    self.minimum_run_time,
+                    self.maximum_run_time))
 
             # record final objective function value
             results_list = []
@@ -385,11 +392,13 @@ def calculate_fd_gradient(
                     m.register_monitors(self.frequencies))
 
             # add monitor used to track dft convergence
-            self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
-                self.decay_by,
-                self.minimum_run_time,
-                self.maximum_run_time
-            ))
+            if any(isinstance(m, LDOS) for m in self.objective_arguments):
+                self.sim.run(mp.dft_ldos(self.frequencies), until_after_sources=mp.stop_when_energy_decayed(dt=1, decay_by=1e-11))
+            else:
+                self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
+                    self.decay_by,
+                    self.minimum_run_time,
+                    self.maximum_run_time))
 
             # record final objective function value
             results_list = []

python/simulation.py

@@ -5213,6 +5213,7 @@ def _ldos(sim, todo):
             sim.ldos_data = mp._dft_ldos_ldos(ldos)
             sim.ldos_Fdata = mp._dft_ldos_F(ldos)
             sim.ldos_Jdata = mp._dft_ldos_J(ldos)
+            sim.ldos_scale = ldos.overall_scale()
             if verbosity.meep > 0:
                 display_csv(sim, 'ldos', zip(mp.get_ldos_freqs(ldos), sim.ldos_data))
     return _ldos

python/tests/test_adjoint_solver.py

@@ -10,7 +10,7 @@
 from enum import Enum
 from utils import ApproxComparisonTestCase
 
-MonitorObject = Enum('MonitorObject', 'EIGENMODE DFT')
+MonitorObject = Enum('MonitorObject', 'EIGENMODE DFT LDOS')
 
 resolution = 30
 
@@ -57,6 +57,11 @@
                        center=mp.Vector3(-0.5*sxy+dpml,0),
                        size=mp.Vector3(),
                        component=mp.Ez)]
+
+line_source = [mp.Source(src=mp.GaussianSource(fcen,fwidth=df),
+                       center=mp.Vector3(-0.85,0),
+                       size=mp.Vector3(0,sxy-2*dpml),
+                       component=mp.Ez)]
 
 k_point = mp.Vector3(0.23,-0.38)
 
@@ -81,6 +86,12 @@ def forward_simulation(design_params, mon_type, frequencies=None, mat2=silicon):
 
     if not frequencies:
         frequencies = [fcen]
+
+    if mon_type.name == 'LDOS':
+        sim.change_sources(line_source)
+        sim.run(mp.dft_ldos(frequencies), until_after_sources=mp.stop_when_energy_decayed(dt=50.0, decay_by = 1e-11))
+        sim.reset_meep()
+        return np.array(sim.ldos_data)
 
     if mon_type.name == 'EIGENMODE':
         mode = sim.add_mode_monitor(frequencies,
@@ -159,6 +170,12 @@ def J(mode_mon):
 
         def J(mode_mon):
             return npa.power(npa.abs(mode_mon[:,4,10]),2)
+    elif mon_type.name == 'LDOS':
+        sim.change_sources(line_source)
+        obj_list = [mpa.LDOS(sim)]
+
+        def J(mode_mon):
+            return npa.array(mode_mon)
 
     opt = mpa.OptimizationProblem(
         simulation=sim,
@@ -410,6 +427,34 @@ def test_adjoint_solver_eigenmode(self):
             print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
             tol = 0.04 if mp.is_single_precision() else 0.01
             self.assertClose(adj_scale,fd_grad,epsilon=tol)
+
+
+    def test_adjoint_solver_LDOS(self):
+        print("*** TESTING LDOS ADJOINT ***")
+
+        ## test the single frequency and multi frequency cases
+        for frequencies in [[fcen], [1/1.58, fcen, 1/1.53]]:
+            ## compute gradient using adjoint solver
+            adjsol_obj, adjsol_grad = adjoint_solver(p, MonitorObject.LDOS, frequencies)
+
+            ## compute unperturbed |Ez|^2
+            Ez2_unperturbed = forward_simulation(p, MonitorObject.LDOS, frequencies)
+
+            ## compare objective results
+            print("|Ez|^2 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,Ez2_unperturbed))
+            self.assertClose(adjsol_obj,Ez2_unperturbed,epsilon=2e-5)
+
+            ## compute perturbed Ez2
+            Ez2_perturbed = forward_simulation(p+dp, MonitorObject.LDOS, frequencies)
+
+            ## compare gradients
+            if adjsol_grad.ndim < 2:
+                adjsol_grad = np.expand_dims(adjsol_grad,axis=1)
+            adj_scale = (dp[None,:]@adjsol_grad).flatten()
+            fd_grad = Ez2_perturbed-Ez2_unperturbed
+            print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
+            tol = 0.07 if mp.is_single_precision() else 0.006
+            self.assertClose(adj_scale,fd_grad,epsilon=tol)
 
 
     def test_gradient_backpropagation(self):

src/dft_ldos.cpp

@@ -29,6 +29,7 @@ dft_ldos::dft_ldos(double freq_min, double freq_max, int Nfreq) {
   for (int i = 0; i < Nfreq; ++i)
     Fdft[i] = Jdft[i] = 0.0;
   Jsum = 1.0;
+  saved_overall_scale = 1.0;
 }
 
 dft_ldos::dft_ldos(const std::vector<double> freq_) {
@@ -39,6 +40,7 @@ dft_ldos::dft_ldos(const std::vector<double> freq_) {
   for (size_t i = 0; i < Nfreq; ++i)
     Fdft[i] = Jdft[i] = 0.0;
   Jsum = 1.0;
+  saved_overall_scale = 1.0;
 }
 
 dft_ldos::dft_ldos(const double *freq_, size_t Nfreq) : freq(Nfreq) {
@@ -49,27 +51,28 @@ dft_ldos::dft_ldos(const double *freq_, size_t Nfreq) : freq(Nfreq) {
   for (size_t i = 0; i < Nfreq; ++i)
     Fdft[i] = Jdft[i] = 0.0;
   Jsum = 1.0;
+  saved_overall_scale = 1.0;
 }
 
 // |c|^2
 static double abs2(complex<double> c) { return real(c) * real(c) + imag(c) * imag(c); }
 
-double *dft_ldos::ldos() const {
+double *dft_ldos::ldos() {
   // we try to get the overall scale factor right (at least for a point source)
   // so that we can compare against the analytical formula for testing
   // ... in most practical cases, the scale factor won't matter because
   //     the user will compute the relative LDOS of 2 cases (e.g. LDOS/vacuum)
 
   // overall scale factor
   double Jsum_all = sum_to_all(Jsum);
-  double scale = 4.0 / pi                 // from definition of LDOS comparison to power
+  saved_overall_scale = 4.0 / pi                 // from definition of LDOS comparison to power
                  * -0.5                   // power = -1/2 Re[E* J]
                  / (Jsum_all * Jsum_all); // normalize to unit-integral current
 
   const size_t Nfreq = freq.size();
   double *sum = new double[Nfreq];
   for (size_t i = 0; i < Nfreq; ++i) /* 4/pi * work done by unit dipole */
-    sum[i] = scale * real(Fdft[i] * conj(Jdft[i])) / abs2(Jdft[i]);
+    sum[i] = saved_overall_scale * real(Fdft[i] * conj(Jdft[i])) / abs2(Jdft[i]);
   double *out = new double[Nfreq];
   sum_to_all(sum, out, Nfreq);
   delete[] sum;

src/meep.hpp

@@ -1412,15 +1412,17 @@ class dft_ldos {
   }
 
   void update(fields &f);          // to be called after each timestep
-  double *ldos() const;            // returns array of Nomega values (after last timestep)
+  double *ldos();            // returns array of Nomega values (after last timestep)
   std::complex<double> *F() const; // returns Fdft
   std::complex<double> *J() const; // returns Jdft
   std::vector<double> freq;
-
+  double overall_scale() const { return saved_overall_scale; }
+
 private:
   std::complex<double> *Fdft;  // Nomega array of field * J*(x) DFT values
   std::complex<double> *Jdft;  // Nomega array of J(t) DFT values
   double Jsum;                 // sum of |J| over all points
+  double saved_overall_scale;  // saved overall scale for adjoint calculation
 };
 
 // dft.cpp (normally created with fields::add_dft_fields)