.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/partitioning.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_partitioning.py: Partitioning ============ Grid partitioning, or domain decomposition, is an important step in setting up parallellized simulations. Xugrid provides utilities for partitioning a grid and its associated data, and for merging partitions back into a single whole. .. GENERATED FROM PYTHON SOURCE LINES 10-16 .. code-block:: Python import matplotlib.pyplot as plt import numpy as np import xugrid as xu .. GENERATED FROM PYTHON SOURCE LINES 17-26 Create partitions ----------------- Xugrid wraps the well known `METIS library`_ via the `pymetis bindings`_. METIS is generally used to partition a grid in such a manner that communication between parallel processes is minimized. We'll demonstrate the functionality by diving the elevation example into several parts. .. GENERATED FROM PYTHON SOURCE LINES 26-30 .. code-block:: Python uda = xu.data.elevation_nl() uda.ugrid.plot(vmin=-20, vmax=90, cmap="terrain") .. image-sg:: /examples/images/sphx_glr_partitioning_001.png :alt: partitioning :srcset: /examples/images/sphx_glr_partitioning_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 31-32 Let's start by dividing the grid into four parts: .. GENERATED FROM PYTHON SOURCE LINES 32-39 .. code-block:: Python partitions = uda.ugrid.partition(n_part=4) fig, axes = plt.subplots(2, 2, figsize=(12.6, 10)) for partition, ax in zip(partitions, axes.ravel()): partition.ugrid.plot(ax=ax, vmin=-20, vmax=90, cmap="terrain") .. image-sg:: /examples/images/sphx_glr_partitioning_002.png :alt: partitioning :srcset: /examples/images/sphx_glr_partitioning_002.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 40-53 Partition the grid ------------------ Calling ``.partition`` on a UgridDataArray or UgridDataset will automatically partition the grid topology, select all associated data, and create a new UgridDataArray or UgridDataset for each partition. However, in some case, we might prefer to pre-compute the labels, and then apply them multiple datasets. To do so, we compute the partition labels from the grid. ``label_partitions`` returns a UgridDataArray, with every cell given its partition label number. We can easily plot this data to visualize the partitions: .. GENERATED FROM PYTHON SOURCE LINES 53-57 .. code-block:: Python labels = uda.ugrid.grid.label_partitions(n_part=12) labels.ugrid.plot() .. image-sg:: /examples/images/sphx_glr_partitioning_003.png :alt: partitioning :srcset: /examples/images/sphx_glr_partitioning_003.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 58-61 Not quite the twelve provinces of the Netherlands! However, we may use the labels to partition the data nonetheless: .. GENERATED FROM PYTHON SOURCE LINES 61-68 .. code-block:: Python partitions = uda.ugrid.partition_by_label(labels) fig, axes = plt.subplots(4, 3, figsize=(15, 15)) for partition, ax in zip(partitions, axes.ravel()): partition.ugrid.plot(ax=ax, vmin=-20, vmax=90, cmap="terrain") .. image-sg:: /examples/images/sphx_glr_partitioning_004.png :alt: partitioning :srcset: /examples/images/sphx_glr_partitioning_004.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 69-78 Since the labels are an ordinary UgridDataArray, we can easily store them in a netCDF file and re-use them in another part of a workflow. Merging partitions ------------------ Generally, after partitioning the data we write it as model input and run a model in parallel. Many model codes produce output per process. Xugrid can merge these partitions back into one whole for post-processing: .. GENERATED FROM PYTHON SOURCE LINES 78-83 .. code-block:: Python merged = xu.merge_partitions(partitions)["elevation"] merged.ugrid.plot(vmin=-20, vmax=90, cmap="terrain") .. image-sg:: /examples/images/sphx_glr_partitioning_005.png :alt: partitioning :srcset: /examples/images/sphx_glr_partitioning_005.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 84-89 Partitioning grids without data ------------------------------- Of course, we can also partition the grid topology without any associated data: .. GENERATED FROM PYTHON SOURCE LINES 89-97 .. code-block:: Python grid = uda.ugrid.grid grid_parts = grid.partition(n_part=4) fig, axes = plt.subplots(2, 2, figsize=(12.6, 10)) for part, ax in zip(grid_parts, axes.ravel()): part.plot(ax=ax) .. image-sg:: /examples/images/sphx_glr_partitioning_006.png :alt: partitioning :srcset: /examples/images/sphx_glr_partitioning_006.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 98-99 ... and merge them back into one whole: .. GENERATED FROM PYTHON SOURCE LINES 99-103 .. code-block:: Python merged_grid, _ = xu.Ugrid2d.merge_partitions(grid_parts) merged_grid.plot() .. image-sg:: /examples/images/sphx_glr_partitioning_007.png :alt: partitioning :srcset: /examples/images/sphx_glr_partitioning_007.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 104-108 Preserving order ---------------- Note that partioning and merging does not preserve order! .. GENERATED FROM PYTHON SOURCE LINES 108-111 .. code-block:: Python uda == merged .. raw:: html 
<xarray.DataArray 'elevation' (mesh2d_nFaces: 5248)> Size: 5kB
    array([False, False, False, ..., False, False, False])
    Coordinates:
      * mesh2d_nFaces  (mesh2d_nFaces) int64 42kB 0 1 2 3 4 ... 5244 5245 5246 5247


