R/PFilter-methods.R

In ONSdigital/Bayesian-demographic-accounts: Bayesian Demographic Accounts

## draw_proposal_init ---------------------------------------------------------

draw_proposal_init_noreg <- function() {
    initial_stock_fixed <- self$initial_stock_fixed
    if (!initial_stock_fixed) {
        n_particle <- self$n_particle
        cdms <- self$cdms_stock
        self$counts_stock[ , 1L] <- draw_counts_true(cdms = cdms,
                                                     i_interval = 1L,
                                                     n_particle = n_particle)
        has_been_filled <- !is.na(self$counts_stock[1L, 1L])
        if (!has_been_filled)
            stop("problem with cohort '", self$cohort, "' and sex/gender '",
                 self$sexgender, "' : no data to estimate initial stock")
    }
    invisible(self)
}

PFilterNoReg$set(which = "public",
                 name = "draw_proposal_init",
                 value = draw_proposal_init_noreg)


draw_proposal_init_withreg <- function() {
    initial_stock_fixed <- self$initial_stock_fixed
    n_particle <- self$n_particle
    if (!initial_stock_fixed) {
        cdms <- self$cdms_stock
        self$counts_stock[ , , 1L] <- draw_counts_true(cdms = cdms,
                                                       i_interval = 1L,
                                                       n_particle = n_particle)
        all_have_been_filled <- !anyNA(self$counts_stock[1L, , 1L])
        if (!all_have_been_filled)
            stop("problem with cohort '", self$cohort, "' and sex/gender '",
                 self$sexgender, "' : no data to estimate initial stock in one or more regions")
    }
    invisible(self)
}

PFilterWithReg$set(which = "public",
                   name = "draw_proposal_init",
                   value = draw_proposal_init_withreg)


## calc_loglik_init -----------------------------------------------------------

calc_loglik_init_noreg <- function() {
    initial_stock_fixed <- self$initial_stock_fixed
    if (initial_stock_fixed) {
        n_particle <- self$n_particle
        self$loglik <- rep(0, times = n_particle)
    }
    else {
        cdms <- self$cdms_stock
        stock <- self$counts_stock[ , 1L]
        self$loglik <- calc_loglik(x = cdms,
                                   counts_true = stock,
                                   i_interval = 1L)
    }
    invisible(self)
}

PFilterNoReg$set(which = "public",
                 name = "calc_loglik_init",
                 value = calc_loglik_init_noreg)


calc_loglik_init_withreg <- function() {
    initial_stock_fixed <- self$initial_stock_fixed
    if (initial_stock_fixed) {
        n_particle <- self$n_particle
        self$loglik <- rep(0, times = n_particle)
    }
    else {
        cdms <- self$cdms_stock
        stock <- self$counts_stock[ , , 1L]
        self$loglik <- calc_loglik(x = cdms,
                                   counts_true = stock,
                                   i_interval = 1L)
    }
    invisible(self)
}

PFilterWithReg$set(which = "public",
                   name = "calc_loglik_init",
                   value = calc_loglik_init_withreg)


## calc_logratio_init ----------------------------------------------------------

calc_logratio_init_noreg <- function() {
    initial_stock_fixed <- self$initial_stock_fixed
    if (initial_stock_fixed) {
        n_particle <- self$n_particle
        self$logratio <- rep(0, times = n_particle)
    }
    else {
        cdms <- self$cdms_stock
        stock <- self$counts_stock[ , 1L]
        logimp <- calc_logimp(cdms = cdms, 
                              counts_true = stock,
                              i_interval = 1L)
        self$logratio <- -1 * logimp
    }
    invisible(self)
}

PFilterNoReg$set(which = "public",
                 name = "calc_logratio_init",
                 value = calc_logratio_init_noreg)
    
calc_logratio_init_withreg <- function() {
    initial_stock_fixed <- self$initial_stock_fixed
    if (initial_stock_fixed) {
        n_particle <- self$n_particle
        self$logratio <- rep(0, times = n_particle)
    }
    else {        
        cdms <- self$cdms_stock
        stock <- self$counts_stock[ , , 1L]
        logimp <- calc_logimp(cdms = cdms,
                              counts_true = stock,
                              i_interval = 1L)
        self$logratio <- -1 * logimp
    }
    invisible(self)
}

