Changeset 682

Timestamp:

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

Author:

dubos

Message:

devel/unstructured : XIOS output with MPI

Location:

codes/icosagcm/devel/Python

Files:

• dynamico/meshes.py (modified) (7 diffs)
• dynamico/parallel.py (modified) (1 diff)
• test/fig_DCMIP2008c5 (modified) (1 prop)
• test/fig_morton (modified) (1 prop)
• test/py/test_xios.py (modified) (8 diffs)
• test/xml/field_def_simple.xml (modified) (1 diff)
• test/xml/filedef_simple.xml (modified) (1 diff)

codes/icosagcm/devel/Python/dynamico/meshes.py

@@ -367 +367 @@
         self.edge2cell   = Stencil_glob(edge_deg, left_right)
         self.cell2edge   = Stencil_glob(primal_deg, primal_edge)
+        self.cell2vertex = Stencil_glob(primal_deg, primal_vertex)
         self.edge2vertex = Stencil_glob(edge_deg, down_up)
         self.vertex2edge = Stencil_glob(dual_deg, dual_edge)
@@ -406 +407 @@
 class Local_Mesh(Abstract_Mesh):
     def __init__(self, pmesh, llm, nqdyn, radius, f):
-        dim_primal, primal_owner, dim_edge, edge_owner, dim_dual = pmesh.members(
-            'dim_primal', 'primal_owner', 'dim_edge', 'edge_owner', 'dim_dual')
+        comm, dim_primal, primal_owner, dim_edge, edge_owner, dim_dual = pmesh.members(
+            'comm', 'dim_primal', 'primal_owner', 'dim_edge', 'edge_owner', 'dim_dual')
         print 'Constructing halos ...'
         edges_E0 = find_my_cells(edge_owner)
@@ -417 +418 @@
         get_own_cells, get_all_cells = dim_primal.getter(cells_C0), dim_primal.getter(cells_C1)
         get_all_duals, get_all_edges = dim_dual.getter(vertices_V1), dim_edge.getter(edges_E2)
-        self.com_primal = parallel.Halo_Xchange(
-            42, dim_primal, cells_C1, get_all_cells(primal_owner))
-        self.com_edges = parallel.Halo_Xchange(
-            73, dim_edge, edges_E2, get_all_edges(edge_owner))
         
         self.lon_i, self.lat_i, self.Ai = pmesh.get(get_all_cells, 'lon_i', 'lat_i', 'Ai')
@@ -428 +425 @@
             get_all_edges, 'le', 'de', 'lon_e', 'lat_e', 'angle_e')
         primal_num, dual_num, edge_num = map(len, (cells_C1, vertices_V1, edges_E2))
-        
+
+        # construct local stencils from global stencils
         dict_E2, dict_C1, dict_V1 = map(inverse_list, (edges_E2, cells_C1, vertices_V1))
-        down_up = pmesh.edge2vertex(edges_E2, dict_V1)
-        left_right = pmesh.edge2cell(edges_E2, dict_C1)
-        primal_edge = pmesh.cell2edge(cells_C1, dict_E2)
-        dual_edge = pmesh.vertex2edge(vertices_V1, dict_E2)
-        trisk = pmesh.edge2edge(edges_E2, dict_E2)
-        Riv2  = pmesh.vertex2cell(vertices_V1, dict_C1)
+        down_up       = pmesh.edge2vertex( edges_E2,    dict_V1)
+        left_right    = pmesh.edge2cell(   edges_E2,    dict_C1)
+        primal_edge   = pmesh.cell2edge(   cells_C1,    dict_E2)
+        dual_edge     = pmesh.vertex2edge( vertices_V1, dict_E2)
+        trisk         = pmesh.edge2edge(   edges_E2,    dict_E2)
+        Riv2          = pmesh.vertex2cell( vertices_V1, dict_C1)
         print 'Edge-vertex degree min/max %d %d'%(down_up.degree.min(), down_up.degree.max() )
         print 'Edge-cell degree min/max %d %d'%(left_right.degree.min(), left_right.degree.max() )
@@ -442 +440 @@
         print 'TRiSK degree min/max : %d %d'%(trisk.degree.min(), trisk.degree.max() )
         print 'Riv2 degree min/max : %d %d'%(Riv2.degree.min(), Riv2.degree.max() )
-        primal_deg, primal_edge = primal_edge.degree, primal_edge.neigh_loc
-        dual_deg, dual_edge = dual_edge.degree, dual_edge.neigh_loc
-        trisk_deg, trisk, wee = trisk.degree, trisk.neigh_loc, trisk.weight
-        dual_deg, dual_vertex, Riv2 = Riv2.degree, Riv2.neigh_loc, Riv2.weight
+        primal_deg, primal_edge     = primal_edge.degree,   primal_edge.neigh_loc
+        dual_deg,   dual_edge       = dual_edge.degree,     dual_edge.neigh_loc
+        trisk_deg,  trisk, wee      = trisk.degree,         trisk.neigh_loc, trisk.weight
+        dual_deg, dual_vertex, Riv2 = Riv2.degree,          Riv2.neigh_loc,  Riv2.weight
         edge_deg, left, right = left_right.degree, left_right.neigh_loc[:,0], left_right.neigh_loc[:,1]
