-
Notifications
You must be signed in to change notification settings - Fork 2
/
heatbudget_snaps.py
152 lines (125 loc) · 6.54 KB
/
heatbudget_snaps.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
from netCDF4 import *
from numpy import *
# For the given simulation, calculate heat budget terms for the volume enclosed
# by the given bounds on longitude, latitude, and depth. Output as NetCDF.
# Input: base_dir = full path to simulation output directory in single quotes
# min_year, max_year = of n in folders "outputn" inside base_dir
# min_lon, max_lon = bounds on longitude, from -180 to 180
# min_lat, max_lat = bounds on latitude, from -90 to 90
# min_depth = shallowest depth in m
# max_depth = deepest depth in m
# Example: from heatbudget import *
# heatbudget('/short/v45/kaa561/mom/archive/gfdl_nyf_1080_AntRunoff_v2/', 101, 110, -70, -65, -75, -65, 500, 1000)
def heatbudget (base_dir, min_year, max_year, min_lon, max_lon, min_lat, max_lat, min_depth, max_depth):
time = range(min_year, max_year+1)
# Read the area of tracer grid cells
input_file = base_dir + 'output' + str(min_year) + '/ocean_grid.nc'
input = Dataset(input_file, 'r')
area = input.variables['area_t'][:,:]
input.close()
print input_file
# Read the longitude, latitude, and depth
input_file = base_dir + 'output' + str(min_year) + '/ocean__' + str(min_year) + '_003' + '.nc'
print 'Input file ' + input_file
input = Dataset(input_file, 'r')
lon = input.variables['xt_ocean'][:]
lat = input.variables['yt_ocean'][:]
depth = input.variables['st_ocean'][:]
# Copy the depth vector into a 3D array with the same dimensions as the heat budget terms; note this requires a transpose (.T)
depth_3d = tile(depth, (len(lon), len(lat), 1)).T
input.close()
# Convert the bounds on longitude so they work with grid (-280 to 80)
if min_lon > 80:
min_lon = min_lon - 360
if max_lon > 80:
max_lon = max_lon - 360
if max_lon < min_lon:
print 'Invalid longitude bounds; grid splits at 80E'
# Find the smallest value of i such that lon[i] > min_lon
min_i = nonzero(lon > min_lon)[0][0]
# Find the smallest value of i such that lon[i] > max_lon
# In practice the actual maximum index selected by min_i:max_i will be max_i - 1
max_i = nonzero(lon > max_lon)[0][0]
# Similarly for j (latitude) and k (depth)
min_j = nonzero(lat > min_lat)[0][0]
max_j = nonzero(lat > max_lat)[0][0]
min_k = nonzero(depth > min_depth)[0][0]
max_k = nonzero(depth > max_depth)[0][0]
# Apply these bounds to the depth_3d and area arrays
depth_3d = depth_3d[min_k:max_k, min_j:max_j, min_i:max_i]
area = area[min_j:max_j, min_i:max_i]
# Calculate volume of the region by summing depth*area over all 3 dimensions
# Note that the 2D area array is automatically broadcast to be 3D
volume = sum(depth_3d*area)
# Set up arrays to hold the heat budget terms (scalars) for each year
Tadv_avg = zeros(len(time)*73)
TKPP_avg = zeros(len(time)*73)
Tsub_avg = zeros(len(time)*73)
Ttend_avg = zeros(len(time)*73)
Tvdiff_avg = zeros(len(time)*73)
swheat_avg = zeros(len(time)*73)
Ttend_theoretical = zeros(len(time)*73)
# Loop through the output
for n in range(min_year, max_year+1):
for d in range(3,364,5):
if d < 10:
input_file= base_dir + 'output' + str(n) + '/ocean__' + str(n) + '_00' + str(d) + '.nc'
elif d < 100:
input_file= base_dir + 'output' + str(n) + '/ocean__' + str(n) + '_0' + str(d) + '.nc'
else:
input_file= base_dir + 'output' + str(n) + '/ocean__' + str(n) + '_' + str(d) + '.nc'
print 'Reading output' + input_file
input = Dataset(input_file, 'r')
# Read the heat budget component terms and apply the bounds on i, j, k
Tadv = input.variables['temp_advection'][0, min_k:max_k, min_j:max_j, min_i:max_i]
TKPP = input.variables['temp_nonlocal_KPP'][0, min_k:max_k, min_j:max_j, min_i:max_i]
Tsub = input.variables['temp_submeso'][0, min_k:max_k, min_j:max_j, min_i:max_i]
Ttend = input.variables['temp_tendency'][0, min_k:max_k, min_j:max_j, min_i:max_i]
Tvdiff = input.variables['temp_vdiffuse_diff_cbt'][0, min_k:max_k, min_j:max_j, min_i:max_i]
swheat = input.variables['sw_heat'][0, min_k:max_k, min_j:max_j, min_i:max_i]
input.close()
# Find the total of each term by summing term*area over all 3 dimensions
# Find the average of each term by dividing that result by the volume
Tadv_avg[n-min_year] = sum(Tadv*area)/volume
TKPP_avg[n-min_year] = sum(TKPP*area)/volume
Tsub_avg[n-min_year] = sum(Tsub*area)/volume
Tvdiff_avg[n-min_year] = sum(Tvdiff*area)/volume
swheat_avg[n-min_year] = sum(swheat*area)/volume
Ttend_avg[n-min_year] = sum(Ttend*area)/volume
# Find theoretical temperature tendency: sum of previous terms (except actual temperature tendency)
Ttend_theoretical[n-min_year] = Tadv_avg[n-min_year] + TKPP_avg[n-min_year] + Tsub_avg[n-min_year] + Tvdiff_avg[n-min_year] + swheat_avg[n-min_year]
# End of loop
# Convert longitude back to normal for output
if min_lon < -180:
min_lon = min_lon + 360
if max_lon < -180:
max_lon = max_lon + 360
# Output to NetCDF file
output_file = base_dir + 'heatbudget_' + str(min_year) +'to'+ str(max_year) + '_' + str(min_lon) + 'to' + str(max_lon) + '_' + str(min_lat) + 'to' + str(max_lat) + '_' + str(min_depth) + 'to' + str(max_depth) + '.nc'
output = Dataset(output_file, 'w', format='NETCDF4')
output.createDimension('time', None)
time_var = output.createVariable('time', 'f8', ('time',))
time_var[:] = time
time_var.units = 'years'
adv_var = output.createVariable('temp_advection', 'f8', ('time',))
adv_var[:] = Tadv_avg
adv_var.units = 'W/m^2'
kpp_var = output.createVariable('temp_nonlocal_KPP', 'f8', ('time',))
kpp_var[:] = TKPP_avg
kpp_var.units = 'W/m^2'
sub_var = output.createVariable('temp_submeso', 'f8', ('time',))
sub_var[:] = Tsub_avg
sub_var.units = 'W/m^2'
vdiff_var = output.createVariable('temp_vdiffuse_diff_cbt', 'f8', ('time',))
vdiff_var[:] = Tvdiff_avg
vdiff_var.units = 'W/m^2'
swheat_var = output.createVariable('sw_heat', 'f8', ('time',))
swheat_var[:] = swheat_avg
swheat_var.units = 'W/m^2'
tend_var = output.createVariable('temp_tendency', 'f8', ('time',))
tend_var[:] = Ttend_avg
tend_var.units = 'W/m^2'
tendtheor_var = output.createVariable('theoretical_temp_tendency', 'f8', ('time',))
tendtheor_var[:] = Ttend_theoretical
tendtheor_var.units = 'W/m^2'
output.close()