Using Walnuts from Python

This notebook will show to how run the same model (a simple standard normal) implemented in Python, Numba (a just-in-time compiler package), and Stan.

import walnutpie

def summarize(name, fit):
    summarizer = walnutpie.Summarizer(fit)
    mean = summarizer.mean()
    std = summarizer.standard_deviation()
    ess = summarizer.ess()
    r_hat = summarizer.r_hat()
    draws = summarizer._num_draws
    print(f"{name}\tdim\tmean\tstd\tess\trhat\tdraws")
    for i in range(len(mean)):
        print(
            f"\t{i}\t{mean[i]:.4f}\t{std[i]:.4f}\t{ess[i]:.2f}\t{r_hat[i]:.4f}\t{draws}"
        )

import os
import bridgestan

stan_code = os.path.join(
    bridgestan.compile.get_bridgestan_path(), "test_models/multi/multi.stan"
)
with open(stan_code, 'r') as f:
    print(f.read())

m = bridgestan.StanModel(
    stan_code,
    {"M": 2, "N": 0, "P": 0},
    make_args=["STAN_THREADS=1"],
)

data {
  int<lower=0> M;
  int<lower=0> N;
  int<lower=0> P;
}
parameters {
  vector[M] alpha;
}
model {
  alpha ~ normal(0, 1);
}

%%time
summarize("stan", walnutpie.walnuts_stan(m, seed=1234))

stan    dim     mean    std     ess     rhat    draws
        0       0.0632  0.9929  973.07  1.0032  1293
        1       0.0506  0.9751  1155.42 1.0043  1293
CPU times: user 17.9 ms, sys: 0 ns, total: 17.9 ms
Wall time: 5.62 ms

<timed eval>:1: UserWarning: Setting 'seed' without also disabling adaptive stopping (by setting min and max number of iterations to the same value, for both warmup and sampling) will not lead to reproducible sampling due to thread scheduling!

Python

Defining a pure-python log density is simple and highly flexible, but will usually be slower than the other options due to the extra overhead of the Python language

import numpy as np
import scipy.stats


def logp(x):
    return np.sum(scipy.stats.norm.logpdf(x)), -x

%%time
summarize("pyfunc", walnutpie.walnuts_pyfunc(logp, num_params=2))

pyfunc  dim     mean    std     ess     rhat    draws
        0       -0.0224 1.0080  718.11  1.0016  831
        1       0.0231  0.9421  691.55  1.0024  831
CPU times: user 1.1 s, sys: 181 ms, total: 1.28 s
Wall time: 772 ms

Numba

If we are willing to use `numba <https://numba.pydata.org/>`__, we can get much faster!

import numba
from numba import types
from numba_stats import norm


@numba.cfunc(
    types.intc(
        types.size_t,
        types.CPointer(types.double),
        types.CPointer(types.double),
        types.CPointer(types.double),
        types.voidptr,
    ),
    nopython=True,
)
def logp_numba(size, x_, grad_, lp, _):
    x = numba.carray(x_, size)
    lp[0] = norm.logpdf(x, 0.0, 1.0).sum()
    grad = numba.carray(grad_, size)
    grad[:] = -x
    return 0

%%time
summarize("numba", walnutpie.walnuts_pyfunc(logp_numba, num_params=2))

numba   dim     mean    std     ess     rhat    draws
        0       -0.0054 0.9677  3392.59 1.0002  3954
        1       -0.0031 1.0057  3231.41 1.0003  3954
CPU times: user 18.9 ms, sys: 0 ns, total: 18.9 ms
Wall time: 6.86 ms