Changeset 680

Timestamp:

02/09/18 16:24:33 (6 years ago)

Author:

dubos

Message:

devel/unstructured : moved mesh partitioning code into meshes.py

Location:

codes/icosagcm/devel/Python

Files:

• dynamico/meshes.py (modified) (8 diffs)
• dynamico/util.py (added)
• src/unstructured.pyx (modified) (3 diffs)
• test/py/DCMIP2008c5.py (modified) (2 diffs)
• test/py/NH_3D_DCMIP31.py (modified) (2 diffs)
• test/py/RSW2_MPAS_W02.py (modified) (2 diffs)
• test/py/RSW_MPAS_W02.py (modified) (2 diffs)
• test/py/partition.py (modified) (3 diffs)
• test/py/test_xios.py (modified) (2 diffs)

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

@@ -5 +5 @@
 import matplotlib.tri as tri
 import matplotlib.pyplot as plt
-from unstructured import init_mesh, list_stencil, ker
+from matplotlib.patches import Polygon
+from matplotlib.collections import PatchCollection
+
+import unstructured as unst
+import parallel
+from util import list_stencil, Base_class
 
 radian=180/math.pi
@@ -35 +40 @@
                 mass_dbk[i,j,:] = dbk_ij
                 mass_bl[i,j,1:]= 1.-dbk_ij.cumsum()
-
+ 
     return mass_bl, mass_dak, mass_dbk
+
+# from theta.k90       : rhodz(l,ij) = mass_dak(l,ij) + mass_col(ij)*mass_dbk(l,ij)
+# from caldyn_wflux.k90: wflux(l,ij) = mass_bl(l,ij) * dmass_col(ij) - convm(l,ij)
+def mass_coefs_from_pressure_coefs(g, ap_il, bp_il):
+    # p = Mg + ptop = a + b.ps
+    # ps = Mcol.g+ptop
+    # M = mass_a + mass_b.Mcol
+    # M = (a+b.ps-ptop)/g
+    #   = (a+b.ptop-ptop)/g + b.Mcol
+    # mass_a  = (a+b.ptop)/g
+    # mass_da = (da+db.ptop)/g
+    # mass_b  = b
+    # mass_db = db
+    n, llm, ptop = ap_il.shape[0], ap_il.shape[1]-1, ap_il[0,-1]
+    print 'hybrid ptop : ', ptop 
+    mass_al = (ap_il+ptop*bp_il)/g
+    mass_dak, mass_dbk = np.zeros((n,llm)), np.zeros((n,llm))
+    for k in range(llm):
+        mass_dak[:,k] = mass_al[:,k]-mass_al[:,k+1]
+        mass_dbk[:,k] = bp_il[:,k]-bp_il[:,k+1]
+    print mass_al[0,:]
+    print mass_dak[0,:]
+    print mass_dbk[0,:]
+    return [ np.asarray(x) for x in bp_il, mass_dak, mass_dbk ]
 
 #----------------------- Cartesian mesh -----------------------
@@ -142 +171 @@
                     (indexu(x-1,y+1),.25)))
         Aiv=np.zeros(nx*ny)+dx*dy # Ai=Av=dx*dy
