From d114b770ddb2f31bee46f08838db99dd9ac9b632 Mon Sep 17 00:00:00 2001
From: Nicholas Tierney <nicholas.tierney@gmail.com>
Date: Tue, 21 Jun 2022 15:29:25 +0800
Subject: [PATCH 1/7] minor changes from rebase

---
 tests/spelling.R | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

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
+}

From a90080bfe40af12b1748431065280e22cbacb32e Mon Sep 17 00:00:00 2001
From: Nicholas Tierney <nicholas.tierney@gmail.com>
Date: Tue, 21 Jun 2022 15:31:10 +0800
Subject: [PATCH 2/7] document conditional-bernoulli

---
 NAMESPACE                    |   2 +
 R/conditional-bernoulli.R    | 142 +++++++++++++++++++++++++++++++++++
 man/conditional_bernoulli.Rd |  53 +++++++++++++
 3 files changed, 197 insertions(+)
 create mode 100644 R/conditional-bernoulli.R
 create mode 100644 man/conditional_bernoulli.Rd

diff --git a/NAMESPACE b/NAMESPACE
index b5e6eab..563c9ed 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -3,6 +3,8 @@
 export(zero_inflated_negative_binomial)
 export(zero_inflated_poisson)
 importFrom(R6,R6Class)
+export(conditional_bernoulli)
+import(R6)
 importFrom(greta,.internals)
 importFrom(tensorflow,tf)
 importFrom(utils,globalVariables)
