|
| 1 | +# Introduction |
| 2 | + |
| 3 | +The purpose of this example is to illustrate Bayesian parameter estimation with a neural estimator called normalizing flows. Normalizing flows are a special type of invertible neural network which can learn the posterior distribution by learning the mapping between samples from the prior distribution and simulated data generated from the model. |
| 4 | + |
| 5 | +In the example below, we estimate the parameters of the [lognormal race model](lnr.md) (LNR; ). In pratice, one is unlikely to estimate the parameters of the LRN with neural networks because it has an analytic likelihood function. However, we use the LNR for illustration because it is fast to stimulate and the estimation properties of the LRN are known. You can reveal copy-and-pastable version of the full code by clicking the ▶ below. |
| 6 | + |
| 7 | +```@raw html |
| 8 | +<details> |
| 9 | +<summary><b>Show Full Code</b></summary> |
| 10 | +``` |
| 11 | +```julia |
| 12 | +using AlgebraOfGraphics |
| 13 | +using CairoMakie |
| 14 | +using Distributions |
| 15 | +using Flux |
| 16 | +using NeuralEstimators |
| 17 | +using Random |
| 18 | +using Plots |
| 19 | +using SequentialSamplingModels |
| 20 | + |
| 21 | +Random.seed!(544) |
| 22 | + |
| 23 | +n = 2 # dimension of each data replicate |
| 24 | +m = 100 # number of independent replicates |
| 25 | +d = 4 # dimension of the parameter vector θ |
| 26 | +w = 128 # width of each hidden layer |
| 27 | + |
| 28 | +function sample_prior(K) |
| 29 | + ν = rand(Normal(-2, 3), K, 2) |
| 30 | + σ = rand(truncated(Normal(1, 3), 0, Inf), K) |
| 31 | + τ = rand(Uniform(0.100, 0.300), K) |
| 32 | + θ = vcat(ν', σ', τ') |
| 33 | + return θ |
| 34 | +end |
| 35 | + |
| 36 | +to_array(x) = Float32[x.choice'; x.rt'] |
| 37 | +simulate(θ, m) = [to_array(rand(LNR(ϑ[1:2], ϑ[3], ϑ[4]), m)) for ϑ ∈ eachcol(θ)] |
| 38 | + |
| 39 | +# Approximate distribution |
| 40 | +approx_dist = NormalisingFlow(d, 2d) |
| 41 | + |
| 42 | +# Neural network mapping data to summary statistics (of the same dimension used in the approximate distribution) |
| 43 | +ψ = Chain(x -> log.(x), Dense(n, w, relu), Dense(w, w, relu)) |
| 44 | +ϕ = Chain(Dense(w, w, relu), Dense(w, 2d)) |
| 45 | +network = DeepSet(ψ, ϕ) |
| 46 | + |
| 47 | +# Initialise a neural posterior estimator |
| 48 | +estimator = PosteriorEstimator(approx_dist, network) |
| 49 | + |
| 50 | +# Train the estimator |
| 51 | +estimator = train( |
| 52 | + estimator, |
| 53 | + sample_prior, |
| 54 | + simulate; |
| 55 | + m, |
| 56 | + K = 25_000 |
| 57 | +) |
| 58 | + |
| 59 | +# Assess the estimator |
| 60 | +θ_test = sample_prior(1000) |
| 61 | +data_test = simulate(θ_test, m) |
| 62 | +assessment = assess(estimator, θ_test, data_test; parameter_names = ["ν₁", "ν₂", "σ", "τ"]) |
| 63 | +bias(assessment) |
| 64 | +rmse(assessment) |
| 65 | +recovery_plot = AlgebraOfGraphics.plot(assessment) |
| 66 | + |
| 67 | +# perform Bayesian parameter estimation on simulated data |
| 68 | +θ = [-1.5, 0, 0.75, 0.250] |
| 69 | +data = simulate(θ, m) |
| 70 | +post_samples = sampleposterior(estimator, data) |
| 71 | +Plots.histogram( |
| 72 | + post_samples', |
| 73 | + layout = (4, 1), |
| 74 | + color = :grey, |
| 75 | + norm = true, |
| 76 | + leg = false, |
| 77 | + grid = false, |
| 78 | + xlabel = ["ν₁" "ν₂" "σ" "τ"] |
| 79 | +) |
| 80 | +vline!([θ'], color = :darkred, linewidth = 2) |
| 81 | +``` |
| 82 | +```@raw html |
| 83 | +</details> |
| 84 | +``` |
| 85 | + |
| 86 | +# Load Dependencies |
| 87 | + |
| 88 | +The first step is to load the dependencies. `NeuralEstimators` and `Flux` are the primary packages for performing Bayesian parameter estimation with normalizing flows. We will also load `AlgebraOfGraphics`, `CairoMakie`, and `Plots` to visualize the parameter recovery and posterior distributions. |
| 89 | + |
| 90 | +```julia |
| 91 | +using AlgebraOfGraphics |
| 92 | +using CairoMakie |
| 93 | +using Distributions |
| 94 | +using Flux |
| 95 | +using NeuralEstimators |
| 96 | +using Random |
| 97 | +using Plots |
| 98 | +using SequentialSamplingModels |
| 99 | +``` |
| 100 | + |
| 101 | +In the code block below, we set the seed for the random number generator so that the results are reproducible. |
| 102 | +```julia |
| 103 | +Random.seed!(544) |
| 104 | +``` |
| 105 | + |
| 106 | +# Simulation Functions |
| 107 | + |
| 108 | +As previously noted, normalizing flow neural networks learn the mapping between the prior distribution and simulated data. Once that mapping is learned, the network is inverted to allow one to sample from posterior distribution. We define two functions to generate training data---one to sample from the prior distribution, and another to sample data from the model, given a sampled parameter vector from the prior distribution. In the code block below, the $K$ samples are generated from each prior and concatonated into a $4 \times K$ array. |
| 109 | + |
| 110 | +```julia |
| 111 | +function sample_prior(K) |
| 112 | + ν = rand(Normal(-2, 3), K, 2) |
| 113 | + σ = rand(truncated(Normal(1, 3), 0, Inf), K) |
| 114 | + τ = rand(Uniform(0.100, 0.300), K) |
| 115 | + θ = vcat(ν', σ', τ') |
| 116 | + return θ |
| 117 | +end |
| 118 | +``` |
| 119 | + |
| 120 | +The code block below specifies the function `simulate` to sample simulated data from the model. In this function, $\theta$ is a $4 \times K$ array, with each column representing an independent sample from the prior distribution, and $m$ is the number of trials sampled from the model per sample from the prior. The helper function `to_array` transforms the data into the required format: an $m \times 2$ array in which the first column consists of choice indices, and the second column consists of reaction times. |
| 121 | + |
| 122 | +```julia |
| 123 | +to_array(x) = Float32[x.choice'; x.rt'] |
| 124 | +simulate(θ, m) = [to_array(rand(LNR(ϑ[1:2], ϑ[3], ϑ[4]), m)) for ϑ ∈ eachcol(θ)] |
| 125 | +``` |
| 126 | +# Configure Neural Network |
| 127 | + |
| 128 | +In this section, we will configure the neural network to perform Bayesian parameter estimation. At a high level, the neural network has two main components. The first component is the `NormalisingFlow`, which learns the density of the model. The second component is the `DeepSet` which learns the summary statistics undelying the distributions. |
| 129 | +```julia |
| 130 | +# Approximate distribution |
| 131 | +approx_dist = NormalisingFlow(d, 2d) |
| 132 | + |
| 133 | +# Neural network mapping data to summary statistics (of the same dimension used in the approximate distribution) |
| 134 | +ψ = Chain(x -> log.(x), Dense(n, w, relu), Dense(w, w, relu)) |
| 135 | +ϕ = Chain(Dense(w, w, relu), Dense(w, 2d)) |
| 136 | +network = DeepSet(ψ, ϕ) |
| 137 | + |
| 138 | +# Initialise a neural posterior estimator |
| 139 | +estimator = PosteriorEstimator(approx_dist, network) |
| 140 | +``` |
| 141 | + |
| 142 | +# Train the Neural Network |
| 143 | + |
| 144 | +```julia |
| 145 | +estimator = train( |
| 146 | + estimator, |
| 147 | + sample_prior, |
| 148 | + simulate; |
| 149 | + m, |
| 150 | + K = 25_000 |
| 151 | +) |
| 152 | +``` |
| 153 | + |
| 154 | +# Assess the Accuracy of the Neural Network |
| 155 | + |
| 156 | +As shown in the code block below, the package NeuralEstimators provides three ways to assess the accuracy of the neural network: `bias`, `rmse`, and scatter plots of the parameer recovery. The parameter recovery plots below indicate that the neural network learned the mapping well. |
| 157 | + |
| 158 | +```julia |
| 159 | +θ_test = sample_prior(1000) |
| 160 | +data_test = simulate(θ_test, m) |
| 161 | +assessment = assess(estimator, θ_test, data_test; parameter_names = ["ν₁", "ν₂", "σ", "τ"]) |
| 162 | +bias(assessment) |
| 163 | +rmse(assessment) |
| 164 | +recovery_plot = AlgebraOfGraphics.plot(assessment) |
| 165 | +``` |
| 166 | + |
| 167 | + |
| 168 | +# Perform Bayesian Parameter Estimation |
| 169 | + |
| 170 | +Now that the neural network has been trained, we can perform Bayesian parameter estimation. In the example below, we simulate data from the model using parameters defined in the vector $\theta$. The estimator and data are passed to `sampleposterior`, which generates samples from the posterior distribution of parameters. As expected, the histogram shows that the posterior distributions are near the true parameter values, displayed as red vertical lines. |
| 171 | + |
| 172 | +```julia |
| 173 | +θ = [-1.5, 0, 0.75, 0.150] |
| 174 | +data = simulate(θ, m) |
| 175 | +post_samples = sampleposterior(estimator, data) |
| 176 | +Plots.histogram( |
| 177 | + post_samples', |
| 178 | + layout = (4, 1), |
| 179 | + color = :grey, |
| 180 | + norm = true, |
| 181 | + leg = false, |
| 182 | + grid = false, |
| 183 | + xlabel = ["ν₁" "ν₂" "σ" "τ"] |
| 184 | +) |
| 185 | +vline!([θ'], color = :darkred, linewidth = 2) |
| 186 | +``` |
| 187 | + |
| 188 | + |
| 189 | +# Save the Trained Neural Network |
| 190 | + |
| 191 | +```julia |
| 192 | +using BSON: @save |
| 193 | +using Flux |
| 194 | + |
| 195 | +model_state = Flux.state(estimator) |
| 196 | +@save "lnr_estimator.bson" model_state |
| 197 | +``` |
| 198 | + |
| 199 | +You can load the trained neural network into a new Julia session with the `@load` macro from `BSON`. In order to successfully reuse the trained neural network, you will need to initialize the neural network before passing the trained parameters. You can reveal copy-and-pastable version of the full code by clicking the ▶ below. |
| 200 | + |
| 201 | +```@raw html |
| 202 | +<details> |
| 203 | +<summary><b>Show Full Code</b></summary> |
| 204 | +``` |
| 205 | +```julia |
| 206 | +using BSON: @load |
| 207 | +using Distributions |
| 208 | +using Flux |
| 209 | +using NeuralEstimators |
| 210 | +using SequentialSamplingModels |
| 211 | + |
| 212 | +Random.seed!(544) |
| 213 | + |
| 214 | +n = 2 # dimension of each data replicate |
| 215 | +m = 100 # number of independent replicates |
| 216 | +d = 4 # dimension of the parameter vector θ |
| 217 | +w = 128 # width of each hidden layer |
| 218 | + |
| 219 | +# Approximate distribution |
| 220 | +approx_dist = NormalisingFlow(d, 2d) |
| 221 | + |
| 222 | +# Neural network mapping data to summary statistics (of the same dimension used in the approximate distribution) |
| 223 | +ψ = Chain(x -> log.(x), Dense(n, w, relu), Dense(w, w, relu)) |
| 224 | +ϕ = Chain(Dense(w, w, relu), Dense(w, 2d)) |
| 225 | +network = DeepSet(ψ, ϕ) |
| 226 | + |
| 227 | +# Initialise a neural posterior estimator |
| 228 | +estimator = PosteriorEstimator(approx_dist, network) |
| 229 | + |
| 230 | +# load the weights |
| 231 | +@load "lnr_estimator.bson" model_state |
| 232 | +Flux.loadmodel!(estimator, model_state) |
| 233 | +``` |
| 234 | +```@raw html |
| 235 | +</details> |
| 236 | +``` |
0 commit comments