WIP:MaterialGrid

python/adjoint/__init__.py

@@ -10,7 +10,7 @@
 
 from .basis import BilinearInterpolationBasis
 
-from .optimization_problem import (OptimizationProblem, Grid)
+from .optimization_problem import (OptimizationProblem, Grid, DesignRegion)
 
 from .filter_source import FilteredSource
 

python/adjoint/optimization_problem.py

@@ -6,6 +6,28 @@
 
 Grid = namedtuple('Grid', ['x', 'y', 'z', 'w'])
 
+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_parameters(design_parameters)
+    def get_gradient(self,fields_a,fields_f,frequencies,geom_list,f):
+        # sanitize the input
+        if (fields_a.ndim < 5) or (fields_f.ndim < 5):
+            raise ValueError("Fields arrays must have 5 dimensions (x,y,z,frequency,component)")
+        num_freqs = np.array(frequencies).size
+
+        grad = np.zeros((self.num_design_params*num_freqs,)) # preallocate
+
+        # compute the gradient
+        mp._get_gradient(grad,fields_a,fields_f,self.volume,np.array(frequencies),geom_list,f) 
+
+        return grad
+
 class OptimizationProblem(object):
     """Top-level class in the MEEP adjoint module.
 
@@ -24,7 +46,7 @@ def __init__(self,
                 simulation,
                 objective_functions,
                 objective_arguments,
-                design_variables,
+                design_regions,
                 frequencies=None,
                 fcen=None,
                 df=None,
@@ -44,13 +66,12 @@ def __init__(self,
         self.objective_arguments = objective_arguments
         self.f_bank = [] # objective function evaluation history
 
-        if isinstance(design_variables, list):
-            self.design_variables = design_variables
+        if isinstance(design_regions, list):
+            self.design_regions = design_regions
         else:
-            self.design_variables = [design_variables]
+            self.design_regions = [design_regions]
 
-        self.num_design_params = [ni.num_design_params for ni in self.design_variables]
-        self.design_regions = [dr.volume for dr in self.design_variables]
+        self.num_design_params = [ni.num_design_params for ni in self.design_regions]
         self.num_design_regions = len(self.design_regions)
 
         # TODO typecheck frequency choices
@@ -158,7 +179,7 @@ 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([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr,yee_grid=False) for dr in self.design_regions]
+        self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=False) for dr in self.design_regions]
 
         # store design region voxel parameters
         self.design_grids = [Grid(*self.sim.get_array_metadata(dft_cell=drm)) for drm in self.design_region_monitors]
@@ -181,11 +202,11 @@ def forward_run(self):
             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((len(dg.x),len(dg.y),len(dg.z),3,self.nf),dtype=np.complex128) for dg in self.design_grids]
+        self.d_E = [np.zeros((len(dg.x),len(dg.y),len(dg.z),self.nf,3),dtype=np.complex128) for dg in self.design_grids]
         for nb, dgm in enumerate(self.design_region_monitors):
             for f in range(self.nf):
                 for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
-                    self.d_E[nb][:,:,:,ic,f] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
+                    self.d_E[nb][:,:,:,f,ic] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
 
         # store objective function evaluation in memory
         self.f_bank.append(self.f0)
@@ -209,24 +230,24 @@ def adjoint_run(self,objective_idx=0):
 
         # register design flux
         # TODO use yee grid directly 
-        self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr,yee_grid=False) for dr in self.design_regions]
+        self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=False) 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, self.minimum_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((len(dg.x),len(dg.y),len(dg.z),3,self.nf),dtype=np.complex128) for dg in self.design_grids])
+        self.a_E.append([np.zeros((len(dg.x),len(dg.y),len(dg.z),self.nf,3),dtype=np.complex128) for dg in self.design_grids])
         for nb, dgm in enumerate(self.design_region_monitors):
             for f in range(self.nf):
                 for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
-                    self.a_E[objective_idx][nb][:,:,:,ic,f] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
+                    self.a_E[objective_idx][nb][:,:,:,f,ic] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
 
         # update optimizer's state
         self.current_state = "ADJ"
 
     def calculate_gradient(self):
-        # Iterate through all design region bases and store gradient w.r.t. permittivity
-        self.gradient = [[2*np.sum(np.real(self.a_E[ar][nb]*self.d_E[nb]),axis=(3)) for nb in range(self.num_design_regions)] for ar in range(len(self.objective_functions))]
+        # Iterate through all design regions and calculate gradient
+        self.gradient = [[dr.get_gradient(self.a_E[ar][dri],self.d_E[dri],self.frequencies,self.sim.geometry,self.sim.fields) for dri, dr in enumerate(self.design_regions)] for ar in range(len(self.objective_functions))]
 
         # Cleanup list of lists
         if len(self.gradient) == 1:
@@ -276,15 +297,15 @@ def calculate_fd_gradient(self,num_gradients=1,db=1e-4,design_variables_idx=0,fi
         for k in fd_gradient_idx:
 
             b0 = np.ones((self.num_design_params[design_variables_idx],))
-            b0[:] = (self.design_variables[design_variables_idx].rho_vector)
+            b0[:] = (self.design_regions[design_variables_idx].design_parameters.design_parameters)
             # -------------------------------------------- #
             # left function evaluation
             # -------------------------------------------- #
             self.sim.reset_meep()
 
             # assign new design vector
             b0[k] -= db
-            self.design_variables[design_variables_idx].set_rho_vector(b0)
+            self.design_regions[design_variables_idx].update_design_parameters(b0)
 
             # initialize design monitors
             self.forward_monitors = []
@@ -306,7 +327,7 @@ def calculate_fd_gradient(self,num_gradients=1,db=1e-4,design_variables_idx=0,fi
 
             # assign new design vector
             b0[k] += 2*db # central difference rule...
-            self.design_variables[design_variables_idx].set_rho_vector(b0)
+            self.design_regions[design_variables_idx].update_design_parameters(b0)
 
             # initialize design monitors
             self.forward_monitors = []
@@ -338,10 +359,10 @@ def update_design(self, rho_vector):
 
         rho_vector ....... a list of numpy arrays that maps to each design region
         """