-        init_mesh(llm,nqdyn,2*nx*ny,nx*ny,nx*ny,4,4,4,
+        unst.init_mesh(llm,nqdyn,2*nx*ny,nx*ny,nx*ny,4,4,4,
                   primal_nb,primal_edge,primal_ne,
                   dual_nb,dual_edge,dual_ne,dual_vertex,
@@ -164 +193 @@
 #---------------------- MPAS fully unstructured mesh ------------------------
 
-def compute_ne(num,deg,edges,left,right):
-    ne = np.zeros_like(edges)
-    for cell in range(num):
-        for iedge in range(deg[cell]):
-            edge = edges[cell,iedge]-1
-            if left[edge]==cell+1: ne[cell,iedge]+=1
-            if right[edge]==cell+1: ne[cell,iedge]-=1
-            if ne[cell,iedge]==0 : print 'error at cell,iedge', cell, iedge
-    return ne
-
-def plot(tri,data, **kwargs):
-    plt.figure(figsize=(12,4))
-    plt.gca().set_aspect('equal')
-    plt.tricontourf(tri, data, 20, **kwargs)
-    plt.colorbar()
-    plt.ylim((-90,90))
-    plt.xlim((0,360))
-
-class MPAS_Mesh:
-    def __init__(self, gridfile, llm, nqdyn, radius, f):
-        self.gridfile, self.llm, self.nqdyn = gridfile,llm,nqdyn
-        self.radius, self.f = radius, f
-        # open mesh file, get main dimensions
+class Mesh_Format(Base_class):
+    def getdims(self,*names): return [len(self.nc.dimensions[name]) for name in names]
+    def get_pdims(self,comm,*names): return [self.get_pdim(comm,name) for name in names]
+    def get_pvars(self,pdim,*names): return [self.get_pvar(pdim,name) for name in names]
+    def getvars(self,*names): 
+            for name in names : print "getvar %s ..."%name
+            time1=time.time()
+            ret=[self.getvar(name) for name in names]
+            print "... Done in %f seconds"%(time.time()-time1)
+            return ret
+    
+class MPAS_Format(Mesh_Format):
+    start_index=1 # indexing starts at 1 as in Fortran
+    translate= {
+        'primal_num':'nCells',
+        'edge_num':'nEdges',
+        'dual_num':'nVertices',
+        'primal_deg':'nEdgesOnCell', 'trisk_deg':'nEdgesOnEdge',
+        'primal_edge':'edgesOnCell', 'left_right':'cellsOnEdge',
+        'dual_edge':'edgesOnVertex', 'down_up':'verticesOnEdge',
+        'primal_vertex':'verticesOnCell', 'dual_vertex':'cellsOnVertex',
+        'trisk':'edgesOnEdge', 'down_up':'verticesOnEdge', 'left_right':'cellsOnEdge',
+        'le':'dvEdge', 'de':'dcEdge', 'Ai':'areaCell', 'Av':'areaTriangle',
+        'lat_i':'latCell','lon_i':'lonCell','lat_v':'latVertex','lon_v':'lonVertex',
+        'lat_e':'latEdge','lon_e':'lonEdge','angle_e':'angleEdge',
+        'wee':'weightsOnEdge','Riv2':'kiteAreasOnVertex'}
+    def __init__(self,gridfilename):
         try:
-            nc = cdf.Dataset(gridfile, "r")
+            self.nc = cdf.Dataset(gridfilename, "r")
         except RuntimeError:
             print """
@@ -195 +228 @@
             """ % gridfile
             raise
-
-        def getdims(*names): return [len(nc.dimensions[name]) for name in names]
-        def getvars(*names): 
-            for name in names : print "getvar %s ..."%name
-            time1=time.time()
-            ret=[nc.variables[name][:] for name in names]
-            print "... Done in %f seconds"%(time.time()-time1)
-            return ret
-        primal_num, edge_num, dual_num = getdims('nCells','nEdges','nVertices')                                   
-        print 'Number of primal cells, dual cells and edges : %d, %d, %d '%(primal_num,dual_num,edge_num)
-        primal_deg, trisk_deg = getvars('nEdgesOnCell','nEdgesOnEdge')
-        dual_deg = [3 for i in range(dual_num) ]
-        dual_deg = np.ascontiguousarray(dual_deg,dtype=np.int32) # NB : Fortran code expects 32-bit ints
-
-        # get indices for stencils 
-        # primal <-> edges
-        primal_edge, left_right = getvars('edgesOnCell','cellsOnEdge')
-        left,right = left_right[:,0], left_right[:,1]
-        primal_ne = compute_ne(primal_num,primal_deg,primal_edge,left,right)
-        # dual <-> edges
-        dual_edge, down_up = getvars('edgesOnVertex','verticesOnEdge')
-        down,up = down_up[:,0], down_up[:,1]
-        dual_ne = compute_ne(dual_num,dual_deg,dual_edge,up,down)
-        # primal <-> dual, edges <-> edges (TRiSK)
-        primal_vertex, dual_vertex, trisk = getvars('verticesOnCell','cellsOnVertex','edgesOnEdge')
-        # get positions, lengths, surfaces and weights
-        le,de,Ai,Av = getvars('dvEdge','dcEdge','areaCell','areaTriangle')
-        lat_i,lon_i = getvars('latCell','lonCell')
-        lat_v,lon_v = getvars('latVertex','lonVertex')
-        lat_e,lon_e,angle_e = getvars('latEdge','lonEdge','angleEdge')
-        wee,Riv2 = getvars('weightsOnEdge','kiteAreasOnVertex')
-
+    def get_pdim(self,comm,name): 
+        name = self.translate[name]
+        return parallel.PDim(self.nc.dimensions[name], comm)
+    def getvar(self,name):
+        # first deal with special cases
+        if name == 'dual_deg':            
+            dual_deg = [3 for i in range(self.dual_num) ]
+            dual_deg = np.ascontiguousarray(dual_deg,dtype=np.int32) 
+            # NB : Fortran code expects 32-bit ints
+            return dual_deg
+        # general case
+        name=self.translate[name]
+        return self.nc.variables[name][:]
+    def get_pvar(self,pdim,name): 
+        # provides parallel access to a NetCDF variable, with name translation
+        # first deal with special cases
+        if name == 'dual_deg': return parallel.CstPArray1D(pdim, np.int32, 3)
+        if name == 'edge_deg': return parallel.CstPArray1D(pdim, np.int32, 2)
+        # general case
+        name=self.translate[name]
+        return parallel.PArray(pdim, self.nc.variables[name])
+    def normalize(self, mesh, radius):
+        (trisk_deg, trisk, Ai,
+         de, le, wee, Av, Riv2) = (mesh.trisk_deg, mesh.trisk, mesh.Ai, 
+                                       mesh.de, mesh.le, mesh.wee, mesh.Av, mesh.Riv2)
+        edge_num, dual_num, r2 = de.size, Av.size, radius**2
         # fix normalization of wee and Riv2 weights
         for edge1 in range(edge_num):
@@ -238 +265 @@
             r=sum(r)
             if abs(r-1.)>1e-6 : print 'error with Riv2 at vertex ', ivertex
-            
-        max_primal_deg, max_dual_deg, max_trisk_deg = getdims('maxEdges','vertexDegree','maxEdges2')
-
-        # CRITICAL : force arrays left, etc. to be contiguous in memory:
-        left,right,up,down = [np.ascontiguousarray(x) for x in (left,right,up,down)]
-        trisk,wee,primal_edge = [np.ascontiguousarray(x) for x in (trisk,wee,primal_edge)]
+        # scale lengths and areas
+        mesh.le, mesh.de, mesh.Av, mesh.Ai = radius*le, radius*de, r2*Av, r2*Ai
         
-        print ('Max stencil sizes (div,curl,trisk) : %d, %d, %d' 
-               % (max_primal_deg, max_dual_deg, max_trisk_deg) )
-
-        r2 = radius**2
-        Av = r2*Av
-        fv = f(lon_v,lat_v) # Coriolis parameter
-        self.Ai, self.Av, self.fv = r2*Ai,Av,fv
-        self.le, self.de, self.le_de = radius*le, radius*de, le/de
-        self.trisk_deg, self.trisk, self.wee = trisk_deg, trisk, wee
-        init_mesh(llm, nqdyn, edge_num, primal_num,dual_num,
-                  max_trisk_deg, max_primal_deg, max_dual_deg,
-                  primal_deg,primal_edge,primal_ne,
-                  dual_deg,dual_edge,dual_ne,dual_vertex,
-                  left,right,down,up,trisk_deg,trisk,
-                  self.Ai,self.Av,self.fv,le/de,Riv2,wee)
-        
-        for edge in range(edge_num):
-            iedge = trisk[edge,0:trisk_deg[edge]]
-            if iedge.min()<1 :
-                print 'error'
-                          
-        self.primal_num, self.edge_num, self.dual_num = primal_num, edge_num, dual_num
-        def period(x) : return (x+2*math.pi)%(2*math.pi)
-        lon_i, lon_v, lon_e = map(period, (lon_i,lon_v,lon_e))
-        self.lon_i, self.lat_i = lon_i, lat_i
-        self.lon_v, self.lat_v = lon_v, lat_v
-        self.lon_e, self.lat_e, self.angle_e = lon_e, lat_e, angle_e
-        self.primal_deg, self.primal_vertex = primal_deg, primal_vertex
-        self.primal = tri.Triangulation(lon_i*180./math.pi, lat_i*180./math.pi)
-        self.dual_deg, self.dual_vertex = dual_deg, dual_vertex
-        self.dual = tri.Triangulation(lon_v*180./math.pi, lat_v*180./math.pi)
-        self.triedge = tri.Triangulation(lon_e*180./math.pi, lat_e*180./math.pi)
-        self.dx = de.min()
-        self.lon3D_i, self.ll3D = np.meshgrid(lon_i, range(llm))
-        self.lat3D_i, self.ll3D = np.meshgrid(lat_i, range(llm))
+def compute_ne(num,deg,edges,left,right):
+    ne = np.zeros_like(edges)
+    for cell in range(num):
+        for iedge in range(deg[cell]):
+            edge = edges[cell,iedge]-1
+            if left[edge]==cell+1: ne[cell,iedge]+=1
+            if right[edge]==cell+1: ne[cell,iedge]-=1
+            if ne[cell,iedge]==0 : print 'error at cell,iedge', cell, iedge
+    return ne
+
+class Abstract_Mesh(Base_class):
     def field_theta(self,n=1): return squeeze((n,self.nqdyn,self.primal_num,self.llm))
     def field_mass(self,n=1):  return squeeze((n,self.primal_num,self.llm))
@@ -294 +293 @@
             ucov[edge,:] = self.de[edge]*(ulon[edge,:]*math.cos(angle)+ulat[edge,:]*math.sin(angle))
         return ucov
+    def plot(self, tri,data, **kwargs):
+        plt.figure(figsize=(12,4))
+        plt.gca().set_aspect('equal')
+        plt.tricontourf(tri, data, 20, **kwargs)
+        plt.colorbar()
+        plt.ylim((-90,90))
+        plt.xlim((0,360))
     def plot_i(self,data, **kwargs):
-        plot(self.primal,data, **kwargs)
+        self.plot(self.primal,data, **kwargs)
     def plot_v(self,data, **kwargs):
-        plot(self.dual,data, **kwargs)
+        self.plot(self.dual,data, **kwargs)
     def plot_e(self,data, **kwargs):
-        plot(self.triedge,data, **kwargs)
+        self.plot(self.triedge,data, **kwargs)
+    
+class Unstructured_Mesh(Abstract_Mesh):
+    def __init__(self, gridfile, llm, nqdyn, radius, f):
+        # open mesh file, get main dimensions
+        getvars = gridfile.getvars
+        # get positions, lengths, surfaces and weights
+        le,de,Ai,Av,wee,Riv2 = getvars('le','de','Ai','Av','wee','Riv2')
+        lat_i,lon_i,lat_v,lon_v = getvars('lat_i','lon_i','lat_v','lon_v')
+        lat_e,lon_e,angle_e = getvars('lat_e','lon_e','angle_e')
+        fv = f(lon_v,lat_v) # Coriolis parameter
+        primal_num, dual_num, edge_num = [ x.size for x in Ai, Av, le ]
+        gridfile.primal_num, gridfile.dual_num = primal_num, dual_num
+
+        primal_deg, dual_deg, trisk_deg = getvars('primal_deg','dual_deg','trisk_deg')
+        
+        # get indices for stencils 
+        # primal <-> edges
+        primal_edge, left_right = getvars('primal_edge','left_right')
+        left,right = left_right[:,0], left_right[:,1]
+        primal_ne = compute_ne(primal_num,primal_deg,primal_edge,left,right)
+        # dual <-> edges
+        dual_edge, down_up = getvars('dual_edge','down_up')
+        down,up = down_up[:,0], down_up[:,1]
+        dual_ne = compute_ne(dual_num,dual_deg,dual_edge,up,down)
+        # primal <-> dual, edges <-> edges (TRiSK)
+        primal_vertex, dual_vertex, trisk = getvars('primal_vertex','dual_vertex','trisk')
+        # CRITICAL : force arrays left, etc. to be contiguous in memory:
+        left,right,up,down = [np.ascontiguousarray(x) for x in (left,right,up,down)]
+        trisk,wee,primal_edge = [np.ascontiguousarray(x) for x in (trisk,wee,primal_edge)]
+        self.set_members(locals(), 'gridfile','llm','nqdyn','radius','f',
+                         'primal_num', 'dual_num', 'edge_num', 'primal_deg', 'primal_vertex',
+                         'lat_i', 'lon_i', 'lat_v', 'lon_v', 'lat_e', 'lon_e', 'angle_e',
+                         'Ai', 'Av', 'fv', 'le', 'de',
+                         'trisk_deg', 'trisk', 'wee', 'Riv2')
+        gridfile.normalize(self, radius)
+        self.le_de = self.le/self.de
+        max_primal_deg, max_dual_deg, max_trisk_deg = [ x.shape[1] for x in primal_edge, dual_edge, trisk]
+        unst.init_mesh(llm, nqdyn, edge_num, primal_num,dual_num,
+                  max_trisk_deg, max_primal_deg, max_dual_deg,
+                  primal_deg,primal_edge,primal_ne,
+                  dual_deg,dual_edge,dual_ne,dual_vertex,
+                  left,right,down,up,trisk_deg,trisk,
+                  self.Ai,self.Av,self.fv,le/de,Riv2,wee)
+        self.primal  = tri.Triangulation(lon_i*180./math.pi, lat_i*180./math.pi)
+        self.dual    = tri.Triangulation(lon_v*180./math.pi, lat_v*180./math.pi)
+        self.triedge = tri.Triangulation(lon_e*180./math.pi, lat_e*180./math.pi)        
+        
+class Unstructured_PMesh(Abstract_Mesh): # Mesh data distributed across MPI processes
+    def __init__(self,comm, gridfile):
+        self.comm, self.gridfile, get_pvars = comm, gridfile, gridfile.get_pvars
+        dim_primal, dim_edge, dim_dual = gridfile.get_pdims(comm,
+            'primal_num','edge_num','dual_num')
+        (primal_deg, primal_vertex, primal_edge, lon_i, 
+         lat_i) = get_pvars(dim_primal, 'primal_deg', 'primal_vertex', 'primal_edge', 'lon_i', 'lat_i' )
+        (edge_deg, trisk_deg, left_right, down_up, 
+         trisk) = get_pvars(dim_edge, 'edge_deg', 'trisk_deg', 'left_right', 'down_up', 'trisk')
+        dual_deg, dual_vertex, dual_edge, Riv2, lon_v, lat_v = get_pvars(dim_dual, 
+            'dual_deg', 'dual_vertex', 'dual_edge', 'Riv2', 'lon_v', 'lat_v')
+        # Let indices start at 0
+        for x in (primal_vertex, primal_edge, dual_vertex,dual_edge,
+            left_right, down_up, trisk ) : x.data = x.data-gridfile.start_index
+        self.edge2cell   = Stencil_glob(edge_deg, left_right)
+        self.cell2edge   = Stencil_glob(primal_deg, primal_edge)
+        self.edge2vertex = Stencil_glob(edge_deg, down_up)
+        self.vertex2edge = Stencil_glob(dual_deg, dual_edge)
+        self.edge2edge   = Stencil_glob(trisk_deg, trisk)
+        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', 'edge_owner', 
+                         'primal_owner', 'primal_deg', 'primal_vertex', 
+                         'lon_v', 'lat_v', 'lon_i', 'lat_i')
+        #        self.dim_primal, self.edge_owner, self.primal_owner = dim_primal, edge_owner, primal_owner
+
+def chain(start, links):
+    for link in links:
+        start = link(start).neigh_set
+        yield start
+
+def patches(degree, bounds, lon, lat):
+    for i in range(degree.size):
+        nb_edge=degree[i]
+        bounds_cell = bounds[i,0:nb_edge]
+        lat_cell    = lat[bounds_cell]
+        lon_cell    = lon[bounds_cell]
+        orig=lon_cell[0]
+        lon_cell    = lon_cell-orig+180.
+        lon_cell    = np.mod(lon_cell,360.)
+        lon_cell    = lon_cell+orig-180.
+#        if np.abs(lon_cell-orig).max()>10. :
+#            print '%d patches :'%mpi_rank, lon_cell
+        lonlat_cell = np.zeros((nb_edge,2))
+        lonlat_cell[:,0],lonlat_cell[:,1] = lon_cell,lat_cell
+        polygon = Polygon(lonlat_cell, True)
+        yield polygon
+
+class Local_Mesh(Base_class):
+    def __init__(self, pmesh):
+        dim_primal, edge_owner, primal_owner = pmesh.dim_primal, pmesh.edge_owner, pmesh.primal_owner
+        print 'Constructing halos ...'
+        edges_E0 = find_my_cells(edge_owner)
+        cells_C0, edges_E1, vertices_V1, edges_E2, cells_C1 = chain(
+            edges_E0, ( pmesh.edge2cell, pmesh.cell2edge, pmesh.edge2vertex, pmesh.vertex2edge, pmesh.edge2cell) )
+        edges_E0, edges_E1, edges_E2 = progressive_list(edges_E0, edges_E1, edges_E2)
+        cells_C0, cells_C1 = progressive_list(cells_C0, cells_C1)
+        print 'E2,E1,E0 ; C1,C0 : ', map(len, (edges_E2, edges_E1, edges_E0, cells_C1, cells_C0))
+        mycells, halo_cells = cells_C0, cells_C1
+        get_mycells, get_halo_cells = dim_primal.getter(mycells), dim_primal.getter(halo_cells)
+        self.com_primal = parallel.Halo_Xchange(42, dim_primal, halo_cells, get_halo_cells(primal_owner))
+    def plot(self, ax, clim, degree, bounds, lon, lat, data):
+        nb_vertex = lon.size # global
+        p = list(patches(degree, bounds, lon, lat))
+        p = PatchCollection(p, linewidth=0.01)
+        p.set_array(data) # set values at each polygon (cell)
+        p.set_clim(clim)
+        ax.add_collection(p)
 
 #--------------------------------------  Mesh reordering  ------------------------------------------#
@@ -341 +465 @@
     nn, xx, yy, zz = x.size, toint(x), toint(y), toint(z)
     idx=np.zeros(nn, dtype=np.int64)
-    ker.dynamico_morton_encode(nn, xx,yy,zz, idx)
+    unst.ker.dynamico_morton_encode(nn, xx,yy,zz, idx)
     return idx
 

codes/icosagcm/devel/Python/src/unstructured.pyx

@@ -1 +1 @@
 import time
 import math
-import dynamico.wrap as wrap
-from dynamico.libs import libicosa
+import wrap
+from libs import libicosa
+from util import list_stencil
 
 from ctypes import c_void_p, c_int, c_double, c_bool
@@ -289 +290 @@
 
 # Helper functions and interface to ParMETIS
-# list_stencil converts an adjacency graph from array format index[num_cells, MAX_EDGES] to compressed format
 # loc_stencil returns the start/end indices (vtxdist) expected by ParMETIS
 # i.e. index[start:end] with start=vtxdist[cell], end=vtxdist[cell+1] lists the edges of cell 'cell'
@@ -295 +295 @@
 import numpy as np
 
-def list_stencil(degree, stencil, cond=lambda x:True):
-    for i in range(degree.size):
-        for j in range(degree[i]):
-            s=stencil[i,j]
-            if cond(s): yield stencil[i,j]
-                
 def loc_stencil(degree, stencil):
     loc=0

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

@@ -4 +4 @@
 from dynamico import DCMIP
 from dynamico import meshes
-from dynamico.meshes import MPAS_Mesh as Mesh
 import math as math
 import matplotlib.pyplot as plt
@@ -45 +44 @@
     nqdyn, preff, Treff = 1, 1e5, 300.
     thermo = dyn.Ideal_perfect(Cpd, Rd, preff, Treff)
-    mesh = Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f)
+    meshfile = meshes.MPAS_Format('grids/x1.%d.grid.nc'%grid)
+    mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
     lev,levp1 = np.arange(llm),np.arange(llm+1)
     lon_i, lat_i, lon_e, lat_e =  mesh.lon_i, mesh.lat_i, mesh.lon_e, mesh.lat_e

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

