|
11 | 11 | _has_varname_to_symbol(info::NamedTuple{names}) where {names} = :varname_to_symbol in names |
12 | 12 | function _check_varname_indexing(c::MCMCChains.Chains) |
13 | 13 | return DynamicPPL.supports_varname_indexing(c) || |
14 | | - error("Chains do not support indexing using $vn.") |
| 14 | + error("$(typeof(c)) do not support indexing using varnmes.") |
15 | 15 | end |
16 | 16 |
|
17 | 17 | # Load state from a `Chains`: By convention, it is stored in `:samplerstate` metadata |
@@ -41,21 +41,152 @@ function DynamicPPL.varnames(c::MCMCChains.Chains) |
41 | 41 | return keys(c.info.varname_to_symbol) |
42 | 42 | end |
43 | 43 |
|
| 44 | +""" |
| 45 | + generated_quantities(model::Model, chain::MCMCChains.Chains) |
| 46 | +
|
| 47 | +Execute `model` for each of the samples in `chain` and return an array of the values |
| 48 | +returned by the `model` for each sample. |
| 49 | +
|
| 50 | +# Examples |
| 51 | +## General |
| 52 | +Often you might have additional quantities computed inside the model that you want to |
| 53 | +inspect, e.g. |
| 54 | +```julia |
| 55 | +@model function demo(x) |
| 56 | + # sample and observe |
| 57 | + θ ~ Prior() |
| 58 | + x ~ Likelihood() |
| 59 | + return interesting_quantity(θ, x) |
| 60 | +end |
| 61 | +m = demo(data) |
| 62 | +chain = sample(m, alg, n) |
| 63 | +# To inspect the `interesting_quantity(θ, x)` where `θ` is replaced by samples |
| 64 | +# from the posterior/`chain`: |
| 65 | +generated_quantities(m, chain) # <= results in a `Vector` of returned values |
| 66 | + # from `interesting_quantity(θ, x)` |
| 67 | +``` |
| 68 | +## Concrete (and simple) |
| 69 | +```julia |
| 70 | +julia> using DynamicPPL, Turing |
| 71 | +
|
| 72 | +julia> @model function demo(xs) |
| 73 | + s ~ InverseGamma(2, 3) |
| 74 | + m_shifted ~ Normal(10, √s) |
| 75 | + m = m_shifted - 10 |
| 76 | +
|
| 77 | + for i in eachindex(xs) |
| 78 | + xs[i] ~ Normal(m, √s) |
| 79 | + end |
| 80 | +
|
| 81 | + return (m, ) |
| 82 | + end |
| 83 | +demo (generic function with 1 method) |
| 84 | +
|
| 85 | +julia> model = demo(randn(10)); |
| 86 | +
|
| 87 | +julia> chain = sample(model, MH(), 10); |
| 88 | +
|
| 89 | +julia> generated_quantities(model, chain) |
| 90 | +10×1 Array{Tuple{Float64},2}: |
| 91 | + (2.1964758025119338,) |
| 92 | + (2.1964758025119338,) |
| 93 | + (0.09270081916291417,) |
| 94 | + (0.09270081916291417,) |
| 95 | + (0.09270081916291417,) |
| 96 | + (0.09270081916291417,) |
| 97 | + (0.09270081916291417,) |
| 98 | + (0.043088571494005024,) |
| 99 | + (-0.16489786710222099,) |
| 100 | + (-0.16489786710222099,) |
| 101 | +``` |
| 102 | +""" |
44 | 103 | function DynamicPPL.generated_quantities( |
45 | 104 | model::DynamicPPL.Model, chain_full::MCMCChains.Chains |
46 | 105 | ) |
47 | 106 | chain = MCMCChains.get_sections(chain_full, :parameters) |
48 | 107 | varinfo = DynamicPPL.VarInfo(model) |
49 | 108 | iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3)) |
50 | 109 | return map(iters) do (sample_idx, chain_idx) |
51 | | - # Update the varinfo with the current sample and make variables not present in `chain` |
52 | | - # to be sampled. |
53 | | - DynamicPPL.setval_and_resample!(varinfo, chain, sample_idx, chain_idx) |
| 110 | + if DynamicPPL.supports_varname_indexing(chain) |
| 111 | + varname_pairs = _varname_pairs_with_varname_indexing( |
| 112 | + chain, varinfo, sample_idx, chain_idx |
| 113 | + ) |
| 114 | + else |
| 115 | + varname_pairs = _varname_pairs_without_varname_indexing( |
| 116 | + chain, varinfo, sample_idx, chain_idx |
| 117 | + ) |
| 118 | + end |
| 119 | + fixed_model = DynamicPPL.fix(model, Dict(varname_pairs)) |
| 120 | + return fixed_model() |
| 121 | + end |
| 122 | +end |
| 123 | + |
| 124 | +""" |
| 125 | + _varname_pairs_with_varname_indexing( |
| 126 | + chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx |
| 127 | + ) |
54 | 128 |
|
55 | | - # TODO: Some of the variables can be a view into the `varinfo`, so we need to |
56 | | - # `deepcopy` the `varinfo` before passing it to `model`. |
57 | | - model(deepcopy(varinfo)) |
| 129 | +Get pairs of `VarName => value` for all the variables in the `varinfo`, picking the values |
| 130 | +from the chain. |
| 131 | +
|
| 132 | +This implementation assumes `chain` can be indexed using variable names, and is the |
| 133 | +preffered implementation. |
| 134 | +""" |
| 135 | +function _varname_pairs_with_varname_indexing( |
| 136 | + chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx |
| 137 | +) |
| 138 | + vns = DynamicPPL.varnames(chain) |
| 139 | + vn_parents = Iterators.map(vns) do vn |
| 140 | + # The call nested_setindex_maybe! is used to handle cases where vn is not |
| 141 | + # the variable name used in the model, but rather subsumed by one. Except |
| 142 | + # for the subsumption part, this could be |
| 143 | + # vn => getindex_varname(chain, sample_idx, vn, chain_idx) |
| 144 | + # TODO(mhauru) This call to nested_setindex_maybe! is unintuitive. |
| 145 | + DynamicPPL.nested_setindex_maybe!( |
| 146 | + varinfo, DynamicPPL.getindex_varname(chain, sample_idx, vn, chain_idx), vn |
| 147 | + ) |
58 | 148 | end |
| 149 | + varname_pairs = Iterators.map(Iterators.filter(!isnothing, vn_parents)) do vn_parent |
| 150 | + vn_parent => varinfo[vn_parent] |
| 151 | + end |
| 152 | + return varname_pairs |
| 153 | +end |
| 154 | + |
| 155 | +""" |
| 156 | +Check which keys in `key_strings` are subsumed by `vn_string` and return the their values. |
| 157 | +
|
| 158 | +The subsumption check is done with `DynamicPPL.subsumes_string`, which is quite weak, and |
| 159 | +won't catch all cases. We should get rid of this if we can. |
| 160 | +""" |
| 161 | +# TODO(mhauru) See docstring above. |
| 162 | +function _vcat_subsumed_values(vn_string, values, key_strings) |
| 163 | + indices = findall(Base.Fix1(DynamicPPL.subsumes_string, vn_string), key_strings) |
| 164 | + return !isempty(indices) ? reduce(vcat, values[indices]) : nothing |
| 165 | +end |
| 166 | + |
| 167 | +""" |
| 168 | + _varname_pairs_without_varname_indexing( |
| 169 | + chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx |
| 170 | + ) |
| 171 | +
|
| 172 | +Get pairs of `VarName => value` for all the variables in the `varinfo`, picking the values |
| 173 | +from the chain. |
| 174 | +
|
| 175 | +This implementation does not assume that `chain` can be indexed using variable names. It is |
| 176 | +thus not guaranteed to work in cases where the variable names have complex subsumption |
| 177 | +patterns, such as if the model has a variable `x` but the chain stores `x.a[1]`. |
| 178 | +""" |
| 179 | +function _varname_pairs_without_varname_indexing( |
| 180 | + chain::MCMCChains.Chains, varinfo, sample_idx, chain_idx |
| 181 | +) |
| 182 | + values = chain.value[sample_idx, :, chain_idx] |
| 183 | + keys = Base.keys(chain) |
| 184 | + keys_strings = map(string, keys) |
| 185 | + varname_pairs = [ |
| 186 | + vn => _vcat_subsumed_values(string(vn), values, keys_strings) for |
| 187 | + vn in Base.keys(varinfo) |
| 188 | + ] |
| 189 | + return varname_pairs |
59 | 190 | end |
60 | 191 |
|
61 | 192 | end |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?v=4&size=40){kind=avatar}
0 commit comments