def eigh_krylov(Afunc, vstart, numiter, numeig):
"""Compute Krylov subspace approximation of eigenvalues and vectors."""
alpha, beta, V = lanczos_iteration(Afunc, vstart, numiter)
# diagonalize Hessenberg matrix
w_hess, u_hess = eigh_tridiagonal(alpha, beta)
# compute Ritz eigenvectors
u_ritz = np.dot(V, u_hess[:, 0:numeig])
return (w_hess[0:numeig], u_ritz)
def GaussLaguerre(self, n, odd=False):
assert n > 0
assert isinstance(n, int)
if odd:
n = 2**n - 1
d = [i for i in range(1, n)]
alpha = [2 * i + 1 for i in range(n)]
#H = np.diag(alpha) + np.diag(d, -1) + np.diag(d, 1)
#[x, v] = np.linalg.eigh(H)
[x, v] = eigh_tridiagonal(alpha, d)
w = v[0, :]**2
rule = {'x': x, 'w': w}
return rule
def Gauss(self, n, pdf, supp, odd=False):
assert n > 0
assert isinstance(n, int)
if odd:
n = 2**n - 1
monic = MonicPoly(n, pdf, supp)
alpha, beta = monic.recurr_coeffs()
#H = np.diag(alpha) + np.diag(np.sqrt(beta), -1) + np.diag(np.sqrt(beta), 1)
#[x, v] = np.linalg.eigh(H)
[x, v] = eigh_tridiagonal(alpha, np.sqrt(beta))
w = v[0, :]**2
rule = {'x': x, 'w': w}
return rule
def convert_chain_to_star(c0,
omega,
t,
get_trafo=False,
force_sp=False,
mp_dps=30,
sort_by=None):
"""
Converts chain coefficients in the form c0, omega, t (system to bath coupling, bath energies,
bath-bath couplings) into the equivalent star geometry coefficients gamma, xi (star system to bath coupling,
star bath energies) by using diagonalization with either arbitrary precision mpmath if the library is installed
or scipy eigh_tridiagonal in float precision
:param c0: System to bath coupling float
:param omega: Bath energies (numpy array)
:param t: Bath-bath couplings (numpy array)
:param get_trafo: If the transformation between the chain and the star should be returned or not
This matrix is only for the omega/t coefficients
:param force_sp: Force the use of the scipy method eigh_tridiagonal, even if mpmath is installed
:param mp_dps: Decimals, which mpmath uses for the computation
:return: gamma (star system to bath coupling), xi (star bath energies),
info dict with the keys: 'trafo': Contains the transformation Matrix between the geometries
"""
assert len(omega) - 1 == len(t)
info = dict()
info['trafo'] = None
if mp is None or force_sp:
w, v = eigh_tridiagonal(omega, t)
gamma = c0 * np.abs(v[0, :])
xi = w
if get_trafo:
info['trafo'] = v
else:
mp.set_dps = mp_dps
nof_coefficients = len(omega)
A = np.zeros((nof_coefficients, nof_coefficients))
drows, dcols = np.diag_indices_from(A)
A[drows[:nof_coefficients], dcols[:nof_coefficients]] = omega
rng = np.arange(nof_coefficients - 1)
A[rng + 1, rng] = t
A[rng, rng + 1] = t
E, Q = mp.eigsy(mp.matrix(A.tolist()))
xi = np.empty(nof_coefficients)
gamma = np.empty(nof_coefficients)
for i in range(A.shape[1]):
xi[i] = float(E[i])
gamma[i] = c0 * np.abs(float(Q[0, i]))
if get_trafo:
Q = np.array(Q.tolist(), dtype=np.float64)
info['trafo'] = Q
gamma, xi = sort_star_coefficients(gamma, xi, sort_by)
return gamma, xi, info
def eigs_fh(lin_map,
d,
n=6,
k=-1,
v0=np.zeros(0),
tol=1e-15,
max_it=1,
which='lm'):
if k < 0:
k = d
else:
k = min(max(k, n + 1), d) # set minimal number of Lanczos basis
if k == d:
max_it = 1 # when constructing the full tri-diagonal matrix, iterate only once
v0 = set_initial_v(v0, d) # handle initial vector
vs = set_random_initial_vectors(v0, k)
alpha = np.zeros((k, ))
beta = np.zeros((k - 1, ))
w = 1
error = np.ones((1, n))
info = dict()
lm = np.zeros((k, 1))
v0_new = np.zeros((d, n))
ind = np.zeros((d, n))
for t in range(0, max_it):
for k_now in range(0, k):
if k_now == 0:
w, alpha[0] = update_w(lin_map, vs[:, [0]])
else:
vs[:, [k_now]], beta[k_now - 1] = update_v(
w, vs[:, :k_now], vs[:, [k_now]])
w, alpha[k_now] = update_w(lin_map, vs[:, [k_now]],
beta[k_now - 1], vs[:, [k_now - 1]])
lm, u = eigh_tridiagonal(
d=alpha, e=beta) # diagonalize the k-by-k tri-diagonal matrix
ind = handle_which(which, lm, n)
v0_new = vs.dot(u[:, ind])
error = np.linalg.norm(x=vs[:, 0:n] - v0_new, axis=0)
error_tot = np.sum(error)
if error_tot > tol:
if t == max_it - 1 and t != 0:
print('Not convergent! error = %g' % error_tot)
info['it_time'] = max_it
else:
vs[:, [0]] = v0_new[:, [0]]
else:
info['it_time'] = t + 1
break
info['error'] = error
return lm[ind], v0_new, info
def numSchrodinger(V, size = size): #V må være en array
H_dia = np.power(hbar, 2) / (m * np.power(Dx, 2)) + V
H_subDia = np.array([-np.power(hbar, 2) / (2 * m * np.power(Dx, 2))] * (size - 1))
# H = diags([H_subDia, H_dia, H_subDia], [-1, 0, 1], shape=(size, size)).toarray() #Trengs kanskje ikke
egenverdi, egenvektor = la.eigh_tridiagonal(H_dia, H_subDia) #Gir oss egenverdiene og egenvektorene til matrisen H
egenverdi *= 1 / sc.eV #Gjør om til elektronvolt
for i in range(len(egenvektor)):
egenvektor_abs_pow = np.power(np.abs(egenvektor[i]), 2)
integral = np.sum(egenvektor_abs_pow) # "Integralet" av absoluttverdien til energiegenfunksjonene i andre
normert_konst = np.sqrt(1 / la.norm(integral)) # Normeringskonstanten
egenvektor[i] *= normert_konst
return egenverdi, egenvektor
def GaussHermite(self, n, odd=False):
assert n > 0
assert isinstance(n, int)
if odd:
n = 2**n - 1
d = np.sqrt(np.arange(n))[1:]
#H = np.diag(d, -1) + np.diag(d, 1)
#[x, v] = np.linalg.eigh(H)
[x, v] = eigh_tridiagonal(np.zeros(n), d)
if odd:
x[int((n - 1) / 2)] = 0.
w = v[0, :]**2
rule = {'x': x, 'w': w}
return rule
def GaussLegendre(self, n, odd = False):
assert n > 0
assert isinstance(n, int)
if odd:
n = 2 ** n - 1
d = np.sqrt([i**2/((2.*i+1.)*(2*i-1.)) for i in range(1,n)])
#H = np.diag(d, -1) + np.diag(d, 1)
#[x, v] = np.linalg.eigh(H)
[x, v] = eigh_tridiagonal(np.zeros(n), d)
w = v[0,:] ** 2
if (x.shape[0]-1) % 2 == 0:
x[(x.shape[0]-1) / 2] = 0.
rule = {'x': x, 'w': w}
return rule
def _regen_dpss(self):
'''
The generated tapers in 2D array [[taper_0], [taper_1], ...] are ordered decreasingly by their respective eigenvalues.
'''
ww = self.nw / self.nn
diag_main = (
(self.nn - 1) / 2 - np.arange(self.nn))**2 * np.cos(2 * np.pi * ww)
diag_off = np.arange(1, self.nn) * np.arange(self.nn - 1, 0, -1) / 2
vecs = eigh_tridiagonal(diag_main,
diag_off,
select='i',
select_range=(self.nn - self.kk,
self.nn - 1))[1]
# polarity follows Slepian convention
return (vecs * np.where(vecs[0, :] > 0, 1, -1)).T[::-1]
def gen_sequences(self, eigenvalue):
'''
generate the discrete prolate spheroidal sequences in the time domain
'''
try:
time_0 = time.time()
self.vecs, self.vals = windows.dpss(self.N, self.N*self.W, self.K, return_ratios=True)
except:
diag_main = ((self.N-1)/2-np.arange(self.N))**2 * np.cos(2*np.pi*self.W)
diag_off = np.arange(1, self.N) * np.arange(self.N-1, 0, -1) / 2
vecs = eigh_tridiagonal(diag_main, diag_off, select='i', select_range=(self.N-self.K,self.N-1))[1]
self.vecs = (vecs * np.where(vecs[0,:]>0, 1, -1)).T[::-1] # normalized energy, polarity follows Slepian convention
if eigenvalue:
A = toeplitz(np.insert( np.sin(2*np.pi*self.W*np.arange(1,self.N))/(np.pi*np.arange(1,self.N)), 0, 2*self.W ))
self.vals = np.diag(self.vecs @ A @ self.vecs.T) # @ is matrix multiplication
def run_test(self, alpha, beta):
n = alpha.shape[0]
# scipy.linalg.eigh_tridiagonal doesn't support complex inputs, so for
# this we call the slower numpy.linalg.eigh.
if np.issubdtype(alpha.dtype, np.complexfloating):
tridiagonal = np.diag(alpha) + np.diag(beta, 1) + np.diag(
np.conj(beta), -1)
eigvals_expected, _ = np.linalg.eigh(tridiagonal)
else:
eigvals_expected = scipy.linalg.eigh_tridiagonal(alpha,
beta,
eigvals_only=True)
eigvals = linalg.eigh_tridiagonal(alpha, beta)
eps = np.finfo(alpha.dtype).eps
atol = 2 * n * eps * np.amax(np.abs(eigvals_expected))
self.assertAllClose(eigvals_expected, eigvals, atol=atol)
def calculate_and_save_results(inputpath, outputpath):
"""
Docstring
"""
rd.load_data(inputpath)
mass = rd.particle_mass()
x_min = rd.x_minimum()
x_max = rd.x_maximum()
npoint = rd.n_point()
fval = rd.first_eigenvalue()
lval = rd.last_eigenvalue()
inter_type = rd.interpolation_type()
x_pot = rd.x_potential()
y_pot = rd.y_potential()
delta = (x_max-x_min)/npoint
const_a = 1/(mass*delta**2)
x_axis = np.linspace(x_min, x_max, npoint)
#interpolating
y_pot_inter = []
if inter_type == "polynomial":
y_pot_inter = KroghInterpolator(x_pot, y_pot)
else:
y_pot_inter = interp1d(x_pot, y_pot, kind=inter_type)
#calculating
expected_x = np.array([])
uncertainties = np.array([])
energies = np.array([])
norm_eigenvecs = np.array([])
maindiag = y_pot_inter(x_axis)+const_a
seconddiag = np.full(npoint-1, -1/2*const_a)
selectrange = (fval-1, lval-1)
ev = eigh_tridiagonal(maindiag, seconddiag, select="i", select_range=selectrange)
(energies, eigenvecs) = ev
norm_eigenvecs = np.array([normalize(eigenvec, delta) for eigenvec in eigenvecs.T])
expected_x = np.array([exp_x(eigenvec, delta, x_axis) for eigenvec in norm_eigenvecs])
uncertaintylist = [uncertainty(eigenvec, delta, x_axis) for eigenvec in norm_eigenvecs]
uncertainties = np.array(uncertaintylist)
rd.save_xyformat(outputpath + "/potential.dat", x_axis, y_pot_inter(x_axis))
rd.save_nxyformat(outputpath + "/wavefuncs.dat", x_axis, norm_eigenvecs.T)
rd.save_xyformat(outputpath + "/energies.dat", energies, ["" for _ in energies])
rd.save_xyformat(outputpath + "/expvalues.dat", expected_x, uncertainties)
def expm_krylov(Afunc, v, dt, numiter):
"""
Compute Krylov subspace approximation of the matrix exponential
applied to input vector: `expm(dt*A)*v`.
Reference:
M. Hochbruck and C. Lubich
On Krylov subspace approximations to the matrix exponential operator
SIAM J. Numer. Anal. 34, 1911 (1997)
"""
alpha, beta, V = lanczos_iteration(Afunc, v, numiter)
# diagonalize Hessenberg matrix
w_hess, u_hess = eigh_tridiagonal(alpha, beta)
return np.dot(
V, np.dot(u_hess,
np.linalg.norm(v) * np.exp(dt * w_hess) * u_hess[0]))
def eigen(vx, vy, xx, inttype, fvalue, lvalue, aa, nn):
"""The function is calculating the eigenvalues and and eigenvectors.
:type vx: [float]
:param vx: x-values of the interpolation points
:type vy: [float]
:param vy: y-values of the interpolation points
:type xx: np.array(float)
:param xx: x-values
:type inttype: string
:param inttype: type of interpolation
:type fvalue: int
:param fvalue: first eigenvalue to be calculated
:type lvalue: int
:param lvalue: last eigenvalue to be calculated
:type aa: float
:param aa: constant for calculation
:type nn: int
:param nn: Number of values to be calculated
:rtype: (np.array(float),np.array(float))
:returns: Returns a tuple with the the eigenvalues and and eigenvectors
"""
if inttype == 'cubic' or inttype == 'linear':
grid = griddata(vx, vy, xx, method=inttype)
elif inttype == 'polynomial':
coefficients = np.polyfit(vx, vy, 2)
grid = np.polyval(coefficients, xx)
dd = grid + aa
ee = np.zeros(nn) + (-0.5) * aa
aa = sclin.eigh_tridiagonal(dd,
ee,
select='i',
select_range=(fvalue, lvalue))
return aa
def get_eig_vectors(self):
"""
:return: v: (M,M) ndarray
The normalized eigenvectors corresponding to the eigenvalues, v[:, i] is corresponding to the w[i].
In each eigenvector v[:, i], v[argmax(abs(v[:, i])), i] >= 0.
w: (M,) ndarray
The eigenvalues in descending order.
"""
w, v = eigh_tridiagonal(self.diagonal, self.off_diagonal, select='a')
sorted_idx = np.argsort(-w)
v = v[:, sorted_idx]
w = w[sorted_idx]
a = np.argmax(np.absolute(v), axis=0)
b = np.array([np.sign(v[a[k], k]) for k in range(len(v))])
v = v * b
return v, w
def eigen_tridiagonal(alpha, beta, maximum=True):
"""Computes eigenvalues of a tridiagonal matrix.
Args:
alpha: vector of diagonal elements
beta: vector of off-diagonal elements
max: whether to compute the max or min magnitude eigenvalue
Returns:
eig: eigenvalue corresponding to max or min magnitude eigenvalue
eig_vector: eigenvalue corresponding to eig
eig_vectors: all eigenvectors
eig_values: all eigenvalues
"""
eig_values, eig_vectors = eigh_tridiagonal(alpha, beta)
if maximum:
ind_eig = np.argmax(np.abs(eig_values))
else:
ind_eig = np.argmin(np.abs(eig_values))
eig = eig_values[ind_eig]
eig_vector = eig_vectors[:, ind_eig]
return eig, eig_vector, eig_vectors, eig_values
def set_eigvalsvecs(self):
"""
Hamiltonian can be represented similar to:
-- --
| 2+V -1 |
| -1 2+V -1 |
| -1 . . |
| . . . |
| . . . |
| . . |
-- --
"""
# Computing eigvals and vecs without the boundaries to avoid singular matrix.
# Array of egeinvectors, 2nd axis specifies which eigenvalue is used
psi = np.zeros((self.N, self.NUM_EIGVALS), dtype=np.complex_)
# numpy.eigh_tridiagonal computes eigvals and eigvecs using a symmetric tridiagonal matrix
la, psi[1:-1] = eigh_tridiagonal(self.d[1:-1], self.e[1:-1], select='i', select_range=(0, self.NUM_EIGVALS-1))
for i in range(self.NUM_EIGVALS):
psi[:,i] = psi[:,i] / np.sqrt(trapz(abs_squared(psi[:,i]), dx=self.dx)) # Normalize each function
return la, psi
def sv_solver(alpha, beta, L, N):
# Input: values of
# N: Number of internal points
# Returns: Eigenvalues and eigenfunctions for the operator
dx = L / (N + 1)
subp = (1 / (dx * dx)) * np.ones(N - 1)
mid = (1 / (dx * dx)) * np.ones(N) * (-2)
eigs, eig_vecs = eigh_tridiagonal(mid, subp)
# Add boundry values to solutions
left_bvs = np.ones(N) * alpha
right_bvs = np.ones(N) * beta
eig_vecs = np.vstack((left_bvs, eig_vecs))
eig_vecs = np.vstack((eig_vecs, right_bvs))
return eigs, eig_vecs
def min_eigvec_Lanczos(A, max_iters=30, tol=1.0e-6, shift=None):
"""Finds approximate minimum eigenvalue/vector using Lanczos method."""
n = A.shape[1]
max_iters = min(max_iters, n - 1)
# Random initialization
Q = np.zeros((n, max_iters + 1))
Q[:, 0] = np.random.randn(n)
Q[:, 0] /= la.norm(Q[:, 0])
# Diagonal and off-diagonal elements
alpha = np.zeros(max_iters)
beta = np.zeros(max_iters)
# Lanczos iteration
matmultcnt = 0
for ii in range(max_iters):
Q[:, ii + 1] = A @ Q[:, ii]
matmultcnt += 1
alpha[ii] = np.vdot(Q[:, ii], Q[:, ii + 1])
if ii == 0:
Q[:, 1] -= alpha[0] * Q[:, 0]
else:
Q[:, ii + 1] -= alpha[ii] * Q[:, ii] + beta[ii - 1] * Q[:, ii - 1]
beta[ii] = la.norm(Q[:, ii + 1])
if abs(beta[ii]) < np.sqrt(n) * np.spacing(1.0):
break
Q[:, ii + 1] /= beta[ii]
# Compute approximate eigenvalues
if ii == 0:
return Q[:, :1], alpha[0], matmultcnt
else:
d, q = la.eigh_tridiagonal(alpha[:ii + 1],
beta[:ii],
select='i',
select_range=(0, 0))
return Q[:, :ii + 1] @ q, d[0], matmultcnt
def solver(potential, mass, x_min, x_max, n_point):
"""
Script to solve the 1-dimensional, stationary schroedinger equation for a given potential. Returns eigenvalues (energielevels) and normalised wavefunctions.
"""
#setting x-axis
xnew = np.linspace(x_min, x_max, n_point)
#delta x
delta = xnew[1] - xnew[0]
#diagonals of coeff matrix, eigenvalues w and eigenvectors v
sub_diagonal = np.array(
[-1 / (2 * mass * delta**2) for i in range(len(xnew) - 1)])
main_diagonal = np.array(
[1 / (mass * delta**2) + potential(i) for i in xnew])
eigenvalue, eigenvector = LA.eigh_tridiagonal(main_diagonal, sub_diagonal)
#normalised wavefunction
wavefunc = np.array(
[eigenvector[:, i] / np.sqrt(delta) for i in range(0, len(xnew))])
return eigenvalue, wavefunc
def __init__(self, m, k0, dx, V0, N, sigma, x_offset=0):
self.m = m
self.k0 = k0
self.dx = dx
self.V0 = V0
self.N = N
self.sigma = sigma
V = [V0] * 4 * N + [
V0 * ((n - N) / (N * 1.0))**2 for n in range(2 * N + 1)
] + [V0] * 4 * N
V = np.asarray(V)
self.V = V
self.Ntot = len(V)
self.d = np.array([(v + hbar**2 / (m * dx**2)) for v in V])
self.e = -hbar**2 / (2 * m * dx**2)
energy, psi_matrix = la.eigh_tridiagonal(
self.d, np.array([self.e] * (self.Ntot - 1)))
self.energy = energy
self.psi_matrix = psi_matrix
self.psi_matrix_complex = psi_matrix * (1.0 + 0.0j)
self.x = np.asarray([dx * n for n in range(self.Ntot)])
self.x0 = self.x[len(self.x) // 2] + x_offset
normfactor = (1.0 + 0.0j) * (2 * np.pi * sigma**2)**(-0.25)
gaussinit = (1.0 + 0.0j) * np.exp(-(self.x - self.x0)**2 /
(4 * sigma**2))
planewavefactor = np.exp(1j * k0 * self.x)
self.Psi0 = normfactor * gaussinit * planewavefactor
self.c = np.zeros(self.Ntot, dtype=np.complex128)
for n in range(self.Ntot):
print("Dotting: {0} out of {1}".format(n + 1, self.Ntot))
self.c[n] = np.vdot(self.psi_matrix_complex[:, n], self.Psi0)
def run_test(self, alpha, beta, eigvals_only=True):
n = alpha.shape[0]
matrix = np.diag(alpha) + np.diag(beta, 1) + np.diag(np.conj(beta), -1)
# scipy.linalg.eigh_tridiagonal doesn't support complex inputs, so for
# this we call the slower numpy.linalg.eigh.
if np.issubdtype(alpha.dtype, np.complexfloating):
eigvals_expected, _ = np.linalg.eigh(matrix)
else:
eigvals_expected = scipy.linalg.eigh_tridiagonal(alpha,
beta,
eigvals_only=True)
eigvals = linalg.eigh_tridiagonal(alpha,
beta,
eigvals_only=eigvals_only)
if not eigvals_only:
eigvals, eigvectors = eigvals
eps = np.finfo(alpha.dtype).eps
atol = n * eps * np.amax(np.abs(eigvals_expected))
self.assertAllClose(eigvals_expected, eigvals, atol=atol)
if not eigvals_only:
self.check_orthogonality(eigvectors, 2 * np.sqrt(n) * eps)
self.check_residual(matrix, eigvals, eigvectors, atol)
def _solve_schrodinger(npoint, eigen, mass, pot):
delta_sq = (abs(pot[1, 0]) - abs(pot[0, 0]))**2
add = 1 / (mass * delta_sq)
# Filling matrix with data to solve with linalg.eigh_tridiagonal()
diag = np.zeros(shape=npoint, dtype=float)
off_diag = np.zeros(shape=npoint - 1, dtype=float)
diag = add + pot[:, 1]
off_diag[:] = -0.5 * add
energies, wavefuncs = linalg.eigh_tridiagonal(diag,
off_diag,
select="i",
select_range=eigen - 1)
wave_new = np.zeros(shape=(npoint, eigen[1] + 1), dtype=float)
wave_new[:, 0] = pot[:, 0]
for mm in range(eigen[1]):
wave_new[:, mm + 1] = wavefuncs[:, mm]
return energies, wave_new
def get_eig_vectors(self):
"""
Calculated the eigenvalues and eigenvectors of B_N matrix.
:return: v: (M,M) ndarray
The normalized eigenvectors corresponding to the eigenvalues, v[:, i] is corresponding to the w[i].
In each eigenvector v[:, i], v[argmax(abs(v[:, i])), i] >= 0.
w: (M,) ndarray
The eigenvalues in descending order.
"""
w, v = eigh_tridiagonal(self.diagonal, self.off_diagonal, select='a')
sorted_idx = np.argsort(-w)
v = v[:, sorted_idx]
w = w[sorted_idx]
# We need to rescale the eigenvectors to fix the sign problem and make consistent with
# the Matlab version.
a = np.argmax(np.absolute(v), axis=0)
b = np.array([np.sign(v[a[k], k]) for k in range(len(v))])
v = v * b
return v, w
def wheeler(m):
"""
David Lignell
Wheeler algorithm for computing weights and abscissas from moments.
From Marchisio and Fox (2013) Computational Models for Polydisperse and Multiphase systems.
input m array of moments (size = 2N)
returns w, x (weights and abscissas)
"""
N2 = len(m)
N = int(N2/2)
sigma = np.zeros((N+1, N2))
a = np.zeros(N)
b = np.zeros(N)
j_diag = np.zeros(N)
j_ldiag = np.zeros(N)
sigma[1,:N2] = m
a[0] = m[1]/m[0]
for k in range(1, N):
l = np.arange(k,N2-k)
sigma[k+1,l] = sigma[k,l+1]-a[k-1]*sigma[k,l]-b[k-1]*sigma[k-1,l]
a[k] = -sigma[k,k]/sigma[k,k-1]+ sigma[k+1,k+1]/sigma[k+1,k]
b[k] = sigma[k+1,k]/sigma[k,k-1]
j_diag = a
j_ldiag = -np.sqrt(np.abs(b[1:]))
x, v = eigh_tridiagonal(j_diag, j_ldiag)
w = v[0,:]**2 * m[0]
return w, x
def lanczos(self, max_num_steps, random_state=None):
assert (self.nhil >= 2)
num_steps = min(max_num_steps, self.nhil)
np.random.seed(random_state)
d = np.zeros(num_steps) # diagonal part of tridiagonal matrix
e = np.zeros(num_steps) # off-diagonal part of tridigaonal matrix
e[0] = 1
v0 = np.zeros(self.nhil)
v1 = np.random.rand(self.nhil)
v1 /= np.linalg.norm(v1)
for j in range(2, num_steps + 1):
d[j - 2] = np.dot(v1, np.dot(self.H, v1))
wj = np.dot(self.H, v1) - d[j - 2] * v1 - e[j - 2] * v0
e[j - 1] = np.linalg.norm(wj)
v0 = v1
v1 = wj / e[j - 1]
# import pdb; pdb.set_trace()
eigvals, eigvecs = eigh_tridiagonal(d, e[1:])
return eigvals
def eigenparam(self):
"""
Compute a vector with the eigenvalues(eV) and another with the
eigenvectors ((Aº)**-1/2) of the quantum hamiltonian with potential
[self.potential [eV]] (each column is an eigenvector with [N]+1
components).
H · phi = E · phi ; H = -(1/2m)(d**2/dx**2) + poten(x)
m := mass / hbar**2 [(eV·Aº**2)**-1]
It solves the 1D time-independent Schrödinger equation for the given
potential (self.potential) inside of a box [a(Aº), b(Aº)], with
[N] intervals.
"""
#Dividing the ab segment in N intervals leave us with a (N+1)x(N+1)
#hamiltonian, where indices 0 and N correspond to the potentials
#barriers. The hamiltonian operator has 3 non-zero diagonals (the main
#diagonal, and the ones next to it), with the following elements.
semi_diag = np.full(self.N, -1. / (2. * self.m * self.deltax**2))
main_diag = self.potential + 1. / (self.m * self.deltax**2)
#Although we keep these walls here, no change seems to happen if we
#remove them (if we don't assign these values)
main_diag[0] += 1000000000 #Potentials barriers
main_diag[self.N] += 1000000000
self.evals, self.evect = spLA.eigh_tridiagonal(main_diag,
semi_diag,
check_finite=False)
#Normalization. Used trapezoids method formula and that sum(evect**2)=1
factors =1/np.sqrt(self.deltax * \
(1. - np.abs(self.evect[0,:])**2/2. - \
np.abs(self.evect[-1,:])**2/2.))
#Normalized vectors (* here multiplies each factor element by each
#evect column)
self.evect = self.evect * factors
def comp_modes(dh, N2, f0=1.0, eivec=False, wmode=False, diag=False):
'''
Compute eigenvalues (and eigenvectors) of the sturm-liouville
equation
d ( f^2 d ) 1
-- ( --- -- psi) + ---- psi = 0
dz ( N^2 dz ) Rd^2
for a given stratification
The eigenvectors correspond to the matrices for the mode/layer
conversion
mod2lay[:,0] is the barotropic mode: should be 1..1
mod2lay[:,i] is the ith baroclinic mode
-To convert from physical to modal:
u_mod = np.dot(lay2mod[:,:],u_lev) # if u_lev is 1D
u_mod = np.einsum('ij,jkl->ikl',lay2mod,u_lev) # if u_lev is 3D
u_mod = np.einsum('ijkl,jkl->ikl',lay2mod,u_lev) #if u_lev is 3D and N2 variable
-To go back to the physical space:
u_lev = np.dot(mod2lay[:,:],u_mod)
u_lev = np.einsum('ij,jkl->ikl',mod2lay,u_mod) # if u_mod is 3D
u_lev = np.einsum('ijkl,jkl->ikl',mod2lay,u_mod) #if u_mod is 3D and N2 variable
the w_modes are related to the p_modes by
w_modes = -1/N2 d p_modes/dz
Parameters
----------
dh : array [nz]
N2 : array [nz (,ny,nx)]
f0 : scalar or array [(ny,nx)]
eivec : Bool
wmode : Bool
diag : Bool
Use transformation matrix to solve a symetric matrix
Returns
-------
if eivec == T
Rd: array [nz (,ny,nx)]
lay2mod: array [nz,nz (,ny,nx)]
mod2lay: array [nz,nz (,ny,nx)]
if eivec == F
Rd: array [nz (,ny,nx)]
'''
N2,f0 = reshape3d(dh,N2,f0)
nl,si_y,si_x = N2.shape
mat_format = "dense"
if diag:
mat_format = "sym_diag"
S = gamma_stretch(dh,N2,f0,wmode=wmode,squeeze=False,mat_format=mat_format)
nlt = (N2 == 0).argmax(axis=0)
nlt = np.where(nlt == 0,nl,nlt)
# put variables in right format
Ht = np.cumsum(dh)
# Ht = np.sum(dh)
dhi = 0.5*(dh[1:] + dh[:-1])
dhcol = dh[:,None]
dhicol = dhi[:,None]
if wmode:
Rd = np.zeros((nl,si_y,si_x))
if eivec:
mod2lay = np.zeros((nl,nl,si_y,si_x))
lay2mod = np.zeros((nl,nl,si_y,si_x))
else:
nlt = nlt + 1
Rd = np.zeros((nl+1,si_y,si_x))
if eivec:
mod2lay = np.zeros((nl+1,nl+1,si_y,si_x))
lay2mod = np.zeros((nl+1,nl+1,si_y,si_x))
for j,i in np.ndindex((si_y,si_x)):
if eivec:
if diag:
iRd2, eigs = la.eigh_tridiagonal(S[1,:nlt[j,i],j,i], S[0,1:nlt[j,i],j,i])
eigr = S[2,:nlt[j,i],j,i,None]*eigs # D*w
eigl = eigs/S[2,:nlt[j,i],j,i,None] # w*D^-1 if eigenvectors are stored in lines but eigl is eigl.T so we do D^-1*w
else:
iRd2, eigl,eigr= la.eig(S[:nlt[j,i],:nlt[j,i],j,i],left=True)
else:
if diag:
iRd2 = la.eigvalsh_tridiagonal(S[1,:nlt[j,i],j,i], S[0,1:nlt[j,i],j,i])
else:
iRd2 = la.eig(S[:nlt[j,i],:nlt[j,i],j,i],right=False)
iRd2 = -iRd2.real
idx = np.argsort(iRd2)
iRd2 = iRd2[idx]
with np.errstate(divide='ignore', invalid='ignore'):
Rd_loc = 1./np.sqrt(iRd2)
Rd[:nlt[j,i],j,i] = Rd_loc
if eivec:
eigl = eigl[:,idx]
eigr = eigr[:,idx]
# Normalize eigenvectors
N2col = N2[:nlt[j,i],j,i][:,None]
cm = Rd_loc[:nlt[j,i],None]*f0[j,i]
if wmode:
scap = np.sum(dhi[:nlt[j,i],None]*eigr*eigr*N2col*cm.T**2,0)
Htt = Ht[nlt[j,i]]
else:
scap = np.sum(dh[:nlt[j,i],None]*eigr*eigr,0)
Htt = Ht[nlt[j,i]-1]
flip = np.sign(eigr[0,:])
eigr = eigr*np.sqrt(Htt/scap)*flip
# # scalar product
# if wmode:
# check = np.sum(N2col.T*eigr[:,1]*eigr[:,1]*dhicol.T*(Rd_loc[1]*f0[j,i])**2)
# else:
# check = np.sum(dhcol.T*eigr[:,2]*eigr[:,2])/Ht
if diag:
eigl = eigl/np.sqrt(Htt/scap)*flip
else:
scap2 = np.sum(eigl*eigr,0)
eigl = eigl/scap2
lay2mod[:nlt[j,i],:nlt[j,i],j,i] = eigl.T
mod2lay[:nlt[j,i],:nlt[j,i],j,i] = eigr
if eivec:
return Rd.squeeze(), lay2mod.squeeze(), mod2lay.squeeze()
else:
return Rd.squeeze()
def _scheme_from_rc_numpy(alpha, beta):
alpha = alpha.astype(numpy.float64)
beta = beta.astype(numpy.float64)
x, V = eigh_tridiagonal(alpha, numpy.sqrt(beta[1:]))
w = beta[0] * V[0, :]**2
return x, w
def run_eigh_tridiagonal(self, alpha, beta, **kwargs):
return linalg.eigh_tridiagonal(alpha, beta, **kwargs)
