diff --git a/python/adjoint/objective.py b/python/adjoint/objective.py index 07aeea618..fd5d5fcd5 100644 --- a/python/adjoint/objective.py +++ b/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 diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py index 707eb8beb..0b06af20b 100644 --- a/python/adjoint/optimization_problem.py +++ b/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 diff --git a/python/tests/test_adjoint_solver.py b/python/tests/test_adjoint_solver.py index 56cee371a..14cb6654f 100644 --- a/python/tests/test_adjoint_solver.py +++ b/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)