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

print 'Starting'

from dynamico import meshes
from dynamico import parallel
from dynamico import maps
from dynamico import partition

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)

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];

    partition.partition_graph(comm, vtxdist, xadj_loc, adjncy_loc, part, nparts=4)

    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 local_mesh(get_mycells):
    #    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(partition.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

def members(struct, *names): return [struct.__dict__ [name] for name in names]

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

grid = 'x1.2562'
#grid = 'x1.10242'
#grid = 'x4.163842'
print 'Reading MPAS file %s ...'%grid

meshfile = meshes.MPAS_Format('grids/%s.grid.nc'%grid)
pmesh = meshes.Unstructured_PMesh(comm, meshfile)
pmesh.partition_metis()

def coriolis(lon,lat): return 0.*lat
llm, nqdyn, radius = 1,1,1.
planet = maps.SphereMap(radius, 0.)
lmesh = meshes.Local_Mesh(pmesh, llm, nqdyn, planet)

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

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)
if(mpi_rank==0): print 'total num :', total_num[0], dim_cell.n

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

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
lmesh.plot_patches(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.png'%mpi_rank, dpi=160)

fig, ax = plt.subplots()
buf.read_own(cell_owner)
buf.update()
lmesh.plot_patches(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.png'%mpi_rank, dpi=160)