def dcst2(f):
""" Takes DCT along x, then DST along y
IN: f, the input 2D numpy array
OUT: b, the 2D transformed array """
M = f.shape[0] # Number of rows
N = f.shape[1] # Number of columns
a = np.zeros((M, N)) # Intermediate array
b = np.zeros((M, N)) # Final array
# Take transform along x
for j in range(N):
a[:, j] = dct(f[:, j])
# Take transform along y
for i in range(M):
b[i, :] = dst(a[i, :])
return b
#Initial wavefunction (at t=0):
#This has the shape of a gaussian function
###########################################
def si_initial(x):
values = np.zeros(len(x),complex)
value = np.exp(-((x-x0)**2.)/(2.*(sigma**2.))) * np.exp(1j*kappa*x)
values[0] = 0
values[-1] = 0
return value
#separating real and imaginary parts of initial wave function(si)
b_real = si_initial(x).real
b_imag = si_initial(x).imag
#calculating discrete sine transform
alpha = dst(b_real)
eta = dst(b_imag)
b_k = alpha + 1j*eta #bk
#k range
k = np.linspace(1,len(b_k),N)
#the arguement inside cos and sin function on page 443
vals = -((np.pi**2.)*(hbar)*(k**2.))/(2.*(M)*(L**2.))
#plotting for animation inside the while loop
plt.ion()
fig = plt.figure()
ax = plt.axes(xlim=(0,1e-8),ylim=(-5e-8,5e-8))
line = ax.plot(x,si_initial(x),'-b')
plt.show()
repsi = zeros(N, float)
impsi = zeros(N, float)
x = linspace(0, L, N)
#Set initial conditions
for n in range(1, N):
xn = n * a
x[n] = xn
gauss = exp(-(xn - x0)**2 / (2 * (sigma**2)))
repsi[n] = gauss * cos(kappa * xn)
impsi[n] = gauss * sin(kappa * xn)
lines = plt.plot(x, ndarray.tolist(repsi))
#determine fourier coefficients
alpha = dst(repsi)
eta = dst(impsi)
b = empty(N, float)
i = 0
for t in arange(0.0, tmax, tstep):
i += 1
for k in range(1, N):
angle = C * k * k * t
b[k] = alpha[k] * cos(angle) + eta[k] * sin(angle)
repsi = idst(b)
if (i % 20 == 0):
lines[0].set_ydata(ndarray.tolist(repsi[:]))
plt.pause(0.0000000001)
plt.draw()
from pylab import plot, show, legend, subplot, xlabel, ylabel
from numpy import zeros, empty, linspace, exp, arange,minimum, pi, sin,cos, array
from dcst import dst, idst, dct, idct
N = 256
# x = pi*n/N
x = arange(N)*pi/N
#function is a sine series
f = sin(x)-2*sin(4*x)+3*sin(5*x)-4*sin(6*x)
#do fourier sine series
fCoeffs = dst(f)
print('Original series: f = sin(x)-2sin(4x)+3sin(5x)-4sin(6x)')
for j in range(7):
print( 'Coefficient of sin(%ix):'%j, fCoeffs[j]/N)
print('See Figure for calculating second derivative')
#second derivative is also a sine series
d2f_dx2_a = -sin(x)+32*sin(4*x)-75*sin(5*x)+144*sin(6*x)
#second derivative using fourier transform
DerivativeCoeffs = -arange(N)**2*fCoeffs
d2f_dx2_b = idst(DerivativeCoeffs)
subplot(2,1,1)
plot(x,d2f_dx2_a, x, d2f_dx2_b,'+')
xlabel('x')
legend(('analytic','using DST'))
subplot(2,1,2)
plot(x,abs(d2f_dx2_a-d2f_dx2_b))
xlabel('x')
ylabel('abs(diff)')
show()
psi = np.zeros(len(x),complex)
# SET BOUNDARY INITIAL CONDITIONS FOR PSI
psi0 = 0.
psi1 = 0.
psi_init = np.exp(-(x[1:len(x)-2] - x0)**2/(2*sig**2))*np.exp(1j*kap*x[1:len(x)-2])
# CREATE INITAL PSI ARRAY
psi[0] = psi0
psi[-1] = psi1
psi[1:len(x)-2] = psi_init
# DIVIDE INTO REAL AND IMAGINARY PARTS
psiR = psi.real
psiI = psi.imag
# PERFORM DISCRETE SINE TRANSFORMS
alpha = dst(psiR)
eta = dst(psiI)
################################## PLOT ##################################
plt.plot(np.abs(alpha + 1j*eta))
plt.ylabel('$|b_k|$',fontsize = 20)
plt.xlabel('$k$',fontsize = 20)
plt.title('Fourier Coefficients')
a = L / num
# initial velocity psi0
def psi0(x):
return C * (x * (L - x) / L**2) * exp(-(x - d)**2 / (2 * theta**2))
# the array of the piano string
pstr = np.arange(0, L + a, a)
#set the initial condition
psi = psi0(pstr)
phi = np.zeros(num + 1)
phik = dst(phi) / num
psik = dst(psi) / num
omegak = pstr / a * np.pi * v / L
n = 0
N = 10000
plt.figure(figsize=(20, 4))
t = 2e-3 #Time spot for plots
phi_all = np.zeros([num + 1, num + 1], float)
for i in range(num):
phi_all[i, :] = np.sin(
(i + 1) * np.pi * pstr /
L) * (phik[i + 1] * np.cos(omegak[i + 1] * t) +
psik[i + 1] / omegak[i + 1] * np.sin(omegak[i + 1] * t))
phi = phi_all.sum(axis=0)
def psi0(x):
return C * x * (L - x) * np.exp(-(x - d)**2 / (2 * sigma**2)) / L**2
#grid points
N = 101
#grid spacing
a = L / N
#create arrays containing the initial conditions
x = np.linspace(0, L, N)
psi_0 = psi0(x)
phi_0 = np.zeros(N, float)
#compute the coefficients of the intial state
psi_coeff = dcst.dst(psi_0)
phi_coeff = dcst.dst(phi_0)
#containers for the coefficients returned by the discrete sine transform
sol_coeff = []
#array of times to be plotted
t = np.array([2, 4, 6, 12, 100]) / 1000
#container for the solutions at the desired times
solutions = []
#loops over the desired times and compute the solution to the wave equation
for t_i in t:
#compute the omega values for each coefficient
k = np.arange(0, N, 1)