-        for bi, b in enumerate(self.design_variables):
+        for bi, b in enumerate(self.design_regions):
             if np.array(rho_vector[bi]).ndim > 1:
                 raise ValueError("Each vector of design variables must contain only one dimension.")
-            b.set_rho_vector(rho_vector[bi])
+            b.update_design_parameters(rho_vector[bi])
 
         self.sim.reset_meep()
         self.current_state = "INIT"

python/geom.py

@@ -278,6 +278,41 @@ def _get_epsmu(self, diag, offdiag, susceptibilities, conductivity_diag, conduct
         # Convert list matrix to 3D numpy array size [freqs,3,3]
         return np.squeeze(epsmu)
 
+class MaterialGrid(object):
+    def __init__(self,grid_size,medium1,medium2,design_parameters=None,grid_type="U_DEFAULT"):
+        self.grid_size = grid_size
+        self.medium1 = medium1
+        self.medium2 = medium2
+        if self.grid_size.x == 0:
+            self.grid_size.x = 1
+        elif self.grid_size.y == 0:
+            self.grid_size.y = 1
+        elif self.grid_size.z == 0:
+            self.grid_size.z = 1
+        self.num_params=int(self.grid_size.x*self.grid_size.y*self.grid_size.z)
+
+        if design_parameters is None:
+            self.design_parameters = np.zeros((self.num_params,))
+        elif design_parameters.size != self.num_params:
+            raise ValueError("design_parameters of shape {} do not match user specified grid dimension: {}".format(design_parameters.size,grid_size))
+        else:
+            self.design_parameters = design_parameters.flatten().astype(np.float64)
+
+        grid_type_dict = {
+            "U_MIN":0,
+            "U_PROD":1,
+            "U_SUM":2,
+            "U_DEFAULT":3
+        }
+        if grid_type not in grid_type_dict:
+            raise ValueError("Invalid grid_type: {}. Must be either U_MIN, U_PROD, U_SUM, or U_DEFAULT".format(grid_type_dict))
+        self.grid_type = grid_type_dict[grid_type]
+
+        self.swigobj = None
+    def update_parameters(self,x):
+        if x.size != self.num_params:
+            raise ValueError("design_parameters of shape {} do not match user specified grid dimension: {}".format(design_parameters.size,grid_size))
+        self.design_parameters[:]=x.flatten().astype(np.float64)
 
 class Susceptibility(object):
 

python/meep.i

@@ -679,6 +679,7 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
         if (((material_data *)$1.material)->medium.H_susceptibilities.items) {
             delete[] ((material_data *)$1.material)->medium.H_susceptibilities.items;
         }
+        delete[] ((material_data *)$1.material)->design_parameters;
         delete[] ((material_data *)$1.material)->epsilon_data;
         delete (material_data *)$1.material;
         geometric_object_destroy($1);
@@ -720,6 +721,7 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
             delete[] ((material_data *)$1.items[i].material)->medium.H_susceptibilities.items;
         }
         delete[] ((material_data *)$1.items[i].material)->epsilon_data;
+        delete[] ((material_data *)$1.items[i].material)->design_parameters;
         delete (material_data *)$1.items[i].material;
         geometric_object_destroy($1.items[i]);
     }
@@ -750,6 +752,77 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
     delete[] $1.items;
 }
 
