#!/usr/bin/env python
# coding: utf-8

# # Week 2: Monte Carlo
# 
# ## Objectives
# 
# * Work through two examples using MC
#     * Random walk
#     * Nuclear decay

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'notebook')


# <font color="red">
# 
# > Warning: if you try to run the cell following an interactive Matplotlib plot very quickly, Jupyter might not actually run the cell. Just try again.
# 
# </font>

# Standard imports

# In[ ]:


import matplotlib.pyplot as plt
import time
import numpy as np


# ## Random walk
# 
# Example in book, but without VPython.

# We first will need to set up a plot. Since we are using an interactive backend for Matplotlib, we need to set up the plot first. We'll plot on point in the center. We'll also set (and not update) the plot limits.

# In[ ]:


fig, ax = plt.subplots()
(line,) = plt.plot([0], [0], ".-")
ax.set_xlim(-5, 5)
ax.set_ylim(-5, 5)
plt.show()


# Now, we loop and generate random numbers for the walk. We are currently ignoring the step length (it will not be exactly one). We'll also simply get the current plot arrays, make new arrays with one extra point, then put them in. A small sleep parameter will help the plot be slow enough to see.

# In[ ]:


jmax = 100
x = 0.0
y = 0.0

for i in range(jmax + 1):
    x += (np.random.random() - 0.5) * 2
    y += (np.random.random() - 0.5) * 2
    xs, ys = line.get_data()
    xs = np.append(xs, [x])
    ys = np.append(ys, [y])
    line.set_data(xs, ys)
    fig.canvas.draw()
    time.sleep(0.1)


# In[ ]:


print("r actual =", np.sqrt(x ** 2 + y ** 2))
print("r expected =", np.sqrt(jmax) * 0.75)


# Exercises:
# * Modify the code so that the step length is one.
# * Modify the code so that the edges of the plot change. Use `np.min` and `np.max` (and maybe add a small scaling factor so keep the plot from touching the edges).

# Let's try again, but this time using numpy syntax to avoid the explicit loop. We'll add the radius correction here.

# In[ ]:


fig, ax = plt.subplots()
(line,) = plt.plot([0], [0], ".-")
ax.set_xlim(-5, 5)
ax.set_ylim(-5, 5)
plt.show()


# In[ ]:


def rwalk(size):
    xs = (np.random.random(size) - 0.5) * 2
    ys = (np.random.random(size) - 0.5) * 2

    r = np.sqrt(xs ** 2 + ys ** 2)
    xs /= r
    ys /= r

    # cumulative sum. [n0, n0+n1, n0+n1+n2, ...]
    xloc = np.cumsum(xs)
    yloc = np.cumsum(ys)

    return xloc, yloc


# In[ ]:


xloc, yloc = rwalk(jmax)
for i in range(jmax):
    line.set_data(xloc[:i], yloc[:i])
    fig.canvas.draw()
    time.sleep(0.1)


# Exercises:
# 
# * Change this to use a single 2D array, of shape `(2,N)`, where `array[0]` is x and `array[1]` is y. You'll need to specify an axis for the cumsum operation.

# We'll plot five of these random walks. Let's rely on Matplotlib's color cycling to keep this clear. We'll use the `*` to expand the two items returned from the `rwalk` into the two items expected by `plot`.

# In[ ]:


fig, ax = plt.subplots()
for i in range(5):
    plt.plot(*rwalk(jmax), ".-")

plt.show()


# Let's redo that function and just compute the final sum - we'll then compute the radius and return that.

# In[ ]:


def rwalksum(size):
    xs = (np.random.random(size) - 0.5) * 2
    ys = (np.random.random(size) - 0.5) * 2

    r = np.sqrt(xs ** 2 + ys ** 2)
    xs /= r
    ys /= r

    xloc = np.sum(xs)
    yloc = np.sum(ys)

    return np.sqrt(xloc ** 2 + yloc ** 2)


# Now, we'll fill in 100,000 entries of an empty array with the values from 100,000 trials.

# In[ ]:


vals = np.empty(100_000, dtype=np.float64)
for i in range(len(vals)):
    vals[i] = rwalksum(jmax)


# Now, let's make a histogram. Letting Matplotlib (really Numpy under the hood) pick the binning is fine:

# In[ ]:


fig, ax = plt.subplots()
ax.hist(vals, bins="auto")
plt.show()


# In[ ]:


np.mean(vals)


# ## Radioactive decay
# 
# Now, let's look at radioactive decay:

# In[ ]:


particles = 200
decay_prob = 0.001
decays = []  # list of decay times

for time in range(500):
    for particle in range(particles):
        decay = np.random.random()
        if decay < decay_prob:
            particles -= 1
            decays.append(time)


# In[ ]:


fig, ax = plt.subplots()
ax.hist(decays, bins=range(500))
# ax.set_yscale('log') # Uncomment if you have lots of particles
plt.show()


# We can build a sound array and play it:

# In[ ]:


from IPython.display import Audio


# We will use the norm from scipy instead of manually computing a gaussian.

# In[ ]:


from scipy.stats import norm


# In[ ]:


ar = np.linspace(-10, 10, 100)
g = norm(0, 3).pdf(ar) * np.sin(ar * 4)


# Let's look at our little sound.

# In[ ]:


plt.figure()
plt.plot(g)
plt.plot(norm(0, 3).pdf(ar))
plt.show()


# We'll spread this out into 100 values per time bin, then we'll convolve this sound shape with each values (these act like delta functions).

# In[ ]:


d = np.histogram(np.array(decays) * 100, bins=range(50000))[0]
d = np.convolve(d, g)
d = d / d.sum()

Audio(data=d, rate=3000)


# Let's look at one of these decays, just to make sure our convolution worked.

# In[ ]:


plt.figure()
plt.plot(d[: decays[0] * 100 + 500])
plt.show()