# Project 5 for the University of Tulsa's CS-7863 Sci-Stat Course
# Systems of Equations and Finite Difference Methods
# Professor: Dr. McKinney, Spring 2023
# Noah L. Schrick - 1492657

## 1. Systems of Equations in Matrix Form
# a. Use built-in R solve
A <- matrix(c(2,3,4,3,2,-3,4,4,2),nrow=3)
b <- matrix(c(3,5,9), nrow=3)

x <- solve(A,b)
x

# b. Verify with matrix mult
all.equal(b, as.matrix(A %*% x))

## 2. Matrix from polynomial
# a. Create vectors for x and y
xvec <- c(-3,-1.5,0.5,2,5)
yvec <- c(6.8,15.2,14.5,-21.2,10)

# b. Create coefficient matrix
A.2 <- matrix(ncol=length(xvec), nrow=length(xvec))
for (i in 1:length(xvec)){
  A.2[,i] = xvec^(i-1)
}
A.2

# c. Solve 
if (!require("matlib")) install.packages("matlib")
library(matlib)

# verify
bvec <- solve(A.2, yvec)
bvec
all.equal(as.matrix(yvec), A.2 %*% bvec) # use all.equal instead of == (tols and storage type)
# EX: identical(as.double(8), as.integer(8)) returns FALSE

# d. Plot
poly_predict <- function(x){ 
  # given input x, returns polynomial prediction 
  sum(bvec*c(1,x,x^2,x^3,x^4))
}

xdomain <- seq(min(xvec),max(xvec),.2)    # predicted domain
y.predict <- sapply(xdomain,poly_predict) # predicted range
plot(xvec,yvec,ylim=c(-50,20))    # plot data
lines(xdomain,y.predict,type="l") # overlay solid line

## 3. Kirchoff
# a. Create vectors
# Resistance vec
Rvec <- c(5,10,5,15,10,20)
# Matrix A and vector b
A.3a <- matrix(c(Rvec[1],0,0,0,Rvec[5],Rvec[6],
                0,Rvec[2],Rvec[3],Rvec[4],-Rvec[5],0,
                1,-1,0,0,-1,0,
                0,1,-1,0,0,0,
                0,0,1,-1,0,0,
                0,0,0,1,1,-1)
              ,nrow=6)
A.3a <- t(A.3a) # easier to conf that A is correct if def matrix as above then t()
V=200
b.3vec <- c(V,0,0,0,0,0)

# Solve
i.a <- solve(A.3a,b.3vec)
i.a

# b. Repeat a. Use V=200, and R=(5,10,5,15,0,20)
Rvec.b <- c(5,10,5,15,0,20)
A.3b <- matrix(c(Rvec.b[1],0,0,0,Rvec.b[5],Rvec.b[6],
                0,Rvec.b[2],Rvec.b[3],Rvec.b[4],-Rvec.b[5],0,
                1,-1,0,0,-1,0,
                0,1,-1,0,0,0,
                0,0,1,-1,0,0,
                0,0,0,1,1,-1)
              ,nrow=6)
A.3b <- t(A.3b) # easier to conf that A is correct if def matrix as above then t()

i.b <- solve(A.3b,b.3vec)
i.b

## 4. Schrodinger eq
if (!require("pracma")) install.packages("pracma")
library(pracma)

#a. Solve and plot the 1d quantum harmonic oscillator wave function
quantumHO <- function(x,y,E){
  psi     <- y[1]
  y1p <- y[2]   # y1’ fill in blank
  y2p <- (x^2)*y[1]-2*E*y[1]  # y2’ fill in blank

  as.matrix(c(y1p,y2p)) 
}

if (!require("reshape")) install.packages("reshape")
library(reshape)
if (!require("ggplot2")) install.packages("ggplot2")
library(ggplot2)
plotQHO <- function(E){
  # plot the wavefunction for a given energy
  QHO.sol <- ode45(function(x,y){quantumHO(x,y,E)}, 
                   y0=c(1,0),t0=0, tfinal=5)
  QHO.melted <- melt(data.frame(position=QHO.sol$t,
                                psi=QHO.sol$y[,1],
                                psi_p=QHO.sol$y[,2]), 
                     id = "position")
  colnames(QHO.melted)[2] <- "Solutions"  # rename "variable" 
  
  g <- ggplot(data = QHO.melted, aes(x=position, y=value, color=Solutions))
  g <- g + geom_line(size=1) + geom_abline(slope=0, intercept=0)
  g <- g + xlab("position") +  ylab("Amplitude") 
  g <- g + ggtitle("Quantum Harmonic Oscillator")  
  show(g)
}

plotQHO(0.5)

# b. Solve with Shooting Method
## Shooting Method Objective Function
QHO_shoot <- function(E){
  # E is an energy guess
  ho.sol <- ode45(function(x,y){quantumHO(x,y,E)}, 
                  y0=c(1,0),t0=0, tfinal=5)
  last_idx <- length(ho.sol$t)
  psi_right <- ho.sol$y[,1][last_idx] # psi at last x
  # We want the wave function at the right boundary to be 
  # exponentially small. So, the log-abs value will be large 
  # and negative for the correct solution.
  return(log(abs(psi_right)))  
}

