Use fragment_stats to split chunks by cost (#681)

* Utilize fragment_stats to determine chunks

* reformat

* split_by_cost method working with factors of 3 and 2

* Move static functions into grid_volume

* Better split measure in split_into_three

libpympb/pympb.cpp

@@ -2287,7 +2287,7 @@ double mode_solver::compute_energy_in_objects(geometric_object_list objects) {
     return 0.0;
   }
 
-  geom_fix_objects0(objects);
+  geom_fix_object_list(objects);
 
   int n1 = mdata->nx;
   int n2 = mdata->ny;

python/meep.i

@@ -787,6 +787,9 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
     delete $1;
 }
 
+// For some reason SWIG needs the namespaced version too
+%apply material_type { meep_geom::material_type };
+
 // Typemap suite for get_array_metadata
 
 %typecheck(SWIG_TYPECHECK_POINTER, fragment="NumPy_Fragments") double* xtics {
@@ -1111,6 +1114,9 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
     }
 }
 
+// For some reason SWIG needs the namespaced version too
+%apply material_type_list { meep_geom::material_type_list };
+
 // Typemap suite for custom_src_time
 
 %typecheck(SWIG_TYPECHECK_POINTER) (std::complex<double> (*func)(double t, void *), void *data) {
@@ -1479,3 +1485,56 @@ PyObject *_get_array_slice_dimensions(meep::fields *f, const meep::volume &where
     import atexit
     atexit.register(report_elapsed_time)
 %}
+
+%inline %{
+meep::structure *create_structure_and_set_materials(vector3 cell_size,
+                                                    std::vector<meep_geom::dft_data> dft_data_list_,
+                                                    std::vector<meep::volume> pml_1d_vols_,
+                                                    std::vector<meep::volume> pml_2d_vols_,
+                                                    std::vector<meep::volume> pml_3d_vols_,
+                                                    std::vector<meep::volume> absorber_vols_,
+                                                    meep::grid_volume &gv,
+                                                    const meep::boundary_region &br,
+                                                    const meep::symmetry &sym,
+                                                    int num_chunks,
+                                                    double Courant,
+                                                    bool use_anisotropic_averaging,
+                                                    double tol,
+                                                    int maxeval,
+                                                    geometric_object_list gobj_list,
+                                                    vector3 center,
+                                                    bool _ensure_periodicity,
+                                                    bool verbose,
+                                                    meep_geom::material_type _default_material,
+                                                    meep_geom::absorber_list alist,
+                                                    meep_geom::material_type_list extra_materials,
+                                                    bool split_chunks_evenly) {
+    // Initialize fragment_stats static members (used for creating chunks in choose_chunkdivision)
+    meep_geom::fragment_stats::geom = gobj_list;
+    meep_geom::fragment_stats::dft_data_list = dft_data_list_;
+    meep_geom::fragment_stats::pml_1d_vols = pml_1d_vols_;
+    meep_geom::fragment_stats::pml_2d_vols = pml_2d_vols_;
+    meep_geom::fragment_stats::pml_3d_vols = pml_3d_vols_;
+    meep_geom::fragment_stats::absorber_vols = absorber_vols_;
+    meep_geom::fragment_stats::tol = tol;
+    meep_geom::fragment_stats::maxeval = maxeval;
+    meep_geom::fragment_stats::resolution = gv.a;
+    meep_geom::fragment_stats::dims = gv.dim;
+    meep_geom::fragment_stats::split_chunks_evenly = split_chunks_evenly;
+    meep_geom::fragment_stats::init_libctl(_default_material, _ensure_periodicity,
+                                           &gv, cell_size, center, &gobj_list);
+
+    meep::structure *s = new meep::structure(gv, NULL, br, sym, num_chunks, Courant,
+                                             use_anisotropic_averaging, tol, maxeval);
+    s->shared_chunks = true;
+
+    meep_geom::set_materials_from_geometry(s, gobj_list, center, use_anisotropic_averaging, tol, maxeval,
+                                           _ensure_periodicity, verbose, _default_material, alist, extra_materials);
+
+    // Return params to default state
+    meep_geom::fragment_stats::resolution = 0;
+    meep_geom::fragment_stats::split_chunks_evenly = false;
+
+    return s;
+}
+%}

python/simulation.py

@@ -471,7 +471,8 @@ def __init__(self,
                  output_single_precision=False,
                  load_structure='',
                  geometry_center=mp.Vector3(),
-                 force_all_components=False):
+                 force_all_components=False,
+                 split_chunks_evenly=True):
 
         self.cell_size = cell_size
         self.geometry = geometry