@@ -4 +4 @@
 from dynamico import DCMIP
 from dynamico import meshes
-from dynamico.meshes import MPAS_Mesh as Mesh
 import math as math
 import matplotlib.pyplot as plt
@@ -49 +48 @@
     nqdyn, preff, Treff = 1, 1e5, 300.
     thermo = dyn.Ideal_perfect(Cpd, Rd, preff, Treff)
-    mesh = Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f)
+    meshfile = meshes.MPAS_Format('grids/x1.%d.grid.nc'%grid)
+    mesh = meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
     lev,levp1 = np.arange(llm),np.arange(llm+1)
     lat_ik,k_i = np.meshgrid(mesh.lat_i,lev, indexing='ij')

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

@@ -1 +1 @@
 print 'Loading DYNAMICO modules ...'
 from dynamico import unstructured as unst
-from dynamico.meshes import MPAS_Mesh as Mesh
+from dynamico.meshes import MPAS_Format, Unstructured_Mesh as Mesh
 from dynamico import time_step
 print '...Done'
@@ -19 +19 @@
 def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter
 print 'Reading MPAS mesh ...'
-mesh = Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f)
+meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
+mesh = Mesh(meshfile, llm, nqdyn, radius, f)
 print '...Done'
 lon, lat = mesh.lon_i, mesh.lat_i

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

