source: codes/icosagcm/devel/Python/test/py/bubble.py @ 631

import math as math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as manimation

from dynamico.meshes import Cartesian_mesh as Mesh
import dynamico.dyn as dyn
import dynamico.time_step as time_step
import dynamico.DCMIP as DCMIP
import dynamico.advect as advect
from dynamico import unstructured as unst

def reldiff(a,b):
    return (a-b)/(np.absolute(a).max()+np.absolute(b).max()+1e-20)

def go(model,courant,scheme,start):
    def plotter(metric, flow, scalars, t):
        mik, gas_ik, ujk, Phi_il, Wil = model.diagnose(*flow)
        Q1, Q2 = scalars
        
        mil = metric.ml_ext(mik)
        wil = Wil/mil*metric.g0
        wik = metric.mk_ext(Wil)/mik*metric.g0
        xik = metric.xik/1000
        xjl = metric.xjl/1000
        xil = metric.xil/1000
        zil = Phi_il/metric.g0/1000
        zik = metric.mk_int(zil)
        zjl = metric.mj(zil)
        print 's ',np.max(gas_ik.s),np.min(gas_ik.s)
        print 'ujk ',np.max(ujk),np.min(ujk)
        print 'wil ',np.max(wil),np.min(wil)
        sik = gas_ik.s
        sik = .5*(abs(sik)+sik)

        plt.clf()
        plt.contourf(xik, zik, sik, 20 )
#        plt.contourf(xik, zik, wik, 20 )
        plt.xlim((-1,1)), plt.xlabel('x (km)')
        plt.ylim((0,ztop/1000.)), plt.ylabel('z (km)')
        plt.colorbar()
        plt.title("t = %f s" % (t))
        plt.savefig('fig_bubble/bubble.png',bbox_inches='tight')
        writer.grab_frame()
        
    transport = advect.SplitTransport(model.metric, advect.HorizVanLeer, advect.VertVanLeer)
    replace = lambda fd,fv : fv
    return DCMIP.run_with_scalars(model,scheme,transport,replace, plotter, courant, start,scalars, T1, N1, N2)

Lx, ztop, nx, llm = 2000.,1000., 100, 50
nqdyn, ny, dx = 1, 1, Lx/nx
T1, N1, N2 = 5., 1, 200
#T1, N1, N2 = 5., 2, 5

metric, thermo, BC, m0ik, gas0, Phi0_il, u0_jk = DCMIP.thermal_bubble(Lx,nx,llm, ztop=ztop, zc=350.)
Sik = m0ik*gas0.s
start = (m0ik, Sik, u0_jk, Phi0_il, 0*Phi0_il)
scalars = (Sik, m0ik*metric.mk_int(Phi0_il))

model = dyn.Nonhydro(thermo,metric,BC)
solver, caldyn_eta = dyn.MassBasedNH(model), unst.eta_mass
#solver, caldyn_eta = dyn.LagrangianNH(model), unst.eta_lag

rho_bot=BC.rho_bot
if False: # explicit time-stepping with rigid lid
    def rigid(flow,tau): # impose rigid BCs at top and bottom
        flow, fast, slow = solver.bwd_fast_slow(flow,0.)
        fast[4][:,0]=0.
        fast[4][:,-1]=0.
        return flow,fast,slow
    BC.rho_bot = 100*rho_bot # moderate stiffness
    courant, time_scheme = 2.0, lambda dt : time_step.RK4(rigid, dt)
else: # HEVI time-stepping with pressure BC
    # Note : parameters BC.rigid_top and BC.rigid_bottom are True by default, but are ignored by the Python code...
    BC.rigid_top=False
    BC.rigid_bottom=False
    BC.rho_bot = 1e6*rho_bot # large stiffness
#    BC.rho_bot = 100*rho_bot # moderate stiffness
    courant, time_scheme = 3.0, lambda dt : time_step.ARK2(solver.bwd_fast_slow, dt, a32=0.7)

FFMpegWriter = manimation.writers['ffmpeg']
metadata = dict(title='Aquaplanet', artist='DYNAMICO_LMDZ5',
                comment='Movie support!')
writer = FFMpegWriter(fps=15, metadata=metadata, bitrate=4000, codec="h263p")
fig = plt.figure(figsize=(12,5))
with writer.saving(fig, "fig_bubble/bubble.avi", 100):
    m0ik, gas0_ik, u0jk, Phi0_il, W0il = go(model, courant, time_scheme, start)

#-------------------------------------------------------------------------------------------------------
# Now use final state to check that Fortran and Python implementations produce the exact same tendencies
#-------------------------------------------------------------------------------------------------------

#W0il[:,0]=0.
#W0il[:,-1]=0.

