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._stacked.shape[0]
    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.0370  0.9487  1003.89 1.0058  949
        1       0.0289  0.9732  943.73  1.0016  949
CPU times: user 12 ms, sys: 993 μs, total: 13 ms
Wall time: 4.71 ms

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.0127  1.0221  1576.80 1.0016  2090
        1       -0.0161 0.9791  1877.34 1.0021  2090
CPU times: user 1.63 s, sys: 494 ms, total: 2.12 s
Wall time: 1.52 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.0142  1.0107  2833.27 1.0002  3985
        1       -0.0218 0.9991  2675.36 1.0002  3985
CPU times: user 20.4 ms, sys: 24 μs, total: 20.4 ms
Wall time: 6.97 ms