from __future__ import absolute_import 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 from dynamico import parallel from dynamico.util import list_stencil, BaseClass, inverse_list from dynamico import getargs from dynamico.partition import partition_mesh log_master, log_world = getargs.getLogger(__name__) INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error INFO_ALL, DEBUG_ALL = log_world.info, log_world.debug radian=180/math.pi # convert from radians to degrees def Triangulation(lon,lat): return tri.Triangulation((lon*radian+360.)%360., radian*lat) #------------------- 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] DEBUG('hybrid ptop : %f' % 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] DEBUG('%s'%mass_al[0,:]) DEBUG('%s'%mass_dak[0,:]) DEBUG('%s'%mass_dbk[0,:]) return [ np.asarray(x) for x in bp_il, mass_dak, mass_dbk ] #----------------------- Base class for "user" meshes --------- class BaseMesh(BaseClass): def zeros(self,dims): return np.zeros([n for n in dims if n>1]) # overriden by meshes.dev.DevMesh #----------------------- 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 def concat(x,y): z = np.asarray([x,y]) return z.transpose() class Cartesian_Mesh(BaseMesh): 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.yy,self.xx,self.ll = np.meshgrid(y,x,lev, indexing='ij') self.yyp1,self.xxp1,self.llp1 = np.meshgrid(y,x,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 def periodize(z,nz): return (z+2*nz)%nz def index(x,y): return 1+periodize(x,nx)+nx*periodize(y,ny) def indexu(x,y): return 2*index(x,y)-1 def indexv(x,y): return 2*index(x,y) def zeros(shape,n=1): return [np.zeros(shape) for i in range(n)] if n>1 else np.zeros(shape) def indices(shape,n=1): return [np.zeros(shape,np.int32) for i in range(n)] if n>1 else np.zeros(shape,np.int32) primal_num, dual_num, edge_num = nx*ny, nx*ny, 2*nx*ny Aiv, lon_i,lon_v,lat_i,lat_v = zeros( nx*ny, 5) lon_e,lat_e,le,de,angle_e = zeros( 2*nx*ny, 5) Riv2 = zeros((nx*ny,4) ) wee = zeros((2*nx*ny,4)) primal_bounds_lon, primal_bounds_lat, dual_bounds_lon, dual_bounds_lat = zeros( (nx*ny,4), 4) primal_deg,dual_deg = indices( nx*ny, 2) primal_edge,primal_ne,primal_vertex = indices((nx*ny,4),3) dual_edge,dual_ne,dual_vertex = indices((nx*ny,4),3) trisk_deg,left,right,up,down = indices( 2*nx*ny, 5) trisk = indices((2*nx*ny,4)) # 4 TRiSK coefs per edge def putval(ij, arrays, vals): for array,v in zip(arrays,vals): array[ij]=v for x in range(nx): for y in range(ny): # NB : Fortran indices start at 1 # primal cells ij=index(x,y)-1 lon_i[ij], lat_i[ij] = x, y put(ij, primal_deg,(primal_edge,primal_vertex,primal_ne, primal_bounds_lon, primal_bounds_lat), ((indexu(x,y), index(x,y), 1, x+.5, y+.5), (indexv(x,y), index(x-1,y), 1, x-.5, y+.5), (indexu(x-1,y), index(x-1,y-1), -1, x-.5, y-.5), (indexv(x,y-1), index(x,y-1), -1, x+.5, y-.5) )) # dual cells ij=index(x,y)-1 lon_v[ij], lat_v[ij] = x+.5, y+.5 put(ij,dual_deg,(dual_edge,dual_vertex,dual_ne,Riv2, dual_bounds_lon, dual_bounds_lat), ((indexv(x+1,y),index(x,y), 1, .25, x, y), (indexu(x,y+1),index(x+1,y), -1, .25, x+1., y), (indexv(x,y), index(x+1,y+1), -1, .25, x+1., y+1.), (indexu(x,y), index(x,y+1), 1, .25, x, y+1.) )) # 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 putval(ij, (left, right, down, up, le, de, angle_e, lon_e, lat_e), (index(x,y), index(x+1,y), index(x,y-1), index(x,y), dy, dx, 0., x+.5, y ) ) 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 putval(ij, (left, right, down, up, le, de, angle_e, lon_e, lat_e), (index(x,y), index(x,y+1), index(x,y), index(x-1,y), dx, dy, .5*math.pi, x, y+.5 ) ) 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))) primal_i, primal_j = [ x.astype(np.int32) for x in lon_i, lat_i] edge_i, edge_j = [ x.astype(np.int32) for x in lon_e, lat_e] lon_i, lon_v, lon_e = [(x+.5)*dx-.5*Lx for x in lon_i, lon_v, lon_e] lat_i, lat_v, lat_e = [(y+.5)*dy-.5*Ly for y in lat_i, lat_v, lat_e] Aiv[:]=dx*dy # Ai=Av=dx*dy le_de, fv = le/de, f+0.*Aiv self.set_members(locals(), 'primal_num', 'dual_num', 'edge_num', 'primal_deg', 'primal_edge', 'primal_ne', 'primal_vertex', 'primal_bounds_lon', 'primal_bounds_lat', 'dual_deg', 'dual_edge', 'dual_ne', 'dual_vertex', 'dual_bounds_lon', 'dual_bounds_lat', 'left','right','down','up', 'trisk_deg','trisk','Riv2','wee', 'lon_i','lat_i','lon_v','lat_v', 'fv', 'le_de','le','de','lon_e','lat_e','angle_e', 'primal_i','primal_j','edge_i','edge_j') self.Ai, self.Av = Aiv, Aiv self.to_dynamico() def to_dynamico(self): pass def ncwrite(self, name): """The method writes Cartesian mesh on the disc. Args: Mesh parameters""" #----------------OPENING NETCDF OUTPUT FILE------------ try: f = cdf.Dataset(name, 'w', format='NETCDF4') except: ERROR("Cartesian_mesh.ncwrite : Error occurred while opening new netCDF mesh file") raise #----------------DEFINING DIMENSIONS-------------------- for dimname, dimsize in [("primal_cell", self.primal_deg.size), ("dual_cell", self.dual_deg.size), ("edge", self.left.size), ("primal_edge_or_vertex", self.primal_edge.shape[1]), ("dual_edge_or_vertex", self.dual_edge.shape[1]), ("TWO", 2), ("trisk_edge", 4)]: f.createDimension(dimname,dimsize) #----------------DEFINING VARIABLES--------------------- f.description = "Cartesian_mesh" f.meshtype, f.nx, f.ny = 'curvilinear', self.nx, self.ny def create_vars(dimname, info): for vname, vtype, vdata in info: var = f.createVariable(vname,vtype,dimname) var[:] = vdata create_vars("primal_cell", [("primal_deg","i4",self.primal_deg), ("primal_i","i4",self.primal_i), ("primal_j","i4",self.primal_j), ("Ai","f8",self.Ai), ("lon_i","f8",self.lon_i), ("lat_i","f8",self.lat_i)] ) create_vars("dual_cell", [("dual_deg","i4",self.dual_deg), ("Av","f8",self.Av), ("lon_v","f8",self.lon_v), ("lat_v","f8",self.lat_v), ] ) create_vars( ("primal_cell","primal_edge_or_vertex"), [("primal_edge", "i4", self.primal_edge), ("primal_ne", "i4", self.primal_ne), ("primal_vertex", "i4", self.primal_vertex), ("primal_bounds_lon", "f8", self.primal_bounds_lon), ("primal_bounds_lat", "f8", self.primal_bounds_lat)] ) create_vars( ("dual_cell","dual_edge_or_vertex"), [("dual_edge", "i4", self.dual_edge), ("dual_vertex","i4",self.dual_vertex), ("dual_bounds_lon", "f8", self.dual_bounds_lon), ("dual_bounds_lat", "f8", self.dual_bounds_lat), ("dual_ne","i4",self.dual_ne), ("Riv2","f8",self.Riv2)] ) create_vars("edge", [("trisk_deg","i4",self.trisk_deg), ("le","f8",self.le), ("de","f8",self.de), ("lon_e","f8", self.lon_e), ("lat_e","f8", self.lat_e), ("angle_e","f8", self.angle_e), ("edge_i","i4",self.edge_i), ("edge_j","i4",self.edge_j)] ) create_vars( ("edge","TWO"), [("edge_left_right","i4", concat(self.left,self.right)), ("edge_down_up","i4", concat(self.down,self.up))] ) create_vars( ("edge","trisk_edge"), [("trisk","i4",self.trisk), ("wee","f8",self.wee)] ) f.close() def field_theta(self,n=1): return self.zeros((n,self.nqdyn,self.ny,self.nx,self.llm)) def field_mass(self,n=1): return self.zeros((n,self.ny,self.nx,self.llm)) def field_z(self,n=1): return self.zeros((n,self.ny,self.nx,self.llm)) def field_w(self,n=1): return self.zeros((n,self.ny,self.nx,self.llm+1)) def field_u(self,n=1): return self.zeros((n,self.ny,2*self.nx,self.llm)) def field_ps(self,n=1): return self.zeros((n,self.ny,self.nx)) def field_ucomp(self,n=1): return self.zeros((n,self.ny,self.nx,self.llm)) 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),:] if self.ny==1 else u[:,range(1,2*self.nx,2),:] def set_vcomp(self,uv,v): if self.ny==1 : uv[range(1,2*self.nx,2),:]=v else : uv[:,range(1,2*self.nx,2),:]=v #---------------------- DYNAMICO format for fully unstructured mesh and curvilinear meshes ---------------------- class Mesh_Format(BaseMesh): 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 : DEBUG("getvar %s ..."%name) time1=time.time() ret=[self.getvar(name) for name in names] DEBUG("... Done in %f seconds"%(time.time()-time1)) return ret class DYNAMICO_Format(Mesh_Format): """ Helps class Unstructured_Mesh to read a DYNAMICO mesh file.""" start_index=1 # indexing starts at 1 as in Fortran def __init__(self,gridfilename): nc = cdf.Dataset(gridfilename, "r") self.nc, self.meshtype = nc, nc.meshtype if self.meshtype == 'curvilinear': self.nx, self.ny = nc.nx, nc.ny def get_pdim(self,comm,name): return parallel.PDim(self.nc.dimensions[name], comm) def getvar(self,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 == 'edge_deg': return parallel.CstPArray1D(pdim, np.int32, 2) # general case return parallel.PArray(pdim, self.nc.variables[name]) def normalize(self, mesh): pass #---------------------- MPAS fully unstructured mesh ------------------------ class MPAS_Format(Mesh_Format): """ Helps class Unstructured_Mesh to read a MPAS mesh file.""" start_index=1 # indexing starts at 1 as in Fortran meshtype='unstructured' translate= { 'primal_cell':'nCells', 'edge':'nEdges', 'dual_cell':'nVertices', 'primal_deg':'nEdgesOnCell', 'trisk_deg':'nEdgesOnEdge', 'primal_edge':'edgesOnCell', 'dual_edge':'edgesOnVertex', 'primal_vertex':'verticesOnCell', 'dual_vertex':'cellsOnVertex', 'trisk':'edgesOnEdge', 'edge_down_up':'verticesOnEdge', 'edge_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', 'primal_ne':'signOnCell', 'dual_ne':'signOnVertex', 'primal_bounds_lon':None, 'primal_bounds_lat':None, 'dual_bounds_lon':None, 'dual_bounds_lat':None } def __init__(self,gridfilename): try: self.nc = cdf.Dataset(gridfilename, "r") except RuntimeError: ERROR(""" 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 None if name is None else parallel.PArray(pdim, self.nc.variables[name]) def normalize(self, mesh): (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 = de.size, Av.size # 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 : WARNING('error with Riv2 at vertex %d'%ivertex) # scale lengths and areas # this is now done by apply_map # mesh.le, mesh.de, mesh.Av, mesh.Ai = radius*le, radius*de, r2*Av, r2*Ai def get_bounds(self, mesh): def bounds(deg, vertex, lon, lat): num, maxdeg = deg.size, deg.max() vertex_lon, vertex_lat = np.zeros((num, maxdeg)), np.zeros((num, maxdeg)) for cell in range(num): cdeg=deg[cell] for ivertex in range(maxdeg): if ivertex edges primal_edge, left_right = getvars('primal_edge','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','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', 'primal_ne', 'lat_i', 'lon_i', 'lat_v', 'lon_v', 'lat_e', 'lon_e', 'angle_e', 'Ai', 'Av', 'fv', 'le', 'de', 'down', 'up', 'le_de', 'trisk_deg', 'trisk', 'wee', 'Riv2', 'dual_ne', 'left', 'right', 'primal_edge') gridfile.normalize(self) self.to_dynamico() self.primal = Triangulation(lon_i, lat_i) self.dual = Triangulation(lon_v, lat_v) self.triedge = Triangulation(lon_e, lat_e) 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_cell','edge','dual_cell') primal_deg, primal_vertex, primal_edge, primal_ne, lon_i, lat_i, Ai = get_pvars( dim_primal, 'primal_deg', 'primal_vertex', 'primal_edge', 'primal_ne', 'lon_i', 'lat_i', 'Ai' ) primal_bounds_lon, primal_bounds_lat = get_pvars(dim_primal, 'primal_bounds_lon', 'primal_bounds_lat') 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', 'edge_left_right', 'edge_down_up', 'trisk', 'le', 'de', 'wee', 'lon_e', 'lat_e', 'angle_e') dual_deg, dual_vertex, dual_edge, dual_ne, Riv2, lon_v, lat_v, Av = get_pvars( dim_dual, 'dual_deg', 'dual_vertex', 'dual_edge', 'dual_ne', 'Riv2', 'lon_v', 'lat_v', 'Av') dual_bounds_lon, dual_bounds_lat = get_pvars(dim_dual, 'dual_bounds_lon', 'dual_bounds_lat') if gridfile.meshtype == 'curvilinear': self.primal_i, self.primal_j = get_pvars(dim_primal, 'primal_i','primal_j') self.dual_i, self.dual_j = get_pvars(dim_dual, 'primal_i','primal_j') self.edge_i, self.edge_j = get_pvars(dim_edge, 'edge_i','edge_j') # 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, primal_ne) self.cell2vertex = Stencil_glob(primal_deg, primal_vertex) self.edge2vertex = Stencil_glob(edge_deg, down_up) self.vertex2edge = Stencil_glob(dual_deg, dual_edge, dual_ne) self.vertex2cell = Stencil_glob(dual_deg, dual_vertex, Riv2) self.edge2edge = Stencil_glob(trisk_deg, trisk, wee) self.set_members(locals(), 'dim_primal', 'dim_dual', 'dim_edge', 'primal_deg', 'primal_vertex', 'primal_edge', 'primal_bounds_lon', 'primal_bounds_lat', 'trisk_deg', 'trisk', 'dual_deg', 'dual_edge', 'dual_bounds_lon', 'dual_bounds_lat', 'lon_v', 'lat_v', 'Av', 'lon_i', 'lat_i', 'Ai', 'le', 'de', 'lon_e', 'lat_e', 'angle_e') def partition_metis(self): INFO('Partitioning with ParMetis...') # TODO edge_owner = partition_mesh(self.comm, self.trisk_deg, self.trisk) edge_owner = partition_mesh(self.trisk_deg, self.trisk, self.comm.Get_size()) self.edge_owner = parallel.LocPArray1D(self.dim_edge, edge_owner) primal_owner = partition_from_stencil(self.edge_owner, self.primal_deg, self.primal_edge) self.primal_owner = parallel.LocPArray1D(self.dim_primal, primal_owner) dual_owner = partition_from_stencil(self.edge_owner, self.dual_deg, self.dual_edge) self.dual_owner = parallel.LocPArray1D(self.dim_dual, dual_owner) def partition_curvilinear(self, ni,nj): def owner(dim,i,j): owner_i, owner_j = (ni*i.data)/nx, (nj*j.data)/ny return parallel.LocPArray1D(dim, owner_i + ni*owner_j) nx, ny = self.gridfile.nx, self.gridfile.ny INFO('Partitioning %d x %d cells in %d x %d blocks ...'%(nx,ny,ni,nj)) n = self.comm.Get_size() assert n == ni*nj, 'Mismatch between ni x nj = %d and MPI_SIZE = %d.'%(ni*nj, n) self.edge_owner = owner(self.dim_edge, self.edge_i, self.edge_j) self.primal_owner = owner(self.dim_primal, self.primal_i, self.primal_j) self.dual_owner = self.primal_owner DEBUG('partition_curvilinear %d : %s'%(self.comm.Get_rank(), self.primal_owner.data) ) 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): Halo_Xchange = parallel.Halo_Xchange # overriden in dev.meshes def __init__(self, pmesh, llm, nqdyn, mapping): comm, dim_primal, primal_owner, dim_edge, edge_owner, dim_dual, dual_owner = pmesh.members( 'comm', 'dim_primal', 'primal_owner', 'dim_edge', 'edge_owner', 'dim_dual', 'dual_owner' ) INFO('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) DEBUG('E2,E1,E0 ; C1,C0 : %s' % 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) primal_num, dual_num, edge_num = map(len, (cells_C1, vertices_V1, edges_E2)) self.meshtype = pmesh.gridfile.meshtype if self.meshtype == 'curvilinear' : self.nx, self.ny = pmesh.gridfile.nx, pmesh.gridfile.ny self.primal_i, self.primal_j = pmesh.get(get_all_cells, 'primal_i', 'primal_j') self.dual_i, self.dual_j = pmesh.get(get_all_duals, 'dual_i', 'dual_j') # 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) DEBUG('Edge-vertex degree min/max %d %d'%(down_up.degree.min(), down_up.degree.max()) ) DEBUG('Edge-cell degree min/max %d %d'%(left_right.degree.min(), left_right.degree.max()) ) DEBUG('Primal degree min/max : %d %d'%(primal_edge.degree.min(), primal_edge.degree.max()) ) DEBUG('Dual degree min/max : %d %d'%(dual_edge.degree.min(), dual_edge.degree.max()) ) DEBUG('TRiSK degree min/max : %d %d'%(trisk.degree.min(), trisk.degree.max()) ) DEBUG('Riv2 degree min/max : %d %d'%(Riv2.degree.min(), Riv2.degree.max()) ) primal_deg, primal_edge, primal_ne = primal_edge.degree, primal_edge.neigh_loc, primal_edge.weight dual_deg, dual_edge, dual_ne = dual_edge.degree, dual_edge.neigh_loc, dual_edge.weight 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] # 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 # construct own dual mesh for XIOS output dual_own_glo = find_my_cells(dual_owner) dual_own_glo = np.asarray(dual_own_glo, dtype=np.int32) dual_own_loc = [dict_V1[i] for i in dual_own_glo] # compute_ne(...) and normalize(...) expect indices starting at 1 for ind in (trisk, primal_vertex, dual_vertex, primal_edge, dual_edge) : ind+=1 left, right, down, up = left+1, right+1, down+1, up+1 # these ones are subarrays, we need deep copies # read from mesh file : coordinates, lengths and areas in reference domain lon_i, lat_i, Ai = pmesh.get(get_all_cells, 'lon_i', 'lat_i', 'Ai') lon_v, lat_v, Av = pmesh.get(get_all_duals, 'lon_v', 'lat_v', 'Av') le, de, lon_e, lat_e, angle_e = pmesh.get(get_all_edges, 'le', 'de', 'lon_e', 'lat_e', 'angle_e') le_de = le/de # store what we have read/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', 'dual_own_loc', 'trisk_deg', 'trisk', 'wee', 'dual_deg', 'dual_edge', 'dual_ne', 'dual_vertex', 'Riv2', 'left', 'right', 'down', 'up', 'primal_deg', 'primal_edge', 'primal_ne', 'vertices_V1', 'edges_E2', 'lon_i', 'lat_i', 'Ai', 'lon_v', 'lat_v', 'Av', 'le', 'de', 'le_de', 'lon_e', 'lat_e', 'angle_e') # obtain bounds of primal and dual cells (for XIOS) if pmesh.primal_bounds_lon is not None: INFO('Reading cell bounds from mesh file.') self.primal_bounds_lon, self.primal_bounds_lat = pmesh.get(get_all_cells, 'primal_bounds_lon', 'primal_bounds_lat') self.dual_bounds_lon, self.dual_bounds_lat = pmesh.get(get_all_duals, 'dual_bounds_lon', 'dual_bounds_lat') else: INFO('Mesh file does not contain information abound cell bounds.') pmesh.gridfile.get_bounds(self) # normalize data read from file depending on file format pmesh.gridfile.normalize(self) # map from reference domain to physical domain self.apply_map(mapping) # now latitudes and longitudes refer to the physical domain self.fv = mapping.coriolis(self.lon_v,self.lat_v) # Coriolis parameter # propagate info to Fortran side self.to_dynamico() # setup halo transfers - NB : llm must be set before we call set_dynamico_transfer self.com_primal = self.Halo_Xchange( 42, dim_primal, cells_C1, get_all_cells(primal_owner)) self.com_edges = self.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 = Triangulation(lon_i, self.lat_i) self.dual = Triangulation(lon_v, self.lat_v) self.triedge = Triangulation(lon_e, self.lat_e) 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 #-------------------------------------- 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(object): 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(object): 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))