## ----setup, include=FALSE-----------------------------------------------------
library(INLA)
set.seed(123)
if (file.exists("myinit.R")) source("myinit.R")
inla.setOption(num.threads="1:1")
inla.setOption(smtp="taucs")
library(knitr)
library(rmarkdown)
knitr::opts_chunk$set(echo=TRUE, cache=FALSE, message=FALSE, warning=FALSE)
knitr::opts_chunk$set(fig.path="figures/rgeneric/")

## ----eval=FALSE---------------------------------------------------------------
#  model = inla.rgeneric.define(rmodel, ...)

## ----eval=FALSE---------------------------------------------------------------
#  y ~ ... + f(idx, model=model, ...)

## ----eval=FALSE---------------------------------------------------------------
#  inla.rgeneric.ar1.model = function(
#          cmd = c("graph", "Q", "mu", "initial", "log.norm.const",
#                  "log.prior", "quit"),
#          theta = NULL)
#  {
#      # for reference and potential storage for objects to
#      # cache, this is the environment of this function
#      # which holds arguments passed as `...` in
#      # `inla.rgeneric.define()`.
#  	envir = parent.env(environment())
#  
#      graph = function(){ <to be completed> }
#      Q = function() { <to be completed> }
#      mu = function() { <to be completed> }
#      log.norm.const = function() { <to be completed> }
#      log.prior = function() { <to be completed> }
#      initial = function() { <to be completed> }
#      quit = function() { <to be completed> }
#  
#      # sometimes this is useful, as argument 'graph' and 'quit'
#      # will pass theta=numeric(0) (or NULL in R-3.6...) as
#  	# the values of theta are NOT
#  	# required for defining the graph. however, this statement
#      # will ensure that theta is always defined.
#      if (!length(theta)) theta = initial()
#  
#      val = do.call(match.arg(cmd), args = list())
#      return (val)
#  }

## ----eval=FALSE---------------------------------------------------------------
#  model = inla.rgeneric.define(inla.rgeneric.ar1.model, n = 100)

## ----eval=FALSE---------------------------------------------------------------
#  interpret.theta = function() {
#      return(list(prec = exp(theta[1L]),
#                  rho = 2 * exp(theta[2L])/(1 + exp(theta[2L])) - 1))
#  }

## ----eval=FALSE---------------------------------------------------------------
#  graph = function() {
#      return (Q())
#  }

## ----eval=FALSE---------------------------------------------------------------
#  Q = function() {
#      p = interpret.theta()
#      Q = p$prec/(1 - p$rho^2) *
#  	    toeplitz(c(1 + p$rho^2, -p$rho, rep(0, n - 2L)))
#      Q[1, 1] = Q[n, n] = p$prec/(1 - p$rho^2)
#      return (inla.as.sparse(Q))
#  }

## ----eval=FALSE---------------------------------------------------------------
#  Q = function() {
#      p = interpret.theta()
#      i = c(1L, n, 2L:(n - 1L), 1L:(n - 1L))
#      j = c(1L, n, 2L:(n - 1L), 2L:n)
#      x = p$prec/(1 - p$rho^2) *
#  	    c(1L, 1L, rep(1 + p$rho^2, n - 2L),
#  	      rep(-p$rho, n - 1L))
#       return (sparseMatrix(i = i, j = j, x = x, giveCsparse = FALSE))
#  }

## ----eval=FALSE---------------------------------------------------------------
#  mu = function() {
#      return(numeric(0))
#  }

## ----eval=FALSE---------------------------------------------------------------
#  log.norm.const = function() {
#     p = interpret.theta()
#     prec.i  = p$prec / (1.0 - p$rho^2)
#     val = n * (- 0.5 * log(2*pi) + 0.5 * log(prec.i)) +
#           0.5 * log(1.0 - p$rho^2)
#     return (val)
#  }