@@ -1 +1 @@
 print 'Loading DYNAMICO modules ...'
 from dynamico import unstructured as unst
-from dynamico.meshes import MPAS_Mesh as Mesh
+from dynamico.meshes import MPAS_Format, Unstructured_Mesh
 from dynamico import time_step
 print '...Done'
@@ -19 +19 @@
 def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter
 print 'Reading MPAS mesh ...'
-mesh = Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f)
+meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
+mesh=Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
 print '...Done'
 lon, lat = mesh.lon_i, mesh.lat_i

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

@@ -57 +57 @@
 # Helper functions to plot unstructured graph
 
-def patches(degree, bounds, lon, lat):
-    for i in range(degree.size):
-        nb_edge=degree[i]
-        bounds_cell = bounds[i,0:nb_edge]
-        lat_cell    = lat[bounds_cell]
-        lon_cell    = lon[bounds_cell]
-        orig=lon_cell[0]
-        lon_cell    = lon_cell-orig+180.
-        lon_cell    = np.mod(lon_cell,360.)
-        lon_cell    = lon_cell+orig-180.
-#        if np.abs(lon_cell-orig).max()>10. :
-#            print '%d patches :'%mpi_rank, lon_cell
-        lonlat_cell = np.zeros((nb_edge,2))
-        lonlat_cell[:,0],lonlat_cell[:,1] = lon_cell,lat_cell
-        polygon = Polygon(lonlat_cell, True)
-        yield polygon
-
-def plot_mesh(ax, clim, degree, bounds, lon, lat, data):
-    nb_vertex = lon.size # global
-    p = list(patches(degree, bounds, lon, lat))
-    print '%d : plot_mesh %d %d %d'%( mpi_rank, degree.size, len(p), len(data) ) 
-    p = PatchCollection(p, linewidth=0.01)
-    p.set_array(data) # set values at each polygon (cell)
-    p.set_clim(clim)
-    ax.add_collection(p)
-
 def local_mesh(get_mycells):
