<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2020 The TensorFlow Authors.
Licensed under the Apache License, Version 2.0 (the "License");
Step1:  Chapter 2 - Small Worlds and Large Worlds
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https
Step2:  2.1.3. From counts to probability
Suppose we have a bag of four marbles, each of which can be either blue or white. There are five possible compositions of the bag, ranging from all white marbles to all blue marbles. Now, suppose we make three draws with replacement from the bag, and end up with blue, then white, then blue. If we were to count the ways that each bag composition could produce this combination of samples, we would find the vector in ways below. For example, there are zero ways that four white marbles could lead to this sample, three ways that a bag composition of one blue marble and three white marbles could lead to the sample, as so on.
We can convert these counts to probabilities by simply dividing the counts by the sum of all the counts.
Code 2.1
Step3:  2.3.2.1. Observed variables
Consider a globe of the earth that we toss into the air nine times. We want to compute the probability that our right index finger will land on water six times out of nine. We can use the binomial distribution to compute this, using a probability of 0.5 of landing on water on each toss.
Code 2.2
Step4:  2.4.3.Grid Approximation
Create a grid approximation for the globe tossing model, for the scenario described in the section above. In the grid approximation technique, we approximate a continuous posterior distribution by using a discrete grid of parameter values.
Code 2.3
Step5:  Code 2.4
Step6:  Code 2.3 and 2.4 using TFP distributions (OPTIONAL)
To replicate the other two priors depicted in Figure 2.6 in the text, try using the following priors, one at time in the Code 2.3 sample above.
Code 2.5
Step7:  2.4.4. Quadratic Approximation
Code 2.6
TFP doesn't natively provide a quap function, like in the book. But it has a lot of optimization tools. We'll use bfgs, which gives us a somewhat straightforward way to get both the measurement and the standard error.
Step8:  The results object itself has a lot of information about the optimization process. The estimate is called the position, and we get the standard error from the inverse_hessian_estimate.
Step9:  Let's compare the mean and standard deviation from the quadratic approximation with an analytical approach based on the Beta distribution.
Code 2.7
Step10:  2.4.5. Markov chain Monte Carlo
We can estimate the posterior using a Markov chain Monte Carlo (MCMC) technique, which will be explained further in Chapter 9. An outline of the algorithm follows. It is written primarily in numpy to illustrate the steps you would take.
n_samples = 1000
p = np.zeros(n_samples)
p[0] = 0.5
W = 6
L = 3
for i in range(1, n_samples)
Step11:  Code 2.9
With the samples from the posterior distribution in p, we can compare the results of the MCMC approximation to the analytical posterior.
<ASSISTANT_TASK:>
Python Code:
#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Explanation: Copyright 2020 The TensorFlow Authors.
Licensed under the Apache License, Version 2.0 (the "License");
End of explanation
#@title Install { display-mode: "form" }
TF_Installation = 'System' #@param ['TF Nightly', 'TF Stable', 'System']
if TF_Installation == 'TF Nightly':
  !pip install -q --upgrade tf-nightly
  print('Installation of `tf-nightly` complete.')
elif TF_Installation == 'TF Stable':
  !pip install -q --upgrade tensorflow
  print('Installation of `tensorflow` complete.')
elif TF_Installation == 'System':
  pass
else:
  raise ValueError('Selection Error: Please select a valid '
                   'installation option.')
#@title Install { display-mode: "form" }
TFP_Installation = "System" #@param ["Nightly", "Stable", "System"]
if TFP_Installation == "Nightly":
  !pip install -q tfp-nightly
  print("Installation of `tfp-nightly` complete.")
elif TFP_Installation == "Stable":
  !pip install -q --upgrade tensorflow-probability
  print("Installation of `tensorflow-probability` complete.")
elif TFP_Installation == "System":
  pass
else:
  raise ValueError("Selection Error: Please select a valid "
                   "installation option.")
