From 842feb6f70f92193ba968c5f3e261c81988dca2b Mon Sep 17 00:00:00 2001
From: athowes <adamthowes@gmail.com>
Date: Thu, 13 Jul 2023 12:47:49 +0100
Subject: [PATCH] Create PCA-AGHQ grid for Naomi using optresults #58

---
 src/naomi-simple_fit/dev/pca-scaled.R | 56 ++++++++++++++++++++++--
 src/naomi-simple_fit/functions.R      | 63 +++++++++++----------------
 src/naomi-simple_fit/script.R         | 10 +----
 3 files changed, 80 insertions(+), 49 deletions(-)

diff --git a/src/naomi-simple_fit/dev/pca-scaled.R b/src/naomi-simple_fit/dev/pca-scaled.R
index 35c2d71..9123e77 100644
--- a/src/naomi-simple_fit/dev/pca-scaled.R
+++ b/src/naomi-simple_fit/dev/pca-scaled.R
@@ -4,7 +4,57 @@
 #' 4. Run fit_aghq with s = 7
 #' 5. Check point estimates against PCA-AGHQ in naomi-simple_point-estimates
 
-levels <- c(rep(k, s), rep(1, n_hyper - s))
-prod(levels)
+tmb_input <- tmb_inputs_simple
+k <- 1
+inner_verbose <- FALSE
+progress <- NULL
+map <- NULL
+DLL <- "naomi_simple"
 
-pca_scaled_base_grid <- mvQuad::createNIGrid(dim = n_hyper, type = "GHe", level = levels)
+if (DLL == "naomi_simple") {
+  stopifnot(inherits(tmb_input, "naomi_simple_tmb_input"))
+} else {
+  stopifnot(inherits(tmb_input, "naomi_tmb_input"))
+}
+
+obj <- local_make_tmb_obj(tmb_input$data, tmb_input$par_init, calc_outputs = 0L, inner_verbose, progress, map, DLL = DLL)
+optresults <- optimize_theta(obj, startingvalue = obj$par, control = default_control_tmb())
+
+#' Can now create whatever grid I'd like using the optresults
+
+#' Create a PCA-AGHQ grid
+#'
+#' @param m Mode vector
+#' @param C Covariance matrix
+#' @param s Small grid dimension
+#' @param k Number of points per small grid dimension
+create_pca_grid <- function(m, C, s, k) {
+  d <- nrow(C)
+  stopifnot(d == length(m))
+  eigenC <- eigen(C)
+  lambda <- eigenC$values
+  E <- eigenC$vectors
+  E_s <- E[, 1:s]
+  gg_s <- mvQuad::createNIGrid(dim = s, type = "GHe", level = k)
+  nodes_out <- t(E_s %*% diag(sqrt(lambda[1:s]), ncol = s) %*% t(mvQuad::getNodes(gg_s)))
+  for(j in 1:d) nodes_out[, j] <- nodes_out[, j] + m[j] # Adaption
+  weights_out <- mvQuad::getWeights(gg_s) * as.numeric(mvQuad::getWeights(mvQuad::createNIGrid(dim = d - s, type = "GHe", level = 1)))
+  weights_out <- det(chol(C)) * weights_out # Adaption
+
+  # Putting things into a mvQuad format manually
+  gg <- mvQuad::createNIGrid(dim = d, type = "GHe", level = 1)
+  gg$level <- rep(NA, times = d)
+  gg$ndConstruction <- "PCA"
+  gg$nodes <- nodes_out
+  gg$weights <- weights_out
+  return(gg)
+}
+
+pca_grid <- create_pca_grid(m = optresults$mode, C = solve(optresults$hessian), s = s, k = k)
+
+pca_grid
+
+# quad <- local_marginal_laplace_tmb(obj, optresults = optresults, ...)
+# objout <- local_make_tmb_obj(tmb_input$data, tmb_input$par_init, calc_outputs = 1L, inner_verbose, progress, map, DLL = DLL)
+# quad$obj <- objout
+# quad
diff --git a/src/naomi-simple_fit/functions.R b/src/naomi-simple_fit/functions.R
index 1b49edd..79e2378 100644
--- a/src/naomi-simple_fit/functions.R
+++ b/src/naomi-simple_fit/functions.R
@@ -184,10 +184,18 @@ local_sample_tmb <- function(fit, M = 1000, rng_seed = NULL, random_only = TRUE,
 }
 
 #' A local version of aghq::normalize_logpost
