#############################################################################
## Uncertainty Analysis in Log-Linear Regression
##
## Joel Sholtes, March 2015
## Dept. Civil and Env. Engr., Colorado State University
##
## This code imports sediment load-discharge bivariate data to create
## log-linear regression models, confidence intervals for the mean response,
## and prediction intervals for predicted responses.
## 
## I use three methods:
##
## 1) Standard OLS techniques with the normally-distributed, equal residual 
##    variance  assumption.
## 2) Robust linear regression using the "MM-estimation" method for mean 
##    response, and bootstrap techniques to calculate the confidence and 
##    prediction intervals. This method does not assume equal residual variance 
##    for the CI but does for the PI. Note that robust regression methods do 
##    assume constant variance of the residuals to fit the regression line.
## 3) Robust Bayesian log-linear model which utilizes a Student's t-distribution 
##    to model the variance of the residuals. This is also an equal variance,
##     symmetricresidual model.
##
## With this code one may load either suspended or bed load sediment transport
## measurements and conduct regression and uncertainty analysis on a site-by-site
## basis.
## 
#############################################################################

#################################
### Pre-Amble: Load Functions ###
#################################

rm(list = ls())
setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Uncertainty")
# path <- "bedload_sites.csv"
#library(Hmisc) # for HD quantile estimator
library(boot) # bootstrap
library(MASS) # RLM
library(snow) # parallel processing
library(Hmisc) # for hdquantile
library(rjags) # for Bayesian Analysis, also loads coda package
source("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/minor.ticks.axis.R")

# Choose whether to analyze Suspended or Bed Load Data
sed.type = "bed"
sed.type = "susp"

############################
### LOAD DATA FOR A SITE ###
############################

# Pick One:

## Suspended Load ##
k = 6 # Potomac
k = 12 # Neuse River
k = 13 # Pee Dee River
k = 16 # Spruce Creek
k = 17 # Bullfrog Ck
k = 22 # Nemadji River
k = 28 # Red River
k = 52 # Yazoo River
k = 74 # Little Snake

## Bed Load ##
k = 4 # Battle Ck
k = 5 # Savery Ck
k = 6 # Big Sandstone Ck
k = 26 # Cat Spur Ck
k = 30 # Big Wood Ck
k = 60 # Eggers Ck

if(sed.type == "susp"){
  ###############################################
  ### Import Suspended Site and Sediment Data ###
  ###############################################
  
  ## Sediment Data (Suspended Sites)
  setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/SuspLoad")
  
  ## Import flow file data
  path <- "susp_sites_v10.csv"
  flow.sites <- read.table(path, header = TRUE, sep = ",", 
                           stringsAsFactors = F)
  flow.sites <- flow.sites[order(flow.sites$site.no),]
  sites.flow <- flow.sites$site.no
  
  # Read in Susp. Sed Data
  # Import Cleaned NWIS Suspended Sed Data (SEE NWIS_Susp_Data_Cleaning.R)
  sed.data <- read.table("susp_data_clean_outliers.csv", stringsAsFactors = F,header = TRUE, sep = ",",
                         colClasses = c("character", "character", "character", "numeric", 
                                        "numeric", "numeric","numeric")
  )
  
  sed.data$SITE_NO <- as.integer(sed.data$SITE_NO)
  
  # match with most recent list of sites to use
  sed.data <- subset(sed.data, sed.data$SITE_NO %in% sites.flow)
  
  # extracts unique site.id's
  sites.sed <- sort(as.integer(unique(sed.data$SITE_NO)))
  
  # Check that flow and sed sites match up
  sum(flow.sites$site.no == sites.sed) == nrow(flow.sites)
  
  # Extract sed data for regression
  site.data <- subset(sed.data, SITE_NO == sites.sed[k])
  n <- nrow(site.data)
  
  x <- site.data$INSTFLOW * 0.0283168466 # convert CFS to CMS
  q.min.sed <- signif(min(x, na.rm= T),2)
  q.max.sed <- round(max(x, na.rm= T),2)
  
  # here I am converting from [mg/L] SSC to mg/L of susp. sand (P63 is % finer than 0.0625 mm)
  # 0.001 converts to kg/s: [mg/L*m^3/s * 1000L/m^3 * 1kg/1E*6mg = kg/s] 
  y <- site.data$SSC * (1 - site.data$P63/100) * x * 0.001
  
  # Write data to table along with CI for regression line
  sed.site <- data.frame(cms = x, kg.sec = y, log.q = log10(x), log.qs = log10(y))
  
  # now remove zero values for log-linear regression
  if (min(y) == 0) {
    sed.site <- sed.site[which(y != 0),]
  }

  ###############################################
  ### Import Bedload Site and Sediment Data ###
  ###############################################
  
  ## Sediment Data (Bedload Sites)
  setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/Bedload/")
  bed.site.info <- read.table("site_summary.csv", stringsAsFactors = F, header = TRUE, sep = ",")
  
  #compiled bedload data
  sed.data <- read.table("BedLoad_R_Data.csv", stringsAsFactors = F,header = TRUE, sep = ",")
  
  # Remove zero discharge load observations for log-transformation
  sed.data <- subset(sed.data, !sed.data$cms == 0)
  
  # extracts unique site.id's
  sites <- bed.site.info$site.no
  
  ############################
  ### LOAD DATA FOR A SITE ###
  ############################

  # Extract sed data for regression
  site.data <- subset(sed.data, site.id == sites[k])
  n <- nrow(site.data)
  
  # extract flow
  x <- site.data$cms
  q.min.sed <- signif(min(x, na.rm= T),2)
  q.max.sed <- round(max(x, na.rm= T),2)
  
  # Extract sediment load 
  y <- site.data$kg.sec 
  
  # Write data to table along with CI for regression line
  sed.site <- data.frame(cms = x, kg.sec = y, log.q = log10(x), log.qs = log10(y))
  
  # now remove zero values for log-linear regression
  if (min(y) == 0) {
    sed.site <- sed.site[which(y != 0),]
  }
}


