Anisotropic material gradient support

python/adjoint/__init__.py

@@ -5,18 +5,19 @@
 
 from .objective import *
 
+from . import utils
+from .utils import DesignRegion
+
 from .basis import BilinearInterpolationBasis
 
-from .optimization_problem import (OptimizationProblem, Grid, DesignRegion)
+from .optimization_problem import (OptimizationProblem)
 
 from .filter_source import FilteredSource
 
 from .filters import *
 
 from .connectivity import *
 
-from . import utils
-
 try:
     from .wrapper import MeepJaxWrapper
 except ModuleNotFoundError as _:

python/adjoint/objective.py

@@ -4,9 +4,10 @@
 import numpy as np
 import meep as mp
 from .filter_source import FilteredSource
-from .optimization_problem import Grid
 from meep.simulation import py_v3_to_vec
+from collections import namedtuple
 
+Grid = namedtuple('Grid', ['x', 'y', 'z', 'w'])
 
 class ObjectiveQuantity(abc.ABC):
     """A differentiable objective quantity.

python/adjoint/optimization_problem.py

@@ -3,48 +3,7 @@
 from autograd import grad, jacobian
 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(),
-            MaterialGrid=None,
-    ):
-        self.volume = volume if volume else mp.Volume(center=center, size=size)
-        self.size = self.volume.size
-        self.center = self.volume.center
-        self.design_parameters = design_parameters
-        self.num_design_params = design_parameters.num_params
-        self.MaterialGrid = MaterialGrid
-
-    def update_design_parameters(self, design_parameters):
-        self.design_parameters.update_weights(design_parameters)
-
-    def get_gradient(self, sim, fields_a, fields_f, frequencies):
-        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((num_freqs, self.num_design_params))  # preallocate
-
-        geom_list = sim.geometry
-        f = sim.fields
-        vol = sim._fit_volume_to_simulation(self.volume)
-        # compute the gradient
-        sim_is_cylindrical = (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical
-        mp._get_gradient(grad,fields_a,fields_f,vol,np.array(frequencies),geom_list,f, sim_is_cylindrical)
-
-        return np.squeeze(grad).T
-
+from . import utils, DesignRegion
 
 class OptimizationProblem(object):
     """Top-level class in the MEEP adjoint module.
@@ -76,9 +35,6 @@ def __init__(
     ):
 
         self.sim = simulation
-        self.components = [mp.Ex,mp.Ey,mp.Ez]
-        if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
-            self.components = [mp.Er,mp.Ep,mp.Ez]
 
         if isinstance(objective_functions, list):
             self.objective_functions = objective_functions
@@ -161,13 +117,13 @@ def __call__(self, rho_vector=None, need_value=True, need_gradient=True):
                 print("Starting forward run...")
                 self.forward_run()
                 print("Starting adjoint run...")
-                self.a_E = []
+                self.D_a = []
                 self.adjoint_run()
                 print("Calculating gradient...")
                 self.calculate_gradient()
             elif self.current_state == "FWD":
                 print("Starting adjoint run...")
-                self.a_E = []
+                self.D_a = []
                 self.adjoint_run()
                 print("Calculating gradient...")
                 self.calculate_gradient()
@@ -210,28 +166,9 @@ 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(
-                self.components,
-                self.frequencies,
-                where=dr.volume,
-                yee_grid=True,
-                decimation_factor=self.decimation_factor,
-            ) for dr in self.design_regions
-        ]
-
-        # store design region voxel parameters
-        self.design_grids = []
-        if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
-            YeeDims = namedtuple('YeeDims', ['Er','Ep','Ez'])
-        else:
-            YeeDims = namedtuple('YeeDims', ['Ex','Ey','Ez'])
-        for drm in self.design_region_monitors:
-            s = [
-                self.sim.get_array_slice_dimensions(c, vol=drm.where)[0]
-                for c in self.components
-            ]
-            self.design_grids += [YeeDims(*s)]
+        self.design_region_monitors = utils.install_design_region_monitors(
+            self.sim, self.design_regions, self.frequencies, self.decimation_factor
+        )
 
     def forward_run(self):
         # set up monitors
