user-specified chunk layouts (for parallelization)

[stevengj]

Perhaps the easiest way to specify this would just be as a binary tree of "cuts". i.e. you'd say "cut the cell in half at x=3", and then for first have "cut it in half at y=7", etcetera. (Meep would then convert this into grid points given your resolution and cell size.) You could also optionally specify MPI ranks at the leaves of the tree to control parallelization.

[oskooi]

For reference, here is a 2d example of how this would be set up. The left figure shows the chunk division for a 10 × 5 cell comprised of 5 chunks. The red dots identify the coordinates of various points in the cell with the origin at lower left. The right figure is a binary-tree representation of the left figure. The nodes of the tree indicate the partition point in either the x or y direction (e.g., x = 3, y = 4, etc) and the leaves indicate the chunk id (which can be replaced with the MPI rank).

Regarding the actual implementation, there are at least two possible approaches: (1) create a new class for the binary tree in Python which is used to generate the chunk vertices that are then passed to C++ or (2) have the user specify the binary tree in Python as e.g. a list which is then passed to C++ where the "tree" is traversed recursively while calling the existing grid_volume::split_at_fraction to form the chunks. (2) seems more in line with split_by_cost which also recursively divides the cell into chunks.

[stevengj]

I would just define a trivial binary-tree class in Python.

[oskooi]

Here is a demonstration of a binary-tree class for the 2d cell which consists of a constructor and a single member function traverse_tree that initially given two vertices (min_corner,max_corner) specifying the cell area traverses the tree recursively starting at its root and computes the two corner vertices of all chunks:

import meep as mp
import copy

class BinaryTree:
    def __init__(self,key,val):
        self.left = None
        self.right = None
        if ((key in ['X','Y'] and isinstance(val, (int,float))) or
            (key is 'chunk_id' and isinstance(val, int))):
            self.data = {key: val}
        else:
            raise ValueError("BinaryTree: key/val must be ('X','Y')/number or 'chunk_id'/int.")

    def traverse_tree(self,min_corner=None,max_corner=None):
        if ((min_corner.x > max_corner.x) or (min_corner.y > max_corner.y)):
            raise RuntimeError("BinaryTree has been incorrectly defined.")

        ## reached a leaf
        if (self.left is None and self.right is None):
            try:
                print("chunk{}:, min_corner: ({},{}), max_corner: ({},{})".format(self.data['chunk_id'],
                                                                                  min_corner.x,
                                                                                  min_corner.y,
                                                                                  max_corner.x,
                                                                                  max_corner.y))
            except:
                raise ValueError("expected key 'chunk_id' but got {}".format(self.data.keys()))

        ## traverse the left branch                                      
        if (self.left is not None):
            new_max_corner = copy.deepcopy(max_corner)
            if 'X' in self.data.keys():
                if new_max_corner.x > self.data['X']:
                    new_max_corner.x = self.data['X']
            elif 'Y' in self.data.keys():
                if new_max_corner.y > self.data['Y']:
                    new_max_corner.y = self.data['Y']
            else:
                raise ValueError("expected key 'X' or 'Y' but got {}".format(self.data.keys()))
            self.left.traverse_tree(min_corner,new_max_corner)

        ## traverse the right branch                                       
        if (self.right is not None):
            new_min_corner = copy.deepcopy(min_corner)
            if 'X' in self.data.keys():
                if new_min_corner.x < self.data['X']:
                    new_min_corner.x = self.data['X']
            elif 'Y' in self.data.keys():
                if new_min_corner.y < self.data['Y']:
                    new_min_corner.y = self.data['Y']
            else:
                raise ValueError("expected key 'X' or 'Y' but got {}".format(self.data.keys()))
            self.right.traverse_tree(new_min_corner,max_corner)

The resolution of the Meep grid is not necessary for computing the chunk vertices.

The tree shown in the figure above is set up as follows:

>>> root = BinaryTree('X',3.0)
>>> root.left = BinaryTree('chunk_id',1)
>>> root.right = BinaryTree('Y',4.0)
>>> root.right.right = BinaryTree('chunk_id',3)
>>> root.right.left = BinaryTree('X',8.0)
>>> root.right.left.left = BinaryTree('chunk_id',2)
>>> root.right.left.right = BinaryTree('Y',2.0)
>>> root.right.left.right.left = BinaryTree('chunk_id',5)
>>> root.right.left.right.right = BinaryTree('chunk_id',4)

It will be somewhat tedious to specify the nodes this way (i.e., one-by-one on separate lines) for a large and complicated tree. Thus, it might be nicer to come up with a more compact representation.

The traverse_tree function is then called with the two vertices arguments defining the cell area.

>>> root.traverse_tree(mp.Vector3(0,0),mp.Vector3(10.0,5.0))
chunk1:, min_corner: (0.0,0.0), max_corner: (3.0,5.0)
chunk2:, min_corner: (3.0,0.0), max_corner: (8.0,4.0)
chunk5:, min_corner: (8.0,0.0), max_corner: (10.0,2.0)
chunk4:, min_corner: (8.0,2.0), max_corner: (10.0,4.0)
chunk3:, min_corner: (3.0,4.0), max_corner: (10.0,5.0)

The corner vertices computed for each chunk agree with the figure above.

With this approach, a list of pairs of vertices for each chunk can be generated in Python and then passed into C++ to structure::choose_chunkdivision to create the actual chunks as a std::vector<grid_volume> bypassing the call to meep::choose_chunkdivision:

meep/src/structure.cpp

Lines 112 to 113 in 5c858e3

	// create the chunks:
	std::vector<grid_volume> chunk_volumes = meep::choose_chunkdivision(gv, v, desired_num_chunks, s);

[stevengj]

>>> root = BinaryTree('X',3.0)
>>> root.left = BinaryTree('chunk_id',1)
>>> root.right = BinaryTree('Y',4.0)
>>> root.right.right = BinaryTree('chunk_id',3)
>>> root.right.left = BinaryTree('X',8.0)
>>> root.right.left.left = BinaryTree('chunk_id',2)
>>> root.right.left.right = BinaryTree('Y',2.0)
>>> root.right.left.right.left = BinaryTree('chunk_id',5)
>>> root.right.left.right.right = BinaryTree('chunk_id',4)

Equivalently, we could just have a constructor from a list of lists, where each list entry is either [ (dim,cut), left, right ] or id:

tree = BinaryTree( [ ('X',3.0),
                     [ 1, [ ('Y',4.0), 
                            [ ('X',8.0), 2, [ ('Y',2.0), 5,4 ] ],
                            3 ]
                     ]
                    ])

(In this sense, you don't even need a BinaryTree type at all … you could just pass this nested list directly to a glue routine that converts it into a C++ list.)

[stevengj]

The constructor might look like:

class BinaryTree:
    def __init__(self, data=None, dir=None, cut=None, id=None, left=None, right=None):
        if data is not None:
            .... construct from nested lists ...
        elif dir is not None:
            self.data = (dir, cut, left, right)
        else:
            self.data = id

so that the tree data is either a tuple (dir, cut, left, right) or an id.

And then we would have a typemap to convert BinaryTree into a corresponding C++ tree.

(Maybe call it BinaryPartition.)