WIP: Speed up choose_chunkdivision with geom_boxt_tree

python/meep.i

@@ -1660,8 +1660,7 @@ meep::structure *create_structure_and_set_materials(vector3 cell_size,
     }
 
     // Return params to default state
-    meep_geom::fragment_stats::resolution = 0;
-    meep_geom::fragment_stats::split_chunks_evenly = false;
+    meep_geom::fragment_stats::reset();
 
     return s;
 }

python/tests/fragment_stats.py

@@ -307,6 +307,20 @@ def test_3d_cell_smaller_than_minimum_fragment_size(self):
         self.assertEqual(fs.box.high.z, 0.5)
         self.assertEqual(fs.num_pixels_in_box, 1000)
 
+    def test_geom_box_tree(self):
+
+        block_size = mp.Vector3(3, 3)
+        mat = mp.Medium(index=13)
+        geometry = [
+            mp.Block(size=block_size, center=mp.Vector3(-3, 3), material=mat),
+            mp.Block(size=block_size, center=mp.Vector3(3, 3), material=mat),
+            mp.Block(size=block_size, center=mp.Vector3(3, -3), material=mat),
+            mp.Block(size=block_size, center=mp.Vector3(-3, -3), material=mat)
+        ]
+        fs = self.get_fragment_stats(mp.Vector3(), mp.Vector3(9, 9), 2, geom=geometry)
+        self.check_stats(fs, 32400, 0, 0, 0, 0)
+        self.assertEqual(fs.num_pixels_in_box, 8100)
+
 
 class TestPMLToVolList(unittest.TestCase):
 

src/meepgeom.cpp

@@ -1639,6 +1639,7 @@ int fragment_stats::maxeval = 0;
 int fragment_stats::resolution = 0;
 meep::ndim fragment_stats::dims = meep::D1;
 geometric_object_list fragment_stats::geom = {};
+geom_box_tree fragment_stats::geom_tree = NULL;
 std::vector<dft_data> fragment_stats::dft_data_list;
 std::vector<meep::volume> fragment_stats::pml_1d_vols;
 std::vector<meep::volume> fragment_stats::pml_2d_vols;
@@ -1706,6 +1707,7 @@ fragment_stats compute_fragment_stats(
   geom_box box = make_box_from_cell(cell_size);
   fragment_stats stats = init_stats(box, tol, maxeval, gv);
   stats.compute();
+  fragment_stats::reset();
   return stats;
 }
 
@@ -1717,16 +1719,34 @@ fragment_stats::fragment_stats(geom_box &bx)
   num_pixels_in_box = get_pixels_in_box(&bx);
 }
 
+void fragment_stats::reset() {
+  resolution = 0;
+  split_chunks_evenly = false;
+  destroy_geom_box_tree(geom_tree);
+  geom_tree = NULL;
+}
+
 void fragment_stats::init_libctl(material_type default_mat, bool ensure_per, meep::grid_volume *gv,
                                  vector3 cell_size, vector3 cell_center,
                                  geometric_object_list *geom_) {
   geom_initialize();
   default_material = default_mat;
-  ensure_periodicity = ensure_per;
   geometry_center = cell_center;
   dimensions = meep::number_of_directions(gv->dim);
   geometry_lattice.size = cell_size;
   geom_fix_object_list(*geom_);
+  geom_box cell_box =  make_box_from_cell(cell_size);
+  int temp_dims = dimensions;
+  int temp_periodicity = ensure_periodicity;
+  dimensions = 3;
+  ensure_periodicity = 0;
+  geom_tree = create_geom_box_tree0(*geom_, cell_box);
+  // if (meep::am_master()) display_geom_box_tree(2, geom_tree);
+  // int depth, nobjects;
+  // geom_box_tree_stats(geom_tree, &depth, &nobjects);
+  // master_printf("depth: %d, nobjects: %d\n", depth, nobjects);
+  dimensions = temp_dims;
+  ensure_periodicity = temp_periodicity;
 }
 
 bool fragment_stats::has_non_medium_material() {
@@ -1757,44 +1777,80 @@ void fragment_stats::update_stats_from_material(material_type mat, size_t pixels
   }
 }
 
