Changeset 760 for codes/icosagcm/devel/Python/test/py/test_xios.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.py (modified) (1 diff)

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

@@ -1 @@
-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)