.. GENERATED FROM PYTHON SOURCE LINES 112-119 The topology is equivalent, but the nodes, edges, and faces are in a different order. This is because ``merge_partitions`` simply concatenates the partitions. The easiest way to restore the order is by providing an example of the original topology. ``reindex_like`` looks at the coordinates of both (equivalent!) grids and automatically determines how to reorder: .. GENERATED FROM PYTHON SOURCE LINES 119-123 .. code-block:: Python reordered = merged.ugrid.reindex_like(uda) uda == reordered .. raw:: html 
<xarray.DataArray 'elevation' (mesh2d_nFaces: 5248)> Size: 5kB
    array([ True,  True,  True, ...,  True,  True,  True])
    Coordinates:
        mesh2d_face_x  (mesh2d_nFaces) float64 42kB 2.388e+04 1.86e+05 ... 3.03e+04
        mesh2d_face_y  (mesh2d_nFaces) float64 42kB 3.648e+05 ... 3.964e+05
      * mesh2d_nFaces  (mesh2d_nFaces) int64 42kB 0 1 2 3 4 ... 5244 5245 5246 5247


.. GENERATED FROM PYTHON SOURCE LINES 124-126 Alternatively, we can also assign IDs, carry these along, and use these to reorder the data after merging. .. GENERATED FROM PYTHON SOURCE LINES 126-138 .. code-block:: Python uds = xu.UgridDataset(grids=[uda.ugrid.grid]) uds["elevation"] = uda uds["cell_id"] = ("mesh2d_nFaces", np.arange(len(uda))) partitions = uds.ugrid.partition(n_part=4) merged = xu.merge_partitions(partitions) order = np.argsort(merged["cell_id"].values) reordered = merged.isel(mesh2d_nFaces=order) uds["elevation"] == reordered["elevation"] .. raw:: html 
<xarray.DataArray 'elevation' (mesh2d_nFaces: 5248)> Size: 5kB
    array([ True,  True,  True, ...,  True,  True,  True])
    Coordinates:
        mesh2d_face_x  (mesh2d_nFaces) float64 42kB 2.388e+04 1.86e+05 ... 3.03e+04
        mesh2d_face_y  (mesh2d_nFaces) float64 42kB 3.648e+05 ... 3.964e+05
      * mesh2d_nFaces  (mesh2d_nFaces) int64 42kB 0 1 2 3 4 ... 5244 5245 5246 5247


.. GENERATED FROM PYTHON SOURCE LINES 139-145 This is required if results are compared with the input, or with results stemming from another partitioning, e.g. one with a different number of partitions. .. _METIS library: https://github.com/KarypisLab/METIS .. _pymetis bindings: https://github.com/inducer/pymetis .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 3.733 seconds) .. _sphx_glr_download_examples_partitioning.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: partitioning.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: partitioning.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: partitioning.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_