1+ import argparse
2+ import os
13
24import numpy as np
3- import os
45import xarray
5- import argparse
66from scipy .sparse import csr_matrix
77from scipy .sparse .csgraph import reverse_cuthill_mckee
88
9+ from mpas_tools .io import write_netcdf
10+
911
1012def sort_mesh (mesh ):
1113 """
@@ -24,9 +26,9 @@ def sort_mesh(mesh):
2426 """
2527 # Authors: Darren Engwirda
2628
27- ncells = mesh .sizes [" nCells" ]
28- nedges = mesh .sizes [" nEdges" ]
29- nduals = mesh .sizes [" nVertices" ]
29+ ncells = mesh .sizes [' nCells' ]
30+ nedges = mesh .sizes [' nEdges' ]
31+ nduals = mesh .sizes [' nVertices' ]
3032
3133 cell_fwd = np .arange (0 , ncells ) + 1
3234 cell_rev = np .arange (0 , ncells ) + 1
@@ -42,23 +44,20 @@ def sort_mesh(mesh):
4244 cell_rev = np .zeros (ncells , dtype = np .int32 )
4345 cell_rev [cell_fwd - 1 ] = np .arange (ncells ) + 1
4446
45- mesh ["cellsOnCell" ][:] = \
46- _sort_rev (mesh ["cellsOnCell" ], cell_rev )
47- mesh ["cellsOnEdge" ][:] = \
48- _sort_rev (mesh ["cellsOnEdge" ], cell_rev )
49- mesh ["cellsOnVertex" ][:] = \
50- _sort_rev (mesh ["cellsOnVertex" ], cell_rev )
47+ mesh ['cellsOnCell' ][:] = _sort_rev (mesh ['cellsOnCell' ], cell_rev )
48+ mesh ['cellsOnEdge' ][:] = _sort_rev (mesh ['cellsOnEdge' ], cell_rev )
49+ mesh ['cellsOnVertex' ][:] = _sort_rev (mesh ['cellsOnVertex' ], cell_rev )
5150
5251 for var in mesh .keys ():
5352 dims = mesh .variables [var ].dims
54- if ( " nCells" in dims ) :
53+ if ' nCells' in dims :
5554 mesh [var ][:] = _sort_fwd (mesh [var ], cell_fwd )
5655
57- mesh [" indexToCellID" ][:] = np .arange (ncells ) + 1
56+ mesh [' indexToCellID' ][:] = np .arange (ncells ) + 1
5857
5958 # sort duals via pseudo-linear cell-wise ordering
6059
61- dual_fwd = np .ravel (mesh [" verticesOnCell" ].values )
60+ dual_fwd = np .ravel (mesh [' verticesOnCell' ].values )
6261 dual_fwd = dual_fwd [dual_fwd > 0 ]
6362
6463 __ , imap = np .unique (dual_fwd , return_index = True )
@@ -68,22 +67,20 @@ def sort_mesh(mesh):
6867 dual_rev = np .zeros (nduals , dtype = np .int32 )
6968 dual_rev [dual_fwd - 1 ] = np .arange (nduals ) + 1
7069
71- mesh ["verticesOnCell" ][:] = \
72- _sort_rev (mesh ["verticesOnCell" ], dual_rev )
70+ mesh ['verticesOnCell' ][:] = _sort_rev (mesh ['verticesOnCell' ], dual_rev )
7371
74- mesh ["verticesOnEdge" ][:] = \
75- _sort_rev (mesh ["verticesOnEdge" ], dual_rev )
72+ mesh ['verticesOnEdge' ][:] = _sort_rev (mesh ['verticesOnEdge' ], dual_rev )
7673
7774 for var in mesh .keys ():
7875 dims = mesh .variables [var ].dims
79- if ( " nVertices" in dims ) :
76+ if ' nVertices' in dims :
8077 mesh [var ][:] = _sort_fwd (mesh [var ], dual_fwd )
8178
82- mesh [" indexToVertexID" ][:] = np .arange (nduals ) + 1
79+ mesh [' indexToVertexID' ][:] = np .arange (nduals ) + 1
8380
8481 # sort edges via pseudo-linear cell-wise ordering
8582
86- edge_fwd = np .ravel (mesh [" edgesOnCell" ].values )
83+ edge_fwd = np .ravel (mesh [' edgesOnCell' ].values )
8784 edge_fwd = edge_fwd [edge_fwd > 0 ]
8885
8986 __ , imap = np .unique (edge_fwd , return_index = True )
@@ -93,60 +90,75 @@ def sort_mesh(mesh):
9390 edge_rev = np .zeros (nedges , dtype = np .int32 )
9491 edge_rev [edge_fwd - 1 ] = np .arange (nedges ) + 1
9592
96- mesh ["edgesOnCell" ][:] = \
97- _sort_rev (mesh ["edgesOnCell" ], edge_rev )
93+ mesh ['edgesOnCell' ][:] = _sort_rev (mesh ['edgesOnCell' ], edge_rev )
9894
99- mesh ["edgesOnEdge" ][:] = \
100- _sort_rev (mesh ["edgesOnEdge" ], edge_rev )
95+ mesh ['edgesOnEdge' ][:] = _sort_rev (mesh ['edgesOnEdge' ], edge_rev )
10196
102- mesh ["edgesOnVertex" ][:] = \
103- _sort_rev (mesh ["edgesOnVertex" ], edge_rev )
97+ mesh ['edgesOnVertex' ][:] = _sort_rev (mesh ['edgesOnVertex' ], edge_rev )
10498
10599 for var in mesh .keys ():
106100 dims = mesh .variables [var ].dims
107- if ( " nEdges" in dims ) :
101+ if ' nEdges' in dims :
108102 mesh [var ][:] = _sort_fwd (mesh [var ], edge_fwd )
109103
110- mesh [" indexToEdgeID" ][:] = np .arange (nedges ) + 1
104+ mesh [' indexToEdgeID' ][:] = np .arange (nedges ) + 1
111105
112106 return mesh
113107
114108
115109def main ():
116110 parser = argparse .ArgumentParser (
117- description = __doc__ ,
118- formatter_class = argparse .RawTextHelpFormatter )
111+ description = __doc__ , formatter_class = argparse .RawTextHelpFormatter
112+ )
113+
114+ parser .add_argument (
115+ '--mesh-file' ,
116+ dest = 'mesh_file' ,
117+ type = str ,
118+ required = True ,
119+ help = 'Path+name to unsorted mesh file.' ,
120+ )
119121
120122 parser .add_argument (
121- "--mesh-file" , dest = "mesh_file" , type = str ,
122- required = True , help = "Path+name to unsorted mesh file." )
123+ '--sort-file' ,
124+ dest = 'sort_file' ,
125+ type = str ,
126+ required = True ,
127+ help = 'Path+name to sorted output file.' ,
128+ )
123129
124130 parser .add_argument (
125- "--sort-file" , dest = "sort_file" , type = str ,
126- required = True , help = "Path+name to sorted output file." )
131+ '--format' ,
132+ dest = 'format' ,
133+ type = str ,
134+ required = False ,
135+ default = 'NETCDF4' ,
136+ help = 'Output format for the sorted mesh file.' ,
137+ )
127138
128139 args = parser .parse_args ()
129140
130141 mesh = xarray .open_dataset (args .mesh_file )
131142
132143 sort_mesh (mesh )
133144
134- with open (os .path .join (os .path .dirname (
135- args .sort_file ), "graph.info" ), "w" ) as fptr :
136- cellsOnCell = mesh ["cellsOnCell" ].values
145+ with open (
146+ os .path .join (os .path .dirname (args .sort_file ), 'graph.info' ), 'w'
147+ ) as fptr :
148+ cellsOnCell = mesh ['cellsOnCell' ].values
137149
138- ncells = mesh .sizes [" nCells" ]
150+ ncells = mesh .sizes [' nCells' ]
139151 nedges = np .count_nonzero (cellsOnCell ) // 2
140152
141- fptr .write (f" { ncells } { nedges } \n " )
153+ fptr .write (f' { ncells } { nedges } \n ' )
142154 for cell in range (ncells ):
143155 data = cellsOnCell [cell , :]
144156 data = data [data > 0 ]
145157 for item in data :
146- fptr .write (f" { item } " )
147- fptr .write (" \n " )
158+ fptr .write (f' { item } ' )
159+ fptr .write (' \n ' )
148160
149- mesh . to_netcdf ( args .sort_file , format = "NETCDF4" )
161+ write_netcdf ( mesh , args .sort_file , format = args . format )
150162
151163
152164def _sort_fwd (data , fwd ):
@@ -212,11 +224,10 @@ def _cell_del2(mesh):
212224 ivec = np .array ([], dtype = np .int32 )
213225 jvec = np .array ([], dtype = np .int32 )
214226
215- topolOnCell = mesh [" nEdgesOnCell" ].values
216- cellsOnCell = mesh [" cellsOnCell" ].values
227+ topolOnCell = mesh [' nEdgesOnCell' ].values
228+ cellsOnCell = mesh [' cellsOnCell' ].values
217229
218230 for edge in range (np .max (topolOnCell )):
219-
220231 # cell-to-cell pairs, if edges exist
221232 mask = topolOnCell > edge
222233 idx_self = np .argwhere (mask ).ravel ()
@@ -237,5 +248,5 @@ def _cell_del2(mesh):
237248 return csr_matrix ((xvec , (ivec , jvec )))
238249
239250
240- if ( __name__ == " __main__" ):\
251+ if __name__ == ' __main__' :
241252 main ()
0 commit comments