######################
### Log-Linear OLS ###
######################

# OLS Log-Linear Regression, needs to call from a data.frame for predict() to work
model.ols <- lm(log.qs ~ log.q, data = sed.site)
b0 <- as.numeric(round(coef(model.ols)[1], 3))
b1 <- as.numeric(round(coef(model.ols)[2], 3))
bcf.qmle <- signif(exp(sum( residuals(model.ols)^2 )/(n-2) /2), 3)
r2 <- signif(summary(model.ols)$r.squared, 2)
sigma.e <- signif(summary(model.ols)$sigma, 2)

# Create confidence and Prediction Intervals
x.pred1 <- data.frame(log.q = seq(min(sed.site$log.q), max(sed.site$log.q), 0.01))

y.conf <- predict(model.ols, newdata = x.pred1, interval = "confidence") 
y.pred <- predict(model.ols, x.pred1, interval = "prediction") 


######################
### Log-Linear RLM ###
######################

data <- data.frame(log.q = sed.site$log.q, log.qs = sed.site$log.qs)
R = 1000 # Number of Bootstrap samples
x.range = range(data$log.q)

# x-vector for linear regression mean response
x1 <- seq(x.range[1], x.range[2], length.out = 100) 

# original RLM model
model.rlm <- rlm(log.qs ~ log.q, data = data, method = "MM", maxit = 100)
bhat = coef(model.rlm)

## Set up Bootstrap Sampling for CI creation: Case Re-sampling for the mean response
# Davison and Hinkley: 6.5 p: 307-315
# Create function for bootstrap sampling
bs <- function(data, indices, x) {
  # Allow boot to select sample 
  data.boot <- data[indices,]
  fit.boot <- rlm(data.boot[,2] ~ data.boot[,1], method = "MM", maxit = 100)
  coef.boot <- coef(fit.boot)
  y.boot <- coef.boot[1] + x*coef.boot[2]
  
  return(c(y.boot)) 
} 

boot.out1 <- boot(data, statistic = bs, R = R, parallel = "snow", x = x1)

# loop for calc BCa for pointwise CI of mean response
n = length(x1) # length of mean response vector in bs function
CI.rlm = matrix(ncol = 2, nrow = n)
for(i in 1:n){
  # calc BCa CI for each boot strapped set of y's at each x in regression interval
  # boot.ci returns an object (list) with a bca vector as a member
  # the 4th and 5th elements of this vector are the lower and upper bounds, respectively.
  CI.rlm[i,] <- boot.ci(boot.out1, conf = 0.95, type = "bca", index = i)$bca[4:5]
}

## Set up Bootstrap Sampling for PI creation: Model Re-sampling for the predicted response
# Davison and Hinkley: 6.5 p: 307-315

# 1) Calc standardized, centered residuals as per Davison and Hinkley 1997, Section 6.5
e <- resid(model.rlm)

# Leverage Correction for Raw Residuals
s <- summary(model.rlm)$sigma
# square root of residual square mean? s = sqrt( sum(e^2) / (n-2) ) ?
# sigma.e?  sigma? stdev?
d <- 2*sum(e/s * psi.bisquare(e/s, c = 4.685, deriv = 0) ) / 
     sum( psi.bisquare(e/s, c = 4.685, deriv = 1) ) -
     sum( psi.bisquare(e/s, c = 4.685, deriv = 0)^2 ) / 
     sum( psi.bisquare(e/s, c = 4.685, deriv = 1) )^2

#leverage
h <- hat(model.matrix(model.rlm)) # leverages for each x.j
r <- e / sqrt(1-d*h) #Leverage corrected residuals
r.cent <- r - mean(r) # center the modified residuals

