source: trunk/LMDZ.PLUTO/util/script_figures/movie_winds.py @ 3868

#! /usr/bin/env python
import os, sys
from    netCDF4               import    Dataset
from    numpy                 import    *
import  numpy                 as        np
import  matplotlib.pyplot     as        mpl
from matplotlib.cm import get_cmap
import pylab
from matplotlib import ticker
import matplotlib.colors as colors
import datetime
from mpl_toolkits.basemap import Basemap, shiftgrid
from FV3_utils import *

############################
filename1=name+"_A.nc"
filename2=name+".nc"
filename3=name+".nc"
# filename3="../../phisinit.nc"
var="tsurf" #variable
phisinit="phisinit" #variable
vari="u"
varj="v"
tint=[30,32] #Time must be as written in the input file
# tint=None
#tintstr=["03:00","09:00","15:00","21:00"] #Time must be as written in the input file

altitude = 10 # default value, in km ! WARNING: only integers accepted for animation script yet
if len(sys.argv)>1:
    altitude=float(sys.argv[1])
    print("Setting altitude to ",altitude, "km")

font=26

nc1=Dataset(filename1)
nc2=Dataset(filename2)
nc3=Dataset(filename3)

lat=getvar(nc1,"latitude")
lon=getvar(nc1,"longitude")
alt=getvar(nc1,"altitude")
tim=getvar(nc1,"Time")
############################

def getwinds(lon,lat,vecx,vecy):
          svx='None'   # arrow every svx box
          svy='None'
          svx=1
          svy=1
          angle='uv'       # 'xy'
          color='black'    # arrow color
          pivot='mid'      # arrow around middle of box. Alternative : tip
          scale=100     # scale arrow : plus grand = fleche plus petite
          width=0.002      # width arrow
          linewidths=0.5   # epaisseur contour arrow
          edgecolors='k'   # couleur contour arrow

    #  *scale*: [ *None* | float ]
    #  Data units per arrow length unit, e.g., m/s per plot width; a smaller
    #  scale parameter makes the arrow longer.  If *None*, a simple
    #  autoscaling algorithm is used, based on the average vector length
    #  and the number of vectors.  The arrow length unit is given by
    #  the *scale_units* parameter

    #  *scale_units*: *None*, or any of the *units* options.
    #  For example, if *scale_units* is 'inches', *scale* is 2.0, and
    #  ``(u,v) = (1,0)``, then the vector will be 0.5 inches long.
    #  If *scale_units* is 'width', then the vector will be half the width
    #  of the axes.

    #  If *scale_units* is 'x' then the vector will be 0.5 x-axis
    #  units.  To plot vectors in the x-y plane, with u and v having
    #  the same units as x and y, use
    #  "angles='xy', scale_units='xy', scale=1".

          x, y = np.meshgrid(lon,lat)
          q = mpl.quiver( x[::svy,::svx],y[::svy,::svx],\
           vecx[::svy,::svx],vecy[::svy,::svx],\
          angles=angle,color=color,pivot=pivot,\
          scale=scale,width=width,linewidths=linewidths,edgecolors=edgecolors)

          # make vector key.
          #keyh = 1.025 ; keyv = 1.05 # upper right corner over colorbar
          keyh = 0.97 ; keyv = 1.06
          keyh = 0.03 ; keyv = 1.07
          #keyh = -0.03 ; keyv = 1.08 # upper left corner
          labelpos='E'    # position label compared to arrow : N S E W
          p = mpl.quiverkey(q,keyh,keyv,\
          5.0,r'$5 m/s$',\
          fontproperties={'size': font,'weight': 'bold'},\
          color='black',labelpos=labelpos,labelsep = 0.07)

def getfigvar(i):
    pal=get_cmap(name="OrRd")
    lev=np.linspace(37,51,15)
    newlon=lon+180
    CF=mpl.contourf(newlon, lat, myvarbis,lev,cmap=pal,extend='both',alpha=0.7)
    yticks=[-90,-60,-30,0,30,60,90]
    xticks=[0,60,120,180,240,300,360]
    cbar=mpl.colorbar(CF,shrink=1, format="%1.0f")
    cbar.ax.set_title("Tsurf [K]",y=1.04,fontsize=font, pad=20)

    for t in cbar.ax.get_yticklabels():
        t.set_fontsize(font)

    c=mpl.contour(newlon, lat, myvar2bis, 10,levels=np.linspace(-4,4,8), colors = 'k', linewidths = 3.5)
    mpl.clabel(c, fmt='%2.1f',inline=1, colors='k', fontsize=23,inline_spacing=1)
    mpl.title('Winds & Surface Temperature',fontsize=font, pad=20)
    # mpl.title('Local Time at Sputnik Planum='+str(i*3)+'H00',fontsize=font)
    mpl.ylabel('Latitude (deg)',labelpad=10,fontsize=font)
    mpl.xlabel('West Longitude (deg)',labelpad=10, fontsize=font)
    mpl.xticks(xticks,fontsize=font)
    mpl.yticks(yticks,fontsize=font)
    getwinds(newlon,lat,u,v)

def getnumalt(choicealt,alt):
    numalt=np.where(abs(alt-choicealt)==min(abs(alt-choicealt)))
    numalt=numalt[0][0]
    return numalt

#######################
numalt=getnumalt(altitude,alt)
print(('numalt =',numalt,'altitude=',alt[numalt]))
uini=getvar(nc1,vari,times=tint)[:,numalt]
vini=getvar(nc1,varj,times=tint)[:,numalt]
myvar=getvar(nc2,var,times=tint)
myvar2=getvar(nc3,phisinit) # phisinitmyvar2=myvar2/0.6169/1000.  # altitude km
nbfig=uini.shape[0]
print(("nbfig=",nbfig))
myvar2bis=switchlon(myvar2, lon)

os.makedirs("movie_winds", exist_ok=True)

for i in range(nbfig):
   mpl.figure(figsize=(30, 15))
   u2=uini[i,:,:]
   v2=vini[i,:,:]
   myv=myvar[i,:,:]
   u=switchlon(u2, lon)
   v=switchlon(v2, lon)
   myvarbis=switchlon(myv, lon)
   print(i,"/",nbfig)
   getfigvar(i)
#    mpl.tight_layout()
   mpl.savefig(f"movie_winds/mapwinds_{altitude:.0f}km_{i:03}.eps",dpi=200)
   mpl.savefig(f"movie_winds/mapwinds_{altitude:.0f}km_{i:03}.png",dpi=200)

left  = None  # the left side of the subplots of the figure
right = None   # the right side of the subplots of the figure
bottom = None  # the bottom of the subplots of the figure
top = None     # the top of the subplots of the figure
wspace = None  # the amount of width reserved for blank space between subplots
hspace = None  # the amount of height reserved for white space between subplots
#mpl.subplots_adjust(left, bottom, right, top, wspace, hspace)
#mpl.subplots_adjust(hspace = .1)

#mpl.show()