-#' I have added an argument whereby the basegrid provided can not be adapted (if it is adapted already)
-local_normalize_logpost <- function(optresults, k, whichfirst = 1, basegrid = NULL, adapt = TRUE, ndConstruction = "product", ...) {
+#' I have added an argument whereby the adaption can be chosen not to occur
+#' This is useful if providing a basegrid which is already adapted by some other method
+local_normalize_logpost <- function(optresults, k = NULL, whichfirst = 1, basegrid = NULL, adapt = TRUE, ndConstruction = "product", ...) {
   S <- length(optresults$mode)
-  thegrid <- basegrid
+
+  if (!is.null(basegrid)) {
+    thegrid <- basegrid
+  } else {
+    if(is.null(k)) stop("You must either provide k or a basegrid")
+    thegrid <- mvQuad::createNIGrid(dim = S, type = "GHe", level = k, ndConstruction = ndConstruction, ...)
+  }
+
   idxorder <- c(whichfirst, (1:S)[-whichfirst])
 
   if(adapt) {
@@ -215,25 +223,12 @@ local_normalize_logpost <- function(optresults, k, whichfirst = 1, basegrid = NU
 
 #' A local version of aghq::aghq
 #' Compatible with local_normalize_logpost -- i.e. allowing no adaption
-local_aghq <- function(ff, k, startingvalue, transformation = default_transformation(), optresults = NULL, basegrid = NULL, adapt = TRUE, control = default_control(), ...) {
+local_aghq <- function(ff, k = NULL, startingvalue, transformation = default_transformation(), optresults = NULL, basegrid = NULL, adapt = TRUE, control = default_control(), ...) {
 
   validate_control(control)
   validate_transformation(transformation)
   transformation <- make_transformation(transformation)
 
-  # If they provided a basegrid, get the k from that. If they also provided a k, compare them and issue a warning
-  if (!is.null(basegrid)) {
-    if (missing(k)) {
-      k <- max(as.numeric(basegrid$level))
-    } else {
-      k2 <- max(as.numeric(basegrid$level))
-      if (k != k2) {
-        warning(paste0("You provided a basegrid and a specified number of quadrature points k. You do not need to specify k if you supply a basegrid. Further, they don't match: your grid has k = ",k2,", but you specified k = ",k,". Proceeding with k = ",k2,", from the supplied grid.\n"))
-        k <- k2
-      }
-    }
-  }
-
   # Optimization
   if (is.null(optresults)) utils::capture.output(optresults <- optimize_theta(ff, startingvalue, control, ...))
 
@@ -253,21 +248,20 @@ local_aghq <- function(ff, k, startingvalue, transformation = default_transforma
   class(out) <- "aghq"
 
   # Marginals
-  d <- length(startingvalue)
-  marginals <- vector(mode = "list", length = d)
-
-  if (control$method_summaries[1] == "correct") {
-    for (j in 1:d) marginals[[j]] <- aghq:::marginal_posterior.aghq(out, j, method = "correct")
-  } else {
-    for (j in 1:d) marginals[[j]] <- aghq:::marginal_posterior.aghq(out, j, method = "reuse")
-  }
-
-  out$marginals <- marginals
+  # d <- length(startingvalue)
+  # marginals <- vector(mode = "list", length = d)
+  #
+  # if (control$method_summaries[1] == "correct") {
+  #   for (j in 1:d) marginals[[j]] <- aghq:::marginal_posterior.aghq(out, j, method = "correct")
+  # } else {
+  #   for (j in 1:d) marginals[[j]] <- aghq:::marginal_posterior.aghq(out, j, method = "reuse")
+  # }
+  #
+  # out$marginals <- marginals
   out
 }
 
 local_marginal_laplace_tmb <- function(ff, k, startingvalue, transformation = default_transformation(), optresults = NULL, basegrid = NULL, adapt = TRUE, control = default_control_tmb(), ...) {
-
   validate_control(control, type = "tmb")
   validate_transformation(transformation)
   transformation <- make_transformation(transformation)
@@ -280,12 +274,11 @@ local_marginal_laplace_tmb <- function(ff, k, startingvalue, transformation = de
   if (control$numhessian) {
     ff$he <- function(theta) numDeriv::jacobian(ff$gr, theta, method = "Richardson")
   }
-  ## Do aghq ##
+
   # The aghq
   quad <- local_aghq(ff = ff, k = k, transformation = transformation, startingvalue = startingvalue, optresults = optresults, basegrid = basegrid, adapt = adapt, control = control, ...)
   if (control$onlynormconst) return(quad) # NOTE: control was passed to aghq here so quad should itself just be a number
 
-  ## Add on the info needed for marginal Laplace ##
   # Add on the quad point modes and curvatures
   distinctthetas <- quad$normalized_posterior$nodesandweights[, grep("theta", colnames(quad$normalized_posterior$nodesandweights))]
   if (!is.data.frame(distinctthetas)) distinctthetas <- data.frame(theta1 = distinctthetas)
@@ -300,13 +293,6 @@ local_marginal_laplace_tmb <- function(ff, k, startingvalue, transformation = de
   modesandhessians$mode <- vector(mode = "list", length = nrow(distinctthetas))
   modesandhessians$H <- vector(mode = "list", length = nrow(distinctthetas))
 
-  # if (is.null(thetanames)) {
-  #   thetanames <- colnames(distinctthetas)
-  # } else {
-  #   colnames(modesandhessians)[colnames(modesandhessians) == colnames(distinctthetas)] <- thetanames
-  #   colnames(quad$normalized_posterior$nodesandweights)[grep('theta',colnames(quad$normalized_posterior$nodesandweights))] <- thetanames
-  # }
-
   for (i in 1:nrow(distinctthetas)) {
     # Get the theta
     theta <- as.numeric(modesandhessians[i,thetanames])
@@ -336,7 +322,8 @@ fit_aghq <- function(tmb_input, inner_verbose = FALSE, progress = NULL, map = NU
   }
 
   obj <- local_make_tmb_obj(tmb_input$data, tmb_input$par_init, calc_outputs = 0L, inner_verbose, progress, map, DLL = DLL)
-  quad <- aghq::marginal_laplace_tmb(obj, startingvalue = obj$par, ...)
+  optresults <- optimize_theta(obj, startingvalue = obj$par, control = default_control_tmb())
+  quad <- local_marginal_laplace_tmb(obj, optresults = optresults, ...)
   objout <- local_make_tmb_obj(tmb_input$data, tmb_input$par_init, calc_outputs = 1L, inner_verbose, progress, map, DLL = DLL)
   quad$obj <- objout
   quad
diff --git a/src/naomi-simple_fit/script.R b/src/naomi-simple_fit/script.R
index 84cc1bb..619b250 100755
--- a/src/naomi-simple_fit/script.R
+++ b/src/naomi-simple_fit/script.R
@@ -144,17 +144,11 @@ if(aghq) {
   if(grid_type == "pca") {
     levels <- c(rep(k, s), rep(1, n_hyper - s))
     pca_base_grid <- mvQuad::createNIGrid(dim = n_hyper, type = "GHe", level = levels)
-
-    #' The k argument here shouldn't be doing anything: this should (in future) be fixed in aghq::aghq
-    quad <- fit_aghq(tmb_inputs_simple, k = 1, basegrid = pca_base_grid, dec.type = 1)
+    quad <- fit_aghq(tmb_inputs_simple, basegrid = pca_base_grid, dec.type = 1)
   }
 
   if(grid_type == "pca-scaled") {
-    levels <- c(rep(k, s), rep(1, n_hyper - s))
-    pca_base_grid <- mvQuad::createNIGrid(dim = n_hyper, type = "GHe", level = levels)
-
-    #' The k argument here shouldn't be doing anything: this should (in future) be fixed in aghq::aghq
-    quad <- fit_aghq(tmb_inputs_simple, k = 1, basegrid = pca_base_grid, dec.type = 1)
+    stop("Not implemented yet!")
   }
 
   if(sample && grid_type != "sparse") {