-    mydegree, mybounds = [get_mycells(x) for x in nEdgesOnCell, verticesOnCell]
+    #    mydegree, mybounds = [get_mycells(x) for x in nEdgesOnCell, verticesOnCell]
+    mydegree, mybounds = [get_mycells(x) for x in primal_deg, primal_vertex]
     print '%d : len(mydegree)=%d'%(mpi_rank, len(mydegree))
     vertex_list = sorted(set(unst.list_stencil(mydegree,mybounds))) 
@@ -94 +69 @@
     return vertex_list, mydegree, mybounds, mylon, mylat
 
+def members(struct, *names): return [struct.__dict__ [name] for name in names]
+
 #--------------- read MPAS grid file ---------------#
 
-#grid = 'x1.2562'
-grid = 'x1.10242'
+grid = 'x1.2562'
+#grid = 'x1.10242'
 #grid = 'x4.163842'
 print 'Reading MPAS file %s ...'%grid
 
-nc = cdf.Dataset('grids/%s.grid.nc'%grid, "r")
-dim_cell, dim_edge, dim_vertex = [
-    parallel.PDim(nc.dimensions[name], comm) 
-    for name in 'nCells','nEdges','nVertices']
-edge_degree   = parallel.CstPArray1D(dim_edge, np.int32, 2)
-vertex_degree = parallel.CstPArray1D(dim_vertex, np.int32, 3)
-nEdgesOnCell, verticesOnCell, edgesOnCell, cellsOnCell, latCell = [
-    parallel.PArray(dim_cell, nc.variables[var])
-    for var in 'nEdgesOnCell', 'verticesOnCell', 'edgesOnCell', 'cellsOnCell', 'latCell' ]
-cellsOnVertex, edgesOnVertex, kiteAreasOnVertex, lonVertex, latVertex = [
-    parallel.PArray(dim_vertex, nc.variables[var])
-    for var in 'cellsOnVertex', 'edgesOnVertex', 'kiteAreasOnVertex', 'lonVertex', 'latVertex']
-nEdgesOnEdge, cellsOnEdge, edgesOnEdge, verticesOnEdge, weightsOnEdge = [
-    parallel.PArray(dim_edge, nc.variables[var])
-    for var in 'nEdgesOnEdge', 'cellsOnEdge', 'edgesOnEdge', 'verticesOnEdge', 'weightsOnEdge']
+meshfile = meshes.MPAS_Format('grids/%s.grid.nc'%grid)
+pmesh = meshes.Unstructured_PMesh(comm, meshfile)
+lmesh = meshes.Local_Mesh(pmesh)
 
