yo yo yo

sdf.py

__author__ = "Chester Balingit"

import numpy as np
from scipy import optimize
import matplotlib.pyplot as plt


## Part 1 - Find all the zeros of the function: 
## f(x) = cos(x) - xsin(x) using the optimize.brentq
## function of the scipy library

def func1(x):
    return np.cos(x) - x*np.sin(x)


def rooter(func,a,b,spacing):
    root = []
    z = 0
    for all in np.arange(a,b,spacing):
        try: z = optimize.brentq(func,all,all+spacing)
        except: pass
        if z not in root and z != 0:   root.append(z)
    return root

## Part 2 - Given parameters, plot the given function 
## below as a function of E ; f = f(E)

# Initial Parameters
a = 0.3e-9
mass = 9.11e-31
v0 = 10*1.6e-19
hbar = 1.05e-34

# Main Functions

def beta(E):
    return np.sqrt(2*mass*(v0-E)/hbar**2)
    
def alpha(E):
    return np.sqrt(2*mass*E/hbar**2)

def func_even(E):
    return beta(E)*np.cos(alpha(E)*a) - alpha(E)*np.sin(alpha(E)*a)

def func_odd(E):
    return alpha(E)*np.cos(alpha(E)*a) + beta(E)*np.sin(alpha(E)*a)

## Condition checks to verify values of alpha and beta
## by relating the plot of tan(alpha*a) and (beta/alpha)
## versus alpha. Plot should be similar to Fig. 2.9
## plt.plot(alpha,alpha*a)
## plt.plot(alpha,beta/alpha)

E = np.linspace(0,v0,1000)
plt.plot(E, func_even(E))
plt.show()

## Problem 3 - Using the parameters from the problem above, 
## find the lowest even, odd solutions to the square well
## problem.

# Evaluate first the zero values for for the function  
# f(E) = beta*cos(a*alpha) - alpha*sin(a*alpha) = 0
# using optimize.brentq

def psyfunc1(x, beta, C):
    return C*np.exp(beta*x)

def psyfunc2(x, alpha, A, B):
    return A*np.sin(alpha*x) + B*np.cos(alpha*x)

def psyfunc3(x, beta, F):
    return F*np.exp(-beta*x)
    

roots_odd = rooter(func_odd,0,10*1.6e-19,1e-21)
roots_even = rooter(func_even,0,10*1.6e-19,1e-21)


#############

plt.subplot(411)
alphax, betax = alpha(roots_odd[0]), beta(roots_odd[0])
C = 1
A = np.exp(-betax*a)/np.cos(alphax*a)
B = 0
F = -1

x1 = np.linspace(-2*a,-a,10000)
plt.plot(x1, psyfunc1(x1, betax, C),"k")

x2 = np.linspace(-a,a,10000)
plt.plot(x2, psyfunc2(x2,alphax,A,B),"k")

x3 = np.linspace(a,2*a,10000)
plt.plot(x3, psyfunc3(x3, betax, F),"k")

############

plt.subplot(412)
alphax, betax = alpha(roots_odd[1]), beta(roots_odd[1])
C = 1
A = np.exp(-betax*a)/np.cos(alphax*a)
B = 0
F = -1

x1 = np.linspace(-2*a,-a,10000)
plt.plot(x1, psyfunc1(x1, betax, C),"k")

x2 = np.linspace(-a,a,10000)
plt.plot(x2, psyfunc2(x2,alphax,A,B),"k")

x3 = np.linspace(a,2*a,10000)
plt.plot(x3, psyfunc3(x3, betax, F),"k")

############

plt.subplot(413)
alphax, betax = alpha(roots_even[0]), beta(roots_even[0])
C = 0.5
A = 0
B = np.exp(-betax*a)/np.sin(alphax*a)
F = 0.5

x1 = np.linspace(-2*a,-a,10000)
plt.plot(x1, psyfunc1(x1, betax, C),"k")

x2 = np.linspace(-a,a,10000)
plt.plot(x2, psyfunc2(x2,alphax,A,B),"k")

x3 = np.linspace(a,2*a,10000)
plt.plot(x3, psyfunc3(x3, betax, F),"k")
############

plt.subplot(414)
alphax, betax = alpha(roots_even[1]), beta(roots_even[1])
C = -1
A = 0
B = -np.exp(-betax*a)/np.sin(alphax*a)
F = -1

x1 = np.linspace(-2*a,-a,10000)
plt.plot(x1, psyfunc1(x1, betax, C),"k")

x2 = np.linspace(-a,a,10000)
plt.plot(x2, psyfunc2(x2,alphax,A,B),"k")

x3 = np.linspace(a,2*a,10000)
plt.plot(x3, psyfunc3(x3, betax, F),"k")

plt.show()