Changeset 680 for codes/icosagcm/devel/Python/test/py/partition.py

Timestamp:

02/09/18 16:24:33 (6 years ago)

Author:

dubos

Message:

devel/unstructured : moved mesh partitioning code into meshes.py

File:

• codes/icosagcm/devel/Python/test/py/partition.py (modified) (3 diffs)

codes/icosagcm/devel/Python/test/py/partition.py

@@ -57 +57 @@
 # Helper functions to plot unstructured graph
 
-def patches(degree, bounds, lon, lat):
-    for i in range(degree.size):
-        nb_edge=degree[i]
-        bounds_cell = bounds[i,0:nb_edge]
-        lat_cell    = lat[bounds_cell]
-        lon_cell    = lon[bounds_cell]
-        orig=lon_cell[0]
-        lon_cell    = lon_cell-orig+180.
-        lon_cell    = np.mod(lon_cell,360.)
-        lon_cell    = lon_cell+orig-180.
-#        if np.abs(lon_cell-orig).max()>10. :
-#            print '%d patches :'%mpi_rank, lon_cell
-        lonlat_cell = np.zeros((nb_edge,2))
-        lonlat_cell[:,0],lonlat_cell[:,1] = lon_cell,lat_cell
-        polygon = Polygon(lonlat_cell, True)
-        yield polygon
-
-def plot_mesh(ax, clim, degree, bounds, lon, lat, data):
-    nb_vertex = lon.size # global
-    p = list(patches(degree, bounds, lon, lat))
-    print '%d : plot_mesh %d %d %d'%( mpi_rank, degree.size, len(p), len(data) ) 
-    p = PatchCollection(p, linewidth=0.01)
-    p.set_array(data) # set values at each polygon (cell)
-    p.set_clim(clim)
-    ax.add_collection(p)
-
 def local_mesh(get_mycells):
-    mydegree, mybounds = [get_mycells(x) for x in nEdgesOnCell, verticesOnCell]
+    #    mydegree, mybounds = [get_mycells(x) for x in nEdgesOnCell, verticesOnCell]
+    mydegree, mybounds = [get_mycells(x) for x in primal_deg, primal_vertex]
     print '%d : len(mydegree)=%d'%(mpi_rank, len(mydegree))
     vertex_list = sorted(set(unst.list_stencil(mydegree,mybounds))) 
@@ -94 +69 @@
     return vertex_list, mydegree, mybounds, mylon, mylat
 
+def members(struct, *names): return [struct.__dict__ [name] for name in names]
+
 #--------------- read MPAS grid file ---------------#
 
-#grid = 'x1.2562'
-grid = 'x1.10242'
+grid = 'x1.2562'
+#grid = 'x1.10242'
 #grid = 'x4.163842'
 print 'Reading MPAS file %s ...'%grid
 
-nc = cdf.Dataset('grids/%s.grid.nc'%grid, "r")
-dim_cell, dim_edge, dim_vertex = [
-    parallel.PDim(nc.dimensions[name], comm) 
-    for name in 'nCells','nEdges','nVertices']
-edge_degree   = parallel.CstPArray1D(dim_edge, np.int32, 2)
-vertex_degree = parallel.CstPArray1D(dim_vertex, np.int32, 3)
-nEdgesOnCell, verticesOnCell, edgesOnCell, cellsOnCell, latCell = [
-    parallel.PArray(dim_cell, nc.variables[var])
-    for var in 'nEdgesOnCell', 'verticesOnCell', 'edgesOnCell', 'cellsOnCell', 'latCell' ]
-cellsOnVertex, edgesOnVertex, kiteAreasOnVertex, lonVertex, latVertex = [
-    parallel.PArray(dim_vertex, nc.variables[var])
-    for var in 'cellsOnVertex', 'edgesOnVertex', 'kiteAreasOnVertex', 'lonVertex', 'latVertex']
-nEdgesOnEdge, cellsOnEdge, edgesOnEdge, verticesOnEdge, weightsOnEdge = [
-    parallel.PArray(dim_edge, nc.variables[var])
-    for var in 'nEdgesOnEdge', 'cellsOnEdge', 'edgesOnEdge', 'verticesOnEdge', 'weightsOnEdge']
+meshfile = meshes.MPAS_Format('grids/%s.grid.nc'%grid)
+pmesh = meshes.Unstructured_PMesh(comm, meshfile)
+lmesh = meshes.Local_Mesh(pmesh)
 
