celltree locate_points: use point_in_polygon_or_on_edge predicate ins…

…tead

Otherwise points falling on internal edges may be identified as falling outside the grid

examples/spatial_indexing.py

@@ -24,7 +24,6 @@
 
 We'll start by importing the required packages with matplotlib for plotting.
 """
-
 import os
 
 import matplotlib.pyplot as plt
@@ -34,27 +33,30 @@
 os.environ["NUMBA_DISABLE_JIT"] = "1"  # small examples, avoid JIT overhead
 from numba_celltree import CellTree2d, demo  # noqa E402
 
-###############################################################################
+# %%
 # Let's start with a rectangular mesh:
+
 nx = ny = 10
 x = y = np.linspace(0.0, 10.0, nx + 1)
 vertices = np.array(np.meshgrid(x, y, indexing="ij")).reshape(2, -1).T
 a = np.add.outer(np.arange(nx), nx * np.arange(ny)) + np.arange(ny)
 faces = np.array([a, a + 1, a + nx + 2, a + nx + 1]).reshape(4, -1).T
 
-###############################################################################
+# %%
 # Determine the edges of the cells, and plot them.
+
 node_x, node_y = vertices.transpose()
 edges = demo.edges(faces, -1)
 
 fig, ax = plt.subplots()
 demo.plot_edges(node_x, node_y, edges, ax, color="black")
 
-###############################################################################
+# %%
 # Locating points
 # ---------------
 #
 # We'll build a cell tree first, then look for some points.
+
 tree = CellTree2d(vertices, faces, -1)
 points = np.array(
     [
@@ -66,21 +68,22 @@
 i = tree.locate_points(points)
 i
 
-###############################################################################
+# %%
 # These numbers are the cell numbers in which we can find the points.
 #
 # A value of -1 means that a point is not located in any cell.
 #
 # Let's get rid of the -1 values, and take a look which cells have been found.
 # We'll color the found cells blue, and we'll draw the nodes to compare.
+
 i = i[i != -1]
 
 fig, ax = plt.subplots()
 ax.scatter(*points.transpose())
 demo.plot_edges(node_x, node_y, edges, ax, color="black")
 demo.plot_edges(node_x, node_y, demo.edges(faces[i], -1), ax, color="blue", linewidth=3)
 
-###############################################################################
+# %%
 # Now let's try a more exotic example.
 vertices, faces = demo.generate_disk(5, 5)
 vertices += 1.0
@@ -91,10 +94,11 @@
 fig, ax = plt.subplots()
 demo.plot_edges(node_x, node_y, edges, ax, color="black")
 
-###############################################################################
+# %%
 # There are certainly no rows or columns to speak of!
 #
 # Let's build a new tree, and look for the same points as before.
+
 tree = CellTree2d(vertices, faces, -1)
 i = tree.locate_points(points)
 i = i[i != -1]
@@ -104,7 +108,7 @@
 demo.plot_edges(node_x, node_y, edges, ax, color="black")
 demo.plot_edges(node_x, node_y, demo.edges(faces[i], -1), ax, color="blue", linewidth=3)
 
-###############################################################################
+# %%
 # It should be clear by now that a point may only fall into a single cell. A
 # point may also be out of bounds. If a cell falls exactly on an edge, one of the
 # two neighbors will be chosen arbitrarily. At any rate, we can always expect
@@ -120,6 +124,7 @@
 # A search of N points will yield N answers (cell numbers). A search of N boxes
 # may yield M answers. To illustrate, let's look for all the cells inside of
 # a box.
+
 box_coords = np.array(
     [
         [4.0, 8.0, 4.0, 6.0],  # xmin, xmax, ymin, ymax
@@ -134,7 +139,7 @@
 )
 demo.plot_boxes(box_coords, ax, color="red", linewidth=3)
 
-###############################################################################
+# %%
 # We can also search for multiple boxes:
 box_coords = np.array(
     [
@@ -146,12 +151,13 @@
 box_i, cell_i = tree.locate_boxes(box_coords)
 box_i, cell_i
 
-###############################################################################
+# %%
 # Note that this method returns two arrays of equal length. The second array
 # contains the cell numbers, as usual. The first array contains the index of
 # the bounding box in which the respective cells fall. Note that there are only
 # two numbers in ``box_i``: there are no cells located in the third box, as we
 # can confirm visually:
+
 cells_0 = cell_i[box_i == 0]
 cells_1 = cell_i[box_i == 1]
 
@@ -165,7 +171,7 @@
 )
 demo.plot_boxes(box_coords, ax, color="red", linewidth=3)
 
-###############################################################################
+# %%
 # Locating cells
 # --------------
 #
@@ -177,6 +183,7 @@
 # * the index of the face to locate
 # * the index of the face in the celtree
 # * the area of the intersection
+
 triangle_vertices = np.array(
     [
         [5.0, 3.0],
@@ -209,11 +216,12 @@
 )
 demo.plot_edges(tri_x, tri_y, demo.edges(triangles, -1), ax, color="red", linewidth=3)
 
-###############################################################################
+# %%
 # Let's color the faces of the mesh by their ratio of overlap. Because our
 # mesh is triangular, we can represent the triangles as two collections of
 # vectors (V, U). Then the area is half of the absolute value of the cross
 # product of U and V.
+
 intersection_faces = faces[cell_i]
 intersection_vertices = vertices[intersection_faces]
 U = intersection_vertices[:, 1] - intersection_vertices[:, 0]
@@ -232,7 +240,7 @@
 demo.plot_edges(node_x, node_y, edges, ax, color="black")
 demo.plot_edges(tri_x, tri_y, demo.edges(triangles, -1), ax, color="red", linewidth=3)
 
-###############################################################################
+# %%
 # ``CellTree2d`` also provides a method to compute overlaps between boxes and a
 # mesh. This may come in handy to compute overlap with a raster, for example to
 # rasterize a mesh.
@@ -256,17 +264,18 @@
 )
 raster_i, cell_i, raster_overlap = tree.intersect_boxes(coords)
 
-###############################################################################
+# %%
 # We can construct a weight matrix with these arrays. This weight matrix stores
 # for every raster cell (row) the area of overlap with a triangle (column).
+
 weight_matrix = np.zeros((ny * nx, len(faces)))
 weight_matrix[raster_i, cell_i] = raster_overlap
 
 fig, ax = plt.subplots()
 colored = ax.imshow(weight_matrix)
 _ = fig.colorbar(colored)
 
-###############################################################################
+# %%
 # This weight matrix can be used for translating data from one mesh to another.
 # Let's generate some mock elevation data for a valley. Then, we'll compute the
 # area weighted mean for every raster cell.
@@ -296,7 +305,7 @@ def saddle_elevation(x, y):
 ax1.imshow(mean_elevation.reshape(ny, nx), extent=(xmin, xmax, ymin, ymax))
 demo.plot_edges(node_x, node_y, edges, ax1, color="white")
 
-###############################################################################
+# %%
 # Such a weight matrix doesn't apply to just boxes and triangles, but to every
 # case of mapping one mesh to another by intersecting cell areas. Note however
 # that the aggregation above is not very efficient. Most of the entries in the
@@ -330,7 +339,7 @@ def saddle_elevation(x, y):
 edge_i, cell_i, intersections = tree.intersect_edges(edge_coords)
 edge_i, cell_i
 
-###############################################################################
+# %%
 # To wrap up, we'll color the intersect faces with the length of the
 # intersected line segments. We can easily compute the length of each segment
 # with the Euclidian norm (Pythagorean distance):
@@ -346,3 +355,5 @@ def saddle_elevation(x, y):
 fig.colorbar(colored)
 ax.add_collection(LineCollection(edge_coords, color="red", linewidth=3))
 demo.plot_edges(node_x, node_y, edges, ax, color="black")
+
+# %%

numba_celltree/celltree.py

@@ -148,6 +148,9 @@ def locate_points(self, points: FloatArray) -> IntArray:
         """
         Find the index of a face that contains a point.
 
