# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Simulating Language, Lab 7, Hierarchical Bayesian Learning

# This week we are working with a very simple hierarchical learning model, applying Bayes rule to learn a high-order proprty of the data. 
#
# This model is built on the frequency learning model that we used in labs 3 and 4. Hopefully you remember from those labs that learners were exposed to data consisting of a sequence of productions of two variants, 0 and 1, and their task was to infer the probability of variant 1, which we sometimes refer to as $\theta$ and sometimes as `pW1` ("probability of word 1"). To do Bayesian inference we need a likelihood function and a prior. The likelihood from our frequency learning model was quite simple and gives you the probability of a particular sequence of 1s and 0s being generated given a particular value of `pW1` - when `pW1` is low then sequences with lots of 0s are more probable than sequences with lots of 1s, when `pW1` is high the reverse is true, and when `pW1` is around 0.5 you expect to see an even mix of 0s and 1s. For our prior we used a *beta distribution*, which is a family of probability distributions that can capture a uniform prior (representing an unbiased learner), a prior favouring regularity, or a prior favouring variability. We specify a single value for our beta distribution which by convention is called $\alpha$ (or `alpha`), which determines the shape of the distibution (NB `alpha` is not a probability, it's just a number we can manipulate to get different shaped priors); when `alpha` is less than 1 you have a prior favouring regularity (very low or very high values of `pW1` have the highest prior probability, and the prior assigns low prior probability to values of `pW1` around 0.5); when `alpha` is exactly 1 you have a flat prior (all values of `pW1` have equal prior probability); when `alpha` is greater than 1 you have a prior that favours variability (values of `pW1` around 0.5 have the highest prior probability, and extreme values of `pW1` have lower prior probability). 
#
# In that frequency learning model, the `alpha` parameter therefore determines the learners' expectations about variability: we can build learners who expect low variability (by setting `alpha` to less than 1), or high variability (by setting `alpha` to greater than 1).  But the `alpha` parameter was set by *us*, the modeller. This week we are going to have learners learn the `alpha` value (rather than learning `pW1`); so they'll learn the higher-level property of whether they should expect variable or regular data.

# ## Old code
#
# The first few chunks of code are borrowed from lab 4 - we'll need the usual stuff for dealing with log probabilities, and we'll also need some components of our lab 4 frequency-learning model. I am going to re-name some of those functions since we will now be dealing with likelihoods and priors for `pW1` *and* for `alpha`, so I want to name the functions so that it's clear whether they are e.g. the prior for `pW1` or for `alpha`.
#
# ### Libraries etc
#
# First we'll load our library functions, and define the functions we need for doing some arithmetic with log probabilities. As usual, *it is not important that you understand exactly how these work* - they're just doing  book-keeping stuff to allow us to use log probabilities rather than probabilities.

# +
import random
from scipy.stats import beta
from scipy.special import logsumexp
from math import log, log1p, exp

# %matplotlib inline
import matplotlib.pyplot as plt
from IPython.display import set_matplotlib_formats
set_matplotlib_formats('svg', 'pdf')


# +
def log_subtract(x,y):
    return x + log1p(-exp(y - x))

def normalize_logprobs(logprobs):
    logtotal = logsumexp(logprobs) #calculates the summed log probabilities
    normedlogs = []
    for logp in logprobs:
        normedlogs.append(logp - logtotal) #normalise - subtracting in the log domain
                                           #is equivalent to dividing in the normal domain
    return normedlogs

def log_roulette_wheel(normedlogs):
    r=log(random.random()) #generate a random number between 0 and 1, then convert to log
    accumulator = normedlogs[0] 
    for i in range(len(normedlogs)):
        if r < accumulator:
            return i
        accumulator = logsumexp([accumulator, normedlogs[i + 1]])


# -

# ### The hypothesis space, prior and likelihoods for pW1 

# As in labs 3 and 4, the main part of the code starts by laying out our hypothesis space for `pW1`, our grid of possible values of `pW1` (which is the estimated probability of using variant 1). 

# +
pW1_grid_granularity = 100
pW1_grid_increment = 1 / pW1_grid_granularity

# sets up the grid of possible probabilities to consider
possible_pW1 = []
for i in range(pW1_grid_granularity):
    possible_pW1.append(pW1_grid_increment / 2 + (pW1_grid_increment * i))

# sets up the grid of the corresponding log probabilities
possible_logpW1 = []
for pW1 in possible_pW1:
    possible_logpW1.append(log(pW1))


# -

# Next we need our prior over possible values of `pW1` - as per labs 3 and 4, this prior is a beta distribution, where the shape of that distribution is determined by the parameter `alpha`. In labs 3 and 4 this function was called `calculate_logprior(alpha)`, but here I am going to rename it `pW1_calculate_logprior(alpha)` to make it clear that this is the prior for values of `pW1`.

def pW1_calculate_logprior(alpha):
    logprior_pW1 = []
    for pW1 in possible_pW1:
        logprior_pW1.append(beta.logpdf(pW1, alpha, alpha)) 
    return normalize_logprobs(logprior_pW1)


# We will also need a likelihood function, which tells us the likelihood of some data (a sequence of variants 0 and 1, e.g. `[0,1,1,0,1]`) given a particular (log) probability of `pW1`. This is the same as the `likelihood(data,logpW1)` function from lab 4, but again I have renamed it so it's clear that this is the likelihood for `pW1`.

def pW1_loglikelihood(data, logpW1):
    logpW0 = log_subtract(log(1), logpW1) #probability of w0 is 1-prob of w1
    logprobs = [logpW0, logpW1]
    loglikelihoods = []
    for d in data:
        loglikelihood_this_item = logprobs[d] #d will be either 0 or 1, 
                                              #so can use as index
        loglikelihoods.append(loglikelihood_this_item)
    return sum(loglikelihoods) #summing log probabilities = 
                               #multiply non-log probabilities


# Finally we are going to want to get our model to produce some data (a sequence of 0s and 1s) given a particular value of (log) `pW1` - this was called `logproduce(logpW1, n_productions)` in lab 4, again I am renaming to make clear this is to do with `pW1`.

# +
def pW1_logproduce(logpW1, n_productions):
    logpW0 = log_subtract(log(1), logpW1)
    logprobs = [logpW0, logpW1]
    data = []
    for p in range(n_productions):
        data.append(log_roulette_wheel(logprobs))
    return data


# -

# It might be worth checking at this point that you are happy with the behaviour of this part of the model - you could inspect the shape of the prior produced by `pW1_calculate_logprior(alpha)` for various values of `alpha`, check that `pW1_loglikelihood(data, logpW1)` generates sensible values, and that `pW1_logproduce(logpW1, n_productions)` behaves as expected. For instance, if you want to see that `pW1_logproduce` produces regular data when `pW1` is low or high, but variable data when `pW1` is around 0.5, you can do something like:
#
# ```python
# print(pW1_logproduce(log(0.05), 20)) #try with pW1 = 0.1 (low)
# print(pW1_logproduce(log(0.95), 20)) #try with pW1 = 0.9 (high)
# print(pW1_logproduce(log(0.5), 20)) #try with pW1 = 0.5 (in the middle)
# ```



# ## New code for higher-level inference

# OK, now we are set up with our lower-level model which gives us priors and likelihoods for our low-level property of the data, `pW1`, we are ready to set up our model which is going to learn the higher-level property of the data, `alpha`. 
#
# Our lower-level model for learning `pW1` assumes that your data comes from a single source - i.e. in lab 4 our learners were inferring `pW1` based on a sequence of data like this:
#
# ```python
# some_data = [0,1,1,0,1,0]
# ```
#
# In this model we are going to allow the possibility that the data comes from *multiple* sources, who might differ in their `pW1`, and our learner's task is to learn how regular or variable those sources tend to be. For example, imagine that you are learning how frequently different speakers use two linguistic variants, variant 0 and variant 1. We could represent the data produced by two individual (highly individually-consistent) speakers like this:
#
# ```python
# speaker0_data = [0,0,0,0,0,0]
# speaker1_data = [1,1,1,1,1,1]
# ```
#
# Then their combined data is represented as a list of the data produced by each speaker, like this:
#
# ```python
# combined_data = [speaker0_data,speaker1_data]
# ```
# which is the same as 
# ```python
# combined_data = [[0,0,0,0,0,0],[1,1,1,1,1,1]]
# ```
# i.e. our data from multiple sources is a *list of lists*, where each list is the data from a single source. Of course we could still represent data from a single source in this representation, e.g. 
#
# ```python
# single_source_data = [speaker0_data]
# ```
# which is the same as 
# ```python
# single_source_data = [[0,0,0,0,0,0]]
# ```
# i.e. a list of lists which happens to only contain a single sub-list.

