source: codes/icosagcm/devel/Python/test/py/RSW2_MPAS_W02.py @ 825

Last change on this file since 825 was 825, checked in by dubos, 5 years ago

devel/Python : moved Fortran bindings and *.pyx to dynamico/dev module + necessary changes to test/py/*.py

File size: 6.2 KB
Line 
1from dynamico import getargs
2getargs.add("--LAM", action='store_true')
3# Args for global mesh
4getargs.add("--grid", type=int, help='Number of hexagons', default=2562)
5# Args for LAM
6getargs.add("--nx", type=int, help='Zonal dimension of LAM mesh', default=100)
7getargs.add("--ny", type=int, help='Meridional dimension of LAM mesh', default=100)
8getargs.add("--dx", type=float, help='Resolution at center of LAM domain', default=1e5)
9getargs.add("--center_lat", type=float, help='Latitude in degrees of LAM center', default=0.)
10getargs.add("--Davies_N1", type=int, default=3)
11getargs.add("--Davies_N2", type=int, default=3)
12
13args = getargs.parse()
14
15log_master, log_world = getargs.getLogger()
16INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error
17
18INFO('Starting')
19
20from mpi4py import MPI
21comm = MPI.COMM_WORLD
22mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
23INFO('%d/%d starting'%(mpi_rank,mpi_size))
24prefix='fig_RSW2_MPAS_W02/%02d'%mpi_rank
25
26INFO('Loading DYNAMICO modules ...')
27from dynamico.dev import unstructured as unst
28from dynamico.dev.meshes import MPAS_Format, Unstructured_PMesh as PMesh, Local_Mesh as Mesh
29from dynamico.dev import meshes
30from dynamico import time_step
31from dynamico import maps
32from dynamico.LAM import Davies
33
34print '...Done'
35
36print 'Loading modules ...'
37import math as math
38import matplotlib.pyplot as plt
39import numpy as np
40print '...Done'
41
42#--------------------------- functions and classes -----------------------------
43
44class myDavies(Davies):
45    def mask(self,X,Y):
46        # X and Y are coordinates in the reference domain (cell = unit square)
47        # numerical domain extends from -nx/2 ... nx/2 and -ny/2 ... ny/2
48        # useful domain extends from -X0 ... X0 and -Y0 ... Y0
49        N3 = args.Davies_N1+args.Davies_N2
50        X0 = args.nx/2. - N3
51        Y0 = args.ny/2. - N3
52        mask  = self.mask0( X,X0,1.)  # Western boundary
53        mask *= self.mask0(-X,X0,1.) # Eastern boundary
54        mask *= self.mask0( Y,Y0,1.) # Northern boundary
55        mask *= self.mask0(-Y,Y0,1.) # Southern boundary
56        return mask
57
58#----------------------------- main program --------------------------------
59
60llm, nqdyn = 1,1 # 2562, 10242, 40962
61Omega, radius, g, gh0 = 2.*np.pi/86400., 6.4e6, 1., 2.94e4
62N, T, courant = 40, 10400., 1.2 # simulation length = N*T
63
64print 'Omega, planetary PV', Omega, 2*Omega/gh0
65
66if args.LAM:
67    nx, ny = args.nx, args.ny
68    filename = 'cart_%03d_%03d.nc'%(nx,ny)
69    INFO('Reading Cartesian mesh ...')
70    meshfile = meshes.DYNAMICO_Format(filename)
71    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
72    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
73    planet = maps.PolarStereoMap(radius,Omega, args.dx, args.center_lat*np.pi/180.)
74else:
75    planet = maps.SphereMap(radius, Omega)
76    print 'Reading MPAS mesh ...'
77    meshfile = MPAS_Format('grids/x1.%d.grid.nc'%args.grid)
78    pmesh = PMesh(comm,meshfile)
79    pmesh.partition_metis()
80
81mesh = Mesh(pmesh, llm, nqdyn, planet)
82
83print '...Done'
84lon, lat = mesh.lon_i, mesh.lat_i
85x,y,z = np.cos(lat)*np.cos(lon), np.cos(lat)*np.sin(lon), np.sin(lat)
86
87unst.setvar('g',g)
88
89c0 = math.sqrt(gh0) # phase speed of barotropic mode
90dx = mesh.de.min()
91dt = courant*dx/c0
92print dx, dt
93nt = int(math.ceil(T/dt))
94dt = T/nt
95print dx, dt, dt*c0/dx
96print T, nt
97
98scheme = time_step.RK4(None, dt)
99print dt, scheme.csjl, scheme.cfjl
100step = unst.caldyn_step_TRSW(mesh,scheme,1)
101
102u0 = Omega*radius/12. # cf Williamson (1991), p.13
103gh1 = radius*Omega*u0+.5*u0**2
104print 'Williamson (1991) test 2, u0=', u0
105ulon = u0*np.cos(mesh.lat_e)
106
107gh_ini = gh0 - gh1*(np.sin(mesh.lat_i)**2)
108u_ini = mesh.ucov2D(ulon,0.*ulon)
109
110# unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, u.shape[0], u)
111
112if args.LAM:
113    davies = myDavies(args.Davies_N1, args.Davies_N2, 
114                      mesh.ref_lon_i, mesh.ref_lat_i, mesh.ref_lon_e,mesh.ref_lat_e)
115    def relax():
116        davies.relax_RSW(llm, step, (gh_ini, u_ini) )
117else:
118    def relax(): pass
119
120def next_flow(Phi, u):
121    step.mass[:], step.theta_rhodz[:], step.u[:] = Phi, Phi, u
122    for i in range(nt):
123        step.next()
124        relax()
125    return step.mass.copy(), step.u.copy()
126
127gh, u = gh_ini, u_ini
128print gh.shape, gh.min(), gh.max(), u.min(), u.max()
129
130for i in range(20):
131    if True:
132        gh, u = next_flow(gh,u)
133        # step.next()
134        # gh,u = step.mass, step.u
135        print i, gh.shape, gh.min(), gh.max(), u.min(), u.max()
136        mesh.plot_i(gh-gh_ini) ; plt.title('err(gh)');
137        plt.savefig('%s_err_gh_%02d.png'%(prefix,i))
138        plt.close()
139    else:
140        # advance by one time step using dynamico_ARK
141        gh,u = step.mass.copy(), step.u.copy()
142        print i, gh.shape, gh.min(), gh.max(), u.min(), u.max()
143        step.next()
144        gh_now,u_now = step.mass.copy(), step.u.copy()
145        dmass,dhs,du_slow,du_fast = step.drhodz[:], step.dtheta_rhodz[:], step.du_slow[:], step.du_fast[:]
146        du=du_slow+du_fast
147
148        # do the same using caldyn, obtain fast/slow tendencies first
149        gh_ref, uref=flow
150        junk, fast, slow = caldyn.bwd_fast_slow(flow,0.)                                                                     
151        dmass_ref, duslow_ref = slow
152        junk, dufast_ref = fast
153        du_ref = duslow_ref+dufast_ref
154        flow=scheme.next(flow)
155        gh_nowref, u_nowref=flow
156#    mesh.plot_i(gh) ; plt.title('gh');
157#    mesh.plot_i(Phi0) ; plt.title('gh');
158        fig = plt.figure(figsize=(6, 8)) 
159        f, ((ax1, ax2), (ax3, ax4), (ax5,ax6), (ax7,ax8)) = plt.subplots(4,2)
160#        mesh.plot_i(dmass-dhs) ; plt.title('dt*d(gh)/dt');
161        ax1.scatter(duslow_ref,du_slow-duslow_ref) ; ax1.set_title('du_slow');
162        ax2.scatter(dufast_ref,du_fast-dufast_ref) ; ax2.set_title('du_fast');
163        ax3.scatter(du_ref,du-du_ref) ; ax3.set_title('du');
164        ax4.scatter(gh_ref,gh-gh_ref) ; ax4.set_title('gh');
165        ax5.scatter(uref, u-uref) ; ax5.set_title('u');
166        ax6.scatter(dmass_ref,dmass-dmass_ref) ; ax6.set_title('dmass');
167        ax7.scatter(gh_nowref,gh_now-gh_nowref) ; ax7.set_title('gh_now');
168        ax8.scatter(u_nowref,u_now-u_nowref) ; ax8.set_title('u_now');
169        f.subplots_adjust(hspace=1.)
170        plt.savefig('%s_scatter_%02d.png'%(prefix,i))
171        plt.close()
172
173print 'Time spent in DYNAMICO (s) : ', unst.getvar('elapsed')
Note: See TracBrowser for help on using the repository browser.