Group usage tutorial: The rank seems to be overlooked in the model graph

[jobrachem]

Originally posted by @Seb-Lorek in https://github.com/liesel-devs/liesel/discussions/105#discussioncomment-7443219

Actually as far as I can see, I need to add the group. Otherwise the rank of the smooth is not added to the model.state dictionary, which I need for the tau2Gibbs Kernel. Here the code (excluding the group set up see here for details)

# Set up model in liesel
# Fixed parameter prior
beta_loc = lsl.Var(0.0, name="beta_loc")
beta_scale = lsl.Var(100.0, name="beta_scale")

# Set up the fixed parameters
beta_dist = lsl.Dist(tfjd.Normal, loc=beta_loc, scale=beta_scale)
beta = lsl.Param(value=np.array([0.0]), distribution=beta_dist, name="beta")

# Set up the smooth parameters
tau2_group = VarianceIG(name="tau2", a=1.0, b=0.0005)

penalty = X.smooth_pen_mat_cent[0]
smooth_group_1 = PSpline(name="smooth_1", basis_matrix=X.design_mat_cent[1], penalty=penalty, tau2_group=tau2_group)

# Set up the scale 
sigma_a = lsl.Var(0.01, name="sigma_a")
sigma_b = lsl.Var(0.01, name="sigma_b")

sigma_dist = lsl.Dist(tfjd.InverseGamma, concentration=sigma_a, scale=sigma_b)
sigma = lsl.Param(value=10.0, distribution=sigma_dist, name="sigma")

Z = lsl.Obs(X.fixed_data, name="Z")

lpred_loc_fn = lambda z, beta, smooth_1: jnp.dot(z, beta) + smooth_1
lpred_loc_calc = lsl.Calc(lpred_loc_fn, z=Z, beta=beta, smooth_1=smooth_group_1["smooth"])

lpred_loc = lsl.Var(lpred_loc_calc, name="lpred_loc")

response_dist = lsl.Dist(tfjd.Normal, loc=lpred_loc, scale=sigma)
response = lsl.Var(y.to_numpy(), distribution=response_dist, name="response")

Building the model without gb.add_groups(smooth_group_1)

gb = lsl.GraphBuilder().add(response)
gb.transform(sigma, tfjb.Exp)

model = gb.build_model()
model

Setting up the EngineBuilder ( for the setup of the Gibbs Kernel tau2_gibbs_kernel for tau2 see here)

builder = gs.EngineBuilder(seed=1337, num_chains=4)

builder.set_model(lsl.GooseModel(model))
builder.set_initial_values(model.state)

builder.add_kernel(tau2_gibbs_kernel(smooth_group_1))
builder.add_kernel(gs.IWLSKernel(["smooth_1_coef"]))
builder.add_kernel(gs.IWLSKernel(["beta"]))
builder.add_kernel(gs.NUTSKernel(["sigma_transformed"]))

builder.set_duration(warmup_duration=1000, posterior_duration=1000)

builder.positions_included = ["sigma"]

engine = builder.build()

Finally calling

engine.sample_all_epochs()

results in an error due to not finding the rank of the smooth coefficients in the model_state object in the EngineBuilder.

[jobrachem]

Hi @Seb-Lorek, I think I found the mistake and was able to fix it. Unfortunately, I cannot run your code example, because X is undefined (and possibly more). Could you post an updated version that I can run for testing?

[Seb-Lorek]

Sure here we go. Here the code (excluding the group set up see here for details)

# Dependencies
import numpy as np
import jax
import jax.numpy as jnp
from scipy.interpolate import BSpline as bs
import liesel.goose as gs
import liesel.model as lsl
from liesel.distributions.mvn_degen import MultivariateNormalDegenerate
from liesel.goose.types import Array
import tensorflow_probability.substrates.jax.distributions as tfjd
import tensorflow_probability.substrates.jax.bijectors as tfjb

rng = np.random.default_rng(42)
# Sample size and true parameters
n = 1000
true_beta = np.array([3.0, 0.2, -0.5])
true_sigma = 1.0

# Data-generating process
x0 = rng.uniform(size=n, low=-3, high=3)
X_mat = np.column_stack((np.ones(n), x0, x0**2))
eps = rng.normal(scale=true_sigma, size=n)
response_vec = X_mat @ true_beta + eps

# Set up model in liesel
# Fixed parameter prior
beta_loc = lsl.Var(0.0, name="beta_loc")
beta_scale = lsl.Var(100.0, name="beta_scale")

# Set up the fixed parameters
beta_dist = lsl.Dist(tfjd.Normal, loc=beta_loc, scale=beta_scale)
beta = lsl.Param(value=np.array([0.0]), distribution=beta_dist, name="beta")

