.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/spatial_indexing_2d_grids.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_spatial_indexing_2d_grids.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_spatial_indexing_2d_grids.py:


Spatial indexing of 2D grids
============================

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-37

.. 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 38-39

Let's start with a rectangular mesh:

.. GENERATED FROM PYTHON SOURCE LINES 39-46

.. 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 47-48

Determine the edges of the cells, and plot them.

.. GENERATED FROM PYTHON SOURCE LINES 48-55

.. 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_2d_grids_001.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 56-60

Locating points
---------------

We'll build a cell tree first, then look for some points.

.. GENERATED FROM PYTHON SOURCE LINES 60-72

.. 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 73-79

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 79-87

.. 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_2d_grids_002.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_002.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 88-89

Now let's try a more exotic example.

.. GENERATED FROM PYTHON SOURCE LINES 89-98

.. 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_2d_grids_003.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_003.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 99-102

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 102-112

.. 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_2d_grids_004.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_004.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 113-128

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 128-143

.. 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_2d_grids_005.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_005.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 144-145

We can also search for multiple boxes:

.. GENERATED FROM PYTHON SOURCE LINES 145-155

.. 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 156-161

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 161-175

.. 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_2d_grids_006.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_006.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 176-187

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 187-220

.. 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_2d_grids_007.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_007.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 221-225

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 225-244

.. 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(U[:, 0] * V[:, 1] - U[:, 1] * V[:, 0])
    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_2d_grids_008.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_008.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 245-248

``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 248-268

.. 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 269-271

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 271-279

.. 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_2d_grids_009.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_009.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 280-283

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 283-309

.. 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_2d_grids_010.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_010.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 310-334

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 <https://en.wikipedia.org/wiki/Sparse_matrix>`_ (see also Scipy
`sparse <https://docs.scipy.org/doc/scipy/reference/sparse.html>`_). 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 <https://pandas.pydata.org/docs/index.html>`_, as Pandas provides
very efficient aggregation in the form of `groupby operations
<https://pandas.pydata.org/docs/reference/groupby.html>`_.

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 334-343

.. 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 344-347

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 347-360

.. 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_2d_grids_011.png
   :alt: spatial indexing 2d grids
   :srcset: /examples/images/sphx_glr_spatial_indexing_2d_grids_011.png
   :class: sphx-glr-single-img






.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 0.764 seconds)


.. _sphx_glr_download_examples_spatial_indexing_2d_grids.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_2d_grids.ipynb <spatial_indexing_2d_grids.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: spatial_indexing_2d_grids.py <spatial_indexing_2d_grids.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: spatial_indexing_2d_grids.zip <spatial_indexing_2d_grids.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_