Filter non-finite de_sample() draws in sample() for unbounded priors too

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: neuralsbi
 Title: Neural Simulation-Based Inference
-Version: 0.6.20
+Version: 0.6.21
 Authors@R:
     person("Pedro", "Nascimento de Lima", email = "plima@rand.org",
            role = c("aut", "cre"), comment = c(ORCID = "0000-0001-9057-198X"))

NEWS.md

@@ -1,4 +1,6 @@
-# neuralsbi 0.6.20
+# neuralsbi 0.6.21
+
+* **`sample()` on an unbounded posterior (e.g. `prior_normal()`) no longer lets a non-finite draw from the density estimator through.** `sample.nsbi_posterior()`'s non-finite-row filter ran only inside `if (bounded)`, since it was added by #234/#236 to work around `within_support()` returning `NA` (not `FALSE`) for a NaN row -- a problem specific to the bounded rejection-sampling path. That left `bounded <- !is.null(prior$lower) || !is.null(prior$upper)` `FALSE` for an unbounded prior, the common case in the package's own NPE examples, with no filter at all: a NaN/Inf row that an under-trained MAF/NSF/MDN occasionally produces was `rbind`'d straight into the returned draws matrix, `attr(draws, "acceptance_rate")` still reported `1.0`, and the corruption propagated silently into `summary()`, `pairplot()`, and `sbc()`/`tarp()` diagnostics -- or surfaced downstream in `map_estimate()` as a confusing "`theta` contains non-finite value" error blaming the seed draw. `sample.nsbi_posterior()` now drops a non-finite row from `de_sample()`'s output unconditionally, before the `bounded` branch's `within_support()` check runs, so `acceptance_rate` reflects the drop for every prior (#244) (#245).
 
 * **`log_lik()`/`log_ratio()`'s `max_batch` now bounds memory for an MDN-based fit even when `theta`, not `x`, is the large dimension.** For MAF/NSF/`linear_gaussian`/NRE, `cross_iid()` (`R/likelihood.R`) chunks the `(theta, x)` cross product by `theta` rows, so `max_batch` bounds memory regardless of which side is large. The MDN's fast i.i.d. path, `mdn_iid_blocks()`, only chunked `x`: it ran `mdn_mixture()` -- the MLP forward pass and Cholesky assembly -- over the whole of `theta` in one call before any chunking happened, so scoring a dense `theta` grid (a profile-likelihood plot, say) against a handful of observations materialized a `(n_theta, K, dim_theta, dim_theta)` tensor however small `max_batch` was set. `mdn_iid_blocks()` now blocks `theta` first, the same way `cross_iid()` does, and chunks observations within each block as before; `mdn_trace_cache()`'s TorchScript shortcut only fires when a single call would cover both dimensions, since a trace recorded at one shape can't stand in for the chunked path (#240) (#243).
 

R/posterior.R

@@ -165,13 +165,23 @@ sample.nsbi_posterior <- function(x, size = 1000, n = size, obs = NULL,
     draw_std <- de_sample(fit$de, xo_std, n_needed)
     draw <- invert_standardizer(fit$std_theta, draw_std)
     n_tried <- n_tried + n_needed
+    # A non-finite row from de_sample() (an under-trained MAF/NSF/MDN can
+    # produce one) needs dropping regardless of whether the prior is bounded
+    # -- unlike within_support()'s NA-vs-FALSE issue below, this filter has
+    # nothing to do with support and ran only inside `if (bounded)` before
+    # #244, so an unbounded prior (prior_normal(), the common case) let a
+    # NaN/Inf draw straight into the returned matrix with acceptance_rate
+    # still reporting 1.0.
+    draw <- draw[apply(is.finite(draw), 1, all), , drop = FALSE]
     if (bounded) {
       # within_support() returns NA for a NaN/NA row, and R's matrix indexing
       # keeps (rather than drops) a row selected by an NA logical index and
       # fills it with NA -- so a NaN draw from the density estimator would
       # otherwise survive as an all-NA row counted toward n (#234). Coerce NA
       # to FALSE so a non-finite draw is rejected the same way an
-      # out-of-bounds one is.
+      # out-of-bounds one is. (The finite-row filter above already removes
+      # non-finite rows, so this NA never actually arises here anymore, but
+      # within_support() is still the source of truth for the bound itself.)
       ok <- within_support(prior, draw)
       ok[is.na(ok)] <- FALSE
       draw <- draw[ok, , drop = FALSE]

tests/testthat/test-posterior-nonfinite-de-draw.R

@@ -9,6 +9,14 @@
 # NaN draw the way an under-trained MAF/NSF/MDN might, and check that a
 # non-finite draw is rejected rather than silently accepted or crashing with
 # base R's unrelated "missing value where TRUE/FALSE needed".
+#
+# GitHub #244: the fix above only ran inside `if (bounded)`, since #234's
+# NA-vs-FALSE problem is specific to within_support(). That left the common
+# case -- an unbounded prior, e.g. prior_normal() -- with no filter at all: a
+# NaN/Inf row from de_sample() went straight into sample()'s returned matrix,
+# and attr(draws, "acceptance_rate") still reported 1.0. The tests below
+# repeat the sample() case with prior_normal() instead of prior_uniform() to
+# cover the unbounded path.
 
 test_that("sample() rejects a NaN draw from the density estimator instead of returning it", {
   set.seed(30)
@@ -87,3 +95,62 @@ test_that("map_estimate()'s objective does not crash when queried at a non-finit
   expect_equal(queried, Inf)
   expect_true(within_support(prior, matrix(map, nrow = 1)))
 })
+
+test_that("sample() rejects a NaN draw from the density estimator with an unbounded prior (#244)", {
+  set.seed(33)
+  prior <- prior_normal(mean = 0.5, sd = 1)
+  simulator <- function(theta) theta + stats::rnorm(1, sd = 0.05)
+  fit <- npe(prior, simulator, n_simulations = 500,
+             density_estimator = "linear_gaussian")
+  post <- posterior(fit, x_obs = 0.5)
+
+  # Same technique as the bounded-prior test above, but prior_normal() has no
+  # `lower`/`upper`, so bounded is FALSE and the fix has to filter without
+  # relying on within_support() at all.
+  real_de_sample <- de_sample
+  call_count <- 0L
+  local_mocked_bindings(
+    de_sample = function(de, x, n) {
+      call_count <<- call_count + 1L
+      draw <- real_de_sample(de, x, n)
+      if (call_count == 1L) draw[1, ] <- NaN
+      draw
+    }
+  )
+
+  draws <- sample(post, n = 50, max_sampling_batches = 10)
+  expect_equal(nrow(draws), 50L)
+  expect_false(anyNA(draws))
+  # without the fix, the NaN row is kept (not dropped), so round one alone
+  # would already satisfy n and no second round would run
+  expect_gt(call_count, 1L)
+})
+
+test_that("sample()'s acceptance_rate reflects a dropped non-finite row with an unbounded prior (#244)", {
+  set.seed(34)
+  prior <- prior_normal(mean = 0.5, sd = 1)
+  simulator <- function(theta) theta + stats::rnorm(1, sd = 0.05)
+  fit <- npe(prior, simulator, n_simulations = 500,
+             density_estimator = "linear_gaussian")
+  post <- posterior(fit, x_obs = 0.5)
+
+  # Every draw in the one and only batch is corrupted, so n_needed rows are
+  # tried and none survive the filter -- acceptance_rate must report that
+  # rather than the pre-fix 1.0, and sample() should warn about the shortfall
+  # the same way it does for a bounded prior leaking mass.
+  real_de_sample <- de_sample
+  local_mocked_bindings(
+    de_sample = function(de, x, n) {
+      draw <- real_de_sample(de, x, n)
+      draw[] <- NaN
+      draw
+    }
+  )
+
+  expect_warning(
+    draws <- sample(post, n = 20, max_sampling_batches = 1),
+    "0/20 samples inside prior support"
+  )
+  expect_equal(nrow(draws), 0L)
+  expect_equal(attr(draws, "acceptance_rate"), 0)
+})