def _calculate_dvr(self):
# calculate grid points
step_length = self.length / (self.n + 1)
self.grid_points = np.array([self.a + step_length * i for i in range(1, self.n + 1)])
# calculate U matrix
j = np.arange(1, self.n + 1)[:, None]
a = np.arange(1, self.n + 1)[None, :]
self._u_mat = (np.sqrt(2 / (self.n + 1)) * np.sin(j * a * np.pi / (self.n + 1)))
return self.grid_points, self._u_mat
def plot_dvr(self, x_min, x_max, npts=None, indices=None):
r"""Plot DVR basis.
Parameters
----------
x_min : float
x_max : float
npts : int, optional
Number of points to calculate on a single state.
indices : (int), optional
An iterable object containing the indices of DVR to be plotted.
"""
if npts is None:
npts = self.n
if indices is None:
indices = np.arange(self.n)
x = np.linspace(x_min, x_max, npts)
y_min = 0.
y_max = 0.
with figure() as fig:
plt.subplots_adjust(left=0.1, right=0.9, bottom=0.1, top=0.9)
for i in indices:
x_i = self.grid_points[i]
plt.plot([x_i, x_i], [-10., 10.], '--', color='gray')
chi = (self.dvr_func(i))(x)
if (self.dvr_func(i))(x_i) < 0.:
chi = -1. * chi
y_max = max(y_max, max(chi))
y_min = min(y_min, min(chi))
plt.plot(x, chi)
plt.plot([x_min, x_max], [0., 0.], 'k-')
plt.xlim(x_min, x_max)
plt.ylim(y_min * 1.05, y_max * 1.05)
plt.savefig('dvr_functions-{}.pdf'.format(self.comment))
return
def spectrum(self, init=None, length=5., max_inter=0.01, window=None, **kwargs):
"""Power spectrum.
Parameters
----------
init : (N,) ndarray, optional
cut : float
max_inter : float
window : float -> float
Window function with length.
Returns
-------
freq : [float]
sigma : [complex]
"""
t, auto = zip(*self.autocorrelation(init=init, stop=length, max_inter=max_inter, **kwargs))
n = len(t)
tau = (t[-1] - t[0]) / (n - 1)
uniform = all(abs(t_i - i * tau) < 1.e-8 for i, t_i in enumerate(t))
if window is not None:
zipped = zip(t, auto)
auto = [a_i * window(t) for t, a_i in zipped]
if uniform:
omega = 2 * np.pi / n / tau
freq = np.arange(n) * omega
sigma = fftpack.ifft(auto)
return freq, sigma
else:
raise NotImplementedError("Need NUDFT, or use smaller max_inter.")
def number_operator(self):
if self.dense:
ans = np.diag(np.arange(self.dim))
else:
def matvec(x):
return self.raising(self.lowering(x))
ans = self.operator(matvec=matvec)
return ans
def hamiltonian(self):
coeff = self.hbar * self.omega
if self.dense:
ans = np.diag(coeff * (0.5 + np.arange(self.dim)))
else:
def matvec(x):
return coeff * (self.raising(self.lowering(x)) + 0.5 * x)
ans = self.operator(matvec=matvec)
return ans
def t_mat(self):
"""Return the kinetic energy matrix in DVR.
Returns
-------
(n, n) np.ndarray
A 2-d matrix.
"""
factor = -self.hbar**2 / (2 * self.m_e)
j = np.arange(1, self.n + 1)
t_matrix = np.diag(-(j * np.pi / self.length)**2)
t_matrix = factor * self.fbr2dvr_mat(t_matrix)
return t_matrix
def h_mat(self):
"""Return the Hamiltonian energy matrix in DVR.
Returns
-------
Hamitonian : LinearOperator
A 2-d matrix.
"""
class _Hamiltonian(LinearOperator):
"""
Parameters
----------
v_diag : (n,) ndarray
In DVR
t_diag : (n,) ndarray
In FBR
"""
def __init__(self, v_diag, t_diag):
n = len(v_diag)
self.v = v_diag
self.t = t_diag / (2. * (n + 1.))
super(_Hamiltonian, self).__init__('d', (n, n))
def _matvec(self, vec):
vec1 = self.v * vec
vec2 = fftpack.dst(self.t * fftpack.dst(vec, type=1), type=1)
return vec1 + vec2
def _rmatvec(self, vec):
return self._matvec(vec)
if self._h_mat is not None:
return self._h_mat
v = self.v_func(self.grid_points)
j = np.arange(1, self.n + 1)
t = (self.hbar * j * np.pi / self.length)**2 / (2 * self.m_e)
return _Hamiltonian(v, t)
def annihilation_operator(self):
if self.dense:
ans = np.diag(np.sqrt(np.arange(1, self.dim)), 1)
else:
ans = self.operator(matvec=self.lowering, rmatvec=self.raising)
return ans
def numberer(self, k):
"""Acting on 0-th index"""
return np.diag(np.arange(self.n_dims[k]))
def annihilator(self, k):
"""Acting on 0-th index"""
dim = self.n_dims[k]
lower = np.eye(dim, k=-1)
sqrt_n = np.diag(np.sqrt(np.arange(dim)))
return sqrt_n @ lower
def creator(self, k):
"""Acting on 0-th index"""
dim = self.n_dims[k]
raiser = np.eye(dim, k=1)
sqrt_n = np.diag(np.sqrt(np.arange(dim)))
return raiser @ sqrt_n
def _numberer(self, k):
return np.diag(np.arange(self.n_dims[k], dtype=DTYPE))
def _lower(self, k):
"""Acting on 0-th index"""
dim = self.n_dims[k]
sqrt_n = np.diag(np.sqrt(np.arange(dim, dtype=DTYPE)))
return sqrt_n @ np.eye(dim, k=-1, dtype=DTYPE)
def _sqrt_numberer(self, k, start=0):
return np.diag(
np.sqrt(np.arange(start, start + self.n_dims[k], dtype=DTYPE)))
def _numberer(self, k, start=0):
return np.diag(np.arange(start, start + self.n_dims[k]))