## The objective function plot informs the energy guesses below.
energy_vals <- seq(0,5.,len=100)
objective_vals <- sapply(energy_vals, QHO_shoot)
plot(energy_vals,objective_vals, xlab="Energy", ylab="Objective")
grid()

# “optimize” finds min or max of function in an interval.
# Find ground state energy E0.
E0<-optimize(QHO_shoot,interval=c(0,1),tol=1e-8)$minimum
E0

# Find first excited state E1.
E1<-optimize(QHO_shoot,interval=c(2,3),tol=1e-8)$minimum
E1
plotQHO(E1)

## 5. Solve harmonic oscillator Schrodinger with finite difference
# a. Solve with R and Julia
if (!require("tictoc")) install.packages("tictoc")
library(tictoc)

solve_QHO <- function(n, xL, xR){
  #n<-1000
  #xL <- -10  # need to change for very excited state
  #xR <- 10
  h<- (xR-xL)/(n-2)  # step size
  x <- seq(xL, xR, by=h) # grid of x values
  
  H <- diag(1/h^2 + .5*x^2) # makes the diagnonal
  dim(H)
  #   diag       upper diag              lower diag
  H <- H + Diag(rep(-1/(2*h^2),n-2),1) + Diag(rep(-1/(2*h^2),n-2),-1)
  
  H.eigs <- eigen(H)
  vals <- H.eigs$values
  vecs <- H.eigs$vectors
  
  return(list(x=x,vecs=vecs,vals=vals))
}

print_QHO_sol <- function(x,vecs,vals,n){
  print("Smallest eigvalue, lowest energy:")
  print(vals[n-1])
  #vals[n-1]  # smallest eigvalue, lowest energy
  plot(x,vecs[,n-1], type="l", main="Ground state wave function") # ground state wave function
  
  print("2nd smallest eigvalue, first excited energy:")
  print(vals[n-2])
  #vals[n-2]  # 2nd smallest eigvalue, first excited energy
  plot(x,vecs[,n-2], type="l", main="First excited wave function") # first excited wave function
  
  print("3rd smallest eigvalue, second excited energy:")
  print(vals[n-3])
  #vals[n-3]  # 3rd smallest eigvalue, second excited energy
  plot(x,vecs[,n-3], type="l", main="Second excited wave function") # second excited wave function
}

# data for timing comparison
timing.table <- matrix(nrow = 0, ncol = 4)  
for (n in seq(500, 3500, by=500)){
  for (x in seq(1,10)){
    tic(quiet=TRUE)
    QHO_sol <- solve_QHO(n, -x, x)
    t <- toc(quiet=TRUE)
    timing.table <- rbind(timing.table, c(-x, x, n, t$callback_msg)) 
  }
}

# answer to main question
QHO_sol <- solve_QHO(1000, -5, 5)
print_QHO_sol(QHO_sol$x, QHO_sol$vecs,QHO_sol$vals,1000)

# b. Solve damped oscillator with perturbation param
yo <- 0
y1 <- 1
tmin <- 0
tfinal <- 2

dho.pert.encap.f <- function(ep){
  dho.pert.f <- function(t, y){
    yn <- y[1]
    ydot <- y[2] - yn
    ypp <- (-(1+ep)*yn - ydot)/ep
    as.matrix(c(ydot,ypp))
  }
  
  proj.obj <- function(v0, y0=yo, tfinal){
    # minimize w.r.t. v0
    proj.sol <- ode45(dho.pert.f, 
                      y=c(yo, y1), t0=tmin, tfinal=tfinal)
    final_index <- length(proj.sol$t)
    yf <- proj.sol$y[final_index,1]  # want equal to right boundary 
    log(abs(yf)) # minimize this
  }
  
  # user specifies tfinal and yfinal for BVP
  ydot_best <- optimize(proj.obj, 
                        interval=c(1,100), #bisect-esque interval
                        tol=1e-10, 
                        y0=0, tfinal=2) # un-optimized obj params
  ydot_best$minimum  # best ydot
  
  best.sol <- rk4sys(dho.pert.f, a=0, b=2, y0=c(0, ydot_best$minimum),
                     n=20)  # 20 integration stepstmax 
}

approx <- function(x){
  exp(1-(x/ep)) + exp(1-x)
}

ep <- 0.5
ep.a <- dho.pert.encap.f(ep)
plot(ep.a$x, ep.a$y[,1], type="l", main="Epsilon=0.5")
par(new=T)
lines(ep.a$x, approx(seq(0,2,len=21)), type="l", col="red")

ep <- 0.05
ep.b <- dho.pert.encap.f(ep)
plot(ep.b$x, ep.b$y[,1], type="l", main="Epsilon=0.05")
par(new=T)
lines(ep.b$x, approx(seq(0,2,len=21)), type="l", col="red")