-# Indices start at 0 on the C/Python side and at 1 on the Fortran/MPAS side
-# hence an offset of 1 is added/substracted where needed.
-for x in (verticesOnCell, edgesOnCell, cellsOnCell, cellsOnVertex, edgesOnVertex,
-          cellsOnEdge, edgesOnEdge, verticesOnEdge) : x.data = x.data-1
-edge2cell, cell2edge, edge2vertex, vertex2edge, cell2cell, edge2edge = [
-    meshes.Stencil_glob(a,b) for a,b in 
-    (edge_degree, cellsOnEdge), (nEdgesOnCell, edgesOnCell),
-    (edge_degree, verticesOnEdge), (vertex_degree, edgesOnVertex),
-    (nEdgesOnCell, cellsOnCell), (nEdgesOnEdge, edgesOnEdge) ]
+(primal_deg, primal_vertex, dim_vertex, dim_cell, cell_owner, 
+ lonVertex, latVertex, lonCell, latCell) = members(
+    pmesh, 'primal_deg', 'primal_vertex', 'dim_dual', 'dim_primal', 'primal_owner', 
+    'lon_v', 'lat_v', 'lon_i', 'lat_i')
 
-#---------------- partition edges and cells ------------------#
-
-print 'Partitioning ...'
-
-edge_owner = unst.partition_mesh(nEdgesOnEdge, edgesOnEdge, mpi_size)
-edge_owner = parallel.LocPArray1D(dim_edge, edge_owner)
-cell_owner = meshes.partition_from_stencil(edge_owner, nEdgesOnCell, edgesOnCell)
-cell_owner = parallel.LocPArray1D(dim_cell, cell_owner)
-
-#--------------------- construct halos  -----------------------#
-
-print 'Constructing halos ...'
-
-def chain(start, links):
-    for link in links:
-        start = link(start).neigh_set
-        yield start
-
-edges_E0 = meshes.find_my_cells(edge_owner)
-cells_C0, edges_E1, vertices_V1, edges_E2, cells_C1 = chain(
-    edges_E0, ( edge2cell, cell2edge, edge2vertex, vertex2edge, edge2cell) )
-
-edges_E0, edges_E1, edges_E2 = meshes.progressive_list(edges_E0, edges_E1, edges_E2)
-cells_C0, cells_C1 = meshes.progressive_list(cells_C0, cells_C1)
-
-print 'E2,E1,E0 ; C1,C0 : ', map(len, (edges_E2, edges_E1, edges_E0, cells_C1, cells_C0))
-
-#com_edges = parallel.Halo_Xchange(24, dim_edge, edges_E2, dim_edge.get(edges_E2, edge_owner))
-
-mycells, halo_cells = cells_C0, cells_C1
-get_mycells, get_halo_cells = dim_cell.getter(mycells), dim_cell.getter(halo_cells)
-com_cells = parallel.Halo_Xchange(42, dim_cell, halo_cells, get_halo_cells(cell_owner))
-
-local_num, total_num = np.zeros(1), np.zeros(1)
+local_num, total_num, com_cells = np.zeros(1), np.zeros(1), lmesh.com_primal
 local_num[0]=com_cells.own_len
 comm.Reduce(local_num, total_num, op=MPI.SUM, root=0)