#plot resids
plot(density(e),lwd = 4, col = "grey")
lines(density(r)$x,density(r)$y)
lines(density(r.cent)$x,density(r.cent)$y, col = "blue")
abline(v=0)

# create data frame with x vals y vals, fitted y vals, and resids
model.data <- data.frame(log.q = data$log.q, log.qs = data$log.qs, 
                         fit = predict(model.rlm), r = r)

# create function to predict new values based on orig model fit
x2 <- seq(x.range[1], x.range[2], length.out = 25) # shorter length b/c of comp time
x.pred <- data.frame(log.q = x2) 
# x.pred needs to be in a dataframe for predict()
d.pred <- data.frame(log.q = x2, fit = predict(model.rlm, newdata = x.pred)) 

# Bootstrap function
bs <- function(data, i, i.pred, d.pred) {
  
  # Allow boot to select sample from residuals and create new mean resp. val
  rlm.b <- rlm(log.qs + r[i] ~ log.q, data = data, method = "MM", maxit = 100)
  
  # Predict mean response based on booted model, rlm.b
  pred.b <- predict(rlm.b, d.pred)
  
  # Calculate predicted response error: 
  # e*rm = (gridded mean response from booted model) - (orig model gridded mean resp + booted orig resid)
  pred.b - (d.pred$fit + data$r[i.pred])
} 

# Bootstrap for Prediction Interval by sampling residuals (same resids twice!)
boot.out2 <- boot(data = model.data, statistic = bs, 
                 R = 50000, 
                 m = nrow(x.pred), # number of prediction to make (for sampling e*_rm)
                 d.pred = d.pred,
                 parallel = "snow") # gridded x vals for prediction

# loop for calc BCa for pointwise PI of predicted response
n = nrow(x.pred) # length of mean response vector in bs function
PI.rlm = matrix(data = d.pred$fit, ncol = 2, nrow = n)

for(i in 1:n){
  #PI.rlm[i,] <-  PI.rlm[i,] + boot.ci(boot.out2, conf = 0.95, type = "bca", index = i)$bca[4:5]
  PI.rlm[i,] <- PI.rlm[i,] + hdquantile(boot.out2$t[,i], probs = c(0.025, 0.975))
}


##################################
### Bayesian Robust Log-Linear ###
##################################
# Following: Kruschke, J. K. (2011). Doing Bayesian Data Analysis:
# A Tutorial with R and BUGS. Academic Press / Elsevier.

setwd("C:/Users/jsholtes/Dropbox/Bayesian Stats/Kruschke Book/R Code")
source("openGraphSaveGraph.R")
source("plotPost.R")
source("HDIofMCMC.R")
# setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Uncertainty/Bayesian")

## THE MODEL.
modelstring = "
model {
for( i in 1 : Ndata ) {
y[i] ~ dt( mu[i] , tau , tdf )
mu[i] <- beta0 + beta1 * x[i]
}
beta0 ~ dnorm( 0 , 1.0E-6 )
beta1 ~ dnorm( 0 , 1.0E-6 )
# beta0 ~ dunif( 0 , 1 )
# beta1 ~ dunif( 0 , 10 )
tau ~ dgamma( 0.001 , 0.001 )
udf ~ dunif(0,1)
tdf <- 1 - tdfGain * log(1-udf) # tdf in [1,Inf).
# tdfGain specified in data section
}
" # close quote for modelstring
writeLines(modelstring,con="model.txt")

## THE DATA.
nSubj = NROW(sed.site)
x = sed.site$log.q
xName="Q"
y = sed.site$log.qs
yName="Qs"

# Re-center data at mean, to reduce autocorrelation in MCMC sampling.
# Standardize (divide by SD) to make initialization easier.
xM = mean( x ) ; xSD = sd( x )
yM = mean( y ) ; ySD = sd( y )
zx = ( x - xM ) / xSD
zy = ( y - yM ) / ySD

# Specify data, as a list.
tdfGain = 50 # 1 for low-baised tdf, 100 for high-biased tdf
dataList = list(
  x = zx ,
  y = zy ,
  Ndata = nSubj ,
  tdfGain = tdfGain
)

## INTIALIZE THE CHAINS.
r = cor(x,y)
initsList = list(
  beta0 = 0 ,         # because data are standardized
  beta1 = r ,         # because data are standardized
  tau = 1 / (1-r^2) , # because data are standardized
  udf = 0.95          # tdf = 4
)

## RUN THE CHAINS
parameters = c("beta0" , "beta1" , "tau" , "tdf" ) 
adaptSteps = 1000              # Number of steps to "tune" the samplers.
burnInSteps = 1000            # Number of steps to "burn-in" the samplers.
nChains = 3                   # Number of chains to run.
numSavedSteps=50000           # Total number of steps in chains to save.
thinSteps=1                   # Number of steps to "thin" (1=keep every step).
nPerChain = ceiling( ( numSavedSteps * thinSteps ) / nChains ) # Steps per chain.

