/
aux_func_3dfields.py
166 lines (130 loc) · 4.62 KB
/
aux_func_3dfields.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
154
155
156
157
158
159
160
161
162
163
164
# Auxiliary functions for analysis of
# ll_4320 mitgcm simulation
# crocha, sio summer 2014
import numpy as np
def rmean(A):
""" Removes time-mean of llc_4320 3d fields; axis=2 is time"""
ix,jx,kx = A.shape
Am = np.repeat(A.mean(axis=2),kx)
Am = Am.reshape(ix,jx,kx)
return A-Am
def spec_est_meridional(U,dx):
""" Computes 1d (meridional) spectral estimates of 3d llc_4320 fields"""
ix,jx,kx = U.shape
N = ix # record length
df = 1./(N*dx) # frequency resolution [cycles / (unit time)]
fNy = 1./(2*dx) # Nyquist frequency
an = np.fft.fft(U,axis=0)
an = an[1:N/2-1,:,:]
E = 2*(an*an.conj())/df/(N**2) # spectral estimate
f = np.arange(1,N/2-1)*df
return E.mean(axis=2),f,df,fNy
def spec_est_zonal(U,dx):
""" Computes 1d (zonal) spectral estimates of 3d llc_4320 fields"""
ix,jx,kx = U.shape
N = jx # record length
df = 1./(N*dx) # frequency resolution [cycles / (unit time)]
fNy = 1./(2*dx) # Nyquist frequency
an = np.fft.fft(U,axis=1)
an = an[:,1:N/2-1,:]
E = 2*(an*an.conj())/df/(N**2) # spectral estimate
f = np.arange(1,N/2-1)*df
return E.mean(axis=2),f,df,fNy
def spec_est_time(U,dt):
""" Computes spectral estimate in time (axis=2) """
ix,jx,kx = U.shape
N = kx # record length
df = 1./(N*dt) # frequency resolution [cycles / (unit time)]
fNy = 1./(2*dt) # Nyquist frequency
an = np.fft.fft(U,axis=2)
an = an[:,:,1:N/2-1]
E = 2*(an*an.conj())/df/(N**2) # spectral estimate
f = np.arange(1,N/2-1)*df
return E.mean(axis=0),f,df,fNy
def spec_error(E,sn,ci):
""" Computes confidence interval for spectral
estimate E.
sn is the number of spectral realizations (dof/2)
ci = .95 for 95 % confidence interval
returns lower (El) and upper (Eu) bounds on E
as well as pdf and cdf used to estimate errors """
## params
dbin = .001
yN = np.arange(0,5.+dbin,dbin)
dof = 2.*sn # DOF = 2 x # of spectral estimates
## PDF for E/E0, where E (E0) is the estimate (true)
## process spectrum (basically a chi^2 distribution)
C = dof / ( (2**sn) * np.math.gamma(sn) ) # constant
pdf_yN = C * ( (dof*yN)**(sn-1) ) * np.exp( -(sn*yN) ) # chi^2(E/E0)
## CDF
cdf_yN = np.cumsum(pdf_yN*dbin) # trapezoidal-like integration
## compute confidence limits
# lower
el = ci
fl = np.where( np.abs(cdf_yN - el) == np.abs(cdf_yN - el).min())
El = E/yN[fl]
# upper
eu = 1 - ci
fu = np.where( np.abs(cdf_yN - eu) == np.abs(cdf_yN - eu).min())
Eu = E/yN[fu]
return El, Eu, cdf_yN, pdf_yN
# if sn larger than 150, assume it is normally-distributed (e.g., Bendat and Piersol)
def spec_error2(E,sn):
std_E = (1/np.sqrt(sn))
El = E/(1 + 2*std_E)
Eu = E/(1 - 2*std_E)
return El, Eu, std_E
def spectral_slope(k,E,kmin,kmax,stdE):
''' compute spectral slope in log space in
a wavenumber subrange [kmin,kmax],
m: spectral slope; mm: uncertainty'''
fr = np.where((k>=kmin)&(k<=kmax))
ki = np.matrix((np.log10(k[fr]))).T
Ei = np.matrix(np.log10(np.real(E[fr]))).T
dd = np.matrix(np.eye(ki.size)*((np.abs(np.log10(stdE)))**2))
G = np.matrix(np.append(np.ones((ki.size,1)),ki,axis=1))
Gg = ((G.T*G).I)*G.T
m = Gg*Ei
mm = np.sqrt(np.array(Gg*dd*Gg.T)[1,1])
yfit = np.array(G*m)
m = np.array(m)[1]
return m, mm
def leg_width(lg,fs):
"""" Sets the linewidth of each legend object """
for legobj in lg.legendHandles:
legobj.set_linewidth(fs)
def auto_corr(x):
""" Computes auto-correlation of 1d array """
a = np.correlate(x,x,mode='full')
a = a[x.size-1:]
a = a/a[0]
return a
def fit_gauss(x,y):
""" Estimate characteristic scale of a auto-correlation function
by fitting a Gaussian to auto_corr"""
y = np.matrix(np.log(y)).T
A1 = np .matrix(np.ones((x.size,1)))
A2 = np.matrix(-(x**2)).T
A = A2
xmax = 650
we = np.float(xmax) - x
we = (we/(xmax))**2
We = np.matrix(np.diag(we))
Gg = ((A.T*We*A).I)*A.T
c = Gg*y
Lfit = np.sqrt(1/c[-1])
yfit = np.exp( A*c )
return Lfit
def block_ave(dist,U,dx):
""" Block-averages the array u onto grid with resolution dx """
ix,jx,kx = U.shape
disti = np.arange(dx/2.,dist[-1]+dx/2.,dx)
Ui = np.zeros((disti.size,jx,kx))
for i in range(0,disti.size):
fn = ((dist >= disti[i]-dx/2.) & (dist <= disti[i]+dx/2.))
fns = np.sum(fn)
if fns>0:
Ui[i,:,:] = np.nansum(U[fn,:,:],axis=0)/fns
else:
Ui[i,:,:] = np.nan
return Ui