PFilterWithReg$set(which = "public",
                   name = "calc_logratio_init",
                   value = calc_logratio_init_withreg)



## draw_proposal --------------------------------------------------------------

draw_proposal_noreg <- function(i_interval) {
    ## inputs
    n_particle <- self$n_particle
    stock_start <- self$counts_stock[ , i_interval] # vector with length 'n_particle'
    count_deaths <- self$counts_deaths[[i_interval]] # scalar
    rate_im1 <- self$rates_immigration1[[i_interval]]
    rate_im2 <- self$rates_immigration2[[i_interval]]
    rate_em1 <- self$rates_emigration1[[i_interval]]
    rate_em2 <- self$rates_emigration2[[i_interval]]
    ## calculate combined immigration parameter
    rate_cim <- rate_im1 + rate_im2
    rate_cem <- rate_em1 + rate_em2
    ## generate stock at end of interval, and net migration
    stock_end <- draw_counts_true(cdms = self$cdms_stock,
                                  i_interval = i_interval,
                                  n_particle = n_particle)
    have_stock_data <- !is.na(stock_end[[1L]])
    if (have_stock_data)
        net_mig <- stock_end - stock_start + count_deaths
    else {
        expose_tmp <- max(stock_start / 2, rate_cim / 4)
        lower_stock <- count_deaths - stock_start
        net_mig <- rskeltr(n = n_particle,
                           mu1 = rate_cim,
                           mu2 = rate_cem * expose_tmp,
                           lower = lower_stock)
        stock_end <- stock_start - count_deaths + net_mig
    }
    ## calculate exposure
    expose <- (stock_start + stock_end) / 4
    ## split net migration into combined immigration and combined emigration
    lower_gross <- abs(net_mig)
    gross_mig <- rpoistr(n = n_particle,
                         lambda = rate_cim + rate_cem * expose,
                         lower = lower_gross)
    is_diff_parity <- (gross_mig %% 2L) != (lower_gross %% 2L)
    gross_mig[is_diff_parity] <- gross_mig[is_diff_parity] + 1L
    count_cim <- (gross_mig + net_mig) %/% 2L
    count_cem <- (gross_mig - net_mig) %/% 2L
    ## split combined immigration and combined emigration
    prob_im1 <- rate_im1 / rate_cim
    prob_em1 <- rate_em1 / rate_cem
    count_im1 <- stats::rbinom(n = n_particle,
                               size = count_cim,
                               prob = prob_im1)
    count_em1 <- stats::rbinom(n = n_particle,
                               size = count_cem,
                               prob = prob_em1)
    count_im2 <- count_cim - count_im1
    count_em2 <- count_cem - count_em1
    ## record results
    self$have_stock_data <- have_stock_data
    self$counts_stock[ , i_interval + 1L] <- stock_end
    self$counts_immigration1[ , i_interval] <- count_im1
    self$counts_emigration1[ , i_interval] <- count_em1
    self$counts_immigration2[ , i_interval] <- count_im2
    self$counts_emigration2[ , i_interval] <- count_em2
    ## return updated object
    invisible(self)
}

PFilterNoReg$set(which = "public",
                 name = "draw_proposal",
                 value = draw_proposal_noreg)


