Skip to content

Commit

Permalink
finish initial dsemRTMB
Browse files Browse the repository at this point in the history
  • Loading branch information
James-Thorson committed Jan 1, 2025
1 parent f374447 commit 9c4598d
Show file tree
Hide file tree
Showing 6 changed files with 267 additions and 78 deletions.
87 changes: 66 additions & 21 deletions R/compute_nll.R
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,24 @@ function( parlist,
"c" <- ADoverload("c")
"[<-" <- ADoverload("[<-")

# Temporary fix for solve(adsparse) returning dense-matrix
sparse_solve = function(x){
invx = solve(x)
if( RTMB:::ad_context() ){
out = sparseMatrix(
i = row(invx),
j = col(invx),
x = 1,
)
out = AD(out)
out@x = invx
#out = drop0(out) # drop0 doesn't work
return(out)
}else{
return(invx)
}
}

# Unpack parameters explicitly
beta_z = parlist$beta_z
delta0_j = parlist$delta0_j
Expand Down Expand Up @@ -46,6 +64,47 @@ function( parlist,
Gamma_kk = AD(ram$G_kk)
V_kk = t(Gamma_kk) %*% Gamma_kk

# Rescale I-Rho and Gamma if using constant marginal variance options
if( (options[2]==1) || (options[2]==2) ){
invIminusRho_kk = sparse_solve(IminusRho_kk) # solve(adsparse) returns dense-matrix
#print(class(invIminusRho_kk))

# Hadamard squared LU-decomposition
# See: https://eigen.tuxfamily.org/dox/group__QuickRefPage.html
squared_invIminusRho_kk = invIminusRho_kk
#print(class(squared_invIminusRho_kk))
#print(invIminusRho_kk)
#REPORT( invIminusRho_kk )
squared_invIminusRho_kk@x = squared_invIminusRho_kk@x^2

if( options[2] == 1 ){
# 1-matrix
ones_k1 = matrix(1, nrow=n_k, ncol=1)

# Calculate diag( t(Gamma) * Gamma )
squared_Gamma_kk = Gamma_kk
squared_Gamma_kk@x = squared_Gamma_kk@x^2
sigma2_k1 = t(squared_Gamma_kk) %*% ones_k1;

# Rowsums
margvar_k = solve(squared_invIminusRho_kk, sigma2_k1)

# Rescale IminusRho_kk and Gamma
invmargsd_kk = invsigma_kk = AD(Diagonal(n_k))
invmargsd_kk@x = 1 / sqrt(margvar_k)
invsigma_kk@x = 1 / sqrt(sigma2_k1)
IminusRho_kk = invmargsd_kk %*% IminusRho_kk;
Gamma_kk = invsigma_kk %*% Gamma_kk;
}else{
# calculate diag(Gamma)^2
targetvar_k = diag(Gamma_kk)^2

# Rescale Gamma
margvar_k = solve(squared_invIminusRho_kk, targetvar_k)
diag(Gamma_kk) = sqrt(margvar_k)
}
}

# Calculate effect of initial condition -- SPARSE version
delta_tj = matrix( 0, nrow=n_t, ncol=n_j )
if( length(delta0_j)>0 ){
Expand All @@ -62,34 +121,20 @@ function( parlist,
# Full rank GMRF
z_tj = x_tj

# Doesn't work ... mat2triplet not implemented
#V_kk = matrix(0, nrow=n_k, ncol=n_k)
#REPORT( V_kk )
#V_kk = as.matrix(V_kk)
#invV_kk = solve(V_kk)
#Qtmp_kk = invV_kk
#Qtmp_kk = AD(Qtmp_kk)
#REPORT( Qtmp_kk )
#Q_kk = Qtmp_kk

# Works
invV_kk = AD(ram$G_kk)
invV_kk@x = 1 / Gamma_kk@x^2
Q_kk = t(IminusRho_kk) %*% invV_kk %*% IminusRho_kk
# Works for diagonal
#invV_kk = AD(ram$G_kk)
#invV_kk@x = 1 / Gamma_kk@x^2
#Q_kk = t(IminusRho_kk) %*% invV_kk %*% IminusRho_kk

# Experiment
#IminusRho_dense = as.matrix( IminusRho_kk )
#V_dense = as.matrix( V_kk )
#Q_RHS = solve(V_kk, IminusRho_kk)
#Q_kk = t(IminusRho_kk) %*% Q_RHS
# Works in general
invV_kk = sparse_solve( V_kk )
Q_kk = t(IminusRho_kk) %*% invV_kk %*% IminusRho_kk

# Fine from here
jnll_gmrf = -1 * dgmrf( as.vector(z_tj), mu=as.vector(xhat_tj + delta_tj), Q=Q_kk, log=TRUE )
#jnll_gmrf = 0
REPORT( Q_kk )
}else{
# Reduced rank projection .. dgmrf is lower precision than GMRF in CPP
#jnll_gmrf = -1 * dgmrf( as.vector(x_tj), mu=rep(1,n_k), Q=Diagonal(n_k), log=TRUE )
jnll_gmrf = -1 * sum( dnorm(x_tj, mean=0, sd=1, log=TRUE) )

#
Expand Down
55 changes: 0 additions & 55 deletions R/dsemRTMB.R
Original file line number Diff line number Diff line change
Expand Up @@ -69,11 +69,6 @@ function( sem,
delta0_j = rep(0, n_j),
x_tj = ifelse( is.na(y_tj), 0, y_tj )
)
#if( control$gmrf_parameterization=="separable" ){
# Params$x_tj = ifelse( is.na(tsdata), 0, tsdata )
#}else{
# Params$eps_tj = ifelse( is.na(tsdata), 0, tsdata )
#}

# Turn off initial conditions
if( estimate_delta0==FALSE ){
Expand Down Expand Up @@ -113,8 +108,6 @@ function( sem,
}

# Build object
#source( file.path(R'(C:\Users\James.Thorson\Desktop\Git\dsem\R)', "make_matrices.R") )
#source( file.path(R'(C:\Users\James.Thorson\Desktop\Git\dsem\R)', "get_jnll.R") )
cmb <- function(f, ...) function(p) f(p, ...) ## Helper to make closure
#f(parlist, model, tsdata, family)
obj = RTMB::MakeADFun(
Expand All @@ -130,36 +123,6 @@ function( sem,
profile = control$profile,
silent = TRUE )

#
if( FALSE ){
repRTMB = obj$report()
range(Rep$Gamma_kk - repRTMB$Gamma_kk)
range(Rep$Rho_kk - repRTMB$Rho_kk)
range(Rep$Q_kk - repRTMB$Q_kk)
range(Rep$mu_tj - repRTMB$mu_tj)
range(Rep$delta_tj - repRTMB$delta_tj)
range(Rep$xhat_tj - repRTMB$xhat_tj)
range(Rep$jnll_gmrf - repRTMB$jnll_gmrf)

#
dgmrf( repRTMB$z_tj, mu=repRTMB$xhat_tj + repRTMB$delta_tj, Q=repRTMB$Q_kk, log=TRUE )
dgmrf( Rep$z_tj, mu=Rep$xhat_tj + Rep$delta_tj, Q=Rep$Q_kk, log=TRUE )
dgmrf( as.vector(Rep$z_tj), mu=as.vector(Rep$xhat_tj + Rep$delta_tj), Q=Rep$Q_kk, log=TRUE )
mvtnorm::dmvnorm( as.vector(Rep$z_tj - (Rep$xhat_tj + Rep$delta_tj)), sigma=as.matrix(solve(Rep$Q_kk)), log=TRUE )
dGMRF( as.vector(Rep$z_tj), mu=as.vector(Rep$xhat_tj + Rep$delta_tj), Q=Rep$Q_kk )

#
Rep = fit$obj$report( fit$obj$env$last.par )

#
example = list(
y = repRTMB$z_tj,
mu = repRTMB$xhat_tj + repRTMB$delta_tj,
Q = Rep$Q_kk
)
saveRDS( example, file=file.path(R'(C:\Users\James.Thorson\Desktop\Git\dsem\scratch)','example.RDS') )
}

if(control$quiet==FALSE) list_parameters(obj)
# bundle
internal = list(
Expand All @@ -171,24 +134,6 @@ function( sem,
covs = covs
)

# Parse priors
# if( !is.null(prior_negloglike) ){
# # prior_negloglike = \(obj) -dnorm(obj$par[1],0,1,log=TRUE)
# prior_value = tryCatch( expr = prior_negloglike(obj) )
# if( is.na(prior_value) ) stop("Check `prior_negloglike(obj$par)`")
# obj$fn_orig = obj$fn
# obj$gr_orig = obj$gr
#
# # BUild prior evaluator
# requireNamespace(RTMB)
# priors_obj = RTMB::MakeADFun( func = prior_negloglike,
# parameters = list(par=obj$par),
# silent = TRUE )
# obj$fn = \(pars) obj$fn_orig(pars) + priors_obj$fn(pars)
# obj$gr = \(pars) obj$gr_orig(pars) + priors_obj$gr(pars)
# internal$priors_obj = priors_obj
# }

# Further bundle
out = list( "obj" = obj,
"ram" = ram,
Expand Down
Binary file modified scratch/example.RDS
Binary file not shown.
81 changes: 81 additions & 0 deletions scratch/example2.R
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
library(Matrix)
library(RTMB)

# Load minimal example
example = readRDS( file=file.path(R'(C:\Users\James.Thorson\Desktop\Git\dsem\scratch)',"example.RDS"))
attach(example)
set.seed(101)
z = rnorm(nrow(Rho_kk))

# Define function with both versions
f = function(p){
#
sparse_solve = function(x){
invx = solve(x)
if( RTMB:::ad_context() ){
out = sparseMatrix(
i = row(invx),
j = col(invx),
x = 1,
)
out = AD(out)
out@x = invx
#out = drop0(out) # drop0 doesn't work
return(out)
}else{
return(invx)
}
}

Gamma_kk = AD( p[1] * Gamma_kk )
Rho_kk = AD( p[2] * Rho_kk )
V_kk = t(Gamma_kk) %*% Gamma_kk
IminusRho_kk = Diagonal(nrow(Rho_kk)) - Rho_kk

# Option-1 ... works
invV_kk = Gamma_kk
invV_kk@x = 1 / Gamma_kk@x^2
Q1_kk = t(IminusRho_kk) %*% invV_kk %*% IminusRho_kk
REPORT( Q1_kk )

# Option-2
invV_kk = solve(V_kk)
Q2_kk = t(IminusRho_kk) %*% invV_kk %*% IminusRho_kk
REPORT( Q2_kk )

# Option-3
invV_kk = solve(V_kk)
#if( RTMB:::ad_context() ){
# temp_kk = sparseMatrix(
# i = row(invV_kk),
# j = col(invV_kk),
# x = 1,
# )
# temp_kk = AD(temp_kk)
# temp_kk@x = invV_kk
# invV_kk = temp_kk
#}
invV_kk = sparse_solve(V_kk)
Q3_kk = t(IminusRho_kk) %*% invV_kk %*% IminusRho_kk
REPORT( Q3_kk )

# Option-1 WORKS for diagonal matrix only
#jnll_gmrf = -1 * dgmrf( z, mu=rep(0,nrow(Rho_kk)), Q=Q1_kk, log=TRUE )

# Option-2 DOES NOT WORK
#jnll_gmrf = -1 * dgmrf( z, mu=rep(0,nrow(Rho_kk)), Q=Q2_kk, log=TRUE )

# Option-3 WORKS generally, but requires a manual coersion from dense to sparse matrix
jnll_gmrf = -1 * dgmrf( z, mu=rep(0,nrow(Rho_kk)), Q=Q3_kk, log=TRUE )
return(jnll_gmrf)
}

# Works for Option-1
f(c(1,1))
obj = MakeADFun( f, c(1,1) )
obj$fn(obj$par)

# Show they're identical if compiled using Option-1
Rep = obj$report()
identical( Rep$Q1_kk, Rep$Q2_kk )
identical( Rep$Q1_kk, Rep$Q3_kk )
102 changes: 102 additions & 0 deletions scratch/example2_tmp.R
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
library(Matrix)
library(RTMB)

# Load minimal example
example = readRDS( file=file.path(R'(C:\Users\James.Thorson\Desktop\Git\dsem\scratch)',"example.RDS"))
attach(example)
set.seed(101)
z = rnorm(nrow(Rho_kk))

# Define function with both versions
f = function(p){
Gamma_kk = AD( p[1] * Gamma_kk )
Rho_kk = AD( p[2] * Rho_kk )
V_kk = t(Gamma_kk) %*% Gamma_kk
IminusRho_kk = Diagonal(nrow(Rho_kk)) - Rho_kk

# Option-1 ... works
invV_kk = Gamma_kk
invV_kk@x = 1 / Gamma_kk@x^2
Q1_kk = t(IminusRho_kk) %*% invV_kk %*% IminusRho_kk
REPORT( Q1_kk )

# Option-2
invV_kk = solve(V_kk)
#trip = mat2triplet(invV_kk)
#print(trip)
#print(invV_kk@i)
#print(invV_kk@p)
#print(invV_kk@x)
#p = invV_kk@p
#n.p = length(p)
#dp <- p[-1L] - p[-n.p]
#n.dp = length(dp)
#j = rep.int(seq.int(from = 0L, length.out = n.dp), dp) + 1
#print(j)
#invV_kk = sparseMatrix(
# i = invV_kk@i + 1,
# p = invV_kk@p,
# #j = j,
# x = invV_kk@x
# )
if( RTMB:::ad_context() ){
#print(invV_kk[1:10,1:10])
#class(invV_kk)
#print(row(invV_kk))

#vec = as.matrix(invV_kk)[1,1]
#print(vec)

#Dim = rep(nrow(V_kk),2)
#invV2_kk = drop0(sparseMatrix( i=1, j=1, x=0, dims=Dim )) # Make with a zero
#for(i in seq_len(nrow(invV_kk)) ){
#for(j in seq_len(ncol(invV_kk)) ){
# Ind_kk = sparseMatrix(
# i = i,
# j = j,
# x = 1,
# dims = Dim
# )
# invV2_kk = invV2_kk + invV_kk[i,j] * Ind_kk
#}}
invV2_kk = sparseMatrix(
i = row(invV_kk),
j = col(invV_kk),
x = 1,
)
invV2_kk = AD(invV2_kk)
invV2_kk@x = invV_kk
REPORT(invV2_kk)
}else{
invV2_kk = invV_kk
}
#invV2_kk@x = invV_kk@x
#invV_kk = AD(invV_kk)
#print(length(invV2_kk@x))
#print(length(invV_kk@x))
#print(invV2_kk)
Q2_kk = t(IminusRho_kk) %*% invV2_kk %*% IminusRho_kk
REPORT( Q2_kk )
#print(Q2_kk)

# Option-3 ... try rebuilding
#Q3_kk = AD(sparseMatrix(i=Q2_kk@i, p=Q2_kk@p, x=Q2_kk@x))
#REPORT( Q3_kk )

# Option-1 WORKS
#jnll_gmrf = -1 * dgmrf( z, mu=rep(0,nrow(Rho_kk)), Q=Q1_kk, log=TRUE )

# Option-2 DOES NOT WORK
jnll_gmrf = -1 * dgmrf( z, mu=rep(0,nrow(Rho_kk)), Q=Q2_kk, log=TRUE )
#return(jnll_gmrf)
}

# Works for Option-1
f(c(1,1))
obj = MakeADFun( f, c(1,1) )
# obj$fn(obj$par)

# Show they're identical if compiled using Option-1
# Rep = obj$report()
# identical( Rep$Q1_kk, Rep$Q2_kk )
# attr(Rep$Q1_kk,"Dim")
Loading

0 comments on commit 9c4598d

Please sign in to comment.