pass binary-tree partition to structure constructor

python/meep.i

@@ -1456,9 +1456,6 @@ void _get_gradient(PyObject *grad, PyObject *fields_a, PyObject *fields_f, PyObj
 
 %typemap(in) meep::binary_partition * {
     $1 = py_bp_to_bp($input);
-    if(!$1) {
-        SWIG_fail;
-    }
 }
 
 %typemap(arginit) meep::binary_partition * {
@@ -1891,7 +1888,8 @@ meep::structure *create_structure_and_set_materials(vector3 cell_size,
                                                     bool split_chunks_evenly,
                                                     bool set_materials,
                                                     meep::structure *existing_s,
-                                                    bool output_chunk_costs) {
+                                                    bool output_chunk_costs,
+                                                    meep::binary_partition *my_bp) {
     // 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_;
@@ -1909,7 +1907,11 @@ meep::structure *create_structure_and_set_materials(vector3 cell_size,
 
     if (output_chunk_costs) {
          meep::volume thev = gv.surroundings();
-         std::vector<grid_volume> chunk_vols = meep::choose_chunkdivision(gv, thev, num_chunks, sym);
+         meep::binary_partition *bp = NULL;
+         if (!my_bp) bp = meep::choose_chunkdivision(gv, thev, num_chunks, sym);
+         std::vector<grid_volume> chunk_vols;
+         std::vector<int> ids;
+         meep::split_by_binarytree(gv, chunk_vols, ids, (!my_bp) ? bp : my_bp);
          for (size_t i = 0; i < chunk_vols.size(); ++i)
               master_printf("CHUNK:, %2zu, %f, %zu\n",i,chunk_vols[i].get_cost(),chunk_vols[i].surface_area());
          return NULL;
@@ -1921,7 +1923,7 @@ meep::structure *create_structure_and_set_materials(vector3 cell_size,
     }
     else {
       s = new meep::structure(gv, NULL, br, sym, num_chunks, Courant,
-                              use_anisotropic_averaging, tol, maxeval);
+                              use_anisotropic_averaging, tol, maxeval, my_bp);
     }
     s->shared_chunks = true;
 

python/simulation.py

@@ -1693,15 +1693,16 @@ def _init_structure(self, k=False):
             absorbers,
             self.extra_materials,
             self.split_chunks_evenly,
-            False if self.chunk_layout else True,
+            False if self.chunk_layout and not isinstance(self.chunk_layout,mp.BinaryPartition) else True,
             None,
-            True if self._output_stats is not None else False
+            True if self._output_stats is not None else False,
+            self.chunk_layout if self.chunk_layout and isinstance(self.chunk_layout,mp.BinaryPartition) else None
         )
 
         if self._output_stats is not None:
             sys.exit(0)
 
-        if self.chunk_layout:
+        if self.chunk_layout and not isinstance(self.chunk_layout,mp.BinaryPartition):
             self.load_chunk_layout(br, self.chunk_layout)
             self.set_materials()
 
@@ -1846,7 +1847,8 @@ def set_materials(self, geometry=None, default_material=None):
             self.split_chunks_evenly,
             True,
             self.structure,
-            False
+            False,
+            None
         )
 
     def dump_structure(self, fname):
@@ -1883,7 +1885,7 @@ def load_chunk_layout(self, br, source):
             ids = source.structure.get_chunk_owners()
             self.structure.load_chunk_layout(vols, [int(f) for f in ids], br)
         else:
-            ## source is either filename (string) or BinaryPartition class object
+            ## source is either filename (string)
             self.structure.load_chunk_layout(source, br)
 
     def init_sim(self):

src/meep.hpp

@@ -729,11 +729,13 @@ class structure {
   structure(const grid_volume &gv, material_function &eps,
             const boundary_region &br = boundary_region(), const symmetry &s = meep::identity(),
             int num_chunks = 0, double Courant = 0.5, bool use_anisotropic_averaging = false,
-            double tol = DEFAULT_SUBPIXEL_TOL, int maxeval = DEFAULT_SUBPIXEL_MAXEVAL);
+            double tol = DEFAULT_SUBPIXEL_TOL, int maxeval = DEFAULT_SUBPIXEL_MAXEVAL,
+            const binary_partition *bp = NULL);
   structure(const grid_volume &gv, double eps(const vec &),
             const boundary_region &br = boundary_region(), const symmetry &s = meep::identity(),
             int num_chunks = 0, double Courant = 0.5, bool use_anisotropic_averaging = false,
-            double tol = DEFAULT_SUBPIXEL_TOL, int maxeval = DEFAULT_SUBPIXEL_MAXEVAL);
+            double tol = DEFAULT_SUBPIXEL_TOL, int maxeval = DEFAULT_SUBPIXEL_MAXEVAL,
+            const binary_partition *bp = NULL);
   structure(const structure *);
   structure(const structure &);
 
@@ -776,7 +778,6 @@ class structure {
   void dump_chunk_layout(const char *filename);
   void load(const char *filename);
   void load_chunk_layout(const char *filename, boundary_region &br);
-  void load_chunk_layout(const binary_partition *bp, boundary_region &br);
   void load_chunk_layout(const std::vector<grid_volume> &gvs,
                          const std::vector<int> &ids,
                          boundary_region &br);
@@ -805,7 +806,7 @@ class structure {
   void use_pml(direction d, boundary_side b, double dx);
   void add_to_effort_volumes(const grid_volume &new_effort_volume, double extra_effort);
   void choose_chunkdivision(const grid_volume &gv, int num_chunks, const boundary_region &br,
-                            const symmetry &s);
+                            const symmetry &s, const binary_partition *bp);
   void check_chunks();
   void changing_chunks();
   // Helper methods for dumping and loading susceptibilities
@@ -814,11 +815,16 @@ class structure {
 };
 
 // defined in structure.cpp
-std::vector<grid_volume> choose_chunkdivision(grid_volume &gv,
-                                              volume &v,
-                                              int num_chunks,
-                                              const symmetry &s);
-
+binary_partition *choose_chunkdivision(grid_volume &gv,
+                                       volume &v,
+                                       int num_chunks,
+                                       const symmetry &s);
+
+// defined in structure_dump.cpp
+void split_by_binarytree(grid_volume gvol,
+                         std::vector<grid_volume> &result_gvs,
+                         std::vector<int> &result_ids,
+                         const binary_partition *bp);
 class src_vol;
 class fields;
 class fields_chunk;

src/structure.cpp

@@ -46,58 +46,62 @@ typedef structure_chunk *structure_chunk_ptr;
 
 structure::structure(const grid_volume &thegv, material_function &eps, const boundary_region &br,
                      const symmetry &s, int num, double Courant, bool use_anisotropic_averaging,
-                     double tol, int maxeval)
+                     double tol, int maxeval, const binary_partition *bp)
     : Courant(Courant), v(D1) // Aaack, this is very hokey.
 {
   outdir = ".";
   shared_chunks = false;
   if (!br.check_ok(thegv)) abort("invalid boundary absorbers for this grid_volume");
   double tstart = wall_time();
-  choose_chunkdivision(thegv, num, br, s);
+  choose_chunkdivision(thegv, num, br, s, bp);
   if (verbosity > 0) master_printf("time for choose_chunkdivision = %g s\n", wall_time() - tstart);
   set_materials(eps, use_anisotropic_averaging, tol, maxeval);
 }
 
 structure::structure(const grid_volume &thegv, double eps(const vec &), const boundary_region &br,
                      const symmetry &s, int num, double Courant, bool use_anisotropic_averaging,
-                     double tol, int maxeval)
+                     double tol, int maxeval, const binary_partition *bp)
     : Courant(Courant), v(D1) // Aaack, this is very hokey.
 {
   outdir = ".";
   shared_chunks = false;
   if (!br.check_ok(thegv)) abort("invalid boundary absorbers for this grid_volume");
   double tstart = wall_time();
-  choose_chunkdivision(thegv, num, br, s);
+  choose_chunkdivision(thegv, num, br, s, bp);
   if (verbosity > 0) master_printf("time for choose_chunkdivision = %g s\n", wall_time() - tstart);
   if (eps) {
     simple_material_function epsilon(eps);
     set_materials(epsilon, use_anisotropic_averaging, tol, maxeval);
   }
 }
 