+        Points that are very close near an edge of a face will also be
+        identified as falling within that face.
+
         Parameters
         ----------
         points: ndarray of floats with shape ``(n_point, 2)``

numba_celltree/query.py

@@ -22,7 +22,7 @@
     box_contained,
     boxes_intersect,
     copy_vertices_into,
-    point_in_polygon,
+    point_in_polygon_or_on_edge,
     to_vector,
 )
 from numba_celltree.utils import allocate_polygon, allocate_stack, pop, push
@@ -49,7 +49,7 @@ def locate_point(point: Point, tree: CellTreeData):
                 # Make sure polygons to test is contiguous (stack allocated) array
                 # This saves about 40-50% runtime
                 poly = copy_vertices_into(tree.vertices, face, polygon_work_array)
-                if point_in_polygon(point, poly):
+                if point_in_polygon_or_on_edge(point, poly):
                     return bbox_index
             continue
 

tests/test_celltree.py

@@ -538,12 +538,12 @@ def test_compute_barycentric_weights_triangle_quads():
     )
     face_indices, weights = tree.compute_barycentric_weights(points)
 
-    expected_indices = np.array([0, 0, 1])
+    expected_indices = np.array([0, 0, 0])
     expected_weights = np.array(
         [
             [1.0, 0.0, 0.0, 0.0],
             [0.25, 0.25, 0.25, 0.25],
-            [0.5, 0.0, 0.0, 0.5],
+            [0.0, 0.5, 0.5, 0.0],
         ]
     )
     assert np.array_equal(face_indices, expected_indices)
