Skip to content

Commit

Permalink
BBNB AD missing value support
Browse files Browse the repository at this point in the history
  • Loading branch information
@kenkellner
kenkellner committed Sep 10, 2024
1 parent 23dc177 commit f34b98e
Show file tree
Hide file tree
Showing 3 changed files with 254 additions and 8 deletions.
43 changes: 35 additions & 8 deletions R/dNmixtureAD.R
View file
Original file line number Diff line number Diff line change
Expand Up @@ -462,13 +462,27 @@ dNmixtureAD_BBNB_v <- nimbleFunction(
if (Nmin == -1 | Nmax == -1) {
stop("Dynamic choice of Nmin/Nmax is not supported for beta binomial N-mixtures.")
}
Nmin <- ADbreak(max( max(x), Nmin )) ## set Nmin to at least the largest x
logProb <- dNmixture_BBNB_steps(x, beta-x,lambda,theta,s,Nmin,Nmax,
dBetaBinom_v(x, Nmin, alpha, beta, len = len, log = TRUE))

max_x <- 0
any_not_na <- FALSE
for (i in 1:length(x)){
xi <- ADbreak(x[i])
if(!is.na(xi)){
if(x[i] > max_x) max_x <- x[i]
any_not_na <- TRUE
}
}
Nmin <- ADbreak(max( max_x, Nmin )) ## set Nmin to at least the largest x

logProb <- 0
if(any_not_na){
logProb <- dNmixture_BBNB_steps(x, beta-x, lambda, theta, s, Nmin, Nmax,
dBetaBinom_v(x, Nmin, alpha, beta, len, log = TRUE), usingAD=TRUE)
}
if (log) return(logProb)
else return(exp(logProb))
returnType(double())
}, buildDerivs=list(run=list())
}, buildDerivs=list(run=list(ignore=c("i", "xi")))
)

#' @rdname dNmixtureAD
Expand Down Expand Up @@ -500,13 +514,26 @@ dNmixtureAD_BBNB_s <- nimbleFunction(
}
# Clen <- 0L
# Clen <- ADbreak(len)
Nmin <- ADbreak(max( max(x), Nmin )) ## set Nmin to at least the largest x
logProb <- dNmixture_BBNB_steps(x, beta-x,lambda,theta,s,Nmin,Nmax,
dBetaBinom_s(x, Nmin, alpha, beta, len = len, log = TRUE))
max_x <- 0
any_not_na <- FALSE
for (i in 1:length(x)){
xi <- ADbreak(x[i])
if(!is.na(xi)){
if(x[i] > max_x) max_x <- x[i]
any_not_na <- TRUE
}
}
Nmin <- ADbreak(max( max_x, Nmin )) ## set Nmin to at least the largest x

logProb <- 0
if(any_not_na){
logProb <- dNmixture_BBNB_steps(x, beta-x, lambda, theta, s, Nmin, Nmax,
dBetaBinom_s(x, Nmin, alpha, beta, len, log = TRUE), usingAD=TRUE)
}
if (log) return(logProb)
else return(exp(logProb))
returnType(double())
}, buildDerivs=list(run=list())
}, buildDerivs=list(run=list(ignore=c("i", "xi")))
)

#' @rdname dNmixtureAD
Expand Down
62 changes: 62 additions & 0 deletions tests/testthat/test-AD.R
View file
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,28 @@ test_that ("dNmixture works with AD", {
v1_case1, v2_case1,
0:2, RCrelTol = c(2e-15, 1e-8, 1e-3, 1e-14))

# Missing value
x <- c(7, 7, NA, 9, 10)
Rmodel <- nimbleModel(nc, data = list(x = x),
inits = list(prob = prob,
lambda = lambda,
theta = theta, s = s),
buildDerivs=TRUE)
Rmodel$calculate()

Cmodel <- compileNimble(Rmodel)
Cmodel$calculate()

nodesList_case1 <- setup_update_and_constant_nodes_for_tests(Rmodel, c('prob', 'lambda', 'theta', 's'))
v1_case1 <- list(arg1 = c(prob, lambda, theta, s)) # taping values for prob and lambda
v2_case1 <- list(arg1 = c(prob2, lambda2, theta2, s2)) # testing values for prob and lambda

res <- model_calculate_test_case(Rmodel, Cmodel,
model_calculate_test, nodesList_case1,
v1_case1, v2_case1,
0:2, RCrelTol = c(2e-15, 1e-8, 1e-3, 1e-14))


##########################
#### dNmixture_v case ####

Expand Down Expand Up @@ -580,6 +602,26 @@ test_that ("dNmixture works with AD", {
v1_case1, v2_case1,
0:2, RCrelTol = c(2e-15, 1e-8, 1e-3, 1e-14))

x <- c(7, 7, NA, 9, 10)
Rmodel <- nimbleModel(nc, data = list(x = x),
inits = list(prob = prob,
lambda = lambda,
theta = theta, s = s),
buildDerivs=TRUE)
Rmodel$calculate()

Cmodel <- compileNimble(Rmodel)
Cmodel$calculate()

nodesList_case1 <- setup_update_and_constant_nodes_for_tests(Rmodel, c('prob', 'lambda', 'theta', 's'))
v1_case1 <- list(arg1 = c(prob, lambda, theta, s)) # taping values for prob and lambda
v2_case1 <- list(arg1 = c(prob2, lambda2, theta2, s2)) # testing values for prob and lambda

res <- model_calculate_test_case(Rmodel, Cmodel,
model_calculate_test, nodesList_case1,
v1_case1, v2_case1,
0:2, RCrelTol = c(2e-15, 1e-8, 1e-3, 1e-14))

##########################
#### dNmixture_BNB_oneObs case ####

Expand Down Expand Up @@ -692,6 +734,26 @@ test_that ("dNmixture works with AD", {
model_calculate_test, nodesList_case1,
v1_case1, v2_case1,
0:2, RCrelTol = c(2e-15, 1e-8, 1e-3, 1e-14))

x <- as.numeric(NA)
Rmodel <- nimbleModel(nc, data = list(x = x),
inits = list(prob = prob, theta = theta, s=s,
lambda = lambda),
buildDerivs=TRUE)
Rmodel$calculate()

Cmodel <- compileNimble(Rmodel)
Cmodel$calculate()

nodesList_case1 <- setup_update_and_constant_nodes_for_tests(Rmodel, c('prob', 'lambda', 'theta', 's'))
v1_case1 <- list(arg1 = c(prob, lambda, theta, s2)) # taping values for prob and lambda
v2_case1 <- list(arg1 = c(prob2, lambda2, theta2, s2)) # testing values for prob and lambda

res <- model_calculate_test_case(Rmodel, Cmodel,
model_calculate_test, nodesList_case1,
v1_case1, v2_case1,
0:2, RCrelTol = c(2e-15, 1e-8, 1e-3, 1e-14))

})


Expand Down
157 changes: 157 additions & 0 deletions tests/testthat/test-NmixtureADnoDerivs.R
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1848,6 +1848,84 @@ test_that("dNmixtureAD_BBNB_v works",
CMlProbX <- cm$getLogProb("x")
expect_equal(CMlProbX, lProbX)

# Missing values
# Uncompiled calculation
x <- c(1, 0, NA, 3, 0)
Nmin <- max(x, na.rm=TRUE)
probX <- dNmixtureAD_BBNB_v(x, lambda, theta, prob, s, Nmin, Nmax, len)

# Manually calculate the correct answer
alpha <- prob * s
beta <- s - prob * s
r <- 1 / theta
pNB <- 1 / (1 + theta * lambda)

correctProbX <- 0
for (N in Nmin:Nmax) {
correctProbX <- correctProbX + dnbinom(N, size = r, prob = pNB) *
prod(dBetaBinom_v(x, N, shape1 = alpha, shape2 = beta), na.rm=TRUE)
}

expect_equal(probX, correctProbX)

# Uncompiled log probability
lProbX <- dNmixtureAD_BBNB_v(x, lambda, theta, prob, s, Nmin, Nmax, len, log = TRUE)
lCorrectProbX <- log(correctProbX)
expect_equal(lProbX, lCorrectProbX)

# Other Nmin / Nmax
Nmin <- 3
for(Nmax in 3:6) {
dynProbX <- dNmixtureAD_BBNB_v(x, lambda, theta, prob, s, Nmin = Nmin, Nmax = Nmax, len)
dynCorrectProbX <- 0
for (N in Nmin:Nmax) {
dynCorrectProbX <- dynCorrectProbX + dnbinom(N, size = r, prob = pNB) *
prod(dBetaBinom_v(x, N, shape1 = alpha, shape2 = beta))
}
expect_equal(dynProbX, dynCorrectProbX)
}
Nmin <- 0
Nmax <- 250

CprobX <- CdNmixtureAD_BBNB_v(x, lambda, theta, prob, s, Nmin, Nmax, len)
expect_equal(CprobX, probX)

ClProbX <- CdNmixtureAD_BBNB_v(x, lambda, theta, prob, s, Nmin, Nmax, len, log = TRUE)
expect_equal(ClProbX, lProbX)

# Use in Nimble model
nc <- nimbleCode({
x[1:5] ~ dNmixtureAD_BBNB_v(lambda = lambda, prob = prob[1:5],
theta = theta, s = s,
Nmin = Nmin, Nmax = Nmax, len = len)

})

m <- nimbleModel(code = nc,
data = list(x = x),
inits = list(lambda = lambda,
prob = prob,
s = s, theta = theta),
constants = list(Nmin = Nmin, Nmax = Nmax,
len = len))
m$calculate()
MlProbX <- m$getLogProb("x")
expect_equal(MlProbX, lProbX)

# Compiled model
cm <- compileNimble(m)
cm$calculate()
CMlProbX <- cm$getLogProb("x")
expect_equal(CMlProbX, lProbX)

# All missing
x <- as.numeric(rep(NA, 5))
probX <- dNmixtureAD_BBNB_v(x, lambda, theta, prob, s, Nmin, Nmax, len)
expect_equal(probX, 1)

CprobX <- CdNmixtureAD_BBNB_v(x, lambda, theta, prob, s, Nmin, Nmax, len)
expect_equal(CprobX, probX)

# Test imputing value for all NAs
xNA <- c(NA, NA, NA, NA, NA)
mNA <- nimbleModel(nc, data = list(x = xNA),
Expand Down Expand Up @@ -2017,6 +2095,76 @@ test_that("dNmixtureAD_BBNB_s works",
CMlProbX <- cm$getLogProb("x")
expect_equal(CMlProbX, lProbX)

# Missing values
# Uncompiled calculation
x <- c(1, 0, NA, 3, 0)
Nmin <- max(x, na.rm=TRUE)

probX <- dNmixtureAD_BBNB_s(x, lambda, theta, prob, s, Nmin, Nmax, len)

# Manually calculate the correct answer
alpha <- prob * s
beta <- s - prob * s
r <- 1 / theta
pNB <- 1 / (1 + theta * lambda)

correctProbX <- 0
for (N in Nmin:Nmax) {
correctProbX <- correctProbX + dnbinom(N, size = r, prob = pNB) *
prod(dBetaBinom_s(x, N, shape1 = alpha, shape2 = beta), na.rm=TRUE)
}

expect_equal(probX, correctProbX)

# Uncompiled log probability
lProbX <- dNmixtureAD_BBNB_s(x, lambda, theta, prob, s, Nmin, Nmax, len, log = TRUE)
lCorrectProbX <- log(correctProbX)
expect_equal(lProbX, lCorrectProbX)
# Other Nmin / Nmax
Nmin <- 3
for(Nmax in 3:6) {
dynProbX <- dNmixtureAD_BBNB_s(x, lambda, theta, prob, s, Nmin = Nmin, Nmax = Nmax, len)
dynCorrectProbX <- 0
for (N in Nmin:Nmax) {
dynCorrectProbX <- dynCorrectProbX + dnbinom(N, size = r, prob = pNB) *
prod(dBetaBinom_s(x, N, shape1 = alpha, shape2 = beta))
}
expect_equal(dynProbX, dynCorrectProbX)
}
Nmin <- 0
Nmax <- 250

CprobX <- CdNmixtureAD_BBNB_s(x, lambda, theta, prob, s, Nmin, Nmax, len)
expect_equal(CprobX, probX)

ClProbX <- CdNmixtureAD_BBNB_s(x, lambda, theta, prob, s, Nmin, Nmax, len, log = TRUE)
expect_equal(ClProbX, lProbX)

m <- nimbleModel(code = nc,
data = list(x = x),
inits = list(lambda = lambda,
prob = prob,
s = s, theta = theta),
constants = list(Nmin = Nmin, Nmax = Nmax,
len = len))
m$calculate()
MlProbX <- m$getLogProb("x")
expect_equal(MlProbX, lProbX)

# Compiled model
cm <- compileNimble(m)
cm$calculate()
CMlProbX <- cm$getLogProb("x")
expect_equal(CMlProbX, lProbX)

# All missing
x <- as.numeric(rep(NA,5))
probX <- dNmixtureAD_BBNB_s(x, lambda, theta, prob, s, Nmin, Nmax, len)
expect_equal(probX, 1)

CprobX <- CdNmixtureAD_BBNB_s(x, lambda, theta, prob, s, Nmin, Nmax, len)
expect_equal(CprobX, probX)

# Test imputing value for all NAs
xNA <- c(NA, NA, NA, NA, NA)
mNA <- nimbleModel(nc, data = list(x = xNA),
Expand Down Expand Up @@ -2186,6 +2334,15 @@ test_that("dNmixtureAD_BBNB_oneObs works",
CMlProbX <- cm$getLogProb("x")
expect_equal(CMlProbX, lProbX)

# Missing value
x <- as.numeric(NA)
probX <- dNmixtureAD_BBNB_oneObs(x, lambda, theta, prob, s, Nmin, Nmax)
expect_equal(probX, 1)

CprobX <- CdNmixtureAD_BBNB_oneObs(x, lambda, theta, prob, s, Nmin, Nmax)
expect_equal(CprobX, probX)


# Test imputing value for all NAs
xNA <- NA
mNA <- nimbleModel(nc, data = list(x = xNA),
Expand Down

0 comments on commit f34b98e

Please sign in to comment.