source: codes/icosagcm/devel/Python/test/py/write_Cartesian_mesh.py @ 750

Last change on this file since 750 was 750, checked in by jisesh, 6 years ago

commit message
File size: 10.8 KB
Line 
1from dynamico.meshes import zeros
2import numpy as np
3import netCDF4 as cdf
4import argparse
5
6#----------------------- Cartesian mesh -----------------------
7
8def concat(x,y):
9    z = np.asarray([x,y])
10    return z.transpose()
11
12# arrays is a list of arrays
13# vals is a list of tuples
14# each tuple is stored in each array
15def put(ij, deg, arrays, vals):
16    k=0
17    for vv in vals: # vv is a tuple of values to be stored in arrays
18        for array,v in zip(arrays,vv):
19            array[ij,k]=v
20        k=k+1
21    deg[ij]=k
22
23class Cartesian_mesh:
24    def __init__(self,nx,ny,llm,nqdyn,Lx,Ly):
25        dx,dy = Lx/nx, Ly/ny
26        self.dx, self.dy = dx,dy
27        self.nx, self.ny, self.llm, self.nqdyn = nx,ny,llm,nqdyn
28        self.field_z = self.field_mass
29        # 1D coordinate arrays
30        x=(np.arange(nx)-nx/2.)*Lx/nx
31        y=(np.arange(ny)-ny/2.)*Ly/ny
32        lev=np.arange(llm)
33        levp1=np.arange(llm+1)
34        self.x, self.y, self.lev, self.levp1 = x,y,lev,levp1
35        # 3D coordinate arrays
36        self.xx,self.yy,self.ll = np.meshgrid(x,y,lev, indexing='ij')
37        self.xxp1,self.yyp1,self.llp1 = np.meshgrid(x,y,levp1, indexing='ij')
38        # beware conventions for indexing
39        # Fortran order : llm,nx*ny,nqdyn  / indices start at 1
40        # Python order : nqdyn,ny,nx,llm   / indices start at 0
41        # indices below follow Fortran while x,y follow Python/C
42        index=lambda x,y : ((x+(nx*(y+2*ny)))%(nx*ny))+1
43        indexu=lambda x,y : 2*index(x,y)-1
44        indexv=lambda x,y : 2*index(x,y)
45        indices = lambda shape : np.zeros(shape,np.int32)
46
47        primal_deg = indices(nx*ny)
48        primal_edge = indices((nx*ny,4))
49        primal_ne = indices((nx*ny,4))
50
51        dual_deg = indices(nx*ny)
52        dual_edge = indices((nx*ny,4))
53        dual_ne = indices((nx*ny,4))
54        primal_vertex = indices((nx*ny,4))
55        dual_vertex = indices((nx*ny,4))
56        Riv2 = np.zeros((nx*ny,4))
57
58        left = indices(2*nx*ny)
59        right = indices(2*nx*ny)
60        up = indices(2*nx*ny)
61        down = indices(2*nx*ny)
62        le_de = np.zeros(2*nx*ny)
63        le = np.zeros(2*nx*ny)
64        de = np.zeros(2*nx*ny)
65        angle_e = np.zeros(2*nx*ny)
66
67        trisk_deg = indices(2*nx*ny)
68        trisk = indices((2*nx*ny,4)) # 4 TRiSK coefs per edge
69        wee = np.zeros((2*nx*ny,4))
70
71        lon_i = indices(nx*ny)
72        lon_v = indices(nx*ny)
73        lon_e = indices(2*nx*ny)
74        lat_i = indices(nx*ny)
75        lat_v = indices(nx*ny)
76        lat_e = indices(2*nx*ny)
77
78        for x in range(nx):
79            for y in range(ny):
80                # NB : Fortran indices start at 1
81                # primal cells
82                put(index(x,y)-1, primal_deg,(primal_edge,primal_vertex,primal_ne),
83                    ((indexu(x,y),index(x,y),1), 
84                    (indexv(x,y),index(x-1,y),1),
85                    (indexu(x-1,y),index(x-1,y-1),-1),
86                    (indexv(x,y-1),index(x,y-1),-1) ))
87                # dual cells
88                put(index(x,y)-1,dual_deg,(dual_edge,dual_vertex,dual_ne,Riv2),
89                   ((indexv(x+1,y),index(x,y),1,.25),
90                    (indexu(x,y+1),index(x+1,y),-1,.25),
91                    (indexv(x,y),index(x+1,y+1),-1,.25),
92                    (indexu(x,y),index(x,y+1),1,.25) ))
93                # edges :
94                # left and right are adjacent primal cells
95                # flux is positive when going from left to right
96                # up and down are adjacent dual cells (vertices)
97                # circulation is positive when going from down to up
98                # u-points
99                ij =indexu(x,y)-1
100                left[ij]=index(x,y)
101                left[ij]=index(x,y)
102                right[ij]=index(x+1,y)
103                down[ij]=index(x,y-1)
104                up[ij]=index(x,y)
105                #le_de[ij]=dy/dx
106                le[ij]=dy
107                de[ij]=dx
108                le_de[ij]=le[ij]/de[ij]
109                angle_e[ij]=0.
110                put(ij,trisk_deg,(trisk,wee),(
111                    (indexv(x,y),.25),
112                    (indexv(x+1,y),.25),
113                    (indexv(x,y-1),.25),
114                    (indexv(x+1,y-1),.25)))
115                # v-points
116                ij = indexv(x,y)-1
117                left[ij]=index(x,y)
118                right[ij]=index(x,y+1)
119                down[ij]=index(x,y)
120                up[ij]=index(x-1,y)
121                #le_de[ij]=dx/dy
122                le[ij]=dx
123                de[ij]=dy
124                le_de[ij]=le[ij]/de[ij]
125                angle_e[ij]=.5*np.pi
126                put(ij,trisk_deg,(trisk,wee),(
127                    (indexu(x,y),-.25),
128                    (indexu(x-1,y),-.25),
129                    (indexu(x,y+1),-.25),
130                    (indexu(x-1,y+1),-.25)))
131                ij = index(x,y)-1
132                lon_i[ij], lat_i[ij] = x, y
133                ij = index(x,y)-1
134                lon_v[ij], lat_v[ij] = x+.5, y+.5
135                ij = indexu(x,y)-1
136                lon_e[ij], lat_e[ij] = x+.5, y
137                ij = indexv(x,y)-1
138                lon_e[ij], lat_e[ij] = x, y+.5
139
140        Aiv=np.zeros(nx*ny)+dx*dy 
141        Ai=Av=np.zeros(nx*ny)+dx*dy
142
143        self.llm=llm
144        self.nqdyn=nqdyn
145        self.nx=nx
146        self.ny=ny
147        self.primal_deg=primal_deg
148        self.primal_edge=primal_edge
149        self.primal_ne=primal_ne
150        self.dual_deg=dual_deg
151        self.dual_edge=dual_edge
152        self.dual_ne=dual_ne
153        self.dual_vertex=dual_vertex
154        self.primal_vertex=primal_vertex
155        self.left=left
156        self.right=right
157        self.down=down
158        self.up=up
159        self.trisk_deg=trisk_deg
160        self.trisk=trisk
161        self.Aiv=Aiv
162        self.Ai=Ai
163        self.Av=Av
164        self.le=le
165        self.de=de
166        self.angle_e=angle_e
167        self.Riv2=Riv2
168        self.wee=wee
169        self.lon_i = lon_i
170        self.lon_v = lon_v
171        self.lon_e = lon_e
172        #self.lon_ev = indices(2*nx*ny)
173        self.lat_i = lat_i
174        self.lat_v = lat_v
175        self.lat_e = lat_e
176        #self.lat_ev = indices(2*nx*ny)
177     
178
179    def field_theta(self,n=1): return zeros((n,self.nqdyn,self.ny,self.nx,self.llm))
180    def field_mass(self,n=1): return zeros((n,self.ny,self.nx,self.llm))
181    def field_z(self,n=1): return zeros((n,self.ny,self.nx,self.llm))
182    def field_w(self,n=1): return zeros((n,self.ny,self.nx,self.llm+1))
183    def field_u(self,n=1):
184        if n==1:
185            return np.zeros((self.ny,2*self.nx,self.llm))
186        else:
187            return np.zeros((n,self.ny,2*self.nx,self.llm))
188    def field_ps(self,n=1): return zeros((n,self.ny,self.nx))
189    def ucomp(self,u): return u[:,range(0,2*self.nx,2),:]
190    def set_ucomp(self,uv,u): uv[:,range(0,2*self.nx,2),:]=u
191    def vcomp(self,u): return u[:,range(1,2*self.nx,2),:]
192
193
194    #----------------------------WRITING MESH----------------------
195    def ncwrite(self, name):
196        """The method writes Cartesian mesh on the disc.
197            Args: Mesh parameters"""
198
199        #----------------OPENING NETCDF OUTPUT FILE------------
200
201        try:
202            f = cdf.Dataset(name, 'w', format='NETCDF4')
203        except:
204            print("CartesianMesh.ncwrite : Error occurred while opening new netCDF mesh file")
205            raise
206
207        #----------------DEFINING DIMENSIONS--------------------
208
209        for dimname, dimsize in [("primal_cell", self.primal_deg.size),
210                                 ("dual_cell", self.dual_deg.size),
211                                 ("edge", self.left.size),
212                                 ("primal_edge_or_vertex", self.primal_edge.shape[1]),
213                                 ("dual_edge_or_vertex", self.dual_edge.shape[1]),
214                                 ("TWO", 2),
215                                 ("trisk_edge", 4)]:
216            f.createDimension(dimname,dimsize)
217
218        #----------------DEFINING VARIABLES---------------------
219       
220        f.description = "Cartesian_mesh"
221        f.nx, f.ny = self.nx, self.ny
222
223        def create_vars(dimname, info):
224            for vname, vtype, vdata in info:
225#                print('create_vars', dimname, vname, vdata.shape)
226                var = f.createVariable(vname,vtype,dimname)
227                var[:] = vdata
228               
229        create_vars("primal_cell", 
230                    [("primal_deg","i4",self.primal_deg), 
231                     ("Ai","f8",self.Aiv),
232                     ("lon_i","f8",self.lon_i),
233                     ("lat_i","f8",self.lat_i)] )
234        create_vars("dual_cell", 
235                    [("dual_deg","i4",self.dual_deg), 
236                     ("Av","f8",self.Aiv),
237                     ("lon_v","f8",self.lon_v),
238                     ("lat_v","f8",self.lat_v),
239                     ] )
240
241        create_vars( ("primal_cell","primal_edge_or_vertex"), 
242                     [("primal_edge", "i4", self.primal_edge),
243                      ("primal_ne", "i4", self.primal_ne),
244                      ("primal_vertex", "i4", self.primal_vertex)] )
245
246        create_vars( ("dual_cell","dual_edge_or_vertex"), 
247                     [("dual_edge", "i4", self.dual_edge),
248                      ("dual_vertex","i4",self.dual_vertex),
249                      ("dual_ne","i4",self.dual_ne),
250                      ("Riv2","f8",self.Riv2)] )
251
252        create_vars("edge",
253                    [("trisk_deg","i4",self.trisk_deg),
254                     ("le","f8",self.le),
255                     ("de","f8",self.de),
256                     ("lon_e","f8", self.lon_e),
257                     ("lat_e","f8", self.lat_e),
258                      ("angle_e","f8", self.angle_e)] )
259
260        create_vars( ("edge","TWO"),
261                     [("edge_left_right","i4", concat(self.left,self.right)),
262                      ("edge_down_up","i4", concat(self.down,self.up))] )
263                     
264
265        create_vars( ("edge","trisk_edge"),
266                    [("trisk","i4",self.trisk),
267                     ("wee","f8",self.wee)] )
268
269        f.close()
270
271
272#--------------------------------- Main program -------------------------------------
273
274parser = argparse.ArgumentParser()
275
276parser.add_argument("-nx", type=int,
277                    default=128, choices=None,
278                    help="number of x points")
279parser.add_argument("-ny", type=int,
280                    default=128, choices=None,
281                    help="number of y points")
282parser.add_argument("-Lx", type=float,
283                    default=8., choices=None,
284                    help="Lx")
285parser.add_argument("-Ly", type=float,
286                    default=8., choices=None,
287                    help="Ly")
288parser.add_argument("-llm", type=int,
289                    default=1, choices=[1],
290                    help="number of vertical levels")
291args = parser.parse_args()
292nx, ny, Lx, Ly, llm, nqdyn = args.nx, args.ny,args.Lx, args.Ly, args.llm, 1
293
294dx,dy=Lx/nx,Ly/ny
295mesh = Cartesian_mesh(nx,ny,llm,nqdyn,Lx,Ly)
296print('Successfully initialized Cartesian mesh')
297mesh.ncwrite('cart_'+str(nx)+'.nc')
298print('Successfully written Cartesian mesh to NetCDF File')
Note: See TracBrowser for help on using the repository browser.