#Example Heat Equation (Implicit PDE Heat Equation)
library(plot3D)
library(MASS)
library(matlib)

#Eq: du/dt = -0.5 d2u/dx2 on [0,T]x[0,L]

########################################################
#Domain
T <- 1
L <- 1

#step size
dt <- 0.1
dx <- 0.01
delta <- dt/(dx)^2
#grid
time <- seq(0,T,by=dt)
space <- seq(0,L,by=dx)
Nt <- length(time)-1
Nx <- length(space)-1

#PDE approximation functions
Aij <- function(t,x) -0.5*delta
Bij <- function(t,x) 1+delta
Cij <- function(t,x) -0.5*delta

#initial and boundary conditions
i.cond.x <- function(x) x^2
b.cond.d <- function(t) sqrt(t)
b.cond.u <- function(t) cos(10*t)

#matrix where we save the solution points
known <- unknown <- matrix(rep(0,Nx-1), nrow=Nx-1)
sol <- matrix(rep(0,(Nt+1)*(Nx+1)), nrow=Nt+1)

sol[1,] <- sapply(space,i.cond.x) #rep(E,Nt+1)
sol[,1] <- sapply(time,b.cond.d)
sol[,Nx+1] <- sapply(time,b.cond.u)

unknown <- matrix(sol[1,2:Nx], nrow=Nx-1)

for(i in 2:(Nt+1)){
  print(i)
  known <- unknown
  MS <- matrix(rep(0,(Nx-1)*(Nx-1)), nrow=Nx-1)
  for(j in 1:(Nx-1)){
    for(k in 1:(Nx-1)){
      if(k==j-1){ MS[j,k] <- Aij(time[i],space[j+1]) }
      if(k==j){ MS[j,k] <- Bij(time[i],space[j+1]) }
      if(k==j+1){ MS[j,k] <- Cij(time[i],space[j+1]) } 
    }
  }
  E <- matrix( c(Aij(time[i],space[2])*b.cond.d(time[i]),rep(0,Nx-3),Cij(time[i],space[Nx])*b.cond.u(time[i])) , nrow=Nx-1)
  unknown <- inv(MS)%*%(known-E)
  sol[i,2:Nx] <- unknown
}

persp3D(time,space,sol, theta=310, phi=30, expand=0.5, ticktype="detailed", clab = c("Value"), xlab="time", ylab="space", zlab="", axes=TRUE)