##=================================================
##  PRACTICUM:  Assessing Potential Nonstationarity 
##  		in Spatial Data
##  PART 1:  Simulating Nonstationary Spatial Data
##
##  Pan-American Advanced Study on Spatio-Temporal 
##  Statistics Institute (PASI) 
##  Buzios, RJ, Brazil,
##  June 16-26, 2014
##
##  Kate Calder (calder@stat.osu.edu) and 
##		and Mark Risser (risser.13@osu.edu)
##  The Ohio State University, Columbus, OH
##  June 2014
##
##=================================================


##=================================================
## The model used for this simulation is a version 
## of the general nonstationary covariance function
## given by Stein, 2005 (unpublished technical 
## report). See PDF accompaniment (overview.pdf)
## for full details.
##
## As the focus is on the covariance structure of
## the model, the mean structure will be assumed
## zero.  We also assume that the nugget, spatial 
## variance, and smoothness parameter are constant 
## over space.
##=================================================

## Necessary libraries.
library(MASS)
library(fields)
library(ellipse)
library(geoR)

##=================================================
## Necessary functions.

# Kernel matrix function
Sigma_x <- function( psi_mean, psi_var, psi_range, locations ){
	
	# kern_mean are the mean parameters for the kernel matrix stationary GP
	# kern_var are the variance parameters for the kernel matrix stationary GP
	# kern_range are the range parameters for the kernel matrix stationary GP
	# locations is a Nx2 vector of observation sites
	
	N = dim(locations)[1]	
	Sigma = array(rep(NA,2*2*N),dim=c(2,2,N))
	
	# Exponential covariances for each component of the decomposition
	C1 = matrix(rep(0,N*N),nrow=N)
	C2 = matrix(rep(0,N*N),nrow=N)
	C3 = matrix(rep(0,N*N),nrow=N)
	
	# Calculate the covariance matrix for each GP
	for(i in 1:N){	
		C1[i,i] = psi_var[1]
		C2[i,i] = psi_var[2]
		C3[i,i] = psi_var[3]
			
		if(i < N){
			for(j in (i+1):N){ 		# Off-diagonal elements
				
				dist_ij = sqrt(sum((locations[i,]-locations[j,])^2))
				
				C1[i,j] = psi_var[1]*exp(-dist_ij/psi_range[1])
				C2[i,j] = psi_var[2]*exp(-dist_ij/psi_range[2])
				C3[i,j] = psi_var[3]*exp(-dist_ij/psi_range[3])
				# Matrices are symmetric
				C1[j,i] = C1[i,j] 
				C2[j,i] = C2[i,j] 
				C3[j,i] = C3[i,j] 
			}
		}
	}
	
	# Draw components of the kernel matrices
	p1 = mvrnorm(1,rep(psi_mean[1],N),C1)
	p2 = mvrnorm(1,rep(psi_mean[2],N),C2)
	p3 = mvrnorm(1,rep(psi_mean[3],N),C3)
	
	# Convert to the kernel matrix parameteriation
	for(k in 1:N){
		
		# Eigenvalues
		Sigma12 = (2/pi)*atan(p3[k])*(sqrt(exp(p1[k]))*sqrt(exp(p2[k])))
		Sigma[,,k] = matrix(c(exp(p1[k]), Sigma12, Sigma12, exp(p2[k])), 2, 2)	
	}
	
	return(Sigma)
		
}

# Covariance matrix for two-dimensional coordinates
NScov <- function( Kern_matrices, sigmasq, locations ){
	
	# Kern_matrices is a 2x2xN array of kernel matrices 
	# sigmasq is proportional to the process variance
	# locations is a Nx2 vector of observation sites
	# A smoothness of nu = 0.5 results in the 
	#	exponential correlation function
	
	N = dim(locations)[1]
	NS.cov = matrix(rep(NA, N^2), nrow=N)
	
	# Calculate the elements of the NxN covariance matrix.
	for(i in 1:N){
		
		# Diagonal elements
		Kerneli = Kern_matrices[,,i]
		det_i = Kerneli[1,1]*Kerneli[2,2] - Kerneli[1,2]*Kerneli[2,1]

		NS.cov[i,i] = sigmasq
		
		Ui = chol(Kerneli)
		
		if(i < N){
		for(j in (i+1):N){ # Off-diagonal elements
			
			Kernelj = Kern_matrices[,,j]
			det_j = Kernelj[1,1]*Kernelj[2,2] - Kernelj[1,2]*Kernelj[2,1]

			avg_ij = 0.5*(Kerneli + Kernelj)
			Uij = chol(avg_ij)			
			det_ij = avg_ij[1,1]*avg_ij[2,2] - avg_ij[1,2]*avg_ij[2,1]				
			vec_ij = backsolve(Uij, (locations[i,]-locations[j,]), transpose = TRUE)
		
			NS.cov[i,j] = sigmasq*(((det_i*det_j)^0.25) / sqrt(det_ij)) * 
							exp(-sqrt( sum(vec_ij^2)))
			
			# NS.cov is a symmetric matrix							
			NS.cov[j,i] = NS.cov[i,j]
		}
		}		
	}
	
	return(NS.cov)
}

