source: codes/icosagcm/devel/Python/test/py/partition.py @ 672

print 'Starting'

from mpi4py import MPI
comm = MPI.COMM_WORLD
mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
print '%d/%d starting'%(mpi_rank,mpi_size)

#from dynamico import partition
from dynamico import parallel, meshes
from dynamico import unstructured as unst

import math as math
import numpy as np
import netCDF4 as cdf

import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection

print 'Done loading modules'

#----------------- partition hand-written 15-cell mesh ------------------#

if mpi_size<15:
    send=np.random.randn(mpi_size)
    recv=np.zeros(mpi_size)
    comm.Alltoall(send,recv)
    #time.sleep(mpi_rank)
#    print mpi_rank, send, recv
    
    adjncy=[1, 5, 0, 2, 6, 1, 3, 7, 2, 4, 8, 3, 9, 
              0, 6, 10, 1, 5, 7, 11, 2, 6, 8, 12, 3, 7, 9, 13, 4, 8, 14,
              5, 11, 6, 10, 12, 7, 11, 13, 8, 12, 14, 9, 13 ]
    xadj=[0, 2, 5, 8, 11, 13, 16, 20, 24, 28, 31, 33, 36, 39, 42, 44]
    
    nb_vert = len(xadj)-1
    vtxdist = [i*nb_vert/mpi_size for i in range(mpi_size+1)]
    xadj, adjncy, vtxdist = [np.asarray(x,np.int32) for x in xadj,adjncy,vtxdist]
    
    idx_start = vtxdist[mpi_rank]
    idx_end = vtxdist[mpi_rank+1]
    nb_vert = idx_end - idx_start
    
    xadj_loc = xadj[idx_start:idx_end+1]-xadj[idx_start]
    adjncy_loc = adjncy[ xadj[idx_start]:xadj[idx_end] ]
    part = 0*xadj_loc[0:-1];

    unst.ker.dynamico_partition_graph(mpi_rank, mpi_size, vtxdist, xadj_loc, adjncy_loc, 4, part)

#    for i in range(len(part)):
#        print 'vertex', i+idx_start, 'proc', part[i]

#-----------------------------------------------------------------------------#
#---------------         partition and plot MPAS mesh       ------------------#
#-----------------------------------------------------------------------------#

# 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]
    print '%d : len(mydegree)=%d'%(mpi_rank, len(mydegree))
    vertex_list = sorted(set(unst.list_stencil(mydegree,mybounds))) 
    print '%d : len(vertex_list))=%d'%(mpi_rank, len(vertex_list))
    get_myvertices = parallel.Get_Indices(dim_vertex, vertex_list)
    mylon, mylat = [get_myvertices(x)*180./math.pi for x in lonVertex, latVertex]
    vertex_dict = parallel.inverse_list(vertex_list)
    meshes.reindex(vertex_dict, mydegree, mybounds)
    return vertex_list, mydegree, mybounds, mylon, mylat

#--------------- read MPAS grid file ---------------#

#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']

# 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) ]

#---------------- 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[0]=com_cells.own_len
comm.Reduce(local_num, total_num, op=MPI.SUM, root=0)
if(mpi_rank==0): print 'total num :', total_num[0], dim_cell.n

#---------------------------- plot -----------------------------#

if True:
    print 'Plotting ...'

    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(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)