Changeset 682 for codes/icosagcm/devel/Python/test/py/test_xios.py

Timestamp:

02/16/18 18:06:01 (6 years ago)

Author:

dubos

Message:

devel/unstructured : XIOS output with MPI

File:

• codes/icosagcm/devel/Python/test/py/test_xios.py (modified) (8 diffs)

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

@@ -3 +3 @@
 import dynamico.unstructured as unst
 import dynamico.xios as xios
-from dynamico.meshes import MPAS_Format, Unstructured_Mesh, radian
+#from dynamico.meshes import MPAS_Format, Unstructured_Mesh, radian
+from dynamico.meshes import radian, MPAS_Format, Unstructured_Mesh, Unstructured_PMesh as PMesh, Local_Mesh as Mesh
 
 from mpi4py import MPI
@@ -11 +12 @@
 
 import numpy as np
+import cProfile
+import re
 
 def boundaries(degree,points,lon,lat):
@@ -28 +31 @@
 #----------------------------- read MPAS mesh --------------------------------
 
-grid, llm, nqdyn, nqtot = 10242, 1,1,1 # 2562, 10242, 40962  
+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
@@ -36 +39 @@
 print 'Reading MPAS mesh ...'
 meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
-mesh = Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
+#mesh = Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
+pmesh = PMesh(comm,meshfile)
+mesh = Mesh(pmesh, llm, nqdyn, radius, f)
+ni_glo = pmesh.dim_primal.n
+ni_loc = len(mesh.primal_own_glo)
 print '...Done'
 
@@ -43 +50 @@
 #--------------------------------- initialize XIOS --------------------------
 
-def setup_domain(cat,id, degree,vertex,lon,lat,lon_v,lat_v): # cat in 'domain','domaingroup'
+def setup_domain(cat,id, degree,vertex,lon,lat,lon_v,lat_v, ncell_glo=None, index_glo=None, displ=None): # cat in 'domain','domaingroup'
     bnd_lon, bnd_lat = boundaries(degree, vertex, lon_v, lat_v)
-    ncell_tot, nvertex = bnd_lon.shape
+    ncell, nvertex = bnd_lon.shape
+    if ncell_glo is None :
+        ncell_glo = ncell
+        index_glo = np.arange(ncell, dtype=np.int32)
+        displ = 0
     mesh = xios.Handle(cat,id)
     mesh.set_attr(type='unstructured')
-    mesh.set_attr(ni_glo=ncell_tot, ibegin=0, ni=ncell_tot)
-    mesh.set_attr(i_index=np.arange(ncell_tot,dtype=np.int32))
+    mesh.set_attr(ni_glo=ncell_glo, ni=ncell, ibegin=displ, i_index=index_glo)
     mesh.set_attr(lonvalue_1d=lon, latvalue_1d=lat)
     mesh.set_attr(nvertex=nvertex,bounds_lon_1d=bnd_lon, bounds_lat_1d=bnd_lat)
@@ -60 +70 @@
     levp1.set_attr(n_glo=llm+1, value=np.arange(llm+1))
     nq.set_attr(n_glo=nqtot, value=np.arange(nqtot))
-    mesh_i=setup_domain('domaingroup','i', mesh.primal_deg, mesh.primal_vertex-1, lon_i, lat_i, lon_v, lat_v) # primal mesh
+     # primal mesh
+    mesh_i=setup_domain('domaingroup','i', mesh.primal_own_deg, mesh.primal_vertex,
+                            lon_i[mesh.primal_own_loc], lat_i[mesh.primal_own_loc],
+                            lon_v, lat_v,
+                            pmesh.dim_primal.n, mesh.primal_own_glo, mesh.displ)
+#    mesh_i=setup_domain('domaingroup','i', mesh.primal_deg, mesh.primal_vertex-1, lon_i, lat_i, lon_v, lat_v) # primal mesh
 #    mesh_v=setup_domain('domain','v', mesh.dual_deg, mesh.dual_vertex-1, lon_v, lat_v, lon_i, lat_i) # dual mesh
 
@@ -92 +107 @@
     xios.context_close_definition()
 
+lat_ik, junk  = np.meshgrid(lat_i, np.arange(llm), indexing='xy')
+
+def send_field1_primal(name, data):
+#    print 'send_field1_primal :', data.shape
+    data = data[mesh.primal_own_loc]
+    data = np.ascontiguousarray(data)
+    print 'send_field1_primal :', data.shape
+    xios.send_field(name, data)
+def send_field2_primal(name, data):
+#    print 'send_field2_primal :', data.shape
+    data = data[:,mesh.primal_own_loc]
+    data = np.ascontiguousarray(data)
+    print 'send_field2_primal :', data.shape
+    xios.send_field(name, data)
+    
 no_error=True
 try:
@@ -100 +130 @@
             xios.update_calendar(i)
         print 'send_field', i
-        xios.send_field('ps', lat_i)
+        send_field1_primal('ps', lat_i)
+        send_field2_primal('theta', lat_ik)
 except:
     print 'An error has occured, trying to close XIOS properly'