Merge pull request #382 from jcapriot/bugfixes

Outstanding bugfixes.

Author: jcapriot

discretize/tree_mesh.py

@@ -1,5 +1,7 @@
 """Module containing the TreeMesh implementation."""
 
+import warnings
+
 #      ___          ___       ___          ___          ___          ___
 #     /\  \        /\  \     /\  \        /\  \        /\  \        /\  \
 #    /::\  \      /::\  \    \:\  \      /::\  \      /::\  \      /::\  \
@@ -169,7 +171,7 @@ class TreeMesh(
 
     diagonal_balance : bool, optional
         Whether to balance cells along the diagonal of the tree during construction.
-        This will effect all calls to refine the tree.
+        This will affect all calls to refine the tree.
 
     Examples
     --------
@@ -233,9 +235,20 @@ class TreeMesh(
     _items = {"h", "origin", "cell_state"}
 
     # inheriting stuff from BaseTensorMesh that isn't defined in _QuadTree
-    def __init__(self, h=None, origin=None, diagonal_balance=False, **kwargs):
+    def __init__(self, h=None, origin=None, diagonal_balance=None, **kwargs):
         if "x0" in kwargs:
             origin = kwargs.pop("x0")
+
+        if diagonal_balance is None:
+            diagonal_balance = False
+            warnings.warn(
+                "In discretize v1.0 the TreeMesh will change the default value of "
+                "diagonal_balance to True, which will likely slightly change meshes you have"
+                "previously created. If you need to keep the current behavoir, explicitly set "
+                "diagonal_balance=False.",
+                FutureWarning,
+                stacklevel=2,
+            )
         super().__init__(h=h, origin=origin, diagonal_balance=diagonal_balance)
 
         cell_state = kwargs.pop("cell_state", None)
@@ -926,7 +939,7 @@ def face_z_divergence(self):  # NOQA D102
 
     def point2index(self, locs):  # NOQA D102
         # Documentation inherited from discretize.base.BaseMesh
-        return self.get_containing_cell_indexes(locs)
+        return self.get_containing_cells(locs)
 
     def cell_levels_by_index(self, indices):
         """Fast function to return a list of levels for the given cell indices.

tests/tree/test_tree.py

@@ -573,5 +573,22 @@ def test_insert_point(self):
         self.assertEqual(mesh1.nC, mesh2.nC)
 
 
+@pytest.mark.parametrize("dim", [2, 3])
+def test_cell_locator(dim):
+    mesh = discretize.TreeMesh((16, 16, 16)[:dim])
+    mesh.insert_cells([0.5, 0.5, 0.5][:dim], levels=-1)
+
+    query_point = [0.21, 0.42, 0.22][:dim]
+    identified_cell = mesh.point2index(query_point)
+    found = False
+    for i, cell in enumerate(mesh):
+        bounds = cell.bounds.reshape((-1, 2))
+
+        if np.all((bounds[:, 0] <= query_point) & (bounds[:, 1] >= query_point)):
+            found = True
+            assert i == identified_cell
+    assert found
+
+
 if __name__ == "__main__":
     unittest.main()