# Set up the smooth parameters
tau2_group = VarianceIG(name="tau2", a=1.0, b=0.0005)

# Create the centered B-spline basis matrix 
n_knots = 20 
degree = 2
knots = np.linspace(x0.min(), x0.max(), num=n_knots+degree+1)
smooth_mat = bs.design_matrix(x0, t=knots, k=degree, extrapolate=True).toarray()

# Create the penalty 
diff_mat = np.diff(np.eye(smooth_mat.shape[-1]), n=degree, axis=0)
pen = np.dot(diff_mat.T, diff_mat) 

# Include the centering constraints 
c = np.mean(smooth_mat, axis=0)
c = np.expand_dims(c, axis=1)
Q, _ = np.linalg.qr(c, mode="complete")
Tb = Q[:,1:]
smooth_mat_cent = np.dot(smooth_mat, Tb)
pen_cent = Tb.T @ pen @ Tb

# Construct the group 
smooth_group_1 = PSpline(name="smooth_1", basis_matrix=smooth_mat_cent, penalty=pen_cent, tau2_group=tau2_group)

# Set up the scale 
sigma_a = lsl.Var(0.01, name="sigma_a")
sigma_b = lsl.Var(0.01, name="sigma_b")

sigma_dist = lsl.Dist(tfjd.InverseGamma, concentration=sigma_a, scale=sigma_b)
sigma = lsl.Param(value=10.0, distribution=sigma_dist, name="sigma")

X = lsl.Obs(np.ones((response_vec.shape[-1],1)), name="X")

lpred_loc_fn = lambda x, beta, smooth_1: jnp.dot(x, beta) + smooth_1
lpred_loc_calc = lsl.Calc(lpred_loc_fn, x=X, beta=beta, smooth_1=smooth_group_1["smooth"])

lpred_loc = lsl.Var(lpred_loc_calc, name="lpred_loc")

response_dist = lsl.Dist(tfjd.Normal, loc=lpred_loc, scale=sigma)
response = lsl.Var(response_vec, distribution=response_dist, name="response")

Building the model without gb.add_groups(smooth_group_1)

gb = lsl.GraphBuilder().add(response)
gb.transform(sigma, tfjb.Exp)

model = gb.build_model()
model

Setting up the EngineBuilder ( for the setup of the Gibbs Kernel tau2_gibbs_kernel for tau2 see here)

builder = gs.EngineBuilder(seed=1337, num_chains=4)

builder.set_model(lsl.GooseModel(model))
builder.set_initial_values(model.state)

builder.add_kernel(tau2_gibbs_kernel(smooth_group_1))
builder.add_kernel(gs.IWLSKernel(["smooth_1_coef"]))
builder.add_kernel(gs.IWLSKernel(["beta"]))
builder.add_kernel(gs.NUTSKernel(["sigma_transformed"]))

builder.set_duration(warmup_duration=1000, posterior_duration=1000)

builder.positions_included = ["sigma"]

engine = builder.build()

Finally calling

engine.sample_all_epochs()

results in an error due to not finding the rank of the smooth coefficients in the model_state object in the EngineBuilder.

[jobrachem]

Thanks a lot @Seb-Lorek ! 🙏

[jobrachem]

Thank you again for the testing code. The group tutorial gets fixed with #135 .

For the record, here is what went wrong: In the definition of SplineCoef, we created a lsl.Data node for the rank and added it to the group - but we did not use it in the lsl.Dist of the coefficient. For that reason, it was not formally connected to the rest of the graph, and the lsl.GraphBuilder could not retrieve it.

Defining the SplineCoef like this fixes the problem:

class SplineCoef(lsl.Group):
    def __init__(self, name: str, penalty: Array, tau2: lsl.param) -> None:
        penalty_var = lsl.Var(penalty, name=f"{name}_penalty")

        evals = jax.numpy.linalg.eigvalsh(penalty)
        rank = lsl.Data(jnp.sum(evals > 0.0), _name=f"{name}_rank")
        _log_pdet = jnp.log(jnp.where(evals > 0.0, evals, 1.0)).sum()
        log_pdet = lsl.Data(_log_pdet, _name=f"{name}_log_pdet")

        prior = lsl.Dist(
            MultivariateNormalDegenerate.from_penalty,
            loc=0.0,
            var=tau2,
            pen=penalty_var,
            rank=rank,
            log_pdet=log_pdet
        )
        start_value = np.zeros(np.shape(penalty)[-1], np.float32)

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

        super().__init__(name, coef=coef, penalty=penalty_var, tau2=tau2, rank=rank)