## ----eval=FALSE---------------------------------------------------------------
#  log.norm.const = function() {
#     return (numeric(0))
#  }

## ----eval=FALSE---------------------------------------------------------------
#  log.prior = function() {
#      p = interpret.theta()
#      val = dgamma(p$prec, shape = 1, rate = 1, log=TRUE) + theta[1L] +
#            dnorm(theta[2L], mean = 0, sd = 1, log=TRUE)
#      return (val)
#  }

## ----eval=FALSE---------------------------------------------------------------
#  initial = function() {
#      return (rep(1, 2))
#  }

## ----eval=FALSE---------------------------------------------------------------
#  quit = function() {
#      return (invisible())
#  }

## -----------------------------------------------------------------------------
n = 100
rho=0.9
x = arima.sim(n, model = list(ar = rho)) * sqrt(1-rho^2)
y = x + rnorm(n, sd = 0.1)
model = inla.rgeneric.define(inla.rgeneric.ar1.model, n=n)
formula = y ~ -1 + f(idx, model=model)
r = inla(formula, data = data.frame(y, idx = 1:n))

## -----------------------------------------------------------------------------
fformula = y ~ -1 +
	f(idx, model = "ar1",
	  hyper = list(prec = list(prior = "loggamma", param = c(1,1)),
      rho = list(prior = "normal", param = c(0,1))))
rr = inla(fformula, data = data.frame(y, idx = 1:n))

## -----------------------------------------------------------------------------
plot(inla.smarginal(rr$marginals.hyperpar[[2]]), 
	 type="l", lwd=5, col="red", xlab="stdev", ylab="density")
lines(inla.tmarginal(exp, r$internal.marginals.hyperpar[[2]]), 
     col="yellow")

## -----------------------------------------------------------------------------
plot(inla.smarginal(rr$marginals.hyperpar[[3]]), 
	 type="l", lwd=5, col="red", xlab="rho", ylab="density")
lines(inla.tmarginal(function(x) 2*exp(x)/(1+exp(x))-1,
                     r$internal.marginals.hyperpar[[3]]), 
     col="yellow")

## -----------------------------------------------------------------------------
round(rbind(native = rr$cpu.used,  
            rgeneric = r$cpu.used), digits = 3)

## -----------------------------------------------------------------------------
inla.rgeneric.iid.model

## ---- eval=FALSE--------------------------------------------------------------
#  ## In this example we do linear regression using 'rgeneric'.
#  ## The regression model is y = a + b*x + noise,  and we
#  ## define 'a + b*x + tiny.noise' as a latent model.
#  ## The dimension is length(x) and number of hyperparameters
#  ## is 2 ('a' and 'b').

## -----------------------------------------------------------------------------
rgeneric.linear.regression =
    function(cmd = c("graph", "Q", "mu", "initial", "log.norm.const", 
                     "log.prior", "quit"),
             theta = NULL)
{
	envir = parent.env(environment())

    ## artifical high precision to be added to the mean-model
    prec.high = exp(15)
    
    interpret.theta = function() {
        return(list(a = theta[1L], b = theta[2L]))
    }
    
    graph = function() {
        G = Diagonal(n = length(x), x=1)
        return(G)
    } 
    
    Q = function() {
        Q = prec.high * graph()
        return(Q)
    }
    
    mu = function() {
        par = interpret.theta()
        return(par$a + par$b * x)
    }

    log.norm.const = function() {
        return(numeric(0))
    }

    log.prior = function() {
        par = interpret.theta()
        val = (dnorm(par$a, mean=0, sd=1, log=TRUE) +
               dnorm(par$b, mean=0, sd=1, log=TRUE))
        return(val)
    }

    initial = function() {
        return(rep(0, 2))
    }
   
    quit = function() {
        return(invisible())
    }

    val = do.call(match.arg(cmd), args = list())
    return(val)
}

## -----------------------------------------------------------------------------
a = 1
b = 2
n = 50
x = rnorm(n)
eta = a + b*x
s = 0.25
y = eta + rnorm(n, sd=s)

