haines2019_to_nimble/m2_nimble_gen.R at main · dapperstats/haines2019_to_nimble · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

##########################################
# Model 2: site covariates for capture   #
##########################################

# and

# not assuming a rectangular survy design but rather
#   using a vector of observations with indexing based on length vector


dM2_nb_vec <- nimbleFunction (
  run = function (x   = double(1),
                  mut = double(),
                  g0  = double(),
                  g1  = double(),
                  rt  = double(),
                  J_i = double(1),
                  R   = integer(),
                  w   = double(1),
                  log = logical(0, default = 0)) {

    mu   <- exp(mut)
    pt   <- g0 + g1 * w
    p    <- expit(pt)
    r    <- exp(rt)


  xtot       <- sum(x)
  ptot       <- double(R)
  x_row      <- integer(R)
  x_miss_row <- integer(R)

  for (i in 1:R) {

    spots_in      <- (sum(J_i[1:i]) - J_i[i] + 1):(sum(J_i[1:i]) - J_i[i] + J_i[i])
    ptot[i]       <- 1 - (1 - p[i])^J_i[i]
    x_row[i]      <- sum(x[spots_in])
    x_miss_row[i] <- x[spots_in] %*% seq(0, J_i[i] - 1)

  }


  x_sumj <- sum(x_miss_row[1:R])
  x_logfact     <- sum(lfactorial(x))

  logProb <- (-x_logfact)

  for (i in 1:R) {

    spots_in      <- (sum(J_i[1:i]) - J_i[i] + 1):(sum(J_i[1:i]) - J_i[i] + J_i[i])

    x_vec  <- seq(0, J_i[i] - 1)


    prob <- double(1)

    for (j in 1:J_i[i]) {
      prob[j] <- pow(1 - p[i], j - 1) * p[i]
    }

    term1   <- lgamma(r + x_row[i]) - lgamma(r)
    term2   <- r * log(r) + x_row[i] * log(mu)
    term3   <- sum(x[spots_in] * log(prob))
    term4   <- -(x_row[i] + r) * log(r + mu * ptot[i])
    logProb <- logProb + term1 + term2 + term3 + term4

  }
  if (log) return(logProb)
  else return(exp(logProb))
  returnType(double())
})


rM2_nb_vec <- nimbleFunction(
  run = function(n   = integer(),
                 mut = double(),
                 g0  = double(),
                 g1  = double(),
                 rt  = double(),
                 J_i = double(1),
                 R   = integer(),
                 w   = double(1)) {

    mu   <- exp(mut)
    pt   <- g0 + g1 * w
    p    <- expit(pt)
    r    <- exp(rt)

    J_tot <- sum(J_i)

    prob <- double(J_tot + 1 * R)
    retain <- logical(J_tot + 1 * R)
    ans <- integer(J_tot + 1 * R)

    for (i in 1:R) {

      for (j in 1:J_i[i]) {

        prob[sum(J_i[1:i]) - J_i[i] + i - 1 + j] <- pow(1 - p[i], j - 1) * p[i]
        retain[sum(J_i[1:i]) - J_i[i] + i - 1 + j] <- TRUE

      }
      prob[sum(J_i[1:i]) - J_i[i] + i - 1 + j + 1] <- 1 - sum(prob[(sum(J_i[1:i]) - J_i[i] + i - 1 + 1):(sum(J_i[1:i]) - J_i[i] + i - 1 + J_i[i])])
    }


    n <- integer(R)

    for (i in 1:R) {

      n[i] <- rnbinom(n = 1, size = r, mu = mu)

      if (n[i] > 0) {

        ans[(sum(J_i[1:i]) - J_i[i] + i - 1):(sum(J_i[1:i]) - J_i[i] + i - 1 + J_i[i]) + 1] <- rmulti(n = 1, size = n[i], prob = prob[(sum(J_i[1:i]) - J_i[i] + i - 1):(sum(J_i[1:i]) - J_i[i] + i - 1 + J_i[i]) + 1])

      }

    }

    return(ans[retain])
    returnType(double(1))
})

registerDistributions(list(
  dM2_nb_vec = list(
    BUGSdist = "dM2_nb_vec(mut, g0, g1, rt, J_i, R, z)",
    Rdist = "dM2_nb_vec(mut, g0, g1, rt, J_i, R, z)",
    discrete = TRUE,
    types = c('value = double(1)',
              'mut = double()',
              'g0  = double()',
              'g1  = double()',
              'rt  = double()',
              'J_i = double(1)',
              'R = integer()',
              'z = double(1)'
              ),
    mixedSizes = FALSE,
    pqAvail = FALSE
  )), verbose = FALSE
)