draw_proposal_withreg <- function(i_interval) {
    ## inputs
    n_region <- self$n_region
    n_particle <- self$n_particle
    stock_start <- self$counts_stock[ , , i_interval]  # matrix n_particle x n_region
    count_deaths <- self$counts_deaths[ , i_interval]  # vector n_region
    rate_in <- self$rates_internal_in[ , i_interval]   # vector n_region
    rate_out <- self$rates_internal_out[ , i_interval] # vector n_region
    rate_im1 <- self$rates_immigration1[ , i_interval] # vector n_region
    rate_im2 <- self$rates_immigration2[ , i_interval] # vector n_region 
    rate_em1 <- self$rates_emigration1[ , i_interval]  # vector n_region
    rate_em2 <- self$rates_emigration2[ , i_interval]  # vector n_region
    ## expand inputs into matrices
    count_deaths <- matrix(rep(count_deaths, each = n_particle),
                           nrow = n_particle,
                           ncol = n_region)
    rate_in <- matrix(rep(rate_in, each = n_particle),
                      nrow = n_particle,
                      ncol = n_region)
    rate_out <- matrix(rep(rate_out, each = n_particle),
                       nrow = n_particle,
                       ncol = n_region)
    rate_im1 <- matrix(rep(rate_im1, each = n_particle),
                       nrow = n_particle,
                       ncol = n_region)
    rate_im2 <- matrix(rep(rate_im2, each = n_particle),
                       nrow = n_particle,
                       ncol = n_region)
    rate_em1 <- matrix(rep(rate_em1, each = n_particle),
                       nrow = n_particle,
                       ncol = n_region)
    rate_em2 <- matrix(rep(rate_em2, each = n_particle),
                       nrow = n_particle,
                       ncol = n_region)    
    ## calculate combined migration parameters
    rate_cim <- rate_im1 + rate_im2
    rate_cem <- rate_em1 + rate_em2
    rate_cin <- rate_in + rate_cim
    rate_cout <- rate_out + rate_cem
    ## Generate stock at end of interval, and net migration.
    ## Some regions may have stock data, while
    ## others do not.
    stock_end <- draw_counts_true(cdms = self$cdms_stock,  # 'stock_end' is matrix,
                                  i_interval = i_interval, # n_particle x n_region
                                  n_particle = n_particle)
    net_mig <- matrix(nrow = n_particle, ncol = n_region)
    for (i_region in seq_len(n_region)) {
        region_has_stock_data <- is.na(stock_end[1L , i_region])
        self$have_stock_data[[i_region]] <- region_has_stock_data
        if (region_has_stock_data)
            net_mig[ , i_region]  <- (stock_end[ , i_region]
                - stock_start[ , i_region]
                + count_deaths[ , i_region])
        else {
            expose_tmp <- max(stock_start[ , i_region] / 2, rate_cim[ , i_region] / 4)
            lower_stock <- count_deaths[ , i_region] - stock_start[ , i_region]
            net_mig[ , i_region] <- rskeltr(n = n_particle,
                                            mu1 = rate_cim[ , i_region],
                                            mu2 = rate_cem[i_region, ] * expose_tmp,
                                            lower = lower_stock)
        }
        stock_end[ , i_region] <- (stock_start[ , i_region]
            - count_deaths[ , i_region]
            + net_mig)
    }
    ## calculate exposure
    expose <- (stock_start + stock_end) / 4
    ## generate internal migration
    count_out <- stats::rpois(n = n_particle * n_region, lambda = rate_out * expose)
    count_out <- matrix(count_out, nrow = n_particle, ncol = n_region)
    count_out_total <- rowSums(count_out)
    count_in <- matrix(nrow = n_particle, ncol = n_region)
    for (i_particle in seq_len(n_particle))
        count_in[i_particle, ] <- stats::rmultinom(n = 1L,
                                                   size = count_out_total[[i_particle]],
                                                   prob = rate_in[i_particle, ])
    ## calculate combined net international migration
    net_mig_int <- net_mig - count_in + count_out
    ## split net international migration into combined immigration and combined emigration
    lower_gross_int <- abs(net_mig_int)
    gross_mig_int <- rpoistr(n = n_particle * n_region,
                             lambda = rate_cim + rate_cem * expose,
                             lower = lower_gross_int)
    gross_mig_int <- matrix(gross_mig_int,              
                            nrow = n_particle,    
                            ncol = n_region)      
    is_diff_parity <- (gross_mig_int %% 2L) != (lower_gross_int %% 2L)
    gross_mig_int[is_diff_parity] <- gross_mig_int[is_diff_parity] + 1L
    count_cim <- (gross_mig_int + net_mig_int) %/% 2L
    count_cem <- (gross_mig_int - net_mig_int) %/% 2L
    ## split combined immigration and combined emigration
    prob_im1 <- rate_im1 / rate_cim
    prob_em1 <- rate_em1 / rate_cem
    count_im1 <- stats::rbinom(n = n_particle * n_region,
                               size = count_cim,
                               prob = prob_im1)
    count_em1 <- stats::rbinom(n = n_particle * n_region,
                               size = count_cem,
                               prob = prob_em1)
    count_im2 <- count_cim - count_im1
    count_em2 <- count_cem - count_em1
    ## record results
    self$counts_stock[ , , i_interval + 1L] <- stock_end
    self$counts_immigration1[ , , i_interval] <- count_im1
    self$counts_emigration1[ , , i_interval] <- count_em1
    self$counts_immigration2[ , , i_interval] <- count_im2
    self$counts_emigration2[ , , i_interval] <- count_em2
    ## return updated object
    invisible(self)
}