-static void split_by_cost(int n, grid_volume gvol,
-                          std::vector<grid_volume> &result,
-                          bool fragment_cost) {
+static binary_partition *split_by_cost(int n, grid_volume gvol, bool fragment_cost) {
   if (n == 1) {
-    result.push_back(gvol);
-    return;
+    binary_partition *bp_leaf = new binary_partition(-1);
+    return bp_leaf;
   }
   else {
     int best_split_point;
     direction best_split_direction;
+    double best_split_position;
     double left_effort_fraction;
     gvol.find_best_split(n, fragment_cost, best_split_point, best_split_direction, left_effort_fraction);
+    best_split_position = gvol.surroundings().get_min_corner().in_direction(best_split_direction) +
+      (gvol.surroundings().in_direction(best_split_direction) * best_split_point) /
+      gvol.num_direction(best_split_direction);
+    binary_partition *bp_split = new binary_partition(best_split_direction, best_split_position);
     grid_volume left_gvol = gvol.split_at_fraction(false, best_split_point, best_split_direction);
     const int num_left = (size_t)(left_effort_fraction * n + 0.5);
-    split_by_cost(num_left, left_gvol, result, fragment_cost);
+    bp_split->left = split_by_cost(num_left, left_gvol, fragment_cost);
     grid_volume right_gvol = gvol.split_at_fraction(true, best_split_point, best_split_direction);
-    split_by_cost(n - num_left, right_gvol, result, fragment_cost);
-    return;
+    bp_split->right = split_by_cost(n - num_left, right_gvol, fragment_cost);
+    return bp_split;
   }
 }
 
 void structure::choose_chunkdivision(const grid_volume &thegv, int desired_num_chunks,
-                                     const boundary_region &br, const symmetry &s) {
+                                     const boundary_region &br, const symmetry &s,
+                                     const binary_partition *bp) {
 
   if (thegv.dim == Dcyl && thegv.get_origin().r() < 0) abort("r < 0 origins are not supported");
 
@@ -108,8 +112,13 @@ void structure::choose_chunkdivision(const grid_volume &thegv, int desired_num_c
   a = gv.a;
   dt = Courant / a;
 
+  binary_partition *my_bp = NULL;
+  if (!bp) my_bp = meep::choose_chunkdivision(gv, v, desired_num_chunks, s);
+
   // create the chunks:
-  std::vector<grid_volume> chunk_volumes = meep::choose_chunkdivision(gv, v, desired_num_chunks, s);
+  std::vector<grid_volume> chunk_volumes;
+  std::vector<int> ids;
+  split_by_binarytree(gv, chunk_volumes, ids, (!bp) ? my_bp : bp);
 
   // initialize effort volumes
   num_effort_volumes = 1;
@@ -125,7 +134,7 @@ void structure::choose_chunkdivision(const grid_volume &thegv, int desired_num_c
   num_chunks = 0;
   chunks = new structure_chunk_ptr[chunk_volumes.size() * num_effort_volumes];
   for (size_t i = 0, stop = chunk_volumes.size(); i < stop; ++i) {
-    const int proc = i * count_processors() / chunk_volumes.size();
+    const int proc = (!bp) ? i * count_processors() / chunk_volumes.size() : ids[i] % count_processors();
     for (int j = 0; j < num_effort_volumes; ++j) {
       grid_volume vc;
       if (chunk_volumes[i].intersect_with(effort_volumes[j], &vc)) {
@@ -136,7 +145,6 @@ void structure::choose_chunkdivision(const grid_volume &thegv, int desired_num_c
   }
 
   check_chunks();
-
   if (meep_geom::fragment_stats::resolution != 0) {
     // Save cost of each chunk's grid_volume
     for (int i = 0; i < num_chunks; ++i) {
@@ -146,8 +154,8 @@ void structure::choose_chunkdivision(const grid_volume &thegv, int desired_num_c
 
 }
 
-std::vector<grid_volume> choose_chunkdivision(grid_volume &gv, volume &v, int desired_num_chunks,
-                                              const symmetry &S) {
+binary_partition *choose_chunkdivision(grid_volume &gv, volume &v, int desired_num_chunks,
+                                       const symmetry &S) {
 
   if (desired_num_chunks == 0) desired_num_chunks = count_processors();
   if (gv.dim == Dcyl && gv.get_origin().r() < 0) abort("r < 0 origins are not supported");
@@ -184,22 +192,18 @@ std::vector<grid_volume> choose_chunkdivision(grid_volume &gv, volume &v, int de
       if (break_this[d]) gv = gv.pad((direction)d);
   }
 
-  // Finally, create the chunks:
-  std::vector<grid_volume> chunk_volumes;
-
   if (meep_geom::fragment_stats::resolution == 0 ||
       meep_geom::fragment_stats::split_chunks_evenly) {
     if (verbosity > 0 && desired_num_chunks > 1)
       master_printf("Splitting into %d chunks by voxels\n", desired_num_chunks);
-    split_by_cost(desired_num_chunks, gv, chunk_volumes, false);
+    return split_by_cost(desired_num_chunks, gv, false);
   }
   else {
     if (verbosity > 0 && desired_num_chunks > 1)
       master_printf("Splitting into %d chunks by cost\n", desired_num_chunks);
-    split_by_cost(desired_num_chunks, gv, chunk_volumes, true);
+    return split_by_cost(desired_num_chunks, gv, true);
   }
 
-  return chunk_volumes;
 }
 
 double structure::estimated_cost(int process) {

src/structure_dump.cpp

@@ -430,10 +430,10 @@ binary_partition::binary_partition(direction _split_dir, double _split_pos) {
   right = NULL;
 }
 
-static void split_by_binarytree(grid_volume gvol,
-                                std::vector<grid_volume> &result_gvs,
-                                std::vector<int> &result_ids,
-                                const binary_partition *bp) {
+void split_by_binarytree(grid_volume gvol,
+                         std::vector<grid_volume> &result_gvs,
+                         std::vector<int> &result_ids,
+                         const binary_partition *bp) {
   // reached a leaf
   if ((bp->left == NULL) && (bp->right == NULL)) {
     result_gvs.push_back(gvol);
@@ -456,13 +456,6 @@ static void split_by_binarytree(grid_volume gvol,
   }
 }
 
-void structure::load_chunk_layout(const binary_partition *bp, boundary_region &br) {
-  std::vector<grid_volume> gvs;
-  std::vector<int> ids;
-  split_by_binarytree(gv, gvs, ids, bp);
-  load_chunk_layout(gvs, ids, br);
-}
-
 void structure::load_chunk_layout(const char *filename, boundary_region &br) {
   // Load chunk grid_volumes from a file
   h5file file(filename, h5file::READONLY, true);