deep_atmosphere.py

import numpy as np
import math
import time, os, sys
import gmshPartition
from initialCondition import *

#### Begin and end time ###
Ti=0.0
Tf=30*24*3600 / X
###########################


hasCircleTransform = True

checkPointInterv=5000
nbStart=0
reloadSuffix='0'

dimension=3
order=3
if (deepAtm):
  exportBaseDir = '/SCRATCH/ucl-mema/sblaise/'
else:
  exportBaseDir = '/SCRATCH/ucl-mema/sblaise/'

if (hasCircleTransform == False):
  exportBaseDir = exportBaseDir + '_noCircle'

exportBaseDir = './' 

expDirName = 'deepAtm'
if (hasCircleTransform == False):
  expDirName = expDirName + '_noCircle'
if (deepAtm == False):
  expDirName = expDirName + '_shallow'
loadDirName = expDirName + reloadSuffix
expDirName = expDirName + str(nbStart)

integOrder=2*order+1+2
dataCacheMap.setDefaultIntegrationOrder(integOrder)

model = GModel()

name='aquaplanet'
if Msg.GetCommSize()>1:
  partStr='_part_%i' % Msg.GetCommSize()
else:
  partStr=''
model.load(name + partStr + '_3d.msh')

groups = dgGroupCollection(model, dimension, order)
groups.splitGroupsByPhysicalTag();
if (hasCircleTransform):
  circleTransform = dgCircleTransform(groups)
  st = dgSpaceTransformSpherical(R, groups)
else:
  st = dgSpaceTransformSpherical(R)


groups._mesh.setSpaceTransform(st, groups)
groupsH = dgGroupCollection.newByTag(model, dimension-1, order, ["bottom_Top","bottom_Bottom"])
#To compute the surface pressure. Du to the ocean-base extrusion, top is actually bottom
extrusion = dgExtrusion(groups, groupsH, ["bottom_Top","bottom_Bottom"])
groups.splitFaceGroupsByOrientation(extrusion) 
#extrusion = dgExtrusion(groups, groupsH, ["bottom_Top","bottom_Bottom"])

XYZ = groups.getFunctionCoordinates();
if (hasCircleTransform):
  xyzCircle = transformToCircle(circleTransform)
else:
  xyzCircle = XYZ

xyzCircleInteg = functionPrecomputed(groups, integOrder, 3)
xyzCircleInteg.compute(xyzCircle)

#def fCoriolis(cmap, FCT, stereo):
#  xyz = stNoCirc.stereo2CartMat(stereo, cmap)
#  lonlat =  stNoCirc.cart2LonLatMat(xyz)
#  FCT[:] = 2 * Omega * np.sin(lonlat[:,1])


#fCor_slow=functionNumpy(1, fCoriolis, [xyzCircle])
#fCor = functionPrecomputed(groups, integOrder, 1)
#fCor.compute(fCor_slow)


claw = dgEulerAtmLaw(dimension)
claw.setFilterMode(FILTER_LINEARVERTICAL);
#claw.setFilterMode(FILTER_NONE);
#claw.setCoriolisFactor(fCor)
claw.setCoordinatesFunction(xyzCircleInteg)
sigma_ch=1#Coriolis
sigma_cv=1#Coriolis
sigma_zh=0#Centrifugal
sigma_zv=0#Centrifugal
sigma_d=0#Irrelevant - should be removed
claw.setOmega(Omega, st, sigma_ch, sigma_cv, sigma_zh, sigma_zv, sigma_d, 0);


c0 = functionConstant([0])
c1 = functionConstant([1])
init=functionNumpy(claw.getNbFields(), initialCondition, [xyzCircle, c0])
initHs=functionNumpy(claw.getNbFields(), initialCondition, [xyzCircle, c1])
solution = dgDofContainer(groups, claw.getNbFields())
solution.interpolate(init)
initDof = dgDofContainer(groups, claw.getNbFields())
initDof.interpolate(init)
initDofHs = dgDofContainer(groups, claw.getNbFields())
initDofHs.interpolate(initHs)

