Index: codes/icosagcm/devel/Python/dynamico/meshes.py
===================================================================
--- codes/icosagcm/devel/Python/dynamico/meshes.py	(revision 759)
+++ codes/icosagcm/devel/Python/dynamico/meshes.py	(revision 760)
@@ -167,4 +167,5 @@
 
         primal_i, primal_j = [ x.astype(np.int32) for x in lon_i, lat_i]
+        edge_i, edge_j = [ x.astype(np.int32) for x in lon_e, lat_e]
         lon_i, lon_v, lon_e = [x*dx-Lx/2 for x in lon_i, lon_v, lon_e]
         lat_i, lat_v, lat_e = [y*dy-Ly/2 for y in lat_i, lat_v, lat_e]
@@ -183,5 +184,5 @@
                          'Aiv','lon_i','lat_i','lon_v','lat_v',
                          'le_de','le','de','lon_e','lat_e','angle_e',
-                         'primal_i','primal_j')
+                         'primal_i','primal_j','edge_i','edge_j')
     def ncwrite(self, name):
         """The method writes Cartesian mesh on the disc.
@@ -249,5 +250,7 @@
                      ("lon_e","f8", self.lon_e),
                      ("lat_e","f8", self.lat_e),
-                      ("angle_e","f8", self.angle_e)] )
+                     ("angle_e","f8", self.angle_e),
+                     ("edge_i","i4",self.edge_i),
+                     ("edge_j","i4",self.edge_j)] )
 
         create_vars( ("edge","TWO"),
@@ -473,4 +476,5 @@
         if gridfile.meshtype == 'curvilinear':
             self.primal_i, self.primal_j = get_pvars(dim_primal, 'primal_i','primal_j')
+            self.edge_i, self.edge_j = get_pvars(dim_edge, 'edge_i','edge_j')
 
         # Let indices start at 0
@@ -484,16 +488,29 @@
         self.vertex2cell = Stencil_glob(dual_deg, dual_vertex, Riv2)
         self.edge2edge   = Stencil_glob(trisk_deg, trisk, wee)
-        print 'Partitioning ...'
-        edge_owner = unst.partition_mesh(trisk_deg, trisk, comm.Get_size())
-        edge_owner = parallel.LocPArray1D(dim_edge, edge_owner)
-        primal_owner = partition_from_stencil(edge_owner, primal_deg, primal_edge)
-        primal_owner = parallel.LocPArray1D(dim_primal, primal_owner)
         self.set_members(locals(), 'dim_primal', 'dim_dual', 'dim_edge',
-                         'edge_owner', 'primal_owner', 'primal_deg', 'primal_vertex', 
+                         'primal_deg', 'primal_vertex', 'primal_edge', 'trisk_deg', 'trisk',
                          'lon_v', 'lat_v', 'Av', 'lon_i', 'lat_i', 'Ai', 
                          'le', 'de', 'lon_e', 'lat_e', 'angle_e')
-#        if meshtype == 'curvilinear' : # read extra information from mesh file
-#            self.primal_i, self.primal_j = get_pvars(dim_primal, 'primal_i','primal_j')
-            
+    def partition(self, edge_owner):
+        self.edge_owner = parallel.LocPArray1D(self.dim_edge, edge_owner)
+        primal_owner = partition_from_stencil(self.edge_owner, self.primal_deg, self.primal_edge)
+        self.primal_owner = parallel.LocPArray1D(self.dim_primal, primal_owner)
+
+    def partition_metis(self):
+        print 'Partitioning with ParMetis...'
+        self.partition(unst.partition_mesh(self.trisk_deg, self.trisk, self.comm.Get_size()))
+
+    def partition_curvilinear(self, ni,nj):
+        def owner(dim,i,j):
+            owner_i, owner_j =  (ni*i.data)/nx, (nj*j.data)/ny
+            return parallel.LocPArray1D(dim, owner_i + ni*owner_j)
+        nx, ny = self.gridfile.nx, self.gridfile.ny
+        print 'Partitioning %d x %d cells in %d x %d blocks ...'%(nx,ny,ni,nj)
+        n = self.comm.Get_size()
+        assert n == ni*nj, 'Mismatch between ni x nj = %d and MPI_SIZE = %d.'%(ni*nj, n)
+        self.edge_owner   = owner(self.dim_edge,   self.edge_i,   self.edge_j)
+        self.primal_owner = owner(self.dim_primal, self.primal_i, self.primal_j)
+        print 'partition_curvilinear %d :'%(self.comm.Get_rank()), self.primal_owner.data
+
     def get(self, getter, *names): return [ getter(self.__dict__[x]) for x in names ]
 
Index: codes/icosagcm/devel/Python/dynamico/xios.py
===================================================================
--- codes/icosagcm/devel/Python/dynamico/xios.py	(revision 759)
+++ codes/icosagcm/devel/Python/dynamico/xios.py	(revision 760)
@@ -3,5 +3,5 @@
 from unstructured import ker
 from dynamico.meshes import radian
-from cxios import finalize
+from mpi4py import MPI
 
 #-----------------------------------------------------------------------------
@@ -9,6 +9,9 @@
 def setup_curvilinear(cat,id,nx,ny,displ,index_own_i,index_own_j,lon,lat):
     mesh = cxios.Handle(cat,id)
-    print 'setup_curvilinear : index_own_i.shape', index_own_i.shape
-    print 'setup_curvilinear : index_own_j.shape', index_own_j.shape
+    def log(name,data) : print 'setup_curvilinear : %s shape min max'%name, data.shape, data.min(), data.max()
+    log('index_own_i',index_own_i)
+    log('index_own_j',index_own_j)
+    log('lon',lon)
+    log('lat',lat)
     mesh.set_attr(type='curvilinear',
                   ni_glo=nx, i_index=index_own_i,
@@ -44,8 +47,23 @@
 #-----------------------------------------------------------------------------
 
-class XIOS_Context:
+class Client:
+    def __enter__(self):
+        ker.dynamico_setup_xios()
+        self.comm=MPI.COMM_WORLD # FIXME : we should use the communicator returned by XIOS
+        return self
+    def __exit__(self, type, value, traceback):
+        print 'xios_finalize()'
+        cxios.finalize()
+
+class Context:
+    def __enter__(self):
+        print 'cxios.context_close_definition()'
+        cxios.context_close_definition()
+        return self
+    def __exit__(self, type, value, traceback):
+        print 'xios_context_finalize()'
+        cxios.context_finalize()
     def init_llm(self, mesh, nqtot):
         self.mesh=mesh
-        ker.dynamico_setup_xios()
         # vertical axis, nq
         llm = mesh.llm
@@ -65,30 +83,27 @@
     def update_calendar(self, i):
         ker.dynamico_xios_update_calendar(i)
-    def finalize(self):
-        cxios.context_finalize()
+    def send_field_primal(self,name, data):
+        own_loc = self.mesh.primal_own_loc
+        if data.ndim==1:
+            data = data[own_loc]
+        if data.ndim==2:
+            data = data[own_loc,:]
+            data = data.transpose() # XIOS expects contiguous horizontal slices
+        data = np.ascontiguousarray(data)
+        cxios.send_field(name, data)
 
-class XIOS_Context_Curvilinear(XIOS_Context):
+class Context_Curvilinear(Context):
     def __init__(self, mesh, nqtot, step):
         self.init_llm(mesh, nqtot)
         # primal mesh
-#        lat_i,lon_i = np.meshgrid(mesh.y,mesh.x, indexing='ij')
-#        ny,nx=lat_i.shape
         own = mesh.primal_own_loc
         lon_i, lat_i = [ x[own]*radian for x in mesh.lon_i, mesh.lat_i ]
         primal_i, primal_j = [ x[own] for x in mesh.primal_i, mesh.primal_j ]
-        print 'setup_curvilinear : ', lon_i.size, lat_i.size, primal_i.size, primal_j.size
         setup_curvilinear('domaingroup','i', mesh.nx, mesh.ny, mesh.displ,
                           primal_i, primal_j, lon_i, lat_i)
         # calendar
         self.init_calendar(step)
-        print 'cxios.context_close_definition()'
-        cxios.context_close_definition()
-    def send_field_primal(self,name, data):
-        if data.ndim==2:
-            data = data.transpose() # XIOS expects contiguous horizontal slices
-        data = np.ascontiguousarray(data)
-        cxios.send_field(name, data)
 
-class XIOS_Context_Unstructured(XIOS_Context):
+class Context_Unstructured(Context):
     def __init__(self, pmesh, mesh,nqtot, step):
         self.init_llm(mesh, nqtot)
@@ -101,13 +116,2 @@
         # calendar
         self.init_calendar(step)
-        print 'cxios.context_close_definition()'
-        cxios.context_close_definition()
-    def send_field_primal(self,name, data):
-        own_loc = self.mesh.primal_own_loc
-        if data.ndim==1:
-            data = data[own_loc]
-        if data.ndim==2:
-            data = data[own_loc,:]
-            data = data.transpose() # XIOS expects contiguous horizontal slices
-        data = np.ascontiguousarray(data)
-        cxios.send_field(name, data)
Index: codes/icosagcm/devel/Python/src/cxios.pyx
===================================================================
--- codes/icosagcm/devel/Python/src/cxios.pyx	(revision 759)
+++ codes/icosagcm/devel/Python/src/cxios.pyx	(revision 760)
@@ -28,5 +28,5 @@
     cdef char* idptr = id
     cdef int idlen = len(id)
-    print 'cxios.write_data', id, data.shape
+#    print 'cxios.write_data', id, data.shape
     ndim=data.ndim
     if ndim==1: send_field1(idptr,idlen,data)
@@ -61,5 +61,5 @@
             create_fun(byref(self.handle), id, c_int(len(id)))
         self.cat, self.prefix_set, self.id = cat, 'set_'+cat+'_', id
-        print 'Handle', cat, self.id, self.handle
+#        print 'Handle', cat, self.id, self.handle
     def set_attr(self, **kwargs):
         prefix, handle = self.prefix_set, self.handle
Index: codes/icosagcm/devel/Python/test/py/DCMIP2008c5_MPI.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/DCMIP2008c5_MPI.py	(revision 759)
+++ codes/icosagcm/devel/Python/test/py/DCMIP2008c5_MPI.py	(revision 760)
@@ -58,4 +58,5 @@
 #    mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
     pmesh = meshes.Unstructured_PMesh(comm,meshfile)
+    pmesh.partition_metis()
     mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, f)
     mesh.pmesh=pmesh
Index: codes/icosagcm/devel/Python/test/py/RSW2_MPAS_W02.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/RSW2_MPAS_W02.py	(revision 759)
+++ codes/icosagcm/devel/Python/test/py/RSW2_MPAS_W02.py	(revision 760)
@@ -29,4 +29,5 @@
 meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
 pmesh = PMesh(comm,meshfile)
+pmesh.partition_metis()
 mesh = Mesh(pmesh, llm, nqdyn, radius, f)
 print '...Done'
Index: codes/icosagcm/devel/Python/test/py/partition.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/partition.py	(revision 759)
+++ codes/icosagcm/devel/Python/test/py/partition.py	(revision 760)
@@ -80,4 +80,5 @@
 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
Index: codes/icosagcm/devel/Python/test/py/test_xios.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/test_xios.py	(revision 759)
+++ codes/icosagcm/devel/Python/test/py/test_xios.py	(revision 760)
@@ -1,52 +1,36 @@
-import dynamico.unstructured as unst
-import dynamico.xios as xios
+from dynamico import unstructured as unst
+from dynamico import xios
 from dynamico.meshes import radian, MPAS_Format, Unstructured_PMesh as PMesh, Local_Mesh as Mesh
+import numpy as np
 
-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 numpy as np
-import cProfile
-import re
+with xios.Client() as client:    
+    comm = client.comm
+    mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
+    print '%d/%d starting'%(mpi_rank,mpi_size)
 
 #----------------------------- read MPAS mesh --------------------------------
 
-grid, llm, nqdyn, nqtot = 2562, 4,1,1 # 2562, 10242, 40962  
-Omega, radius, g, gh0 = 2.*np.pi/86400., 6.4e6, 1., 2.94e4
-N, T, courant = 40, 10800., 1.2 # simulation length = N*T
-print 'Omega, planetary PV', Omega, 2*Omega/gh0
+    grid, llm, nqdyn, nqtot = 2562, 4,1,1 # 2562, 10242, 40962  
+    Omega, radius, g, gh0 = 2.*np.pi/86400., 6.4e6, 1., 2.94e4
+    N, T, courant = 40, 10800., 1.2 # simulation length = N*T
+    print 'Omega, planetary PV', Omega, 2*Omega/gh0
 
-def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter                                                                                                                
-print 'Reading MPAS mesh ...'
-meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
-pmesh = PMesh(comm,meshfile)
-mesh = Mesh(pmesh, llm, nqdyn, radius, f)
-print '...Done'
+    def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter                                                                                                                
+    print 'Reading MPAS mesh ...'
+    meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
+    pmesh = PMesh(comm,meshfile)
+    pmesh.partition_metis()
+    mesh = Mesh(pmesh, llm, nqdyn, radius, f)
+    print '...Done'
 
 #--------------------------------- write some data ----------------------------------------
 
-context=xios.XIOS_Context_Unstructured(pmesh,mesh,nqtot, 3600)
+    with xios.Context_Unstructured(pmesh,mesh,nqtot, 3600) as context:
+        lat_i = radian*mesh.lat_i
+        lat_ik, junk  = np.meshgrid(lat_i, np.arange(llm), indexing='ij')
 
-lat_i = radian*mesh.lat_i
-lat_ik, junk  = np.meshgrid(lat_i, np.arange(llm), indexing='ij')
-
-no_error=True
-try:
-    for i in range(100):
-        context.update_calendar(i)
-        print 'send_field', i, lat_ik.shape
-        context.send_field_primal('ps', lat_i)
-        context.send_field_primal('theta', lat_ik)
-except:
-    print 'An error has occured, trying to close XIOS properly'
-    no_error=False
-
-print 'xios.context_finalize()'
-context.finalize()
-print 'xios.finalize()'
-xios.finalize()
-print 'Done'
-
-if not no_error : raise
+        for i in range(100):
+            context.update_calendar(i)
+            print 'send_field', i, lat_ik.shape
+            context.send_field_primal('ps', lat_i)
+            context.send_field_primal('theta', lat_ik)
Index: codes/icosagcm/devel/Python/test/py/test_xios_cartesian.py
===================================================================
--- codes/icosagcm/devel/Python/test/py/test_xios_cartesian.py	(revision 759)
+++ codes/icosagcm/devel/Python/test/py/test_xios_cartesian.py	(revision 760)
@@ -3,55 +3,45 @@
 from dynamico import unstructured as unst
 from dynamico import meshes
-ker=unst.ker
-
-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 numpy as np
 
-nqdyn,nqtot,llm = 1,1,4
+# client.__exit__() and context.__exit__() are called at the end of the 'with' blocks
+# even if an exception is raised in the body of the block
 
-if False:
-    nx,ny,llm = 64,64
-    Lx,Ly,f = 1.,1.,1.
-    mesh = meshes.Cartesian_Mesh(nx,ny,llm,nqdyn,Lx,Ly,f)
-else:
-    filename, g, f, radius = 'cart_128_128.nc', 1., 1., None
-    filename = 'cart_004_004.nc'
-    unst.setvar('g',g)
-    print 'Reading Cartesian mesh ...'
-    def coriolis(lon,lat):
-        return f+0.*lon
-    meshfile = meshes.DYNAMICO_Format(filename)
-    nx,ny=meshfile.nx,meshfile.ny
+with xios.Client() as client:
+    comm = client.comm
+    mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
+    print '%d/%d starting'%(mpi_rank,mpi_size)
 
-    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
-    mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)
+    nqdyn,nqtot,llm = 1,1,4
+
+    if False:
+        nx,ny,llm = 64,64
+        Lx,Ly,f = 1.,1.,1.
+        mesh = meshes.Cartesian_Mesh(nx,ny,llm,nqdyn,Lx,Ly,f)
+    else:
+        filename, g, f, radius = 'cart_128_128.nc', 1., 1., None
+    #    filename = 'cart_008_008.nc'
+        unst.setvar('g',g)
+        print 'Reading Cartesian mesh ...'
+        def coriolis(lon,lat):
+            return f+0.*lon
+        meshfile = meshes.DYNAMICO_Format(filename)
+        nx,ny=meshfile.nx,meshfile.ny
+
+        pmesh = meshes.Unstructured_PMesh(comm,meshfile)
+    #    pmesh.partition_metis()
+        pmesh.partition_curvilinear(2,2)
+        mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)
 
 #--------------------------------- write some data ----------------------------------------
 
-context=xios.XIOS_Context_Curvilinear(mesh,nqtot, 3600)
-#lon_i,lat_i = mesh.xx[:,:,0].flatten(), mesh.yy[:,:,0].flatten()
-lon_i,lat_i = mesh.lon_i, mesh.lat_i
-lat_ik, junk  = np.meshgrid(lat_i, np.arange(llm), indexing='ij')
+    with xios.Context_Curvilinear(mesh,nqtot, 3600) as context:
+        lon_i,lat_i = mesh.lon_i, mesh.lat_i
+        lat_ik, junk  = np.meshgrid(lat_i, np.arange(llm), indexing='ij')
+        lon_ik, junk  = np.meshgrid(lon_i, np.arange(llm), indexing='ij')
 
-no_error=True
-try:
-    for i in range(3):
-        context.update_calendar(i)
-        print 'send_field', i, lat_ik.shape
-        context.send_field_primal('ps', lat_i)
-        context.send_field_primal('theta', lat_ik)
-except:
-    print 'An error has occured, trying to close XIOS properly'
-    no_error=False
-
-print 'xios.context_finalize()'
-context.finalize()
-print 'xios.finalize()'
-xios.finalize()
-print 'Done'
-
-if not no_error : raise
+        for i in range(48):
+            context.update_calendar(i)
+#            print 'send_field', i, lat_ik.shape
+            context.send_field_primal('ps', lat_i+lon_i)
+            context.send_field_primal('theta', lat_ik+lon_ik)
Index: codes/icosagcm/devel/Python/test/xml/iodef.xml
===================================================================
--- codes/icosagcm/devel/Python/test/xml/iodef.xml	(revision 759)
+++ codes/icosagcm/devel/Python/test/xml/iodef.xml	(revision 760)
@@ -3,5 +3,5 @@
 	<context  id="xios">
 		<variable_definition>
-                <variable id="print_file" type="bool"> false </variable> 
+                <variable id="print_file" type="bool"> true </variable> 
        		<variable_group id="buffer">
          	</variable_group>