# Create, initialize, and adapt the model
jagsModel = jags.model( "model.txt" , data=dataList , inits=initsList , 
                        n.chains=nChains , n.adapt=adaptSteps )
# Burn-in
update( jagsModel , n.iter=burnInSteps )

# The saved MCMC chain
codaSamples = coda.samples( jagsModel , variable.names=parameters , 
                            n.iter=nPerChain , thin=thinSteps )
# resulting codaSamples object has these indices: 
#   codaSamples[[ chainIdx ]][ stepIdx , paramIdx ]

## EXAMINE THE RESULTS
checkConvergence = TRUE
if ( checkConvergence ) {
  openGraph(width=7,height=7)
  autocorr.plot( codaSamples[[1]] , ask=FALSE ) 
  show( gelman.diag( codaSamples ) )
  effectiveChainLength = effectiveSize( codaSamples ) 
  show( effectiveChainLength )
}

# Convert coda-object codaSamples to matrix object for easier handling.
# But note that this concatenates the different chains into one long chain.
# Result is mcmcChain[ stepIdx , paramIdx ]
mcmcChain = as.matrix( codaSamples )

# Extract chain values:
tdfSamp = mcmcChain[, "tdf" ]
tdfM = mean( tdfSamp )
z0 = mcmcChain[, "beta0" ]
z1 = mcmcChain[, "beta1" ]
zTau = mcmcChain[, "tau" ]
zSigma = 1 / sqrt( zTau ) # Convert precision to SD

# Convert to original log scale:
b1 = z1 * ySD / xSD
b0 = ( z0 * ySD + yM - z1 * ySD * xM / xSD )
sigma = zSigma * ySD

# Plot Traceplots for each parameter
par(mfrow = c(2,2), pty = "m", mar = c(4,4,1,1))
plot(10000:20000,exp(b0)[10000:20000], type = "l", ylab = "beta0", xlab = "")
plot(10000:20000,b1[10000:20000], type = "l", ylab = "beta1", xlab = "")
plot(10000:20000,sigma[10000:20000], type = "l", ylab = "sigma", xlab = "Sample Number")
plot(10000:20000,log(tdfSamp)[10000:20000], type = "l", ylab = "log(tdf)", xlab = "Sample Number")

# Plot coefficient values to make sure not too correlated.
par(mfrow = c(1,2), mar = c(4,5,0,0), oma = c(1,1,1,1))
z.1 <- kde2d(z0, z1, n=50)
plot(z0,z1, pch = ".",
     xlab = expression(z[0]),
     ylab = expression(z[1]),
     cex.lab = 1.2)
contour(z.1, drawlabels=FALSE, nlevels=15, col="royalblue", add=TRUE)

z.2 <- kde2d(exp(b0), b1, n=50)
plot(exp(b0),b1, pch = ".",
     xlab = expression(beta[0]),
     ylab = expression(beta[1]),
     cex.lab = 1.2)
contour(z.2, drawlabels=FALSE, nlevels=15, col="royalblue", add=TRUE)

##############################
### HDPI for Mean Response ###
##############################

xLim= range(sed.site$log.q)
yLim= range(sed.site$log.qs)
# yLim = c(-10,35)
xMeanResp = seq( xLim[1], xLim[2], length = 50 )

# Define matrix for recording posterior predicted y values at each x value.
# One row per x value, with each row holding random predicted y values.
SampSize = length(b1)
yMeanResp = matrix( 0 , nrow=length(xMeanResp) , ncol=SampSize )
# Define matrix for recording HDI limits of posterior predicted y values:
yHDI.CI = matrix( 0 , nrow=length(xMeanResp) , ncol=2 )

# Generate posterior predicted y values: using model, generate 50001 samples
# of 20 y values at each xPostPredict. This is used to generate the posterior
# distribution at each value of X, hence the median or mean value as well
# as the 95% crediility interval (quantile)
# This gets only one y value, at each x, for each step in the chain.
for ( chainIdx in 1:SampSize ) {
  yMeanResp[,chainIdx] = ( b0[chainIdx] + b1[chainIdx] * xMeanResp )
}

# Calc HDI for Prediction Interval
for ( xIdx in 1:length(xMeanResp) ) {
  yHDI.CI[xIdx,] = HDIofMCMC( yMeanResp[xIdx,] )
}

#############################
### Posterior prediction: ###
#############################

xPostPred = seq( xLim[1] , xLim[2] , length=20)
# Define matrix for recording posterior predicted y values at each x value.
# One row per x value, with each row holding random predicted y values.
postSampSize = length(b1)
yPostPred = matrix( 0 , nrow=length(xPostPred) , ncol=postSampSize )
# Define matrix for recording HDI limits of posterior predicted y values:
yHDI.PI = matrix( 0 , nrow=length(xPostPred) , ncol=2 )

