Rliesel vs. Pure Python

[jeli2103]

Hello,
I have just discovered this nice discussion category. Actually I do not think you need to provide more tutorials. You can refer to discussions here and people could post some example codes? For example I had the same issue @Seb-Lorek with the rank not being a model node (and changed the gibbs sampler function, but now there is no need for this anymore).
So let me do a first step. I started with rliesel and now switched to pure Python. In the following I want to provide some code for both using an additive model with a linear effect and a P-Spline. The results are more or less the same (although the pure python sampling runs a lot into the maximum tree depth...):

library(rliesel)
library(mgcv)

library(reticulate)
reticulate::use_condaenv("---")

# simulate
set.seed(42)
n <- 500
x1 <- runif(n)
x2 <- runif(n)

y <- 2 + 5 * x1 + 5 * sin(x2 * 5) + rnorm(n)

df <- data.frame(y=y, x1=x1, x2=x2)

## Specify mean predictor
mu <- rliesel::predictor(~ x1 + s(x2, bs="ps", k=30), 
                         inverse_link = "Identity")

model <- rliesel::liesel(
  response = df$y,
  distribution = "Normal",
  predictors = list(
    loc = mu,
    scale = predictor(~1, inverse_link = "Exp")
  ),
  data=df
)

plot(model)

import jax.numpy as jnp
import liesel.model as lsl
import liesel.goose as gs

import time 

model = r.model  # not needed if own penalty is used

# set up kernels
builder = gs.EngineBuilder(seed=42, num_chains=3)
builder.set_model(lsl.GooseModel(model))

# set up gibbs kernel for smoothing param
for key in model.groups().keys():
  if 'np' in key:
    spline = model.groups()[key]
    tau2_kernels = lsl.tau2_gibbs_kernel(spline)
    builder.add_kernel(tau2_kernels)

# use nuts for coefficients
parameters = [name for name, var in model.vars.items() if var.parameter]
coef = []
for p in parameters:
  if 'tau2' not in p:
    coef.append(p)

builder.add_kernel(gs.NUTSKernel(coef))

# init
builder.set_initial_values(model.state)

# mcmc setting
builder.set_duration(warmup_duration=1000, posterior_duration=2000)

# start sampling
engine = builder.build()
start_time = time.time()
engine.sample_all_epochs()

rliesel_time = time.time() - start_time

# check
results = engine.get_results()
summary = gs.Summary(results).to_dataframe().reset_index()

# set up everything in python
import tensorflow_probability.substrates.jax.distributions as tfd
import tensorflow_probability.substrates.jax.bijectors as tfb
from liesel.distributions.mvn_degen import MultivariateNormalDegenerate
from liesel.goose.types import Array
import numpy as np

# define classes for psplines
class VarianceIG(lsl.Group):
    def __init__(
        self, name: str, a: float, b: float, start_value: float = 1000.0
    ) -> None:
        a_var = lsl.Var(a, name=f'{name}_a')
        b_var = lsl.Var(b, name=f'{name}_b')

        prior = lsl.Dist(tfd.InverseGamma, concentration=a_var, scale=b_var)
        s2 = lsl.Param(start_value, distribution=prior, name=name)
        super().__init__(name=name, a=a_var, b=b_var, tau2=s2)
        
        

class Linear(lsl.Group):
    def __init__(
        self, name: str, loc: float, scale: float, design: Array
    ) -> None:
        beta_loc = lsl.Var(loc, name=f'beta_{name}_loc')
        beta_scale = lsl.Var(scale, name=f'beta_{name}_scale')  # scale=sqrt(100^2)
        beta_dist = lsl.Dist(tfd.Normal, loc=beta_loc, scale=beta_scale)
        beta_li = lsl.Param(
            value=np.zeros([design.shape[1]]), distribution=beta_dist,
            name=f'beta_{name}'
        )
        design_node = lsl.Obs(design, name=f'X_{name}')
        eta_li_calc = lsl.Calc(jnp.dot, design_node, beta_li)
        eta_li = lsl.Var(eta_li_calc, name=f'eta_{name}')

        super().__init__(
            name=name, beta=beta_li, design_matrix=design_node, eta=eta_li
        )


class GaussCoef(lsl.Group):
    def __init__(
            self, name: str, penalty: Array, tau2: lsl.Param
    ) -> None:
        penalty_var = lsl.Var(penalty, name=f'{name}_K')

        # we save the rank of the penalty matrix here to use
        # it later in our  gibbs kernel
        rank = lsl.Data(np.linalg.matrix_rank(penalty),
                        _name=f'{name}_rank')

        prior = lsl.Dist(
            MultivariateNormalDegenerate.from_penalty,
            loc=0.0,
            var=tau2,
            pen=penalty_var,
            rank=rank
        )

        start_value = np.zeros(np.shape(penalty)[-1], np.float32)

        gamma = lsl.Param(start_value, distribution=prior, name=name)

        super().__init__(
            name, beta=gamma, K=penalty_var, tau2=tau2, rank=rank
        )
        

