def linear_discretization(spec_func, stop, num, start=0.0):
"""A simple linear method to discretize a spectral density.
Parameters
----------
spec_func : float -> float
Offen denoted as J(w).
start : float, optional
Start point of the spectrum, defalut is 0.0.
stop : float
End point of the spectrum.
num : int
Number of modes to be given.
Returns
-------
ans : [(float, float)]
`ans[i][0]` is the omega of one mode and `ans[i][1]` is the
corrospoding coupling in second quantization for all `i` in
`range(0, num)`.
"""
def direct_quad(a, b):
density = quad(spec_func, a, b)[0]
omega = quad(lambda x: x * spec_func(x), a, b)[0] / density
coupling = np.sqrt(density)
return omega, coupling
space = np.linspace(start, stop, num + 1, endpoint=True)
omega_0, omega_1 = space[:-1], space[1:]
ans = list(map(direct_quad, omega_0, omega_1))
return ans
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 plot_func(self, func_list, x_min, x_max, y_min=0., y_max=0., npts=None):
r"""Plot functions.
Parameters
----------
func_list : [float->float]
List of functions to plot.
x_min : float
x_max : float
y_min : float, optional
y_max : float, optional
npts : int, optional
Number of points to calculate on a single state.
"""
if npts is None:
npts = self.n
x = np.linspace(x_min, x_max, npts)
with figure() as fig:
plt.subplots_adjust(left=0.1, right=0.9, bottom=0.1, top=0.9)
for func in func_list:
chi = func(x)
y_max = max(y_max, max(chi))
y_min = min(y_min, min(chi))
plt.plot(x, chi)
plt.xlim(x_min, x_max)
plt.ylim(y_min * 1.05, y_max * 1.05)
plt.savefig('functions-{}.pdf'.format(self.comment))
return
def plot_wf(self, vec, msg=None):
"""Plot 2D wavefunction.
Parameters
----------
vec : (N,) ndarray
msg : string
Notes
-----
This is a generator. Use .next() to plot, and .send(vec, msg) to
plot next wavefunction with the same figure parameters.
"""
x_lim, y_lim = self.bound_list[:2]
x, y = self.grid_points_list[:2]
shape = self.n_list[:2]
x, y = np.meshgrid(x, y)
z_lim = int(np.max(np.abs(vec)) * 15) / 10
bound = np.linspace(-z_lim, z_lim, 100)
while vec is not None:
vec = np.reshape(vec, shape)
if msg is None:
msg = int(time.time())
with figure() as fig:
plt.contourf(x, y, vec, bound, cmap='seismic')
plt.colorbar(boundaries=bound)
plt.savefig('functions-{}.pdf'.format(msg))
received = yield
if received is None:
raise StopIteration
else:
vec, msg = received
def plot_eigen(self, x_min, x_max, npts=None, n_plot=None, scale=2.):
r"""Plot the eigenstate together with energy and potential curve.
Parameters
----------
x_min : float
x_max : float
npts : int, optional
Number of points to calculate on a single state.
n_plot : int, optional
Number of states to calculate.
scale : float, optional
Adjust the scale of wavefunction.
"""
if npts is None:
npts = self.n
if n_plot is None:
n_plot = len(self.energy)
x = np.linspace(x_min, x_max, npts)
vx = [self.v_func(x_) for x_ in x]
with figure() as fig:
plt.subplots_adjust(left=0.1, right=0.9, bottom=0.1, top=0.9)
plt.plot(x, vx, 'k-', lw=2)
y_min = min(vx)
y_max = self.v_func(0)
for i in range(n_plot):
e = self.energy[i]
plt.plot([x[0], x[-1]], [e, e], '--', color='gray')
phi = (self.dvr2cont(self.eigenstates[i]))(x)
plt.plot(x, scale * phi + e)
y_max = min(y_min, e - scale * min(phi))
y_max = max(y_max, e + scale * max(phi))
plt.xlim(x_min, x_max)
plt.ylim(y_min, 1.05 * y_max)
plt.savefig('eigenstates-{}.pdf'.format(self.comment))
return