rgen = inla.rgeneric.define(model = rgeneric.linear.regression, x=x)
r = inla(y ~ -1 + f(idx, model=rgen),
         data = data.frame(y, idx = 1:n))
rr = inla(y ~ 1 + x,
          data = data.frame(y, x),
          control.fixed = list(prec.intercept = 1, prec = 1))

## -----------------------------------------------------------------------------
plot(inla.smarginal(r$marginals.hyperpar[['Theta1 for idx']]),
     type="l", lwd=5, col="red", xlab="Intercept", ylab="density")
lines(inla.smarginal(rr$marginals.fixed$'(Intercept)'), col="yellow")

## -----------------------------------------------------------------------------
plot(inla.smarginal(r$marginals.hyperpar[['Theta2 for idx']]),
     type="l", lwd=5, col="red", xlab="Slope", ylab="density")
lines(inla.smarginal(rr$marginals.fixed$x), col="yellow")

## -----------------------------------------------------------------------------
r = inla.hyperpar(r)

## -----------------------------------------------------------------------------
plot(inla.smarginal(r$marginals.hyperpar[['Theta1 for idx']]),
     type="l", lwd=5, col="red", xlab="Intercept", ylab="density")
lines(inla.smarginal(rr$marginals.fixed$'(Intercept)'), col="yellow")

## -----------------------------------------------------------------------------
plot(inla.smarginal(r$marginals.hyperpar[['Theta2 for idx']]),
     type="l", lwd=5, col="red", xlab="Slope", ylab="density")
lines(inla.smarginal(rr$marginals.fixed$x), col="yellow")

## -----------------------------------------------------------------------------
rgeneric.test = function(
    cmd = c("graph", "Q", "mu", "initial", "log.norm.const", "log.prior", "quit"),
    theta = NULL)
{
    envir = parent.env(environment())

    graph = function() {
        return(matrix(1, n, n))
    }

    Q = function() {
        R <- matrix(sin(1:n^2), n, n)
        R <- R %*% t(R)
        diag(R) <- diag(R)+1
        Q <- exp(theta[1]) * R
        return(Q)
    }

    mu = function() return (numeric(0))

    log.norm.const = function() {
        return (numeric(0))
    }

    log.prior = function() {
        return (dgamma(exp(theta[1]), shape = 1, rate = 1, log=TRUE) + theta[1])
    }

    initial = function() {
        return(4)
    }

    if (!length(theta)) theta = initial()
    val = do.call(match.arg(cmd), args = list())

    return (val)
}

## -----------------------------------------------------------------------------
n = 200
s = .1
Q <- rgeneric.test("Q", theta = 0)
library(mvtnorm)
S <- solve(as.matrix(Q))
S <- (S + t(S))/2
x <- drop(rmvnorm(1, sigma = S))
y <- x + rnorm(n, sd = s)
cont.family = list(hyper = list(prec = list(initial=log(1/s^2), fixed=TRUE)))

r1 = inla(y ~ -1 + f(idx, model="generic", Cmatrix = Q,
                     hyper = list(prec = list(prior = "loggamma", param = c(1, 1)))), 
          data = data.frame(y = y, idx = 1:n), control.family = cont.family)
ld <- 0.5 * log(det(as.matrix(Q)))
r1$mlik <- r1$mlik + ld ## see the documentation for why

model2 = inla.rgeneric.define(rgeneric.test, n=n, optimize = FALSE)
r2 = inla(y ~ -1 + f(idx, model=model2), 
          data = data.frame(y = y, idx = 1:n), control.family = cont.family)

## -----------------------------------------------------------------------------
r2$mlik - r1$mlik

