def get_Lw(rho, z, z0=None, Nz=400, mode=1, omega=2 * np.pi / (12.42 * 3600)):
"""
A function to calculate the wavelength of a specific vertical mode for a specific frequency given stratification conditions.
z0 is the final depth. Will interpolate there if necessary.
"""
if z0 is None:
z0 = max(z)
Z = -np.linspace(0, 1, Nz) * z0
dZ = np.abs(Z[1] - Z[0])
# Interpolate the density profile onto all points
Fi = interp1d(z, rho, axis=0, fill_value='extrapolate')
rhoZ = Fi(Z)
drho_dz = grad_z(rhoZ, Z, axis=0)
N2 = -GRAV * drho_dz / RHO0
phi, c = isw.iwave_modes(N2, dZ)
phi_n = phi[:, mode]
c_n = c[mode]
k = omega / c_n
Lw = 2 * np.pi / k
return Lw
def solve_tg(N2, U, Uzz, dz):
"""
Use the scipy integration function (Uses Newton method)
"""
nx = N2.shape[0]
x = np.linspace(0, nx*dz, nx)
#B = (c - U)**2.
Fn = interp1d(x, N2, kind=2)
Fu = interp1d(x, U, kind=2)
Fuzz = interp1d(x, Uzz, kind=2)
def fun(x, y, p):
uc2 = (Fu(x) - p[0])**2
#A = Fuzz(x)/(Fu(x) - p[0]) - Fn(x)/uc2
A = -Fn(x)/uc2
return np.vstack([y[1], A * y[0] ])
def bc(ya, yb, p):
# BC: y(0) = y(H) = 0 and y'(0) = 0
#return np.array([ ya[0], yb[0], ya[1] ])
# BC: y(0) = y(H) = 0
return np.array([ ya[0], yb[0], yb[1] ])
y = np.zeros((2,x.size))
# Find the initial phase speed without shear
mode = 0
phi, cn = iwave_modes(N2, dz)
y[0,:] = phi[mode,:]
res_a = solve_bvp(fun, bc, x, y, p=[cn[mode]], tol=1e-10)
return res_a.sol(x)[0,:]
def tgsolve(z, U, N2, mode):
"""
Taylor-Goldstein equation solver
"""
status = 0
dz = np.abs(z[1] - z[0]) # Assume constant spacing
# Guess the initial conditions from the shear-free solution
phin, cn = iwave_modes(N2, dz)
cguess = cn[mode]
phi_2_guess = 1.05
# Compute the second derivative
Uz = np.gradient(U, dz)
Uzz = np.gradient(Uz, dz)
# Functions mapping the velocity and N^2 profile to any height
# (Use interpolation here although we could ascribe an analytic expression)
Fn = interp1d(z, N2, kind=2, fill_value='extrapolate')
Fu = interp1d(z, U, kind=2, fill_value='extrapolate')
Fuzz = interp1d(z, Uzz, kind=2, fill_value='extrapolate')
## Optimization step
phi_2_guess = -0.15
#soln = least_squares(objective2, [phi_2_guess, cguess], xtol=1e-12,gtol=1e-9,\
# #bounds=((-2.2,cguess-cguess*0.25), (2.2, cguess+cguess*0.25)),\
# args=(Fn,Fu,Fuzz,z), verbose=0,
# diff_step=1e-4,\
# method='dogbox')
soln = root(objective2, [phi_2_guess, cguess], args=(Fn,Fu,Fuzz, z),\
method='krylov', tol=1e-8)
#method='hybr', tol=1e-8)
# Go back and get the optimal profile using the solution
phi_2, cn = soln['x']
phiall = odeint(odefun2, [0, phi_2], z, args=(cn, Fn, Fu, Fuzz))
phi = phiall[:, 0]
# Normalize the profile
idx = np.where(np.abs(phi) == np.abs(phi).max())[0][0]
phi /= phi[idx]
#print phi[0], phi[-1]
if np.abs(phi[-1]) > 1e-6:
print('Warning: Convergence Issue with TG-solver')
status = -1
return phi, cn, status
def calc_vkdv_params(self, Nz, Nx):
# Initialise arrays
Phi = np.zeros((Nz, Nx))
Cn = np.zeros((Nx, ))
Alpha = np.zeros((Nx, ))
Beta = np.zeros((Nx, ))
Q = np.zeros((Nx, ))
r20 = np.zeros((Nx, ))
# Loop through and compute the eigenfunctions etc at each point
if self.verbose:
print('Calculating eigenfunctions...')
phi0, cn0 = isw.iwave_modes(self.N2[:, 0], self.dZ[0])
phi0 = phi0[:, self.mode]
phi0 = phi0 / np.abs(phi0).max()
phi0 *= np.sign(phi0.sum())
print_n = self.Nsubset * (Nx // self.Nsubset) // (100 /
self.print_freq)
for ii in range(0, Nx, self.Nsubset):
if (ii % (print_n) == 0) and self.verbose:
print('%3.1f %% complete...' %
(self.print_freq * ii / print_n))
#phi, cn = iwave_modes_sparse(N2[:,ii], dZ[ii], h[ii])
#phi, cn = isw.iwave_modes(self.N2[:,ii], self.dZ[ii], h[ii])
phi, cn = isw.iwave_modes(self.N2[:, ii], self.dZ[ii])
# Extract the mode of interest
phi_1 = phi[:, self.mode]
c1 = cn[self.mode]
# Normalize so the max(phi)=1
phi_1 = phi_1 / np.abs(phi_1).max()
phi_1 *= np.sign(phi_1.sum())
# Work out if we need to flip the sign (only really matters for higher modes)
dphi0 = phi0[1] - phi0[0]
dphi1 = phi_1[1] - phi_1[0]
if np.sign(dphi1) != np.sign(dphi0):
phi_1 *= -1
Cn[ii] = c1
Phi[:, ii] = phi_1
if self.ekdv:
# Compute cubic nonlinearity on the subsetted grid then interpolate
r20[ii] = isw.calc_r20(phi_1, c1, self.N2[:, ii], self.dZ[ii])
# Interpolate all of the variables back onto the regular grid
x = self.x
idx = list(range(0, Nx, self.Nsubset))
interpm = 'cubic'
F = interp1d(x[idx], Cn[idx], kind=interpm, fill_value='extrapolate')
Cn = F(x)
F = interp1d(x[idx],
Phi[:, idx],
kind=interpm,
axis=1,
fill_value='extrapolate')
Phi = F(x)
if self.ekdv:
F = interp1d(x[idx],
r20[idx],
kind=interpm,
fill_value='extrapolate')
r20 = F(x)
for ii in range(self.Nx):
phi_1 = Phi[:, ii]
c1 = Cn[ii]
Alpha[ii] = calc_alpha(phi_1, c1, self.dZ[ii])
Beta[ii] = calc_beta(phi_1, c1, self.dZ[ii])
#Q[ii] = calc_Qamp(phi_1, c1, self.dZ[ii])
Q[ii] = calc_Qamp_H97(phi_1, Phi[:,0],\
c1, Cn[0], self.dZ[ii], self.dZ[0])
return Phi, Cn, Alpha, Beta, Q, r20
def calc_nliw_params(infile, zmin, dz, mode=0, samples=500):
h5file = "%s_beta-samples-array-all-data.h5" % infile
csvfile = "%s.csv" % infile
nparams = 6
# Load the data
data = load_density_h5(h5file, samples)
# Load the input data from the csv file (this is currently not saved by the stan/r output)
rho = read_density_csv(csvfile)
time = rho.time.values
nparam, nt, ntrace = data[:].shape
zout = np.arange(0, zmin, -dz)
# Calculate c and alpha
#samples = ntrace
alpha_ens = np.zeros((nt, samples))
c_ens = np.zeros((nt, samples))
rand_loc = np.random.randint(0, samples, samples)
for tstep in tqdm(range(0, nt)):
#if (tstep%20==0):
# print('%d of %d...'%(tstep,nt))
for ii in range(samples):
#elif nparams == 6:
rhotmp = double_tanh(data[:, tstep, rand_loc[ii]], zout)
N2 = -9.81 / 1000 * np.gradient(rhotmp, -dz)
phi, cn = isw.iwave_modes(N2, dz)
phi_1 = phi[:, mode]
phi_1 = phi_1 / np.abs(phi_1).max()
phi_1 *= np.sign(phi_1.sum())
alpha = isw.calc_alpha(phi_1, cn[mode], N2, dz)
alpha_ens[tstep, ii] = alpha
c_ens[tstep, ii] = cn[mode]
#mykdv = kdv.KdV(rhotmp,zout)
# Export to an xarray data set
# Create an xray dataset with the output
dims2 = (
'time',
'ensemble',
)
dims3 = ('params', 'time', 'ensemble')
#time = rho.time.values
#time = range(nt)
coords2 = {'time': time, 'ensemble': range(samples)}
coords3 = {
'time': time,
'ensemble': range(samples),
'params': range(nparams)
}
cn_da = xr.DataArray(
c_ens,
coords=coords2,
dims=dims2,
attrs={
'long_name': '',
'units': ''
},
)
alpha_da = xr.DataArray(
alpha_ens,
coords=coords2,
dims=dims2,
attrs={
'long_name': '',
'units': ''
},
)
beta_da = xr.DataArray(
data,
coords=coords3,
dims=dims3,
attrs={
'long_name': '',
'units': ''
},
)
dsout = xr.Dataset({
'cn': cn_da,
'alpha': alpha_da,
'beta': beta_da,
'rho': rho
})
print(dsout)
return dsout
def calc_nliw_params(h5file, depthtxt, dz, mode=0, samples=None):
# Lload the data
data, time, rho, depth, rho_std, z_std, rho_mu = load_density_h5(h5file)
nparams, nt, ntrace = data[:].shape
# Calculate c and alpha
if samples is None:
samples = ntrace
alpha_ens = np.zeros((nt, samples))
c_ens = np.zeros((nt, samples))
alpha_mu_ens = np.zeros((nt, samples))
c_mu_ens = np.zeros((nt, samples))
# Depth file paramaters
x = depthtxt[:, 0]
zmins = -depthtxt[:, 1]
L = x[-1]
dx = x[1] - x[0]
nx = x.shape[0]
rand_loc = np.random.randint(0, ntrace, samples)
beta_out = np.zeros((nparams, nt, samples))
for tstep in tqdm(range(0, nt)):
#if (tstep%20==0):
# print('%d of %d...'%(tstep,nt))
print(tstep, nt, nx)
for ii in tqdm(range(samples)):
alpha_mu = 0
cn_mu = 0
for kk in range(nx):
zmin = zmins[kk]
zout = np.arange(0, zmin, -dz)
if nparams == 4:
rhotmp = single_tanh(data[:, tstep, rand_loc[ii]],
zout / z_std)
elif nparams == 6:
rhotmp = double_tanh(data[:, tstep, rand_loc[ii]],
zout / z_std)
elif nparams == 7:
rhotmp = double_tanh_7(data[:, tstep, rand_loc[ii]],
zout / z_std)
beta_out[:, tstep, ii] = data[:, tstep, rand_loc[ii]]
# Need to scale rho
rhotmp = rhotmp * rho_std + rho_mu
N2 = -9.81 / 1000 * np.gradient(rhotmp, -dz)
phi, cn = isw.iwave_modes(N2, dz)
phi_1 = phi[:, mode]
phi_1 = phi_1 / np.abs(phi_1).max()
phi_1 *= np.sign(phi_1.sum())
alpha = isw.calc_alpha(phi_1, cn[mode], N2, dz)
alpha_mu += alpha * dx
cn_mu += cn[mode] * dx
# Always leave the last value as alpha and cn
alpha_ens[tstep, ii] = alpha
c_ens[tstep, ii] = cn[mode]
alpha_mu_ens[tstep, ii] = alpha_mu / L
c_mu_ens[tstep, ii] = cn_mu / L
#mykdv = kdv.KdV(rhotmp,zout)
# Export to an xarray data set
# Create an xray dataset with the output
dims2 = (
'time',
'ensemble',
)
#dims2a = ('time','depth',)
dims3 = ('params', 'time', 'ensemble')
#time = rho.time.values
#time = range(nt)
coords2 = {'time': time, 'ensemble': range(samples)}
#coords2a = {'time':time, 'depth':depth[:,0]}
coords3 = {
'time': time,
'ensemble': range(samples),
'params': range(nparams)
}
#rho = xr.DataArray(rho.T,
# coords=coords2a,
# dims=dims2a,
# attrs={'long_name':'', 'units':''},
# )
cn_da = xr.DataArray(
c_ens,
coords=coords2,
dims=dims2,
attrs={
'long_name': '',
'units': ''
},
)
alpha_da = xr.DataArray(
alpha_ens,
coords=coords2,
dims=dims2,
attrs={
'long_name': '',
'units': ''
},
)
cn_mu_da = xr.DataArray(
c_mu_ens,
coords=coords2,
dims=dims2,
attrs={
'long_name': '',
'units': ''
},
)
alpha_mu_da = xr.DataArray(
alpha_mu_ens,
coords=coords2,
dims=dims2,
attrs={
'long_name': '',
'units': ''
},
)
beta_da = xr.DataArray(
beta_out,
coords=coords3,
dims=dims3,
attrs={
'long_name': '',
'units': ''
},
)
dsout = xr.Dataset({'cn':cn_da, 'alpha':alpha_da, 'beta':beta_da,\
'cn_mu':cn_mu_da, 'alpha_mu':alpha_mu_da,})
return dsout
def calc_nliw_params(h5file, zmin, dz, mode=0):
# Lload the data
#data,time, rho, depth, rho_std, z_std, rho_mu = load_density_h5(h5file)
data, time, rho_std, z_std, rho_mu = load_density_h5(h5file)
nparams, nt, ntrace = data[:].shape
zout = np.arange(0, zmin, -dz)
# Calculate c and alpha
samples = ntrace
alpha_ens = np.zeros((nt, samples))
c_ens = np.zeros((nt, samples))
rand_loc = np.random.randint(0, ntrace, samples)
for tstep in tqdm(range(0, nt)):
#if (tstep%20==0):
# print('%d of %d...'%(tstep,nt))
for ii in range(samples):
if nparams == 4:
rhotmp = single_tanh(data[:, tstep, rand_loc[ii]],
zout / z_std)
elif nparams == 6:
rhotmp = double_tanh(data[:, tstep, rand_loc[ii]],
zout / z_std)
elif nparams == 7:
rhotmp = double_tanh_7(data[:, tstep, rand_loc[ii]],
zout / z_std)
# Need to scale rho
rhotmp = rhotmp * rho_std + rho_mu
N2 = -9.81 / 1000 * np.gradient(rhotmp, -dz)
phi, cn = isw.iwave_modes(N2, dz)
phi_1 = phi[:, mode]
phi_1 = phi_1 / np.abs(phi_1).max()
phi_1 *= np.sign(phi_1.sum())
alpha = isw.calc_alpha(phi_1, cn[mode], N2, dz)
alpha_ens[tstep, ii] = alpha
c_ens[tstep, ii] = cn[mode]
#mykdv = kdv.KdV(rhotmp,zout)
# Export to an xarray data set
# Create an xray dataset with the output
dims2 = (
'time',
'ensemble',
)
#dims2a = ('time','depth',)
dims3 = ('params', 'time', 'ensemble')
#time = rho.time.values
#time = range(nt)
coords2 = {'time': time, 'ensemble': range(ntrace)}
#coords2a = {'time':time, 'depth':depth[:,0]}
coords3 = {
'time': time,
'ensemble': range(ntrace),
'params': range(nparams)
}
#rho = xr.DataArray(rho.T,
# coords=coords2a,
# dims=dims2a,
# attrs={'long_name':'', 'units':''},
# )
cn_da = xr.DataArray(
c_ens,
coords=coords2,
dims=dims2,
attrs={
'long_name': '',
'units': ''
},
)
alpha_da = xr.DataArray(
alpha_ens,
coords=coords2,
dims=dims2,
attrs={
'long_name': '',
'units': ''
},
)
beta_da = xr.DataArray(
data,
coords=coords3,
dims=dims3,
attrs={
'long_name': '',
'units': ''
},
)
dsout = xr.Dataset({
'cn': cn_da,
'alpha': alpha_da,
'beta': beta_da,
})
return dsout
Fuzz = interp1d(z, Uzz, kind=2, fill_value='extrapolate')
def odefun2(phi, y, c):
phi1, phi2 = phi
dphi1 = phi2
uc2 = (c - Fu(y))**2
#A = Fuzz(x)/(Fu(x) - p[0]) - Fn(x)/uc2
dphi2 = -Fn(y)/uc2 * phi1
return [dphi1, dphi2]
# Guess initial conditions
phi_1 = 0.
mode = 0
phin, cn = iwave_modes(N2, dz)
cguess = cn[mode]
#
#soln = odeint(odefun, [phi_1, phi_2], z, args=(1.59,))
def objective(args):
u2_0, cn = args
dspan = np.linspace(0, d)
U = odeint(odefun2, [phi_1, u2_0], dspan, args=(cn,))
u1 = U[:,0]
u2 = U[:,1]
# Ensure the top bc is zero
return u1[-1]#, u2[-1]
soln = least_squares(objective, [1.05, cguess], xtol=1e-12, \