[719] | 1 | from dynamico import meshes |
---|
| 2 | from dynamico import unstructured as unst |
---|
| 3 | from dynamico import time_step |
---|
| 4 | |
---|
| 5 | import math as math |
---|
| 6 | import matplotlib.pyplot as plt |
---|
| 7 | import numpy as np |
---|
| 8 | import netCDF4 as cdf |
---|
| 9 | |
---|
| 10 | class DYNAMICO_Format(meshes.Mesh_Format): |
---|
| 11 | """ Netcdf input file handling. |
---|
| 12 | Creates a dictionary of names corresponding to original data file.""" |
---|
| 13 | |
---|
| 14 | start_index=1 # indexing starts at 1 as in Fortran |
---|
| 15 | translate = { 'down_up':'edge_down_up', 'left_right':'edge_left_right' } #new name is edge_left_right |
---|
| 16 | |
---|
| 17 | def __init__(self,gridfilename): |
---|
| 18 | try: |
---|
| 19 | self.nc = cdf.Dataset(gridfilename, "r") |
---|
| 20 | print('#NC4: Opening file:',gridfilename) |
---|
| 21 | except RuntimeError: |
---|
| 22 | print """ Unable to open grid file %s, maybe you forgot to download it ? |
---|
| 23 | To do so, go to the 'Python/' dir and execute './get_MPAS_grids.sh'. |
---|
| 24 | """ % gridfile |
---|
| 25 | raise |
---|
| 26 | def get_pdim(self,comm,name): |
---|
| 27 | if name in self.translate: name = self.translate[name] |
---|
| 28 | return parallel.PDim(self.nc.dimensions[name], comm) |
---|
| 29 | def getvar(self,name): |
---|
| 30 | if name in self.translate: name = self.translate[name] |
---|
| 31 | return self.nc.variables[name][:] |
---|
| 32 | def get_pvar(self,pdim,name): |
---|
| 33 | # provides parallel access to a NetCDF variable, with name translation |
---|
| 34 | # first deal with special cases |
---|
| 35 | if name == 'edge_deg': return parallel.CstPArray1D(pdim, np.int32, 2) |
---|
| 36 | # general case |
---|
| 37 | if name in self.translate: name = self.translate[name] |
---|
| 38 | return parallel.PArray(pdim, self.nc.variables[name]) |
---|
| 39 | def normalize(self, mesh, radius): pass |
---|
| 40 | |
---|
| 41 | #--------------------------------- Main program ----------------------------- |
---|
| 42 | |
---|
| 43 | nx,ny,llm,nqdyn=128,128,1,1 |
---|
| 44 | Lx,Ly,g,f = 8.,8.,1.,1. |
---|
| 45 | dx,dy=Lx/nx,Ly/ny |
---|
| 46 | |
---|
[747] | 47 | filename, llm, nqdyn, g, f, radius = 'cart128.nc', 1, 1, 1., 1., None |
---|
[719] | 48 | unst.setvar('g',g) |
---|
| 49 | |
---|
| 50 | print 'Reading Cartesian mesh ...' |
---|
| 51 | meshfile = DYNAMICO_Format(filename) |
---|
| 52 | #print("#NC4: check point: 0") |
---|
| 53 | #print(meshfile) |
---|
| 54 | |
---|
| 55 | def coriolis(lon,lat): |
---|
| 56 | return f+0.*lon |
---|
| 57 | |
---|
| 58 | mesh=meshes.Unstructured_Mesh(meshfile, llm, nqdyn, radius, coriolis) |
---|
| 59 | caldyn = unst.Caldyn_RSW(mesh) |
---|
| 60 | #print(caldyn.__dict__.keys()) |
---|
| 61 | |
---|
| 62 | xx = (mesh.lat_i-nx/2.)*dx |
---|
| 63 | yy = (mesh.lon_i-ny/2.)*dy |
---|
| 64 | |
---|
| 65 | x1,x2,yy = xx-1., xx+1., yy |
---|
[747] | 66 | u0 = mesh.field_u() # flow initally at rest |
---|
[719] | 67 | h0 = 1+0.1*(np.exp(-2.*(x1*x1+yy*yy))+ |
---|
| 68 | np.exp(-2.*(x2*x2+yy*yy))) |
---|
[747] | 69 | flow0=meshes.asnum([h0,u0]) |
---|
[719] | 70 | |
---|
| 71 | #print('h0 : ', h0[1234]-1.) |
---|
| 72 | #print('h0 : ', h0[2345]-1.) |
---|
| 73 | |
---|
| 74 | cfl = .8 |
---|
| 75 | dt = cfl/math.sqrt((nx/Lx)**2+(ny/Ly)**2) |
---|
| 76 | |
---|
| 77 | T=1. |
---|
| 78 | N=int(T/dt)+1 |
---|
| 79 | dt=T/N |
---|
| 80 | print N,dt,Lx/nx |
---|
[747] | 81 | #scheme = time_step.ARK2(caldyn.bwd_fast_slow, dt, precision=unst.np_num) |
---|
| 82 | scheme = time_step.RK4(caldyn.bwd_fast_slow, dt, precision=unst.np_num) |
---|
[719] | 83 | |
---|
| 84 | flow=flow0 |
---|
| 85 | |
---|
| 86 | def minmax(name, x): print('Min/max %s :'%name, x.min(), x.max() ) |
---|
| 87 | def reshape(data): return data.reshape((nx,ny)) |
---|
| 88 | |
---|
| 89 | x, y = map(reshape, (xx,yy) ) |
---|
| 90 | |
---|
| 91 | for i in range(10): |
---|
| 92 | h,u=flow |
---|
[747] | 93 | flow, fast, slow = caldyn.bwd_fast_slow(flow, unst.zero_num) |
---|
[719] | 94 | junk, du_fast = fast |
---|
| 95 | dh, du_slow = slow |
---|
| 96 | # minmax('lon',mesh.lon_i) |
---|
| 97 | # minmax('lat',mesh.lat_i) |
---|
| 98 | # minmax('x',xx) |
---|
| 99 | # minmax('y',yy) |
---|
| 100 | minmax('PV', caldyn.qv-1.) |
---|
| 101 | # minmax('geopot', caldyn.geopot) |
---|
| 102 | # minmax('du_fast', du_fast) |
---|
| 103 | plt.figure(); plt.pcolor(x,y,reshape(caldyn.qv)) |
---|
| 104 | plt.colorbar(); plt.title('potential vorticity') |
---|
[747] | 105 | plt.savefig('fig_RSW_2D_mesh/%03d.png'%i) |
---|
| 106 | plt.close() |
---|
[719] | 107 | # plt.figure(); plt.pcolor(mesh.x,mesh.y,h) |
---|
| 108 | # plt.colorbar(); plt.title('h') |
---|
| 109 | # plt.show() |
---|
| 110 | # plt.pcolor(x,y,vcomp(u)/dx) |
---|
| 111 | # plt.colorbar(); plt.title('v') |
---|
| 112 | # plt.show() |
---|
| 113 | for j in range(5): |
---|
| 114 | unst.elapsed=0. |
---|
| 115 | flow = scheme.advance(flow,N) |
---|
[747] | 116 | # flow = scheme.advance(flow,5) |
---|
[719] | 117 | # flops |
---|
| 118 | # mass flux : 1add+1mul per edge => 4 |
---|
| 119 | # div U : 4 add per cell => 4 |
---|
| 120 | # KE : 4*(2add+1mul) per cell => 12 |
---|
| 121 | # grad KE : 1 add per edge => 2 |
---|
| 122 | # grad h : 1 add per edge => 2 |
---|
| 123 | # qv : 4+4+1 add +4mul + 1div per cell => 10 |
---|
| 124 | # qu : 1add+1mul per edge => 4 |
---|
| 125 | # TrisK : 4add+4mul+4add+1add per edge => 26 |
---|
| 126 | # Total : 64 FLOPS/cell |
---|
| 127 | print i,j, unst.elapsed, 100.*(1.-unst.getvar('elapsed')/unst.elapsed) |
---|
| 128 | print 'GFlops', 64*(N*nx*ny)/unst.elapsed/1e9 |
---|
| 129 | unst.setvar('elapsed',0.) |
---|