Merge pull request #8 from atsa-es/Uonly

v0.0.7 finished adding U, A, x0, C and tinits

NEWS.md

@@ -1,8 +1,9 @@
 # marssTMB 0.0.7
 
-* updated `marxss.hpp` to be in MARSS format with X and Y as mxT and nxT
-* added Q estimation to `marxss.hpp` (note `MARSS_tmb()` will block this until this is tested)
-* added C estimation to `marxss.hpp` (note `MARSS_tmb()` will block this until this is tested)
+* updated `marxss.hpp` to be in MARSS format with X and Y as mxT and nxT.
+* added Q, C, U, x0 and A estimation to `marxss.hpp`. Minimal testing so far.
+* V0 = 0 is allowed.
+* tinitx=1 or tinitx=0 allowed.
 
 # marssTMB 0.0.6
 

R/MARSS-TMB.R

@@ -34,12 +34,13 @@ MARSS_tmb <- function(y,
                      inits = NULL,
                      miss.value = as.numeric(NA),
                      method = "TMB",
-                     form = "dfa",
+                     form = c("marxss", "dfa"),
                      fit = TRUE,
                      silent = FALSE,
                      control = list(fun.opt="nlminb"),
                      ...) {
   pkg <- "marssTMB"
+  form <- match.arg(form)
   method <- match.arg(method)
   # This section is temporary
   fun.opt <- ifelse(is.null(control$fun.opt), "nlminb", control$fun.opt )
@@ -97,46 +98,42 @@ MARSS_tmb <- function(y,
   x[["method"]] <- paste0(method, " with optimization function ", fun.opt, if(fun.opt=="optim") paste0(" with method ", optim.method))
 
   # Error check for the DFA model
-  # This section is temporary. Currently only DFA model with tinitx=1 is allowed.
-  if(x[["model"]][["tinitx"]] != 1) 
-    stop(paste0(pkg, ": tinitx must be 1. Set tinitx=1 in model list."))
   model.descrip <- MARSS:::describe.marssMODEL(x[["model"]])
-  if(form == "dfa"){
+
   is.unconstrained <- function(elem) substr(model.descrip[[elem]], 1, 5) == "uncon"
-  is.diagonal <- function(elem) substr(model.descrip[[elem]], 1, 8) == "diagonal"
+  is.diagonal <- function(elem){
+    substr(model.descrip[[elem]], 1, 8) == "diagonal" | model.descrip[[elem]] == "scalar (1 x 1)" }
   is.fixed <- function(elem) substr(model.descrip[[elem]], 1, 5) == "fixed"
   is.identity <- function(elem){
     val <- model.descrip[[elem]]
     substr(val, 1, 8) == "identity" | val == "fixed and all one (1 x 1)"
   }
   is.zero <- function(elem) substr(model.descrip[[elem]], 1, 14) == "fixed and zero"
-
-  # Check that R is allowed
-  elem <- "R"
-  ok <- is.diagonal(elem) | is.fixed(elem) | is.unconstrained(elem)
-  if(!ok){
-    stop(paste0(pkg, ": R must be diagonal, fixed or unconstrained"))
+
+  # Check that R and Q are allowed
+  for(elem in c("R", "Q")){
+    ok <- is.diagonal(elem) | is.fixed(elem) | is.unconstrained(elem) | is.identity(elem)
+    if(!ok) stop(paste0(pkg, ": ", elem, " must be diagonal, fixed or unconstrained"))
   }
 
-  # Check that Q and B are identity
-  for(elem in c("B", "Q")){
+  # Check that B is identity
+  for(elem in c("B")){
   ok <- is.identity(elem)
   if(!ok) stop(paste0(pkg, ": ", elem, " must be identity"))
   }
 
-  # Check that u, a, and C are zero
-  for(elem in c("U", "A", "C")){
-  ok <- is.zero(elem)
-  if(!ok) stop(paste0(pkg, ": ", elem, " must be zero"))
-  }
-
-  # Check that x0 and V0 are fixed
-  for(elem in c("x0", "V0")){
-    ok <- is.fixed(elem)
+  # Check that a, and C are zero
+#  for(elem in c("A", "C")){
+#  ok <- is.zero(elem)
+#  if(!ok) stop(paste0(pkg, ": ", elem, " must be zero"))
+#  }
+  
+  # Check that V0 is fixed
+  for(elem in c("V0")){
+    ok <- is.fixed(elem) | is.identity(elem)
     if(!ok) stop(paste0(pkg, ": ", elem, " must be fixed"))
   }
-  }
-
+
   if(fit) return(MARSStmb(x))
   return(x)
 }

R/MARSStmb.R

@@ -44,17 +44,22 @@ MARSStmb <- function(MLEobj) {
   # Expand out to full covariate matrix
   d_Covars <- matrix(d_Covars[,1,], nrow = nrow(d_Covars))
   c_Covars <- matrix(c_Covars[,1,], nrow = nrow(c_Covars))
-  # if user did not pass in c_Covars then c must have T columns
-  no_c_covars <- ncol(c_Covars) == 1
 
   # Set up the initial matrices
   eleminits <- list()
-  for (elem in c("Z", "D", "R", "C", "Q", "V0", "x0")) {
+  model.elem <- attr(MODELobj, "par.names")
+  for (elem in model.elem[!(model.elem %in% c("c", "d"))]) {
     eleminits[[elem]] <- coef(MLEobj, type = "matrix", what = "start")[[elem]]
   }
+  # Check that no 0s on diagonal of V0 unless V0 is all zero
+  # Note V0 is fixed by definition
+  V0_is_zero <- all(eleminits[["V0"]]==0)
+  if(!V0_is_zero && any(diag(eleminits[[elem]])==0))
+    stop(paste0(pkg, ": V0 can only have 0s on the diagonal if it is all zero"))
+
   # Set up the maps
   elemmaps <- list()
-  for (elem in c("Z", "D", "C", "x0")) {
+  for (elem in model.elem) {
     elemmaps[[elem]] <- create.elem.maps(MLEobj, elem = elem)[["map"]]
   }
   # maps for var-cov matrices have diagonal separate from off-diagonal
@@ -72,7 +77,10 @@ MARSStmb <- function(MLEobj) {
     Y = y,
     d_Covar = d_Covars,
     c_Covar = c_Covars,
-    no_c_covars = as.numeric(no_c_covars)
+    has_c_covars = as.numeric(ncol(c_Covars) != 1),
+    has_d_covars = as.numeric(ncol(d_Covars) != 1),
+    V0_is_zero = as.numeric(V0_is_zero),
+    tinitx = MODELobj[["tinitx"]]
   )
 
   # Note x0 and V0 are fixed (stochastic prior) for DFA
@@ -81,18 +89,20 @@ MARSStmb <- function(MLEobj) {
   # Creates the list of initial (start) values of parameter list
   R <- eleminits[["R"]]
   sdR <- sqrt(diag(R))
-  corrR <- diag(1 / sdR) %*% R %*% diag(1 / sdR) # correlation matrix
+  corrR <- diag(1/sdR, n) %*% R %*% diag(1/sdR, n) # correlation matrix
   Q <- eleminits[["Q"]]
   sdQ <- sqrt(diag(Q))
-  corrQ <- diag(1 / sdQ) %*% Q %*% diag(1 / sdQ) # correlation matrix
+  corrQ <- diag(1/sdQ, m) %*% Q %*% diag(1/sdQ, m) # correlation matrix
   parameters <- list(
     X = matrix(0, ncol = TT, nrow = m), # states
     x0 = eleminits[["x0"]],
     V0 = eleminits[["V0"]],
     logsdQ = log(sdQ), # log of sqrt of diagonal of Q
     cholCorrQ = chol(corrQ)[upper.tri(Q)], # off-diagonal of chol of corr Q
+    U = eleminits[["U"]],
     C = eleminits[["C"]],
     Z = eleminits[["Z"]],
+    A = eleminits[["A"]],
     D = eleminits[["D"]],
     logsdR = log(sdR), # log of sqrt of diagonal of R
     cholCorrR = chol(corrR)[upper.tri(R)] # off-diagonal of chol of corr R
@@ -106,12 +116,19 @@ MARSStmb <- function(MLEobj) {
     V0 = factor(matrix(NA, nrow = m, ncol = m)),
     logsdQ = elemmaps[["Q"]][["diag"]],
     cholCorrQ = elemmaps[["Q"]][["offdiag"]],
+    U = elemmaps[["U"]],
     C = elemmaps[["C"]],
     Z = elemmaps[["Z"]],
+    A = elemmaps[["A"]],
     D = elemmaps[["D"]],
     logsdR = elemmaps[["R"]][["diag"]],
     cholCorrR = elemmaps[["R"]][["offdiag"]]
   )
+  if(MODELobj[["tinitx"]]==1 & V0_is_zero){
+    mat <- matrix(1:(m*TT), m, TT)
+    mat[,1] <- NA
+    maplist$X <- mat |> unlist() |> as.factor()
+  }
 
   # Creates the model object and runs the optimization
   obj1 <- TMB::MakeADFun(

R/utils.R

@@ -31,9 +31,10 @@ create.varcov.maps <- function(x, elem){
 create.elem.maps <- function(x, elem="Z"){
     fixed <- x$model$fixed[[elem]]
     free <- x$model$free[[elem]]
+    par.names <- colnames(free)
     var.dim <- attr(x$model, "model.dims")[[elem]][1:2]
     par.as.list <- MARSS:::fixed.free.to.formula(fixed[, , 1, drop = FALSE], free[, , 1, drop = FALSE], var.dim)
     mat <- lapply(par.as.list, function(x){ifelse(!is.character(x), NA, x)}) |>
       matrix(var.dim)
-    return(list(map = mat |> unlist() |> as.factor(), map.matrix=mat, raw.matrix=par.as.list))
+    return(list(map = mat |> unlist() |> factor(levels=par.names), map.matrix=mat, raw.matrix=par.as.list))
 }

README.md

@@ -1,13 +1,16 @@
 # marssTMB
 
-The beginning of a companion R package to MARSS that fits MARSS models with TMB. The package is in very active development. See [changelog](news/index.html) for current status.
+## To use
+
+Fit MARSS models as usual, but use `MARSS_tmb()` instead of `MARSS()`. Not R and Q can only be diagonal, unconstrained or fixed. `equalvarcov` is not available nor are other custom Q matrices available with the EM algorithm in `MARSS()`. But TMB is very fast. 
 
 See the documentation at [marssTMB](https://atsa-es.github.io/marssTMB/)
 
 ## install
 
 ```
-remotes::install_github("atsa-es/marssTMB")
+# Install marssTMB in R:
+install.packages('marssTMB', repos = c('https://atsa-es.r-universe.dev', 'https://cloud.r-project.org'))
 ```
 
 ## Notes

src/TMB/marxss.hpp

@@ -14,15 +14,20 @@ Type marxss(objective_function<Type>* obj) {
   DATA_MATRIX(Y); /* n x T */
   DATA_MATRIX(d_Covar);
   DATA_MATRIX(c_Covar);
-  DATA_INTEGER(no_c_covars);
+  DATA_INTEGER(has_c_covars);
+  DATA_INTEGER(has_d_covars);
+  DATA_INTEGER(V0_is_zero);
+  DATA_INTEGER(tinitx);
   PARAMETER_MATRIX(X); /* State m x T */
   PARAMETER_MATRIX(x0);
   PARAMETER_MATRIX(V0); /* x[1] */
 //  PARAMETER_MATRIX(Q); /* x[t] - x[t-1] */
+  PARAMETER_MATRIX(U);
   PARAMETER_MATRIX(C);
   PARAMETER_VECTOR(logsdQ); /* log of the sqrt of diag Q*/
   PARAMETER_VECTOR(cholCorrQ);
   PARAMETER_MATRIX(Z);
+  PARAMETER_MATRIX(A);
   PARAMETER_MATRIX(D);
   PARAMETER_VECTOR(logsdR); /* log of the sqrt of diag R*/
   PARAMETER_VECTOR(cholCorrR);
@@ -43,30 +48,45 @@ Type marxss(objective_function<Type>* obj) {
 
   matrix<Type> predX(nX,1);  /* m x 1 */
 
-  MVNORM_t<Type> initialState(V0);
-  //MVNORM_t<Type> neg_log_density_process(Q);
-  /* Define likelihood */
-  Type ans=0;
+
+  Type ans=0; /* Define likelihood */
   //ans -= dnorm(vector<Type>(u.row(0)),Type(0),Type(1),1).sum();
-  ans += initialState(X.col(0)); /* tinitx=1 */
+  if(V0_is_zero){
+    if(tinitx){
+      X.col(0) = x0;
+    }else{
+      predX = x0 + U;
+      if(has_c_covars){
+        predX = predX + C * c_Covar.col(0);
+      }
+      vector<Type> differ = X.col(0)-predX;
+      ans += VECSCALE(corMatGenQ,sdQ)(differ);
+    }
+  }else{
+    MVNORM_t<Type> initialState(V0);
+    ans += initialState(X.col(0)-x0); /* tinitx=1 */
+  }
   for(int i=1;i<timeSteps;i++){ 
     //ans+= neg_log_density_process(u.row(i)-u.row(i-1)); // Process likelihood
     //vector<Type> differ = u.row(i)-u.row(i-1);
-    // predX is m x 1 so u must be transposed
     // if statement is temporary until I can figure how create a
     // a diagonal matrix with 1 on the -1 diagonal
     // diag(1:(timeSteps+1))[1:timeSteps, 2:(timeSteps+1)]
-    if(no_c_covars){
-      predX = X.col(i-1) + C * c_Covar;
-    }else{
-      predX = X.col(i-1) + C * c_Covar.col(i);
+    predX = X.col(i-1) + U;
+    if(has_c_covars){
+      predX = predX + C * c_Covar.col(i);
     }
     vector<Type> differ = X.col(i)-predX;
     ans += VECSCALE(corMatGenQ,sdQ)(differ);
   }
 
   matrix<Type> predY(nY, timeSteps);  
-  predY = Z * X + D * d_Covar;
+  matrix<Type> rowOne(1, timeSteps); 
+  rowOne.setOnes();
+  predY = Z * X + A * rowOne;
+  if(has_d_covars){
+    predY = predY + D * d_Covar;
+  }
 
   for(int i=0;i<timeSteps;i++){ //move one time step at a time
     int nonNAcount = 0; //start at zero NA values

vignettes/MARSS_TMB.Rmd

@@ -58,8 +58,8 @@ Fit with TMB.
 ```{r}
 library(marssTMB)
 m3 <- dfaTMB(dat, model=list(m=1, R='unconstrained'))
-m4 <- MARSS_tmb(dat, model=mod.list)
-m5 <- MARSS_tmb(dat, model=mod.list, control=list(fun.opt="optim"))
+m4 <- MARSS_tmb(dat, model=mod.list, form='dfa')
+m5 <- MARSS_tmb(dat, model=mod.list, control=list(fun.opt="optim"), form='dfa')
 ```
 
 ### Log likelihoods