Skip to content

Commit dd294c8

Browse files
committed
Expand on docs
1 parent 9d17281 commit dd294c8

File tree

1 file changed

+70
-16
lines changed

1 file changed

+70
-16
lines changed

usage/tracking-extra-quantities/index.qmd

Lines changed: 70 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,14 @@ using Pkg;
1313
Pkg.instantiate();
1414
```
1515

16-
Often, the most natural parameterization for a model is not the most computationally feasible.
16+
Often, there are quantities in models that we might be interested in viewing the values of, but which are not random variables in the model that are explicitly drawn from a distribution.
17+
18+
As a motivating example, the most natural parameterization for a model might not the most computationally feasible.
1719
Consider the following (efficiently reparametrized) implementation of Neal's funnel [(Neal, 2003)](https://arxiv.org/abs/physics/0009028):
1820

1921
```{julia}
2022
using Turing
23+
setprogress!(false)
2124
2225
@model function Neal()
2326
# Raw draws
@@ -33,9 +36,7 @@ end
3336

3437
In this case, the random variables exposed in the chain (`x_raw`, `y_raw`) are not in a helpful form — what we're after are the deterministically transformed variables `x` and `y`.
3538

36-
More generally, there are often quantities in our models that we might be interested in viewing, but which are not explicitly present in our chain.
37-
38-
There are two ways of tracking such extra quantities.
39+
There are two ways to track these extra quantities in Turing.jl.
3940

4041
## Using `:=` (during inference)
4142

@@ -53,7 +54,7 @@ For example:
5354
x := exp.(y ./ 2) .* x_raw
5455
end
5556
56-
sample(Neal_coloneq(), NUTS(), 1000; progress=false)
57+
sample(Neal_coloneq(), NUTS(), 1000)
5758
```
5859

5960
## Using `returned` (post-inference)
@@ -69,30 +70,83 @@ Alternatively, one can specify the extra quantities as part of the model functio
6970
# Transform and return as a NamedTuple
7071
y = 3 * y_raw
7172
x = exp.(y ./ 2) .* x_raw
72-
return [x; y]
73+
return (x=x, y=y)
7374
end
7475
75-
chain = sample(Neal_return(), NUTS(), 1000; progress=false)
76+
chain = sample(Neal_return(), NUTS(), 1000)
7677
```
7778

78-
This chain does not contain `x` and `y`, but we can extract the values using the `returned` function.
79-
Calling this function outputs an array of values specified in the return statement of the model.
79+
The sampled chain does not contain `x` and `y`, but we can extract the values using the `returned` function.
80+
Calling this function outputs an array:
8081

8182
```{julia}
82-
returned(Neal_return(), chain)
83+
nts = returned(Neal_return(), chain)
8384
```
8485

85-
Each element of this corresponds to an array with the values of `x1, x2, ..., x9, y` for each posterior sample.
86+
where each element of which is a NamedTuple, as specified in the return statement of the model.
87+
88+
```{julia}
89+
nts[1]
90+
```
91+
92+
## Which to use?
93+
94+
There are some pros and cons of using `returned`, as opposed to `:=`.
95+
96+
Firstly, `returned` is more flexible, as it allows you to track any type of object; `:=` only works with variables that can be inserted into an `MCMCChains.Chains` object.
97+
(Notice that `x` is a vector, and in the first case where we used `:=`, reconstructing the vector value of `x` can also be rather annoying as the chain stores each individual element of `x` separately.)
98+
99+
However, if used carelessly, `returned` can lead to unnecessary computation.
100+
For example, in `Neal_return()` above, the `x` and `y` variables are also calculated during the inference process (i.e. the call to `sample()`), but are then thrown away.
101+
They are then calculated _again_ when `returned()` is called.
102+
103+
To avoid this, you will essentially have to create two different models, one for inference and one for post-inference.
104+
The simplest way of doing this is to add a parameter to the model argument:
105+
106+
```{julia}
107+
@model function Neal_coloneq_optional(track::Bool)
108+
# Raw draws
109+
y_raw ~ Normal(0, 1)
110+
x_raw ~ arraydist([Normal(0, 1) for i in 1:9])
111+
112+
if track
113+
y = 3 * y_raw
114+
x = exp.(y ./ 2) .* x_raw
115+
return (x=x, y=y)
116+
else
117+
return nothing
118+
end
119+
end
86120
87-
In this case, it might be useful to reorganize our output into a matrix for plotting:
121+
chain = sample(Neal_coloneq_optional(false), NUTS(), 1000)
122+
```
123+
124+
The above ensures that `x` and `y` are not calculated during inference, but allows us to still use `returned` to extract them:
88125

89126
```{julia}
90-
reparam_chain = reduce(hcat, returned(Neal_return(), chain))'
127+
returned(Neal_coloneq_optional(true), chain)
91128
```
92129

93-
from which we can recover a vector of our samples:
130+
Another equivalent option is to use a submodel:
94131

95132
```{julia}
96-
x1_samples = reparam_chain[:, 1]
97-
y_samples = reparam_chain[:, 10]
133+
@model function Neal()
134+
y_raw ~ Normal(0, 1)
135+
x_raw ~ arraydist([Normal(0, 1) for i in 1:9])
136+
return (x_raw=x_raw, y_raw=y_raw)
137+
end
138+
139+
chain = sample(Neal(), NUTS(), 1000)
140+
141+
@model function Neal_with_extras()
142+
neal ~ to_submodel(Neal(), false)
143+
y = 3 * neal.y_raw
144+
x = exp.(y ./ 2) .* neal.x_raw
145+
return (x=x, y=y)
146+
end
147+
148+
returned(Neal_with_extras(), chain)
98149
```
150+
151+
Note that for the `returned` call to work, the `Neal_with_extras()` model must have the same variable names as stored in `chain`.
152+
This means the submodel `Neal()` must not be prefixed, i.e. `to_submodel()` must be passed a second parameter `false`.

0 commit comments

Comments
 (0)