unst.setvar('nb_threads', 1)
mesh = Mesh(nx,ny,llm,nqdyn,Lx,ny*dx,0.)
xx_ik, xx_il, ll = mesh.xx[:,0,:]/1000, mesh.xxp1[:,0,:]/1000, mesh.ll[:,0,:]

mass_bl,mass_dak,mass_dbk=unst.compute_hybrid_coefs(m0ik)
unst.init_hybrid(mass_bl,mass_dak,mass_dbk)

caldyn_thermo = unst.thermo_entropy
#caldyn_thermo = unst.thermo_theta
g = mesh.dx/metric.dx_g0
caldyn = unst.Caldyn_NH(caldyn_thermo,caldyn_eta, mesh,metric,thermo,BC,g)
unst.setvar('rigid',False)

s0ik = gas0_ik.s
if caldyn_thermo == unst.thermo_theta:
    s0ik = thermo.T0*np.exp(s0ik/thermo.Cpd)

# 2 momentum components in DYNAMICO vs 1 component in slice
u0jk = u0jk + 100*dx # shift horizontal wind by 100 m/s
u0_ker = mesh.field_u()
mesh.set_ucomp(u0_ker,u0jk)

tau=0.1
# impose hybdrid distribution of mass (Fortran)
flow = (m0ik/dx,s0ik*m0ik/dx,u0_ker,Phi0_il,W0il/dx) # NB : different definitions of mass field in DYNAMICO / Python
flow,fast,slow = caldyn.bwd_fast_slow(flow, 0.)
# shift Phi from target value Phi_bot
Phi0_il = Phi0_il+g*100.*np.exp(-10.*xx_il*xx_il)
m,S,u,Phi,W = flow
flow = m,S,u,Phi0_il,W
# compute tendencies (0=Python)
flow0 = (m*dx,S*dx,u0jk,Phi0_il,W0il, 0., 0.) # NB : different definitions of mass field in DYNAMICO / Python
flow0, fast0, slow0 = solver.bwd_fast_slow(flow0,tau,True) 
# compute tendencies (Fortran)
flow,fast,slow = caldyn.bwd_fast_slow(flow, tau) 


zz_il = Phi0_il/metric.g0/1000.
zz_ik=.5*(zz_il[:,0:-1]+zz_il[:,1:])
print
print 'ptop, model top (km) :', unst.getvar('ptop'), zz_il.max()
        
for data0,data,dname in ((flow0,flow,'flow'),(slow0,slow,'slow'),(fast0,fast,'fast')):
    def plot(xx,zz,var,vname):
        title = "%s_%s" % (vname,dname)
        filename = "bubble/bubble_%s.png" % title
        var = 1.*var
        if var.size>1:
            plt.figure(figsize=(12,5))
            plt.contourf(xx,zz,var,20)
#                    plt.pcolor(xx,zz,var)
            plt.xlim((-1.,1.)), plt.xlabel('x (km)')
            plt.ylim((0.,1.)), plt.ylabel('z (km)')        
            plt.colorbar()
            plt.title(title)
            plt.savefig(filename,bbox_inches='tight')
            print title, var.min(), var.max()
            if var.max()>1e100:
                raise OverflowError
                    
    dm,dS,du,dPhi,dW = data    
    du=mesh.ucomp(du)/dx
    du=du[0,:,:]

    dm0, dS0, du0, dPhi0, dW0, junk, junk = data0
    dm0, dS0, du0, dW0 = dm0/dx, dS0/dx, du0/dx, dW0/dx

    if dname == 'flow':
        hflux=mesh.ucomp(caldyn.hflux)
        plot(xx_ik, zz_ik, dm,'m')
        plot(xx_ik, zz_ik, du,'u')
        plot(xx_ik, zz_ik, du-du0,'uu')
        plot(xx_ik, zz_ik, dS/m,'s')
        plot(xx_il, zz_il, dW,'w')
        plot(xx_il, zz_il, dW-dW0,'ww')
    else:
        dS=reldiff(dS,dS0)
        dW=dW-dW0
        dPhi=dPhi-dPhi0
        du=reldiff(du,du0)
        plot(xx_ik, zz_ik, dS,'dS')
        plot(xx_il, zz_il, dPhi,'dPhi')
        plot(xx_il, zz_il, dW,'dW')
        plot(xx_ik, zz_ik, du,'du')

if unst.getvar('rigid'):
    print 'Fortran BCs : rigid'
else:
    print 'Fortran BCs : pressure/soft', unst.getvars('ptop','rho_bot','Phi_bot','pbot')

if BC.rigid_top:
    print 'Top BC : rigid'
else:
    print 'Top BC : pressure', BC.ptop

if BC.rigid_bottom:
    print 'Bottom BC : rigid'
else:
    print 'Bottom BC : soft', BC.rho_bot[0], BC.Phi_bot[0], BC.pbot[0]