From 2956e63ee3f99beff06ee89393bd51ab5b0c8f37 Mon Sep 17 00:00:00 2001 From: dengwirda Date: Thu, 6 Jul 2023 10:25:59 -0600 Subject: [PATCH 1/2] Add bandwidth-minimising mesh sorting routine --- .../mpas_tools/mesh/creation/sort_mesh.py | 187 ++++++++++++++++++ conda_package/setup.py | 1 + 2 files changed, 188 insertions(+) create mode 100644 conda_package/mpas_tools/mesh/creation/sort_mesh.py diff --git a/conda_package/mpas_tools/mesh/creation/sort_mesh.py b/conda_package/mpas_tools/mesh/creation/sort_mesh.py new file mode 100644 index 000000000..f55375a5c --- /dev/null +++ b/conda_package/mpas_tools/mesh/creation/sort_mesh.py @@ -0,0 +1,187 @@ + +import numpy as np +import os +import xarray +import argparse +from scipy.sparse import csr_matrix +from scipy.sparse.csgraph import reverse_cuthill_mckee + + +def sort_fwd(data, fwd): + vals = data.values + vals = vals[fwd - 1] + return vals + + +def sort_rev(data, rev): + vals = data.values + mask = vals > 0 + vals[mask] = rev[vals[mask] - 1] + return vals + + +def sort_mesh(mesh): + """ + SORT-MESH: sort cells, edges and duals in the mesh + to improve cache-locality. + + """ + # Authors: Darren Engwirda + + ncells = mesh.dims["nCells"] + nedges = mesh.dims["nEdges"] + nduals = mesh.dims["nVertices"] + + cell_fwd = np.arange(0, ncells) + 1 + cell_rev = np.arange(0, ncells) + 1 + edge_fwd = np.arange(0, nedges) + 1 + edge_rev = np.arange(0, nedges) + 1 + dual_fwd = np.arange(0, nduals) + 1 + dual_rev = np.arange(0, nduals) + 1 + + # sort cells via RCM ordering of adjacency matrix + + cell_fwd = reverse_cuthill_mckee(cell_del2(mesh)) + 1 + + cell_rev = np.zeros(ncells, dtype=np.int32) + cell_rev[cell_fwd - 1] = np.arange(ncells) + 1 + + mesh["cellsOnCell"][:] = \ + sort_rev(mesh["cellsOnCell"], cell_rev) + mesh["cellsOnEdge"][:] = \ + sort_rev(mesh["cellsOnEdge"], cell_rev) + mesh["cellsOnVertex"][:] = \ + sort_rev(mesh["cellsOnVertex"], cell_rev) + + for var in mesh.keys(): + dims = mesh.variables[var].dims + if ("nCells" in dims): + mesh[var][:] = sort_fwd(mesh[var], cell_fwd) + + mesh["indexToCellID"][:] = np.arange(ncells) + 1 + + # sort duals via pseudo-linear cell-wise ordering + + dual_fwd = np.ravel(mesh["verticesOnCell"].values) + dual_fwd = dual_fwd[dual_fwd > 0] + + __, imap = np.unique(dual_fwd, return_index=True) + + dual_fwd = dual_fwd[np.sort(imap)] + + dual_rev = np.zeros(nduals, dtype=np.int32) + dual_rev[dual_fwd - 1] = np.arange(nduals) + 1 + + mesh["verticesOnCell"][:] = \ + sort_rev(mesh["verticesOnCell"], dual_rev) + + mesh["verticesOnEdge"][:] = \ + sort_rev(mesh["verticesOnEdge"], dual_rev) + + for var in mesh.keys(): + dims = mesh.variables[var].dims + if ("nVertices" in dims): + mesh[var][:] = sort_fwd(mesh[var], dual_fwd) + + mesh["indexToVertexID"][:] = np.arange(nduals) + 1 + + # sort edges via pseudo-linear cell-wise ordering + + edge_fwd = np.ravel(mesh["edgesOnCell"].values) + edge_fwd = edge_fwd[edge_fwd > 0] + + __, imap = np.unique(edge_fwd, return_index=True) + + edge_fwd = edge_fwd[np.sort(imap)] + + edge_rev = np.zeros(nedges, dtype=np.int32) + edge_rev[edge_fwd - 1] = np.arange(nedges) + 1 + + mesh["edgesOnCell"][:] = \ + sort_rev(mesh["edgesOnCell"], edge_rev) + + mesh["edgesOnEdge"][:] = \ + sort_rev(mesh["edgesOnEdge"], edge_rev) + + mesh["edgesOnVertex"][:] = \ + sort_rev(mesh["edgesOnVertex"], edge_rev) + + for var in mesh.keys(): + dims = mesh.variables[var].dims + if ("nEdges" in dims): + mesh[var][:] = sort_fwd(mesh[var], edge_fwd) + + mesh["indexToEdgeID"][:] = np.arange(nedges) + 1 + + return mesh + + +def cell_del2(mesh): + """ + CELL-DEL2: form cell-to-cell sparse adjacency graph + + """ + xvec = np.array([], dtype=np.int8) + ivec = np.array([], dtype=np.int32) + jvec = np.array([], dtype=np.int32) + + topolOnCell = mesh["nEdgesOnCell"].values + cellsOnCell = mesh["cellsOnCell"].values + + for edge in range(np.max(topolOnCell)): + + # cell-to-cell pairs, if edges exist + mask = topolOnCell > edge + idx_self = np.argwhere(mask).ravel() + idx_next = cellsOnCell[mask, edge] - 1 + + # cell-to-cell pairs, if cells exist + mask = idx_next >= 0 + idx_self = idx_self[mask] + idx_next = idx_next[mask] + + # dummy matrix values, just topol. needed + val_edge = np.ones(idx_next.size, dtype=np.int8) + + ivec = np.hstack((ivec, idx_self)) + jvec = np.hstack((jvec, idx_next)) + xvec = np.hstack((xvec, val_edge)) + + return csr_matrix((xvec, (ivec, jvec))) + + +if (__name__ == "__main__"): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawTextHelpFormatter) + + parser.add_argument( + "--mesh-file", dest="mesh_file", type=str, + required=True, help="Path+name to unsorted mesh file.") + + parser.add_argument( + "--sort-file", dest="sort_file", type=str, + required=True, help="Path+name to sorted output file.") + + args = parser.parse_args() + + mesh = xarray.open_dataset(args.mesh_file) + + sort_mesh(mesh) + + with open(os.path.join(os.path.dirname( + args.sort_file), "graph.info"), "w") as fptr: + cellsOnCell = mesh["cellsOnCell"].values + + ncells = mesh.dims["nCells"] + nedges = np.count_nonzero(cellsOnCell) // 2 + + fptr.write("{} {}\n".format(ncells, nedges)) + for cell in range(ncells): + data = cellsOnCell[cell, :] + data = data[data > 0] + for item in data: + fptr.write("{} ".format(item)) + fptr.write("\n") + + mesh.to_netcdf(args.sort_file, format="NETCDF4") diff --git a/conda_package/setup.py b/conda_package/setup.py index 82996002a..5506aadcb 100755 --- a/conda_package/setup.py +++ b/conda_package/setup.py @@ -89,6 +89,7 @@ 'mpas_to_triangle = mpas_tools.mesh.creation.mpas_to_triangle:main', 'triangle_to_netcdf = mpas_tools.mesh.creation.triangle_to_netcdf:main', 'jigsaw_to_netcdf = mpas_tools.mesh.creation.jigsaw_to_netcdf:main', + 'sort_mesh = mpas_tools.mesh.creation.sort_mesh:main', 'scrip_from_mpas = mpas_tools.scrip.from_mpas:main', 'compute_mpas_region_masks = mpas_tools.mesh.mask:entry_point_compute_mpas_region_masks', 'compute_mpas_transect_masks = mpas_tools.mesh.mask:entry_point_compute_mpas_transect_masks', From 42cd3563efdc045c9ce7de06c5249a9a89929afd Mon Sep 17 00:00:00 2001 From: dengwirda Date: Fri, 4 Aug 2023 06:38:00 -0600 Subject: [PATCH 2/2] Clean-up mesh sorting --- .../mpas_tools/mesh/creation/sort_mesh.py | 186 +++++++++++------- conda_package/recipe/meta.yaml | 2 + 2 files changed, 122 insertions(+), 66 deletions(-) diff --git a/conda_package/mpas_tools/mesh/creation/sort_mesh.py b/conda_package/mpas_tools/mesh/creation/sort_mesh.py index f55375a5c..2eb75ad79 100644 --- a/conda_package/mpas_tools/mesh/creation/sort_mesh.py +++ b/conda_package/mpas_tools/mesh/creation/sort_mesh.py @@ -7,24 +7,20 @@ from scipy.sparse.csgraph import reverse_cuthill_mckee -def sort_fwd(data, fwd): - vals = data.values - vals = vals[fwd - 1] - return vals - - -def sort_rev(data, rev): - vals = data.values - mask = vals > 0 - vals[mask] = rev[vals[mask] - 1] - return vals - - def sort_mesh(mesh): """ - 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 """ # Authors: Darren Engwirda @@ -41,22 +37,22 @@ def sort_mesh(mesh): # 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_rev = np.zeros(ncells, dtype=np.int32) cell_rev[cell_fwd - 1] = np.arange(ncells) + 1 mesh["cellsOnCell"][:] = \ - sort_rev(mesh["cellsOnCell"], cell_rev) + _sort_rev(mesh["cellsOnCell"], cell_rev) mesh["cellsOnEdge"][:] = \ - sort_rev(mesh["cellsOnEdge"], cell_rev) + _sort_rev(mesh["cellsOnEdge"], cell_rev) mesh["cellsOnVertex"][:] = \ - sort_rev(mesh["cellsOnVertex"], cell_rev) + _sort_rev(mesh["cellsOnVertex"], cell_rev) for var in mesh.keys(): dims = mesh.variables[var].dims if ("nCells" in dims): - mesh[var][:] = sort_fwd(mesh[var], cell_fwd) + mesh[var][:] = _sort_fwd(mesh[var], cell_fwd) mesh["indexToCellID"][:] = np.arange(ncells) + 1 @@ -73,15 +69,15 @@ def sort_mesh(mesh): dual_rev[dual_fwd - 1] = np.arange(nduals) + 1 mesh["verticesOnCell"][:] = \ - sort_rev(mesh["verticesOnCell"], dual_rev) + _sort_rev(mesh["verticesOnCell"], dual_rev) mesh["verticesOnEdge"][:] = \ - sort_rev(mesh["verticesOnEdge"], dual_rev) + _sort_rev(mesh["verticesOnEdge"], dual_rev) for var in mesh.keys(): dims = mesh.variables[var].dims if ("nVertices" in dims): - mesh[var][:] = sort_fwd(mesh[var], dual_fwd) + mesh[var][:] = _sort_fwd(mesh[var], dual_fwd) mesh["indexToVertexID"][:] = np.arange(nduals) + 1 @@ -98,29 +94,120 @@ def sort_mesh(mesh): edge_rev[edge_fwd - 1] = np.arange(nedges) + 1 mesh["edgesOnCell"][:] = \ - sort_rev(mesh["edgesOnCell"], edge_rev) + _sort_rev(mesh["edgesOnCell"], edge_rev) mesh["edgesOnEdge"][:] = \ - sort_rev(mesh["edgesOnEdge"], edge_rev) + _sort_rev(mesh["edgesOnEdge"], edge_rev) mesh["edgesOnVertex"][:] = \ - sort_rev(mesh["edgesOnVertex"], edge_rev) + _sort_rev(mesh["edgesOnVertex"], edge_rev) for var in mesh.keys(): dims = mesh.variables[var].dims if ("nEdges" in dims): - mesh[var][:] = sort_fwd(mesh[var], edge_fwd) + mesh[var][:] = _sort_fwd(mesh[var], edge_fwd) mesh["indexToEdgeID"][:] = np.arange(nedges) + 1 return mesh + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawTextHelpFormatter) + + parser.add_argument( + "--mesh-file", dest="mesh_file", type=str, + required=True, help="Path+name to unsorted mesh file.") + + parser.add_argument( + "--sort-file", dest="sort_file", type=str, + required=True, help="Path+name to sorted output file.") + + args = parser.parse_args() + + mesh = xarray.open_dataset(args.mesh_file) + + sort_mesh(mesh) + + with open(os.path.join(os.path.dirname( + args.sort_file), "graph.info"), "w") as fptr: + cellsOnCell = mesh["cellsOnCell"].values + + ncells = mesh.dims["nCells"] + nedges = np.count_nonzero(cellsOnCell) // 2 + + fptr.write(f"{ncells} {nedges}\n") + for cell in range(ncells): + data = cellsOnCell[cell, :] + data = data[data > 0] + for item in data: + fptr.write(f"{item} ") + fptr.write("\n") + + mesh.to_netcdf(args.sort_file, format="NETCDF4") + + +def _sort_fwd(data, fwd): + """ + Apply a forward permutation to a mesh array + + Parameters + ---------- + data : array-like + An MPAS mesh array to permute + + fwd : numpy.ndarray + An array of integers defining the permutation + + Returns + ------- + data : numpy.ndarray + The forward permuted MPAS mesh array + """ + vals = data.values + vals = vals[fwd - 1] + return vals -def cell_del2(mesh): +def _sort_rev(data, rev): """ - CELL-DEL2: form cell-to-cell sparse adjacency graph + Apply a reverse permutation to a mesh array + + Parameters + ---------- + data : array-like + An MPAS mesh array to permute + + rev : numpy.ndarray + An array of integers defining the permutation + + Returns + ------- + data : numpy.ndarray + The reverse permuted MPAS mesh array + """ + vals = data.values + mask = vals > 0 + vals[mask] = rev[vals[mask] - 1] + return vals + +def _cell_del2(mesh): """ + Form cell-to-cell sparse adjacency graph + + Parameters + ---------- + mesh : xarray.Dataset + A dataset containing an MPAS mesh to sort + + Returns + ------- + del2 : scipy.sparse.csr_matrix + The cell-to-cell adjacency graph as a sparse matrix + """ xvec = np.array([], dtype=np.int8) ivec = np.array([], dtype=np.int32) jvec = np.array([], dtype=np.int32) @@ -149,39 +236,6 @@ def cell_del2(mesh): return csr_matrix((xvec, (ivec, jvec))) - -if (__name__ == "__main__"): - parser = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.RawTextHelpFormatter) - - parser.add_argument( - "--mesh-file", dest="mesh_file", type=str, - required=True, help="Path+name to unsorted mesh file.") - - parser.add_argument( - "--sort-file", dest="sort_file", type=str, - required=True, help="Path+name to sorted output file.") - - args = parser.parse_args() - - mesh = xarray.open_dataset(args.mesh_file) - - sort_mesh(mesh) - - with open(os.path.join(os.path.dirname( - args.sort_file), "graph.info"), "w") as fptr: - cellsOnCell = mesh["cellsOnCell"].values - - ncells = mesh.dims["nCells"] - nedges = np.count_nonzero(cellsOnCell) // 2 - - fptr.write("{} {}\n".format(ncells, nedges)) - for cell in range(ncells): - data = cellsOnCell[cell, :] - data = data[data > 0] - for item in data: - fptr.write("{} ".format(item)) - fptr.write("\n") - - mesh.to_netcdf(args.sort_file, format="NETCDF4") + +if (__name__ == "__main__"):\ + main() diff --git a/conda_package/recipe/meta.yaml b/conda_package/recipe/meta.yaml index ab3bb74c5..188f16bea 100644 --- a/conda_package/recipe/meta.yaml +++ b/conda_package/recipe/meta.yaml @@ -33,6 +33,7 @@ build: - prepare_seaice_partitions = mpas_tools.seaice.partition:prepare_partitions - create_seaice_partitions = mpas_tools.seaice.partition:create_partitions - simple_seaice_partitions = mpas_tools.seaice.partition:simple_partitions + - sort_mesh = mpas_tools.mesh.creation.sort_mesh:main requirements: build: @@ -101,6 +102,7 @@ test: - planar_hex --nx=20 --ny=40 --dc=1000. --outFileName='periodic_mesh_20x40_1km.nc' - translate_planar_grid -f 'periodic_mesh_10x20_1km.nc' -d 'periodic_mesh_20x40_1km.nc' - MpasMeshConverter.x mesh_tools/mesh_conversion_tools/test/mesh.QU.1920km.151026.nc mesh.nc + - sort_mesh --mesh-file mesh.nc --sort-file sorted_mesh.nc - MpasCellCuller.x mesh.nc culled_mesh.nc -m mesh_tools/mesh_conversion_tools/test/land_mask_final.nc - MpasMaskCreator.x mesh.nc arctic_mask.nc -f mesh_tools/mesh_conversion_tools/test/Arctic_Ocean.geojson - planar_hex --nx=30 --ny=20 --dc=1000. --npx --npy --outFileName='nonperiodic_mesh_30x20_1km.nc'