import time import math import numpy as np import netCDF4 as cdf import matplotlib.tri as tri import matplotlib.pyplot as plt from matplotlib.patches import Polygon from matplotlib.collections import PatchCollection import unstructured as unst import parallel from util import list_stencil, Base_class, inverse_list from precision import zeros radian=180/math.pi #------------------- Hybrid mass-based coordinate ------------- # compute hybrid coefs from distribution of mass def compute_hybrid_coefs(mass): if mass.ndim==2 : nx,llm=mass.shape mass_dak = np.zeros((nx,llm)) mass_dbk = np.zeros((nx,llm)) mass_bl = np.zeros((nx,llm+1))+1. for i in range(nx): m_i = mass[i,:] dbk_i = m_i/sum(m_i) mass_dbk[i,:] = dbk_i mass_bl[i,1:]= 1.-dbk_i.cumsum() if mass.ndim==3 : nx,ny,llm=mass.shape mass_dak = np.zeros((nx,ny,llm)) mass_dbk = np.zeros((nx,ny,llm)) mass_bl = np.zeros((nx,ny,llm+1))+1. for i in range(nx): for j in range(ny): m_ij = mass[i,j,:] dbk_ij = m_ij/sum(m_ij) 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 ----------------------- # arrays is a list of arrays # vals is a list of tuples # each tuple is stored in each array def put(ij, deg, arrays, vals): k=0 for vv in vals: # vv is a tuple of values to be stored in arrays for array,v in zip(arrays,vv): array[ij,k]=v k=k+1 deg[ij]=k class Cartesian_mesh: def __init__(self,nx,ny,llm,nqdyn,Lx,Ly,f): dx,dy = Lx/nx, Ly/ny self.dx, self.dy, self.f = dx,dy,f self.nx, self.ny, self.llm, self.nqdyn = nx,ny,llm,nqdyn self.field_z = self.field_mass # 1D coordinate arrays x=(np.arange(nx)-nx/2.)*Lx/nx y=(np.arange(ny)-ny/2.)*Ly/ny lev=np.arange(llm) levp1=np.arange(llm+1) self.x, self.y, self.lev, self.levp1 = x,y,lev,levp1 # 3D coordinate arrays self.xx,self.yy,self.ll = np.meshgrid(x,y,lev, indexing='ij') self.xxp1,self.yyp1,self.llp1 = np.meshgrid(x,y,levp1, indexing='ij') # beware conventions for indexing # Fortran order : llm,nx*ny,nqdyn / indices start at 1 # Python order : nqdyn,ny,nx,llm / indices start at 0 # indices below follow Fortran while x,y follow Python/C index=lambda x,y : ((x+(nx*(y+2*ny)))%(nx*ny))+1 indexu=lambda x,y : 2*index(x,y)-1 indexv=lambda x,y : 2*index(x,y) indices = lambda shape : np.zeros(shape,np.int32) primal_nb = indices(nx*ny) primal_edge = indices((nx*ny,4)) primal_ne = indices((nx*ny,4)) dual_nb = indices(nx*ny) dual_edge = indices((nx*ny,4)) dual_ne = indices((nx*ny,4)) dual_vertex = indices((nx*ny,4)) Riv2 = np.zeros((nx*ny,4)) left = indices(2*nx*ny) right = indices(2*nx*ny) up = indices(2*nx*ny) down = indices(2*nx*ny) le_de = np.zeros(2*nx*ny) trisk_deg = indices(2*nx*ny) trisk = indices((2*nx*ny,4)) # 4 TRiSK coefs per edge wee = np.zeros((2*nx*ny,4)) for x in range(nx): for y in range(ny): # NB : Fortran indices start at 1 # primal cells put(index(x,y)-1,primal_nb,(primal_edge,primal_ne), ((indexu(x,y),1), (indexv(x,y),1), (indexu(x-1,y),-1), (indexv(x,y-1),-1) )) # dual cells put(index(x,y)-1,dual_nb,(dual_edge,dual_vertex,dual_ne,Riv2), ((indexv(x+1,y),index(x,y),1,.25), (indexu(x,y+1),index(x+1,y),-1,.25), (indexv(x,y),index(x+1,y+1),-1,.25), (indexu(x,y),index(x,y+1),1,.25) )) # edges : # left and right are adjacent primal cells # flux is positive when going from left to right # up and down are adjacent dual cells (vertices) # circulation is positive when going from down to up # u-points ij =indexu(x,y)-1 left[ij]=index(x,y) right[ij]=index(x+1,y) down[ij]=index(x,y-1) up[ij]=index(x,y) le_de[ij]=dy/dx put(ij,trisk_deg,(trisk,wee),( (indexv(x,y),-.25), (indexv(x+1,y),-.25), (indexv(x,y-1),-.25), (indexv(x+1,y-1),-.25))) # v-points ij = indexv(x,y)-1 left[ij]=index(x,y) right[ij]=index(x,y+1) down[ij]=index(x,y) up[ij]=index(x-1,y) le_de[ij]=dx/dy put(ij,trisk_deg,(trisk,wee),( (indexu(x,y),.25), (indexu(x-1,y),.25), (indexu(x,y+1),.25), (indexu(x-1,y+1),.25))) Aiv=np.zeros(nx*ny) Aiv[:]=dx*dy # Ai=Av=dx*dy 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, left,right,down,up,trisk_deg,trisk, Aiv,Aiv,f+0.*Aiv,le_de,Riv2,-wee) def field_theta(self,n=1): return zeros((n,self.nqdyn,self.ny,self.nx,self.llm)) def field_mass(self,n=1): return zeros((n,self.ny,self.nx,self.llm)) def field_z(self,n=1): return zeros((n,self.ny,self.nx,self.llm)) def field_w(self,n=1): return zeros((n,self.ny,self.nx,self.llm+1)) def field_u(self,n=1): return zeros((self.ny,2*self.nx,self.llm)) def field_ps(self,n=1): return zeros((n,self.ny,self.nx)) def ucomp(self,u): return u[range(0,2*self.nx,2),:] if self.ny==1 else u[:,range(0,2*self.nx,2),:] def set_ucomp(self,uv,u): if self.ny==1 : uv[range(0,2*self.nx,2),:]=u else : uv[:,range(0,2*self.nx,2),:]=u def vcomp(self,u): return u[:,range(1,2*self.nx,2),:] #---------------------- MPAS fully unstructured mesh ------------------------ 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: self.nc = cdf.Dataset(gridfilename, "r") except RuntimeError: print """ Unable to open grid file %s, maybe you forgot to download it ? To do so, go to the 'Python/' dir and execute './get_MPAS_grids.sh'. """ % gridfile raise 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): for i in range(trisk_deg[edge1]): edge2=trisk[edge1,i]-1 # NB Fortran vs C/Python indexing wee[edge1,i] = de[edge1]*wee[edge1,i]/le[edge2] for ivertex in range(dual_num): for j in range(3): Riv2[ivertex,j]=Riv2[ivertex,j]/Av[ivertex] r=Riv2[ivertex,:] r=sum(r) if abs(r-1.)>1e-6 : print 'error with Riv2 at vertex ', ivertex # scale lengths and areas mesh.le, mesh.de, mesh.Av, mesh.Ai = radius*le, radius*de, r2*Av, r2*Ai 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 'compute_ne error at cell,iedge', cell, iedge return ne class Abstract_Mesh(Base_class): def field_theta(self,n=1): return zeros((n,self.nqdyn,self.primal_num,self.llm)) def field_mass(self,n=1): return zeros((n,self.primal_num,self.llm)) def field_z(self,n=1): return zeros((n,self.dual_num,self.llm)) def field_w(self,n=1): return zeros((n,self.primal_num,self.llm+1)) def field_u(self,n=1): return zeros((n,self.edge_num,self.llm)) def field_ps(self,n=1): return zeros((n,self.primal_num,)) def ucov2D(self, ulon, ulat): return self.de*(ulon*np.cos(self.angle_e)+ulat*np.sin(self.angle_e)) def ucov3D(self, ulon, ulat): ucov = zeros((self.edge_num,ulon.shape[1])) for edge in range(self.edge_num): angle=self.angle_e[edge] 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=(10,4)) plt.gca().set_aspect('equal') plt.tricontourf(tri, data, 20, **kwargs) plt.colorbar() plt.ylim((-90,90)) plt.xlim((-180,180)) def plot_i(self,data, **kwargs): self.plot(self.primal,data, **kwargs) def plot_v(self,data, **kwargs): self.plot(self.dual,data, **kwargs) def plot_e(self,data, **kwargs): self.plot(self.triedge,data, **kwargs) class Unstructured_Mesh(Abstract_Mesh): def __init__(self, gridfile, llm, nqdyn, radius, f): 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*radian, lat_i*radian) self.dual = tri.Triangulation(lon_v*radian, lat_v*radian) self.triedge = tri.Triangulation(lon_e*radian, lat_e*radian) 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, Ai = get_pvars( dim_primal, 'primal_deg', 'primal_vertex', 'primal_edge', 'lon_i', 'lat_i', 'Ai' ) edge_deg, trisk_deg, left_right, down_up, trisk, le, de, wee, lon_e, lat_e, angle_e = get_pvars( dim_edge, 'edge_deg', 'trisk_deg', 'left_right', 'down_up', 'trisk', 'le', 'de', 'wee', 'lon_e', 'lat_e', 'angle_e') dual_deg, dual_vertex, dual_edge, Riv2, lon_v, lat_v, Av = get_pvars( dim_dual, 'dual_deg', 'dual_vertex', 'dual_edge', 'Riv2', 'lon_v', 'lat_v', 'Av') # 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.cell2vertex = Stencil_glob(primal_deg, primal_vertex) self.edge2vertex = Stencil_glob(edge_deg, down_up) self.vertex2edge = Stencil_glob(dual_deg, dual_edge) self.vertex2cell = Stencil_glob(dual_deg, dual_vertex, Riv2) self.edge2edge = Stencil_glob(trisk_deg, trisk, wee) 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', 'dim_edge', 'edge_owner', 'primal_owner', 'primal_deg', 'primal_vertex', 'lon_v', 'lat_v', 'Av', 'lon_i', 'lat_i', 'Ai', 'le', 'de', 'lon_e', 'lat_e', 'angle_e') def get(self, getter, *names): return [ getter(self.__dict__[x]) for x in names ] 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(Abstract_Mesh): def __init__(self, pmesh, llm, nqdyn, radius, f): 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) 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)) 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.lon_i, self.lat_i, self.Ai = pmesh.get(get_all_cells, 'lon_i', 'lat_i', 'Ai') self.lon_v, self.lat_v, self.Av = pmesh.get(get_all_duals, 'lon_v', 'lat_v', 'Av') self.fv = f(self.lon_v,self.lat_v) # Coriolis parameter self.le, self.de, self.lon_e, self.lat_e, self.angle_e = pmesh.get( 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) 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() ) print 'Primal degree min/max : %d %d'%(primal_edge.degree.min(), primal_edge.degree.max() ) print 'Dual degree min/max : %d %d'%(dual_edge.degree.min(), dual_edge.degree.max() ) 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 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] 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 left, right, down, up = left+1, right+1, down+1, up+1 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', '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, 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,self.le/self.de,Riv2,wee) # 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') self.primal = tri.Triangulation(self.lon_i*radian, self.lat_i*radian) self.dual = tri.Triangulation(self.lon_v*radian, self.lat_v*radian) self.triedge = tri.Triangulation(self.lon_e*radian, self.lat_e*radian) def plot_patches(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 ------------------------------------------# # NB : Indices start at 0 # permutations map an old index to a new index : perm[old_index]=new_index # enumerate(perm) is the list [ (0,perm[0]) , (1,perm[1]), ...] = [(old, new) ... ] def reorder_values(perm, *datas): # data in datas : ndarray containing indices ; perm : permutation to apply to values for data in datas: for x in np.nditer(data, op_flags=['readwrite']): x[...]=perm[x] def reorder_indices(perm, *datas): # data in datas : ndarray containing values ; perm : permutation to apply to indices for data in datas: data_out, ndim = data.copy(), data.ndim if ndim == 1: for old,new in enumerate(perm): data_out[new]=data[old] if ndim == 2: for old,new in enumerate(perm): data_out[new,:]=data[old,:] data[:]=data_out[:] def key_to_perm(keys): # keys : maps an old index to a key # returns : a list mapping old indices to new indices obtained by sorting keys in ascending order ordered = [old for key,old in sorted(zip(keys,range(len(keys))))] # ordered maps new indices to old indices # but we want to map old indices to new indices perm = list(range(len(keys))) for new,old in enumerate(ordered): perm[old]=new return perm # maps old indices to new indices def toint(x): return np.int32((x+1.)*65536) def morton_index(x,y,z): nn, xx, yy, zz = x.size, toint(x), toint(y), toint(z) idx=np.zeros(nn, dtype=np.int64) unst.ker.dynamico_morton_encode(nn, xx,yy,zz, idx) return idx #-------------------------------------- Mesh partitioning ------------------------------------------# # NB : Indices start at 0 def partition_from_stencil(owner2, degree, stencil): # given a stencil dim1->dim2 and owner2 on dim2, define owner[i] on dim1 as min(stencil[i,:] if i is even, max if odd dim1, dim2= degree.dim, owner2.dim degree, stencil, n = degree.data, stencil.data, dim1.end-dim1.start cells2 = list_stencil(degree, stencil) cells2 = sorted(list(set(list(cells2)))) # list of cells for which we need to know owner2 get2 = dim2.getter(cells2) owner1, owner2, glob2loc = np.zeros(n, dtype=np.int32), get2(owner2), get2.dict for i in range(n): owners = [ owner2[glob2loc[stencil[i,j]]] for j in range(degree[i]) ] if i%2 == 0 : owner1[i] = min(owners) else : owner1[i] = max(owners) return owner1 def find_my_cells(owner): # owner : a PArray1D containing the data returned by partition_mesh dim, comm, owner = owner.dim, owner.dim.comm, owner.data mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() cells=[set() for i in range(mpi_size)] for i in range(len(owner)): cells[owner[i]].add(dim.start+i) cells = [sorted(list(cells[i])) for i in range(mpi_size)] mycells = comm.alltoall(cells) mycells = sorted(sum(mycells, [])) # concatenate into a single list return mycells #---------------------------------- Stencil management ---------------------------------------# # NB : Indices start at 0 # Class Stencil represents an adjacency relationship (e.g. cell=>edges) # using adjacency information read from PArrays # It computes a list of "edges" adjacent to a given list of "cells" # This is used to form the sets E0 -> C0 -> E1 -> V1 -> E2 -> C1 # which are then used to form lists of global indices for V1,C1,E2 such that # C0, E0, E1 form contiguous subsets of C1, E2 starting from 0 def reindex(vertex_dict, degree, bounds): for i in range(degree.size): for j in range(degree[i]): bounds[i,j] = vertex_dict[bounds[i,j]] return bounds class Stencil_glob: def __init__(self, degree, neigh, weight=None): self.degree, self.neigh, self.weight = degree, neigh, weight def __call__(self, cells, neigh_dict=None): return Stencil(cells, self.degree, self.neigh, neigh_dict, self.weight) class Stencil: def __init__(self, cells, degree, neigh, neigh_dict, weight=None): get = degree.dim.getter(cells) mydegree, myneigh = [get(x) for x in (degree, neigh)] if not weight is None : myweight = get(weight) if neigh_dict is None : keep = lambda n : True else : # if neigh_dict is present, only neighbours in dict are retained keep = lambda n : n in neigh_dict neigh_set = list_stencil(mydegree, myneigh, keep) self.neigh_set = list(set(list(neigh_set) )) rej=0 for i in range(len(mydegree)): # keep only elements in neigh_dict, in-place k=0 for j in range(mydegree[i]): n=myneigh[i,j] if keep(n): myneigh[i,k]=n if not weight is None : myweight[i,k]=myweight[i,j] k=k+1 else: rej=rej+1 mydegree[i]=k if neigh_dict is None: neigh_dict = {j:i for i,j in enumerate(self.neigh_set)} myneigh_loc = reindex(neigh_dict, mydegree, myneigh) self.degree, self.neigh_glob, self.neigh_loc = mydegree, myneigh, myneigh_loc if weight is not None: self.weight = myweight def progressive_iter(mylist, cell_lists): for thelist in cell_lists: mylist = mylist + list(set(thelist)-set(mylist)) yield mylist def progressive_list(*cell_lists): # cell_lists : a tuple of lists of global indices, with each list a subset of the next # returns : a list 'mylist' such that for each list 'thelist' in cell_lists, thelist = mylist[0:len(thelist)] # example : edge_list = progressive_list(E0,E1,E2) with E0,E1,E2 increasing lists of cell edges return list(progressive_iter([], cell_lists))