# 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)
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

# 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

# b. Solve damped oscillator with perturbation param
# Plot

## 6. Solve heat diffusion equation

# EC. Solve the square plate problem