@@ -599,3 +599,71 @@ def test_to_dict_of_lists(datadir):
     assert isinstance(d, dict)
     assert list(d.keys()) == list(range(len(tree.celltree_data.nodes)))
     assert max(len(v) for v in d.values()) == 2
+
+
+def test_locate_point_on_edge():
+    nodes = np.array(
+        [
+            [0.0, 0.0],
+            [3.0, 0.0],
+            [1.0, 1.0],
+            [0.0, 2.0],
+            [3.0, 2.0],
+        ]
+    )
+    faces = np.array(
+        [
+            [0, 1, 2],
+            [0, 2, 3],
+            [2, 4, 3],
+        ]
+    )
+    fill_value = -1
+    tree = CellTree2d(nodes, faces, fill_value, n_buckets=4)
+    points = np.array(
+        [
+            [0.0, 0.0],
+            [0.01, 0.01],
+            [0.05, 0.05],
+            [0.15, 0.15],
+            [0.25, 0.25],
+            [0.35, 0.35],
+            [0.45, 0.45],
+            [0.55, 0.55],
+            [0.65, 0.65],
+            [0.75, 0.75],
+        ]
+    )
+    result = tree.locate_points(points)
+    assert (result != -1).all()
+
+    nodes = np.array(
+        [
+            [0.0, 0.0],  # 0
+            [2.0, 0.0],  # 1
+            [2.0, 2.0],  # 2
+            [0.0, 2.0],  # 3
+            [4.0, 0.0],  # 4
+            [4.0, 4.0],  # 5
+            [0.0, 4.0],  # 6
+        ]
+    )
+    faces = np.array(
+        [
+            [0, 1, 2, 3],
+            [1, 4, 5, 2],
+            [3, 2, 5, 6],
+        ]
+    )
+    fill_value = -1
+    tree = CellTree2d(nodes, faces, fill_value, n_buckets=4)
+    points = np.array(
+        [
+            [0.0, 0.0],
+            [0.0, 4.0],
+            [2.0, 2.0],
+            [4.0, 0.0],
+            [4.0, 4.0],
+        ]
+    )
+    assert (result != -1).all()