solution.axpy(initDofHs,-1)
initDof.axpy(initDofHs,-1)

rhoHs = dgDofContainer(groups, 1)
rhoHs.interpolate(functionExtractCompNew (initHs, 0));
rhoThetaHs = dgDofContainer(groups, 1)
rhoThetaHs.interpolate(functionExtractCompNew (initHs, 4));
claw.setHydrostaticState(rhoHs,rhoThetaHs)
claw.setPhysicalConstants(gamma,Rd,p0,g,R)

def uvw(cmap, FCT, sol, rhoHs):
  rho = sol[:,0] + rhoHs[:]
  invRho = 1.0 / rho[:]
  FCT[:,0] = sol[:,1] * invRho
  FCT[:,1] = sol[:,2] * invRho  
  FCT[:,2] = sol[:,3] * invRho

trms = claw.getActiveTerms();
trms.setAdvV(True);
trms.setAdvT(True);
trms.setPGrad(True);
trms.setVDiv(True);
trms.setCor(True);
trms.setGrav(True);
trms.setDiff(False);
trms.setLax(True);

toReplace =  VectorFunctorConst([function.getSolution(), function.getSolutionGradient()])
replaceBy =  VectorFunctorConst([initDof.getFunction(), initDof.getFunctionGradient()])
outsideBoundary = claw.newOutsideValueBoundaryGeneric("",toReplace,replaceBy)
boundaryWall = claw.newBoundaryWall()
#BOTTOM BND
claw.addBoundaryCondition('top_Top', boundaryWall)#
claw.addBoundaryCondition('top_Bottom', boundaryWall)#
#TOP BND
#claw.addBoundaryCondition('bottom_Top', outsideBoundary)
#claw.addBoundaryCondition('bottom_Bottom', outsideBoundary)
claw.addBoundaryCondition('bottom_Top', boundaryWall)
claw.addBoundaryCondition('bottom_Bottom', boundaryWall)


def expCoord(cmap, FCT, stereo, amplFactor):
    stereoAmpl = np.empty(stereo.shape)
    stereoAmpl[:,0] = stereo[:,0]
    stereoAmpl[:,1] = stereo[:,1]
    stereoAmpl[:,2] = stereo[:,2] * amplFactor
    xyz = stNoCirc.stereo2CartMat(stereoAmpl, cmap)
    FCT[:,:] = xyz[:,:]
c500 = functionConstant([500/X])
exportCoord=functionNumpy(3, expCoord, [xyzCircle, c500])

def exportFct(FCT, rhoHs, rhoThetaHs, sol):
    for i in range(0,FCT.size1()):
        #rhop
        FCT.set(i,0,sol(i,0))
        #rhouvw
        FCT.set(i,1,sol(i,1))
        FCT.set(i,2,sol(i,2))
        FCT.set(i,3,sol(i,3))
        FCT.set(i,4,sol(i,3))
        #rhothetap
        FCT.set(i,5,sol(i,4))
        FCT.set(i,6,rhoHs(i,0))
        FCT.set(i,7,rhoThetaHs(i,0))
#        FCT.set(i,8,vort(i,0))
#        FCT.set(i,9,vort(i,1))
#        FCT.set(i,10,vort(i,2))


nCompExp=[1,3,1,1,1,1]
namesExp=["rhoPot","rhoUVW","rhoW","rhoThetaP","rhoHs","rhoThetaHs"]
exportFunction=functionPython(sum(nCompExp), exportFct, [rhoHs.getFunction(), rhoThetaHs.getFunction(),solution.getFunction()])#rhoThetaHs.getFunction()])



#timeIter = dgERK(claw, None, DG_ERK_22)
petscIm = dgLinearSystemExtrusionDouble(claw, groups, extrusion)
#petscIm = dgLinearSystemExtrusionBlockDouble(claw, groups, extrusion)

