Index: codes/icosagcm/devel/Python/env/X64_CURIE.env
===================================================================
--- codes/icosagcm/devel/Python/env/X64_CURIE.env	(revision 619)
+++ codes/icosagcm/devel/Python/env/X64_CURIE.env	(revision 620)
@@ -3,3 +3,3 @@
 
 # Needed for Intel OpenMP
-KERNELS_LDFLAGS="$NETCDF_LIBDIR $NETCDF_LIB -liomp5 $MKL_LIBS"
+KERNELS_LDFLAGS="$NETCDF_LIBDIR $NETCDF_LIB -liomp5 $MKL_LIBS $PARMETIS_LDFLAGS"
Index: codes/icosagcm/devel/Python/env/linux.env
===================================================================
--- codes/icosagcm/devel/Python/env/linux.env	(revision 619)
+++ codes/icosagcm/devel/Python/env/linux.env	(revision 620)
@@ -1,7 +1,11 @@
+export CODEGEN_ENV_SET=yes
+
 #----------- Useful directories ---------
+
 BASEDIR=$(dirname ${BASH_SOURCE[0]})
 export BASEDIR=$(cd $BASEDIR/../.. ; pwd)
 
 #--------- To create shared libraries -------
+
 export F90=mpif90
 export LINK=$F90
@@ -21,6 +25,4 @@
 export FFLAGS="$FFLAGS_BASE $FFLAGS_PROD"
 
-#----------- To compile with cython ---------
-
 #---------------- To execute ----------------
 
@@ -31,11 +33,4 @@
 export PYTHON=$(which python)
 
-#--------------- Link --------------
-#export LDFLAGS="$PARMETIS_LDFLAGS"
-#export LIBS="-lparmetis"
-
-
-#-------------- Execute --------------
-
 # disable Intel OpenMP in MKL (loaded by numpy) to avoid conflict with our openMP
 export MKL_THREADING_LAYER=SEQUENTIAL
@@ -43,4 +38,2 @@
 # select a Matplotlib backend that works without an X11 connection
 export MPLBACKEND="agg"
-
-export CODEGEN_ENV_SET=yes
Index: codes/icosagcm/devel/Python/src/partition.c
===================================================================
--- codes/icosagcm/devel/Python/src/partition.c	(revision 620)
+++ codes/icosagcm/devel/Python/src/partition.c	(revision 620)
@@ -0,0 +1,30 @@
+#include <mpi.h>
+#include <stdio.h>
+#include <parmetis.h>
+
+void dynamico_partition_graph(int mpi_rank, int mpi_size, 
+                     idx_t vtxdist[], idx_t xadj_loc[], idx_t adjncy_loc[],
+                     idx_t nparts, idx_t part[])
+{
+  MPI_Comm comm = MPI_COMM_WORLD; /* avoid passing a communicator as argument due to Fortran/C issues */
+  idx_t edgecut;
+  idx_t options[]={0}, wgtflag=0, numflag=0, ncon=1;
+  real_t tpwgts[2*nparts], ubvec[]={1.05};
+
+  /* int ParMETIS V3 PartKway (
+     idx t *vtxdist, idx t *xadj, idx t *adjncy, idx t *vwgt, idx t *adjwgt, 
+     idx t *wgtflag, idx t *numflag, idx t *ncon, idx t *nparts, 
+     real t *tpwgts, real t *ubvec, idx t *options, idx t *edgecut, idx t *part, 
+     MPI Comm *comm )
+  */
+
+/*  printf("sizeof : %d %d \n", sizeof(int), sizeof(idx_t));
+  printf("partition_graph : %d %d %d\n", mpi_rank, mpi_size, nparts); 
+  printf("vtxdist : \n"); 
+  for(int i=0; i<mpi_size+1; i++) printf("%d %d %d\n", mpi_rank, i, vtxdist[i]); */
+      
+  for(int i=0; i<2*nparts; i++) tpwgts[i] = 1./nparts;
+  ParMETIS_V3_PartKway(vtxdist, xadj_loc, adjncy_loc, NULL, NULL, 
+		       &wgtflag, &numflag, &ncon, &nparts, 
+		       tpwgts, ubvec, options, &edgecut, part, &comm); 
+}
Index: codes/icosagcm/devel/Python/src/unstructured.pyx
===================================================================
--- codes/icosagcm/devel/Python/src/unstructured.pyx	(revision 619)
+++ codes/icosagcm/devel/Python/src/unstructured.pyx	(revision 620)
@@ -44,5 +44,5 @@
                      ['dynamico_init_hybrid', c_void_p,3],
                      ['dynamico_caldyn_unstructured', c_double, c_void_p,20],
-#                     ['dynamico_partition_graph', c_int,2, c_void_p,3, c_int, c_void_p],
+                     ['dynamico_partition_graph', c_int,2, c_void_p,3, c_int, c_void_p],
                      ])
 
Index: codes/icosagcm/devel/Python/test/py/partition.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/partition.py	(revision 620)
+++ codes/icosagcm/devel/Python/test/py/partition.py	(revision 620)
@@ -0,0 +1,200 @@
+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)
+
+import sys
+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
+
+#from dynamico import partition
+from dynamico import parallel
+from dynamico import unstructured as unst
+from dynamico.unstructured import list_stencil, reindex
+
+print 'Done loading modules'
+
+sys.stdout.flush()
+
+#----------------- 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(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)
+    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
+sys.stdout.flush()
+
+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 = [
+    unst.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 ...'
+sys.stdout.flush()
+
+edge_owner = unst.partition_mesh(nEdgesOnEdge, edgesOnEdge, mpi_size)
+edge_owner = parallel.LocPArray1D(dim_edge, edge_owner)
+cell_owner = unst.partition_from_stencil(edge_owner, nEdgesOnCell, edgesOnCell)
+cell_owner = parallel.LocPArray1D(dim_cell, cell_owner)
+
+#--------------------- construct halos  -----------------------#
+
+print 'Constructing halos ...'
+sys.stdout.flush()
+
+def chain(start, links):
+    for link in links:
+        start = link(start).neigh_set
+        yield start
+
+edges_E0 = unst.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 = unst.progressive_list(edges_E0, edges_E1, edges_E2)
+cells_C0, cells_C1 = unst.progressive_list(cells_C0, cells_C1)
+
+print 'E2,E1,E0 ; C1,C0 : ', map(len, (edges_E2, edges_E1, edges_E0, cells_C1, cells_C0))
+sys.stdout.flush()
+
+#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
+sys.stdout.flush()
+
+#---------------------------- plot -----------------------------#
+
+if True:
+    print 'Plotting ...'
+    sys.stdout.flush()
+
+    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)
