source: codes/icosagcm/devel/Python/dynamico/meshes.py @ 829

from __future__ import absolute_import     

from dynamico.dev import unstructured as unst

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
# TODO from dynamico.parmetis import partition_mesh
from dynamico.dev.unstructured 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

#------------------- 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_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(index(x,y)-1, primal_deg,(primal_edge,primal_vertex,primal_ne),
                    ((indexu(x,y),index(x,y),1), 
                    (indexv(x,y),index(x-1,y),1),
                    (indexu(x-1,y),index(x-1,y-1),-1),
                    (indexv(x,y-1),index(x,y-1),-1) ))
                # 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),
                   ((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 
                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',
                         'dual_deg', 'dual_edge', 'dual_ne', 'dual_vertex',
                         '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)] )

        create_vars( ("dual_cell","dual_edge_or_vertex"), 
                     [("dual_edge", "i4", self.dual_edge),
                      ("dual_vertex","i4",self.dual_vertex),
                      ("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'}
    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 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 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(BaseMesh):
    def to_dynamico(self): pass
    def field_theta(self,n=1): return self.zeros((n,self.nqdyn,self.primal_num,self.llm))
    def field_mass(self,n=1):  return self.zeros((n,self.primal_num,self.llm))
    def field_z(self,n=1):     return self.zeros((n,self.dual_num,self.llm))
    def field_w(self,n=1):     return self.zeros((n,self.primal_num,self.llm+1))
    def field_u(self,n=1):     return self.zeros((n,self.edge_num,self.llm))
    def field_ps(self,n=1):    return self.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 = self.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)
    def apply_map(self, mapping):
        """Before calling apply_map, lat and lon are coordinates in the reference domain.
        After calling apply_map, lat and lon are coordinates in the physical domain."""
        lat_i, lat_e, lat_v = self.lat_i, self.lat_e, self.lat_v
        lon_i, lon_e, lon_v = self.lon_i, self.lon_e, self.lon_v
        factor_e = mapping.map_factor(lon_e, lat_e)
        factor2_i = mapping.map_factor(lon_i, lat_i)**2
        factor2_v = mapping.map_factor(lon_v, lat_v)**2
        self.le, self.de, self.Av, self.Ai = self.le*factor_e, self.de*factor_e, self.Av*factor2_v, self.Ai*factor2_i
        self.angle_e += mapping.map_angle(lon_e, lat_e)
        self.lon_i, self.lat_i = mapping.map(lon_i, lat_i)
        self.lon_v, self.lat_v = mapping.map(lon_v, lat_v)
        self.lon_e, self.lat_e = mapping.map(lon_e, lat_e)
        self.ref_lon_i, self.ref_lat_i = lon_i, lat_i
        self.ref_lon_e, self.ref_lat_e = lon_e, lat_e
        self.ref_lon_v, self.ref_lat_v = lon_v, lat_v

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')
        le_de = le/de
        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','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  = 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_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' )
        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')
        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', 
                         'trisk_deg', 'trisk', 'dual_deg', 'dual_edge',
                         '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...')
        edge_owner = partition_mesh(self.comm, self.trisk_deg, self.trisk)
#        edge_owner = unst.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
        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

        # 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')

        # 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  = 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(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))