@@ -254,16 +191,8 @@ def forward_run(self):
         if len(self.f0) == 1:
             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((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(self.components):
-                for f in range(self.nf):
-                    self.d_E[nb][ic][f, :, :, :] = atleast_3d(
-                        self.sim.get_dft_array(dgm, c, f))
+        # Store forward fields for each set of design variables in array
+        self.D_f = utils.gather_design_region_fields(self.sim,self.design_region_monitors,self.frequencies)
 
         # store objective function evaluation in memory
         self.f_bank.append(self.f0)
@@ -288,11 +217,14 @@ def adjoint_run(self):
         # set up adjoint sources and monitors
         self.prepare_adjoint_run()
 
-        if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
+        # flip the m number
+        if utils._check_if_cylindrical(self.sim):
             self.sim.m = -self.sim.m
 
+        # flip the k point
         if self.sim.k_point:
             self.sim.k_point *= -1
+
         for ar in range(len(self.objective_functions)):
             # Reset the fields
             self.sim.reset_meep()
@@ -301,15 +233,9 @@ def adjoint_run(self):
             self.sim.change_sources(self.adjoint_sources[ar])
 
             # register design flux
-            self.design_region_monitors = [
-                self.sim.add_dft_fields(
-                    self.components,
-                    self.frequencies,
-                    where=dr.volume,
-                    yee_grid=True,
-                    decimation_factor=self.decimation_factor,
-                ) for dr in self.design_regions
-            ]
+            self.design_region_monitors = utils.install_design_region_monitors(
+            self.sim, self.design_regions, self.frequencies, self.decimation_factor
+            )
 
             # Adjoint run
             self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
@@ -318,34 +244,26 @@ def adjoint_run(self):
                 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(self.components):
-                    for f in range(self.nf):
-                        if (self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL):
-                            # Addtional factor of 2 for cyldrical coordinate
-                            self.a_E[ar][nb][ic][f, :, :, :] = 2 * atleast_3d(self.sim.get_dft_array(dgm, c, f))
-                        else:
-                            self.a_E[ar][nb][ic][f, :, :, :] = atleast_3d(self.sim.get_dft_array(dgm, c, f))
-
-        if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
+            # Store adjoint fields for each design set of design variables
+            self.D_a.append(utils.gather_design_region_fields(self.sim,self.design_region_monitors,self.frequencies))
+
+        # reset the m number
+        if utils._check_if_cylindrical(self.sim):
             self.sim.m = -self.sim.m
 
-        # update optimizer's state
+        # reset the k point
         if self.sim.k_point: self.sim.k_point *= -1
+
+        # update optimizer's state
         self.current_state = "ADJ"
 
     def calculate_gradient(self):
         # Iterate through all design regions and calculate gradient
         self.gradient = [[
             dr.get_gradient(
                 self.sim,
-                self.a_E[ar][dri],
-                self.d_E[dri],
+                self.D_a[ar][dri],
+                self.D_f[dri],
                 self.frequencies,
             ) for dri, dr in enumerate(self.design_regions)
         ] for ar in range(len(self.objective_functions))]
@@ -486,7 +404,7 @@ def calculate_fd_gradient(
 
         return fd_gradient, fd_gradient_idx
 
-    def update_design(self, rho_vector):
+    def update_design(self, rho_vector, beta=None):
         """Update the design permittivity function.
 
         rho_vector ....... a list of numpy arrays that maps to each design region
@@ -497,6 +415,8 @@ def update_design(self, rho_vector):
                     "Each vector of design variables must contain only one dimension."
                 )
             b.update_design_parameters(rho_vector[bi])
+            if beta:
+                b.update_beta(beta)
 
         self.sim.reset_meep()
         self.current_state = "INIT"

python/adjoint/utils.py

@@ -3,14 +3,93 @@
 import meep as mp
 import numpy as onp
 
-from . import ObjectiveQuantity, DesignRegion
+from . import ObjectiveQuantity
 
 # Meep field components used to compute adjoint sensitivities
-_ADJOINT_FIELD_COMPONENTS = [mp.Ex, mp.Ey, mp.Ez]
+_ADJOINT_FIELD_COMPONENTS = [mp.Dx, mp.Dy, mp.Dz]
+_ADJOINT_FIELD_COMPONENTS_CYL = [mp.Dr, mp.Dp, mp.Dz]
 
 # The frequency axis in the array returned by `mp._get_gradient()`
 _GRADIENT_FREQ_AXIS = 1
 
+# The returned axis order from get_dft_array in cylindrical coordinates
+_FREQ_AXIS = 0
+_RHO_AXIS = 2
+_PHI_AXIS = 3
+_Z_AXIS = 1
+
+class DesignRegion(object):
+    def __init__(
+            self,
+            design_parameters,
+            volume=None,
+            size=None,
+            center=mp.Vector3()
+    ):
+        self.volume = volume if volume else mp.Volume(center=center, size=size)
+        self.size = self.volume.size
+        self.center = self.volume.center
+        self.design_parameters = design_parameters
+        self.num_design_params = design_parameters.num_params
+
+    def update_design_parameters(self, design_parameters):
+        self.design_parameters.update_weights(design_parameters)
+
+    def update_beta(self,beta):
+        self.design_parameters.beta=beta
+
+    def get_gradient(self, sim, fields_a, fields_f, frequencies):
+        num_freqs = onp.array(frequencies).size
+        shapes = []
+        '''We have the option to linear scale the gradients up front
+        using the scalegrad parameter (leftover from MPB API). Not
+        currently needed for any existing feature (but available for
+        future use)'''
+        scalegrad = 1
+        for component_idx, component in enumerate(_compute_components(sim)):
+            '''We need to correct for the rare cases that get_dft_array
+            returns a singleton element for the forward and adjoint fields.
+            This only occurs when we are in 2D and only working with a particular
+            polarization (as the other fields are never stored). For example, the
+            2D in-plane polarization consists of a single scalar Ez field 
+            (technically, meep doesn't store anything for these cases, but get_dft_array
+            still returns 0).
+
+            Our get_gradient algorithm, however, requires we pass an array of 
+            zeros with the proper shape as the design_region.'''
+            spatial_shape = sim.get_array_slice_dimensions(component, vol=self.volume)[0]
+            if (fields_a[component_idx][0,...].size == 1):
+                fields_a[component_idx] = onp.zeros(onp.insert(spatial_shape,0,num_freqs))
+                fields_f[component_idx] = onp.zeros(onp.insert(spatial_shape,0,num_freqs))
+            if _check_if_cylindrical(sim):
+                '''For some reason, get_dft_array returns the field
+                components in a different order than the convention used
+                throughout meep. So, we reorder them here so we can use
+                the same field macros later in our get_gradient function.
+                '''
+                fields_a[component_idx] = onp.transpose(fields_a[component_idx],(_FREQ_AXIS,_RHO_AXIS,_PHI_AXIS,_Z_AXIS))
+                fields_f[component_idx] = onp.transpose(fields_f[component_idx],(_FREQ_AXIS,_RHO_AXIS,_PHI_AXIS,_Z_AXIS))
+            shapes.append(fields_a[component_idx].shape)
+            fields_a[component_idx] = fields_a[component_idx].flatten(order='C')
+            fields_f[component_idx] = fields_f[component_idx].flatten(order='C')
+        shapes = onp.asarray(shapes).flatten(order='C')
+        fields_a = onp.concatenate(fields_a)
+        fields_f = onp.concatenate(fields_f)
+
+        grad = onp.zeros((num_freqs, self.num_design_params))  # preallocate
+        geom_list = sim.geometry
+        f = sim.fields
+        vol = sim._fit_volume_to_simulation(self.volume)
+
+        # compute the gradient
+        mp._get_gradient(grad,scalegrad,fields_a,fields_f,sim.gv,vol.swigobj,onp.array(frequencies),sim.geps,shapes)
+        return onp.squeeze(grad).T
+
+def _check_if_cylindrical(sim):
+    return sim.is_cylindrical or (sim.dimensions == mp.CYLINDRICAL)
+
+def _compute_components(sim):
+    return _ADJOINT_FIELD_COMPONENTS_CYL if _check_if_cylindrical(sim) else _ADJOINT_FIELD_COMPONENTS
 
 def _make_at_least_nd(x: onp.ndarray, dims: int = 3) -> onp.ndarray:
     """Makes an array have at least the specified number of dimensions."""
@@ -61,15 +140,16 @@ def install_design_region_monitors(
     simulation: mp.Simulation,
     design_regions: List[DesignRegion],
     frequencies: List[float],
+    decimation_factor: int = 0
 ) -> List[mp.DftFields]:
     """Installs DFT field monitors at the design regions of the simulation."""
     design_region_monitors = [
         simulation.add_dft_fields(
-            _ADJOINT_FIELD_COMPONENTS,
+            _compute_components(simulation),
             frequencies,
             where=design_region.volume,
             yee_grid=True,
-            decimation_factor=0
+            decimation_factor=decimation_factor
         ) for design_region in design_regions
     ]
     return design_region_monitors
@@ -120,7 +200,7 @@ def gather_design_region_fields(
     fwd_fields = []
     for monitor in design_region_monitors:
         fields_by_component = []
-        for component in _ADJOINT_FIELD_COMPONENTS:
+        for component in _compute_components(simulation):
             fields_by_freq = []
             for freq_idx, _ in enumerate(frequencies):
                 fields = simulation.get_dft_array(monitor, component, freq_idx)