# Generate posterior predicted y values: using model, generate 50001 samples
# of 20 y values at each xPostPredict. This is used to generate the posterior
# distribution at each value of X, hence the median or mean value as well
# as the 95% crediility interval (quantile)

# This generates only one y value, at each x, for each sample in the chain.
for ( chainIdx in 1:postSampSize ) {
  # p. 398 of Kruschke book: dt(x|mu,tau,df) in R is sqrt(tau)*dt((x-mu)*sqrt(tau),df).
  # From R's random number generator rt(n,df), need to multiply by SD 
  # and add mean.
  yPostPred[,chainIdx] = ( rt( length(xPostPred) , df=tdfSamp[chainIdx] )
                           * sigma[chainIdx] 
                           + b0[chainIdx] + b1[chainIdx] * xPostPred )
}

# Calc HDI for Prediction Interval
for ( xIdx in 1:length(xPostPred) ) {
  yHDI.PI[xIdx,] = HDIofMCMC( yPostPred[xIdx,] )
}


#####################################
### PLOTS 
#####################################

### Susp Sites ###

# #Little Snake
x.range = c(-.7, 2.5)
y.range = c(-3.5, 3.8)

# #Spruce Creek
x.range = c(-2.1, 0.9)
y.range = c(-6, 0.1)

#BullFrog Creek
x.range = c(-1, 1.5)
y.range = c(-5.1, 0.1)

# Yazoo River
x.range = c(1.6, 3.4)
y.range = c(-2, 2.4)

# Potomac River
x.range = c(1.2, 3.75)
y.range = c(-3.5, 3.2)

# Red River of the North
x.range = c(0.8, 3.25)
y.range = c(-3.5, 2)

### Bed load sites ###

# Big Sandstone Ck
x.range = c(-0.5, 1)
y.range = c(-5.7, -0.6)

# East Fork Savery Creek near Encampment, WY
x.range = c(-0.7, 0.3)
y.range = c(-4.9, -0.2)

# Eggers Creek, ID
x.range = c(-2.0, -0.6)
y.range = c(-4.9, -1.3)

# BigWood
x.range = c(.75,1.5)
y.range = c(-3.4,0.5)

setwd("C:/Users/jsholtes/Dropbox/Research/Qeff/Uncertainty/Figures")
source("C:/Users/jsholtes/Dropbox/Research/Qeff/Analysis/minor.ticks.axis.R")

pdf("ThreeReg_BigWood1.pdf", width = 8, height = 3.5)

par(pty = "s", mfrow = c(1,3), oma = c(1,6,1,1), mar = c(5,0,0,0))

cex.axis = 1.2
cex.lab = 1.2
lwd = 1.5

# Range of log10 values
xat <- floor(min(sed.site$log.q)):ceiling(max(sed.site$log.q))
yat <- floor(min(y.pred[,2])):ceiling(max(y.pred[,3]))

### PLOTS: OLS ###
plot(sed.site$log.q, sed.site$log.qs, 
     xlim = x.range,
     #ylim = y.range,
     axes=F, 
     xlab = "",
     ylab = "",
     cex.lab = cex.lab
)

labels <- sapply(xat, function(i) as.expression(bquote(10^ .(i))) )
#at1[which(at1 %% 2 == 0)] <- NA
minor.ticks.axis(1,9, mn = min(floor(x.range)), mx = max(ceiling(x.range)), cex = cex.axis)
mtext("Sediment Load (kg/s)", side = 2, line = 4, cex = cex.lab) 

axis(2, at = yat, labels = NA, cex.axis = cex.axis)
# yat[which(yat %% 2 == 0)] <- NA
labels <- sapply(yat, function(i) as.expression(bquote(10^ .(i))) )
axis(2, at = yat, labels = labels, cex.axis = cex.axis)

box()

# Add CI and PI Lines
lines(x.pred1$log.q, y.conf[,1], lwd = lwd)
lines(x.pred1$log.q, y.conf[,2], lty = 2, col = "royalblue", lwd = lwd)
lines(x.pred1$log.q, y.conf[,3], lty = 2, col = "royalblue", lwd = lwd)
lines(x.pred1$log.q, y.pred[,2], lty = 6, col = "seagreen", lwd = lwd)
lines(x.pred1$log.q, y.pred[,3], lty = 6, col = "seagreen", lwd = lwd)

### PLOTS: RLM ####

plot(sed.site$log.q, sed.site$log.qs, type = "n",
     xlim = x.range,
     ylim = y.range,
     axes=F, 
     xlab = "",
     ylab = "",
     cex.lab = cex.lab
)