diff --git a/R/conditional-bernoulli.R b/R/conditional-bernoulli.R
new file mode 100644
index 0000000..e64fb9f
--- /dev/null
+++ b/R/conditional-bernoulli.R
@@ -0,0 +1,142 @@
+#' @name conditional_bernoulli
+#' @title conditional bernoulli distribution
+#'
+#' @description greta probability distribution over a K-dimensional vector of
+#'   binary variables, arising from independent Bernoulli draws each conditioned
+#'   on a single draw from another Bernoulli draw.
+#'
+#' @details A compound distribution, where elements of the bernoulli vector
+#'   variable _y_ can only be 1 if a scalar latent bernoulli variable
+#'   _z_ takes value 1, i.e.:
+#'
+#'   \deqn{ y_i ~ bernoulli(z * p_i)}
+#'   \deqn{z ~ bernoulli(psi)}
+#'   where
+#'   \deqn{p_i = p(y_i = 1 | z = 1)}
+#'   \deqn{psi = p(z = 1)}
+#'
+#'   _p_ and _psi_ are distinguishable provided there are multiple
+#'   trials in each observation of _y_. The density of this compound
+#'   distribution can be calculated directly, explicitly integrating over the
+#'   latent variable _z_, as:
+#'
+#'   \deqn{psi * prod((p ^ y) * (1 - p) ^ (1 - y)) + max(y) * (1 - psi)}
+#'
+#'   This formulation underpins the ecological imperfect-detection model of
+#'   MacKenzie et al. where _y_ and _p_ are vectors indicating whether
+#'   a species was detected at each visit, and the probability of detection
+#'   (which may vary between visits), and _z_ and _psi_ are scalars
+#'   indicating whether the species was present (assumed to be the same at all
+#'   visits) and the probability of being present.
+#'
+#' @references MacKenzie, D. I., Nichols, J. D., Lachman, G. B., Droege, S.,
+#'   Andrew Royle, J., & Langtimm, C. A. (2002). Estimating site occupancy rates
+#'   when detection probabilities are less than one. _Ecology_, 83(8),
+#'   2248-2255.
+#'
+#' @param p matrix (of dimension `dim` x K) of (conditional) probabilities
+#'   of success
+#' @param psi scalar or column vector (of length `dim`) of probabilities
+#'   for the latent bernoulli variable
+#' @param dim a scalar giving the number of rows in the resulting greta array
+#'
+#' @export
+conditional_bernoulli <- function(p, psi, dim = 1) {
+  distrib("conditional_bernoulli", p, psi, dim)
+}
+
+# multivariate probit distribution
+conditional_bernoulli_distribution <- R6::R6Class(
+  classname = "conditional_bernoulli_distribution",
+  inherit = distribution_node,
+  public = list(
+    initialize = function(p, psi, dim) {
+
+      # coerce to greta arrays
+      p <- as.greta_array(p)
+      psi <- as.greta_array(psi)
+
+      # check dimensions of p
+      if (ncol(p) < 2 | length(dim(p)) != 2) {
+        stop("p must be a 2D greta array with at least two columns, ",
+          "but has dimensions ",
+          paste(dim(p), collapse = " x "),
+          call. = FALSE
+        )
+      }
+
+      # check dimensions of psi
+      if (ncol(psi) != 1 | length(dim(psi)) != 2) {
+        stop("psi must be a 2D greta array with one column, ",
+          "but has dimensions ",
+          paste(dim(psi), collapse = " x "),
+          call. = FALSE
+        )
+      }
+
+      # compare possible dimensions
+      dim_p <- nrow(p)
+      dim_psi <- nrow(psi)
+
+      if (dim_p != dim_psi) {
+        stop("p and psi have different dimensions, ",
+          dim_p, " vs ", dim_psi,
+          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("conditional_bernoulli",
+        dim = c(dim, ncol(p)),
+        discrete = TRUE
+      )
+      self$add_parameter(p, "p")
+      self$add_parameter(psi, "psi")
+    },
+    tf_distrib = function(parameters) {
+      p <- parameters$p
+      psi <- parameters$psi
+
+      # return a tf function, taking the binary vector and returning the density
+
+      log_prob <- function(x) {
+
+        # for each row, were all elements 0?
+        none <- tf_as_float(tf_rowsums(x, 1L) == 0)
+
+        one <- fl(1)
+
+        # log conditional probability
+        # cp <- (p ^ x) * (one - p) ^ (one - x)
+        # log_cp <- log(cp)
+        log_cp <- x * log(p) + (one - x) * log(one - p)
+
+        # log probability
+        # log(rowProds(cp) * psi + nd * (1 - psi))
+        prob <- exp(tf_rowsums(log_cp, 1L) + log(psi)) + none * (one - psi)
+        log(prob)
+      }
+
+      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/conditional_bernoulli.Rd b/man/conditional_bernoulli.Rd
new file mode 100644
index 0000000..091dcc1
--- /dev/null
+++ b/man/conditional_bernoulli.Rd
@@ -0,0 +1,53 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/conditional-bernoulli.R
+\name{conditional_bernoulli}
+\alias{conditional_bernoulli}
+\title{conditional bernoulli distribution}
+\usage{
+conditional_bernoulli(p, psi, dim = 1)
+}
+\arguments{
+\item{p}{matrix (of dimension \code{dim} x K) of (conditional) probabilities
+of success}
+
+\item{psi}{scalar or column vector (of length \code{dim}) of probabilities
+for the latent bernoulli variable}
+
+\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, arising from independent Bernoulli draws each conditioned
+on a single draw from another Bernoulli draw.
+}
+\details{
+A compound distribution, where elements of the bernoulli vector
+variable \emph{y} can only be 1 if a scalar latent bernoulli variable
+\emph{z} takes value 1, i.e.:
+
+\deqn{ y_i ~ bernoulli(z * p_i)}
+\deqn{z ~ bernoulli(psi)}
+where
+\deqn{p_i = p(y_i = 1 | z = 1)}
+\deqn{psi = p(z = 1)}
+
+\emph{p} and \emph{psi} are distinguishable provided there are multiple
+trials in each observation of \emph{y}. The density of this compound
+distribution can be calculated directly, explicitly integrating over the
+latent variable \emph{z}, as:
+
+\deqn{psi * prod((p ^ y) * (1 - p) ^ (1 - y)) + max(y) * (1 - psi)}
+
+This formulation underpins the ecological imperfect-detection model of
+MacKenzie et al. where \emph{y} and \emph{p} are vectors indicating whether
+a species was detected at each visit, and the probability of detection
+(which may vary between visits), and \emph{z} and \emph{psi} are scalars
+indicating whether the species was present (assumed to be the same at all
+visits) and the probability of being present.
+}
+\references{
+MacKenzie, D. I., Nichols, J. D., Lachman, G. B., Droege, S.,
+Andrew Royle, J., & Langtimm, C. A. (2002). Estimating site occupancy rates
+when detection probabilities are less than one. \emph{Ecology}, 83(8),
+2248-2255.
+}

From 67f0efde9da8bdcdbed624b5e21768311f249288 Mon Sep 17 00:00:00 2001
From: njtierney <nicholas.tierney@gmail.com>
Date: Fri, 10 Jun 2022 18:14:22 +0800
Subject: [PATCH 3/7] import R6 officially, apply style formatting

---
 NAMESPACE                  | 3 +--
 R/internals.R              | 1 -
 man/greta.distributions.Rd | 1 +
 3 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/NAMESPACE b/NAMESPACE
index 563c9ed..60b7099 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -1,10 +1,9 @@
 # Generated by roxygen2: do not edit by hand
 
+export(conditional_bernoulli)
 export(zero_inflated_negative_binomial)
 export(zero_inflated_poisson)
 importFrom(R6,R6Class)
-export(conditional_bernoulli)
-import(R6)
 importFrom(greta,.internals)
 importFrom(tensorflow,tf)
 importFrom(utils,globalVariables)
diff --git a/R/internals.R b/R/internals.R
index eb7b30c..3c38463 100644
--- a/R/internals.R
+++ b/R/internals.R
@@ -1,5 +1,4 @@
 # need some internal greta functions accessible
-
 check_dims <- .internals$checks$check_dims
 check_in_family <- .internals$checks$check_in_family
 check_positive <- .internals$checks$check_positive
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}
 }
 

From 55708be1b2a4c478a960bb239d84d26ee3cf25be Mon Sep 17 00:00:00 2001
From: Nicholas Tierney <nicholas.tierney@gmail.com>
Date: Tue, 21 Jun 2022 16:41:23 +0800
Subject: [PATCH 4/7] first pass at creating some basic checks for
 conditional_bernoulli

---
 NEWS.md                                     |  1 +
 R/check_valid_probability.R                 | 23 ++++++
 R/conditional-bernoulli.R                   | 60 +++++++++++++---
 tests/testthat/test-conditional_bernoulli.R | 80 +++++++++++++++++++++
 4 files changed, 153 insertions(+), 11 deletions(-)
 create mode 100644 R/check_valid_probability.R
 create mode 100644 tests/testthat/test-conditional_bernoulli.R

diff --git a/NEWS.md b/NEWS.md
index 18794df..b1489d2 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,3 +1,4 @@
 # greta.distributions 0.0.0.9000
 
 * Added a `NEWS.md` file to track changes to the package.
+* Added `conditional_bernoulli` distribution (#5)
\ No newline at end of file
diff --git a/R/check_valid_probability.R b/R/check_valid_probability.R
new file mode 100644
index 0000000..0e9fa72
--- /dev/null
+++ b/R/check_valid_probability.R
@@ -0,0 +1,23 @@
+check_valid_probability <- function(x, var_name = "x") {
+  if (any(x < 0) | any(x > 1)) {
+    
+    first_line <- glue::glue(
+      "{.var [var_name]} must be a valid probability - between 0 and 1",
+      .open = "[",
+      .close = "]"
+    )
+    second_line <- glue::glue(
+      "We see {.var [var_name]} = {x}",
+      .open = "[",
+      .close = "]"
+    )
+    msg <- cli::format_error(
+      c(
+        first_line,
+        second_line
+      )
+    )
+    stop(msg,
+         call. = FALSE)
+  }
+}
\ No newline at end of file
diff --git a/R/conditional-bernoulli.R b/R/conditional-bernoulli.R
index e64fb9f..3560f7b 100644
--- a/R/conditional-bernoulli.R
+++ b/R/conditional-bernoulli.R
@@ -41,6 +41,10 @@
 #' @param dim a scalar giving the number of rows in the resulting greta array
 #'
 #' @export
+#' @examples 
+#' \dontrun{
+#' 
+#' }
 conditional_bernoulli <- function(p, psi, dim = 1) {
   distrib("conditional_bernoulli", p, psi, dim)
 }
@@ -51,25 +55,39 @@ conditional_bernoulli_distribution <- R6::R6Class(
   inherit = distribution_node,
   public = list(
     initialize = function(p, psi, dim) {
+      
+      # check that p and psi  are between 0 and
+      p_val <- p
+      psi_val <- psi
 
       # coerce to greta arrays
       p <- as.greta_array(p)
       psi <- as.greta_array(psi)
-
+      
       # check dimensions of p
       if (ncol(p) < 2 | length(dim(p)) != 2) {
-        stop("p must be a 2D greta array with at least two columns, ",
-          "but has dimensions ",
-          paste(dim(p), collapse = " x "),
+        msg <- cli::format_error(
+          c(
+            "{.var p} must be a 2D array with at least two columns",
+            "but {.var p} has dimensions {paste(dim(p), collapse = 'x')}"
+          )
+        ) 
+        stop(
+          msg,
           call. = FALSE
         )
       }
 
       # check dimensions of psi
       if (ncol(psi) != 1 | length(dim(psi)) != 2) {
-        stop("psi must be a 2D greta array with one column, ",
-          "but has dimensions ",
-          paste(dim(psi), collapse = " x "),
+        msg <- cli::format_error(
+          c(
+            "{.var psi} must be a 2D greta array with one column",
+            "but {.var psi} has dimensions {paste(dim(psi), collapse = 'x')}"
+          )
+        )
+        stop(
+          msg,
           call. = FALSE
         )
       }
@@ -79,8 +97,18 @@ conditional_bernoulli_distribution <- R6::R6Class(
       dim_psi <- nrow(psi)
 
       if (dim_p != dim_psi) {
-        stop("p and psi have different dimensions, ",
-          dim_p, " vs ", dim_psi,
+        msg <- cli::format_error(
+          c(
+            "{.var p} and {.var psi} must have the same number of rows",
+            "But we see {.var p} and {.var psi} have:",
+            "{.var p}: {dim_p} {?row/rows}",
+            "{.var psi}: {dim_psi} {?row/rows}",
+            "Perhaps you need to coerce {.var p} or {.var psi} to an \\
+            appropriate matrix?"
+          )
+        )
+        stop(
+          msg,
           call. = FALSE
         )
       }
@@ -89,11 +117,21 @@ conditional_bernoulli_distribution <- R6::R6Class(
       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)),
+        msg <- cli::format_error(
+          c(
+            "{.var dim} must be a scalar positive integer, but was:",
+            "{capture.output(dput(dim_old))}"
+          )
+        )
+        stop(
+          msg,
           call. = FALSE
         )
       }
+      
+      # check p and psi are between 0 and 1
+      check_valid_probability(p_val, var_name = "p")
+      check_valid_probability(psi_val, var_name = "psi")
 
       # coerce the parameter arguments to nodes and add as children and
       # parameters
diff --git a/tests/testthat/test-conditional_bernoulli.R b/tests/testthat/test-conditional_bernoulli.R
new file mode 100644
index 0000000..2b57d50
--- /dev/null
+++ b/tests/testthat/test-conditional_bernoulli.R
@@ -0,0 +1,80 @@
+# test_that("conditional_bernoulli distribution has correct density", {
+#   skip_if_not(check_tf_version())
+#   
+#   
+#   compare_distribution(greta.distributions::conditional_bernoulli,
+#                        stats::dnorm,
+#                        parameters = list(mean = -2, sd = 3),
+#                        x = rnorm(100, -2, 3)
+#   )
+# })
+
+test_that(
+  "conditional_bernoulli fails when given the wrong argument dimensions", {
+    # check p is a 2d array with at least 2 columns
+    expect_snapshot_error(
+      conditional_bernoulli(p = 1, psi = 1)
+    )
+    # check p and psi have the same number of rows
+    expect_snapshot_error(
+      conditional_bernoulli(
+        p = matrix(c(1,2), ncol = 2), 
+        psi = c(1,2)
+        )
+    )
+    expect_snapshot_error(
+      conditional_bernoulli(
+        p = matrix(c(0.1,0.9), ncol = 2), 
+        psi = matrix(c(1,2), nrow = 2)
+        )
+    )
+    expect_snapshot_error(
+      conditional_bernoulli(
+        p = matrix(c(0.1,0.9), ncol = 2), 
+        psi = c(0.9,0.1)
+        )
+    )
+  
+})
+#' 
+#' #' @param p matrix (of dimension `dim` x K) of (conditional) probabilities
+#' #'   of success
+#' #' @param psi scalar or column vector (of length `dim`) of probabilities
+#' #'   for the latent bernoulli variable
+#' 
+# check dimensions of psi
+# if (ncol(psi) != 1 | length(dim(psi)) != 2) {
+#   msg <- cli::format_error(
+#     c(
+#       "{.var psi} must be a 2D greta array with one column",
+#       "but {.var psi} has dimensions {paste(dim(psi), collapse = 'x')}"
+#     )
+#   )
+
+
+# compare possible dimensions
+# dim_p <- nrow(p)
+# dim_psi <- nrow(psi)
+# 
+# if (dim_p != dim_psi) {
+#   msg <- cli::format_error(
+#     c(
+#       "{.var p} and {.var psi} must have the same number of rows",
+#       "But we see {.var p} and {.var psi} have:",
+#       "{.var p}: {dim_p} {?row/rows}",
+#       "{.var psi}: {dim_psi} {?row/rows}",
+#       "Perhaps you need to coerce {.var p} or {.var psi} to an \\
+#             appropriate matrix?"
+#     )
+
+# # check dim is a positive scalar integer
+# dim_old <- dim
+# dim <- as.integer(dim)
+# if (length(dim) > 1 || dim <= 0 || !is.finite(dim)) {
+#   msg <- cli::format_error(
+#     c(
+#       "{.var dim} must be a scalar positive integer, but was:",
+#       "{capture.output(dput(dim_old))}"
+#     )
+#   )
+# 

From 213bca21ee987d8e88e90b280f3ac9340b0d7827 Mon Sep 17 00:00:00 2001
From: Nicholas Tierney <nicholas.tierney@gmail.com>
Date: Tue, 21 Jun 2022 16:51:27 +0800
Subject: [PATCH 5/7] add snapshot tests, and also a failing test to calculate
 values from the distribution

---
 R/conditional-bernoulli.R                     | 11 ++++--
 .../testthat/_snaps/conditional_bernoulli.md  | 39 +++++++++++++++++++
 tests/testthat/test-calculate-distribution.R  | 11 ++++++
 tests/testthat/test-conditional_bernoulli.R   | 14 +++++++
 4 files changed, 72 insertions(+), 3 deletions(-)
 create mode 100644 tests/testthat/_snaps/conditional_bernoulli.md
 create mode 100644 tests/testthat/test-calculate-distribution.R

diff --git a/R/conditional-bernoulli.R b/R/conditional-bernoulli.R
index 3560f7b..4eeef67 100644
--- a/R/conditional-bernoulli.R
+++ b/R/conditional-bernoulli.R
@@ -43,7 +43,12 @@
 #' @export
 #' @examples 
 #' \dontrun{
-#' 
+#' cb <- conditional_bernoulli(
+#'   p = matrix(c(0.1,0.9), ncol = 2), 
+#'   psi = c(0.9),
+#'   dim = 1
+#'   )
+#' cb
 #' }
 conditional_bernoulli <- function(p, psi, dim = 1) {
   distrib("conditional_bernoulli", p, psi, dim)
@@ -82,7 +87,7 @@ conditional_bernoulli_distribution <- R6::R6Class(
       if (ncol(psi) != 1 | length(dim(psi)) != 2) {
         msg <- cli::format_error(
           c(
-            "{.var psi} must be a 2D greta array with one column",
+            "{.var psi} must be a 2D array with one column",
             "but {.var psi} has dimensions {paste(dim(psi), collapse = 'x')}"
           )
         )
@@ -120,7 +125,7 @@ conditional_bernoulli_distribution <- R6::R6Class(
         msg <- cli::format_error(
           c(
             "{.var dim} must be a scalar positive integer, but was:",
-            "{capture.output(dput(dim_old))}"
+            "{dim_old}"
           )
         )
         stop(
diff --git a/tests/testthat/_snaps/conditional_bernoulli.md b/tests/testthat/_snaps/conditional_bernoulli.md
new file mode 100644
index 0000000..773138e
--- /dev/null
+++ b/tests/testthat/_snaps/conditional_bernoulli.md
@@ -0,0 +1,39 @@
+# conditional_bernoulli fails when given the wrong argument dimensions
+
+    `p` must be a 2D array with at least two columns
+    but `p` has dimensions 1x1
+
+---
+
+    `p` and `psi` must have the same number of rows
+    But we see `p` and `psi` have:
+    `p`: 1 row
+    `psi`: 2 rows
+    Perhaps you need to coerce `p` or `psi` to an appropriate matrix?
+
+---
+
+    `p` and `psi` must have the same number of rows
+    But we see `p` and `psi` have:
+    `p`: 1 row
+    `psi`: 2 rows
+    Perhaps you need to coerce `p` or `psi` to an appropriate matrix?
+
+---
+
+    `p` and `psi` must have the same number of rows
+    But we see `p` and `psi` have:
+    `p`: 1 row
+    `psi`: 2 rows
+    Perhaps you need to coerce `p` or `psi` to an appropriate matrix?
+
+---
+
+    `dim` must be a scalar positive integer, but was:
+    0.1
+
+---
+
+    `dim` must be a scalar positive integer, but was:
+    0.1
+
diff --git a/tests/testthat/test-calculate-distribution.R b/tests/testthat/test-calculate-distribution.R
new file mode 100644
index 0000000..f56dca2
--- /dev/null
+++ b/tests/testthat/test-calculate-distribution.R
@@ -0,0 +1,11 @@
+test_that("calculate works with conditional_bernouilli", {
+  cb <- conditional_bernoulli(
+    p = matrix(c(0.1,0.9), ncol = 2),
+    psi = c(0.9),
+    dim = 1
+  )
+  expect_success(
+    calculate(cb, nsim = 10)
+  )
+})
+
diff --git a/tests/testthat/test-conditional_bernoulli.R b/tests/testthat/test-conditional_bernoulli.R
index 2b57d50..349440b 100644
--- a/tests/testthat/test-conditional_bernoulli.R
+++ b/tests/testthat/test-conditional_bernoulli.R
@@ -34,6 +34,20 @@ test_that(
         psi = c(0.9,0.1)
         )
     )
+    expect_snapshot_error(
+      conditional_bernoulli(
+        p = matrix(c(0.1,0.9), ncol = 2), 
+        psi = c(0.9),
+        dim = 0.1
+        )
+    )
+    expect_snapshot_error(
+      conditional_bernoulli(
+        p = matrix(c(0.1,0.9), ncol = 2), 
+        psi = c(0.9),
+        dim = 0.1
+        )
+    )
   
 })
 #' 

From 1c501ca6158862eccfd8f1cabe476aaf43f88a81 Mon Sep 17 00:00:00 2001
From: njtierney <nicholas.tierney@gmail.com>
Date: Thu, 4 Aug 2022 14:40:16 +0800
Subject: [PATCH 6/7] adding 2d array checker

---
 R/conditional-bernoulli.R |  1 +
 R/utils.R                 | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+)
 create mode 100644 R/utils.R

diff --git a/R/conditional-bernoulli.R b/R/conditional-bernoulli.R
index 4eeef67..e845e27 100644
--- a/R/conditional-bernoulli.R
+++ b/R/conditional-bernoulli.R
@@ -70,6 +70,7 @@ conditional_bernoulli_distribution <- R6::R6Class(
       psi <- as.greta_array(psi)
       
       # check dimensions of p
+      check_if_2d_array(p)
       if (ncol(p) < 2 | length(dim(p)) != 2) {
         msg <- cli::format_error(
           c(
diff --git a/R/utils.R b/R/utils.R
new file mode 100644
index 0000000..cdf6a51
--- /dev/null
+++ b/R/utils.R
@@ -0,0 +1,32 @@
+
+n_dim <- function(x) length(dim(x))
+
+check_if_2d_array <- function(x){
+  if (n_dim(x) != 2) {
+    msg <- cli::format_error(
+      c(
+        "{.var x} must be a 2D array",
+        "but {.var x} has dimensions {paste(dim(x), collapse = 'x')}"
+      )
+    ) 
+    stop(
+      msg,
+      call. = FALSE
+    )
+  }
+}
+
+check_n_col_gte <- function(x, n_col){
+  if (ncol(x) < n_col) {
+    msg <- cli::format_error(
+      c(
+        "{.var x} must have at least two columns",
+        "but {.var x} has {ncol(x)} columns"
+      )
+    ) 
+    stop(
+      msg,
+      call. = FALSE
+    )
+  }
+}

From f7f622794319e62780f72d5188f26bd751cc3693 Mon Sep 17 00:00:00 2001
From: njtierney <nicholas.tierney@gmail.com>
Date: Wed, 11 Dec 2024 13:10:08 +1100
Subject: [PATCH 7/7] remove inline checking code and move into utils

---
 R/check_valid_probability.R  |  23 -------
 R/conditional-bernoulli.R    |  66 ++------------------
 R/utils.R                    | 117 +++++++++++++++++++++++++++++++----
 man/conditional_bernoulli.Rd |  10 +++
 4 files changed, 119 insertions(+), 97 deletions(-)
 delete mode 100644 R/check_valid_probability.R

diff --git a/R/check_valid_probability.R b/R/check_valid_probability.R
deleted file mode 100644
index 0e9fa72..0000000
--- a/R/check_valid_probability.R
+++ /dev/null
@@ -1,23 +0,0 @@
-check_valid_probability <- function(x, var_name = "x") {
-  if (any(x < 0) | any(x > 1)) {
-    
-    first_line <- glue::glue(
-      "{.var [var_name]} must be a valid probability - between 0 and 1",
-      .open = "[",
-      .close = "]"
-    )
-    second_line <- glue::glue(
-      "We see {.var [var_name]} = {x}",
-      .open = "[",
-      .close = "]"
-    )
-    msg <- cli::format_error(
-      c(
-        first_line,
-        second_line
-      )
-    )
-    stop(msg,
-         call. = FALSE)
-  }
-}
\ No newline at end of file
diff --git a/R/conditional-bernoulli.R b/R/conditional-bernoulli.R
index e845e27..8d9d0f9 100644
--- a/R/conditional-bernoulli.R
+++ b/R/conditional-bernoulli.R
@@ -71,69 +71,12 @@ conditional_bernoulli_distribution <- R6::R6Class(
       
       # check dimensions of p
       check_if_2d_array(p)
-      if (ncol(p) < 2 | length(dim(p)) != 2) {
-        msg <- cli::format_error(
-          c(
-            "{.var p} must be a 2D array with at least two columns",
-            "but {.var p} has dimensions {paste(dim(p), collapse = 'x')}"
-          )
-        ) 
-        stop(
-          msg,
-          call. = FALSE
-        )
-      }
-
-      # check dimensions of psi
-      if (ncol(psi) != 1 | length(dim(psi)) != 2) {
-        msg <- cli::format_error(
-          c(
-            "{.var psi} must be a 2D array with one column",
-            "but {.var psi} has dimensions {paste(dim(psi), collapse = 'x')}"
-          )
-        )
-        stop(
-          msg,
-          call. = FALSE
-        )
-      }
-
+      check_if_2d_gte_two_col(p)
+      check_if_2d_one_col(psi)
       # compare possible dimensions
-      dim_p <- nrow(p)
-      dim_psi <- nrow(psi)
+      check_params_same_rows(p, psi)
 
-      if (dim_p != dim_psi) {
-        msg <- cli::format_error(
-          c(
-            "{.var p} and {.var psi} must have the same number of rows",
-            "But we see {.var p} and {.var psi} have:",
-            "{.var p}: {dim_p} {?row/rows}",
-            "{.var psi}: {dim_psi} {?row/rows}",
-            "Perhaps you need to coerce {.var p} or {.var psi} to an \\
-            appropriate matrix?"
-          )
-        )
-        stop(
-          msg,
-          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)) {
-        msg <- cli::format_error(
-          c(
-            "{.var dim} must be a scalar positive integer, but was:",
-            "{dim_old}"
-          )
-        )
-        stop(
-          msg,
-          call. = FALSE
-        )
-      }
+      check_dim_positive_scalar_int(dim)
       
       # check p and psi are between 0 and 1
       check_valid_probability(p_val, var_name = "p")
@@ -174,6 +117,7 @@ conditional_bernoulli_distribution <- R6::R6Class(
 
       list(
         log_prob = log_prob,
+        sample = NULL,
         cdf = NULL,
         log_cdf = NULL
       )
diff --git a/R/utils.R b/R/utils.R
index cdf6a51..ef86dc8 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -1,32 +1,123 @@
 
 n_dim <- function(x) length(dim(x))
 
-check_if_2d_array <- function(x){
+check_if_2d_array <- function(x,
+                              call = rlang::caller_env()){
   if (n_dim(x) != 2) {
-    msg <- cli::format_error(
+    cli::cli_abort(
       c(
         "{.var x} must be a 2D array",
         "but {.var x} has dimensions {paste(dim(x), collapse = 'x')}"
-      )
+      ),
+      call = call
     ) 
-    stop(
-      msg,
-      call. = FALSE
-    )
   }
 }
 
-check_n_col_gte <- function(x, n_col){
+check_n_col_gte <- function(x,
+                            n_col,
+                            call = rlang::caller_env()){
   if (ncol(x) < n_col) {
-    msg <- cli::format_error(
+    cli::cli_abort(
       c(
         "{.var x} must have at least two columns",
         "but {.var x} has {ncol(x)} columns"
-      )
+      ),
+      call = call
+    ) 
+  }
+}
+
+check_valid_probability <- function(x, 
+                                    var_name = "x",
+                                    call = rlang::caller_env()) {
+  if (any(x < 0) | any(x > 1)) {
+    
+    first_line <- glue::glue(
+      "{.var [var_name]} must be a valid probability - between 0 and 1",
+      .open = "[",
+      .close = "]"
+    )
+    second_line <- glue::glue(
+      "We see {.var [var_name]} = {x}",
+      .open = "[",
+      .close = "]"
+    )
+    cli::cli_abort(
+      c(
+        first_line,
+        second_line
+      ),
+      call = call
+    )
+  }
+}
+
+check_if_2d_gte_two_col <- function(p,
+                                 call = rlang::caller_env()){
+  does_not_have_at_least_two_cols <- ncol(p) < 2 | length(dim(p)) != 2
+  if (does_not_have_at_least_two_cols) {
+    cli::cli_abort(
+      c(
+        "{.var p} must be a 2D array with at least two columns",
+        "but {.var p} has dimensions {paste(dim(p), collapse = 'x')}"
+      ),
+      call = call
     ) 
-    stop(
-      msg,
-      call. = FALSE
+  }
+}
+
+check_if_2d_one_col <- function(psi,
+                                call = rlang::caller_env()){
+  # check dimensions of psi
+  not_2d_or_one_col <- ncol(psi) != 1 | length(dim(psi)) != 2
+  if (not_2d_or_one_col) {
+    cli::cli_abort(
+      c(
+        "{.var psi} must be a 2D array with one column",
+        "but {.var psi} has dimensions {paste(dim(psi), collapse = 'x')}"
+      ),
+      call = call
+    )
+  }
+}
+
+check_params_same_rows <- function(p, 
+                                   psi,
+                                   call = rlang::caller_env()){
+  dim_p <- nrow(p)
+  dim_psi <- nrow(psi)
+  
+  not_same_nrows <- dim_p != dim_psi
+  
+  if (not_same_nrows) {
+    cli::cli_abort(
+      c(
+        "{.var p} and {.var psi} must have the same number of rows",
+        "But we see {.var p} and {.var psi} have:",
+        "{.var p}: {dim_p} {?row/rows}",
+        "{.var psi}: {dim_psi} {?row/rows}",
+        "Perhaps you need to coerce {.var p} or {.var psi} to an \\
+            appropriate matrix?"
+      ),
+      call = call
+    )
+  }
+}
+
+check_dim_positive_scalar_int <- function(dim,
+                                          call = rlang::caller_env()){
+  # check dim is a positive scalar integer
+  dim_old <- dim
+  dim <- as.integer(dim)
+  not_scalar_positive_integer <- length(dim) > 1 || dim <= 0 || !is.finite(dim)
+  if (not_scalar_positive_integer) {
+    cli::cli_abort(
+      c(
+        "{.var dim} must be a scalar positive integer, but was:",
+        "{dim_old}"
+      ),
+      call = call
     )
   }
 }
diff --git a/man/conditional_bernoulli.Rd b/man/conditional_bernoulli.Rd
index 091dcc1..0c0699d 100644
--- a/man/conditional_bernoulli.Rd
+++ b/man/conditional_bernoulli.Rd
@@ -45,6 +45,16 @@ a species was detected at each visit, and the probability of detection
 indicating whether the species was present (assumed to be the same at all
 visits) and the probability of being present.
 }
+\examples{
+\dontrun{
+cb <- conditional_bernoulli(
+  p = matrix(c(0.1,0.9), ncol = 2), 
+  psi = c(0.9),
+  dim = 1
+  )
+cb
+}
+}
 \references{
 MacKenzie, D. I., Nichols, J. D., Lachman, G. B., Droege, S.,
 Andrew Royle, J., & Langtimm, C. A. (2002). Estimating site occupancy rates