+//--------------------------------------------------
+// typemaps needed for material grid
+//--------------------------------------------------
+
+%inline %{
+void _get_gradient(PyObject *grad, PyObject *fields_a, PyObject *fields_f, PyObject *grid_volume, PyObject *frequencies, PyObject *py_geom_list, PyObject *f) {
+    // clean the gradient array
+    PyArrayObject *pao_grad = (PyArrayObject *)grad;
+    if (!PyArray_Check(pao_grad)) meep::abort("grad parameter must be numpy array.");
+    if (!PyArray_ISCARRAY(pao_grad)) meep::abort("Numpy grad array must be C-style contiguous.");
+    meep::realnum * grad_c = (meep::realnum *)PyArray_DATA(pao_grad);
+
+    // clean the adjoint fields array
+    PyArrayObject *pao_fields_a = (PyArrayObject *)fields_a;
+    if (!PyArray_Check(pao_fields_a)) meep::abort("fields parameter must be numpy array.");
+    if (!PyArray_ISCARRAY(pao_fields_a)) meep::abort("Numpy fields array must be C-style contiguous.");
+    std::complex<double> * fields_a_c = (std::complex<double> *)PyArray_DATA(pao_fields_a);
+
+    // clean the forward fields array
+    PyArrayObject *pao_fields_f = (PyArrayObject *)fields_f;
+    if (!PyArray_Check(pao_fields_f)) meep::abort("fields parameter must be numpy array.");
+    if (!PyArray_ISCARRAY(pao_fields_f)) meep::abort("Numpy fields array must be C-style contiguous.");
+    std::complex<double> * fields_f_c = (std::complex<double> *)PyArray_DATA(pao_fields_f);
+
+    // scalegrad not currently used
+    double scalegrad = 1.0;
+
+    // clean the volume object
+    void* where;
+
+    PyObject* swigobj = PyObject_GetAttrString(grid_volume, "swigobj");
+    SWIG_ConvertPtr(swigobj,&where,NULL,NULL);
+    const meep::volume* where_vol = (const meep::volume*)where;
+
+    // clean the frequencies array
+    PyArrayObject *pao_freqs = (PyArrayObject *)frequencies;
+    if (!PyArray_Check(pao_freqs)) meep::abort("frequencies parameter must be numpy array.");
+    if (!PyArray_ISCARRAY(pao_freqs)) meep::abort("Numpy fields array must be C-style contiguous.");
+    meep::realnum* frequencies_c = (meep::realnum *)PyArray_DATA(pao_freqs);
+
+    // get the proper dimensions of the fields array
+    if (PyArray_NDIM(pao_fields_a) !=5) {meep::abort("Fields array must have 5 dimensions.");}
+    int fdims_c[4];
+    for (int i = 0; i < PyArray_NDIM(pao_fields_a); ++i) {
+      fdims_c[i] = PyArray_DIMS(pao_fields_a)[i];
+    }
+
+    // prepare a geometry_tree
+    geometric_object_list *l;
+    l = new geometric_object_list();
+    if (!py_list_to_gobj_list(py_geom_list,l)) meep::abort("Unable to convert geometry tree.");
+
+    geom_box_tree geometry_tree = calculate_tree(*where_vol,*l);
+
+    // clean the fields pointer
+    void* f_v;
+    SWIG_ConvertPtr(f,&f_v,NULL,NULL);
+    meep::fields* f_c = (meep::fields *)f_v;
+
+    // calculate the gradient
+    meep_geom::material_grids_addgradient(grad_c,fields_a_c,fields_f_c,frequencies_c,fdims_c,scalegrad,*where_vol,geometry_tree,f_c);
+
+    destroy_geom_box_tree(geometry_tree);
+    delete[] l;
+
+}
+%}
+//--------------------------------------------------
+// end typemaps needed for material grid
+//--------------------------------------------------
+
 // Typemap suite for sources
 
 %typecheck(SWIG_TYPECHECK_POINTER) const meep::src_time & {
@@ -841,6 +914,7 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
     if ($1->medium.H_susceptibilities.items) {
         delete[] $1->medium.H_susceptibilities.items;
     }
+    delete[] $1->design_parameters;
     delete[] $1->epsilon_data;
     delete $1;
 }
@@ -1190,6 +1264,7 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
             if ($1.items[i]->medium.H_susceptibilities.items) {
                 delete[] $1.items[i]->medium.H_susceptibilities.items;
             }
+            delete[] $1.items[i]->design_parameters;
             delete[] $1.items[i]->epsilon_data;
         }
         delete[] $1.items;
@@ -1404,6 +1479,7 @@ PyObject *_get_array_slice_dimensions(meep::fields *f, const meep::volume &where
         GyrotropicSaturatedSusceptibility,
         Lattice,
         LorentzianSusceptibility,
+        MaterialGrid,
         Matrix,
         Medium,
         MultilevelAtom,

python/simulation.py

@@ -662,6 +662,7 @@ def __init__(self,
         self.collect_stats = collect_stats
         self.fragment_stats = None
         self._output_stats = os.environ.get('MEEP_STATS', None)
+        self.material_grids = [] # Needed to store array of c pointers for each grid
 
         self.special_kz = False
         if self.cell_size.z == 0 and self.k_point and self.k_point.z != 0: