trustL1 update for new dMod stucture

Author: simonbeyer1

R/objClass.R

@@ -310,156 +310,124 @@ normL2 <- function(data, x, errmodel = NULL, times = NULL,
 #' and its derivatives with respect to p and sigma. Sigma parameters being
 #' passed to the function are ALWAYS assumed to be on a log scale, i.e. internally
 #' sigma parameters are converted by \code{exp()}.
-#' @examples
-#' mu <- c(A = 0, B = 0)
-#' sigma <- c(A = 0.1, B = 1)
-#' myfn <- constraintL2(mu, sigma)
-#' myfn(pars = c(A = 1, B = -1))
-#' 
-#' # Introduce sigma parameter but fix them (sigma parameters
-#' # are assumed to be passed on log scale)
-#' mu <- c(A = 0, B = 0)
-#' sigma <- paste("sigma", names(mu), sep = "_")
-#' myfn <- constraintL2(mu, sigma)
-#' pars <- c(A = .8, B = -.3, sigma_A = -1, sigma_B = 1)
-#' myfn(pars = pars[c(1, 3)], fixed = pars[c(2, 4)])
-#' 
-#' # Assume same sigma parameter for both A and B
-#' # sigma is assumed to be passed on log scale
-#' mu <- c(A = 0, B = 0)
-#' myfn <- constraintL2(mu, sigma = "sigma")
-#' pars <- c(A = .8, B = -.3, sigma = 0)
-#' myfn(pars = pars)
-#' 
 #' @export
 constraintL2 <- function(mu, sigma = 1, attr.name = "prior", condition = NULL) {
 
+  estimateSigma <- is.character(sigma)
 
-  # Aktuell zu kompliziert aufgesetzt. Man sollte immer die komplette Hessematrix/Gradient
-  # auswerten und dann die Elemente streichen, die in fixed sind!
-  
-  
-  estimateSigma <- ifelse(is.character(sigma), TRUE, FALSE)
-  if (length(sigma) > 1 & length(sigma) < length(mu))
-    stop("sigma must either have length 1 or at least length equal to length of mu.")
-  
-  ## Augment sigma if length = 1
   if (length(sigma) == 1) 
     sigma <- structure(rep(sigma, length(mu)), names = names(mu))
   if (is.null(names(sigma)))
     names(sigma) <- names(mu)
-  if (!is.null(names(sigma)) & !all(names(mu) %in% names(sigma)))
-    stop("Names of sigma and names of mu do not match.")
-  
-  ## Bring sigma in correct order (no matter if character or numeric)
   sigma <- sigma[names(mu)]
 
   controls <- list(mu = mu, sigma = sigma, attr.name = attr.name)
 
-  myfn <- function(..., fixed = NULL, deriv = TRUE, conditions = condition, env = NULL) {
+  myfn <- function(..., fixed = NULL, deriv = TRUE, deriv2 = FALSE, conditions = condition, env = NULL) {
 
-    arglist <- list(...)
-    arglist <- arglist[match.fnargs(arglist, "pars")]
-    pouter <- arglist[[1]]
-    
-    # Import from controls
+    pouter <- list(...)[[match.fnargs(list(...), "pars")]]
     mu <- controls$mu
     sigma <- controls$sigma
-    attr.name <- controls$attr.name
-    nmu <- length(mu)
 
-    # pouter can be a list (if result from a parameter transformation)
-    # In this case match with conditions and evaluate only those
-    # If there is no overlap, return NULL
-    # If pouter is not a list, evaluate the constraint function 
-    # for this pouter.
+    # Handle list input (multiple conditions)
     if (is.list(pouter) && !is.null(conditions)) {
       available <- intersect(names(pouter), conditions)
-      defined <- ifelse(is.null(condition), TRUE, condition %in% conditions)
-      
-      if (length(available) == 0 | !defined) return()
-      pouter <- pouter[intersect(available, condition)]
+      if (length(available) == 0) return()
+      pouter <- pouter[available]
     }
     if (!is.list(pouter)) pouter <- list(pouter)
 
-    
     outlist <- lapply(pouter, function(p) {
 
+      # Get deriv attributes before any manipulation
+      dP <- attr(p, "deriv", exact = TRUE)
+      dP2 <- attr(p, "deriv2", exact = TRUE)
 
-      pars <- c(p, fixed)[names(mu)]
-      p1 <- setdiff(intersect(names(mu), names(p)), names(fixed))
+      # Combine and extract available parameters
+      all_pars <- c(as.numeric(p), fixed)
+      names(all_pars) <- c(names(p), names(fixed))
+      avail <- intersect(names(mu), names(all_pars))
 
-      # if estimate sigma, produce numeric sigma vector from the parameters provided in p and fixed
-      if (estimateSigma) {
-        sigmapars <- sigma
-        sigma <- exp(c(p, fixed)[sigma])
-        names(sigma) <- names(mu)
-        Jsigma <- do.call(cbind, lapply(unique(sigmapars), function(s) {
-          (sigmapars == s)*sigma
-        }))
-        colnames(Jsigma) <- unique(sigmapars)
-        rownames(Jsigma) <- names(sigma)
-        p2 <- setdiff(intersect(unique(sigmapars), names(p)), names(fixed))
-      }
+      if (length(avail) == 0)
+        return(objlist(value = 0, 
+                       gradient = if(deriv) setNames(rep(0, length(p)), names(p)) else NULL,
+                       hessian = if(deriv) matrix(0, length(p), length(p), dimnames = list(names(p), names(p))) else NULL))
 
-      # Compute constraint value and derivatives
-      val <- sum((pars - mu)^2/sigma^2) + estimateSigma * sum(log(sigma^2))
-      val.p <- 2*(pars - mu)/sigma^2
-      val.sigma <- -2*(pars-mu)^2/sigma^3 + 2/sigma
-      val.p.p <- diag(2/sigma^2, nmu, nmu); colnames(val.p.p) <- rownames(val.p.p) <- names(mu)
-      val.p.sigma <- diag(-4*(pars-mu)/sigma^3, nmu, nmu); colnames(val.p.sigma) <- rownames(val.p.sigma) <- names(mu)
-      val.sigma.sigma <- diag(6*(pars-mu)^2/sigma^4 - 2/sigma^2, nmu, nmu); colnames(val.sigma.sigma) <- rownames(val.sigma.sigma) <- names(mu)
+      pars <- all_pars[avail]
+      mu_a <- mu[avail]
+      sig_a <- sigma[avail]
+      n_a <- length(avail)
 
-      # Multiply with Jacobian of sigma vector if estimate sigma
+      p1 <- intersect(setdiff(names(mu), names(fixed)), names(p))
+      p2 <- character(0)
+      
+      # Handle sigma estimation
       if (estimateSigma) {
-        val.sigma.sigma <- t(Jsigma) %*% val.sigma.sigma %*% Jsigma + diag((t(val.sigma) %*% Jsigma)[1,], ncol(Jsigma), ncol(Jsigma))
-        val.sigma <- (val.sigma %*% Jsigma)[1,]
-        val.p.sigma <- (val.p.sigma %*% Jsigma)
+        sig_a <- exp(all_pars[sig_a])
+        names(sig_a) <- avail
+        p2 <- intersect(setdiff(unique(sigma[avail]), names(fixed)), names(p))
       }
 
+      # Compute value
+      res <- pars - mu_a
+      val <- sum(res^2 / sig_a^2) + estimateSigma * sum(2 * log(sig_a))
 
       gr <- hs <- NULL
       if (deriv) {
-        # Produce output gradient and hessian
-        gr <- rep(0, length(p)); names(gr) <- names(p)
+        gr <- setNames(rep(0, length(p)), names(p))
         hs <- matrix(0, length(p), length(p), dimnames = list(names(p), names(p)))
 
-        # Set values in gradient and hessian
-        gr[p1] <- val.p[p1]
-        hs[p1, p1] <- val.p.p[p1, p1]
-        if (estimateSigma) {
-          gr[p2] <- val.sigma[p2]
-          hs[p1, p2] <- val.p.sigma[p1, p2]
-          hs[p2, p1] <- t(val.p.sigma)[p2, p1]
-          hs[p2, p2] <- val.sigma.sigma[p2, p2]
+        p1_a <- intersect(p1, avail)
+        if (length(p1_a) > 0) {
+          gr[p1_a] <- 2 * res[p1_a] / sig_a[p1_a]^2
+          diag(hs)[p1_a] <- 2 / sig_a[p1_a]^2
         }
 
-        # Multiply with derivatives of incoming parameter
-        dP <- attr(p, "deriv")
+        if (estimateSigma && length(p2) > 0) {
+          # Aggregate sigma derivatives by sigma parameter name
+          for (sp in p2) {
+            idx <- which(sigma[avail] == sp)
+            gr[sp] <- sum(-2 * res[idx]^2 / sig_a[idx]^2 + 2)
+            hs[sp, sp] <- sum(4 * res[idx]^2 / sig_a[idx]^2)
+          }
+          # Cross terms p1 x p2
+          for (sp in p2) {
+            idx <- names(sigma[avail])[sigma[avail] == sp]
+            common <- intersect(idx, p1_a)
+            if (length(common) > 0) {
+              hs[common, sp] <- -4 * res[common] / sig_a[common]^2
+              hs[sp, common] <- hs[common, sp]
+            }
+          }
+        }
+        
+        # Chain rule
         if (!is.null(dP)) {
-          gr <- as.vector(gr %*% dP); names(gr) <- colnames(dP)
-          hs <- t(dP) %*% hs %*% dP; colnames(hs) <- colnames(dP); rownames(hs) <- colnames(dP)
+          gr_inner <- gr
+          gr <- as.vector(gr_inner %*% dP)
+          names(gr) <- colnames(dP)
+          hs <- t(dP) %*% hs %*% dP
+          if (!is.null(dP2)) {
+            hs <- hs + colSums(gr_inner * matrix(dP2, nrow = length(gr_inner)))
+            dim(hs) <- c(ncol(dP), ncol(dP))
+          }
+          dimnames(hs) <- list(colnames(dP), colnames(dP))
         }
       }
 
       objlist(value = val, gradient = gr, hessian = hs)
-      
-      
     })
 
     out <- Reduce("+", outlist)
     attr(out, controls$attr.name) <- out$value
     attr(out, "env") <- env
-    return(out)
-    
-    
+    out
   }
+  
   class(myfn) <- c("objfn", "fn")
   attr(myfn, "conditions") <- condition
   attr(myfn, "parameters") <- names(mu)
-  return(myfn)
-  
-  
+  myfn
 }