Add bandwidth-minimising mesh sorting routine #515

dengwirda · 2023-07-06T19:34:37Z

This PR adds a mesh sorting routine mpas_tools.mesh.creation.sort_mesh that can be used to reindex an MPAS mesh file to bring the cells, edges and vertices into a more cache-local sequence. This undoes the somewhat 'random' ordering introduced by the mesh conversion tools, and may be desirable re: efficiency for downstream dycore + coupler applications.

The reordering is implemented via a reverse cuthill-mckee approach --- minimising the bandwidth of the cell-to-cell adjacency graph and then permuting the edges and vertices to match. Sorting both culled and unculled meshes is supported, as well as meshes containing additional variables, which are reindexed according to their underlying nCells, nEdges or nVertices dimensions.

The mesh sorter should be callable as a cmd-line tool:

python sort_mesh.py \
--mesh-file=path+name-to-mpas-file-for-sort \
--sort-file=path+name-to-sorted-output-file

and/or inline:

from mpas_tools.mesh.creation.sort_mesh import sort_mesh
mesh = xarray.open_dataset("mpas-file-name")
sort_mesh(mesh)

Tested on pm-cpu using a QU30km MPAS-O standalone config., which successfully spun-up to 90 days after a sort of the culled mesh. Sorting improves the locality of the various remapping arrays (verticesOnCell, edgesOnVertex, etc, etc) with all showing smooth distributions.

edgesOnCell=unsorted:

edgesOnCell=issorted:

dengwirda · 2023-07-06T19:41:55Z

Just keeping this in draft for now @xylar --- before going too far I think we should check whether this reindexing will solve the locality issues wrt. the E3SM coupler and I've reached out to the team re: benchmark cases.

xylar · 2023-07-06T19:52:43Z

@dengwirda. beautiful! Keep me posted.

xylar · 2023-07-06T19:54:16Z

Picasso, by the way, is turning over in his grave. The unsorted way it clearly more cubist.

xylar

@dengwirda, I have some suggestions that are purely aesthetic to do with integrating this code a bit more with the style of mpas_tools overall. None of this has any bearing on functionality.

I'm going to cherry-pick the current commit onto a branch I have for making release candidates for testing in compass. I'll let you know how that testing goes.

xylar · 2023-07-26T14:03:21Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+    """
+    SORT-MESH: sort cells, edges and duals in the mesh
+    to improve cache-locality.
+
+    """


Here's the docstring format we use in MPAS-Tools:

Suggested change

"""

SORT-MESH: sort cells, edges and duals in the mesh

to improve cache-locality.

"""

"""

Sort cells, edges and duals in the mesh

to improve cache-locality.

Parameters

----------

mesh : xarray.Dataset

A dataset containing an MPAS mesh to sort

Returns

-------

mesh : xarray.Dataset

A dataset containing the sorted MPAS mesh

"""

xylar · 2023-07-26T14:03:53Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+from scipy.sparse.csgraph import reverse_cuthill_mckee
+
+
+def sort_fwd(data, fwd):


Let's indicate that this is "private":

Suggested change

def sort_fwd(data, fwd):

def _sort_fwd(data, fwd):

Let's also move it after sort_mesh(). Similarly for other "private" functions.

Maybe add a very brief docstring to at least say what the function does.

xylar · 2023-07-26T14:05:18Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+    return vals
+
+
+def sort_rev(data, rev):


Suggested change

def sort_rev(data, rev):

def _sort_rev(data, rev):

Make this private, move it below the "public" functions and add a brief docstring.

xylar · 2023-07-26T14:06:00Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+    for var in mesh.keys():
+        dims = mesh.variables[var].dims
+        if ("nCells" in dims):
+            mesh[var][:] = sort_fwd(mesh[var], cell_fwd)


Suggested change

mesh[var][:] = sort_fwd(mesh[var], cell_fwd)

mesh[var][:] = _sort_fwd(mesh[var], cell_fwd)

xylar · 2023-07-26T14:06:08Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+    for var in mesh.keys():
+        dims = mesh.variables[var].dims
+        if ("nVertices" in dims):
+            mesh[var][:] = sort_fwd(mesh[var], dual_fwd)


Suggested change

