-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathccsm4_oceanhelper.py
137 lines (102 loc) · 5.13 KB
/
ccsm4_oceanhelper.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
# ccsm4 ocean helper functions
import numpy as np
import numpy.ma as ma
def mask_seafloor(input, kmt, axis=0):
""" def mask_seafloor(input, kmt):
default is that input has depth (zt) as first dimension.
axis is axis of depth dimension.
Change axis=1 if time is first and depth is second.
returns input as masked array
"""
tmp = ma.zeros(input.shape)
ndepth = input.shape[axis]
if axis==1:
# need to tile kmt
kmtt = np.tile(kmt,(input.shape[0],1,1))
for lii in np.arange(0,ndepth): # loop through each depth layer
# mask out levels below sea floor
if axis==0:
tmp[lii,...] = ma.masked_where(kmt <= lii, input[lii,...])
elif axis==1:
tmp[:,lii,...] = ma.masked_where(kmtt <= lii, input[:,lii,...])
else:
print 'axis must = 1 or 0'
return tmp
def ocn_zonal_mean(input, kmt, axis=2):
""" def ocn_zonal_mean(input, kmt, axis=2):
Take zonal mean of ocean data. To do this, must know sea floor depth.
Default is that longitude is axis=2 (ie variable depth x lat x lon).
Change to axis=3 if time dimension is first.
returns zonal mean of size dim(input) - 1
"""
if axis==3:
# there is a time dimension, so depth is axis 1
tmp = mask_seafloor(input, kmt, axis=1)
else:
# no time dimension, depth is axis 0
tmp = mask_seafloor(input, kmt, axis=0)
zm = ma.mean(tmp,axis=axis)
return zm
def ocn_regional_average(input, tlat,tlon, tarea, limsdict):
""" def ocn_regional_average(input, tlat,tlon, tarea, limsdict):
horizontal regional average, given lat and lon limits.
For now this is only good in the southern hem because of
the more regular grid.
limsdict = ('latlims': (slim,nlim), 'lonlims': (wlim,elim))
@@@ not done yet 1/17/2015
"""
print 'ocn_regional_average() not implemented yet @@@@@'
latlims=limsdict['latlims']
slim = latlims[0]; nlim = latlims[1]
onelat=tlat[:,0]
# pick sub-region in tarea
# @@@@ I am not sure what to do about the weird grid in the northern hem. prob
# @@@@ need to remap first.
tareasub = tarea[np.logical_and(onelat<=nlim,onelat>slim),:] # sublat x lon
"""
Nlim=-72 # how close to "ice shelf"?
Slim=-74
# control: ==============================
# pick sublat range from area
tareapigcsub = tempareapigc[:,np.logical_and(onelon<= Nlim,onelon>Slim),:] # depth x sublat x lon
# total area in the sublat x lon box
totareapigxyc = ma.sum(ma.sum(tareapigcsub,axis=2),axis=1) # depth
# total area in the zonal dir
totareapigxc = ma.sum(tareapigcsub,axis=2) # depth x sublat
# tile sublat/lon box tot area with sublat range
totareapigxytc = np.tile(totareapigxyc,(tareapigcsub.shape[1],1))
totareapigxytc = np.transpose(totareapigxytc,(1,0))
# Now weights for averaging a zonal mean, with latitude
wgtspigc = totareapigxc / totareapigxytc
# Average meridionally
avgpigc,ret1 = ma.average(tempfldcpigzm[:,np.logical_and(onelon<= Nlim,onelon>Slim)],
axis=1,weights=wgtspigc,returned=True)
# === and with time dim
# Pert runs, with time =================
# Now pick the sublat range from tareapig and create weights
tareapigsub1 = tempareapig1[:,:,np.logical_and(onelon<= Nlim,onelon>Slim),:] # time x depth x sublat x lon
tareapigsub2 = tempareapig2[:,:,np.logical_and(onelon<= Nlim,onelon>Slim),:] # time x depth x sublat x lon
## total area in sublat x lon box
totareapigxy1 = ma.sum(ma.sum(tareapigsub1,axis=2),axis=2) # time x depth
totareapigxy2 = ma.sum(ma.sum(tareapigsub2,axis=2),axis=2) # time x depth
totareapigx1 = ma.sum(tareapigsub1,axis=3) # total area in the lon range for e/ time x depth x lat
totareapigx2 = ma.sum(tareapigsub2,axis=3) # total area in the lon range for e/ time x depth x lat
# should be able to take zonal mean no prob
# then average in meridional direction with weights totareapigx / totareapigxy per latitude
# zm already done: tempfldpigzm1,2
# first tile over sublat to create weights
totareapigxyt2 = np.tile(totareapigxy2,(tareapigsub2.shape[2],1,1))
totareapigxyt2=np.transpose(totareapigxyt2,(1,2,0)) # time x depth x lat?
wgtspig2 = totareapigx2 / totareapigxyt2
totareapigxyt1 = np.tile(totareapigxy1,(tareapigsub1.shape[2],1,1))
totareapigxyt1=np.transpose(totareapigxyt1,(1,2,0)) # time x depth x lat?
wgtspig1 = totareapigx1 / totareapigxyt1
print 'wgtspig1.shape: ' + str(wgtspig1.shape)
print 'wgtspig2.shape: ' + str(wgtspig2.shape)
print 'tempfldpigzm1.shape: ' + str(tempfldpigzm1.shape)
# then average meridionally
avgpig1,ret1 = ma.average(tempfldpigzm1[:,:,np.logical_and(onelon<= Nlim,onelon>Slim)],
axis=2,weights=wgtspig1,returned=True)
avgpig2,ret2 = ma.average(tempfldpigzm2[:,:,np.logical_and(onelon<= Nlim,onelon>Slim)],
axis=2,weights=wgtspig2,returned=True)
"""