# Linear regression
In this tutorial, we build a linear regression model with Liesel and
estimate it with Goose. Our goal is to illustrate the most important
features of the software in a straightforward context.
## Model building with Liesel
Liesel is based on the concept of probabilistic graphical models (PGMs)
to represent (primarily Bayesian) statistical models, so let us start
with a very brief look at what PGMs are and how they are implemented in
Liesel.
### Probabilistic graphical models
In a PGM, each variable is represented as a node. There are two basic
types of nodes in Liesel: strong and weak nodes. A strong node is a node
whose value is defined “outside” of the model, for example, if the node
represents some observed data or a parameter (parameters are usually set
by an inference algorithm such as an optimizer or sampler). In contrast,
a weak node is a node whose value is defined “within” the model, that
is, it is a deterministic function of some other nodes. An
exp-transformation mapping a real-valued parameter to a positive number,
for example, would be a weak node.
In addition, each node can have an optional probability distribution.
The probability density or mass function of the distribution evaluated
at the value of the node gives its log-probability. In a typical
Bayesian regression model, the response node would have a normal
distribution and the parameter nodes would have some prior distribution
(for example, a normal-inverse-gamma prior). The following table shows
the different node types and some examples of their use cases.
| | **Strong node** | **Weak node** |
|--------------------------|------------------------------|----------------------------------------------------|
| **With distribution** | Response, parameter, … | Copula, … |
| **Without distribution** | Covariate, hyperparameter, … | Inverse link function, parameter transformation, … |
A PGM is essentially a collection of connected nodes. Two nodes can be
connected through a directed edge, meaning that the first node is an
input for the value or the distribution of the second node. Nodes
*without* an edge between them are assumed to be conditionally
independent, allowing us to factorize the model log-probability as
$$
\log p(\text{Model}) = \sum_{\text{Node $\in$ Model}} \log p(\text{Node} \mid \text{Inputs}(\text{Node})).
$$
### Imports
Before we can generate the data and build the model graph, we need to
load Liesel and a number of other packages. We usually import the model
building library `liesel.model` as `lsl`, and the MCMC library
`liesel.goose` as `gs`.
``` python
import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
# We use distributions and bijectors from tensorflow probability
import tensorflow_probability.substrates.jax.distributions as tfd
import tensorflow_probability.substrates.jax.bijectors as tfb
import liesel.goose as gs
import liesel.model as lsl
rng = np.random.default_rng(42)
```
### Generating the data
Now we can simulate 500 observations from the linear regression model
$y_i \sim \mathcal{N}(\beta_0 + \beta_1 x_i, \;\sigma^2)$ with the true
parameters $\boldsymbol{\beta} = (\beta_0, \beta_1)' = (1, 2)'$ and
$\sigma = 1$. The relationship between the response $y_i$ and the
covariate $x_i$ is visualized in the following scatterplot.
``` python
# sample size and true parameters
n = 500
true_beta = np.array([1.0, 2.0])
true_sigma = 1.0
# data-generating process
x0 = rng.uniform(size=n)
X_mat = np.column_stack([np.ones(n), x0])
eps = rng.normal(scale=true_sigma, size=n)
y_vec = X_mat @ true_beta + eps
# plot the simulated data
plt.scatter(x0, y_vec)
plt.title("Simulated data from the linear regression model")
plt.xlabel("Covariate x")
plt.ylabel("Response y")
plt.show()
```
### Building the model graph
The graph of a Bayesian linear regression model is a tree, where the
hyperparameters of the prior are the leaves and the response is the
root. To build this tree in Liesel, we need to start from the leaves and
work our way down to the root. As the most basic building blocks of a
model, Liesel provides the {class}`.Var` class for instantiating
variables and the {class}`.Dist`, class for wrapping probability
distributions. The {class}`.Var` class comes with four constructors,
namely {meth}`.Var.new_param` for parameters, {meth}`.Var.new_obs` for
observed data, {meth}`.Var.new_calc` for variables that are
deterministic functions of other variables in the model, and
{meth}`.Var.new_value` for fixed values. See also [Model Building
(liesel.model)](model_overview).
#### The regression coefficients
Let’s assume the weakly informative prior
$\beta_0, \beta_1 \sim \mathcal{N}(0, 100^2)$ for the regression
coefficients. Now, let us create the node for the regression
coefficients.
To do so, we need to define its initial value and its node distribution
using the {class}`.Dist` class. This class wraps distribution classes
with the TensorFlow Probability (TFP) API to connect them to our node
classes. Here, the node distribution is initialized with three
arguments: the TFP distribution object
[(`tfd.Normal`)](https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/Normal),
and the two hyperparameter nodes representing the parameters of the
distribution. TFP uses the names `loc` for the mean and `scale` for the
standard deviation, so we have to use the same names here. This is a
general feature of {class}`.Dist`, you can always use the parameter
names from TFP to refer to the parameters of your distribution.
``` python
beta_prior = lsl.Dist(tfd.Normal, loc=0.0, scale=100.0)
```
Note that you could also provide a {class}`.Node` or {class}`.Var`
instance for the `loc` and `scale` argument - this fact allows you to
build hierarchical models. If you provide floats like we do here, Liesel
will turn them into {class}`.Value` nodes under the hood.
With this distribution object, we can now create the node for our
regression coefficient with the {meth}`.Var.new_param` constructor:
``` python
beta = lsl.Var.new_param(value=np.array([0.0, 0.0]), distribution=beta_prior, name="beta")
```
#### The standard deviation
The second branch of the tree contains the residual standard deviation.
We build it in a similar way, now using the weakly informative prior
$\sigma^2 \sim \text{InverseGamme}(a, b)$ with $a = b = 0.01$ on the
squared standard deviation, i.e. the variance. Again, we use the
parameter names based on TFP. This time, we supply the hyperparameters
as {class}`.Var` instances.
``` python
a = lsl.Var.new_param(0.01, name="a")
b = lsl.Var.new_param(0.01, name="b")
sigma_sq_prior = lsl.Dist(tfd.InverseGamma, concentration=a, scale=b)
sigma_sq = lsl.Var.new_param(value=10.0, distribution=sigma_sq_prior, name="sigma_sq")
```
Since we need to work not only with the variance, but with the scale, we
initialize the scale using {meth}`.Var.new_calc`, to compute the square
root.
We can use this variable constructor to include computations based on
our nodes. It always takes a function as its first argument, and the
nodes to be used as function inputs as the following arguments. This is
the first weak node that we are setting up - all previous nodes have
been strong.
``` python
sigma = lsl.Var.new_calc(jnp.sqrt, sigma_sq, name="sigma").update()
```
#### Design matrix, fitted values, and response
To compute the matrix-vector product $\mathbf{X}\boldsymbol{\beta}$, we
use another variable instantiated via {meth}`.Var.new_calc`. We can view
our model as $y_i \sim \mathcal{N}(\mu_i, \;\sigma^2)$ with
$\mu_i = \beta_0 + \beta_1 x_i$, so we use the name `mu` for this
product.
``` python
X = lsl.Var.new_obs(X_mat, name="X")
mu = lsl.Var.new_calc(jnp.dot, X, beta, name="mu")
```
Finally, we can connect the branches of the tree in a response node. The
value of the node are our observed response values. And since we assumed
the model $y_i \sim \mathcal{N}(\beta_0 + \beta_1 x_i, \;\sigma^2)$, we
also need to specify the response’s distribution. We use our previously
created nodes `sigma` and `mu` to specify this distribution:
``` python
y_dist = lsl.Dist(tfd.Normal, loc=mu, scale=sigma)
y = lsl.Var.new_obs(y_vec, distribution=y_dist, name="y")
```
#### Bringing the model together
Now, to construct a full-fledged Liesel model from our individual node
objects, we can set up the {class}`.Model`. Here, we will only add the
response node.
``` python
model = lsl.Model([y])
model
```
Model(24 nodes, 8 vars)
Since all other nodes are directly or indirectly connected to this node,
the Model will add those nodes automatically when it builds the model.
The model provides a couple of convenience features, for example, to
evaluate the model log-probability, or to update the nodes in a
topological order.
The {func}`.plot_vars()` function visualizes the graph of a model.
Strong nodes are shown in blue, weak nodes in red. Nodes with a
probability distribution are highlighted with a star. In the figure
below, we can see the tree-like structure of the graph and identify the
two branches for the mean and the standard deviation of the response. If
the layout of the graph looks messy for you, please make sure you have
the `pygraphviz` package installed.
``` python
lsl.plot_vars(model)
```
### Node and model log-probabilities
The log-probability of the model, which can be interpreted as the
(unnormalized) log-posterior in a Bayesian context, can be accessed with
the `log_prob` property.
``` python
model.log_prob
```
Array(-1179.656, dtype=float32)
The individual nodes also have a `log_prob` property. In fact, because
of the conditional independence assumption of the model, the
log-probability of the model is given by the sum of the
log-probabilities of the nodes with probability distributions. We take
the sum for the `.log_prob` attributes of `beta` and `y` because, per
default, the attributes return the individual log-probability
contributions of each element in the values of the nodes. So for `beta`
we would get two log-probability values, and for `y` we would get 500.
``` python
beta.log_prob.sum() + sigma_sq.log_prob + y.log_prob.sum()
```
Array(-1179.656, dtype=float32)
Nodes without a probability distribution return a log-probability of
zero.
``` python
sigma.log_prob
```
0.0
The log-probability of a node depends on its value and its inputs. Thus,
if we change the variance of the response from 10 to 1, the
log-probability of the corresponding node, the log-probability of the
response node, and the log-probability of the model change as well. Note
that, since the actual input to the response distribution is the
standard deviation $\sigma$, we have to update its value after changing
the value of $\sigma^2$.
``` python
print(f"Old value of sigma_sq: {sigma_sq.value}")
```
Old value of sigma_sq: 10.0
``` python
print(f"Old log-prob of sigma_sq: {sigma_sq.log_prob}")
```
Old log-prob of sigma_sq: -6.972140312194824
``` python
print(f"Old log-prob of y: {y.log_prob.sum()}\n")
```
Old log-prob of y: -1161.6356201171875
``` python
sigma_sq.value = 1.0
sigma.update()
```
Var(name="sigma")
``` python
print(f"New value of sigma_sq: {sigma_sq.value}")
```
New value of sigma_sq: 1.0
``` python
print(f"New log-prob of sigma_sq: {sigma_sq.log_prob}")
```
New log-prob of sigma_sq: -4.655529975891113
``` python
print(f"New log-prob of y: {y.log_prob.sum()}\n")
```
New log-prob of y: -1724.6702880859375
``` python
print(f"New model log-prob: {model.log_prob}")
```
New model log-prob: -1740.3740234375
For most inference algorithms, we need the gradient of the model
log-probability with respect to the parameters. Liesel uses [the JAX
library for numerical computing and machine
learning](https://github.com/google/jax) to compute gradients using
automatic differentiation.
## MCMC inference with Goose
This section illustrates the key features of Liesel’s MCMC framework
Goose. To use Goose, the user needs to select one or more sampling
algorithms, called (transition) kernels, for the model parameters. Goose
comes with a number of standard kernels such as Hamiltonian Monte Carlo
({class}`~.goose.HMCKernel`) or the No U-Turn Sampler
({class}`~.goose.NUTSKernel`). Multiple kernels can be combined in one
sampling scheme and assigned to different parameters, and the user can
implement their own problem-specific kernels, as long as they are
compatible with the {class}`.Kernel` protocol. In any case, the user is
responsible for constructing a mathematically valid algorithm. Refer to
[MCMC Sampling (liesel.goose)](goose_overview) for an overview of
important Goose functionality.
We start with a very simple sampling scheme, keeping $\sigma^2$ fixed at
the true value and using a NUTS sampler for $\boldsymbol{\beta}$. The
kernels are added to a {class}`~.goose.Engine`, which coordinates the
sampling, including the kernel tuning during the warmup, and the MCMC
bookkeeping. The engine can be configured step by step with a
{class}`.EngineBuilder`. We need to inform the builder about the model,
the initial values, the kernels, and the sampling duration. Finally, we
can call the {meth}`.EngineBuilder.build` method, which returns a fully
configured engine.
``` python
sigma_sq.value = true_sigma**2 # setting sigma_sq to the true value
builder = gs.EngineBuilder(seed=1337, num_chains=4)
interface = gs.LieselInterface(model)
builder.set_model(interface)
builder.set_initial_values(model.state)
builder.add_kernel(gs.NUTSKernel(["beta"]))
builder.set_duration(warmup_duration=1000, posterior_duration=1000)
engine = builder.build()
```
Now we can run the MCMC algorithm for the specified duration by calling
the {meth}`~.goose.Engine.sample_all_epochs` method. In a first step,
the model and the sampling algorithm are compiled, so don’t worry if you
don’t see an output right away. The subsequent samples will be generated
much faster. Finally, we can extract the results and print a summary
table.
``` python
engine.sample_all_epochs()
```
0%| | 0/3 [00:00
Parameter summary:
|
|
|
kernel
|
mean
|
sd
|
q_0.05
|
q_0.5
|
q_0.95
|
sample_size
|
ess_bulk
|
ess_tail
|
rhat
|
|
parameter
|
index
|
|
|
|
|
|
|
|
|
|
|
|
beta
|
(0,)
|
kernel_00
|
0.980
|
0.089
|
0.833
|
0.981
|
1.127
|
4000
|
1023.503
|
1223.167
|
1.001
|
|
(1,)
|
kernel_00
|
1.915
|
0.154
|
1.660
|
1.914
|
2.170
|
4000
|
996.507
|
1195.113
|
1.001
|
Error summary:
|
|
|
|
|
count
|
relative
|
|
kernel
|
error_code
|
error_msg
|
phase
|
|
|
|
kernel_00
|
1
|
divergent transition
|
warmup
|
61
|
0.015
|
|
posterior
|
0
|
0.000
|
If we need more samples, we can append another epoch to the engine and
sample it by calling either the {meth}`~.goose.Engine.sample_next_epoch`
or the {meth}`~.goose.Engine.sample_all_epochs` method. The epochs are
described by {class}`.EpochConfig` objects.
``` python
engine.append_epoch(
gs.EpochConfig(gs.EpochType.POSTERIOR, duration=1000, thinning=1, optional=None)
)
engine.sample_next_epoch()
```
0%| | 0/40 [00:00
Parameter summary:
|
|
|
kernel
|
mean
|
sd
|
q_0.05
|
q_0.5
|
q_0.95
|
sample_size
|
ess_bulk
|
ess_tail
|
rhat
|
|
parameter
|
index
|
|
|
|
|
|
|
|
|
|
|
|
beta
|
(0,)
|
kernel_00
|
0.983
|
0.089
|
0.836
|
0.983
|
1.129
|
4000
|
1178.689
|
1425.977
|
1.003
|
|
(1,)
|
kernel_00
|
1.911
|
0.157
|
1.655
|
1.911
|
2.177
|
4000
|
1234.211
|
1393.635
|
1.004
|
|
sigma_sq
|
()
|
kernel_01
|
1.044
|
0.067
|
0.939
|
1.040
|
1.161
|
4000
|
3857.532
|
3510.714
|
1.001
|
Error summary:
|
|
|
|
|
count
|
relative
|
|
kernel
|
error_code
|
error_msg
|
phase
|
|
|
|
kernel_00
|
1
|
divergent transition
|
warmup
|
59
|
0.015
|
|
posterior
|
0
|
0.000
|
We can plot the trace plots of the chains with
{func}`~.goose.plot_trace()`.
``` python
g = gs.plot_trace(results)
```
We could also take a look at a kernel density estimator with
{func}`~.goose.plot_density()` and the estimated autocorrelation with
{func}`~.goose.plot_cor()`. Alternatively, we can output all three
diagnostic plots together with {func}`~.goose.plot_param()`. The
following plot shows the parameter $\beta_0$.
``` python
gs.plot_param(results, param="beta", param_index=0)
```
Here, we end this first tutorial. We have learned about a lot of
different classes and we have seen how we can flexibly use different
Kernels for drawing MCMC samples - that is quite a bit for the start.
Now, have fun modelling with Liesel!