+void fragment_stats::compute_overlaps_from_tree(geom_box_tree t, std::vector<double> &overlaps) {
+  if (!t || !geom_boxes_intersect(&t->b, &box))
+    return;
+
+  if (t->nobjects > 0) {
+    geom_box intersection;
+    geom_box_intersection(&intersection, &t->b, &box);
+
+    for (int i = 0; i < t->nobjects; ++i) {
+      if (geom_boxes_intersect(&intersection, &t->objects[i].box)) {
+        geom_box shifted_box = {vector3_minus(intersection.low, t->objects[i].shiftby),
+                                vector3_minus(intersection.high, t->objects[i].shiftby)};
+        double overlap = box_overlap_with_object(shifted_box, *t->objects[i].o, tol, maxeval);
+        if (overlap > 0.0) {
+          for (int j = 0; j < geom.num_items; ++j) {
+            if (t->objects[i].o == geom.items + j) {
+              overlaps[j] += overlap;
+              break;
+            }
+          }
+        }
+      }
+    }
+  }
+
+  compute_overlaps_from_tree(t->t1, overlaps);
+  compute_overlaps_from_tree(t->t2, overlaps);
+}
+
 void fragment_stats::compute_stats() {
 
   if (geom.num_items == 0) {
     // If there is no geometry, count the default material for the whole fragment
     update_stats_from_material((material_type)default_material, num_pixels_in_box);
   }
+  else {
+    std::vector<double> overlaps(geom.num_items, 0.0);
+    compute_overlaps_from_tree(geom_tree, overlaps);
 
-  for (int i = 0; i < geom.num_items; ++i) {
-    geometric_object *go = &geom.items[i];
-    double overlap = box_overlap_with_object(box, *go, tol, maxeval);
-
+    size_t total_obj_pixels = 0;
     bool anisotropic_pixels_already_added = false;
-    if (eps_averaging) {
-      // If the object doesn't overlap the entire box, that implies that
-      // an object interface intercepts the box, which means we treat
-      // the entire box as anisotropic. This method could give some false
-      // positives if there is another object with the same material behind
-      // the current object, but in practice it is probably reasonable to
-      // assume that there is a material interface somwhere in the box so
-      // we won't worry about fancier edge-detection methods for now.
-      if (overlap != 1.0) {
-        anisotropic_pixels_already_added = true;
-        num_anisotropic_eps_pixels += num_pixels_in_box;
-        if (mu_not_1(go->material)) {
-          num_anisotropic_mu_pixels += num_pixels_in_box;
+
+    for (int i = 0; i < geom.num_items; ++i) {
+      geometric_object *go = &geom.items[i];
+      double overlap = overlaps[i];
+
+      if (eps_averaging && !anisotropic_pixels_already_added) {
+        // If the object doesn't overlap the entire box, that implies that
+        // an object interface intercepts the box, which means we treat
+        // the entire box as anisotropic. This method could give some false
+        // positives if there is another object with the same material behind
+        // the current object, but in practice it is probably reasonable to
+        // assume that there is a material interface somwhere in the box so
+        // we won't worry about fancier edge-detection methods for now.
+        if (overlap != 1.0) {
+          anisotropic_pixels_already_added = true;
+          num_anisotropic_eps_pixels += num_pixels_in_box;
+          if (mu_not_1(go->material)) {
+            num_anisotropic_mu_pixels += num_pixels_in_box;
+          }
         }
       }
-    }
 
-    // Count contributions from material of object
-    size_t pixels = (size_t)ceil(overlap * num_pixels_in_box);
-    if (pixels > 0) {
-      material_type mat = (material_type)go->material;
-      update_stats_from_material(mat, pixels, anisotropic_pixels_already_added);
+      // Count contributions from material of object
+      size_t pixels = (size_t)ceil(overlap * num_pixels_in_box);
+      if (pixels > 0) {
+        total_obj_pixels += pixels;
+        material_type mat = (material_type)go->material;
+        update_stats_from_material(mat, pixels, anisotropic_pixels_already_added);
+      }
     }
-
     // Count contributions from default_material
-    size_t default_material_pixels = num_pixels_in_box - pixels;
+    // TODO: This is wrong if objects overlap each other.
+    size_t default_material_pixels = num_pixels_in_box - total_obj_pixels;
     if (default_material_pixels > 0) {
       update_stats_from_material((material_type)default_material, default_material_pixels,
                                  anisotropic_pixels_already_added);

src/meepgeom.hpp

@@ -67,6 +67,7 @@ struct fragment_stats {
   static int resolution;
   static meep::ndim dims;
   static geometric_object_list geom;
+  static geom_box_tree geom_tree;
   static std::vector<dft_data> dft_data_list;
   static std::vector<meep::volume> pml_1d_vols;
   static std::vector<meep::volume> pml_2d_vols;
@@ -79,6 +80,7 @@ struct fragment_stats {
   static void init_libctl(meep_geom::material_type default_mat, bool ensure_per,
                           meep::grid_volume *gv, vector3 cell_size, vector3 cell_center,
                           geometric_object_list *geom_list);
+  static void reset();
 
   size_t num_anisotropic_eps_pixels;
   size_t num_anisotropic_mu_pixels;
@@ -116,6 +118,7 @@ struct fragment_stats {
   void compute_dft_stats();
   void compute_pml_stats();
   void compute_absorber_stats();
+  void compute_overlaps_from_tree(geom_box_tree t, std::vector<double> &overlaps);
 };
 
 fragment_stats