Propagate chunk_layout to Python meep.Simulation object (#1673)

* upstream sim.chunk_layout propagation into meep

* fixed stateful funkiness with test_chunk_layout.py

* changes for review

.gitignore

@@ -14,6 +14,7 @@
 *.pyc
 _build
 .vscode
+.idea
 
 # autotools stuff
 Makefile

python/meep.i

@@ -1429,6 +1429,10 @@ void _get_gradient(PyObject *grad, PyObject *fields_a, PyObject *fields_f, PyObj
   $1 = temp.get();
 }
 
+%typemap(out) const meep::binary_partition * {
+  $result = bp_to_py_bp($1);
+}
+
 %typemap(arginit) meep::binary_partition * {
     $1 = NULL;
 }

python/simulation.py

@@ -1210,6 +1210,7 @@ def __init__(self,
         self.default_material = default_material
         self.epsilon_input_file = epsilon_input_file
         self.num_chunks = chunk_layout.numchunks() if isinstance(chunk_layout,mp.BinaryPartition) else num_chunks
+        self._num_chunks_original = self.num_chunks
         self.Courant = Courant
         self.global_d_conductivity = 0
         self.global_b_conductivity = 0
@@ -1238,6 +1239,7 @@ def __init__(self,
         self.force_all_components = force_all_components
         self.split_chunks_evenly = split_chunks_evenly
         self.chunk_layout = chunk_layout
+        self._chunk_layout_original = self.chunk_layout
         self.collect_stats = collect_stats
         self.fragment_stats = None
         self._output_stats = os.environ.get('MEEP_STATS', None)
@@ -1706,6 +1708,12 @@ def _init_structure(self, k=False):
             self.load_chunk_layout(br, self.chunk_layout)
             self.set_materials()
 
+        # Update sim.chunk_layout if it is generated internally from Meep
+        if self.chunk_layout is None:
+            self.chunk_layout = self.structure.get_binary_partition()
+            # We need self.num_chunks to be consistent
+            self.num_chunks = self.chunk_layout.numchunks()
+
         if self.load_structure_file:
             self.load_structure(self.load_structure_file)
 
@@ -3608,11 +3616,15 @@ def change_sources(self, new_sources):
     def reset_meep(self):
         """
         Reset all of Meep's parameters, deleting the fields, structures, etcetera, from
-        memory as if you had not run any computations.
+        memory as if you had not run any computations. If the num_chunks or chunk_layout
+        attributes have been modified internally, they are reset to their original
+        values passed in at instantiation.
         """
         self.fields = None
         self.structure = None
         self.dft_objects = []
+        self.num_chunks = self._num_chunks_original
+        self.chunk_layout = self._chunk_layout_original
         self._is_initialized = False
 
     def restart_fields(self):

python/tests/test_chunk_layout.py

@@ -2,36 +2,43 @@
 import copy
 import unittest
 
-process_ids = []
-chunk_areas = []
 
 def traverse_tree(bp=None,min_corner=None,max_corner=None):
-    if ((min_corner.x > max_corner.x) or (min_corner.y > max_corner.y)):
-        raise RuntimeError("min_corner/max_corner have been incorrectly defined.")
-
-    ## reached a leaf
-    if (bp.left is None and bp.right is None):
-        process_ids.append(bp.proc_id)
-        chunk_area = (max_corner.x-min_corner.x)*(max_corner.y-min_corner.y)
-        chunk_areas.append(chunk_area)
-
-    ## traverse the left branch
-    if (bp.left is not None):
-        new_max_corner = copy.deepcopy(max_corner)
-        if bp.split_dir == mp.X:
-            new_max_corner.x = bp.split_pos
-        else:
-            new_max_corner.y = bp.split_pos
-        traverse_tree(bp.left,min_corner,new_max_corner)
-
-    ## traverse the right branch                                      
-    if (bp.right is not None):
-        new_min_corner = copy.deepcopy(min_corner)
-        if bp.split_dir == mp.X:
-            new_min_corner.x = bp.split_pos
-        else:
-            new_min_corner.y = bp.split_pos
-        traverse_tree(bp.right,new_min_corner,max_corner)
+
+    process_ids = []
+    chunk_areas = []
+
+    def _traverse_tree(bp=None,min_corner=None,max_corner=None):
+        if ((min_corner.x > max_corner.x) or (min_corner.y > max_corner.y)):
+            raise RuntimeError("min_corner/max_corner have been incorrectly defined.")
+
+        ## reached a leaf
+        if (bp.left is None and bp.right is None):
+            process_ids.append(bp.proc_id)
+            chunk_area = (max_corner.x-min_corner.x)*(max_corner.y-min_corner.y)
+            chunk_areas.append(chunk_area)
+
+        ## traverse the left branch
+        if (bp.left is not None):
+            new_max_corner = copy.deepcopy(max_corner)
+            if bp.split_dir == mp.X:
+                new_max_corner.x = bp.split_pos
+            else:
+                new_max_corner.y = bp.split_pos
+            _traverse_tree(bp.left,min_corner,new_max_corner)
+
+        ## traverse the right branch
+        if (bp.right is not None):
+            new_min_corner = copy.deepcopy(min_corner)
+            if bp.split_dir == mp.X:
+                new_min_corner.x = bp.split_pos
+            else:
+                new_min_corner.y = bp.split_pos
+            _traverse_tree(bp.right,new_min_corner,max_corner)
+
+    _traverse_tree(bp=bp, min_corner=min_corner, max_corner=max_corner)
+
+    return process_ids, chunk_areas
 
 
 class TestChunkLayoutBinaryPartition(unittest.TestCase):
@@ -49,10 +56,26 @@ def test_chunk_layout_binary_partition(self):
         owners = sim.structure.get_chunk_owners()
         areas = [ v.surroundings().full_volume() for v in sim.structure.get_chunk_volumes() ]
 
-        traverse_tree(chunk_layout,-0.5*cell_size,0.5*cell_size)
+        process_ids, chunk_areas = traverse_tree(chunk_layout,-0.5*cell_size,0.5*cell_size)
+
+        self.assertListEqual([int(f) for f in owners],[f % mp.count_processors() for f in process_ids])
+        self.assertListEqual(areas,chunk_areas)
+
+    def test_meep_default_chunk_layout(self):
+        cell_size = mp.Vector3(10.0,5.0,0)
+        sim = mp.Simulation(cell_size=cell_size,
+                            resolution=10)
+
+        sim.init_sim()
+        owners = sim.structure.get_chunk_owners()
+        areas = [ v.surroundings().full_volume() for v in sim.structure.get_chunk_volumes() ]
+
+        chunk_layout = sim.chunk_layout
+
+        process_ids, chunk_areas = traverse_tree(chunk_layout,-0.5*cell_size,0.5*cell_size)
 
         self.assertListEqual([int(f) for f in owners],[f % mp.count_processors() for f in process_ids])
         self.assertListEqual(areas,chunk_areas)
-        
+
 if __name__ == '__main__':
     unittest.main()

python/typemap_utils.cpp

@@ -1086,3 +1086,31 @@ static PyObject *py_binary_partition_object() {
   }
   return bp_type;
 }
+
+// Converts a meep::binary_partition object into a Python class instance
+static PyObject *bp_to_py_bp(const meep::binary_partition *bp) {
+  PyObject *bp_class = py_binary_partition_object();
+  PyObject *args = PyTuple_New(0);  // no numbered arguments to pass
+  if (bp->is_leaf()) {
+    // leaf nodes will have proc_id and no other properties
+    int proc_id = bp->get_proc_id();
+    PyObject *kwargs = Py_BuildValue("{s:i}", "proc_id", proc_id);
+    PyObject *py_bp = PyObject_Call(bp_class, args, kwargs);
+    Py_DECREF(args);
+    Py_DECREF(kwargs);
+    return py_bp;
+  } else {
+    // other nodes will have left, right, split_dir, split_pos
+    PyObject *left = bp_to_py_bp(bp->left_tree());
+    PyObject *right = bp_to_py_bp(bp->right_tree());
+    meep::direction split_dir = bp->get_plane().dir;
+    double split_pos = bp->get_plane().pos;
+    PyObject *kwargs =
+        Py_BuildValue("{s:O,s:O,s:i,s:d}", "left", left, "right", right,
+                      "split_dir", split_dir, "split_pos", split_pos);
+    PyObject *py_bp = PyObject_Call(bp_class, args, kwargs);
+    Py_DECREF(args);
+    Py_DECREF(kwargs);
+    return py_bp;
+  }
+}