"""
.. _two_poschl_teller:
Poschl-Teller potentials
########################
Summary:
Calculates a system with two Poschl-Teller wells
"""
import sys
import os
currentpath = os.path.abspath('.')
sys.path.insert(0, os.path.dirname(currentpath))
import numpy as np
from ext_potentials import poschl_teller
from non_interacting_solver import SparseEigenSolver
import functools
import matplotlib.pyplot as plt
[docs]def get_plotting_params():
r'''Initialize figure for plots.
Returns:
fig: a Figure object with initialized parameters.
ax: an Axes object with initialized parameters.
'''
# plotting parameters
params = {'mathtext.default': 'default'}
plt.rcParams.update(params)
plt.rcParams['axes.axisbelow'] = True
fig_size = plt.rcParams["figure.figsize"]
fig_size[0] = 9
fig_size[1] = 6
plt.rcParams["figure.figsize"] = fig_size
fig, ax = plt.subplots()
return fig, ax
[docs]def plot_and_save(x_list, y_list, x_label, y_label, name,
xlim=(None, None), ylim=(None, None), folder=None):
r'''Plot a single curve on a graph and save it as .png file.
Args:
x_list: list (or array), specify the x-axis points.
y_list: list (or array), specify the y-axis points.
x_label: label for x-axis.
y_label: label for y-axis.
name: str, name of the image saved (without .png).
folder: str, the path to the folder to save the image in. Does not need
to pre-exist before running.
'''
# initialize figure for plots
fig, ax = get_plotting_params()
# matplotlib trick to obtain same color of a previous plot
ax.plot(x_list, y_list, marker='o', linestyle='solid', color='blue')
# set labels for x and y axis
ax.set_xlabel(x_label, fontsize=18)
ax.set_ylabel(y_label, fontsize=18)
# set x and y limits
xlim_l = xlim[0]
xlim_r = xlim[1]
ylim_b = ylim[0]
ylim_t = ylim[1]
if xlim_l != None:
ax.set_xlim(left=xlim_l)
if xlim_r != None:
ax.set_xlim(right=xlim_r)
if ylim_b != None:
ax.set_ylim(bottom=ylim_b)
if ylim_t != None:
ax.set_ylim(top=ylim_t)
# set grids
plt.grid(alpha=0.4)
plt.gca().xaxis.grid(True, which='minor', alpha=0.4)
plt.gca().yaxis.grid(True, which='minor', alpha=0.4)
# save image
if folder != None:
# create folder if no such directory
if not os.path.isdir(folder):
os.mkdir(folder)
# save fig
plt.savefig(f'{folder}/{name}.png')
plt.close()
else:
# save fig
plt.savefig(f'{name}.png')
plt.close()
[docs]def plot_multiple_and_save(plot_list, name,
xlim=(None, None), ylim=(None, None), folder=None):
r'''Plot mutiple curves on a single graph and save it as .png file.
Args:
plot_list: list of lists, specify plotting parameters for each curve
in the form of [label, x_list, y_list, linestyle, color].
name: str, name of the image saved (without .png).
xlim: tuple, xlim[0] == left limit, xlim[1] == right limit; setting to
None means flexible ranges.
ylim: tuple, ylim[0] == bottom limit, ylim[1] == top limit; setting to
None means flexible ranges.
folder: str, the path to the folder to save the image in. Does not need
to pre-exist before running.
'''
# initialize figure for plots
fig, ax = get_plotting_params()
# matplotlib trick to obtain same color of a previous plot
for data_list in plot_list:
label = data_list[0]
x_list = data_list[1]
y_list = data_list[2]
linestyle = data_list[3]
color = data_list[4]
ax.plot(x_list, y_list, label=label, linestyle=linestyle, color=color)
# set x label
ax.set_xlabel('x', fontsize=18)
# set x and y limits
xlim_l = xlim[0]
xlim_r = xlim[1]
ylim_b = ylim[0]
ylim_t = ylim[1]
if xlim_l != None:
ax.set_xlim(left=xlim_l)
if xlim_r != None:
ax.set_xlim(right=xlim_r)
if ylim_b != None:
ax.set_ylim(bottom=ylim_b)
if ylim_t != None:
ax.set_ylim(top=ylim_t)
# set grids
plt.grid(alpha=0.4)
plt.gca().xaxis.grid(True, which='minor', alpha=0.4)
plt.gca().yaxis.grid(True, which='minor', alpha=0.4)
# set legend
plt.legend(loc=2)
# save image
if folder != None:
# create folder if no such directory
if not os.path.isdir(folder):
os.mkdir(folder)
# save fig
plt.savefig(f'{folder}/{name}.png')
plt.close()
else:
# save fig
plt.savefig(f'{name}.png')
plt.close()
[docs]def two_poschl_teller(grids, d, lam=1., a=1.):
r"""Two Poschl-Teller potential wells seperated by distance d.
Args:
grids: numpy array of grid points for evaluating 1d potential.
(num_grids,)
d: float, distance seperated by two Poschl-Teller potential wells.
lam: float, lambda in the Poschl-Teller potential function.
a: float, coefficient in the Poschl-Teller potential function.
center: float, the center of the potential.
Returns:
Potential on grid with shape (num_grid,)
Raises:
ValueError: If lam is not positive.
"""
potential_grids = poschl_teller(grids, lam=lam, a=a, center=-d / 2)
potential_grids += poschl_teller(grids, lam=lam, a=a, center=d / 2)
return potential_grids
if __name__ == '__main__':
# initialize variables
test_range = (-20, 20)
grids = np.linspace(*test_range, 1000)
boundary_condition = 'open'
n_point_stencil = 5
d_list = []
E_list = []
solver = SparseEigenSolver(grids,
boundary_condition=boundary_condition,
n_point_stencil=n_point_stencil)
# fill d_list and E_list
for d in np.linspace(0, 10, 50):
d_list.append(d)
potential_fn = functools.partial(two_poschl_teller, d=d)
solver.update_potential(potential_fn)
solver.solve_ground_state()
energy = solver.eigenvalues[0]
E_list.append(energy)
# save E vs. d into d_E_table.dat
import csv
with open('d_E_table.csv', 'w', newline='') as f:
writer = csv.writer(f)
writer.writerow(['d', 'E'])
writer.writerows(zip(d_list, E_list))
# plot E vs d
plot_and_save(d_list, E_list, 'd', 'E', 'E vs d')
# plot potential, wave function and eigenvalue at d = test_d
test_d = 2
potential_fn = functools.partial(two_poschl_teller, d=test_d)
solver.update_potential(potential_fn)
solver.solve_ground_state()
plot_list = [['potential', grids, potential_fn(grids), '-', 'orange'],
['wave function', grids, solver.wave_function[0], '-', 'blue'],
['eigenvalue', grids,
np.full(len(grids), solver.eigenvalues[0]), '--', 'green']]
plot_multiple_and_save(plot_list, f'd={test_d}', ylim=(-2, 1))