Changeset 760 for codes/icosagcm/devel/Python/test/py/test_xios_cartesian.py

Timestamp:

10/15/18 01:04:44 (6 years ago)

Author:

dubos

Message:

devel/Python : block-wise partitioning

File:

• codes/icosagcm/devel/Python/test/py/test_xios_cartesian.py (modified) (1 diff)

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

@@ -3 +3 @@
 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)