Plots of prior and posterior distributions for a model I have an exercise that ask to plot the prior and posterior distribution in the Poisson-Gamma model. I did it (I think it's correct) inspired
in the answer of this question https://stats.stackexchange.com/questions/70661/how-does-the-beta-prior-affect-the-posterior-under-a-binomial-likelihood

Is there anything that can improve the code?

colors = c("red","blue","green","orange","purple")

n = 10

lambda = .2

x = rpois(n,lambda)
grid = seq(0,7,.01)


alpha = c(.5,5,1,2,2)
beta = c(.5,1,3,2,5)

plot(grid,grid,type="n",xlim=c(0,5),ylim=c(0,4),xlab="",ylab="Prior Density",
     main="Prior Distributions", las=1)
for(i in 1:length(alpha)){
  prior = dgamma(grid,shape=alpha[i],rate=1/beta[i]) 
  lines(grid,prior,col=colors[i],lwd=2)
}

legend("topleft", legend=c("Gamma(0.5,0.5)", "Gamma(5,1)", "Gamma(1,3)", "Gamma(2,2)", "Gamma(2,5)"),
       lwd=rep(2,5), col=colors, bty="n", ncol=3)

for(i in 1:length(alpha)){
  dev.new()
  plot(grid,grid,type="n",xlim=c(0,5),ylim=c(0,10),xlab="",ylab="Density",xaxs="i",yaxs="i",
       main="Prior and Posterior Distribution")

  alpha.star = alpha[i] + sum(x)
  beta.star = beta[i] + n
  prior = dgamma(grid,shape=alpha[i],rate=1/beta[i])
  post = dgamma(grid,shape=alpha.star,rate=beta.star)

  lines(grid,post,lwd=2)
  lines(grid,prior,col=colors[i],lwd=2)
  legend("topright",c("Prior","Posterior"),col=c(colors[i],"black"),lwd=2)

}
There was a couple of errors in the code as originally posted. The code was OK as far as base R goes.

First thing I did was run styler and then lintr on your code; these two things help clean up the coding style in your scripts.

That does things like this:

colors = c("red","blue","green","orange","purple")

# changed to (spaces / idiomatic assignment):
colors &lt;- c("red", "blue", "green", "orange", "purple")



Then I changed your 1:length(alphas) to seq_along(alphas). The latter is a bit safer since the former can fail with empty input.

Then I replaced your 5-separate prior/posterior plots with a single plot that contains 5 panels. This makes it easier to compare the appropriateness of the different priors. To do this, I removed your dev.new()s, added a call to par(mfrow = c(number_of_rows, number_of_columns)) and obviously, tidied this up afterwards (returning to a 1x1 grid)

par(mfrow = c(2, ceiling(length(alpha) / 2)))

for (i in seq_along(alpha)) {
    # removed dev.new()
   ... plotting code ...
   )
 }

par(mfrow = c(1, 1))


Then I cleaned up your experimental data and your prior-parameters / plotting parameters; they were all winding around each other. I also renamed your alpha / beta vectors - in R, these correspond to the shape and rate parameters that are passed into dgamma:

# ---- experimental data

num_observations &lt;- 10

lambda &lt;- .2

x &lt;- rpois(num_observations, lambda)

# ---- prior parameters

# assumed 'beta' was a rate parameter
# - this, since there was confusion in the parameterisation of dgamma():
#   - early section used rate = 1 / beta[i];
#   - later section used rate = beta[i]; and
#   - definition of beta_star = beta[i] + n; implied beta was definitely a rate

shape &lt;- c(.5, 5, 1, 2, 2)
rate &lt;- c(.5, 1, 3, 2, 5)

# ---- plotting parameters

colors &lt;- c("red", "blue", "green", "orange", "purple")

# ---- search parameters

grid &lt;- seq(0, 2, .01)



Then I made a function to do your prior-comparison stuff (the first set of plots). Any parameters that were to be passed through to plot were passed in using the ... argument.

# ---- comparison of the prior distributions

plot_priors &lt;- function(grid, shapes, rates, colors,
                        legend_text, lwd = 2, ...) {
  plot(grid, grid, type = "n", ...)

  for (i in seq_along(shape)) {
    prior &lt;- dgamma(grid, shape = shape[i], rate = rate[i])
    lines(grid, prior, col = colors[i], lwd = lwd)
  }

  legend(
    "topleft",
    legend = legend_text, lwd = lwd, col = colors, bty = "n", ncol = 2
  )
}



This can be called like:

plot_priors(
  grid, shape, rate, colors,
  legend_text = paste0("Gamma(", c("0.5,0.5", "5,1", "1,3", "2,2", "2,5"), ")"),
  xlim = c(0, 1), ylim = c(0, 4), xlab = "", ylab = "Prior Density",
  main = "Prior Distributions", las = 1
)


It's useful to split your computations away from your plotting code - so I extracted the code you used to compute the posterior params:

compute_posterior_parameters &lt;- function(observations,
                                         prior_shape,
                                         prior_rate) {
  list(
    shape = prior_shape + sum(observations),
    rate = prior_rate + length(observations)
  )
}


Then I pulled the plotting code for your prior/posterior comparisons into a function (similarly to the above)