# site.nm <- paste("0", flow.sites$site.no[k], " - ", flow.sites$site.nm[k], sep = "")
# usr0 <- par("usr")       # get user coordinates
# par(usr = c(0, 1, 0, 1))
# text(.05, .95, labels = site.nm, cex = 0.75, adj = c(0,0))
# par(usr = usr0)

#labels <- sapply(xat, function(i) as.expression(bquote(10^ .(i))) )
#at1[which(at1 %% 2 == 0)] <- NA
#axis(1, at = xat, labels = labels,cex.axis = cex.axis)
minor.ticks.axis(1,9, mn = min(floor(x.range)), mx = max(ceiling(x.range)), cex = cex.axis)

box()

mtext(expression("Discharge ("*m^3/s*")"), side = 1, line = 4, cex = cex.lab)

# Add Bootstrapped lines
for(i in 1:R){
  lines(x1,boot.out1$t[i,], col = "grey")
}

points(data$log.q, data$log.qs)

# Add CI and PI Lines
lines(x1, boot.out1$t0, lwd = lwd)
lines(x1,CI.rlm[,1], lty = 2, col = "royalblue", lwd = lwd)
lines(x1,CI.rlm[,2], lty = 2, col = "royalblue", lwd = lwd)
lines(x2, PI.rlm[,1], lty = 6, col = "seagreen", lwd = lwd)
lines(x2, PI.rlm[,2], lty = 6, col = "seagreen", lwd = lwd)

legend("bottomright",
       c("Bootstrapped Mean Resp."),
       lwd = 1,
       col = c("grey"),
       lty = 1,
       bty = "n",
       cex = .75)

### PLOTS: BAYESIAN ###
plot(sed.site$log.q, sed.site$log.qs, 
     xlim = x.range,
     ylim = y.range,
     axes = F, 
     xlab = "",
     ylab = "",
     cex.lab = cex.lab
)

labels <- sapply(xat, function(i) as.expression(bquote(10^ .(i))) )
#at1[which(at1 %% 2 == 0)] <- NA
minor.ticks.axis(1,9, mn = min(floor(x.range)), mx = max(ceiling(x.range)), cex = cex.axis)
box()

# Superimpose posterior predicted 95% HDIs for Prediction and Mean Response:
lines( xPostPred , apply(yPostPred,1,median) , lwd = lwd)
lines( xMeanResp, yHDI.CI[,1], lwd = lwd, lty = 2, col="royalblue" )
lines( xMeanResp, yHDI.CI[,2], lwd = lwd, lty = 2, col="royalblue" )
lines( xPostPred, yHDI.PI[,1], lwd = lwd, lty = 6, col="seagreen" )
lines( xPostPred, yHDI.PI[,2], lwd = lwd, lty = 6, col="seagreen" )

legend("bottomright",
       c("Mean Response", "95% Conf. Int.", "95% Pred. Int."),
       lwd = 1.5,
       col = c("black", "royalblue", "seagreen"),
       lty = c(1,2,6),
       bty = "n",
       y.intersp = 1,
       cex = .75)

dev.off()


####################################################
## Plot all Confidence / Prediction bands together
###################################################

pdf("ThreeReg_BigWood2.pdf", width = 6, height = 6)
lwd0 = lwd
par(mfrow = c(1,1), mar = c(5,5,0,0), oma = c(1,1,1,1))
plot(sed.site$log.q, sed.site$log.qs,
     xlim = x.range,
     ylim = y.range,
     axes=F, 
     xlab = expression("Discharge ("*m^3/s*")"),
     ylab = "Sediment Load (kg/s)",
     cex.lab = cex.lab
)

### OLS Lines
lwd = 1
lines(x.pred1$log.q, y.conf[,1], lwd = lwd, col = "gray70")
lines(x.pred1$log.q, y.conf[,2], lty = 1, col = "royalblue", lwd = lwd)
lines(x.pred1$log.q, y.conf[,3], lty = 1, col = "royalblue", lwd = lwd)
lines(x.pred1$log.q, y.pred[,2], lty = 1, col = "seagreen", lwd = lwd)
lines(x.pred1$log.q, y.pred[,3], lty = 1, col = "seagreen", lwd = lwd)

### RLMS Lines
lwd = 1
lines(x1, boot.out1$t0, lwd = lwd, col = "gray40", lty = 2)
lines(x1,CI.rlm[,1], lty = 2, col = "royalblue4", lwd = lwd)
lines(x1,CI.rlm[,2], lty = 2, col = "royalblue4", lwd = lwd)
lines(x2, PI.rlm[,1], lty = 2, col = "darkseagreen4", lwd = lwd)
lines(x2, PI.rlm[,2], lty = 2, col = "darkseagreen4", lwd = lwd)

