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.0154 0.9873  2284.23 1.0006  2746
        1       0.0063  1.0297  2435.32 1.0024  2746
CPU times: user 22.2 ms, sys: 0 ns, total: 22.2 ms
Wall time: 6.76 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.0005 1.0106  2796.71 1.0002  3210
        1       0.0199  0.9823  2878.16 1.0003  3210
CPU times: user 1.97 s, sys: 673 ms, total: 2.64 s
Wall time: 1.86 s

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.0094  0.9758  3446.45 1.0002  4000
        1       0.0012  0.9875  3422.25 1.0008  4000
CPU times: user 17.2 ms, sys: 1.88 ms, total: 19 ms
Wall time: 6.78 ms