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

print 'Loading modules ...'
import math as math
import matplotlib.pyplot as plt
import numpy as np
print '...Done'

print 'Loading DYNAMICO modules ...'
from dynamico import unstructured as unst
from dynamico.meshes import MPAS_Mesh as Mesh
from dynamico import time_step
print '...Done'

grid, llm, nqdyn = 10242, 1,1 # 2562, 10242, 40962
Omega, radius, g, gh0 = 2.*np.pi/86400., 6.4e6, 1., 2.94e4
N, T, courant = 40, 10800., 1.2 # simulation length = N*T

print 'Omega, planetary PV', Omega, 2*Omega/gh0

def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter
print 'Reading MPAS mesh ...'
mesh = Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f)
print '...Done'
lon, lat = mesh.lon_i, mesh.lat_i
x,y,z = np.cos(lat)*np.cos(lon), np.cos(lat)*np.sin(lon), np.sin(lat)

unst.setvar('g',g)

c0 = math.sqrt(gh0) # phase speed of barotropic mode
dx = mesh.de.min()
dt = courant*dx/c0
nt = int(math.ceil(T/dt))
dt = T/nt
print dx, dt, dt*c0/dx
print T, nt

caldyn = unst.Caldyn_RSW(mesh)
scheme = time_step.RK4(caldyn.bwd_fast_slow, dt)
#scheme = time_step.RKn_simple(1,caldyn.bwd_fast_slow, dt) # forward Euler scheme
print dt, scheme.csjl, scheme.cfjl
step = unst.caldyn_step_TRSW(mesh,scheme,nt)

u0 = Omega*radius/12. # cf Williamson (1991), p.13
gh1 = radius*Omega*u0+.5*u0**2
print 'Williamson (1991) test 2, u0=', u0
ulon = u0*np.cos(mesh.lat_e)
Phi0 = gh0 - gh1*(np.sin(mesh.lat_i)**2)
zeta0 = (2*u0/radius+2*Omega)*np.sin(mesh.lat_v)
Phi0v = gh0 - (radius*Omega*u0+.5*u0**2)*(np.sin(mesh.lat_v)**2)
q0 = zeta0/Phi0v
fu_perp = mesh.ucov2D(0.*ulon,2*Omega*np.sin(mesh.lat_e)*ulon)

# initial condition
step.mass[:]=Phi0
step.theta_rhodz[:]=Phi0
step.u[:]=mesh.ucov2D(ulon,0.*ulon)

gh,u = step.mass.copy(), step.u.copy()

print gh.shape, gh.min(), gh.max(), u.min(), u.max()
flow=gh,u

for i in range(20):
    if True:
        step.next()
        gh,u = step.mass, step.u
        print i, gh.shape, gh.min(), gh.max(), u.min(), u.max()
        mesh.plot_i(gh-Phi0) ; plt.title('err(gh)');
        plt.savefig('fig_RSW2_MPAS_W02/err_gh_%02d.png'%i)
        plt.close()
    else:
        # advance by one time step using dynamico_ARK
        gh,u = step.mass.copy(), step.u.copy()
        print i, gh.shape, gh.min(), gh.max(), u.min(), u.max()
        step.next()
        gh_now,u_now = step.mass.copy(), step.u.copy()
        dmass,dhs,du_slow,du_fast = step.drhodz[:], step.dtheta_rhodz[:], step.du_slow[:], step.du_fast[:]
        du=du_slow+du_fast

        # do the same using caldyn, obtain fast/slow tendencies first
        gh_ref, uref=flow
        junk, fast, slow = caldyn.bwd_fast_slow(flow,0.)                                                                      
        dmass_ref, duslow_ref = slow
        junk, dufast_ref = fast
        du_ref = duslow_ref+dufast_ref
        flow=scheme.next(flow)
        gh_nowref, u_nowref=flow
#    mesh.plot_i(gh) ; plt.title('gh');
#    mesh.plot_i(Phi0) ; plt.title('gh');
        fig = plt.figure(figsize=(6, 8)) 
        f, ((ax1, ax2), (ax3, ax4), (ax5,ax6), (ax7,ax8)) = plt.subplots(4,2)
#        mesh.plot_i(dmass-dhs) ; plt.title('dt*d(gh)/dt');
        ax1.scatter(duslow_ref,du_slow-duslow_ref) ; ax1.set_title('du_slow');
        ax2.scatter(dufast_ref,du_fast-dufast_ref) ; ax2.set_title('du_fast');
        ax3.scatter(du_ref,du-du_ref) ; ax3.set_title('du');
        ax4.scatter(gh_ref,gh-gh_ref) ; ax4.set_title('gh');
        ax5.scatter(uref, u-uref) ; ax5.set_title('u');
        ax6.scatter(dmass_ref,dmass-dmass_ref) ; ax6.set_title('dmass');
        ax7.scatter(gh_nowref,gh_now-gh_nowref) ; ax7.set_title('gh_now');
        ax8.scatter(u_nowref,u_now-u_nowref) ; ax8.set_title('u_now');
        f.subplots_adjust(hspace=1.)
        plt.savefig('fig_RSW2_MPAS_W02/scatter_%02d.png'%i)
        plt.close()

print 'Time spent in DYNAMICO (s) : ', unst.getvar('elapsed')