Implemented qtrunc.poisson() (#54)

ocbe-uio · Aug 13, 2024 · d64cb14 · d64cb14
1 parent 7cc797f
commit d64cb14
Show file tree

Hide file tree

Showing 5 changed files with 130 additions and 0 deletions.
diff --git a/R/qtrunc.R b/R/qtrunc.R
@@ -142,6 +142,14 @@ qtrunc.normal <- function(
   return(q)
 }
 
+qtrunc.poisson <- function(p, lambda, a = 0, b = Inf, ..., lower.tail, log.p) {
+  F_a <- ppois(a - 1L, lambda, lower.tail, FALSE)
+  F_b <- ppois(b, lambda, lower.tail, FALSE)
+  rescaled_p <- rescale_p(p, F_a, F_b, lower.tail, log.p)
+  q <- qpois(rescaled_p, lambda, lower.tail, FALSE)
+  return(q)
+}
+
 rescale_p <- function(p, F_a, F_b, lower.tail, log.p) {
   if (log.p) {
     p <- exp(p)

diff --git a/tests/testthat/test-qtrunc-truncated-a.R b/tests/testthat/test-qtrunc-truncated-a.R
@@ -303,3 +303,33 @@ test_that("qtrunc() works as expected (normal)", {
     }
   }
 })
+
+test_that("qtrunc() works as expected (poisson)", {
+  fam <- "poisson"
+  for (lt in c(TRUE, FALSE)) {
+    for (lg in c(FALSE, TRUE)) {
+      for (i in seq_len(3L)) {
+        lambda <- sample(1:50, 1L)
+        pt <- runif(i)
+        a <- qtrunc(min(pt) / 2, fam, lambda, lower.tail = TRUE, log.p = FALSE)
+        if (lg) pt <- log(pt)
+        q_trunc <- qtrunc(pt, fam, lambda, a = a, lower.tail = lt, log.p = lg)
+        q_stats <- qpois(pt, lambda, lower.tail = lt, log.p = lg)
+        expect_length(q_trunc, i)
+        for (ii in seq_along(pt)) {
+          expect_gte(q_trunc[ii], q_stats[ii])
+          q_hi <- min(q_trunc[ii] + 1L)
+          q_lo <- max(q_trunc[ii] - 1L, 0L, a)
+          ptr_1 <- ptrunc(q_lo, fam, lambda, a = a, lower.tail = lt, log.p = lg)
+          ptr_2 <- ptrunc(q_hi, fam, lambda, a = a, lower.tail = lt, log.p = lg)
+          # because pt will have been rounded
+          if (lt) {
+            expect_lte(ptr_1, ptr_2)
+          } else {
+            expect_gte(ptr_1, ptr_2)
+          }
+        }
+      }
+    }
+  }
+})
diff --git a/tests/testthat/test-qtrunc-truncated-ab.R b/tests/testthat/test-qtrunc-truncated-ab.R
@@ -301,3 +301,33 @@ test_that("qtrunc() works as expected (normal)", {
     }
   }
 })
+
+test_that("qtrunc() works as expected (poisson)", {
+  fam <- "poisson"
+  for (lt in c(TRUE, FALSE)) {
+    for (lg in c(FALSE, TRUE)) {
+      for (i in seq_len(3L)) {
+        lambda <- sample(1:50, 1L)
+        pt <- runif(i)
+        a <- qtrunc(min(pt) / 2, fam, lambda, lower.tail = TRUE, log.p = FALSE)
+        b <- qtrunc(sqrt(max(pt)), fam, lambda, lower.tail = TRUE, log.p = FALSE)
+        if (lg) pt <- log(pt)
+        q_trunc <- qtrunc(pt, fam, lambda, a = a, b = b, lower.tail = lt, log.p = lg)
+        q_stats <- qpois(pt, lambda, lower.tail = lt, log.p = lg)
+        expect_length(q_trunc, i)
+        for (ii in seq_along(pt)) {
+          q_lo <- max(q_trunc[ii] - 1L, 0L, a)
+          q_hi <- min(q_trunc[ii] + 1L, b)
+          ptr_1 <- ptrunc(q_lo, fam, lambda, a = a, b = b, lower.tail = lt, log.p = lg)
+          ptr_2 <- ptrunc(q_hi, fam, lambda, a = a, b = b, lower.tail = lt, log.p = lg)
+          # because pt will have been rounded
+          if (lt) {
+            expect_lte(ptr_1, ptr_2)
+          } else {
+            expect_gte(ptr_1, ptr_2)
+          }
+        }
+      }
+    }
+  }
+})
diff --git a/tests/testthat/test-qtrunc-truncated-b.R b/tests/testthat/test-qtrunc-truncated-b.R
@@ -308,3 +308,33 @@ test_that("qtrunc() works as expected (normal)", {
     }
   }
 })