plot_prior_post_comparison &lt;- function(
                                       observations,
                                       grid, shapes, rates, colors,
                                       lwd = 2,
                                       ...) {
  # make a grid for plotting
  par(mfrow = c(2, ceiling(length(shapes) / 2)))

  for (i in seq_along(shapes)) {
    # details of the prior and post distributions
    posterior_params &lt;- compute_posterior_parameters(
      observations,
      prior_shape = shapes[i], prior_rate = rates[i]
    )
    prior &lt;- dgamma(
      grid,
      shape = shapes[i],
      rate = rates[i]
    )
    post &lt;- dgamma(
      grid,
      shape = posterior_params$shape,
      rate = posterior_params$rate
    )

    # plotting code
    plot(grid, grid, type = "n", ...)
    lines(grid, post, lwd = lwd)
    lines(grid, prior, col = colors[i], lwd = lwd)
    legend("topright",
      c("Prior", "Posterior"),
      col = c(colors[i], "black"), lwd = lwd
    )
  }

  # revert the plotting grid back to 1x1
  par(mfrow = c(1, 1))
}



note that the par calls, which change the plotting grid are all nested inside the function, so any subsequent plots should be unaffected.

Then I called that function:

# ---- prior/posterior comparison

plot_prior_post_comparison(
  observations = x,
  grid = grid, shapes = shape, rates = rate, colors = colors,
  xlim = c(0, 1), ylim = c(0, 10), xlab = "", ylab = "Density",
  xaxs = "i", yaxs = "i",
  main = "Prior and Posterior Distribution"
)





Then I put all the functions at the start and all the calls at the end of a script:

The full code:

# ---- comparison of the prior distributions

plot_priors &lt;- function(grid, shapes, rates, colors,
                        legend_text, lwd = 2, ...) {
  plot(grid, grid, type = "n", ...)

  for (i in seq_along(shape)) {
    prior &lt;- dgamma(grid, shape = shape[i], rate = rate[i])
    lines(grid, prior, col = colors[i], lwd = lwd)
  }

  legend(
    "topleft",
    legend = legend_text, lwd = lwd, col = colors, bty = "n", ncol = 2
  )
}

# ---- prior:posterior analysis

compute_posterior_parameters &lt;- function(observations,
                                         prior_shape,
                                         prior_rate) {
  list(
    shape = prior_shape + sum(observations),
    rate = prior_rate + length(observations)
  )
}

plot_prior_post_comparison &lt;- function(
                                       observations,
                                       grid, shapes, rates, colors,
                                       lwd = 2,
                                       ...) {
  # make a grid for plotting
  par(mfrow = c(2, ceiling(length(shapes) / 2)))

  for (i in seq_along(shapes)) {
    # details of the prior and post distributions
    posterior_params &lt;- compute_posterior_parameters(
      observations,
      prior_shape = shapes[i], prior_rate = rates[i]
    )
    prior &lt;- dgamma(
      grid,
      shape = shapes[i],
      rate = rates[i]
    )
    post &lt;- dgamma(
      grid,
      shape = posterior_params$shape,
      rate = posterior_params$rate
    )

    # plotting code
    plot(grid, grid, type = "n", ...)
    lines(grid, post, lwd = lwd)
    lines(grid, prior, col = colors[i], lwd = lwd)
    legend("topright",
      c("Prior", "Posterior"),
      col = c(colors[i], "black"), lwd = lwd
    )
  }

  # revert the plotting grid back to 1x1
  par(mfrow = c(1, 1))
}


# ----

# ---- experimental data

num_observations &lt;- 10

lambda &lt;- .2

x &lt;- rpois(num_observations, lambda)

# ---- prior parameters

# assumed 'beta' was a rate parameter
# - this, since there was confusion in the parameterisation of dgamma():
#   - early section used rate = 1 / beta[i];
#   - later section used rate = beta[i]; and
#   - definition of beta_star = beta[i] + n; implied beta was definitely a rate

shape &lt;- c(.5, 5, 1, 2, 2)
rate &lt;- c(.5, 1, 3, 2, 5)

# ---- plotting parameters

colors &lt;- c("red", "blue", "green", "orange", "purple")

# ---- search parameters

grid &lt;- seq(0, 2, .01)

# ---- comparison of priors

plot_priors(
  grid, shape, rate, colors,
  legend_text = paste0("Gamma(", c("0.5,0.5", "5,1", "1,3", "2,2", "2,5"), ")"),
  xlim = c(0, 1), ylim = c(0, 4), xlab = "", ylab = "Prior Density",
  main = "Prior Distributions", las = 1
)

# ---- prior/posterior comparison

plot_prior_post_comparison(
  observations = x,
  grid = grid, shapes = shape, rates = rate, colors = colors,
  xlim = c(0, 1), ylim = c(0, 10), xlab = "", ylab = "Density",
  xaxs = "i", yaxs = "i",
  main = "Prior and Posterior Distribution"
)



I don't thnk its perfect - the plotting and calculation steps are still pretty tied together - but it should be more easy to add extra pairs of rate/shape values into your prior distributions, for example.

One place the code coould be further improved is by passing in a data.frame where each row contains the shape/rate/colour values and any annotations for a given prior. The functions contain too many arguments at the moment, and this would fix that (and guarantee that there is a shape for each rate).