user-specified chunk layouts

doc/docs/Python_User_Interface.md

@@ -320,10 +320,11 @@ Python. `Vector3` is a `meep` class.
   `mp.dump_structure`. Defaults to an empty string. See [Load and Dump
   Structure](#load-and-dump-structure) for more information.
 
-+ **`chunk_layout` [`string` or `Simulation` instance]** — This will cause the
-  `Simulation` to use the chunk layout described by either an h5 file (created by
-  `Simulation.dump_chunk_layout`) or another `Simulation`. See [Load and Dump
-  Structure](#load-and-dump-structure) for more information.
++ **`chunk_layout` [`string` or `Simulation` or `BinaryPartition` instance]** —
+  This will cause the `Simulation` to use the chunk layout described by either
+  (1) an `.h5` file (created using `Simulation.dump_chunk_layout`), (2) another
+  `Simulation` or (3) a [`BinaryPartition`](#binarypartition) object. See
+  [Load and Dump Structure](#load-and-dump-structure) for more information.
 
 The following require a bit more understanding of the inner workings of Meep to
 use. See also [SWIG Wrappers](#swig-wrappers).
@@ -7288,6 +7289,51 @@ former value.
 
 </div>
 
+---
+<a id="BinaryPartition"></a>
+
+### BinaryPartition
+
+```python
+class BinaryPartition(object):
+```
+
+<div class="class_docstring" markdown="1">
+
+Binary tree class used for specifying a cell partition of arbitrary sized chunks for use as the
+`chunk_layout` parameter of the `Simulation` class object.
+
+</div>
+
+
+
+<a id="BinaryPartition.__init__"></a>
+
+<div class="class_members" markdown="1">
+
+```python
+def __init__(self,
+             data=None,
+             split_dir=None,
+             split_pos=None,
+             left=None,
+             right=None,
+             id=None):
+```
+
+<div class="method_docstring" markdown="1">
+
+The constructor accepts three separate groups of arguments: (1) `data`: a list of lists where each
+list entry is either (a) a node defined as `[ (split_dir,split_pos), left, right ]` for which `split_dir`
+and `split_pos` define the splitting direction (i.e., `X`, `Y`, `Z`) and position (e.g., `3.5`,
+`-4.2`, etc.) and `left` and `right` are the two branches (themselves `BinaryPartition` objects)
+or (b) a leaf with integer value `id` for the chunk id, (2) a node defined using `split_dir`, `split_pos`,
+`left`, and `right`, or (3) a leaf with `id`. This input format enables specifying the binary tree using
+either a single list for the entire tree or defining the nodes and leaves individually.
+
+</div>
+
+</div>
 
 Miscellaneous Functions Reference
 ---------------------------------

doc/docs/Python_User_Interface.md.in

@@ -822,6 +822,7 @@ The following classes are available directly via the `meep` package.
 @@ Verbosity.get @@
 @@ Verbosity.set @@
 
+@@ BinaryPartition[methods-with-docstrings] @@
 
 Miscellaneous Functions Reference
 ---------------------------------

python/meep.i

@@ -1439,6 +1439,27 @@ void _get_gradient(PyObject *grad, PyObject *fields_a, PyObject *fields_f, PyObj
   }
 }
 
+// typemaps for binary_partition
+
+%typecheck (SWIG_TYPECHECK_POINTER) binary_partition * {
+    $1 = PyObject_IsInstance($input, py_binary_partition_object());
+}
+
+%typemap(in) meep::binary_partition * {
+    $1 = py_bp_to_bp($input);
+    if(!$1) {
+        SWIG_fail;
+    }
+}
+
+%typemap(arginit) meep::binary_partition * {
+    $1 = NULL;
+}
+
+%typemap(freearg) meep::binary_partition * {
+    delete $1;
+}
+
 // Tells Python to take ownership of the h5file* this function returns so that
 // it gets garbage collected and the file gets closed.
 %newobject meep::fields::open_h5file;
@@ -1659,6 +1680,7 @@ PyObject *_get_array_slice_dimensions(meep::fields *f, const meep::volume &where
     )
     from .simulation import (
         Absorber,
+        BinaryPartition,
         Ldos,
         EnergyRegion,
         FluxRegion,

python/simulation.py

@@ -1162,10 +1162,11 @@ def __init__(self,
           `mp.dump_structure`. Defaults to an empty string. See [Load and Dump
           Structure](#load-and-dump-structure) for more information.
 
-        + **`chunk_layout` [`string` or `Simulation` instance]** — This will cause the
-          `Simulation` to use the chunk layout described by either an h5 file (created by
-          `Simulation.dump_chunk_layout`) or another `Simulation`. See [Load and Dump
-          Structure](#load-and-dump-structure) for more information.
+        + **`chunk_layout` [`string` or `Simulation` or `BinaryPartition` instance]** —
+          This will cause the `Simulation` to use the chunk layout described by either
+          (1) an `.h5` file (created using `Simulation.dump_chunk_layout`), (2) another
+          `Simulation` or (3) a [`BinaryPartition`](#binarypartition) object. See
+          [Load and Dump Structure](#load-and-dump-structure) for more information.
 
         The following require a bit more understanding of the inner workings of Meep to
         use. See also [SWIG Wrappers](#swig-wrappers).
@@ -5171,3 +5172,44 @@ def merge_subgroup_data(data):
     comm.Alltoallv(smsg, rmsg)
 
     return output
+
+class BinaryPartition(object):
+    """
+    Binary tree class used for specifying a cell partition of arbitrary sized chunks for use as the
+    `chunk_layout` parameter of the `Simulation` class object.
+    """
+    def __init__(self, data=None, split_dir=None, split_pos=None, left=None, right=None, id=None):
+        """
+        The constructor accepts three separate groups of arguments: (1) `data`: a list of lists where each
+        list entry is either (a) a node defined as `[ (split_dir,split_pos), left, right ]` for which `split_dir`
+        and `split_pos` define the splitting direction (i.e., `X`, `Y`, `Z`) and position (e.g., `3.5`,
+        `-4.2`, etc.) and `left` and `right` are the two branches (themselves `BinaryPartition` objects)
+        or (b) a leaf with integer value `id` for the chunk id, (2) a node defined using `split_dir`, `split_pos`,
+        `left`, and `right`, or (3) a leaf with `id`. This input format enables specifying the binary tree using
+        either a single list for the entire tree or defining the nodes and leaves individually.
+        """
+        self.split_dir = None
+        self.split_pos = None
+        self.id = None
+        self.left = None
+        self.right = None
+        if data is not None:
+            if isinstance(data,list) and len(data) == 3:
+                if isinstance(data[0],tuple) and len(data[0]) == 2:
+                    self.split_dir = data[0][0]
+                    self.split_pos = data[0][1]
+                else:
+                    raise ValueError("expecting 2-tuple (split_dir,split_pos) but got {}".format(data[0]))
+                self.left = BinaryPartition(data=data[1])
+                self.right = BinaryPartition(data=data[2])
+            elif isinstance(data,int):
+                self.id = data
+            else:
+                raise ValueError("expecting list [(split_dir,split_pos), left, right] or int (id) but got {}".format(data))
+        elif split_dir is not None:
+            self.split_dir = split_dir
+            self.split_pos = split_pos
+            self.left = left
+            self.right = right
+        else:
+            self.id = id

python/typemap_utils.cpp

@@ -1029,3 +1029,49 @@ static PyObject *gobj_list_to_py_list(geometric_object_list *objs) {
 
   return py_res;
 }
+
+static meep::binary_partition *py_bp_to_bp(PyObject *bp) {
+    meep::binary_partition *bp = NULL;
+    if (bp == Py_None) return bp;
+
+    PyObject *id = PyObject_GetAttrString(bp, "id");
+    PyObject *split_dir = PyObject_GetAttrString(bp, "split_dir");
+    PyObject *split_pos = PyObject_GetAttrString(bp, "split_pos");
+    PyObject *left = PyObject_GetAttrString(bp, "left");
+    PyObject *right = PyObject_GetAttrString(bp, "right");
+
+    if (!id || !split_dir || !split_pos || !left || !right) {
+        // error....
+    }
+
+    if (PyLong_Check(id)) {
+         bp = new meep::binary_partition(PyLong_AsLong(id));
+    } else {
+         bp = new meep::binary_partition(direction(PyLong_AsLong(split_dir)), PyFloat_AsDouble(split_pos));
+         bp->left = py_bp_to_bp(left);
+         bp->right = py_bp_to_bp(right);
+    }
+
+    Py_XDECREF(id);
+    Py_XDECREF(split_dir);
+    Py_XDECREF(split_pos);
+    Py_XDECREF(left);
+    Py_XDECREF(right);
+    return bp;
+}
+
+static PyObject *get_meep_mod() {
+  // Return value: Borrowed reference
+  static PyObject *meep_mod = NULL;
+  if (meep_mod == NULL) { meep_mod = PyImport_ImportModule("meep"); }
+  return meep_mod;
+}
+
+static PyObject *py_binary_partition_object() {
+  // Return value: Borrowed reference
+  static PyObject *bp_type = NULL;
+  if (bp_type == NULL) {
+    bp_type = PyObject_GetAttrString(get_meep_mod(), "BinaryPartition");
+  }
+  return bp_type;
+}

src/meep.hpp

@@ -707,6 +707,8 @@ boundary_region pml(double thickness, direction d, double Rasymptotic = 1e-15,
 boundary_region pml(double thickness, double Rasymptotic = 1e-15, double mean_stretch = 1.0);
 #define no_pml() boundary_region()
 
+class binary_partition;
+
 class structure {
 public:
   structure_chunk **chunks;
@@ -773,7 +775,11 @@ 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, boundary_region &br);
+  void load_chunk_layout(const std::vector<grid_volume> &gvs,
+                         const std::vector<int> &chunk_ids,
+                         boundary_region &br);
 
   // monitor.cpp
   std::complex<double> get_chi1inv(component, direction, const ivec &origloc, double frequency = 0,
@@ -2150,6 +2156,17 @@ vec get_k(void *vedata);
 realnum linear_interpolate(realnum rx, realnum ry, realnum rz, realnum *data, int nx, int ny,
                            int nz, int stride);
 
+// binary tree class for importing layout of chunk partition
+class binary_partition {
+public:
+  binary_partition(int _id);
+  binary_partition(direction _split_dir, double _split_pos);
+  direction split_dir;
+  double split_pos;
+  int id;
+  binary_partition *left, *right;
+};
+
 } /* namespace meep */
 
 #endif /* MEEP_H */

src/structure_dump.cpp

@@ -414,6 +414,55 @@ void structure::set_chiP_from_file(h5file *file, const char *dataset, field_type
   }
 }
 
+binary_partition::binary_partition(int _id) {
+  id = _id;
+  split_dir = NO_DIRECTION;
+  split_pos = 0.0;
+  left = NULL;
+  right = NULL;
+}
+
+binary_partition::binary_partition(direction _split_dir, double _split_pos) {
+  split_dir = _split_dir;
+  split_pos = _split_pos;
+  id = -1;
+  left = NULL;
+  right = NULL;
+}
+
+static void split_by_binarytree(grid_volume gvol,
+                                std::vector<grid_volume> &result,
+                                std::vector<int> &chunk_ids,
+                                const binary_partition *bp) {
+  // reached a leaf
+  if ((bp->left == NULL) && (bp->right == NULL)) {
+    result.push_back(gvol);
+    chunk_ids.push_back(bp->id);
+    return;
+  }
+
+  int split_point = (size_t)((bp->split_pos - gvol.surroundings().in_direction_min(bp->split_dir)) /
+                             gvol.surroundings().in_direction(bp->split_dir) *
+                             gvol.num_direction(bp->split_dir) + 0.5);
+  // traverse left branch
+  if (bp->left != NULL) {
+    grid_volume left_gvol = gvol.split_at_fraction(false, split_point, bp->split_dir);
+    split_by_binarytree(left_gvol, result, chunk_ids, bp->left);
+  }
+  // traverse right branch
+  if (bp->right != NULL) {
+    grid_volume right_gvol = gvol.split_at_fraction(true, split_point, bp->split_dir);
+    split_by_binarytree(right_gvol, result, chunk_ids, bp->right);
+  }
+}
+
+void structure::load_chunk_layout(const binary_partition *bp, boundary_region &br) {
+  std::vector<grid_volume> gvs;
+  std::vector<int> chunk_ids;
+  split_by_binarytree(gv, gvs, chunk_ids, bp);
+  load_chunk_layout(gvs, chunk_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);
@@ -476,6 +525,18 @@ void structure::load_chunk_layout(const std::vector<grid_volume> &gvs, boundary_
   check_chunks();
 }
 
+void structure::load_chunk_layout(const std::vector<grid_volume> &gvs,
+                                  const std::vector<int> &chunk_ids,
+                                  boundary_region &br) {
+  // Recreate the chunks with the new grid_volumes
+  for (int i = 0; i < num_chunks; ++i) {
+    if (chunks[i]->refcount-- <= 1) delete chunks[i];
+    chunks[i] = new structure_chunk(gvs[i], v, Courant, chunk_ids[i]);
+    br.apply(this, chunks[i]);
+  }
+  check_chunks();
+}
+
 void structure::load(const char *filename) {
   h5file file(filename, h5file::READONLY, true);