@@ -167 +94 @@
 #---------------------------- plot -----------------------------#
 
-if True:
-    print 'Plotting ...'
+print 'Plotting ...'
 
-    halo_vertex_list, mydegree, mybounds, mylon, mylat = local_mesh(com_cells.get_all)
-    buf = parallel.LocalArray1(com_cells)
+halo_vertex_list, mydegree, mybounds, mylon, mylat = local_mesh(com_cells.get_all)
+buf = parallel.LocalArray1(com_cells)
 
-    fig, ax = plt.subplots()
-    buf.read_own(latCell) # reads only own values
-    buf.data = np.cos(10.*buf.data)
-    buf.update() # updates halo
-    plot_mesh(ax,[-math.pi/2,math.pi/2], mydegree, mybounds, mylon, mylat, buf.data)
-    plt.xlim(-190.,190.)
-    plt.ylim(-90.,90.)
-    plt.savefig('fig_partition/A%03d.pdf'%mpi_rank, dpi=1600)
+fig, ax = plt.subplots()
+buf.read_own(latCell) # reads only own values
+buf.data = np.cos(10.*buf.data)
+buf.update() # updates halo
+lmesh.plot(ax,[-math.pi/2,math.pi/2], mydegree, mybounds, mylon, mylat, buf.data)
+plt.xlim(-190.,190.)
+plt.ylim(-90.,90.)
+plt.savefig('fig_partition/A%03d.pdf'%mpi_rank, dpi=1600)
 