petscIm.setParameter("petscOptions","-ksp_type preonly -pc_type lu")


#dofIm = dgDofManager.newDG(groups, claw, petscIm)
dofIm = dgDofManager.newDGBlock(groups, claw.getNbFields(), petscIm)
timeIter = dgIMEXRK(claw, dofIm, 2)     #timeorder
#timeIter.getNewton().setVerb(1)     #Verbosity


dt = min(claw.getMinOfTimeSteps(solution, extrusion),2000)
if (Msg.GetCommRank()==0):
    print ("Time step:",dt)
#nbSteps = int(math.ceil(Tf/dt))


t=Ti

def energy(cmap, FCT, stereo, rhoHs, rhoThetaHs, sol):
    rho = sol[:,0]+rhoHs[:]
    rhoTheta = sol[:,4]+rhoThetaHs[:]
    p = p0 * (rhoTheta * Rd / p0)**(gamma);
    T = rhoTheta/rho * (p / p0)**(Rd / Cp)
    u_st = np.empty((sol.shape[0],3))
    u_st[:,0] = sol[:,1]/rho
    u_st[:,1] = sol[:,2]/rho
    u_st[:,2] = sol[:,3]/rho
    e = Cv * T + 0.5 * (u_st[:,0]*u_st[:,0] + u_st[:,1]*u_st[:,1] + u_st[:,2]*u_st[:,2]) + g * stereo[:,2]
    FCT[:,0] = rho * e
    FCT[:,1] = rho * Cv * T
    FCT[:,2] = rho * (0.5 * (u_st[:,0]*u_st[:,0] + u_st[:,1]*u_st[:,1] + u_st[:,2]*u_st[:,2]))
    FCT[:,3] = rho * (g * stereo[:,2])
energyF=functionNumpy(4, energy, [xyzCircle, rhoHs.getFunction(), rhoThetaHs.getFunction(), solution.getFunction()])
energyIntegrator =  dgFunctionIntegrator(groups, energyF)
intEnergy = fullMatrixDouble(4,1)

def momentum(cmap, FCT, stereo, sol):
    rhou_st = np.empty((sol.shape[0],3))
    rhou_st[:,0] = sol[:,1]
    rhou_st[:,1] = sol[:,2]
    rhou_st[:,2] = sol[:,3]
    rhou_cart = stNoCirc.stereoToCartVecMat(stereo, rhou_st, cmap)
    FCT[:,0] = rhou_cart[:,0]
    FCT[:,1] = rhou_cart[:,1]
    FCT[:,2] = rhou_cart[:,2]
    FCT[:,3] = (FCT[:,0]**2 + FCT[:,1]**2 + FCT[:,2]**2)**0.5
momentumF=functionNumpy(4, momentum, [xyzCircle, solution.getFunction()])
momentumIntegrator =  dgFunctionIntegrator(groups, momentumF)
intMomentum = fullMatrixDouble(4,1)




def angularMomentum(cmap, FCT, stereo, rhou_cart):
    xyz = stNoCirc.stereo2CartMat(stereo, cmap)
    FCT[:,0] = xyz[:,1]*rhou_cart[:,2] - xyz[:,2]*rhou_cart[:,1]
    FCT[:,1] = xyz[:,2]*rhou_cart[:,0] - xyz[:,0]*rhou_cart[:,2]
    FCT[:,2] = xyz[:,0]*rhou_cart[:,1] - xyz[:,1]*rhou_cart[:,0]
    FCT[:,3] = (R + stereo[:,2])**2 * Omega
angularMomentumF=functionNumpy(4, angularMomentum, [xyzCircle, momentumF])
angularMomentumIntegrator =  dgFunctionIntegrator(groups, angularMomentumF)
intAngularMomentum = fullMatrixDouble(4,1)


def densityOne(cmap, FCT, rhoHs, sol):
    FCT[:,0] = rhoHs[:] + sol[:,0]
    FCT[:,1] = 1