@@ -518,6 +519,7 @@ def __init__(self,
         self._is_initialized = False
         self._fragment_size = 10
         self.force_all_components = force_all_components
+        self.split_chunks_evenly = split_chunks_evenly
 
     # To prevent the user from having to specify `dims` and `is_cylindrical`
     # to Volumes they create, the library will adjust them appropriately based
@@ -823,7 +825,7 @@ def _boundary_layers_to_vol_list(self, boundaries):
 
         return vols1, vols2, vols3
 
-    def _compute_fragment_stats(self, gv):
+    def _make_fragment_lists(self, gv):
 
         def convert_volumes(dft_obj):
             volumes = []
@@ -847,6 +849,12 @@ def convert_volumes(dft_obj):
         absorber_vols1, absorber_vols2, absorber_vols3 = self._boundary_layers_to_vol_list(absorbers)
         absorber_vols = absorber_vols1 + absorber_vols2 + absorber_vols3
 
+        return (dft_data_list, pml_vols1, pml_vols2, pml_vols3, absorber_vols)
+
+    def _compute_fragment_stats(self, gv):
+
+        dft_data_list, pml_vols1, pml_vols2, pml_vols3, absorber_vols = self._make_fragment_lists(gv)
+
         stats = mp.compute_fragment_stats(
             self.geometry,
             gv,
@@ -886,7 +894,6 @@ def _init_structure(self, k=False):
         gv = self._create_grid_volume(k)
         sym = self._create_symmetries(gv)
         br = _create_boundary_region_from_boundary_layers(self.boundary_layers, gv)
-
         absorbers = [bl for bl in self.boundary_layers if type(bl) is Absorber]
 
         if self.material_function:
@@ -899,22 +906,33 @@ def _init_structure(self, k=False):
             self.default_material = self.epsilon_input_file
 
         self.fragment_stats = self._compute_fragment_stats(gv) if isinstance(self.default_material, mp.Medium) else []
+        dft_data_list, pml_vols1, pml_vols2, pml_vols3, absorber_vols = self._make_fragment_lists(gv)
+
+        self.structure = mp.create_structure_and_set_materials(
+            self.cell_size,
+            dft_data_list,
+            pml_vols1,
+            pml_vols2,
+            pml_vols3,
+            absorber_vols,
+            gv,
+            br,
+            sym,
+            self.num_chunks,
+            self.Courant,
+            self.eps_averaging,
+            self.subpixel_tol,
+            self.subpixel_maxeval,
+            self.geometry,
+            self.geometry_center,
+            self.ensure_periodicity and not not self.k_point,
+            self.verbose,
+            self.default_material,
+            absorbers,
+            self.extra_materials,
+            self.split_chunks_evenly
+       )
 
-        self.structure = mp.structure(gv, None, br, sym, self.num_chunks, self.Courant,
-                                      self.eps_averaging, self.subpixel_tol, self.subpixel_maxeval)
-        self.structure.shared_chunks = True
-
-        mp.set_materials_from_geometry(self.structure,
-                                       self.geometry,
-                                       self.geometry_center,
-                                       self.eps_averaging,
-                                       self.subpixel_tol,
-                                       self.subpixel_maxeval,
-                                       self.ensure_periodicity and not not self.k_point,
-                                       False,
-                                       self.default_material,
-                                       absorbers,
-                                       self.extra_materials)
         if self.load_structure_file:
             self.load_structure(self.load_structure_file)
 

scheme/meep.i

@@ -263,6 +263,8 @@ static meep::vec my_kpoint_func(double freq, int mode, void *user_data) {
 
 %warnfilter(302,325,451,503,509);
 
+%ignore meep_geom::fragment_stats;
+
 %include "meep_renames.i"
 %include "meep_enum_renames.i"
 %include "meep_op_renames.i"

src/meep/vec.hpp

@@ -20,6 +20,7 @@
 #define MEEP_VEC_H
 
 #include <complex>
+#include <vector>
 #include <stddef.h>
 
 namespace meep {
@@ -987,7 +988,11 @@ class grid_volume {
   grid_volume split(size_t num, int which) const;
   grid_volume split_by_effort(int num, int which, int Ngv = 0, const grid_volume *v = NULL,
                               double *effort = NULL) const;
-  grid_volume split_at_fraction(bool want_high, int numer) const;
+  grid_volume split_by_cost(int desired_num_chunks, int proc_num) const;
+  std::vector<grid_volume> split_into_n(int n) const;
+  void split_into_three(std::vector<grid_volume> &result) const;
+  grid_volume split_at_fraction(bool want_high, int numer, int bestd = -1, int bestlen = 1) const;
+  double get_cost() const;
   grid_volume halve(direction d) const;
   void pad_self(direction d);
   grid_volume pad(direction d) const;