from PyNIO import nio from spharm import Spharmt, regrid import numpy as np from dateutils import dateshift, daterange import sys, os datapath1 = '/lfs0/tmp/whitaker/ensda/ensrf_test2' runname1 = 'enkf' datapath2 = '/lfs0/tmp/whitaker/ncepgfs/gsi_test' runname2 = 'gdas1' #datapath2 = '/lfs0/tmp/whitaker/ensda/ensrf_test' #runname2 = 'enkf' date1 = sys.argv[1] date2 = sys.argv[2] fhr = sys.argv[3] varshort = sys.argv[4] print datapath1 print datapath2 print 'fhr = ',fhr def getmean(diff,coslats): meancoslats = coslats.mean() diffvar = (coslats*diff).mean() return diffvar/meancoslats def getcorr(f,a): f = f - f.mean(axis=0) a = a - a.mean(axis=0) covfa = (f*a).mean(axis=0) varf = (f**2).mean(axis=0) vara = (a**2).mean(axis=0) return covfa/(np.sqrt(varf)*np.sqrt(vara)) def getmse(f,a): return ((f-a)**2).mean(axis=0) nlonsin = 800; nlatsin = 400 nlons = 360; nlats = 181 ntrunc = 20 latbound = 20. sin = Spharmt(nlonsin,nlatsin) sout = Spharmt(nlons,nlats) delta = 360./nlons lats = 90.-delta*np.arange(nlats) coslats = np.cos((np.pi/180.)*lats) coslats = coslats[:,np.newaxis]*np.ones((nlats,nlons)) latnh = lats.tolist().index(latbound) latsh = lats.tolist().index(-latbound) coslatsnh = coslats[0:latnh+1,:] coslatssh = coslats[latsh:,:] coslatstr = coslats[latnh:latsh+1,:] lons = delta*np.arange(nlons) lons, lats = np.meshgrid(lons,lats[::-1]) filename = varshort if varshort == 'hgt': varname = 'HGT_3_ISBL_10' ecvarname = 'gh_P0_L100_GGA0' elif varshort == 'mslp': varname = 'PRMSL_3_MSL_10' ecvarname = 'msl_P0_L101_GGA0' elif varshort == 'temp': varname = 'TMP_3_ISBL_10' ecvarname = 't_P0_L100_GGA0' elif varshort == 'uwnd': varname = 'U_GRD_3_ISBL_10' ecvarname = 'u_P0_L100_GGA0' filename = 'uv' elif varshort == 'vwnd': varname = 'V_GRD_3_ISBL_10' ecvarname = 'v_P0_L100_GGA0' filename = 'uv' elif varshort in ['chi','psi']: filename = 'uv' varnamev = 'V_GRD_3_ISBL_10' varnameu = 'U_GRD_3_ISBL_10' ecvarnamev = 'v_P0_L100_GGA0' ecvarnameu = 'u_P0_L100_GGA0' else: print 'unknown variable',varshort raise SystemExit if filename in ['mslp']: levels = [None] elif filename == 'hgt': levels = [1000,925,850,700,500,300,250,200,50] else: levels = [1000,925,850,700,500,300,250,200] for level in levels: ecfilename = filename+'.grib' f = nio.open_file('/p72/reanl/whitaker/'+ecfilename) lonsin = f.variables['lon_0'][:] latsin = f.variables['lat_0'][:] datesin = f.variables['initial_time0'][:] dates = [] for date in datesin: d = date.split('/') mm = d[0]; dd = d[1]; yyyy = d[2][0:4] hh = (d[2].split('(')[1]).split(':')[0] dates.append(yyyy+mm+dd+hh) if varshort not in ['psi','chi']: varin = f.variables[ecvarname] else: varinu = f.variables[ecvarnameu] varinv = f.variables[ecvarnamev] datefcsts = daterange(date1,date2,12) fcsts_2 = np.zeros((len(datefcsts),nlats,nlons),np.float32) fcsts_1 = np.zeros((len(datefcsts),nlats,nlons),np.float32) anals = np.zeros((len(datefcsts),nlats,nlons),np.float32) if levels[0] is not None: eclevels = (0.01*f.variables['lv_ISBL0'][:]).tolist() nlevec = eclevels.index(level) for ntime,date in enumerate(datefcsts): dateverif = dateshift(date,int(fhr)) nt = dates.index(dateverif) if varshort in ['chi','psi']: if fhr=='0': grbfile_1 =\ os.path.join(os.path.join(datapath1,date),runname1+'.t'+date[8:10]+'z.pgrbanl') grbfile_2 =\ os.path.join(os.path.join(datapath2,date),runname2+'.t'+date[8:10]+'z.pgrbanl') else: grbfile_1 =\ os.path.join(os.path.join(datapath1,date),runname1+'.t'+date[8:10]+'z.pgrbf'+fhr) grbfile_2 =\ os.path.join(os.path.join(datapath2,date),runname2+'.t'+date[8:10]+'z.pgrbf'+fhr) f2 = nio.open_file(grbfile_2+'.grib') levels = f2.variables['lv_ISBL3'][:].tolist() nlev = levels.index(level) dat2u = f2.variables[varnameu][nlev] dat2v = f2.variables[varnamev][nlev] f2.close() f1 = nio.open_file(grbfile_1+'.grib') levels = f1.variables['lv_ISBL3'][:].tolist() nlev = levels.index(level) dat1u = f1.variables[varnameu][nlev] dat1v = f1.variables[varnamev][nlev] f1.close() datainu = varinu[nt,nlevec,:,:] datainv = varinv[nt,nlevec,:,:] datverifu = regrid(sin,sout,datainu,ntrunc=ntrunc) datverifv = regrid(sin,sout,datainv,ntrunc=ntrunc) if varshort == 'psi': datverif,chi =\ sout.getpsichi(datverifu,datverifv,ntrunc=ntrunc) dat2,chi = sout.getpsichi(dat2u,dat2v,ntrunc=ntrunc) dat1,chi = sout.getpsichi(dat1u,dat1v,ntrunc=ntrunc) else: psi,datverif = sout.getpsichi(datverifu,datverifv,ntrunc=ntrunc) psi,dat2 = sout.getpsichi(dat2u,dat2v,ntrunc=ntrunc) psi,dat1 = sout.getpsichi(dat1u,dat1v,ntrunc=ntrunc) else: if fhr=='0': grbfile_1 =\ os.path.join(os.path.join(datapath1,date),runname1+'.t'+date[8:10]+'z.pgrbanl') grbfile_2 =\ os.path.join(os.path.join(datapath2,date),runname2+'.t'+date[8:10]+'z.pgrbanl') else: grbfile_1 =\ os.path.join(os.path.join(datapath1,date),runname1+'.t'+date[8:10]+'z.pgrbf'+fhr) grbfile_2 =\ os.path.join(os.path.join(datapath2,date),runname2+'.t'+date[8:10]+'z.pgrbf'+fhr) f2 = nio.open_file(grbfile_2+'.grib') levels = f2.variables['lv_ISBL3'][:].tolist() if level is not None: nlev = levels.index(level) dat2 = f2.variables[varname][nlev] else: dat2 = f2.variables[varname][:] if ntrunc is not None: dat2spec = sout.grdtospec(dat2,ntrunc=ntrunc) dat2 = sout.spectogrd(dat2spec) f2.close() f1 = nio.open_file(grbfile_1+'.grib') if level is not None: levels = f1.variables['lv_ISBL3'][:].tolist() nlev = levels.index(level) dat1 = f1.variables[varname][nlev] else: dat1 = f1.variables[varname][:] if ntrunc is not None: dat1spec = sout.grdtospec(dat1,ntrunc=ntrunc) dat1 = sout.spectogrd(dat1spec) f1.close() if level is not None: datain = varin[nt,nlevec,:,:] else: datain = varin[nt,:,:] datverif = regrid(sin,sout,datain,ntrunc=ntrunc) fcsts_2[ntime] = dat2 fcsts_1[ntime] = dat1 anals[ntime] = datverif corr_2 = getcorr(fcsts_2,anals) corr_1 = getcorr(fcsts_1,anals) corrnh = corr_2[0:latnh+1,:] corrsh = corr_2[latsh:,:] corrtr = corr_2[latnh:latsh+1,:] corrg_2 = getmean(corr_2,coslats) corrnh_2 = getmean(corrnh,coslatsnh) corrsh_2 = getmean(corrsh,coslatssh) corrtr_2 = getmean(corrtr,coslatstr) corrnh = corr_1[0:latnh+1,:] corrsh = corr_1[latsh:,:] corrtr = corr_1[latnh:latsh+1,:] corrg_1 = getmean(corr_1,coslats) corrnh_1 = getmean(corrnh,coslatsnh) corrsh_1 = getmean(corrsh,coslatssh) corrtr_1 = getmean(corrtr,coslatstr) mse_2 = getmse(fcsts_2,anals) mse_1 = getmse(fcsts_1,anals) if varshort in ['psi','chi']: mse_2 = mse_2/1.e10 mse_1 = mse_1/1.e10 msenh = mse_2[0:latnh+1,:] msesh = mse_2[latsh:,:] msetr = mse_2[latnh:latsh+1,:] mseg_2 = np.sqrt(getmean(mse_2,coslats)) msenh_2 = np.sqrt(getmean(msenh,coslatsnh)) msesh_2 = np.sqrt(getmean(msesh,coslatssh)) msetr_2 = np.sqrt(getmean(msetr,coslatstr)) msenh = mse_1[0:latnh+1,:] msesh = mse_1[latsh:,:] msetr = mse_1[latnh:latsh+1,:] mseg_1 = np.sqrt(getmean(mse_1,coslats)) msenh_1 = np.sqrt(getmean(msenh,coslatsnh)) msesh_1 = np.sqrt(getmean(msesh,coslatssh)) msetr_1 = np.sqrt(getmean(msetr,coslatstr)) print '%s %6.3f %6.3f %6.3f %6.3f %6.3f %6.3f %6.3f %6.3f' %\ (level,corrnh_1,corrnh_2,corrtr_1,corrtr_2,\ corrsh_1,corrsh_2,corrg_1,corrg_2) print '%s %6.3f %6.3f %6.3f %6.3f %6.3f %6.3f %6.3f %6.3f' %\ (level,msenh_1,msenh_2,msetr_1,msetr_2,\ msesh_1,msesh_2,mseg_1,mseg_2)