Poynting flux Adjoint

python/adjoint/objective.py

@@ -1,7 +1,7 @@
 """Handling of objective functions and objective quantities."""
 import abc
 from collections import namedtuple
-
+from autograd import numpy as npa
 import numpy as np
 from meep.simulation import py_v3_to_vec
 
@@ -11,6 +11,26 @@
 
 Grid = namedtuple("Grid", ["x", "y", "z", "w"])
 
+# 3 possible components for E x n and H x n
+# signs are handled in code
+EH_TRANSVERSE = [
+    [mp.Hz, mp.Hy, mp.Ez, mp.Ey],
+    [mp.Hx, mp.Hz, mp.Ex, mp.Ez],
+    [mp.Hy, mp.Hx, mp.Ey, mp.Ex],
+]
+
+# Holds the components for each current source
+# for the cases of x,y, and z normal vectors.
+# This is the same as swapping H and E in the above list
+JK_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],
+]
+
+# Holds the amplitudes used in Poynting Flux adjoint sources
+FLUX_AMPLITUDES = np.array([1 / 4, -1 / 4, -1 / 4, 1 / 4], dtype=np.complex128)
+
 
 class ObjectiveQuantity(abc.ABC):
     """A differentiable objective quantity.
@@ -483,3 +503,145 @@ def __call__(self):
         self.ldos_scale = self.sim.ldos_scale
         self.ldos_Jdata = self.sim.ldos_Jdata
         return np.array(self._eval)
+
+
+class PoyntingFlux(ObjectiveQuantity):
+    """A frequency-dependent Poynting Flux adjoint source.
+    Attributes:
+        volume: The volume over which the Poynting Flux is calculated.
+        This function currently only works for monitors with a defined
+        normal vector (e.g. planes in 3d or lines in 2d). User supplied
+        normal vectors may be implemented in the future. It also only
+        works with monitors aligned to a coordinate direction.
+        decimation_factor: Whether to skip points in the time series every
+        decimation_factor timesteps. See "add_dft_fields" documentation.
+        The biggest warning there is to be careful to avoid aliasing if
+        the fields vary quickly in time.
+        Note on yee_grid: For the Poynting Flux to work, H and E components
+        must lie at the same points. Therefore, the Yee grid will always be false.
+    """
+
+    def __init__(self, sim, volume, scale=1, decimation_factor=0):
+        super().__init__(sim)
+        # _fit_volume_to_simulation increases the dimensionality of
+        # the volume, so we'll use the user input volume
+        self.volume = sim._fit_volume_to_simulation(volume)
+        self.decimation_factor = decimation_factor
+        self.scale = scale
+        # get_normal returns an index for the two
+        # dictionaries of cross products
+        self.normal = self.get_normal(volume)
+
+    def register_monitors(self, frequencies):
+        self._frequencies = np.asarray(frequencies)
+        self._monitor = []
+        # List to hold FourierFields objects
+        self.F_fields_list = []
+        for comp in EH_TRANSVERSE[self.normal]:
+            # instantiate the FourierFields monitors
+            F_field = FourierFields(self.sim, self.volume, comp)
+            self.F_fields_list.append(F_field)
+            self._monitor.append(F_field.register_monitors(self._frequencies))
+        return self._monitor
+
+    def place_adjoint_source(self, dJ):
+        source = []
+        print("This is dJ[0]'s shape:")
+        print(np.array(dJ[0]).shape)
+        squeezed_dJ_0 = np.array(dJ[0]).squeeze()
+        print("This is squeezed dj0 shape:")
+        print(squeezed_dJ_0.shape)
+
+        print("This is metadata's shape:")
+        print(self.field_component_evaluations[4].shape)
+        for pos, field in enumerate(self.F_fields_list):
+            # Make sure there's a nonzero value in the gradient
+            # (zero sources don't converge)
+            # Check is also in prepare_adjoint_run,
+            # but necessary here too since the source is a vector
+            if np.any(dJ[pos]):
+                reshaped_dJ = np.reshape(
+                    np.array(dJ[pos]).squeeze(),
+                    self.field_component_evaluations[4].shape,
+                )
+                # new_source = field.place_adjoint_source(reshaped_dJ)
+                new_source = field.place_adjoint_source(np.array(dJ[pos]).squeeze())
+                print("This is the new source:")
+                print(new_source)
+                print("new soruce's shape")
+                print(np.array(new_source).shape)
+                source.append(
+                    # field.place_adjoint_source(np.flipud(np.array(dJ[pos]).squeeze()))[
+                    #     0
+                    # ]
+                    new_source
+                )
+        final_array = np.array(source).flatten()
+        print("This is the final array shape:")
+        # print(final_array)
+        print(np.array(final_array).shape)
+        # print("This is the final_array with an extra array")
+
+        # test_arr = [final_array]
+        # print(test_arr)
+        # print(test_arr.shape)
+        return final_array.tolist()
+
+    def __call__(self):
+        self.field_component_evaluations = []
+        # Get integration weights Meep uses
+        self.metadata = self.sim.get_array_metadata(vol=self.volume)
+        for field in self.F_fields_list:
+            # Get the dft evaluation from a call to the underlying
+            # FourierFields object
+            field_here = field()
+            # make sure it isn't a list of scalar zeros equal to the number of
+            # frequencies (usually caused by symmetries making fields 0)
+            # fixes the np.array error in the return
+            # when we give it a "ragged" array
+            if (np.squeeze(field_here).size) == self._frequencies.size:
+                print("does the empty array check work")
+                print(field_here)
+                field_here = np.array([np.zeros(self.metadata[3].shape)])
+            self.field_component_evaluations.append(field_here)
+
+        self.field_component_evaluations.append(
+            np.array([self.metadata[3]]).astype(complex)
+        )
+        [H1, H2, E1, E2, meta] = self.field_component_evaluations
+
+        self._eval = self.field_component_evaluations
+        print("This is meta*E2")
+        print(meta * E2)
+        print("This is the np array")
+        print(np.array([H1, H2, E1, E2, meta]))
+        return np.array([H1, H2, E1, E2, meta])
+
+    # takes in a 1x5xNxM vector where the size five array corresponds to
+    # [H1,H2,E1,E1,meta]
+    # multiple frequencies will be tested later
+    @staticmethod
+    def compute_flux(*inputs):
+        # Check if all the inputs are nonzero
+        flux = npa.sum(
+            npa.real(
+                inputs[0][4]
+                * (
+                    npa.conj(inputs[0][0]) * inputs[0][3]
+                    - npa.conj(inputs[0][1]) * inputs[0][2]
+                )
+            )
+        )
+        return flux
+
+    # returns 0,1, or 2 corresponding to x, y, or z normal vectors
+    # TODO: Handle user-specified normal vectors and cases when 2d
+    # has a zero-size dimension other than z
+    def get_normal(self, volume):
+        # I'll add cylindrical later (since the normal vector gets a little different)
+        if volume.size.x == 0:
+            return 0
+        elif volume.size.y == 0:
+            return 1
+        else:
+            return 2

python/adjoint/optimization_problem.py

@@ -45,6 +45,7 @@ def __init__(
             self.objective_functions = objective_functions
         else:
             self.objective_functions = [objective_functions]
+
         self.objective_arguments = objective_arguments
         self.f_bank = []  # objective function evaluation history
 
@@ -197,6 +198,7 @@ def forward_run(self):
 
         # record objective quantities from user specified monitors
         self.results_list = [m() for m in self.objective_arguments]
+
         # evaluate objectives
         self.f0 = [fi(*self.results_list) for fi in self.objective_functions]
         if len(self.f0) == 1:
@@ -215,6 +217,8 @@ def prepare_adjoint_run(self):
             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
+                print(np.any(dJ))
+                print(dJ)
                 if np.any(dJ):
                     self.adjoint_sources[ar] += m.place_adjoint_source(
                         dJ

python/tests/test_adjoint_solver.py

@@ -13,7 +13,8 @@
 from autograd import tensor_jacobian_product
 from utils import ApproxComparisonTestCase
 
-MonitorObject = Enum("MonitorObject", "EIGENMODE DFT LDOS")
+# MonitorObject = Enum("MonitorObject", "EIGENMODE DFT LDOS POYNTING")
+MonitorObject = Enum("MonitorObject", "POYNTING")
 
 
 class TestAdjointSolver(ApproxComparisonTestCase):
@@ -23,7 +24,7 @@ def setUpClass(cls):
 
         cls.silicon = mp.Medium(epsilon=12)
         cls.sapphire = mp.Medium(
-            epsilon_diag=(10.225, 10.225, 9.95),
+            epsilon_cddiag=(10.225, 10.225, 9.95),
             epsilon_offdiag=(-0.825, -0.55 * np.sqrt(3 / 2), 0.55 * np.sqrt(3 / 2)),
         )
 
@@ -170,6 +171,17 @@ def J(mode_mon):
             def J(mode_mon):
                 return npa.array(mode_mon)
 
+        elif mon_type.name == "POYNTING":
+            obj_list = [
+                mpa.PoyntingFlux(
+                    sim, mp.Volume(center=mp.Vector3(1.25), size=mp.Vector3(0, 1, 0))
+                )
+            ]
+            # inputs contains H1, H2, E1, E2, and metadata (in that order)
+            def J(*inputs):
+                flux = obj_list[0].compute_flux(*inputs)
+                return flux
+
         opt = mpa.OptimizationProblem(
             simulation=sim,
             objective_functions=J,
@@ -332,6 +344,74 @@ def mapping(self, x, filter_radius, eta, beta):
 
         return projected_field.flatten()
 
+    def test_Poynting_Flux(self):
+        print("*** TESTING POYNTING OBJECTIVE ***")
+
+        # test the single frequency and multi frequency cases
+        for frequencies in [[self.fcen], [1 / 1.58, self.fcen, 1 / 1.53]]:
+            # compute objective value and its gradient for unperturbed design
+            unperturbed_val, unperturbed_grad = self.adjoint_solver(
+                self.p, MonitorObject.POYNTING, frequencies, need_gradient=False
+            )
+
+            # compute objective value for perturbed design
+            perturbed_val = self.adjoint_solver(
+                self.p + self.dp,
+                MonitorObject.POYNTING,
+                frequencies,
+                need_gradient=False,
+            )
+
+            # compare directional derivative
+            if unperturbed_grad.ndim < 2:
+                unperturbed_grad = np.expand_dims(unperturbed_grad, axis=1)
+            adj_dd = (self.dp[None, :] @ unperturbed_grad).flatten()
+            fnd_dd = perturbed_val - unperturbed_val
+            print(
+                f"directional derivative:, {adj_dd} (adjoint solver), {fnd_dd} (finite difference)"
+            )
+            print("\n\n\n\n\n")
+            print("")
+            print("This is the unperturbed poynting gradient:")
+            print(unperturbed_grad)
+
+            tol = 0.07 if mp.is_single_precision() else 0.006
+            self.assertClose(adj_dd, fnd_dd, epsilon=tol)
+
+    def test_eigenmode(self):
+        print("*** TESTING EIGENMODE OBJECTIVE ***")
+
+        # test the single frequency and multi frequency case
+        for frequencies in [[self.fcen], [1 / 1.58, self.fcen, 1 / 1.53]]:
+            # compute objective value and its gradient for unperturbed design
+            unperturbed_val, unperturbed_grad = self.adjoint_solver(
+                self.p, MonitorObject.EIGENMODE, frequencies
+            )
+
+            # compute objective for perturbed design
+            perturbed_val = self.adjoint_solver(
+                self.p + self.dp,
+                MonitorObject.EIGENMODE,
+                frequencies,
+                need_gradient=False,
+            )
+
+            # compare directional derivative
+            if unperturbed_grad.ndim < 2:
+                unperturbed_grad = np.expand_dims(unperturbed_grad, axis=1)
+            adj_dd = (self.dp[None, :] @ unperturbed_grad).flatten()
+            fnd_dd = perturbed_val - unperturbed_val
+            print(
+                f"directional derivative:, {adj_dd} (adjoint solver), {fnd_dd} (finite difference)"
+            )
+            print("\n\n\n\n\n")
+            print("")
+            print("This is the unperturbed eigenmode gradient:")
+            print(unperturbed_grad)
+
+            tol = 0.04 if mp.is_single_precision() else 0.01
+            self.assertClose(adj_dd, fnd_dd, epsilon=tol)
+
     def test_DFT_fields(self):
         print("*** TESTING DFT OBJECTIVE ***")
 
@@ -385,6 +465,10 @@ def test_eigenmode(self):
             print(
                 f"directional derivative:, {adj_dd} (adjoint solver), {fnd_dd} (finite difference)"
             )
+            print("\n\n\n\n\n")
+            print("")
+            print("This is the unperturbed eigenmode gradient:")
+            print(unperturbed_grad)
 
             tol = 0.04 if mp.is_single_precision() else 0.01
             self.assertClose(adj_dd, fnd_dd, epsilon=tol)