# - The two speakers in the `combined_data` example above are highly regular. If exposed to this kind of data, intuitively what would you infer about how variable or regular speakers tend to be? How could you capture that expectation about per-speaker variability or regularity in terms of an `alpha` value? Remember, `alpha` is just a way of expressing an expectation about variability.
# - How would you represent the combined data from 3 highly variable speakers, all of whom use both variants roughly equally frequently? And what would that correspond to in terms of an `alpha` value?

# ### Hypothesis space for alpha

# OK, so now we have a way of representing data from multiple sources, and hopefully you've had a think about what kinds of inferences you might make about the regularity or variability of sources in general based on that kind of data. We can now build our model and see if it matches our intuitions. 
#
# As usual we need to lay out our hypothesis space - this time we are trying to learn `alpha`, so we need to lay out a set of possible `alpha` values to consider. Remember, `alpha`<1 indicates an expectation of regularity, favouring extreme (very low or very high) values of `pW1`; `alpha`>1 represents an expectation of variability, where values of `pW1` around 0.5 are favoured. We therefore want to allow a nice range of possible values of `alpha`, including both very low and very high values. Here's a list of 9 possible `alpha` values that covers a good range.

possible_alphas = [0.001,0.005,0.01,0.05,0.1,0.5,1,5,10]


# ### Prior for alpha

# Next we need a prior, which captures our learner's prior expectations about the probability of the various hypotheses it is considering. In this case, it's a prior over possible values of `alpha`, so it's an expectation about expectations of regularity/variability, which is a bit of a head-twister - but remember, it's just a prior as usual, so a set of (log) probabilities, one per candidate hypothesis. 
#
# In this case we'll keep things simple and assume a uniform prior: so all possible values of `alpha` from `possible_alphas` have the same prior probability. The function `alpha_calculate_logprior()` will just lay these out for you - note that the name includes a reminder that this is the prior for `alpha` (and not the prior for `pW1`, which is what `pW1_calculate_logprior` does). 

def alpha_calculate_logprior():
    alpha_logprior = []
    for alpha in possible_alphas:
        alpha_logprior.append(log(1/len(possible_alphas))) #unform prior over alpha values
    return normalize_logprobs(alpha_logprior) 

# Take a look at the values returned by `alpha_calculate_logprior()` and check they make sense. You could also think what you'd want those values to look like if you wanted to build in a prior bias favouring e.g. low values of `alpha`, or high values of `alpha`, or values of `alpha` that were low but not too low.



# ### Likelihood for alpha

# Now the final thing we need is a likelihood function - the probability of some data (sequences of 0s and 1s from one or more data sources) given a particular value of `alpha`. This is the only complicated bit of this model, because the likelihood involves several steps to go from a value of `alpha` to some data. 
#
# First we'll consider each data source (e.g. each independent speaker) individually - for conveniencelet's think about the data from source i which we'll call `data_i`. We can ask two questions about `data_i`:
# 1. What is the probability that `data_i` would be generated by a source with a particular value of `pW1`? 
# 2. What is the probability that source i will have that value of `pW1`, given that the particular value of `alpha` we are considering (which captures our expectation about how likely different values of `pW1` are)?
#
# We can put numbers on both of these quite easily. The number in step 1 is the likelihood of data given a particular value of `pW1`, which `pW1_loglikelihood(data, logpW1)` will calculate for us. And the number in step 2 is just our prior over values of `pW1` - how likely is this value of `pW1` given a particular valiue of `alpha` - which `pW1_calculate_logprior(alpha)` gives us. So to work out the likelihood of a particular speaker with a particular source selecting a particular value of `pW1` and then producing a particular sequence of data `data_i` we can just multiply the numbers from 1 and 2 together.

