def _other_dpss_method(N, NW, Kmax):
"""Returns the Discrete Prolate Spheroidal Sequences of orders [0,Kmax-1]
for a given frequency-spacing multiple NW and sequence length N.
See dpss function that is the official version. This version is indepedant
of the C code and relies on Scipy function. However, it is slower by a factor 3
Tridiagonal form of DPSS calculation from:
"""
# here we want to set up an optimization problem to find a sequence
# whose energy is maximally concentrated within band [-W,W].
# Thus, the measure lambda(T,W) is the ratio between the energy within
# that band, and the total energy. This leads to the eigen-system
# (A - (l1)I)v = 0, where the eigenvector corresponding to the largest
# eigenvalue is the sequence with maximally concentrated energy. The
# collection of eigenvectors of this system are called Slepian sequences,
# or discrete prolate spheroidal sequences (DPSS). Only the first K,
# K = 2NW/dt orders of DPSS will exhibit good spectral concentration
# [see http://en.wikipedia.org/wiki/Spectral_concentration_problem]
# Here I set up an alternative symmetric tri-diagonal eigenvalue problem
# such that
# (B - (l2)I)v = 0, and v are our DPSS (but eigenvalues l2 != l1)
# the main diagonal = ([N-1-2*t]/2)**2 cos(2PIW), t=[0,1,2,...,N-1]
# and the first off-diangonal = t(N-t)/2, t=[1,2,...,N-1]
# [see Percival and Walden, 1993]
from scipy import linalg as la
Kmax = int(Kmax)
W = float(NW)/N
ab = np.zeros((2,N), 'd')
nidx = np.arange(N)
ab[0,1:] = nidx[1:]*(N-nidx[1:])/2.
ab[1] = ((N-1-2*nidx)/2.)**2 * np.cos(2*np.pi*W)
# only calculate the highest Kmax-1 eigenvectors
l,v = la.eig_banded(ab, select='i', select_range=(N-Kmax, N-1))
dpss = v.transpose()[::-1]
# By convention (Percival and Walden, 1993 pg 379)
# * symmetric tapers (k=0,2,4,...) should have a positive average.
# * antisymmetric tapers should begin with a positive lobe
fix_symmetric = (dpss[0::2].sum(axis=1) < 0)
for i, f in enumerate(fix_symmetric):
if f:
dpss[2*i] *= -1
fix_skew = (dpss[1::2,1] < 0)
for i, f in enumerate(fix_skew):
if f:
dpss[2*i+1] *= -1
# Now find the eigenvalues of the original
# Use the autocovariance sequence technique from Percival and Walden, 1993
# pg 390
# XXX : why debias false? it's all messed up o.w., even with means
# on the order of 1e-2
acvs = _autocov(dpss, debias=False) * N
r = 4*W*np.sinc(2*W*nidx)
r[0] = 2*W
eigvals = np.dot(acvs, r)
return dpss, eigvals
def solve(self, n_eigvals=20):
print '[S1d_H] start solving'
t_start = time.clock()
self.eigvals, self.eigvects = eig_banded(self.Hbanded, select='i',
select_range=[0, n_eigvals])
t_end = time.clock()
print '[S1d_H] end solving (%.2f sec)' % (t_end - t_start)
def lanczos_tridiag_eig(alpha, beta, check_finite=True):
"""Find the eigen-values and -vectors of the Lanczos triadiagonal matrix.
Parameters
----------
alpha : array of float
The diagonal.
beta : array of float
The k={-1, 1} off-diagonal. Only first ``len(alpha) - 1`` entries used.
"""
Tk_banded = np.empty((2, alpha.size))
Tk_banded[1, -1] = 0.0 # sometimes can get nan here? -> breaks eig_banded
Tk_banded[0, :] = alpha
Tk_banded[1, :beta.size] = beta
try:
tl, tv = scla.eig_banded(Tk_banded,
lower=True,
check_finite=check_finite)
# sometimes get no convergence -> use dense hermitian method
except scla.LinAlgError: # pragma: no cover
tl, tv = np.linalg.eigh(np.diag(alpha) +
np.diag(beta[:alpha.size - 1], -1),
UPLO='L')
return tl, tv
def _other_dpss_method(N, NW, Kmax):
"""Returns the Discrete Prolate Spheroidal Sequences of orders [0,Kmax-1]
for a given frequency-spacing multiple NW and sequence length N.
See dpss function that is the official version. This version is indepedant
of the C code and relies on Scipy function. However, it is slower by a factor 3
Tridiagonal form of DPSS calculation from:
"""
# here we want to set up an optimization problem to find a sequence
# whose energy is maximally concentrated within band [-W,W].
# Thus, the measure lambda(T,W) is the ratio between the energy within
# that band, and the total energy. This leads to the eigen-system
# (A - (l1)I)v = 0, where the eigenvector corresponding to the largest
# eigenvalue is the sequence with maximally concentrated energy. The
# collection of eigenvectors of this system are called Slepian sequences,
# or discrete prolate spheroidal sequences (DPSS). Only the first K,
# K = 2NW/dt orders of DPSS will exhibit good spectral concentration
# [see http://en.wikipedia.org/wiki/Spectral_concentration_problem]
# Here I set up an alternative symmetric tri-diagonal eigenvalue problem
# such that
# (B - (l2)I)v = 0, and v are our DPSS (but eigenvalues l2 != l1)
# the main diagonal = ([N-1-2*t]/2)**2 cos(2PIW), t=[0,1,2,...,N-1]
# and the first off-diangonal = t(N-t)/2, t=[1,2,...,N-1]
# [see Percival and Walden, 1993]
from scipy import linalg as la
Kmax = int(Kmax)
W = float(NW)/N
ab = np.zeros((2,N), 'd')
nidx = np.arange(N)
ab[0,1:] = nidx[1:]*(N-nidx[1:])/2.
ab[1] = ((N-1-2*nidx)/2.)**2 * np.cos(2*np.pi*W)
# only calculate the highest Kmax-1 eigenvectors
l,v = la.eig_banded(ab, select='i', select_range=(N-Kmax, N-1))
dpss = v.transpose()[::-1]
# By convention (Percival and Walden, 1993 pg 379)
# * symmetric tapers (k=0,2,4,...) should have a positive average.
# * antisymmetric tapers should begin with a positive lobe
fix_symmetric = (dpss[0::2].sum(axis=1) < 0)
for i, f in enumerate(fix_symmetric):
if f:
dpss[2*i] *= -1
fix_skew = (dpss[1::2,1] < 0)
for i, f in enumerate(fix_skew):
if f:
dpss[2*i+1] *= -1
# Now find the eigenvalues of the original
# Use the autocovariance sequence technique from Percival and Walden, 1993
# pg 390
# XXX : why debias false? it's all messed up o.w., even with means
# on the order of 1e-2
acvs = _autocov(dpss, debias=False) * N
r = 4*W*np.sinc(2*W*nidx)
r[0] = 2*W
eigvals = np.dot(acvs, r)
return dpss, eigvals
def groundstate(pm, H):
r"""Calculates the ground-state of the system for a given potential.
.. math::
\hat{H} \phi_{j} = \varepsilon_{j} \phi_{j}
parameters
----------
pm : object
Parameters object
H : array_like
2D array of the Hamiltonian matrix in band form, indexed as H[band,space_index]
returns
-------
density : array_like
1D array of the electron density, indexed as density[space_index]
orbitals : array_like
2D array of the Kohn-Sham orbitals, index as orbitals[space_index,orbital_number]
eigenvalues : array_like
1D array of the Kohn-Sham eigenvalues, indexed as eigenvalues[orbital_number]
"""
# Solve the Kohn-Sham equations
eigenvalues, orbitals = spla.eig_banded(H, lower=True)
# Normalise the orbitals
orbitals /= np.sqrt(pm.space.delta)
# Calculate the electron density
density = electron_density(pm, orbitals)
return density, orbitals, eigenvalues
def non_approx(pm):
r"""Calculates the two lowest non-interacting eigenstates of the system. These can then be
expressed in Slater determinant form as an approximation to the exact many-electron wavefunction.
.. math::
\bigg(-\frac{1}{2} \frac{d^{2}}{dx^{2}} + V_{\mathrm{ext}}(x) \bigg) \phi_{j}(x)
= \varepsilon_{j} \phi_{j}(x)
parameters
----------
pm : object
Parameters object
returns array_like and array_like
1D array of the 1st non-interacting eigenstate, indexed as eigenstate_1[space_index].
1D array of the 2nd non-interacting eigenstate, indexed as eigenstate_2[space_index].
"""
# Construct the single-electron Hamiltonian
K = NON.construct_K(pm)
H = copy.copy(K)
H[0,:] += pm.space.v_ext[:]
# Solve the single-electron TISE
eigenvalues, eigenfunctions = spla.eig_banded(H, lower=True, select='i', select_range=(0,1))
# Take the two lowest eigenstates
eigenstate_1 = eigenfunctions[:,0]
eigenstate_2 = eigenfunctions[:,1]
return eigenstate_1, eigenstate_2
def gauss_legendre(n):
# http://www.scientificpython.net/1/post/2012/04/gausslegendre1.html
k=np.arange(1.0,n)
a_band = np.zeros((2,n))
a_band[1,0:n-1] = k/np.sqrt(4*k*k-1)
x,V = linalg.eig_banded(a_band,lower=True)
w=2*np.real(np.power(V[0,:],2))
return x, w
def computeEigenFunctions(self,jRange=None): if jRange: sel='i' else: sel ='a' #sel = 'a' lambdas,basis = eig_banded(self.bandedMatrix,lower=True, check_finite=False, overwrite_a_band=True, select=sel, select_range = jRange) basis/=np.sqrt(self.dx) #For normalization return lambdas,basis
def gauss_legendre(n):
# http://www.scientificpython.net/1/post/2012/04/gausslegendre1.html
k = np.arange(1.0, n)
a_band = np.zeros((2, n))
a_band[1, 0:n - 1] = k / np.sqrt(4 * k * k - 1)
x, V = linalg.eig_banded(a_band, lower=True)
w = 2 * np.real(np.power(V[0, :], 2))
return x, w
def computeEigenFunctions(self, jRange=None):
if jRange: sel = 'i'
else: sel = 'a'
#sel = 'a'
lambdas, basis = eig_banded(self.bandedMatrix,
lower=True,
check_finite=False,
overwrite_a_band=True,
select=sel,
select_range=jRange)
basis /= np.sqrt(self.dx) #For normalization
return lambdas, basis
def diag_banded(A, n=2):
"""A is in upper banded form.
Returns the smallest n eigenvalue and its corresponding eigenvector.
"""
try:
from scipy.linalg import eig_banded
info("Import of scipy successful", verbosity.high)
except ImportError:
raise ValueError(" ")
d = eig_banded(A, select='i', select_range=(0, n), eigvals_only=True, check_finite=False)
return d
def get_evals_evecs(self):
"""
get bound states
"""
H = self.create_H_banded()
evals, evecs = eig_banded(H,
lower=True,
select="v",
select_range=(self.pot.min(),
self.pot.max()))
# fix the eigenvector norm to the position space integral
evecs /= np.sqrt(self.dx)
return evals, evecs
def compute_spectrum(hess_vec_prod, size, method="exact", max_steps=10):
"""Returns an array containing sorted eigenvalues and eigenvectors.
hess_vec_prod takes a vector only"""
# The slow, N^3 way.
if method == "exact":
hessian = matrix_from_mvp(hess_vec_prod, size)
return np.linalg.eigvals(hessian)
elif method == "approx":
a, b = lanczos_iteration(hess_vec_prod, size, max_steps)
banded = np.concatenate((b[None, :], a[None, :]), axis=0)
eigvals, eigvects = eig_banded(banded)
return subset_to_full(eigvals, eigvects[0, :]**2, size)
else:
raise ValueError("Unrecognized method {0}".format(method))
def h_step(self, H0, H1):
r"""Performs one minimisation step
parameters
----------
H0: array_like
input Hamiltonian to be mixed (banded form)
H1: array_like
output Hamiltonian to be mixed (banded form)
returns
-------
H: array_like
mixed hamiltonian (banded form)
"""
## TODO: implement analytic derivative
#dE_dtheta_0 = 2 * np.sum(self.braket(dirs, H, wfs).real)
# For the moment, we simply fit the parabola through 3 points
npt = 3
lambdas = np.linspace(0, 1, npt)
energies = np.zeros(npt)
# self-consistent variant
for i in range(npt):
l = lambdas[i]
Hl = (1 - l) * H0 + l * H1
enl, wfsl = spla.eig_banded(Hl, lower=True)
energies[i] = self.total_energy(wfsl)
# improve numerical stability
energies -= energies[0]
p = np.polyfit(lambdas, energies, 2)
a, b, c = p
l = -b / (2 * a)
# don't want step to get too large
l = np.minimum(1, l)
#import matplotlib.pyplot as plt
#fig, axs = plt.subplots(1,1)
#plt.plot(lambdas, energies)
##axs.set_ylim(np.min(total_energies), np.max(total_energies))
#plt.show()
Hl = (1 - l) * H0 + l * H1
return Hl
def RealWalk(t=7):
'''Generate the unitary corresponding to our quantum walk evolved to a time t
t : evolution time
returns : 21x21 (unitary) array'''
offdiag=[0.466,0.54,0.504,0.515,0.494,0.451,0.508,0.511,0.483,0.505,0.547,\
0.471,0.493,0.499,0.447,0.467,0.544,0.524,0.51,0.503]
h=numpy.zeros((2,21))
for i in range(21):
h[1,i]=2.122
for i in range(20):
h[0,i+1]=offdiag[i]
evals,evecs=linalg.eig_banded(h)
diag=numpy.zeros((21,21),dtype=complex)
for i,z in enumerate(numpy.exp(-1j*evals*t)):
diag[i,i]=z
return numpy.dot(evecs, numpy.dot(diag, evecs.T.conjugate()))
def solve_eigen_problem(self, A, B, solver):
"""Solve the eigen problem"""
N = A.testfunction[0].N
s = A.testfunction[0].slice()
self.V = np.zeros((N, N))
self.lmbda = np.ones(N)
if solver == 0:
self.lmbda[s], self.V[s, s] = scipy_la.eigh(A.diags().toarray(),
B.diags().toarray())
elif solver == 1:
#self.lmbda[s], self.V[s, s] = scipy_la.eigh(B.diags().toarray())
a = np.zeros((3, N-2))
a[0, :] = B[0]
a[2, :-2] = B[2]
self.lmbda[s], self.V[s, s] = scipy_la.eig_banded(a, lower=True)
def interaction_matrix_div_E0_squared(Jmax,K,M,KMsign,delta_alpha,alpha_perp,diagonalize=True):
''' The interaction matrix <JKM|V|J'KM> divided by E_0^2
(the electric field amplitude squared).
E_0 is just a constant, so the same interaction matrix
can be used for all laser pulses by just scaling this one
with E_0^2 = 2*I/(c*epsilon_0) times a conversion factor
between polarizability volume and polarizability.
Input is the Jmax, the maximum J state,
K and M for the state, along with KMsign, the relative sign between K and M,
and delta_alpha and alpha_perp, the components of the polarizability
tensor.
'''
try:
cache = interaction_matrix_div_E0_squared.cache;
except AttributeError:
cache = dict();
interaction_matrix_div_E0_squared.cache = cache;
try:
return cache[(Jmax,K,M,KMsign,delta_alpha,alpha_perp,diagonalize)];
except KeyError:
pass;
(U, U0, U1, U2) = MeanCos2Matrix(Jmax,K,M,KMsign);
Jmin = max(K,M);
tr = numpy.ones(Jmax+1); # Construct identity matrix,
tr[0:Jmin] = 0; # but with 0...Jmin zeroed
V0 = -(delta_alpha*U0+alpha_perp*tr)/4;
V1 = -delta_alpha*U1/4;
V2 = -delta_alpha*U2/4;
if (diagonalize):
cat = numpy.concatenate;
a = numpy.vstack((cat(([0,0],V2)),cat(([0],V1)),V0));
eig,vec = linalg.eig_banded(a);
vec = numpy.ascontiguousarray(vec);
V = (V0, V1, V2, eig, vec);
else:
V = (V0, V1, V2);
cache[(Jmax,K,M,KMsign,delta_alpha,alpha_perp,diagonalize)] = V;
return V;
def eigen(self, eigvals_only=False, overwrite_a_band=False, select='a',
select_range=None, max_ev=0):
"""
Solve real symmetric or complex hermitian band matrix
eigenvalue problem.
Uses scipy.linalg.eig_banded function.
"""
from scipy.linalg import eig_banded
w, v = eig_banded(self.ab, lower=self.lower,
eigvals_only=eigvals_only,
overwrite_a_band=overwrite_a_band,
select=select,
select_range=select_range,
max_ev=max_ev)
return EigendecompositionOperator(w=w, v=v)
def test_eig_banded(self):
"""Compare eigenvalues and eigenvectors of eig_banded
with those of linalg.eig. """
w_sym, evec_sym = eig_banded(self.bandmat_sym)
evec_sym_ = evec_sym[:, argsort(w_sym.real)]
assert_array_almost_equal(sort(w_sym), self.w_sym_lin)
assert_array_almost_equal(abs(evec_sym_), abs(self.evec_sym_lin))
w_herm, evec_herm = eig_banded(self.bandmat_herm)
evec_herm_ = evec_herm[:, argsort(w_herm.real)]
assert_array_almost_equal(sort(w_herm), self.w_herm_lin)
assert_array_almost_equal(abs(evec_herm_), abs(self.evec_herm_lin))
# extracting eigenvalues with respect to an index range
ind1 = 2
ind2 = 6
w_sym_ind, evec_sym_ind = eig_banded(self.bandmat_sym,
select='i',
select_range=(ind1, ind2))
assert_array_almost_equal(sort(w_sym_ind),
self.w_sym_lin[ind1:ind2 + 1])
assert_array_almost_equal(abs(evec_sym_ind),
abs(self.evec_sym_lin[:, ind1:ind2 + 1]))
w_herm_ind, evec_herm_ind = eig_banded(self.bandmat_herm,
select='i',
select_range=(ind1, ind2))
assert_array_almost_equal(sort(w_herm_ind),
self.w_herm_lin[ind1:ind2 + 1])
assert_array_almost_equal(abs(evec_herm_ind),
abs(self.evec_herm_lin[:, ind1:ind2 + 1]))
# extracting eigenvalues with respect to a value range
v_lower = self.w_sym_lin[ind1] - 1.0e-5
v_upper = self.w_sym_lin[ind2] + 1.0e-5
w_sym_val, evec_sym_val = eig_banded(self.bandmat_sym,
select='v',
select_range=(v_lower, v_upper))
assert_array_almost_equal(sort(w_sym_val),
self.w_sym_lin[ind1:ind2 + 1])
assert_array_almost_equal(abs(evec_sym_val),
abs(self.evec_sym_lin[:, ind1:ind2 + 1]))
v_lower = self.w_herm_lin[ind1] - 1.0e-5
v_upper = self.w_herm_lin[ind2] + 1.0e-5
w_herm_val, evec_herm_val = eig_banded(self.bandmat_herm,
select='v',
select_range=(v_lower, v_upper))
assert_array_almost_equal(sort(w_herm_val),
self.w_herm_lin[ind1:ind2 + 1])
assert_array_almost_equal(abs(evec_herm_val),
abs(self.evec_herm_lin[:, ind1:ind2 + 1]))
def _gauss_nodes_weights(a,b):
r"""Calculate the nodes and weights for given
recursion coefficients assuming a normalized
weights functions.
see Walter Gautschi, Algorithm 726: ORTHPOL;
a Package of Routines for Generating Orthogonal
Polynomials and Gauss-type Quadrature Rules, 1994
"""
assert len(a) == len(b)
a_band = np.vstack((np.sqrt(b),a))
w, v = eig_banded(a_band)
nodes = w # eigenvalues
weights = b[0] * v[0,:]**2 # first component of each eigenvector
# the prefactor b[0] from the original paper
# accounts for the weights of unnormalized weight functions
return nodes, weights
def _gauss_nodes_weights(a,b):
r"""Calculate the nodes and weights for given
recursion coefficients assuming a normalized
weights functions.
see Walter Gautschi, Algorithm 726: ORTHPOL;
a Package of Routines for Generating Orthogonal
Polynomials and Gauss-type Quadrature Rules, 1994
"""
assert len(a) == len(b)
a_band = np.vstack((np.sqrt(b),a))
w, v = eig_banded(a_band)
nodes = w # eigenvalues
weights = b[0] * v[0,:]**2 # first component of each eigenvector
# the prefactor b[0] from the original paper
# accounts for the weights of unnormalized weight functions
return nodes, weights
def test_eig_banded(self):
"""Compare eigenvalues and eigenvectors of eig_banded
with those of linalg.eig. """
w_sym, evec_sym = eig_banded(self.bandmat_sym)
evec_sym_ = evec_sym[:,argsort(w_sym.real)]
assert_array_almost_equal(sort(w_sym), self.w_sym_lin)
assert_array_almost_equal(abs(evec_sym_), abs(self.evec_sym_lin))
w_herm, evec_herm = eig_banded(self.bandmat_herm)
evec_herm_ = evec_herm[:,argsort(w_herm.real)]
assert_array_almost_equal(sort(w_herm), self.w_herm_lin)
assert_array_almost_equal(abs(evec_herm_), abs(self.evec_herm_lin))
# extracting eigenvalues with respect to an index range
ind1 = 2
ind2 = 6
w_sym_ind, evec_sym_ind = eig_banded(self.bandmat_sym,
select='i', select_range=(ind1, ind2) )
assert_array_almost_equal(sort(w_sym_ind),
self.w_sym_lin[ind1:ind2+1])
assert_array_almost_equal(abs(evec_sym_ind),
abs(self.evec_sym_lin[:,ind1:ind2+1]) )
w_herm_ind, evec_herm_ind = eig_banded(self.bandmat_herm,
select='i', select_range=(ind1, ind2) )
assert_array_almost_equal(sort(w_herm_ind),
self.w_herm_lin[ind1:ind2+1])
assert_array_almost_equal(abs(evec_herm_ind),
abs(self.evec_herm_lin[:,ind1:ind2+1]) )
# extracting eigenvalues with respect to a value range
v_lower = self.w_sym_lin[ind1] - 1.0e-5
v_upper = self.w_sym_lin[ind2] + 1.0e-5
w_sym_val, evec_sym_val = eig_banded(self.bandmat_sym,
select='v', select_range=(v_lower, v_upper) )
assert_array_almost_equal(sort(w_sym_val),
self.w_sym_lin[ind1:ind2+1])
assert_array_almost_equal(abs(evec_sym_val),
abs(self.evec_sym_lin[:,ind1:ind2+1]) )
v_lower = self.w_herm_lin[ind1] - 1.0e-5
v_upper = self.w_herm_lin[ind2] + 1.0e-5
w_herm_val, evec_herm_val = eig_banded(self.bandmat_herm,
select='v', select_range=(v_lower, v_upper) )
assert_array_almost_equal(sort(w_herm_val),
self.w_herm_lin[ind1:ind2+1])
assert_array_almost_equal(abs(evec_herm_val),
abs(self.evec_herm_lin[:,ind1:ind2+1]) )
def nnls_solveh_banded(S, B, X0, niter=50):
# Number of timebins and units.
nt = S.shape[1]
nn = B.shape[1]
# On first iteration, warm-start from projected least-squares.
if X0 is None:
X0 = np.maximum(0, solveh_banded(S, B))
# Determine Lipshitz constant
w = eig_banded(S,
lower=False,
select='i',
select_range=(nt - 1, nt - 1),
eigvals_only=True)[0]
# Run projected gradient descent in parallel across neurons.
_parallel_proj_grad(S, B, X0, 1 / w, niter)
return X0
def svd_eye_minus_hat_matrix(self):
"""
Following Craven & Wahba find the singular value decomposition of
F = D Q R^{-1/2}
This function returns the non-zero singular values S and the
corresponding left singular vectors in U, satisfying
I - Hp = D U [ si**2 / ( 6(1-p) si**2 + p ) ] U.T D**(-1)
where si is the ith singular value.
"""
# TODO: is it indeed faster to take the non-sparse inverse?!
method = 4
if method == 0:
sqrt_invR = sqrtm(spla.inv(self.R).A)
elif method == 1:
sqrt_invR = sqrtm(la.inv(self.R.todense()))
elif method == 2:
invR = la.inv(self.R.todense())
eR, oR = la.eigh(invR)
sqrt_invR = oR.dot(np.diag(np.sqrt(eR))).dot(oR.T)
elif method == 3:
eR, oR = la.eigh(self.R.todense())
sqrt_invR = oR.dot(np.diag(1. / np.sqrt(eR))).dot(oR.T)
elif method == 4:
# TODO:
# deal with the error
# File "splines.py", line 378, in svd_eye_minus_hat_matrix
# eR, oR = la.eig_banded(self.R.data[self.R.offsets>=0][::-1])
# File "/ph2users/eldada/lib/anaconda/lib/python2.7/site-packages/scipy/linalg/decomp.py",
# line 563, in eig_banded
# raise LinAlgError("eig algorithm did not converge")
try:
eR, oR = la.eig_banded(self.R.data[self.R.offsets >= 0][::-1])
sqrt_invR = oR.dot(np.diag(1. / np.sqrt(eR))).dot(oR.T)
except LinAlgError:
# if eig_banded fails try the eigh
eR, oR = la.eigh(self.R.todense())
sqrt_invR = oR.dot(np.diag(1. / np.sqrt(eR))).dot(oR.T)
U, S, VT = la.svd(self.D * self.Q * sqrt_invR, full_matrices=False)
return U, S
def _gauss_from_coefficients_numpy(alpha, beta):
# assert isinstance(alpha, numpy.ndarray)
# assert isinstance(beta, numpy.ndarray)
alpha = alpha.astype(numpy.float64)
beta = beta.astype(numpy.float64)
x, V = eig_banded(numpy.vstack((numpy.sqrt(beta), alpha)), lower=False)
w = beta[0] * scipy.real(scipy.power(V[0, :], 2))
# eigh_tridiagonal is only available from scipy 1.0.0, and has problems
# with precision. TODO find out how/why/what
# try:
# from scipy.linalg import eigh_tridiagonal
# except ImportError:
# # Use eig_banded
# x, V = \
# eig_banded(numpy.vstack((numpy.sqrt(beta), alpha)), lower=False)
# w = beta[0]*scipy.real(scipy.power(V[0, :], 2))
# else:
# x, V = eigh_tridiagonal(alpha, numpy.sqrt(beta[1:]))
# w = beta[0] * V[0, :]**2
return x, w
def eigen(self,
eigvals_only=False,
overwrite_a_band=False,
select='a',
select_range=None,
max_ev=0):
"""
Solve real symmetric or complex hermitian band matrix
eigenvalue problem.
Uses scipy.linalg.eig_banded function.
"""
from scipy.linalg import eig_banded
w, v = eig_banded(self.ab,
lower=self.lower,
eigvals_only=eigvals_only,
overwrite_a_band=overwrite_a_band,
select=select,
select_range=select_range,
max_ev=max_ev)
return EigendecompositionOperator(w=w, v=v)
def svd_eye_minus_hat_matrix(self):
"""
Following Craven & Wahba find the singular value decomposition of
F = D Q R^{-1/2}
This function returns the non-zero singular values S and the
corresponding left singular vectors in U, satisfying
I - Hp = D U [ si**2 / ( 6(1-p) si**2 + p ) ] U.T D**(-1)
where si is the ith singular value.
"""
# TODO: is it indeed faster to take the non-sparse inverse?!
method = 4
if method==0:
sqrt_invR = sqrtm(spla.inv(self.R).A)
elif method==1:
sqrt_invR = sqrtm(la.inv(self.R.todense()))
elif method==2:
invR = la.inv(self.R.todense())
eR, oR = la.eigh(invR)
sqrt_invR = oR.dot(np.diag(np.sqrt(eR))).dot(oR.T)
elif method==3:
eR, oR = la.eigh(self.R.todense())
sqrt_invR = oR.dot(np.diag(1./np.sqrt(eR))).dot(oR.T)
elif method==4:
# TODO:
# deal with the error
# File "splines.py", line 378, in svd_eye_minus_hat_matrix
# eR, oR = la.eig_banded(self.R.data[self.R.offsets>=0][::-1])
# File "/ph2users/eldada/lib/anaconda/lib/python2.7/site-packages/scipy/linalg/decomp.py",
# line 563, in eig_banded
# raise LinAlgError("eig algorithm did not converge")
try:
eR, oR = la.eig_banded(self.R.data[self.R.offsets>=0][::-1])
sqrt_invR = oR.dot(np.diag(1./np.sqrt(eR))).dot(oR.T)
except LinAlgError:
# if eig_banded fails try the eigh
eR, oR = la.eigh(self.R.todense())
sqrt_invR = oR.dot(np.diag(1./np.sqrt(eR))).dot(oR.T)
U, S, VT = la.svd(self.D * self.Q * sqrt_invR, full_matrices=False)
return U,S
def solve_gsks_equations(pm, hamiltonian):
r"""Solves the ground-state Kohn-Sham equations to find the ground-state
Kohn-Sham eigenfunctions, energies and electron density.
.. math::
\hat{H}\phi_{j}(x) = \varepsilon_{j}\phi_{j}(x) \\
n_{\mathrm{KS}}(x) = \sum_{j=1}^{N}|\phi_{j}(x)|^{2}
parameters
----------
pm : object
Parameters object
hamiltonian : array_like
2D array of the Hamiltonian matrix, index as
K[band,space_index]
returns array_like and array_like and array_like
2D array of the ground-state Kohn-Sham eigenfunctions, indexed as
wavefunctions_ks[space_index,eigenfunction]. 1D array containing the
ground-state Kohn-Sham eigenenergies, indexed as
energies_ks[eigenenergies]. 1D array of the ground-state Kohn-Sham
electron density, indexed as density_ks[space_index].
"""
# Solve the Kohn-Sham equations
energies_ks, wavefunctions_ks = spla.eig_banded(hamiltonian, lower=True)
# Normalise the wavefunctions
for j in range(pm.space.npt):
norm = np.linalg.norm(wavefunctions_ks[:, j]) * pm.space.delta**0.5
wavefunctions_ks[:, j] /= norm
# Calculate the electron density
density_ks = np.sum(wavefunctions_ks[:, :pm.sys.NE]**2,
axis=1,
dtype=np.float)
return wavefunctions_ks, energies_ks, density_ks
def tensquad(order, dist):
"""Computes the tensor product quadrature rule with recurrence coefficients"""
if not isinstance(dist, Joint): dist = Joint(dist)
nbrPts = order + 1
dim = dist[:].shape[0]
J = np.zeros((2, nbrPts))
points = np.zeros((nbrPts, dim))
weights = np.zeros((nbrPts, dim))
# Integration points and weights of each variable
for i in range(dim):
coef = dist[i].coef(nbrPts)
J[1] = np.append(np.sqrt(coef[1][1:]), [0])
J[0] = coef[0]
val, vec = linalg.eig_banded(J, lower=1)
weights[:, i] = vec[0, :]**2
points[:, i] = val.real
# Places the points and weights in the domain
point = np.zeros((nbrPts**dim, dim))
weight = np.zeros((nbrPts**dim, dim))
for i in range(dim):
v1 = np.repeat(points[:, i], nbrPts**(dim - 1 - i))
v2 = np.repeat(weights[:, i], nbrPts**(dim - 1 - i))
weight[:, i] = np.tile(v2, nbrPts**i)
point[:, i] = np.tile(v1, nbrPts**i)
weight = np.prod(weight, axis=1)
point = np.squeeze(point)
return point, weight
def gauss(alpha, beta):
"""
Compute the Gauss nodes and weights from the recursion
coefficients associated with a set of orthogonal polynomials
Inputs:
alpha - recursion coefficients
beta - recursion coefficients
Outputs:
x - quadrature nodes
w - quadrature weights
Adapted from the MATLAB code by Walter Gautschi
http://www.cs.purdue.edu/archives/2002/wxg/codes/gauss.m
"""
from scipy.linalg import eig_banded
A = np.vstack((np.sqrt(beta), alpha))
x, V = eig_banded(A, lower=False)
w = beta[0] * sp.real(sp.power(V[0, :], 2))
return x, w
def gauss(alpha,beta):
"""
Compute the Gauss nodes and weights from the recursion
coefficients associated with a set of orthogonal polynomials
Inputs:
alpha - recursion coefficients
beta - recursion coefficients
Outputs:
x - quadrature nodes
w - quadrature weights
Adapted from the MATLAB code by Walter Gautschi
http://www.cs.purdue.edu/archives/2002/wxg/codes/gauss.m
"""
from scipy.linalg import eig_banded
A = np.vstack((np.sqrt(beta),alpha))
x,V = eig_banded(A,lower=False)
w = beta[0]*sp.real(sp.power(V[0,:],2))
return x,w
def __init__(self, Nr, dr, mmax, queue, kernel):
r_list = cell_centers(0.0, Nr * dr, Nr)
A1 = 2 * np.pi * (r_list - 0.5 * dr)
A2 = 2 * np.pi * (r_list + 0.5 * dr)
V = np.pi * ((r_list + 0.5 * dr)**2 - (r_list - 0.5 * dr)**2)
self.Lambda = np.sqrt(V)
# Highest negative mode is never needed due to symmetry, so we have 2*mmax modes instead of 2*mmax+1.
if mmax == 0:
self.vals = np.zeros((Nr, 1))
else:
self.vals = np.zeros((Nr, 2 * mmax))
# Save storage by only keeping eigenvectors for positive modes (negative modes are the same)
self.Hi = np.zeros((Nr, Nr, mmax + 1))
for m in range(0, mmax + 1):
T1 = A1 / (dr * V)
T2 = -(A1 + A2) / (dr * V) - (m / r_list)**2
T3 = A2 / (dr * V)
# Boundary conditions
T2[0] += T1[0]
T2[Nr - 1] -= T3[Nr - 1]
# Symmetrize the matrix
# S = Lambda * T * Lambda^-1
# This is the root-volume weighting
T1[1:] *= self.Lambda[1:] / self.Lambda[:-1]
T3[:-1] *= self.Lambda[:-1] / self.Lambda[1:]
a_band_upper = np.zeros((2, Nr))
a_band_upper[
0, :] = T1 # T3->T1 thanks to scipy packing and symmetry
a_band_upper[1, :] = T2
self.vals[:, m], self.Hi[:, :, m] = eig_banded(a_band_upper)
# Set eigenvalues of negative modes (they are the same as corresponding positive modes)
# Modes are packed in usual FFT fashion, so negative indices work as expected.
for m in range(1, mmax):
self.vals[:, -m] = self.vals[:, m]
self.queue = queue
self.kernel = kernel
# Create a new figure
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
frame = 0
while frame < frames:
k2 = k2start + speed*frame
# Create effective Hamiltonian and solve
A = np.zeros([2,levels], dtype=np.complex128)
for n in range(levels):
A[1,n] = V2*np.cos(q*b*k2/p + 2*np.pi*np.float(n)*q/p)
if n < (levels-1):
A[0,n+1] = V1*np.exp(np.complex(0.0,q*a*k1/p))
[d, v] = eig_banded(A)
for i in range(levels):
E[i] = d[i]
State = v[:,i]
StateAbs = abs(State)
ns[i] = StateAbs.argmax()
if frame == 0:
prev_ns[i] = ns[i]
positions[i] = ns[i]
if frame > 0:
ss[i] = ns[i] - prev_ns[i]
if ss[i] > p/2.0:
ss[i] = -(p - ss[i])
elif ss[i] < -p/2.0:
ss[i] = p + ss[i]
def slepian(M, width, sym=True):
"""Return a digital Slepian (DPSS) window.
Used to maximize the energy concentration in the main lobe. Also called
the digital prolate spheroidal sequence (DPSS).
Parameters
----------
M : int
Number of points in the output window. If zero or less, an empty
array is returned.
width : float
Bandwidth
sym : bool, optional
When True (default), generates a symmetric window, for use in filter
design.
When False, generates a periodic window, for use in spectral analysis.
Returns
-------
w : ndarray
The window, with the maximum value always normalized to 1
Examples
--------
Plot the window and its frequency response:
>>> from scipy import signal
>>> from scipy.fftpack import fft, fftshift
>>> import matplotlib.pyplot as plt
>>> window = signal.slepian(51, width=0.3)
>>> plt.plot(window)
>>> plt.title("Slepian (DPSS) window (BW=0.3)")
>>> plt.ylabel("Amplitude")
>>> plt.xlabel("Sample")
>>> plt.figure()
>>> A = fft(window, 2048) / (len(window)/2.0)
>>> freq = np.linspace(-0.5, 0.5, len(A))
>>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max())))
>>> plt.plot(freq, response)
>>> plt.axis([-0.5, 0.5, -120, 0])
>>> plt.title("Frequency response of the Slepian window (BW=0.3)")
>>> plt.ylabel("Normalized magnitude [dB]")
>>> plt.xlabel("Normalized frequency [cycles per sample]")
"""
if M < 1:
return np.array([])
if M == 1:
return np.ones(1, 'd')
odd = M % 2
if not sym and not odd:
M = M + 1
# our width is the full bandwidth
width = width / 2
# to match the old version
width = width / 2
m = np.arange(M, dtype='d')
H = np.zeros((2, M))
H[0, 1:] = m[1:] * (M - m[1:]) / 2
H[1, :] = ((M - 1 - 2 * m) / 2)**2 * np.cos(2 * np.pi * width)
_, win = linalg.eig_banded(H, select='i', select_range=(M-1, M-1))
win = win.ravel() / win.max()
if not sym and not odd:
win = win[:-1]
return win
# Create a new figure
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
frame = 0
while frame < frames:
k2 = k2start + speed * frame
# Create effective Hamiltonian and solve
A = np.zeros([2, levels], dtype=np.complex128)
for n in range(levels):
A[1, n] = V2 * np.cos(q * b * k2 / p + 2 * np.pi * np.float(n) * q / p)
if n < (levels - 1):
A[0, n + 1] = V1 * np.exp(np.complex(0.0, q * a * k1 / p))
[d, v] = eig_banded(A)
for i in range(levels):
E[i] = d[i]
State = v[:, i]
StateAbs = abs(State)
ns[i] = StateAbs.argmax()
if frame == 0:
prev_ns[i] = ns[i]
positions[i] = ns[i]
if frame > 0:
ss[i] = ns[i] - prev_ns[i]
if ss[i] > p / 2.0:
ss[i] = -(p - ss[i])
elif ss[i] < -p / 2.0:
ss[i] = p + ss[i]
def slepian(M, width, sym=True):
"""Return a digital Slepian (DPSS) window.
Used to maximize the energy concentration in the main lobe. Also called
the digital prolate spheroidal sequence (DPSS).
Parameters
----------
M : int
Number of points in the output window. If zero or less, an empty
array is returned.
width : float
Bandwidth
sym : bool, optional
When True (default), generates a symmetric window, for use in filter
design.
When False, generates a periodic window, for use in spectral analysis.
Returns
-------
w : ndarray
The window, with the maximum value always normalized to 1
Examples
--------
Plot the window and its frequency response:
>>> from scipy import signal
>>> from scipy.fftpack import fft, fftshift
>>> import matplotlib.pyplot as plt
>>> window = signal.slepian(51, width=0.3)
>>> plt.plot(window)
>>> plt.title("Slepian (DPSS) window (BW=0.3)")
>>> plt.ylabel("Amplitude")
>>> plt.xlabel("Sample")
>>> plt.figure()
>>> A = fft(window, 2048) / (len(window)/2.0)
>>> freq = np.linspace(-0.5, 0.5, len(A))
>>> response = 20 * np.log10(np.abs(fftshift(A / abs(A).max())))
>>> plt.plot(freq, response)
>>> plt.axis([-0.5, 0.5, -120, 0])
>>> plt.title("Frequency response of the Slepian window (BW=0.3)")
>>> plt.ylabel("Normalized magnitude [dB]")
>>> plt.xlabel("Normalized frequency [cycles per sample]")
"""
if M < 1:
return np.array([])
if M == 1:
return np.ones(1, 'd')
odd = M % 2
if not sym and not odd:
M = M + 1
# our width is the full bandwidth
width = width / 2
# to match the old version
width = width / 2
m = np.arange(M, dtype='d')
H = np.zeros((2, M))
H[0, 1:] = m[1:] * (M - m[1:]) / 2
H[1, :] = ((M - 1 - 2 * m) / 2)**2 * np.cos(2 * np.pi * width)
_, win = linalg.eig_banded(H, select='i', select_range=(M - 1, M - 1))
win = win.ravel() / win.max()
if not sym and not odd:
win = win[:-1]
return win
def golub_welsch(order, dist, acc=100, **kws):
"""
Golub-Welsch algorithm for creating quadrature nodes and weights
Parameters
----------
order : int
Quadrature order
dist : Dist
Distribution nodes and weights are found for with
`dim=len(dist)`
acc : int
Accuracy used in discretized Stieltjes procedure.
Will be increased by one for each itteration.
Returns
-------
x : numpy.array
Optimal collocation nodes with `x.shape=(dim, order+1)`
w : numpy.array
Optimal collocation weights with `w.shape=(order+1,)`
Examples
--------
>>> Z = cp.Normal()
>>> x, w = cp.golub_welsch(3, Z)
>>> print(x)
[[-2.33441422 -0.74196378 0.74196378 2.33441422]]
>>> print(w)
[ 0.04587585 0.45412415 0.45412415 0.04587585]
Multivariate
>>> Z = cp.J(cp.Uniform(), cp.Uniform())
>>> x, w = cp. golub_welsch(1, Z)
>>> print(x)
[[ 0.21132487 0.21132487 0.78867513 0.78867513]
[ 0.21132487 0.78867513 0.21132487 0.78867513]]
>>> print(w)
[ 0.25 0.25 0.25 0.25]
"""
o = np.array(order)*np.ones(len(dist), dtype=int)+1
P, g, a, b = stieltjes(dist, np.max(o), acc=acc, retall=True, **kws)
X, W = [], []
dim = len(dist)
for d in range(dim):
if o[d]:
A = np.empty((2, o[d]))
A[0] = a[d, :o[d]]
A[1, :-1] = np.sqrt(b[d, 1:o[d]])
vals, vecs = eig_banded(A, lower=True)
x, w = vals.real, vecs[0, :]**2
indices = np.argsort(x)
x, w = x[indices], w[indices]
else:
x, w = np.array([a[d, 0]]), np.array([1.])
X.append(x)
W.append(w)
if dim==1:
x = np.array(X).reshape(1,o[0])
w = np.array(W).reshape(o[0])
else:
x = cp.utils.combine(X).T
w = np.prod(cp.utils.combine(W), -1)
assert len(x)==dim
assert len(w)==len(x.T)
return x, w
def golub_welsch(order, dist, acc=100, **kws):
"""
Golub-Welsch algorithm for creating quadrature nodes and weights
Parameters
----------
order : int
Quadrature order
dist : Dist
Distribution nodes and weights are found for with
`dim=len(dist)`
acc : int
Accuracy used in discretized Stieltjes procedure.
Will be increased by one for each itteration.
Returns
-------
x : numpy.array
Optimal collocation nodes with `x.shape=(dim, order+1)`
w : numpy.array
Optimal collocation weights with `w.shape=(order+1,)`
Examples
--------
>>> Z = cp.Normal()
>>> x, w = cp.golub_welsch(3, Z)
>>> print x
[[-2.33441422 -0.74196378 0.74196378 2.33441422]]
>>> print w
[ 0.04587585 0.45412415 0.45412415 0.04587585]
Multivariate
>>> Z = cp.J(cp.Uniform(), cp.Uniform())
>>> x, w = cp. golub_welsch(1, Z)
>>> print x
[[ 0.21132487 0.21132487 0.78867513 0.78867513]
[ 0.21132487 0.78867513 0.21132487 0.78867513]]
>>> print w
[ 0.25 0.25 0.25 0.25]
"""
o = np.array(order) * np.ones(len(dist), dtype=int) + 1
P, g, a, b = stieltjes(dist, np.max(o), acc=acc, retall=True, **kws)
X, W = [], []
dim = len(dist)
for d in xrange(dim):
if o[d]:
A = np.empty((2, o[d]))
A[0] = a[d, :o[d]]
A[1, :-1] = np.sqrt(b[d, 1:o[d]])
vals, vecs = eig_banded(A, lower=True)
x, w = vals.real, vecs[0, :]**2
indices = np.argsort(x)
x, w = x[indices], w[indices]
# p = P[-1][d]
# dp = po.differential(p, po.basis(1,1,dim)[d])
#
# x = x - p(x)/dp(x)
# x = x - p(x)/dp(x)
# x = x - p(x)/dp(x)
#
# z = np.arange(dim)
# arg = np.array([k*(z == d) for k in range(2*o[d]-3)])
# b_ = dist.mom(arg.T)
#
# X_, r = np.meshgrid(x, np.arange(2*o[d]-3))
# X_ = X_**r
# w = np.linalg.lstsq(X_, b_)[0].flatten()
# print "w", w.shape
# print np.linalg.lstsq(X_, b_)[0]
else:
x, w = np.array([a[d, 0]]), np.array([1.])
X.append(x)
W.append(w)
if dim == 1:
x = np.array(X).reshape(1, o[0])
w = np.array(W).reshape(o[0])
else:
x = combine(X).T
w = np.prod(combine(W), -1)
assert len(x) == dim
assert len(w) == len(x.T)
return x, w
def main(parameters):
r"""Calculates the ground-state of the system. If the system is perturbed, the time evolution of
the perturbed system is then calculated.
parameters
----------
parameters : object
Parameters object
returns object
Results object
"""
# Array initialisations
pm = parameters
string = 'EXT: constructing arrays'
pm.sprint(string, 1)
pm.setup_space()
# Construct the kinetic energy matrix
K = construct_K(pm)
# Construct the Hamiltonian matrix
H = np.copy(K)
H[0, :] += pm.space.v_ext[:]
# Solve the Schroedinger equation
string = 'EXT: calculating the ground-state density'
pm.sprint(string, 1)
energies, wavefunctions = spla.eig_banded(H, lower=True)
# Normalise the wavefunctions
wavefunctions /= np.sqrt(pm.space.delta)
# Calculate the ground-state density
density = np.absolute(wavefunctions[:, 0])**2
# Calculate the ground-state energy
energy = energies[0]
string = 'EXT: ground-state energy = {:.5f}'.format(energy)
pm.sprint(string, 1)
# Save the quantities to file
results = rs.Results()
results.add(pm.space.v_ext, 'gs_ext_vxt')
results.add(density, 'gs_ext_den')
results.add(energy, 'gs_ext_E')
results.add(wavefunctions.T, 'gs_ext_eigf')
results.add(energies, 'gs_ext_eigv')
if (pm.run.save):
results.save(pm)
# Propagate through real time
if (pm.run.time_dependence):
# Print to screen
string = 'EXT: constructing arrays'
pm.sprint(string, 1, newline=True)
# Construct the Hamiltonian matrix
H = np.copy(K)
H[0, :] += pm.space.v_ext[:]
if (pm.sys.im == 1):
H = H.astype(np.cfloat)
H[0, :] += pm.space.v_pert[:]
# Construct the sparse matrices used in the Crank-Nicholson method
A = construct_A(pm, H)
C = 2.0 * sps.identity(pm.space.npt, dtype=np.cfloat) - A
# Construct the time-dependent density array
density = np.zeros((pm.sys.imax, pm.space.npt), dtype=np.float)
# Save the ground-state
wavefunction = wavefunctions[:, 0].astype(np.cfloat)
density[0, :] = np.absolute(wavefunction[:])**2
# Print to screen
string = 'EXT: real time propagation'
pm.sprint(string, 1)
# Perform real time iterations
for i in range(1, pm.sys.imax):
# Construct the vector b
b = C * wavefunction
# Solve Ax=b
wavefunction, info = spsla.cg(A,
b,
x0=wavefunction,
tol=pm.ext.rtol_solver)
# Normalise the wavefunction
norm = npla.norm(wavefunction) * np.sqrt(pm.space.delta)
wavefunction /= norm
string = 'EXT: t = {:.5f}, normalisation = {}'.format(
i * pm.sys.deltat, norm**2)
pm.sprint(string, 1, newline=False)
# Calculate the density
density[i, :] = np.absolute(wavefunction[:])**2
# Calculate the current density
current_density = calculate_current_density(pm, density)
# Save the quantities to file
results.add(density, 'td_ext_den')
results.add(current_density, 'td_ext_cur')
results.add(pm.space.v_ext + pm.space.v_pert, 'td_ext_vxt')
if (pm.run.save):
results.save(pm)
return results
def iwave_modes_banded(N2, dz, k=None):
"""
!!! DOES NOT WORK!!!
Calculates the eigenvalues and eigenfunctions to the internal wave eigenvalue problem:
$$
\left[ \frac{d^2}{dz^2} - \frac{1}{c_0} \bar{\rho}_z \right] \phi = 0
$$
with boundary conditions
"""
nz = N2.shape[0] # Remove the surface values
if k is None:
k = nz - 2
dz2 = 1 / dz**2
# Construct the LHS matrix, A
A = np.vstack([-1*dz2*np.ones((nz,)),\
2*dz2*np.ones((nz,)),\
-1*dz2*np.ones((nz,)),\
])
# BC's
#A[0,0] = -1.
#A[0,1] = 0.
#A[-1,-1] = -1.
#A[-1,-2] = 0.
A[1, 0] = -1.
A[2, 0] = 0.
A[1, -1] = -1.
A[0, -1] = 0.
# Now convert from a generalized eigenvalue problem to
# A.v = lambda.B.v
# a standard problem
# A.v = lambda.v
# By multiply the LHS by inverse of B
# (B^-1.A).v = lambda.v
# B^-1 = 1/N2 since B is diagonal
A[0, :] /= N2
A[1, :] /= N2
A[2, :] /= N2
w, phi = linalg.eig_banded(A)
pdb.set_trace()
## Main diagonal
#dd = 2*dz2*np.ones((nz,))
#dd /= N2
#dd[0] = -1
#dd[-1] = -1
## Off diagonal
#ee = -1*dz2*np.ones((nz-1,))
#ee /= N2[0:-1]
#ee[0] = 0
#ee[-1] = 0
## Solve... (use scipy not numpy)
#w, phi = linalg.eigh_tridiagonal(dd, ee )
#####
c = 1. / np.power(w, 0.5) # since term is ... + N^2/c^2 \phi
# Sort by the eigenvalues
idx = np.argsort(c)[::-1] # descending order
## Calculate the actual phase speed
cn = np.real(c[idx])
idxgood = ~np.isnan(cn)
phisort = phi[:, idx]
return np.real(phisort[:, idxgood]), np.real(cn[idxgood])