mesh[var][:] = sort_fwd(mesh[var], dual_fwd)

mesh[var][:] = _sort_fwd(mesh[var], dual_fwd)

xylar · 2023-07-26T14:08:03Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+    return mesh
+
+
+def cell_del2(mesh):


Suggested change

def cell_del2(mesh):

def _cell_del2(mesh):

Also, move this below main() along with the other "private" functions.

Alternatively, if you want to keep this "public", it should have a fleshed out docstring like I suggested for sort_mesh().

xylar · 2023-07-26T14:08:11Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+
+def cell_del2(mesh):
+    """
+    CELL-DEL2: form cell-to-cell sparse adjacency graph


Suggested change

CELL-DEL2: form cell-to-cell sparse adjacency graph

form cell-to-cell sparse adjacency graph

xylar · 2023-07-26T14:08:50Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+
+    # sort cells via RCM ordering of adjacency matrix
+
+    cell_fwd = reverse_cuthill_mckee(cell_del2(mesh)) + 1


Suggested change

cell_fwd = reverse_cuthill_mckee(cell_del2(mesh)) + 1

cell_fwd = reverse_cuthill_mckee(_cell_del2(mesh)) + 1

xylar · 2023-07-26T14:10:59Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+        ncells = mesh.dims["nCells"]
+        nedges = np.count_nonzero(cellsOnCell) // 2
+
+        fptr.write("{} {}\n".format(ncells, nedges))


Going forward, we have a pretty strong preference for f-strings over using the .format() method.

Suggested change

fptr.write("{} {}\n".format(ncells, nedges))

fptr.write(f"{ncells} {nedges}\n")

xylar · 2023-07-26T14:11:17Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+            data = cellsOnCell[cell, :]
+            data = data[data > 0]
+            for item in data:
+                fptr.write("{} ".format(item))


Suggested change

fptr.write("{} ".format(item))

fptr.write(f"{item} ")

xylar · 2023-07-26T14:21:23Z

One more thing: It would be good to add the entry point (redundantly, I know) to this file:
https://github.com/MPAS-Dev/MPAS-Tools/blob/master/conda_package/recipe/meta.yaml
along with a test to make sure you can at least get the help message:
https://github.com/MPAS-Dev/MPAS-Tools/blob/master/conda_package/recipe/meta.yaml#L140
Or if you want, you could test it properly like one of these:
https://github.com/MPAS-Dev/MPAS-Tools/blob/master/conda_package/recipe/meta.yaml#L103

I added the entry point and a test like I'm suggesting to the 0.22.0rc2 recipe on conda-forge:
https://github.com/conda-forge/mpas_tools-feedstock/pull/77/files#diff-f3725a55bf339595bf865fec73bda8ac99f283b0810c205442021f29c06eea9a

xylar · 2023-07-26T15:07:21Z

conda_package/mpas_tools/mesh/creation/sort_mesh.py

+    return csr_matrix((xvec, (ivec, jvec)))
+
+
+if (__name__ == "__main__"):


To work properly as an entry point, this needs to be a function:

Suggested change

if (__name__ == "__main__"):

def main():

Then, at the bottom of the file, you can have:

if (__name__ == "__main__"):\ main()

xylar · 2023-07-27T15:35:52Z

I tested things out in MPAS-Dev/compass#657 and everything worked well. I need to visualize the edgesOnCell, etc. to make sure I'm seeing something similar to the plots above but the mechanics seem to work well.

@dengwirda would you like to make the changes I mentioned above? If so, that would be my preference but I'm also happy to make them if you don't have time.

xylar · 2023-08-01T11:08:06Z

@dengwirda, using Paraview, I looked at all the following fields from the ECwISC mesh I tested earlier:

cellsOnCell, edgesOnCell, verticesOnCell, cellsOnVertex, cellsOnEdge, edgesOnVertex, verticesOnEdge, edgesOnEdge

and they look great to me. I won't post a bunch of identical-looking plots here to prove it.

So I think this is good to go as soon as my earlier clean-up comments are addressed.

dengwirda · 2023-08-04T13:25:31Z

Thanks for the review @xylar, I believe I got all of those changes in but let me know if you need anything further.

xylar · 2023-08-04T13:35:37Z

@dengwirda, it looks excellent! Let me just make another test build with both this and the fixes I have been making in #514. If my testing goes well, I'll approve and merge this.

xylar

After using this sorting in MPAS-Dev/compass#657 and re-testing, I uncovered 2 unrelated issues that my testing of this work happened to expose:
#518
MPAS-Dev/compass#661
After those fixes, I was able to run successfully with this mesh sorting for QU240, QUwISC240, EC30to60 and ECwISC30to60 meshes all the way through files_for_e3sm in compass.

Add bandwidth-minimising mesh sorting routine

2956e63

dengwirda marked this pull request as ready for review July 26, 2023 13:24

dengwirda requested a review from xylar July 26, 2023 13:24

xylar requested changes Jul 26, 2023

View reviewed changes

xylar reviewed Jul 26, 2023

View reviewed changes

This comment was marked as resolved.

Sign in to view

Clean-up mesh sorting

42cd356

xylar mentioned this pull request Aug 7, 2023

Fix MOC edge sign and mask #518

Merged

xylar approved these changes Aug 8, 2023

View reviewed changes

xylar merged commit c0bb8b5 into MPAS-Dev:master Aug 8, 2023
7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add bandwidth-minimising mesh sorting routine #515

Add bandwidth-minimising mesh sorting routine #515

dengwirda commented Jul 6, 2023

dengwirda commented Jul 6, 2023

xylar commented Jul 6, 2023

xylar commented Jul 6, 2023

xylar left a comment

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar Jul 26, 2023

xylar commented Jul 26, 2023 •

edited

Loading

xylar Jul 26, 2023 •

edited

Loading

xylar commented Jul 27, 2023

xylar commented Aug 1, 2023

This comment was marked as resolved.

dengwirda commented Aug 4, 2023

xylar commented Aug 4, 2023

xylar left a comment

-    """
-    SORT-MESH: sort cells, edges and duals in the mesh
-    to improve cache-locality.
-    """
+    """
+    Sort cells, edges and duals in the mesh
+    to improve cache-locality.
+    Parameters
+    ----------
+    mesh : xarray.Dataset
+        A dataset containing an MPAS mesh to sort
+    Returns
+    -------
+    mesh : xarray.Dataset
+        A dataset containing the sorted MPAS mesh
+    """

		from scipy.sparse.csgraph import reverse_cuthill_mckee


		def sort_fwd(data, fwd):

	mesh[var][:] = sort_fwd(mesh[var], cell_fwd)
	mesh[var][:] = _sort_fwd(mesh[var], cell_fwd)

	mesh[var][:] = sort_fwd(mesh[var], dual_fwd)
	mesh[var][:] = _sort_fwd(mesh[var], dual_fwd)

	CELL-DEL2: form cell-to-cell sparse adjacency graph
	form cell-to-cell sparse adjacency graph


		# sort cells via RCM ordering of adjacency matrix

		cell_fwd = reverse_cuthill_mckee(cell_del2(mesh)) + 1

	cell_fwd = reverse_cuthill_mckee(cell_del2(mesh)) + 1
	cell_fwd = reverse_cuthill_mckee(_cell_del2(mesh)) + 1

	fptr.write("{} {}\n".format(ncells, nedges))
	fptr.write(f"{ncells} {nedges}\n")

		return csr_matrix((xvec, (ivec, jvec)))


		if (__name__ == "__main__"):

Add bandwidth-minimising mesh sorting routine #515

Add bandwidth-minimising mesh sorting routine #515

Conversation

dengwirda commented Jul 6, 2023

dengwirda commented Jul 6, 2023

xylar commented Jul 6, 2023

xylar commented Jul 6, 2023

xylar left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xylar commented Jul 26, 2023 • edited Loading

xylar Jul 26, 2023 • edited Loading

Choose a reason for hiding this comment

xylar commented Jul 27, 2023

xylar commented Aug 1, 2023

This comment was marked as resolved.

dengwirda commented Aug 4, 2023

xylar commented Aug 4, 2023

xylar left a comment

Choose a reason for hiding this comment

xylar commented Jul 26, 2023 •

edited

Loading

xylar Jul 26, 2023 •

edited

Loading