import numpy as np
import arviz as az
import pandas as pd
import tensorflow as tf
import tensorflow_probability as tfp
import scipy.stats as stats
# visualization 
import matplotlib.pyplot as plt
# aliases
tfd = tfp.distributions
az.style.use('seaborn-colorblind')
Explanation: Chapter 2 - Small Worlds and Large Worlds
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/probability/examples/statistical_rethinking/notebooks/02_small_worlds_and_large_worlds"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/probability/blob/main/tensorflow_probability/examples/statistical_rethinking/notebooks/02_small_worlds_and_large_worlds.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/probability/blob/main/tensorflow_probability/examples/statistical_rethinking/notebooks/02_small_worlds_and_large_worlds.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/probability/examples/statistical_rethinking/notebooks/02_small_worlds_and_large_worlds.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>
Imports and utility functions
End of explanation
# define a list of constants
# divide by ways each value can occur
ways = tf.constant([0., 3, 8, 9, 0])
new_ways = ways / tf.reduce_sum(ways)
new_ways
Explanation: 2.1.3. From counts to probability
Suppose we have a bag of four marbles, each of which can be either blue or white. There are five possible compositions of the bag, ranging from all white marbles to all blue marbles. Now, suppose we make three draws with replacement from the bag, and end up with blue, then white, then blue. If we were to count the ways that each bag composition could produce this combination of samples, we would find the vector in ways below. For example, there are zero ways that four white marbles could lead to this sample, three ways that a bag composition of one blue marble and three white marbles could lead to the sample, as so on.
We can convert these counts to probabilities by simply dividing the counts by the sum of all the counts.
Code 2.1
End of explanation
# probability of 6 successes in 9 trials with 0.5 probability
tfd.Binomial(total_count=9, probs=0.5).prob(6)
Explanation: 2.3.2.1. Observed variables
Consider a globe of the earth that we toss into the air nine times. We want to compute the probability that our right index finger will land on water six times out of nine. We can use the binomial distribution to compute this, using a probability of 0.5 of landing on water on each toss.
Code 2.2
End of explanation
# define grid
n_points = 20 # change to an odd number for Code 2.5 graphs to 
              # match book examples in Figure 2.6
p_grid = tf.linspace(start=0., stop=1., num=n_points)
#define prior
prior = tf.ones([n_points])
# compute likelihood at each value in grid
likelihood = tfd.Binomial(total_count=9, probs=p_grid).prob(6)
# compute product of likelihood and prior
unstd_posterior = likelihood * prior
# standardize the posterior, so it sums to 1
posterior = unstd_posterior / tf.reduce_sum(unstd_posterior)
posterior
Explanation: 2.4.3.Grid Approximation
Create a grid approximation for the globe tossing model, for the scenario described in the section above. In the grid approximation technique, we approximate a continuous posterior distribution by using a discrete grid of parameter values.
Code 2.3
End of explanation
_, ax = plt.subplots(figsize=(9, 4))
ax.plot(p_grid, posterior, "-o")
ax.set(
    xlabel="probability of water",
    ylabel="posterior probability",
    title="20 points");
Explanation: Code 2.4
End of explanation
first_prior = tf.where(condition=p_grid < 0.5, x=0., y=1)
second_prior = tf.exp(-5 * abs(p_grid - 0.5))
_, axes = plt.subplots(nrows=2, ncols=2, figsize=(12, 4), 
                       constrained_layout=True)
axes[0, 0].plot(p_grid, first_prior)
axes[0, 0].set_title('First prior')
axes[1, 0].plot(p_grid, first_prior * likelihood)
axes[1, 0].set_title('First posterior')
axes[0, 1].plot(p_grid, second_prior)
axes[0, 1].set_title('Second prior')
axes[1, 1].plot(p_grid, second_prior * likelihood)
axes[1, 1].set_title('Second posterior');
Explanation: Code 2.3 and 2.4 using TFP distributions (OPTIONAL)
To replicate the other two priors depicted in Figure 2.6 in the text, try using the following priors, one at time in the Code 2.3 sample above.
Code 2.5
End of explanation
W = 6
L = 3
dist = tfd.JointDistributionNamed({
    "water": lambda probability: tfd.Binomial(total_count=W + L, 
                                              probs=probability),
    "probability": tfd.Uniform(low=0., high=1.)
})
def neg_log_prob(x):
    return tfp.math.value_and_gradient(
          lambda p: -dist.log_prob(
              water=W, 
              probability=tf.clip_by_value(p[-1], 0., 1.)), 
          x,
    )