densityOneF=functionNumpy(2, densityOne, [rhoHs.getFunction(), solution.getFunction()])
massVolIntegrator =  dgFunctionIntegrator(groups, densityOneF)
intMassVol = fullMatrixDouble(2,1)



n_export=0
start = time.clock()

if (Msg.GetCommRank()==0):
  if not os.path.exists(exportBaseDir+'/'+expDirName):
    os.makedirs(exportBaseDir+'/'+expDirName)
  integFile = open(exportBaseDir+'/'+expDirName+'/integrals','w')
  integFile.close()

Msg.Barrier()

solutionExporter = dgIdxExporter(solution,  exportBaseDir + '/' + expDirName + '/solution')

i=0
while (True):
    if (i < nbStart):
        t=t+dt
        i=i+1
        continue
    if (i == nbStart and i > 0):
        loadFile=exportBaseDir + '/' + loadDirName + '/solution/solution-'+ ("%06d" % nbStart)+'.idx'
        Msg.Info('Reloading solution '+loadFile)
        solution.importIdx(loadFile) 
    if (i%checkPointInterv==0 and i > nbStart and checkPointInterv > 0):
      Msg.Info('Writing solution for reload')
      solutionExporter.exportIdx(i, t)                                                                                                       

    if (i%10 == 0):
       massVolIntegrator.compute(intMassVol)
       energyIntegrator.compute(intEnergy)
       momentumIntegrator.compute(intMomentum)
       angularMomentumIntegrator.compute(intAngularMomentum)
       if (Msg.GetCommRank()==0):
         integFile = open(exportBaseDir+'/'+expDirName+'/integrals','a')
         print("Time %e\nMass %.15e Momentum %.15e %.15e %.15e Ang Momentum %.15e %.15e %.15e\nTot Energy %.15e Internal %.15e Kinetic %.15e Potential %.15e" %(t,intMassVol(0,0), intMomentum(0,0), intMomentum(1,0), intMomentum(2,0), intAngularMomentum(0,0), intAngularMomentum(1,0), intAngularMomentum(2,0), intEnergy(0,0), intEnergy(1,0), intEnergy(2,0), intEnergy(3,0)))
         integFile.write("%e %.25e %.25e %.25e %.25e %.25e %.25e %.25e %.25e %.25e %.25e %.25e %.25e %.25e \n" %(t,intMassVol(0,0), intMomentum(0,0), intMomentum(1,0), intMomentum(2,0), intMomentum(3,0), intAngularMomentum(0,0), intAngularMomentum(1,0), intAngularMomentum(2,0), intAngularMomentum(3,0), intEnergy(0,0), intEnergy(1,0), intEnergy(2,0), intEnergy(3,0)))
         integFile.close()

    if (i%10 == 0):
      if (math.isnan(solution.norm())):
        Msg.Fatal("Nan number in solution");
    startIter = time.clock()
    if (i%500 == 0 or (Tf-t)/Tf < 1e-12):
        solution.exportFunctionVtk(exportFunction,exportBaseDir+'/'+expDirName+'/export', t, i,"solution",nCompExp,namesExp,exportCoord)
        if (Msg.GetCommRank()==0):
            print ('\nWriting output',n_export,'at time',t,'over',Tf,'and step',i)
            elapsed = (time.clock() - start)
            print ('Time elapsed: ',elapsed,' s')
            n_export=n_export+1
        if ((Tf-t)/Tf < 1e-12):
          break;
    if (t+dt > Tf):
      dt = Tf-t
      if (Msg.GetCommRank()==0):
        print ("Last time step:",dt)
    norm = timeIter.iterate (solution, dt, t)
    t=t+dt
    if (Msg.GetCommRank()==0):
        sys.stdout.write(('[%.2f]')%(time.clock()-startIter))
        sys.stdout.flush()
    i=i+1
print ('')    

elapsed = (time.clock() - start)
if (Msg.GetCommRank()==0):
  print ('Time elapsed: ',elapsed,' s')


Msg.Exit(0)