### BAYESIAN Lines
lwd = 1
lines( xPostPred , apply(yPostPred,1,median) , lwd = lwd, lty = 6)
lines( xMeanResp, yHDI.CI[,1], lwd = lwd, lty = 6, col="navy" )
lines( xMeanResp, yHDI.CI[,2], lwd = lwd, lty = 6, col="navy" )
lines( xPostPred, yHDI.PI[,1], lwd = lwd, lty = 6, col="darkgreen" )
lines( xPostPred, yHDI.PI[,2], lwd = lwd, lty = 6, col="darkgreen" )

legend("bottomright",
       c("OLS", "RLM", "Bayesian"),
       lwd = 1.5,
       col = c("gray70", "gray40", "black"),
       lty = c(1,2,6),
       bty = "n",
       y.intersp = 1,
       cex = 0.8)

# labels <- sapply(xat, function(i) as.expression(bquote(10^ .(i))) )
# axis(1, at = xat, labels = labels,cex.axis = cex.axis)
minor.ticks.axis(1,9, mn = min(floor(x.range)), mx = max(ceiling(x.range)), cex = cex.axis)

box()

yat <- floor(min(y.pred[,2])):ceiling(max(y.pred[,3]))
axis(2, at = yat, labels = NA, cex.axis = cex.axis)
# yat[which(yat %% 2 == 0)] <- NA
labels <- sapply(yat, function(i) as.expression(bquote(10^ .(i))) )
axis(2, at = yat, labels = labels, cex.axis = cex.axis)

dev.off()





### Plot regression uncertainty with PDF uncertainty to produce sediment yield curve with uncertainty

par(mfrow = c(2,2), oma = c(3,1,0,0), mar = c(2,5,1,1), pty = "s")
x.range = c(.75,1.5)
y.range = c(-3.4,0.5)
cex.axis = 1.2
cex.lab = 1.2
lwd = 1.5
x.range = c(1.4,3.1)
y.range  = c(-2.25, 1.25)

# Range of log10 values
xat <- floor(min(sed.site$log.q)):ceiling(max(sed.site$log.q))
yat <- floor(min(y.pred[,2])):ceiling(max(y.pred[,3]))

plot(sed.site$log.q, sed.site$log.qs, 
     type = "n",
     xlim = x.range,
     ylim = y.range,
     axes=F, 
     xlab = "",
     ylab = "Sediment Load (kg/s)",
     cex.lab = cex.lab
)
text(x=grconvertX(.1, from='npc'), y = grconvertY(.9, from='npc'), labels = "(a)")
#labels <- sapply(xat, function(i) as.expression(bquote(10^ .(i))) )
#at1[which(at1 %% 2 == 0)] <- NA
minor.ticks.axis(1,9, mn = min(floor(x.range)), mx = max(ceiling(x.range)), cex = cex.axis)
minor.ticks.axis(2,9, mn = min(floor(y.range)), mx = max(ceiling(y.range)), cex = cex.axis)

box()
col2 <- col2rgb("royalblue")/255
for(i in 1:R){
  abline(log10(b0.boot[i]), b1.boot[i], col=rgb(t(col2),alpha = 0.5))
}

# Add CI and PI Polygons
col1 <- col2rgb("seagreen")/255
polygon( c(x.pred1$log.q, rev(x.pred1$log.q)), c(y.pred[,2],rev(y.pred[,3])), col=rgb(t(col1),alpha = 0.5), border=F) 
col2 <- col2rgb("royalblue")/255
polygon( c(x.pred1$log.q, rev(x.pred1$log.q)), c(y.conf[,2],rev(y.conf[,3])), col=rgb(t(col2),alpha = 0.5), border=F) 
lines(x.pred1$log.q, y.conf[,1], lwd = lwd)
points(sed.site$log.q, sed.site$log.qs)

## PLOT KDF ###
# Plot asymptotically-normal CBs
plot(kde$x,kde$y, ylim=range(kde.cb), xlim = c(0,2000), type = "n", ylab = "Probability Density", xlab = "", cex.lab = cex.lab, cex.axis = cex.axis) 

col2 <- col2rgb("firebrick")/255
polygon( c(kde$x, rev(kde$x)), c(kde.cb[1,], rev(kde.cb[2,])), col=rgb(t(col2), alpha = 0.5), border=F) 
lines(kde, lty = 1)

text(x=grconvertX(.9, from='npc'), y = grconvertY(.9, from='npc'), labels = "(b)")

## Plot sed yield curve ##
x.pred3 <- if(min(q) < 1) {
  data.frame(log.q = seq(log10(.00000001)*1.1, max(log10(q))*1.1, 0.005))
}else{
  data.frame(log.q = seq(log10(.00000001)*.9, max(log10(q))*1.1, 0.005))
}

y.conf2 <- as.matrix(predict(model.ols, newdata = x.pred3, interval = "confidence"))
y.pred2 <- as.matrix(predict(model.ols, newdata = x.pred3, interval = "prediction"))