-# Indices start at 0 on the C/Python side and at 1 on the Fortran/MPAS side
-# hence an offset of 1 is added/substracted where needed.
-for x in (verticesOnCell, edgesOnCell, cellsOnCell, cellsOnVertex, edgesOnVertex,
-          cellsOnEdge, edgesOnEdge, verticesOnEdge) : x.data = x.data-1
-edge2cell, cell2edge, edge2vertex, vertex2edge, cell2cell, edge2edge = [
-    meshes.Stencil_glob(a,b) for a,b in 
-    (edge_degree, cellsOnEdge), (nEdgesOnCell, edgesOnCell),
-    (edge_degree, verticesOnEdge), (vertex_degree, edgesOnVertex),
-    (nEdgesOnCell, cellsOnCell), (nEdgesOnEdge, edgesOnEdge) ]
+(primal_deg, primal_vertex, dim_vertex, dim_cell, cell_owner, 
+ lonVertex, latVertex, lonCell, latCell) = members(
+    pmesh, 'primal_deg', 'primal_vertex', 'dim_dual', 'dim_primal', 'primal_owner', 
+    'lon_v', 'lat_v', 'lon_i', 'lat_i')
 
-#---------------- partition edges and cells ------------------#
-
-print 'Partitioning ...'
-
-edge_owner = unst.partition_mesh(nEdgesOnEdge, edgesOnEdge, mpi_size)
-edge_owner = parallel.LocPArray1D(dim_edge, edge_owner)
-cell_owner = meshes.partition_from_stencil(edge_owner, nEdgesOnCell, edgesOnCell)
-cell_owner = parallel.LocPArray1D(dim_cell, cell_owner)
-
-#--------------------- construct halos  -----------------------#
-
-print 'Constructing halos ...'
-
-def chain(start, links):
-    for link in links:
-        start = link(start).neigh_set
-        yield start
-
-edges_E0 = meshes.find_my_cells(edge_owner)
-cells_C0, edges_E1, vertices_V1, edges_E2, cells_C1 = chain(
-    edges_E0, ( edge2cell, cell2edge, edge2vertex, vertex2edge, edge2cell) )
-
-edges_E0, edges_E1, edges_E2 = meshes.progressive_list(edges_E0, edges_E1, edges_E2)
-cells_C0, cells_C1 = meshes.progressive_list(cells_C0, cells_C1)
-
-print 'E2,E1,E0 ; C1,C0 : ', map(len, (edges_E2, edges_E1, edges_E0, cells_C1, cells_C0))
-
-#com_edges = parallel.Halo_Xchange(24, dim_edge, edges_E2, dim_edge.get(edges_E2, edge_owner))
-
-mycells, halo_cells = cells_C0, cells_C1
-get_mycells, get_halo_cells = dim_cell.getter(mycells), dim_cell.getter(halo_cells)
-com_cells = parallel.Halo_Xchange(42, dim_cell, halo_cells, get_halo_cells(cell_owner))
-
-local_num, total_num = np.zeros(1), np.zeros(1)
+local_num, total_num, com_cells = np.zeros(1), np.zeros(1), lmesh.com_primal
 local_num[0]=com_cells.own_len
 comm.Reduce(local_num, total_num, op=MPI.SUM, root=0)
@@ -167 +94 @@
 #---------------------------- plot -----------------------------#
 
-if True:
-    print 'Plotting ...'
+print 'Plotting ...'
 
-    halo_vertex_list, mydegree, mybounds, mylon, mylat = local_mesh(com_cells.get_all)
-    buf = parallel.LocalArray1(com_cells)
+halo_vertex_list, mydegree, mybounds, mylon, mylat = local_mesh(com_cells.get_all)
+buf = parallel.LocalArray1(com_cells)
 
-    fig, ax = plt.subplots()
-    buf.read_own(latCell) # reads only own values
-    buf.data = np.cos(10.*buf.data)
-    buf.update() # updates halo
-    plot_mesh(ax,[-math.pi/2,math.pi/2], mydegree, mybounds, mylon, mylat, buf.data)
-    plt.xlim(-190.,190.)
-    plt.ylim(-90.,90.)
-    plt.savefig('fig_partition/A%03d.pdf'%mpi_rank, dpi=1600)
+fig, ax = plt.subplots()
+buf.read_own(latCell) # reads only own values
+buf.data = np.cos(10.*buf.data)
+buf.update() # updates halo
+lmesh.plot(ax,[-math.pi/2,math.pi/2], mydegree, mybounds, mylon, mylat, buf.data)
+plt.xlim(-190.,190.)
+plt.ylim(-90.,90.)
+plt.savefig('fig_partition/A%03d.pdf'%mpi_rank, dpi=1600)
 
-    fig, ax = plt.subplots()
-    buf.read_own(cell_owner)
-    buf.update()
-    plot_mesh(ax,[0,mpi_rank+1], mydegree, mybounds, mylon, mylat, buf.data)
-    plt.xlim(-190.,190.)
-    plt.ylim(-90.,90.)
-    plt.savefig('fig_partition/B%03d.pdf'%mpi_rank, dpi=1600)
+fig, ax = plt.subplots()
+buf.read_own(cell_owner)
+buf.update()
+lmesh.plot(ax,[0,mpi_rank+1], mydegree, mybounds, mylon, mylat, buf.data)
+plt.xlim(-190.,190.)
+plt.ylim(-90.,90.)
+plt.savefig('fig_partition/B%03d.pdf'%mpi_rank, dpi=1600)