-        edge_deg, down, up = down_up.degree, down_up.neigh_loc[:,0], down_up.neigh_loc[:,1]
+        edge_deg, down, up    = down_up.degree,    down_up.neigh_loc[:,0],    down_up.neigh_loc[:,1]
         for edge in range(edge_num): 
             if edge_deg[edge]<2: up[edge]=down[edge]
         max_primal_deg, max_dual_deg, max_trisk_deg = [x.shape[1] for x in primal_edge, dual_edge, trisk]
+
+        # construct own primal mesh for XIOS output
+        primal_own_glo = find_my_cells(primal_owner)
+        primal_vertex = pmesh.cell2vertex(primal_own_glo, dict_V1)
+        primal_own_deg, primal_vertex = primal_vertex.degree, primal_vertex.neigh_loc
+        primal_own_glo = np.asarray(primal_own_glo, dtype=np.int32)
+        primal_own_loc = [dict_C1[i] for i in primal_own_glo]
+        primal_own_num = len(primal_own_glo)
+        primal_own_all = comm.allgather(primal_own_num)
+        displ = sum(primal_own_all[0:comm.Get_rank()]) # required by XIOS
+        print 'local primal mesh :', primal_vertex.min(), primal_vertex.max(), primal_own_all, displ
+        
         # compute_ne(...) and normalize(...) expect indices starting at 1
         trisk, dual_vertex, primal_edge, dual_edge =  trisk+1, dual_vertex+1, primal_edge+1, dual_edge+1
@@ -456 +466 @@
         primal_ne = compute_ne(primal_num,primal_deg,primal_edge,left,right)
         dual_ne = compute_ne(dual_num,dual_deg,dual_edge,up,down)
+
+        # store what we have computed
         self.set_members(locals(), 'nqdyn', 'llm', 'primal_num', 'dual_num', 'edge_num',
-                         'trisk_deg', 'trisk', 'wee', 'dual_deg', 'dual_vertex', 'Riv2')
-
+                         'primal_own_glo', 'primal_own_loc', 'primal_own_deg',
+                         'primal_deg', 'primal_vertex', 'displ',
+                        'trisk_deg', 'trisk', 'wee', 'dual_deg', 'dual_vertex', 'Riv2')
+
+        # normalize and propagate info to Fortran side
         pmesh.gridfile.normalize(self,radius)
         unst.init_mesh(llm, nqdyn, self.edge_num, self.primal_num, self.dual_num,
@@ -466 +481 @@
                   left,right,down,up,trisk_deg,trisk,
                   self.Ai,self.Av,self.fv,self.le/self.de,Riv2,wee)
-        # llm must be set before we call set_dynamico_transfer
+
+        # setup halo transfers - NB : llm must be set before we call set_dynamico_transfer
+        self.com_primal = parallel.Halo_Xchange(
+            42, dim_primal, cells_C1, get_all_cells(primal_owner))
+        self.com_edges = parallel.Halo_Xchange(
+            73, dim_edge, edges_E2, get_all_edges(edge_owner))
         self.com_primal.set_dynamico_transfer('primal')
         self.com_edges.set_dynamico_transfer('edge')

codes/icosagcm/devel/Python/dynamico/parallel.py

@@ -148 +148 @@
     def set_dynamico_transfer(self,name):
         def list_to_args(lst):
-            ranks, lst = zip(*lst) # list of tuples => tuple of lists
+            ranks, lst = ([], []) if lst == [] else zip(*lst) # list of tuples => tuple of lists
             lens = map(len, lst)
             return len(ranks), ranks, lens, sum(lens), sum(lst,[])

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'

codes/icosagcm/devel/Python/test/xml/field_def_simple.xml

@@ -16 +16 @@
       <field id="phis" />
       <field id="phi" />
+
+      <field_group axis_ref="lev"> 
+        <field id="theta" />
+      </field_group>
+     
     </field_group> 
     

codes/icosagcm/devel/Python/test/xml/filedef_simple.xml

@@ -15 +15 @@
     <field_group id="dcmip2016_output_field" ts_enabled="false">
       <field field_ref="ps" name="PS"       standard_name="surface_pressure" long_name="Surface pressure"         unit="Pa" />
+      <field field_ref="theta" />
       <field field_ref="phis" operation="once" freq_offset="0ts" name="PHIS"   standard_name="surface_geopotential" long_name="Surface geopotential" unit="m2/m2"/>      
     </field_group>