update treemesh balancing option in more places.

discretize/mixins/mpl_mod.py

@@ -2149,7 +2149,7 @@ def __plot_slice_tree(
         normal[normalInd] = 1
 
         # create a temporary TreeMesh with the slice through
-        temp_mesh = discretize.TreeMesh(h2d, x2d)
+        temp_mesh = discretize.TreeMesh(h2d, x2d, diagonal_balance=False)
         level_diff = self.max_level - temp_mesh.max_level
 
         # get list of cells which intersect the slicing plane

discretize/tree_mesh.py

@@ -1070,7 +1070,7 @@ def equals(self, other):  # NOQA D102
 
     def __reduce__(self):
         """Return the necessary items to reconstruct this object's state."""
-        return TreeMesh, (self.h, self.origin), self.__getstate__()
+        return TreeMesh, (self.h, self.origin, False), self.__getstate__()
 
     cellGrad = deprecate_property(
         "cell_gradient", "cellGrad", removal_version="1.0.0", error=True

discretize/utils/mesh_utils.py

@@ -401,6 +401,7 @@ def mesh_builder_xyz(
     depth_core=None,
     expansion_factor=1.3,
     mesh_type="tensor",
+    tree_diagonal_balance=None,
 ):
     """Generate a tensor or tree mesh using a cloud of points.
 
@@ -428,6 +429,9 @@ def mesh_builder_xyz(
         Expansion factor for padding cells. Ignored if *mesh_type* = *tree*
     mesh_type : {'tensor', 'tree'}
         Specify output mesh type
+    tree_diagonal_balance : bool, optional
+        Whether to diagonally balance the tree mesh, `None` will use the `TreeMesh`
+        default behavoir.
 
     Returns
     -------
@@ -521,7 +525,7 @@ def expand(dx, pad):
             h_dim += [np.ones(2**maxLevel) * h[ii]]
 
         # Define the mesh and origin
-        mesh = discretize.TreeMesh(h_dim)
+        mesh = discretize.TreeMesh(h_dim, diagonal_balance=tree_diagonal_balance)
 
         for ii, _cc in enumerate(nC):
             core = limits[ii][0] - limits[ii][1]