# This lets us put a number on the likelihood that a particular value of `alpha` will lead to a data source (e.g. a speaker) with a particular value of `pW1` who produces a particular sequence of data. But we are actually not actually trying to infer `pW1` for each speaker, we just want to infer `alpha`. So we don't need to keep track of our estimate of the particular value of `pW1` that this speaker is using - we know they'll have *some* value of `pW1`, but it could be any value from `possible_logpW1`. This means to calculate the likelihood of one speaker's data given a particular value of `alpha` we can just calculate `pW1_loglikelihood(data, logpW1) + pW1_calculate_logprior(alpha)` for each possible value of `pW1` and then sum them all together to get the likelihood of the data abstracting over exactly which value of `pW1` this speaker is using (bearing in mind that we need to use a special function for summing log probabilities).
#
# That allows us to put a number on the likelihood of seeing a particular set of data from *one* source given a particular value of `alpha`. But we want to be able to handle multiple data sources - so what we'll do is just perform this calculation for each data source then multiply the resulting per-source probabilities (add them in the log domain) to arrive at a single number that indicates the (log) likelihood of the combined data over all our sources given a particular value of `alpha`.
#
# [An aside that you can skip if you aren't bothered or don't want any more stuff to think about at this point: the reason we sometimes sum and sometimes multiply the probabilities is to do with whether the events we are considering are mutually exclusive or independent. For candidate values of `pW1` for a single source, the various probabilities are mutually exclusive - they either have `pW1=0.05` OR `pW1=0.1` OR ..., so we sum the probabilities over these mutually exclusive probabilities. For combining data from several independent sources, we want to know the likelihood that a particular value of `alpha` leads to the observed data from source 1 AND the observed data from source 2 AND ..., so we multiply the probabilities.]
#
# Our function `alpha_loglikelihood(data, alpha)` will perform all of these calculations for us, calculating the likelihood of a set of data given a particular value of `alpha`. 

def alpha_loglikelihood(data, alpha):
    likelihood_logprobs = [] #there will be one value per data source
    pW1_logprior = pW1_calculate_logprior(alpha) #our prior over pW1
    for data_i in data: #consider each data source in turn
        data_i_likelihood_logprobs = [] #likilihood of data_i given alpha for each value of pW1
        for i in range(len(possible_logpW1)): #consider each pW1 in turn
            logpW1 = possible_logpW1[i] 
            logp_h = pW1_logprior[i] #prior probability of this pW1 given our alpha value
            logp_d = pW1_loglikelihood(data_i, logpW1) #likelihood of data given this pW1
            #likelihood of data_i given this pW1 = p(pW1)*p(data_i|pW1)
            data_i_likelihood_logprobs.append(logp_h + logp_d) #adding logs = multiplying non-logs
        #sum over values of pW1 to get overall likelihood of data_i
        data_i_loglikelihood = logsumexp(data_i_likelihood_logprobs)
        #add the calculated likelihood for data_i to our building list
        likelihood_logprobs.append(data_i_loglikelihood)
    #then at the end take product of p(data_i|alpha) for the independent data sources
    return sum(likelihood_logprobs) #adding logs = multiplying non-logs

# For example, to calculate the likelihood (log probability) that we'll see two highly regular speakers given an alpha value of 5 (which favours highly *variable* speakers), you would do:
#
# ```python
# alpha_loglikelihood([[0,0,0,0,0,0],[1,1,1,1,1,1]], 5)
# ```
# If you'd rather see a probability than a log probability you can do:
# ```python
# exp(alpha_loglikelihood([[0,0,0,0,0,0],[1,1,1,1,1,1]], 5))
# ```
#
# Play around with different values for alpha here. 
# - Which values of `alpha` result in a high likelihood for this data? Which values of `alpha` result in a low likelihood of this data? Why? 
# - Does the fact that our two example speakers here use different variants (one speaker uses variant 0 consistently, the other uses variant 1 consistently) matter to the likelihood? 
# - What happens to the likelihood if you try different kinds of data, e.g. involving highly variable speakers?



# ### Posterior for alpha
#
# Now we have set up the prior and the likelihood we can do Bayesian inference for `alpha` in the usual way - we'll pass in the data and the prior over values of `alpha`, and it'll tell us the posterior probability of each value from `possible_alphas`.