results = tfp.optimizer.bfgs_minimize(neg_log_prob, initial_position=[0.5])
assert results.converged
Explanation: 2.4.4. Quadratic Approximation
Code 2.6
TFP doesn't natively provide a quap function, like in the book. But it has a lot of optimization tools. We'll use bfgs, which gives us a somewhat straightforward way to get both the measurement and the standard error.
End of explanation
approximate_posterior = tfd.Normal(
    results.position,
    tf.sqrt(results.inverse_hessian_estimate),
)
print(
    "mean:", approximate_posterior.mean(),
    "\nstandard deviation: ", approximate_posterior.stddev(),
)
Explanation: The results object itself has a lot of information about the optimization process. The estimate is called the position, and we get the standard error from the inverse_hessian_estimate.
End of explanation
_, ax = plt.subplots(figsize=(9, 4))
x = tf.linspace(0., 1., num=101)
ax.plot(x, tfd.Beta(W + 1, L + 1).prob(x), label='Analytic posterior')
# values obained from quadratic approximation
ax.plot(x, tf.squeeze(approximate_posterior.prob(x)), "--", 
         label='Quadratic approximation')
ax.set(
    xlabel='Probability of water',
    ylabel='Posterior probability',
    title='Comparing quadratic approximation to analytic posterior'
)
ax.legend();
Explanation: Let's compare the mean and standard deviation from the quadratic approximation with an analytical approach based on the Beta distribution.
Code 2.7
End of explanation
@tf.function
def do_sampling():
  def get_model_log_prob(probs):
    return tfd.Binomial(total_count=W + L, probs=probs).log_prob(W)
  sampling_kernel = tfp.mcmc.RandomWalkMetropolis(get_model_log_prob)
  return tfp.mcmc.sample_chain(
      num_results=5000,
      current_state=.5,
      kernel=sampling_kernel,
      num_burnin_steps=500,
    trace_fn=None, 
  )
samples = do_sampling()
Explanation: 2.4.5. Markov chain Monte Carlo
We can estimate the posterior using a Markov chain Monte Carlo (MCMC) technique, which will be explained further in Chapter 9. An outline of the algorithm follows. It is written primarily in numpy to illustrate the steps you would take.
n_samples = 1000
p = np.zeros(n_samples)
p[0] = 0.5
W = 6
L = 3
for i in range(1, n_samples):
    p_new = tfd.Normal(loc=p[i - 1], scale=0.1).sample(1)
    if p_new &lt; 0:
        p_new = -p_new
    if p_new &gt; 1:
        p_new = 2 - p_new
    q0 = tfd.Binomial(total_count=W+L, probs=p[i - 1]).prob(W)
    q1 = tfd.Binomial(total_count=W+L, probs=p_new).prob(W)
    if stats.uniform.rvs(0, 1) &lt; q1 / q0:
        p[i] = p_new
    else:
        p[i] = p[i - 1]
But to actually use Metropolis Hastings in TFP, you can use the mcmc library. The version of the algorithm that we want is called RandomWalkMetropolis.
Code 2.8
End of explanation
_, ax = plt.subplots(figsize=(9, 4))
az.plot_kde(samples, label="Metropolis approximation", ax=ax)
x = tf.linspace(0., 1., num=100)
ax.plot(x, tfd.Beta(W + 1, L + 1).prob(x), "C1", label="True posterior")
ax.legend();
Explanation: Code 2.9
With the samples from the posterior distribution in p, we can compare the results of the MCMC approximation to the analytical posterior.
End of explanation
<END_TASK>