Add ability to visualize chunks (#671)

* Add visualize_chunks

* Make gv2box public

doc/docs/Python_User_Interface.md

@@ -1361,6 +1361,12 @@ Returns a `mp.Volume` read from a GDSII file `fname` on `layer` with `zmin` and
 fr = mp.FluxRegion(volume=mp.GDSII_vol(fname, layer, zmin, zmax))
 ```
 
+### Data Visualization
+
+**`Simulation.visualize_chunks()`**
+—
+Displays an interactive image of how the cell is divided into chunks. Each rectangular region is a chunk, and each color represents a different processor. Requires [matplotlib](https://matplotlib.org).
+
 Run and Step Functions
 ----------------------
 

python/simulation.py

@@ -1917,6 +1917,66 @@ def get_sfield_r(self):
     def get_sfield_p(self):
         return self.get_array(cmplx=True, component=mp.Sp)
 
+    def visualize_chunks(self):
+        if self.structure is None:
+            self.init_sim()
+        try:
+            import matplotlib.pyplot as plt
+            import matplotlib.cm
+            import matplotlib.colors
+            from mpl_toolkits.mplot3d import Axes3D
+            from mpl_toolkits.mplot3d.art3d import Poly3DCollection, Line3DCollection
+        except ImportError:
+            warnings.warn("matplotlib is required for visualize_chunks", ImportWarning)
+            return
+
+        vols = self.structure.get_chunk_volumes()
+        owners = self.structure.get_chunk_owners()
+
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection='3d')
+        chunk_colors = matplotlib.cm.rainbow(np.linspace(0, 1, mp.count_processors()))
+
+        def plot_box(box, proc):
+            low = mp.Vector3(box.low.x, box.low.y, box.low.z)
+            high = mp.Vector3(box.high.x, box.high.y, box.high.z)
+            points = [low, high]
+
+            x_len = mp.Vector3(high.x) - mp.Vector3(low.x)
+            y_len = mp.Vector3(y=high.y) - mp.Vector3(y=low.y)
+            xy_len = mp.Vector3(high.x, high.y) - mp.Vector3(low.x, low.y)
+
+            points += [low + x_len]
+            points += [low + y_len]
+            points += [low + xy_len]
+            points += [high - x_len]
+            points += [high - y_len]
+            points += [high - xy_len]
+            points = np.array([np.array(v) for v in points])
+
+            edges = [
+                [points[0], points[2], points[4], points[3]],
+                [points[1], points[5], points[7], points[6]],
+                [points[0], points[3], points[5], points[7]],
+                [points[1], points[4], points[2], points[6]],
+                [points[3], points[4], points[1], points[5]],
+                [points[0], points[7], points[6], points[2]]
+            ]
+
+            faces = Poly3DCollection(edges, linewidths=1, edgecolors='k')
+            color_with_alpha = matplotlib.colors.to_rgba(chunk_colors[proc], alpha=0.2)
+            faces.set_facecolor(color_with_alpha)
+            ax.add_collection3d(faces)
+
+            # Plot the points themselves to force the scaling of the axes
+            ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0)
+
+        if mp.am_master():
+            for i, v in enumerate(vols):
+                plot_box(mp.gv2box(v), owners[i])
+            ax.set_aspect('auto')
+            plt.show()
+
 
 def _create_boundary_region_from_boundary_layers(boundary_layers, gv):
     br = mp.boundary_region()

src/meep.hpp

@@ -708,6 +708,8 @@ class structure {
 
   bool equal_layout(const structure &) const;
   void print_layout(void) const;
+  std::vector<volume> get_chunk_volumes() const;
+  std::vector<int> get_chunk_owners() const;
 
   // structure_dump.cpp
   void dump(const char *filename);

src/meepgeom.cpp

@@ -236,7 +236,7 @@ meep::vec vector3_to_vec(const vector3 v3) {
   }
 }
 
-static geom_box gv2box(const meep::volume &v) {
+geom_box gv2box(const meep::volume &v) {
   geom_box box;
   box.low = vec_to_vector3(v.get_min_corner());
   box.high = vec_to_vector3(v.get_max_corner());

src/meepgeom.hpp

@@ -174,6 +174,7 @@ bool is_medium(material_type md, medium_struct **m);
 bool is_medium(void *md, medium_struct **m);
 bool is_metal(meep::field_type ft, const material_type *material);
 void check_offdiag(medium_struct *m);
+geom_box gv2box(const meep::volume &v);
 
 /***************************************************************/
 /* routines in GDSIIgeom.cc ************************************/

src/structure.cpp

@@ -940,4 +940,21 @@ void structure::print_layout(void) const {
   }
 }
 
+std::vector<volume> structure::get_chunk_volumes() const {
+  std::vector<volume> result;
+
+  for (int i = 0; i < num_chunks; ++i) {
+    result.push_back(chunks[i]->gv.surroundings());
+  }
+  return result;
+}
+
+std::vector<int> structure::get_chunk_owners() const {
+  std::vector<int> result;
+
+  for (int i = 0; i < num_chunks; ++i) {
+    result.push_back(chunks[i]->n_proc());
+  }
+  return result;
+}
 } // namespace meep