diff --git a/NAMESPACE b/NAMESPACE index b5e6eab..f3818e6 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -1,5 +1,6 @@ # Generated by roxygen2: do not edit by hand +export(multivariate_probit) export(zero_inflated_negative_binomial) export(zero_inflated_poisson) importFrom(R6,R6Class) diff --git a/R/multivariate-probit.R b/R/multivariate-probit.R new file mode 100644 index 0000000..e2b93c1 --- /dev/null +++ b/R/multivariate-probit.R @@ -0,0 +1,186 @@ +#' @name multivariate_probit +#' @title multivariate probit distribution +#' +#' @description greta probability distribution over a K-dimensional vector of +#' binary variables. +#' +#' @param mean matrix or column vector (of length K) of unconstrained parameters +#' giving the mean of the latent normal parameters +#' @param C K-dimensional correlation matrix +#' @param dim a scalar giving the number of rows in the resulting greta array +#' +#' @importFrom R6 R6Class +#' @export +multivariate_probit <- function(mean, C, dim = 1) { + distrib("multivariate_probit", mean, C, dim) +} + +# multivariate probit distribution +multivariate_probit_distribution <- R6::R6Class( + "multivariate_probit_distribution", + inherit = .internals$nodes$node_classes$distribution_node, + public = list( + initialize = function(mean, C, dim) { + + # coerce to greta arrays + mu <- as.greta_array(mean) + R <- as.greta_array(C) + + # check dimensions of mean + if (ncol(mean) != 1 | length(dim(mean)) != 2) { + stop("mean must be a 2D greta array with one column, but has dimensions ", + paste(dim(mean), collapse = " x "), + call. = FALSE + ) + } + + # check dimensions of C + if (nrow(C) != ncol(C) | length(dim(C)) != 2) { + stop("C must be a square 2D greta array, but has dimensions ", + paste(dim(C), collapse = " x "), + call. = FALSE + ) + } + + # compare possible dimensions + dim_mean <- nrow(mean) + dim_C <- nrow(C) + + if (dim_mean != dim_C) { + stop("mean and C have different dimensions, ", + dim_mean, " vs ", dim_C, + call. = FALSE + ) + } + + if (dim_mean == 1) { + stop("the multivariate probit distribution is for vectors, ", + "but the parameters were scalar", + call. = FALSE + ) + } + + # check dim is a positive scalar integer + dim_old <- dim + dim <- as.integer(dim) + if (length(dim) > 1 || dim <= 0 || !is.finite(dim)) { + stop("dim must be a scalar positive integer, but was: ", + capture.output(dput(dim_old)), + call. = FALSE + ) + } + + # coerce the parameter arguments to nodes and add as children and + # parameters + super$initialize("multivariate_probit", + dim = c(dim, length(mean)), + discrete = TRUE + ) + self$add_parameter(mean, "mean") + self$add_parameter(C, "C") + + # create latent variable and add as a parameter + u <- variable(0, 1, dim = self$dim) + self$add_parameter(u, "u") + # this needs to be overwritted on distribution assignment to get the + # correct dimensions! + }, + tf_distrib = function(parameters) { + mean <- parameters$mean + C <- parameters$C + u <- parameters$u + + # change this to use the cholesky factor if available! + L <- tf$cholesky(C) + + N <- self$dim[1] + K <- self$dim[2] + + stop("not implemented") + + # return a tf function, taking the binary vector and returning the density + # (including the log jacobian transform?) + + log_prob <- function(x) { + + # find the elements of x that are 1, and those that are 0, count them + # and split u too. Do this on the values matrix, rather than the tensor? + + # Add this as a parameter and ignore x? This will make it harder to + # switch to discrete inference though. + + N_pos + N_neg + idx_neg + idx_pos + + # list of length k, giving the index to positive elements in a vector + # representing column k + idx_pos_list + idx_neg_list + + # u_pos <- u[idx_pos] + # u_neg <- u[idx_neg] + + bound <- tf$zeros(self$dim) + lj_pos <- tf$zeros(c(N_pos, K)) + lj_neg <- tf$zeros(c(N_neg, K)) + prev <- tf$zeros(N) + + # replacement indices + dummy_idx <- .internals$utils$dummy_arrays$dummy(N, K) + + for (k in seq_len(K)) { + + # vector index for this column + k_idx <- dummy_idx[, k] + + # pos/neg index relative to this slice + idx_pos_slice <- idx_pos_list[[k]] + idx_neg_slice <- idx_neg_list[[k]] + + # pos/neg index relative to the whole matrix + idx_pos_matrix <- k_idx[idx_pos_slice] + idx_neg_matrix <- k_idx[idx_neg_slice] + + # update the bounds for this slice + bound_k <- tf_iprobit(-(mu[k_idx] + prev) / L) # make this a backsolve + + # positive and negative bounds in this slice + bound_pos_k <- bound_k[idx_pos_slice] + bound_neg_k <- bound_k[idx_neg_slice] + + # positive and negative elements of u in this slice + u_pos_k <- u[k_idx[idx_pos_list[[k]]]] + u_pos_k <- u[k_idx[idx_neg_list[[k]]]] + + z_pos_k <- tf_iprobit(bound_pos_k + (1 - bound_pos_k) * u_pos_k) + z_neg_k <- tf_iprobit(bound_neg_k * u_neg_k) + + # recombine into z + z <- tf_replace(z, z_pos_k, k_idx[idx_pos_list[[k]]], self$dim) + z <- tf_replace(z, z_neg_k, k_idx[idx_neg_list[[k]]], self$dim) + + prev <- L[k, 1:k] %*% z[1:k, ] + + # these can be in lists to combine later + lj_pos[, k] <- tf$log(1 - bound_pos[, k]) + lj_neg[, k] <- tf$log(bound_neg[, k]) + } + + # return the log density (the log jacobian of the transformation) + tf$reduce_sum(lj_pos) + tf$reduce_sum(lj_neg) + } + + list( + log_prob = log_prob, + cdf = NULL, + log_cdf = NULL + ) + }, + + # no CDF for multivariate distributions + tf_cdf_function = NULL, + tf_log_cdf_function = NULL + ) +) diff --git a/man/greta.distributions.Rd b/man/greta.distributions.Rd index 625d4ba..ef3e161 100644 --- a/man/greta.distributions.Rd +++ b/man/greta.distributions.Rd @@ -12,6 +12,7 @@ describe your package here, you can re-use the text from DESCRIPTION Useful links: \itemize{ \item \url{https://github.com/greta-dev/greta.distributions} + \item \url{https://greta-dev.github.io/greta.distributions/} \item Report bugs at \url{https://github.com/greta-dev/greta.distributions/issues} } diff --git a/man/multivariate_probit.Rd b/man/multivariate_probit.Rd new file mode 100644 index 0000000..389d0ff --- /dev/null +++ b/man/multivariate_probit.Rd @@ -0,0 +1,20 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/multivariate-probit.R +\name{multivariate_probit} +\alias{multivariate_probit} +\title{multivariate probit distribution} +\usage{ +multivariate_probit(mean, C, dim = 1) +} +\arguments{ +\item{mean}{matrix or column vector (of length K) of unconstrained parameters +giving the mean of the latent normal parameters} + +\item{C}{K-dimensional correlation matrix} + +\item{dim}{a scalar giving the number of rows in the resulting greta array} +} +\description{ +greta probability distribution over a K-dimensional vector of +binary variables. +} diff --git a/tests/spelling.R b/tests/spelling.R index 2e1ac39..d60e024 100644 --- a/tests/spelling.R +++ b/tests/spelling.R @@ -1,9 +1,7 @@ -if(requireNamespace('spelling', quietly = TRUE)) { - +if (requireNamespace("spelling", quietly = TRUE)) { spelling::spell_check_test( vignettes = TRUE, error = FALSE, skip_on_cran = TRUE ) - -} \ No newline at end of file +} diff --git a/tests/testthat/test-multivariate-probit.R b/tests/testthat/test-multivariate-probit.R new file mode 100644 index 0000000..5959394 --- /dev/null +++ b/tests/testthat/test-multivariate-probit.R @@ -0,0 +1,24 @@ +source("helpers.R") + +test_that("multivariate_probit distribution has correct density", { + skip_if_not(check_tf_version()) + + # compare_distribution( + # greta_fun = multivariate_probit, + # r_fun = extraDistr::m, + # parameters = list(lambda = 2, pi = 0.2), + # x = sample_zero_inflated_pois( + # n = 100, + # lambda = 2, + # pi = 0.2 + # ) + # ) +}) + +test_that("multivariate_probit distribution works",{ + skip_if_not(check_tf_version()) + mp <- multivariate_probit(mean = matrix(rnorm(2), nrow = 2), C = diag(2)) + expect_no_error( + calculate(mp, nsim = 1) + ) +}) \ No newline at end of file