def alpha_logposterior(data, alpha_logprior):
    posterior_logprobs = []
    for i in range(len(possible_alphas)):
        alpha_value = possible_alphas[i] 
        logp_h = alpha_logprior[i] #prior probability of this value of alpha
        logp_d = alpha_loglikelihood(data, alpha_value) #likelihood of data given this alpha value
        posterior_logprobs.append(logp_h + logp_d) #adding logs = multiplying non-logs
    return normalize_logprobs(posterior_logprobs) 

# So if you want to see the posterior probability distribution over our `possible_alphas` you would do:
# ```python
# alpha_logposterior([[0,0,0,0,0,0],[1,1,1,1,1,1]],alpha_calculate_logprior())
# ```
# - Check this runs and you understand the meaning of the values it returns.
# - Try plugging in different kinds of data (e.g. data from three individuaully-variable speakers) to see how that affects the posterior distribution over alpha values. How does the model match up to your intuitions about what we can learn from exposure to individually-variable or individually-consistent speakers?



# Finally, we might also want to have our model predict, based on what it's learned from some data, what *new* data sources (i.e. new speakers) would look like. `alpha_logproduce` will do this for you - you pass it a posterior produced by `alpha_logposterior`, tell it how many sources you want it to generate data for (`n_sources`) and how many data points you want per source (`n_productions_per_source`) and it'll generate some random data for you. 

def alpha_logproduce(alpha_logposterior, n_sources, n_productions_per_source):
    alpha_i_index = log_roulette_wheel(alpha_logposterior) #select a value of alpha based on alpha_logposterior
    alpha = possible_alphas[alpha_i_index]
    #print(alpha) #uncomment this if you want to see the alpha value it has selected from the posterior
    pW1_logprior = pW1_calculate_logprior(alpha) #prior over pW1 given this value of alpha
    produced_data = []
    for i in range(n_sources): #for each source
        pW1_i_index = log_roulette_wheel(pW1_logprior) #select a random pW1 for this source based on the prior
        pW1_i = possible_logpW1[pW1_i_index] 
        data_i = pW1_logproduce(pW1_i, n_productions_per_source) #generate some data using this pW1
        produced_data.append(data_i) #and add it to our building data
    return produced_data

# For example, you can use this code as follows to calculate a posterior based on data from two speakers and then use it to predict the data that would be generated by two new speakers each producing 10 data items:
# ```python
# a_posterior = alpha_logposterior([[0,0,0,0,0,0],[1,1,1,1,1,1]],alpha_calculate_logprior())
# alpha_logproduce(a_posterior, 2, 10)
# ```

# Again, have a little play around with this to see how the data affects the inferred posterior and how that in turn affects the new data generated by `alpha_logproduce`.



# ## Questions

# The priority for this worksheet is to work through the in-text questions above, experimenting with the code and checking in particular that the `alpha_likelihood`, `alpha_logposterior` and `alpha_logproduce` functions make sense to you. Once you are happy with that, try these questions:
#
# 1. Imagine you have data from two speakers each producing 6 data items. How does the data generated by those speakers affect the learner's posterior probability distribution over values of `alpha`? You could see what happens when learning from highly regular speakers who both use the same variant, highly regular speakers who use *different* variants, two highly variable speakers, or two speakers who differ in their variability. How does the behaviour of the model match (or not!) your intuitions about how a learner should respond to this kind of data?
# 2. How do those results change if you have more data from each speaker, e.g. 12 data points per speaker rather than 6?
# 3. How does changing the number of speakers (sources) change the shape of the posterior? For instance, if you have 16 data points split across 1, 2, 4, 8, or 16 speakers, how does it affect the learner's inferred posterior over `alpha`, and why? It might be worth thinking about your intuitions here before you run the model - if you want to infer what speakers are like in general, would you rather have lots of data from very few speakers, very little data from lots of speakers, or something in the middle? Again, see if the model matches your intuitions.
# 4. I was slightly cowardly about the prior over values of `alpha` and just used a uniform prior. Do you think you can mount an argument here that this prior should be something other than uniform? If you are feeling ambitious, plug in a different shape over prior (by creating a new version of `alpha_calculate_logprior`) and see how it affects the model's behaviour. 



StarcoderdataJupyterScriptsDedupFiltered