-    fig, ax = plt.subplots()
-    buf.read_own(cell_owner)
-    buf.update()
-    plot_mesh(ax,[0,mpi_rank+1], mydegree, mybounds, mylon, mylat, buf.data)
-    plt.xlim(-190.,190.)
-    plt.ylim(-90.,90.)
-    plt.savefig('fig_partition/B%03d.pdf'%mpi_rank, dpi=1600)
+fig, ax = plt.subplots()
+buf.read_own(cell_owner)
+buf.update()
+lmesh.plot(ax,[0,mpi_rank+1], mydegree, mybounds, mylon, mylat, buf.data)
+plt.xlim(-190.,190.)
+plt.ylim(-90.,90.)
+plt.savefig('fig_partition/B%03d.pdf'%mpi_rank, dpi=1600)

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

@@ -1 +1 @@
 #print 'import dynamico'
 #import dynamico
-print 'import dynamico.unstructured'
 import dynamico.unstructured as unst
-print 'import dynamico.xios'
 import dynamico.xios as xios
-print 'Done.'
-from dynamico.meshes import MPAS_Mesh as Mesh, radian
+from dynamico.meshes import MPAS_Format, Unstructured_Mesh, radian
 
 from mpi4py import MPI
@@ -38 +35 @@
 def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter                                                                                                                
 print 'Reading MPAS mesh ...'
-mesh = Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f)
+meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
+mesh = Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
 print '...Done'
 
 lon_i, lat_i, lon_v, lat_v = [x*radian for x in mesh.lon_i, mesh.lat_i, mesh.lon_v, mesh.lat_v]
-#lon_i, lat_i, lon_v, lat_v = mesh.lon_i, mesh.lat_i, mesh.lon_v, mesh.lat_v
 
 #--------------------------------- initialize XIOS --------------------------