# Function to plot the covariance/correlation over the spatial region.
plot.dependence <- function( type, ref_x, ref_y, all_x, all_y, locations, Cov_matrix ){
	
	# type = covariance or correlation
	# (ref_x, ref_y) is the (x,y) of the reference point
	# all_x and all_y give all of the (x,y) values
	# locations give the grid of locations
	# Cov_matrix is the covariance matrix for all locations
	
		
	N <- dim(locations)[1]
	index <- (1:N)[(locations[,1] == ref_x) & (locations[,2] == ref_y)]
	if( type == "covariance"){
		
		image.plot(all_x,all_y,matrix(Cov_matrix[index,], length(all_x), length(all_y)), main="Covariance", xlab="x", ylab="y", asp=T)
		
	}
	if( type == "correlation"){

		image.plot(all_x,all_y,matrix(Cov_matrix[index,], length(all_x), length(all_y))/Cov_matrix[index,index], main="Correlation", xlab="x", ylab="y", asp=T)
		
		
	}
}

##=================================================


##=================================================
## Generate data from this model

## True parameter values
mu_true = 0						# Process mean
tausq_true = .1					# Nugget variance
sigmasq_true = 1					# Process variance
# nu_true = 0.5					# Process smoothness (fixed to 0.5)
psi_mean = c(-2.3, -1.9, 1)	    # Mean parameters for Psi (m)
psi_var = c(1, 1, 1)				# Variance parameter for Psi (v)
psi_range = c(4, 4, 4)			# Range parameter for Psi (r)

##==============
## Two dimensional locations: grid
spacing = 0.2					# Desired spacing
x_length = 5					# Maximum longitude
y_length = 5					# Maximum latitude

grid_x = seq(from=0, to=x_length, by=spacing)
grid_y = seq(from=0, to=y_length, by=spacing)
grid_locations = expand.grid(grid_x, grid_y)
grid_locations = matrix(c(grid_locations[,1], grid_locations[,2]), ncol=2, byrow=F)

N = dim(grid_locations)[1]		# Number of observations

# Look at the grid locations
plot(grid_locations[,1],grid_locations[,2],pch="+",xlab="x",ylab="y",main="Spatial Domain", asp=T)

##==============
## Two dimensional locations: irregular spacing

# N and region size will remain as above.

set.seed(2)
irreg_x = x_length*runif(N)
irreg_y = y_length*runif(N)

irreg_locations = matrix(c(irreg_x,irreg_y), ncol=2, byrow=F)

# Look at the irregularly spaced locations
plot(irreg_locations[,1],irreg_locations[,2],pch="+",xlab="x",ylab="y",main="Spatial Domain", asp=T)

##==============
## Calculate kernel matrices and marginal likelihood covariance

#=======
# For gridded locations
set.seed(10)
Kernels_grid = Sigma_x(psi_mean, psi_var,psi_range,grid_locations)
Omega_grid = NScov( Kernels_grid, sigmasq_true, grid_locations )

# Subgrid for plotting ellipses
plot_subgrid <- seq(1, N, by=10)

# Look at the kernel matrices, as represented by their corresponding 
# 50% probability ellipses
par(mfrow=c(1,1))
plot(0, 0,
 		type='n', asp=T,
 		xlim=c(0,x_length),
 		ylim=c(0,y_length),
 		xlab="",
 		ylab="")
for(j in 1:length(plot_subgrid))
	{
		lines(ellipse(Kernels_grid[,,plot_subgrid[j]],
		 	centre=grid_locations[plot_subgrid[j],], level=.5))
	}
	
# Plot the implied spatial dependence over the region of interest for
# three reference points

# Covariance plots
par(mfrow=c(1,3))
plot.dependence( type="covariance", 1, 1, grid_x, grid_y, grid_locations, Omega_grid )
plot.dependence( type="covariance", 4, 2, grid_x, grid_y, grid_locations, Omega_grid )
plot.dependence( type="covariance", 3.2, 4, grid_x, grid_y, grid_locations, Omega_grid )

# Correlation plots
par(mfrow=c(1,3))
plot.dependence( type="correlation", 1, 1, grid_x, grid_y, grid_locations, Omega_grid )
plot.dependence( type="correlation", 4, 2, grid_x, grid_y, grid_locations, Omega_grid )
plot.dependence( type="correlation", 3.2, 4, grid_x, grid_y, grid_locations, Omega_grid )

#=======
# For irregularly spaced locations
set.seed(24)
Kernels_irreg = Sigma_x(psi_mean, psi_var, psi_range, irreg_locations)
Omega_irreg = NScov( Kernels_irreg, sigmasq_true, irreg_locations )

# Subgrid for plotting ellipses
plot_subgrid <- seq(1, N, by=5)

# Look at the kernel matrices, as represented by their corresponding 
# 50% probability ellipses
par(mfrow=c(1,1))
plot(0, 0,
 		type='n', asp=T,
 		xlim=c(0,x_length),
 		ylim=c(0,y_length),
 		xlab="",
 		ylab="")
for(j in 1:length(plot_subgrid))
	{
		lines(ellipse(Kernels_irreg[,,plot_subgrid[j]],
		 	centre=irreg_locations[plot_subgrid[j],], level=.5))
	}


##==============
## Take a draw from the marginal likelihood

# For gridded locations
set.seed(5)
Z_grid = mvrnorm(1, rep(mu_true,N), diag(rep(tausq_true,N)) + Omega_grid )

# For irregularly spaced locations
set.seed(8)
Z_irreg = mvrnorm(1, rep(mu_true,N), diag(rep(tausq_true,N)) + Omega_irreg )

# Look at the sampled data
par(mfrow=c(1,2))
quilt.plot(grid_locations, 
	Z_grid, 
	main="Observed data at gridded locations", 
	asp=T)
quilt.plot(irreg_locations, 
	Z_irreg, 
	main="Observed data at irregularly spaced locations", 
	asp=T)

##==============
## Make geodata objects

Z_grid.geo <- as.geodata( cbind(grid_locations, Z_grid) )
Z_irreg.geo <- as.geodata( cbind(irreg_locations, Z_irreg) )