## -----------------------------------------------------------------------------
rgeneric.test.opt.1 = function(
    cmd = c("graph", "Q", "mu", "initial", "log.norm.const", "log.prior", "quit"),
    theta = NULL)
{
    envir = parent.env(environment())

    if (!exists("cache.done", envir = envir)) {
        R <- matrix(sin(1:n^2), n, n)
        R <- R %*% t(R)
        diag(R) <- diag(R)+1
        R.logdet <- log(det(R))
        R <- inla.as.sparse(R)
        idx <- which(R@i <= R@j)
        R@i <- R@i[idx]
        R@j <- R@j[idx]
        R@x <- R@x[idx]
        assign("R", R, envir = envir)
        norm.const <- -n/2 * log(2*pi) + 0.5 * R.logdet
        assign("norm.const", norm.const, envir = envir)
        assign("cache.done", TRUE, envir = envir)
    }

    graph = function() {
        return (R)
    }

    Q = function() {
        return(exp(theta[1]) * R)
    }

    mu = function() return (numeric(0))

    log.norm.const = function() {
        return (norm.const + n/2 * theta[1])
    }

    log.prior = function() {
        return (dgamma(exp(theta[1]), shape = 1, rate = 1, log=TRUE) + theta[1])
    }

    initial = function() {
        return(4)
    }

    if (!length(theta)) theta = initial()
    val = do.call(match.arg(cmd), args = list())

    return (val)
}

## -----------------------------------------------------------------------------
A=matrix(1:9,3,3)
A
inla.as.sparse(A)@x

## -----------------------------------------------------------------------------
rgeneric.test.opt.2 = function(
    cmd = c("graph", "Q", "mu", "initial", "log.norm.const", "log.prior", "quit"),
    theta = NULL)
{
    envir = parent.env(environment())

    if (!exists("cache.done", envir = envir)) {
        R <- matrix(sin(1:n^2), n, n)
        R <- R %*% t(R)
        diag(R) <- diag(R)+1
        R.logdet <- log(det(R))
        R <- inla.as.sparse(R)
        idx <- which(R@i <= R@j)
        R@i <- R@i[idx]
        R@j <- R@j[idx]
        R@x <- R@x[idx]
        assign("R", R, envir = envir)
        norm.const <- -n/2 * log(2*pi) + 0.5 * R.logdet
        assign("norm.const", norm.const, envir = envir)
        assign("cache.done", TRUE, envir = envir)
    }

    graph = function() {
        return (R)
    }

    Q = function() {
        ## since R was created with 'inla.sparse.matrix' above, the indices are sorted in a
        ## spesific order. This ordering is REQUIRED for R@x to be interpreted correctly.
        return(exp(theta[1]) * R@x)
    }

    mu = function() return (numeric(0))

    log.norm.const = function() {
        return (norm.const + n/2 * theta[1])
    }

    log.prior = function() {
        return (dgamma(exp(theta[1]), shape = 1, rate = 1, log=TRUE) + theta[1])
    }

    initial = function() {
        return(4)
    }

    if (!length(theta)) theta = initial()
    val = do.call(match.arg(cmd), args = list())

    return (val)
}

## -----------------------------------------------------------------------------
model3 = inla.rgeneric.define(rgeneric.test.opt.1, n=n, optimize = FALSE)
r3 = inla(y ~ -1 + f(idx, model=model3), 
          data = data.frame(y = y, idx = 1:n), control.family = cont.family)

model4 = inla.rgeneric.define(rgeneric.test.opt.2, n=n, optimize = TRUE)
r4 = inla(y ~ -1 + f(idx, model=model4), 
          data = data.frame(y = y, idx = 1:n), control.family = cont.family)

## -----------------------------------------------------------------------------
print(r2$mlik - r1$mlik)
print(r3$mlik - r1$mlik)
print(r4$mlik - r1$mlik)

print(rbind(native = r1$cpu[2], 
            rgeneric.plain = r2$cpu[2], 
            rgeneric.cache = r3$cpu[2], 
            rgeneric.optimze = r4$cpu[2]))

