forked from LisaNeef/DARTpy
-
Notifications
You must be signed in to change notification settings - Fork 0
/
TEM.py
154 lines (131 loc) · 4.83 KB
/
TEM.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
# Python module for dealing with TEM diagnostics
# Lisa Neef, 26 Aug 2015
# load the required packages
import numpy as np
#import datetime
import os.path
from netCDF4 import Dataset
import DART_state_space as DSS
#-------- constants
H = 7.0E3 # approximate scale height
Rd = 286.9968933 # Gas constant for dry air J/degree/kg
def load_Wang_TEM_file(E,datetime_in,hostname='taurus',verbose=False):
"""
Load the file of TEM diagnostics computed for a certain experiment and a
given month,
using code by Wuke Wang.
Inputs:
E : experiment dictionary
datetime_in : datetime.datetime object for the file to load
hostname : default is taurus
verbose : default is False
These are the variables that this subroutine can read (and what is allowed
in E['variable']:
float VTy(time, lev, lat, ens) ;
VTy:long_name = "VTy: Vstar*dT/dy" ;
VTy:units = "K/day" ;
float WS(time, lev, lat, ens) ;
WS:long_name = "WS: Wstar*S, S=H*N2/R" ;
WS:units = "K/day" ;
float VSTAR(time, lev, lat, ens) ;
VSTAR:long_name = "VSTAR" ;
VSTAR:units = "M/S" ;
float WSTAR(time, lev, lat, ens) ;
WSTAR:long_name = "WSTAR" ;
WSTAR:units = "M/S" ;
float FPHI(time, lev, lat, ens) ;
FPHI:long_name = "FPHI" ;
FPHI:units = "KG/S2" ;
float FZ(time, lev, lat, ens) ;
FZ:long_name = "FZ" ;
FZ:units = "KG/S2" ;
float DELF(time, lev, lat, ens) ;
DELF:long_name = "DELF" ;
DELF:units = "M/S2" ;
"""
# if the variable given in E isn't a TEM diagnostic, change it to the default variable
# wstar (residual vertical velocity)
tem_variables_list = ['VSTAR','WSTAR','FPHI','FZ','DELF']
dynamical_heating_rates_list = ['VTY','WS']
variable_name = E['variable']
if variable_name.upper() not in tem_variables_list+dynamical_heating_rates_list:
print(variable_name+' is not a valid diagnostic -- retrieving w* instead')
variable_name = 'WSTAR'
# find the file path corresponding to this experiment
import experiment_settings as es
ff = es.exp_paths_TEM(E,datetime_in)
# load the file
if os.path.isfile(ff):
VV = None
if verbose:
print('Loading TEM diagnostics file file '+ff+' and variable '+variable_name)
f = Dataset(ff,'r')
lat = f.variables['lat'][:]
lev = f.variables['lev'][:]
time = f.variables['time'][:]
VV = f.variables[variable_name][:]
if VV is None:
print('Unable to find variable '+E['variable']+' in file '+ff)
f.close()
# bad flag is -999 -- turn it into np.nan
# actually there seem to be other large negative numbers in here that aren't physical -
# maybe they were created by the daysplit step in CDO
VV[np.abs(VV)>900.]=np.nan
#VV[VV==-999]=np.nan
#VV[VV<-500]=np.nan
# select the vertical and lat ranges specified in E
# if only one number is specified, find the lev,lat, or lon closest to it
# TODO: make this an external subroutine
levrange=E['levrange']
if levrange[0] == levrange[1]:
ll = levrange[0]
idx = (np.abs(lev-ll)).argmin()
lev2 = lev[idx]
k1 = idx
k2 = idx
else:
k1 = (np.abs(lev-levrange[1])).argmin()
k2 = (np.abs(lev-levrange[0])).argmin()
lev2 = lev[k1:k2+1]
latrange=E['latrange']
if latrange[0] == latrange[1]:
ll = latrange[0]
idx = (np.abs(lat-ll)).argmin()
lat2 = lat[idx]
k1 = idx
k2 = idx
else:
# selection of the right latitude range depends on whether they go from south to north, or vice versa
nlat = len(lat)
if lat[0] < lat[nlat-1]:
j2 = (np.abs(lat-latrange[1])).argmin()
j1 = (np.abs(lat-latrange[0])).argmin()
lat2 = lat[j1:j2+1]
else:
j1 = (np.abs(lat-latrange[1])).argmin()
j2 = (np.abs(lat-latrange[0])).argmin()
lat2 = lat[j1:j2+1]
# if lat2 or lev2 are single numbers, turn them into length-1 lists. -- this
# is needed to make the plotting codes work
if isinstance(lat2, np.float64) or isinstance(lat2, np.float32) :
lat2=[lat2]
if isinstance(lev2, np.float64) or isinstance(lev2, np.float32) :
lev2=[lev2]
# now select the relevant lat and lev regions
if 'ERA' in E['exp_name']:
# wuke's TEM diagnostics for ERA have shape time x lev x lat
Vout = VV[:,k1:k2+1,j1:j2+1]
else:
# All the variables in Wuke's TEM diagnostics for WACCM have shape time x lev x lat x ensemble_member
Vout = VV[:,k1:k2+1,j1:j2+1,:]
# finally, for dynamical heating due to vertical residual circulation, we are actually interested in -wstar*S,
# whereas Wuke's data just has wstar*S -- so reverse the sign here.
if variable_name == 'WS':
Vout = -Vout
# for file not found
else:
print('Unable to find TEM diagnostic file '+ff)
Vout = None
lat2 = None
lev2 = None
return Vout,lat2,lev2