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.0145  1.0243  2057.93 1.0009  2570
        1       0.0402  1.0227  1749.67 1.0008  2570
CPU times: user 23.3 ms, sys: 793 μs, total: 24.1 ms
Wall time: 7.34 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.0293 1.0075  781.98  1.0087  949
        1       0.0747  1.0543  654.57  1.0047  949
CPU times: user 1.42 s, sys: 258 ms, total: 1.68 s
Wall time: 1.05 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.0103 1.0013  2741.72 1.0004  3952
        1       -0.0114 0.9977  3432.56 1.0007  3952
CPU times: user 23.4 ms, sys: 0 ns, total: 23.4 ms
Wall time: 7.14 ms