def error(x, V, n_basis_vec, exact):
# Initialize the Schrod
eqn = schrod.Schrod(x, V)
# Calculate the error
rel_err=[]
for n in n_basis_vec:
eqn.set_n_basis(n)
eqn.solve()
lam = eqn.eigs
err = (np.abs((lam - exact[0:n])/exact[0:n]))[...,0:n_basis_vec[0]]
rel_err.append(err)
return np.asarray(rel_err)
def timing(x, V, n_basis_vec, n_tests=1):
# Initialize the Schrod
eqn = schrod.Schrod(x, V)
# Time the Schrod for different numbers of basis states
times = []
for n in n_basis_vec:
avg_time = 0
for i in range(n_tests):
eqn.set_n_basis(n)
ti = time.time()
eqn.solve()
tf = time.time()
avg_time += (tf - ti) / n_tests
times.append(avg_time)
return times
Solve Schodinger's equation and plot the probability distributions
Author: D. Hudson Smith
"""
from __future__ import print_function
import matplotlib.pyplot as plt
import numpy as np
import schrod
# Specify the potential
x = np.linspace(-3, 3, 200)
V = x + 0.5 * x**4
# Create and solve Schrodinger's equation
eqn = schrod.Schrod(x, V, n_basis=40)
eqn.solve()
# Get the eigenvalues and eigenvectors
eig_vals = eqn.eigs
probs = eqn.prob_eig_x()
# Plot everything
plt.xlim((x[0], x[-1]))
plt.ylim((-1, 9))
plt.title("Probability distributions for\n$V(x) = x + 0.5 x^4$", fontsize=18)
plt.xlabel("$x$")
plt.ylabel('Energy')
n_eigs_plt = 5
plt.plot(x, V, 'b-', lw=2)
"""
Solve Schodinger's equation up to a specified error
Author: D. Hudson Smith
"""
from __future__ import print_function
import schrod, numpy
# Specify the potential
x = numpy.linspace(-20, 20, 200)
V = 1. / 2. * x**2
# Ask for 10 digits of precision
eqn = schrod.Schrod(x, V)
sol = eqn.solve_to_tol(n_eig=5, tol=1e-10, n_init=50, n_max=150)
# Print the first five eigenvalues
numpy.set_printoptions(15)
print("Eigenvalues: \n", eqn.eigs[0:5], "\nError estimate: \n", sol.eig_errs,
"\nNumber of basis states at convergence: \n", sol.n_basis_converged)
# If <tol> cannot be achieved, a warning is returned:
sol = eqn.solve_to_tol(n_eig=5, tol=1e-15, n_init=100, n_max=150, n_step=10)
# Print the first five eigenvalues
print("Eigenvalues: \n", eqn.eigs[0:5], "\nError estimate: \n", sol.eig_errs,
"\nNumber of basis states at termination: \n", sol.n_basis_converged)
# Sample output
# Eigenvalues:
######################################################################
# Set and solve Schrodinger's Equation
# Specify the potential
def Vfunc(x):
return np.abs(x)
x = np.linspace(-30, 30, 600)
V = Vfunc(x)
Vmax = V[0] / 3
# Create and solve Schrodinger's equation
eqn = schrod.Schrod(x, V, n_basis=100)
sol = eqn.solve()
# Print the first five eigenvalues
print("Eigenvalues: ", eqn.eigs[0:5])
######################################################################
# The initial wave packet
# Helper functions for gaussian wave-packets
def gauss_x(x, a, x0, k0):
"""
a gaussian wave packet of width a, centered at x0, with momentum k0
"""
return ((a * np.sqrt(np.pi))**(-0.5) *