class Spline(lsl.Group):
    def __init__(
        self, name, basis_matrix: Array, penalty: Array, tau2_group: lsl.Group,
        # smooth_type: str
    ) -> None:
        coef_group = GaussCoef(
            name=f'gamma_{name}', penalty=penalty, tau2=tau2_group['tau2'],
        )

        design_matrix = lsl.Obs(basis_matrix, name=f'Z_{name}')
        eta = lsl.Var(
            lsl.Calc(jnp.dot, design_matrix, coef_group['beta']), name=f'eta_{name}'
        )

        group_vars = coef_group.nodes_and_vars | tau2_group.nodes_and_vars

        super().__init__(
            name=name,
            design_matrix=design_matrix,
            eta=eta,
            **group_vars
        )

# retrieve from mgcv
k = r.model.vars['loc_np0_K'].value
z = r.model.vars['loc_np0_X'].value

# Model
# Part 1: Model for the mean
# a) linear part
icept = np.ones([r.df.shape[0], 1])
x_li = np.concatenate([icept, r.df[['x1']]], axis=1)

linear = Linear(
    name='li',
    loc=0.0, scale=1000.0,
    design=x_li
)

# b) nonlinear part
# smoothing parameter
tau2_nl = VarianceIG(name='tau2_nl', a=0.01, b=0.01, start_value=5.0)
# model variable
pspline = Spline(
    name='nl', basis_matrix=z,
    penalty=k, tau2_group=tau2_nl
)


# bringing mean together
calc = lsl.Calc(lambda *args: sum(args), linear['eta'], pspline['eta'])
y_hat = lsl.Var(calc, name='mu')

# Part 2: Model for the standard deviation
sigma = VarianceIG(name='sigma', a=0.01, b=0.01, start_value=5.0)


# Part 3: Model for the response
y_dist = lsl.Dist(tfd.Normal, loc=y_hat, scale=sigma.vars['tau2'])
y = lsl.Var(r.df['y'].to_numpy(), distribution=y_dist, name='y')

# bringing the model together
gb = lsl.GraphBuilder().add(y)
# parameter constraints
gb.transform(sigma.vars['tau2'], tfb.Exp) 
# set up model and see graph
model2 = gb.build_model()
lsl.plot_vars(model2)

# set up kernels
builder2 = gs.EngineBuilder(seed=42, num_chains=3)
builder2.set_model(lsl.GooseModel(model2))
builder2.positions_included = ['sigma']

# gibbs for smoothing param
tau2_kernels = lsl.tau2_gibbs_kernel(pspline)
builder2.add_kernel(tau2_kernels)

# use nuts for coefficients and error variance
parameters = [name for name, var in model2.vars.items() if var.parameter]
param = []
for p in parameters:
    if 'tau2' not in p:
        param.append(p)

builder2.add_kernel(gs.NUTSKernel(param))

# init
builder2.set_initial_values(model2.state)

# mcmc setting
builder2.set_duration(warmup_duration=1000, posterior_duration=2000)

# start sampling
start_time = time.time()
engine2 = builder2.build()
engine2.sample_all_epochs()

# time
pyliesel_time = time.time() - start_time

# check
results2 = engine2.get_results()
summary2 = gs.Summary(results2).to_dataframe().reset_index()

# variable mean comparison
(summary.loc[summary.variable=='loc_p0_beta', 'mean'].round(2).to_numpy() ==
 summary2.loc[summary2.variable=='beta_li', 'mean'].round(2).to_numpy())
(summary.loc[summary.variable=='loc_np0_beta', 'mean'].round(2).to_numpy() ==
 summary2.loc[summary2.variable=='gamma_nl', 'mean'].round(2).to_numpy())
# time comparison
print(rliesel_time)
print(pyliesel_time)


gs.plot_param(results2, param='beta_li', param_index=1)
gs.plot_param(results2, param='gamma_nl', param_index=4)

[jobrachem]

Hi @liesel-groupie, thanks for sharing your code! 🎉 Awesome that you are using pure-Python liesel 😊

Do you have a theory why the sampling behavior of the R-coded and the Python-coded models is different? Without having gone through your code line-by-line, I would expect that they are the same model and therefore should behave the same way in sampling. What I found sometimes helps to make sampling more efficient is splitting up the sampling of regression coefficients for smooth functions into blocks. Then you would have one block for each function. You can see an example for this in the Comparing samplers tutorial (NUTS sampler).

As a side note: If you find the time, the showcase could benefit from just a few lines of verbal commentary on what is achieved in the individual code blocks. That would help readers understand what is going on.

[Ttlia]

If your primary focus is statistical analysis and visualization, and you're working in a field where R is commonly used (like academia or certain areas of industry), R might be a good choice.

If you're looking for a more versatile language that can be applied to various domains (including statistics and data science), and you want to explore machine learning and AI, Python might be a better fit.