+
+test_that("qtrunc() works as expected (poisson)", {
+  fam <- "poisson"
+  for (lt in c(TRUE, FALSE)) {
+    for (lg in c(FALSE, TRUE)) {
+      for (i in seq_len(3L)) {
+        lambda <- sample(1:50, 1L)
+        pt <- runif(i)
+        b <- qtrunc(sqrt(max(pt)), fam, lambda, lower.tail = TRUE, log.p = FALSE)
+        if (lg) pt <- log(pt)
+        q_trunc <- qtrunc(pt, fam, lambda, b = b, lower.tail = lt, log.p = lg)
+        q_stats <- qpois(pt, lambda, lower.tail = lt, log.p = lg)
+        expect_length(q_trunc, i)
+        for (ii in seq_along(pt)) {
+          expect_lte(q_trunc[ii], q_stats[ii])
+          q_lo <- max(q_trunc[ii] - 1L, 0L)
+          q_hi <- min(q_trunc[ii] + 1L, b)
+          ptr_1 <- ptrunc(q_lo, fam, lambda, b = b, lower.tail = lt, log.p = lg)
+          ptr_2 <- ptrunc(q_hi, fam, lambda, b = b, lower.tail = lt, log.p = lg)
+          # because pt will have been rounded
+          if (lt) {
+            expect_lte(ptr_1, ptr_2)
+          } else {
+            expect_gte(ptr_1, ptr_2)
+          }
+        }
+      }
+    }
+  }
+})
diff --git a/tests/testthat/test-qtrunc-untruncated.R b/tests/testthat/test-qtrunc-untruncated.R
@@ -297,3 +297,35 @@ test_that("qtrunc() works as expected (normal)", {
     }
   }
 })
+
+test_that("qtrunc() works as expected (poisson)", {
+  fam <- "poisson"
+  for (lt in c(TRUE, FALSE)) {
+    for (lg in c(FALSE, TRUE)) {
+      for (i in seq_len(3L)) {
+        lambda <- sample(1:50, 1L)
+        pt <- runif(i)
+        if (lg) pt <- log(pt)
+        q_trunc <- qtrunc(pt, fam, lambda, lower.tail = lt, log.p = lg)
+        q_stats <- qpois(pt, lambda, lower.tail = lt, log.p = lg)
+        expect_length(q_trunc, i)
+        for (ii in seq_along(pt)) {
+          expect_equal(q_trunc[ii], q_stats[ii])
+          # Working back to p from q
+          q_lo <- max(q_trunc[ii] - 1L, 0L)
+          q_hi <- min(q_trunc[ii] + 1L)
+          ptr_1 <- ptrunc(q_lo, fam, lambda, lower.tail = lt, log.p = lg)
+          ptr_2 <- ptrunc(q_hi, fam, lambda, lower.tail = lt, log.p = lg)
+          # because pt will have been rounded
+          if (q_trunc[ii] > 0L && lt) {
+            expect_gte(pt[ii], ptr_1)
+            expect_lte(pt[ii], ptr_2)
+          } else if (q_trunc[ii] > 0L && !lt) {
+            expect_lte(pt[ii], ptr_1)
+            expect_gte(pt[ii], ptr_2)
+          }
+        }
+      }
+    }
+  }
+})