# Sediment Yield Record with uncertainty
yield.mean <- approx(x = x.pred3$log.q, y = y.conf2[,1], xout = log10(q))$y*log10(86.4)
yield.low <- approx(x = x.pred3$log.q, y = y.conf2[,2], xout = log10(q))$y*log10(86.4)
yield.upp <- approx(x = x.pred3$log.q, y = y.conf2[,3], xout = log10(q))$y*log10(86.4)
yield.low.p <- approx(x = x.pred3$log.q, y = y.pred2[,2], xout = log10(q))$y*log10(86.4)
yield.upp.p <- approx(x = x.pred3$log.q, y = y.pred2[,3], xout = log10(q))$y*log10(86.4)

# Sediment Yield Curve with Uncertainty
kdf.mean <- approx(x = 10^(x.pred3$log.q), y = 10^(y.conf2[,1]), xout = kde$x)$y*86.4*kde$y
kdf.low <- approx(x = 10^(x.pred3$log.q), y = 10^(y.conf2[,2]), xout = kde$x)$y*86.4*kde.cb[1,]
kdf.high <- approx(x = 10^(x.pred3$log.q), y = 10^(y.conf2[,3]), xout = kde$x)$y*86.4*kde.cb[2,]
kdf.low.p <- approx(x = 10^(x.pred3$log.q), y = 10^(y.pred2[,2]), xout = kde$x)$y*86.4*kde.cb[1,]
kdf.high.p <- approx(x = 10^(x.pred3$log.q), y = 10^(y.pred2[,3]), xout = kde$x)$y*86.4*kde.cb[2,]

# replace first NA with zero
kdf.mean[1] <- 0 
kdf.low[1] <- 0 
kdf.high[1] <- 0 
kdf.low.p[1] <- 0 
kdf.high.p[1] <- 0 

# Plots

col3 <- col2rgb("seagreen")/255
plot(kde$x, kdf.mean, type = "l", ylim = range(c(kdf.low, kdf.high.p), na.rm = T), xlim = c(0,2000),
     xlab = "", ylab = "Sediment Yield Prob. Density", cex.lab = cex.lab, cex.axis = cex.axis)
polygon( c(kde$x, rev(kde$x)), c(kdf.low.p,rev(kdf.high.p)), col=rgb(t(col3),alpha = 0.5), border=F) 
col4 <- col2rgb("royalblue")/255
polygon( c(kde$x, rev(kde$x)), c(kdf.low,rev(kdf.high)), col=rgb(t(col4),alpha =0.5), border=F) 
lines(kde$x, kdf.mean)
text(x=grconvertX(.1, from='npc'), y = grconvertY(.9, from='npc'), labels = "(c)")
mtext(expression("Discharge ("*m^3/s*")"), side = 1, outer = T, line = 1, cex = 1)

segments(kde$x[which.max(kdf.low.p)], 0, kde$x[which.max(kdf.low.p)], max(kdf.low.p), lty = 2, col = "grey30")
segments(kde$x[which.max(kdf.high.p)], 0, kde$x[which.max(kdf.high.p)], max(kdf.high.p), lty = 2, col = "grey30")

## Calc Cum Sed Yield Curve ##
x.c <- rollmean(kde$x,2)
cumYield.mean <- cum.int(kde$x, kdf.mean)
cumYield.low <- cum.int(kde$x, kdf.low)
cumYield.high <- cum.int(kde$x, kdf.high)
cumYield.low.p <- cum.int(kde$x, kdf.low.p)
cumYield.high.p <- cum.int(kde$x, kdf.high.p)

# Plot ##
plot(x.c, cumYield.mean, type = "n", ylim = range(c(cumYield.low.p, cumYield.high.p), na.rm = T), xlim = c(0,2000),
     xlab = "", ylab = "Cum. Sediment Yield",cex.lab = cex.lab, cex.axis = cex.axis)
polygon( c(x.c, rev(x.c)), c(cumYield.low,rev(cumYield.high)), col=rgb(t(col4),alpha = .5), border=F) 
polygon( c(x.c, rev(x.c)), c(cumYield.low.p,rev(cumYield.high.p)), col=rgb(t(col3),alpha = .5), border=F) 
lines(x.c, cumYield.mean)
text(x=grconvertX(.1, from='npc'), y = grconvertY(.9, from='npc'), labels = "(d)")

qs50.l <- approx(cumYield.low, x.c, xout = 0.5)$y
qs50.h <- approx(cumYield.high, x.c, xout = 0.5)$y

segments(qs50.l, -1, qs50.l, 0.5, lty = 2, col = "grey30")
segments(qs50.h, -1, qs50.h, 0.5, lty = 2, col = "grey30")
segments(-500, 0.5, qs50.h, 0.5, lty = 1, col = "grey30")
