Export a learned likelihood to Stan

stan_code() writes the likelihood learned by nle() out as Stan source: a functions block that recomputes \(\log q_\phi(x \mid \theta)\) in Stan's own language, with the trained weights travelling alongside as data. Stan differentiates the generated code itself, so NUTS gets exact gradients and nothing has to be linked against torch at run time.

Usage

stan_code(fit, name = "nsbi_log_lik", model = TRUE)

write_stan_model(fit, file, name = "nsbi_log_lik", model = TRUE)

stan_data(fit, x_obs = NULL, model = TRUE)

Arguments

fit

An nsbi_nle object from nle().

name

Prefix for the generated functions.

model

Generate a complete, runnable model (the default) or only the functions block, for #include-ing into a model of your own. In stan_data(), this must agree with the model a paired stan_code() call used: model = TRUE requires x_obs, because the generated model's data block declares N and x and there is nothing to fill them with otherwise; model = FALSE matches a functions-only export, which has no N/x to fill, so x_obs is optional there.

file

Path to write to.

x_obs

Observation to put in the data list. Rows are independent observations. Required when model = TRUE.

Value

stan_code() returns the Stan program as a single string; write_stan_model() returns file invisibly; stan_data() returns a named list ready for cmdstanr or rstan.

Why this is useful

The exported function is an ordinary Stan function of theta. That means the surrogate likelihood stops being the whole model and becomes one term in a model you write: put a hierarchical prior over the parameters, add covariates, mix in a second data source with a likelihood you do know, and let NUTS sample the lot. None of that is reachable from a posterior estimator, which only ever knows the one conditional it was trained on.

The generated interface

Two entry points, where w is the packed weight vector from stan_data():

real nsbi_log_lik_lpdf(vector x, vector theta, vector w);      // one observation
real nsbi_log_lik_sum_lpdf(matrix x, vector theta, vector w);  // rows are i.i.d.

Both take x and theta in the original units the simulator and prior use; the standardization the estimator trained under is folded into the generated code. Prefer the _sum form for repeated observations: for the MDN and the linear-Gaussian estimator the conditional distribution depends on theta alone, so it is computed once and reused across every row.

Why the code is generated rather than called

Python sbi hands its learned likelihood to Pyro or PyMC as a callable, because everything lives in one process and the sampler can differentiate the same graph the estimator was trained in. Stan cannot work that way: it compiles to C++ and needs the gradient inside its own autodiff. The alternative would be an external C++ header (--allow-undefined plus a USER_HEADER) linking torch into every Stan compile, with the gradients plumbed by hand. Generating source keeps the run-time dependency at zero and makes the result something you can read, edit, and check – which is what the package's own tests do, by evaluating the emitted functions and comparing them against log_lik().

Supported estimators

"mdn", "maf" and "linear_gaussian". "nsf" is not exported – its rational-quadratic spline transform would be a large and fragile block of generated Stan. Refit with "maf" if you need a flow.

nre() fits are not exported at all. What gets transpiled here is a density; a ratio estimator holds a classifier instead, and there is no p(x | theta) in it to write out.

Supported priors

Only model = TRUE needs a prior, since only then is there a model block to put one in. prior_uniform() and prior_normal() travel as data, so their bounds and their mean and standard deviation come from stan_data() and the compiled model does not care what they are. Every other named family (see prior_families), including prior_independent() products and prior_truncated() bounds, is written out as literal sampling statements with T[,] where the support was cut down. A prior_custom() cannot be written out at all: take stan_code(fit, model = FALSE) and write the model block yourself.

Examples

prior <- prior_uniform(c(mu = -3), c(mu = 3))
fit <- nle(prior, function(mu) c(y = rnorm(1, mu, 0.5)),
           n_simulations = 500, density_estimator = "linear_gaussian")

cat(substr(stan_code(fit), 1, 400))
#> // Generated by neuralsbi::stan_code(). Do not edit by hand.
#> // Surrogate likelihood q(x | theta) from an nle() fit (linear_gaussian, 500 simulations).
#> functions {
#>   // conditional mean, given the standardized parameter
#>   vector nsbi_log_lik_mean(vector ts, vector w) {
#>     return (to_matrix(segment(w, 1, 2), 2, 1))' * append_row(1.0, ts);
#>   }
#> 
#>   real nsbi_log_lik_lpdf(vector x, vector theta, vecto
str(stan_data(fit, matrix(rnorm(10), ncol = 1)), max.level = 1)
#> List of 6
#>  $ nsbi_nw  : int 3
#>  $ nsbi_w   : num [1:3] 9.32e-18 9.59e-01 2.83e-01
#>  $ N        : int 10
#>  $ x        : num [1:10, 1] 0.811 -1.586 -0.975 -0.782 1.202 ...
#>  $ nsbi_low : num -3
#>  $ nsbi_high: num 3
str(stan_data(fit, model = FALSE), max.level = 1)
#> List of 4
#>  $ nsbi_nw  : int 3
#>  $ nsbi_w   : num [1:3] 9.32e-18 9.59e-01 2.83e-01
#>  $ nsbi_low : num -3
#>  $ nsbi_high: num 3