How to set up an additive linear predictor in liesel

[Seb-Lorek]

I currently face the issue of defining an additive linear predictor for the location with fixed and smooth effects in liesel. I use the GraphBuilder class to build my model. I created a group for the smooth effect as outlined in the Advanced Group Usage. My question is now how to set up the linear predictor node i currently use the following

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")

where smooth_group_1 is an instance of class PSpline(lsl.Group). Finally I build the graph via

gb = lsl.GraphBuilder().add(response)
gb.add_groups(smooth_group_1)

Is this the intended way of defining an additive linear predictor in liesel? Moreover am I adding the group correctly after calling .add(response)? I also saw that there exists a class to directly define distributional regression models via methods of the class DistRegBuilder. But there are so far no tutorials on how to set up a model with this class. Happy to receive some input.

[jobrachem]

Hi @Seb-Lorek, thanks for your question! 😊

The code you use looks to me like it would work. You are using the lsl.Calc class exactly like it is intended to be used, by providing the function that you need to be executed. It is indeed possible to use the DistRegBuilder for this, but since it is not well-documented at the moment, let's not get into that here.

One thing you could do is more clearly separate the linear part and the nonlinear part with a setup like this:

x_lin = ... # the design matrix for the linear effect

x_lin_node = lsl.Obs(x_lin, name="x_lin")
coef_lin = lsl.Param(..., name="beta_lin")
smooth_lin_calc = lsl.Calc(jnp.dot, x_lin, coef_lin)
smooth_lin = lsl.Var(smooth_lin_calc, name="smooth_lin")

smooth_1 = ... # defined in your code

location_calc = lsl.Calc(lambda *args: sum(args), smooth_lin, smooth_1) # corrected sum(*args) to sum(args)
predictor_location = lsl.Var(location_calc, name="predictor_location")

It will lead to a graph that looks like this:

But that is more of a maintainability / code organization suggestion. As I said, I think your code should work fine.

[Seb-Lorek]

Nice, thanks for the answer. The choice of defining the linear predictor node in the way I defined it above was driven by the fact that such a model set-up mimics the formula object one would define for such a model in R.

Just one more question to better understand how liesel adds a group to the graph. In the code above I added the group after I already built the full graph. Does liesel search through the nodes and adds the group on top of all the nodes that belong to this group, such that I can connect the group after I already built the graph via lsl.GraphBuilder().add(response)? But the ordering does not matter right so calling

gb = lsl.GraphBuilder().add_groups(smooth_group_1)
gb.add(response)

results in the same graph. How would liesel go about building the graph in the this case? This question might go too deep into the inner workings of liesel, but some high level abstraction would be nice for my mental model on how liesel builds the graph.

By the way I really like the flexible model building library in liesel as it is really close to the underlying math (i.e. how I would define such a model on a paper) which I always prefer over more abstract model building libraries that rely on just calling functions. The latter would probably correspond to how the DistRegBuilder works.

[jobrachem]

Actually you don't really need to explicitly add the group to the GraphBuilder - it will take the response as the root node and build the graph up from there. As I see it currently, adding a group is more of a convenience function for when you want to quickly plot the graph of a group. Then you can just throw the group into a GraphBuilder and plot its graph.

By the way I really like the flexible model building library in liesel as it is really close to the underlying math (i.e. how I would define such a model on a paper) which I always prefer over more abstract model building libraries that rely on just calling functions.

That's really nice to hear :)

[Seb-Lorek]

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]

Thanks for the message @Seb-Lorek, nice report! I think this is a bug, because it means the rank is not a formal part of the model graph. I'll need to have a bit of a closer look at that, but my first theory is that the SplineCoef class from the group usage tutorial is to blame.

P.S. you can make your messages even more easier to read by defining your codeblocks as

```python
a = lsl.Var(1.0, name="a")
```

This will activate syntax highlighting:

a = lsl.Var(1.0, name="a")

[jobrachem]

Issue for the rank topic: #131

[Seb-Lorek]

Maybe also one general suggestion would be to create more general tutorials on liesel, i.e. connecting nodes, building the graph and how liesel works under the hood. While model specific tutorials are always helpful and get you going fast, they don't give you a deep understanding on how a software package works. For me, this always becomes prominent as soon as I need to deviate from the model setup in the tutorials, even though your documentation certainly provides assistance. Of course in some of the tutorials you also touch up on these points.

[jobrachem]

I agree, improving the documentation in this regard would be beneficial. For the time being, I created an issue that holds this task: #108

[wiep]

The updated manuscript should be soon on arxix. In v1, the section on liesel model is completely outdated and does not reflect the implementation in liesel 0.2.0 and later.

[jeli2103]

Hey. Thanks for the great discussion. A small correction to your line about building the sum:
location_calc = lsl.Calc(lambda *args: sum(*args), smooth_lin, smooth_1)
should be
location_calc = lsl.Calc(lambda *args: sum(args), smooth_lin, smooth_1)

[jobrachem]

You're right @liesel-groupie , thanks! I corrected it in the answer above.

[jobrachem]

And welcome to the Liesel discussion board! 🎉 😊