.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/spatial_indexing.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_spatial_indexing.py: Spatial indexing ================ The goal of a cell tree is to quickly locate cells of an unstructured mesh. Unstructured meshes are challenging in this regard: for a given point, we cannot simply compute a row and column number as we would for structured data such as rasters. The most straightforward procedure is checking every single cell, until we find the cell which contains the point we're looking for. This is clearly not efficient. A cell tree is bounding volume hierarchy (BVH) which may be used as a spatial index. A spatial index is a data structure to search a spatial object efficiently, without exhaustively checking every cell. The implementation in ``numba_celltree`` provides four ways to search the tree: * Locating single points * Locating bounding boxes * Locating convex polygons (e.g. cells of another mesh) * Locating line segments This example provides an introduction to searching a cell tree for each of these. We'll start by importing the required packages with matplotlib for plotting. .. GENERATED FROM PYTHON SOURCE LINES 27-36 .. code-block:: Python import os import matplotlib.pyplot as plt import numpy as np from matplotlib.collections import LineCollection os.environ["NUMBA_DISABLE_JIT"] = "1" # small examples, avoid JIT overhead from numba_celltree import CellTree2d, demo # noqa E402 .. GENERATED FROM PYTHON SOURCE LINES 37-38 Let's start with a rectangular mesh: .. GENERATED FROM PYTHON SOURCE LINES 38-44 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 45-46 Determine the edges of the cells, and plot them. .. GENERATED FROM PYTHON SOURCE LINES 46-52 .. code-block:: Python 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") .. image-sg:: /examples/images/sphx_glr_spatial_indexing_001.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 53-57 Locating points --------------- We'll build a cell tree first, then look for some points. .. GENERATED FROM PYTHON SOURCE LINES 57-68 .. code-block:: Python tree = CellTree2d(vertices, faces, -1) points = np.array( [ [-5.0, 1.0], [4.5, 2.5], [6.5, 4.5], ] ) i = tree.locate_points(points) i .. rst-class:: sphx-glr-script-out .. code-block:: none array([-1, 24, 46]) .. GENERATED FROM PYTHON SOURCE LINES 69-75 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. .. GENERATED FROM PYTHON SOURCE LINES 75-82 .. code-block:: Python 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) .. image-sg:: /examples/images/sphx_glr_spatial_indexing_002.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_002.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 83-84 Now let's try a more exotic example. .. GENERATED FROM PYTHON SOURCE LINES 84-93 .. code-block:: Python vertices, faces = demo.generate_disk(5, 5) vertices += 1.0 vertices *= 5.0 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") .. image-sg:: /examples/images/sphx_glr_spatial_indexing_003.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_003.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 94-97 There are certainly no rows or columns to speak of! Let's build a new tree, and look for the same points as before. .. GENERATED FROM PYTHON SOURCE LINES 97-106 .. code-block:: Python tree = CellTree2d(vertices, faces, -1) i = tree.locate_points(points) 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) .. image-sg:: /examples/images/sphx_glr_spatial_indexing_004.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_004.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 107-122 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 one answer per cell. This is not the case for line segments, bounding boxes, or convex polygons: a line may intersect multiple cells, and a bounding box or polygon may contain multiple cells. Locating bounding boxes ----------------------- 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. .. GENERATED FROM PYTHON SOURCE LINES 122-136 .. code-block:: Python box_coords = np.array( [ [4.0, 8.0, 4.0, 6.0], # xmin, xmax, ymin, ymax ] ) box_i, cell_i = tree.locate_boxes(box_coords) fig, ax = plt.subplots() demo.plot_edges(node_x, node_y, edges, ax, color="black") demo.plot_edges( node_x, node_y, demo.edges(faces[cell_i], -1), ax, color="blue", linewidth=2 ) demo.plot_boxes(box_coords, ax, color="red", linewidth=3) .. image-sg:: /examples/images/sphx_glr_spatial_indexing_005.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_005.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 137-138 We can also search for multiple boxes: .. GENERATED FROM PYTHON SOURCE LINES 138-148 .. code-block:: Python box_coords = np.array( [ [4.0, 8.0, 4.0, 6.0], [0.0, 8.0, 8.0, 10.0], [10.0, 13.0, 2.0, 8.0], ] ) box_i, cell_i = tree.locate_boxes(box_coords) box_i, cell_i .. rst-class:: sphx-glr-script-out .. code-block:: none (array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]), array([ 36, 37, 33, 34, 35, 84, 32, 89, 87, 90, 39, 20, 97, 99, 98, 101, 102, 88, 44, 100, 31, 92, 91, 56, 55, 57, 58, 63, 62, 64, 74, 23, 24, 76, 75, 25, 26, 29, 30, 79, 80, 27, 28, 81, 82, 70, 69, 16, 17, 68, 14, 15, 65])) .. GENERATED FROM PYTHON SOURCE LINES 149-154 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: .. GENERATED FROM PYTHON SOURCE LINES 154-167 .. code-block:: Python cells_0 = cell_i[box_i == 0] cells_1 = cell_i[box_i == 1] fig, ax = plt.subplots() demo.plot_edges(node_x, node_y, edges, ax, color="black") demo.plot_edges( node_x, node_y, demo.edges(faces[cells_0], -1), ax, color="blue", linewidth=2 ) demo.plot_edges( node_x, node_y, demo.edges(faces[cells_1], -1), ax, color="green", linewidth=2 ) demo.plot_boxes(box_coords, ax, color="red", linewidth=3) .. image-sg:: /examples/images/sphx_glr_spatial_indexing_006.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_006.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 168-179 Locating cells -------------- Similarly, we can look for other cells (convex polygons) and compute the overlap: This returns three arrays of equal length: * the index of the face to locate * the index of the face in the celtree * the area of the intersection .. GENERATED FROM PYTHON SOURCE LINES 179-211 .. code-block:: Python triangle_vertices = np.array( [ [5.0, 3.0], [7.0, 3.0], [7.0, 5.0], [0.0, 6.0], [4.0, 4.0], [6.0, 10.0], ] ) triangles = np.array( [ [0, 1, 2], [3, 4, 5], ] ) tri_x, tri_y = triangle_vertices.transpose() tri_i, cell_i, area = tree.intersect_faces(triangle_vertices, triangles, -1) cells_0 = cell_i[tri_i == 0] cells_1 = cell_i[tri_i == 1] fig, ax = plt.subplots() demo.plot_edges(node_x, node_y, edges, ax, color="black") demo.plot_edges( node_x, node_y, demo.edges(faces[cells_0], -1), ax, color="blue", linewidth=2 ) demo.plot_edges( node_x, node_y, demo.edges(faces[cells_1], -1), ax, color="green", linewidth=2 ) demo.plot_edges(tri_x, tri_y, demo.edges(triangles, -1), ax, color="red", linewidth=3) .. image-sg:: /examples/images/sphx_glr_spatial_indexing_007.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_007.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 212-216 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. .. GENERATED FROM PYTHON SOURCE LINES 216-234 .. code-block:: Python intersection_faces = faces[cell_i] intersection_vertices = vertices[intersection_faces] U = intersection_vertices[:, 1] - intersection_vertices[:, 0] V = intersection_vertices[:, 2] - intersection_vertices[:, 0] full_area = 0.5 * np.abs(np.cross(V, U)) ratio = area / full_area fig, ax = plt.subplots() colored = ax.tripcolor( node_x, node_y, intersection_faces, ratio, ) fig.colorbar(colored) 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) .. image-sg:: /examples/images/sphx_glr_spatial_indexing_008.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_008.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 235-238 ``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. .. GENERATED FROM PYTHON SOURCE LINES 238-258 .. code-block:: Python dx = 1.0 xmin = 0.0 xmax = 10.0 dy = -1.0 ymin = 0.0 ymax = 10.0 y, x = np.meshgrid( np.arange(ymax, ymin + dy, dy), np.arange(xmin, xmax + dx, dx), indexing="ij", ) ny = y.shape[0] - 1 nx = x.shape[1] - 1 coords = np.column_stack( [a.ravel() for a in [x[:-1, :-1], x[1:, 1:], y[1:, 1:], y[:-1, :-1]]] ) raster_i, cell_i, raster_overlap = tree.intersect_boxes(coords) .. GENERATED FROM PYTHON SOURCE LINES 259-261 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). .. GENERATED FROM PYTHON SOURCE LINES 261-268 .. code-block:: Python 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) .. image-sg:: /examples/images/sphx_glr_spatial_indexing_009.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_009.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 269-272 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. .. GENERATED FROM PYTHON SOURCE LINES 272-298 .. code-block:: Python def saddle_elevation(x, y): return np.sin(0.6 * x + 2) + np.sin(0.2 * y) # Generate an elevation for every triangle centroid_x, centroid_y = vertices[faces].mean(axis=1).transpose() face_z = saddle_elevation(centroid_x, centroid_y) # Compute the weighted mean of the triangles per raster cell weighted_elevation = np.dot(weight_matrix, face_z) area_sum = np.dot(weight_matrix, np.ones(len(faces))) mean_elevation = np.full(ny * nx, np.nan) intersects = area_sum > 0 mean_elevation[intersects] = weighted_elevation[intersects] / area_sum[intersects] fig = plt.figure(figsize=(10, 4)) ax0 = fig.add_subplot(1, 2, 1, projection="3d") ax0.plot_trisurf( node_x, node_y, faces, saddle_elevation(node_x, node_y), cmap="viridis" ) ax1 = fig.add_subplot(1, 2, 2) ax1.imshow(mean_elevation.reshape(ny, nx), extent=(xmin, xmax, ymin, ymax)) demo.plot_edges(node_x, node_y, edges, ax1, color="white") .. image-sg:: /examples/images/sphx_glr_spatial_indexing_010.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_010.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 299-323 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 weight matrix are 0; a raster cell only intersects a small number triangles. Such a matrix is much more efficiently stored and processed as a `sparse matrix `_ (see also Scipy `sparse `_). The arrays returned by the ``intersect_`` methods of ``CellTree2d`` form a coordinate list (COO). Such a coordinate list can also be easily used to aggregate values with `Pandas `_, as Pandas provides very efficient aggregation in the form of `groupby operations `_. Locating lines -------------- As a final example, we will compute the intersections with two lines (edges). This once again returns three arrays of equal length: * the index of the line * the index of the cell * the location of the intersection .. GENERATED FROM PYTHON SOURCE LINES 323-332 .. code-block:: Python edge_coords = np.array( [ [[0.0, 0.0], [10.0, 10.0]], [[0.0, 10.0], [10.0, 0.0]], ] ) edge_i, cell_i, intersections = tree.intersect_edges(edge_coords) edge_i, cell_i .. rst-class:: sphx-glr-script-out .. code-block:: none (array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]), array([ 71, 23, 24, 25, 84, 86, 89, 87, 90, 55, 63, 64, 47, 60, 118, 116, 117, 54, 105, 41, 40, 106, 103, 98, 101, 88, 100, 16, 31, 92, 94, 96, 91, 14, 15, 65])) .. GENERATED FROM PYTHON SOURCE LINES 333-336 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): .. GENERATED FROM PYTHON SOURCE LINES 336-348 .. code-block:: Python length = np.linalg.norm(intersections[:, 1] - intersections[:, 0], axis=1) fig, ax = plt.subplots() colored = ax.tripcolor( node_x, node_y, faces[cell_i], length, ) fig.colorbar(colored) ax.add_collection(LineCollection(edge_coords, color="red", linewidth=3)) demo.plot_edges(node_x, node_y, edges, ax, color="black") .. image-sg:: /examples/images/sphx_glr_spatial_indexing_011.png :alt: spatial indexing :srcset: /examples/images/sphx_glr_spatial_indexing_011.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.974 seconds) .. _sphx_glr_download_examples_spatial_indexing.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: spatial_indexing.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: spatial_indexing.py ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_