PFilterWithReg$set(which = "public",
                   name = "draw_proposal",
                   value = draw_proposal_withreg)


## calc_loglik ----------------------------------------------------------------

calc_loglik_noreg <- function(i_interval) {
    ans_stock <- calc_loglik(self$cdms_stock,
                             counts_true = self$counts_stock[ , i_interval + 1L],
                             i_interval = i_interval)
    ans_im1 <- calc_loglik(self$cdms_immigration1,
                           counts_true = self$counts_immigration1[ , i_interval],
                           i_interval = i_interval)
    ans_em1 <- calc_loglik(self$cdms_emigration1,
                           counts_true = self$counts_emigration1[ , i_interval],
                           i_interval = i_interval)
    ans_im2 <- calc_loglik(self$cdms_immigration2,
                           counts_true = self$counts_immigration2[ , i_interval],
                           i_interval = i_interval)
    ans_em2 <- calc_loglik(self$cdms_emigration2,
                           counts_true = self$counts_emigration2[ , i_interval],
                           i_interval = i_interval)
    self$loglik <- ans_stock + ans_im1 + ans_em1 + ans_im2 + ans_em2
    invisible(self)
}

PFilterNoReg$set(which = "public",
                 name = "calc_loglik",
                 value = calc_loglik_noreg)


calc_loglik_withreg <- function(i_interval) {
    ans_stock <- calc_loglik(self$cdms_stock,
                             counts_true = self$counts_stock[ , i_interval + 1L],
                             i_interval = i_interval)
    ans_in <- calc_loglik(self$cdms_internal_in,
                          counts_true = self$counts_internal_in[ , i_interval],
                          i_interval = i_interval)
    ans_out <- calc_loglik(self$cdms_internal_out,
                           counts_true = self$counts_internal_out[ , i_interval],
                           i_interval = i_interval)
    ans_im1 <- calc_loglik(self$cdms_immigration1,
                           counts_true = self$counts_immigration1[ , i_interval],
                           i_interval = i_interval)
    ans_em1 <- calc_loglik(self$cdms_emigration1,
                           counts_true = self$counts_emigration1[ , i_interval],
                           i_interval = i_interval)
    ans_im2 <- calc_loglik(self$cdms_immigration2,
                           counts_true = self$counts_immigration2[ , i_interval],
                           i_interval = i_interval)
    ans_em2 <- calc_loglik(self$cdms_emigration2,
                           counts_true = self$counts_emigration2[ , i_interval],
                           i_interval = i_interval)
    self$loglik <- ans_stock + ans_in + ans_out + ans_im1 + ans_em1 + ans_im2 + ans_em2
    invisible(self)
}

PFilterWithReg$set(which = "public",
                   name = "calc_loglik",
                   value = calc_loglik_withreg)



## calc_logratio --------------------------------------------------------------

