Feat/refine tree

discretize/TreeMesh.py

@@ -346,9 +346,10 @@ def permuteE(self):
 
     def plotSlice(
         self, v, vType='CC',
-        normal='Z', ind=None, grid=True, view='real',
+        normal='Z', ind=None, grid=False, view='real',
         ax=None, clim=None, showIt=False,
-        pcolorOpts=None, streamOpts=None, gridOpts=None
+        pcolorOpts=None, streamOpts=None, gridOpts=None,
+        range_x=None, range_y=None,
     ):
 
         if pcolorOpts is None:
@@ -357,10 +358,33 @@ def plotSlice(
             streamOpts = {'color': 'k'}
         if gridOpts is None:
             gridOpts = {'color': 'k', 'alpha': 0.5}
-        if vType not in ['CC', 'F', 'E']:
-            raise ValueError('vType must be one of CC, F, or E')
+        vTypeOpts = ['CC', 'N', 'F', 'E', 'Fx', 'Fy', 'Fz', 'E', 'Ex', 'Ey', 'Ez']
+        viewOpts = ['real', 'imag', 'abs']
+        normalOpts = ['X', 'Y', 'Z']
+        if vType not in vTypeOpts:
+            raise ValueError(
+                "vType must be in ['{0!s}']".format("', '".join(vTypeOpts))
+            )
         if self.dim == 2:
-            raise Exception('plotSlice not support for 2D, use plotImage')
+            raise NotImplementedError(
+                'Must be a 3D mesh. Use plotImage.'
+            )
+        if view == 'vec':
+            raise NotImplementedError(
+                'Vector view plotting is not implemented for TreeMesh (yet)'
+            )
+        if view not in viewOpts:
+            raise ValueError(
+                "view must be in ['{0!s}']".format("', '".join(viewOpts))
+            )
+        normal = normal.upper()
+        if normal not in normalOpts:
+            raise ValueError(
+                "normal must be in ['{0!s}']".format("', '".join(normalOpts))
+            )
+
+        if not isinstance(grid, bool):
+            raise TypeError('grid must be a boolean')
 
         import matplotlib.pyplot as plt
         import matplotlib
@@ -385,7 +409,7 @@ def plotSlice(
         if type(ind) not in integer_types:
             raise ValueError('ind must be an integer')
 
-        #create a temporary TreeMesh with the slice through
+        # create a temporary TreeMesh with the slice through
         temp_mesh = TreeMesh(h2d, x2d)
         level_diff = self.max_level - temp_mesh.max_level
 
@@ -403,24 +427,41 @@ def plotSlice(
         tm_gridboost[:, antiNormalInd] = temp_mesh.gridCC
         tm_gridboost[:, normalInd] = slice_loc
 
-        #interpolate values to self.gridCC if not 'CC'
-        if vType in ['F', 'E']:
+        # interpolate values to self.gridCC if not 'CC'
+        if vType is not 'CC':
             aveOp = 'ave' + vType + '2CC'
             Av = getattr(self, aveOp)
-            v = Av*v
-
-        #interpolate values from self.gridCC to grid2d
+            if v.size == Av.shape[1]:
+                v = Av*v
+            elif len(vType) == 2:
+                # was one of Fx, Fy, Fz, Ex, Ey, Ez
+                # assuming v has all three components in these cases
+                vec_ind = {'x': 0, 'y': 1, 'z': 2}[vType[1]]
+                if vType[0] == 'E':
+                    i_s = np.cumsum([0, self.nEx, self.nEy, self.nEz])
+                elif vType[0] == 'F':
+                    i_s = np.cumsum([0, self.nFx, self.nFy, self.nFz])
+                v = v[i_s[vec_ind]:i_s[vec_ind+1]]
+                v = Av*v
+
+        # interpolate values from self.gridCC to grid2d
         ind_3d_to_2d = self._get_containing_cell_indexes(tm_gridboost)
         v2d = v[ind_3d_to_2d]
 
         if ax is None:
-            fig = plt.figure()
+            plt.figure()
             ax = plt.subplot(111)
-        else:
-            assert isinstance(ax, matplotlib.axes.Axes), "ax must be an matplotlib.axes.Axes"
-            fig = ax.figure
-
-        out = temp_mesh.plotImage(v2d, ax=ax, grid=grid, showIt=showIt, clim=clim)
+        elif not isinstance(ax, matplotlib.axes.Axes):
+            raise Exception("ax must be an matplotlib.axes.Axes")
+
+        out = temp_mesh.plotImage(
+            v2d, vType='CC',
+            grid=grid, view=view,
+            ax=ax, clim=clim, showIt=False,
+            pcolorOpts=pcolorOpts,
+            gridOpts=gridOpts,
+            range_x=range_x,
+            range_y=range_y)
 
         ax.set_xlabel('y' if normal == 'X' else 'x')
         ax.set_ylabel('y' if normal == 'Z' else 'z')

discretize/tree.cpp

@@ -395,148 +395,160 @@ void Cell::divide(node_map_t& nodes, double* xs, double* ys, double* zs, bool fo
     }else if(force){
         do_splitting = true;
     }else{
-        int_t test_level = (*test_func)(this);
-        if(test_level > level){
-            do_splitting = true;
+        if(is_leaf()){
+            // Only need to call the function if I am a leaf...
+            int_t test_level = (*test_func)(this);
+            do_splitting = test_level > level;
         }
     }
-    if(!do_splitting){
-        return;
-    }
-    //If i haven't already been split...
-    if(children[0] == NULL){
-        spawn(nodes, children, xs, ys, zs);
-
-        //If I need to be split, and my neighbor is below my level
-        //Then it needs to be split
-        //-x,+x,-y,+y,-z,+z
-        if(balance){
-            for(int_t i = 0; i < 2*n_dim; ++i){
-                if(neighbors[i] != NULL && neighbors[i]->level < level){
-                    neighbors[i]->divide(nodes, xs, ys, zs, true);
+    if(do_splitting){
+        //If i haven't already been split...
+        if(is_leaf()){
+            spawn(nodes, children, xs, ys, zs);
+
+            //If I need to be split, and my neighbor is below my level
+            //Then it needs to be split
+            //-x,+x,-y,+y,-z,+z
+            if(balance){
+                for(int_t i = 0; i < 2*n_dim; ++i){
+                    if(neighbors[i] != NULL && neighbors[i]->level < level){
+                        neighbors[i]->divide(nodes, xs, ys, zs, true);
+                    }
                 }
             }
-        }
 
-        //Set children's neighbors (first do the easy ones)
-        // all of the children live next to each other
-        children[0]->set_neighbor(children[1], 1);
-        children[0]->set_neighbor(children[2], 3);
-        children[1]->set_neighbor(children[3], 3);
-        children[2]->set_neighbor(children[3], 1);
-
-        if(n_dim == 3){
-            children[4]->set_neighbor(children[5], 1);
-            children[4]->set_neighbor(children[6], 3);
-            children[5]->set_neighbor(children[7], 3);
-            children[6]->set_neighbor(children[7], 1);
-
-            children[0]->set_neighbor(children[4], 5);
-            children[1]->set_neighbor(children[5], 5);
-            children[2]->set_neighbor(children[6], 5);
-            children[3]->set_neighbor(children[7], 5);
-        }
+            //Set children's neighbors (first do the easy ones)
+            // all of the children live next to each other
+            children[0]->set_neighbor(children[1], 1);
+            children[0]->set_neighbor(children[2], 3);
+            children[1]->set_neighbor(children[3], 3);
+            children[2]->set_neighbor(children[3], 1);
+
+            if(n_dim == 3){
+                children[4]->set_neighbor(children[5], 1);
+                children[4]->set_neighbor(children[6], 3);
+                children[5]->set_neighbor(children[7], 3);
+                children[6]->set_neighbor(children[7], 1);
+
+                children[0]->set_neighbor(children[4], 5);
+                children[1]->set_neighbor(children[5], 5);
+                children[2]->set_neighbor(children[6], 5);
+                children[3]->set_neighbor(children[7], 5);
+            }
 
-        // -x direction
-        if(neighbors[0]!=NULL && !(neighbors[0]->is_leaf())){
-            children[0]->set_neighbor(neighbors[0]->children[1], 0);
-            children[2]->set_neighbor(neighbors[0]->children[3], 0);
-        }
-        else{
-            children[0]->set_neighbor(neighbors[0], 0);
-            children[2]->set_neighbor(neighbors[0], 0);
-        }
-        // +x direction
-        if(neighbors[1]!=NULL && !neighbors[1]->is_leaf()){
-            children[1]->set_neighbor(neighbors[1]->children[0], 1);
-            children[3]->set_neighbor(neighbors[1]->children[2], 1);
-        }else{
-            children[1]->set_neighbor(neighbors[1], 1);
-            children[3]->set_neighbor(neighbors[1], 1);
-        }
-        // -y direction
-        if(neighbors[2]!=NULL && !neighbors[2]->is_leaf()){
-            children[0]->set_neighbor(neighbors[2]->children[2], 2);
-            children[1]->set_neighbor(neighbors[2]->children[3], 2);
-        }else{
-            children[0]->set_neighbor(neighbors[2], 2);
-            children[1]->set_neighbor(neighbors[2], 2);
-        }
-        // +y direction
-        if(neighbors[3]!=NULL && !neighbors[3]->is_leaf()){
-            children[2]->set_neighbor(neighbors[3]->children[0], 3);
-            children[3]->set_neighbor(neighbors[3]->children[1], 3);
-        }else{
-            children[2]->set_neighbor(neighbors[3], 3);
-            children[3]->set_neighbor(neighbors[3], 3);
-        }
-        if(n_dim==3){
             // -x direction
             if(neighbors[0]!=NULL && !(neighbors[0]->is_leaf())){
-                children[4]->set_neighbor(neighbors[0]->children[5], 0);
-                children[6]->set_neighbor(neighbors[0]->children[7], 0);
+                children[0]->set_neighbor(neighbors[0]->children[1], 0);
+                children[2]->set_neighbor(neighbors[0]->children[3], 0);
             }
             else{
-                children[4]->set_neighbor(neighbors[0], 0);
-                children[6]->set_neighbor(neighbors[0], 0);
+                children[0]->set_neighbor(neighbors[0], 0);
+                children[2]->set_neighbor(neighbors[0], 0);
             }
             // +x direction
             if(neighbors[1]!=NULL && !neighbors[1]->is_leaf()){
-                children[5]->set_neighbor(neighbors[1]->children[4], 1);
-                children[7]->set_neighbor(neighbors[1]->children[6], 1);
+                children[1]->set_neighbor(neighbors[1]->children[0], 1);
+                children[3]->set_neighbor(neighbors[1]->children[2], 1);
             }else{
-                children[5]->set_neighbor(neighbors[1], 1);
-                children[7]->set_neighbor(neighbors[1], 1);
+                children[1]->set_neighbor(neighbors[1], 1);
+                children[3]->set_neighbor(neighbors[1], 1);
             }
             // -y direction
             if(neighbors[2]!=NULL && !neighbors[2]->is_leaf()){
-                children[4]->set_neighbor(neighbors[2]->children[6], 2);
-                children[5]->set_neighbor(neighbors[2]->children[7], 2);
+                children[0]->set_neighbor(neighbors[2]->children[2], 2);
+                children[1]->set_neighbor(neighbors[2]->children[3], 2);
             }else{
-                children[4]->set_neighbor(neighbors[2], 2);
-                children[5]->set_neighbor(neighbors[2], 2);
+                children[0]->set_neighbor(neighbors[2], 2);
+                children[1]->set_neighbor(neighbors[2], 2);
             }
             // +y direction
             if(neighbors[3]!=NULL && !neighbors[3]->is_leaf()){
-                children[6]->set_neighbor(neighbors[3]->children[4], 3);
-                children[7]->set_neighbor(neighbors[3]->children[5], 3);
+                children[2]->set_neighbor(neighbors[3]->children[0], 3);
+                children[3]->set_neighbor(neighbors[3]->children[1], 3);
             }else{
-                children[6]->set_neighbor(neighbors[3], 3);
-                children[7]->set_neighbor(neighbors[3], 3);
+                children[2]->set_neighbor(neighbors[3], 3);
+                children[3]->set_neighbor(neighbors[3], 3);
             }
-            // -z direction
-            if(neighbors[4]!=NULL && !neighbors[4]->is_leaf()){
-                children[0]->set_neighbor(neighbors[4]->children[4], 4);
-                children[1]->set_neighbor(neighbors[4]->children[5], 4);
-                children[2]->set_neighbor(neighbors[4]->children[6], 4);
-                children[3]->set_neighbor(neighbors[4]->children[7], 4);
-            }else{
-                children[0]->set_neighbor(neighbors[4], 4);
-                children[1]->set_neighbor(neighbors[4], 4);
-                children[2]->set_neighbor(neighbors[4], 4);
-                children[3]->set_neighbor(neighbors[4], 4);
-            }
-            // +z direction
-            if(neighbors[5]!=NULL && !neighbors[5]->is_leaf()){
-                children[4]->set_neighbor(neighbors[5]->children[0], 5);
-                children[5]->set_neighbor(neighbors[5]->children[1], 5);
-                children[6]->set_neighbor(neighbors[5]->children[2], 5);
-                children[7]->set_neighbor(neighbors[5]->children[3], 5);
-            }else{
-                children[4]->set_neighbor(neighbors[5], 5);
-                children[5]->set_neighbor(neighbors[5], 5);
-                children[6]->set_neighbor(neighbors[5], 5);
-                children[7]->set_neighbor(neighbors[5], 5);
+            if(n_dim==3){
+                // -x direction
+                if(neighbors[0]!=NULL && !(neighbors[0]->is_leaf())){
+                    children[4]->set_neighbor(neighbors[0]->children[5], 0);
+                    children[6]->set_neighbor(neighbors[0]->children[7], 0);
+                }
+                else{
+                    children[4]->set_neighbor(neighbors[0], 0);
+                    children[6]->set_neighbor(neighbors[0], 0);
+                }
+                // +x direction
+                if(neighbors[1]!=NULL && !neighbors[1]->is_leaf()){
+                    children[5]->set_neighbor(neighbors[1]->children[4], 1);
+                    children[7]->set_neighbor(neighbors[1]->children[6], 1);
+                }else{
+                    children[5]->set_neighbor(neighbors[1], 1);
+                    children[7]->set_neighbor(neighbors[1], 1);
+                }
+                // -y direction
+                if(neighbors[2]!=NULL && !neighbors[2]->is_leaf()){
+                    children[4]->set_neighbor(neighbors[2]->children[6], 2);
+                    children[5]->set_neighbor(neighbors[2]->children[7], 2);
+                }else{
+                    children[4]->set_neighbor(neighbors[2], 2);
+                    children[5]->set_neighbor(neighbors[2], 2);
+                }
+                // +y direction
+                if(neighbors[3]!=NULL && !neighbors[3]->is_leaf()){
+                    children[6]->set_neighbor(neighbors[3]->children[4], 3);
+                    children[7]->set_neighbor(neighbors[3]->children[5], 3);
+                }else{
+                    children[6]->set_neighbor(neighbors[3], 3);
+                    children[7]->set_neighbor(neighbors[3], 3);
+                }
+                // -z direction
+                if(neighbors[4]!=NULL && !neighbors[4]->is_leaf()){
+                    children[0]->set_neighbor(neighbors[4]->children[4], 4);
+                    children[1]->set_neighbor(neighbors[4]->children[5], 4);
+                    children[2]->set_neighbor(neighbors[4]->children[6], 4);
+                    children[3]->set_neighbor(neighbors[4]->children[7], 4);
+                }else{
+                    children[0]->set_neighbor(neighbors[4], 4);
+                    children[1]->set_neighbor(neighbors[4], 4);
+                    children[2]->set_neighbor(neighbors[4], 4);
+                    children[3]->set_neighbor(neighbors[4], 4);
+                }
+                // +z direction
+                if(neighbors[5]!=NULL && !neighbors[5]->is_leaf()){
+                    children[4]->set_neighbor(neighbors[5]->children[0], 5);
+                    children[5]->set_neighbor(neighbors[5]->children[1], 5);
+                    children[6]->set_neighbor(neighbors[5]->children[2], 5);
+                    children[7]->set_neighbor(neighbors[5]->children[3], 5);
+                }else{
+                    children[4]->set_neighbor(neighbors[5], 5);
+                    children[5]->set_neighbor(neighbors[5], 5);
+                    children[6]->set_neighbor(neighbors[5], 5);
+                    children[7]->set_neighbor(neighbors[5], 5);
+                }
             }
         }
     }
     if(!force){
-        for(int_t i = 0; i < (1<<n_dim); ++i){
-            children[i]->divide(nodes, xs, ys, zs);
+        if(!is_leaf()){
+            for(int_t i = 0; i < (1<<n_dim); ++i){
+                children[i]->divide(nodes, xs, ys, zs);
+            }
         }
     }
 };
 
+
+void Cell::set_test_function(function func){
+    test_func = func;
+    if(!is_leaf()){
+        for(int_t i = 0; i < (1<<n_dim); ++i){
+            children[i]->set_test_function(func);
+        }
+    }
+}
+
 void Cell::build_cell_vector(cell_vec_t& cells){
     if(this->is_leaf()){
         cells.push_back(this);
@@ -723,7 +735,7 @@ void Tree::build_tree_from_function(function test_func){
     for(int_t iz=0; iz<nz_roots; ++iz)
         for(int_t iy=0; iy<ny_roots; ++iy)
             for(int_t ix=0; ix<nx_roots; ++ix)
-                roots[iz][iy][ix]->test_func = test_func;
+                roots[iz][iy][ix]->set_test_function(test_func);
     //Now we can divide
     for(int_t iz=0; iz<nz_roots; ++iz)
         for(int_t iy=0; iy<ny_roots; ++iy)