calc_logratio_noreg <- function(i_interval) {
    expose <- 0.25 * (self$counts_stock[ , i_interval] + self$counts_stock[ , i_interval + 1L])
    expect_im <- (self$rates_immigration1[[i_interval]]
        + self$rates_immigration2[[i_interval]])
    expect_em <- ((self$rates_emigration1[[i_interval]]
        + self$rates_emigration2[[i_interval]]) * expose)
    expect_gross_mig <- expect_im + expect_em
    combined_im <- self$immigration1[[i_interval]] + self$immigration2[[i_interval]]
    combined_em <- self$emigrration1[[i_interval]] + self$emigration2[[i_interval]]
    net_mig <- combined_im - combined_em
    gross_mig <- combined_im + combined_em
    lower_net_mig <- self$counts_death[[i_interval]] - self$stock[ , i_interval]
    lower_gross_mig <- abs(net_mig)
    if (gross_mig == lower_gross_mig)
        logratio_gross_mig <- 0
    else
        logratio_gross_mig <- log(expect_gross_mig) - log(gross_mig + expect_gross_mig)
    if (self$drew_stock)
        logimp <- calc_logimp(self$cdms_population,
                                counts_true = self$stock[ , i_interval + 1L],
                                i_interval = i_interval,
                                n_particle = self$n_particle)
    else {
        stock_start <- self$counts_stock[[i_interval]]
        expose_tmp <- max(stock_start / 2, combined_im / 4)
        expect_em_tmp <- ((self$rates_emigration1[[i_interval]]
            + self$rates_emigration2[[i_interval]]) * expose_tmp)
        logimp <- dskeltr1(x = self$net_mig,
                            mu1 = expect_im,
                            mu2 = expect_em_tmp,
                            lower = lower_net_mig,
                            log = TRUE)
    }
    ans <- (stats::dpois(x = self$counts_births_to[[i_interval]], ## x always has length 1 - use the length-1 C++ equivalent
                         lambda = self$rates_births * expose,
                         log = TRUE)
        + stats::dpois(x = self$counts_deaths[[i_interval]], ## x always has length 1 - use the length-1 C++ equivalent
                       lambda = self$rates_deaths * expose,
                       log = TRUE)
        + stats::dpois(x = self$net_mig, ## x always has length 1 - use the length-1 C++ equivalent
                       lambda = expect_gross_mig,
                       log = TRUE)
        + logratio_gross_mig + 
        + stats::ppois(q = lower_gross_mig - 1L, ## q always has length 1 - use the length-1 C++ equivalent
                       lambda = expect_gross_mig,
                       lower.tail = FALSE,
                       log = TRUE)
        + stats::dbinom(x = combined_im, ## x always has length 1 - use the length-1 C++ equivalent
                        size = gross_mig, 
                        prob = expect_im / expect_gross_mig,
                        log = TRUE)
        - pskel1(q = lower_net_mig - 1L,
                 mu1 = expect_im,
                 mu2 = expect_em,
                 lower_tail = FALSE,
                 log_p = TRUE))
    ans <- ans - logimp
    ans
}

PFilterNoReg$set(which = "public",
                 name = "calc_logratio",
                 value = calc_logratio_noreg)


calc_logratio_withreg <- function(i_interval) {
    stop("not written yet")
}

PFilterWithReg$set(which = "public",
                   name = "calc_logratio",
                   value = calc_logratio_withreg)


## calc_wt --------------------------------------------------------------------

calc_wt <- function(i_interval) {
    loglik <- self$loglik[ , i_interval + 1L]
    logratio <- self$logratio[ , i_interval + 1L]
    log_wt_unnorm <- loglik + logratio
    max <- max(log_wt_unnorm)
    wt_unnorm <- exp(log_wt_unnorm - max)
    self$wt[, i_interval + 1L] <- wt_unnorm / sum(wt_unnorm)
    invisible(self)
}

PFilter$set(which = "public",
            name = "calc_wt",
            value = calc_wt)



## resample -------------------------------------------------------------------

resample <- function(i_interval) {
    stop("not written yet")
}

PFilter$set(which = "public",
            name = "resample",
            value = resample)


## run ------------------------------------------------------------------------

run <- function(draw_init) {
    n_interval <- self$n_interval
    if (draw_init) {
        self$draw_proposal_init()
        self$calc_loglik_init()
        self$calc_logratio_init()
        self$calc_wt(0L)
        self$resample(0L)
    }
    for (i_interval in seq_len(n_interval)) {
        self$draw_proposal(i_interval)
        self$calc_loglik(i_interval)
        self$calc_logratio(i_interval)
        self$calc_wt(i_interval)
        self$resample(i_interval)
    }
    invisible(self)
}

PFilter$set(which = "public",
            name = "run",
            value = run)



## make_output -------------------------------------------------------------

make_output_noreg <- function() {
    stop("not written yet")
}


make_output_withreg <- function() {
    stop("not written yet")
}