forked from probml/pyprobml
-
Notifications
You must be signed in to change notification settings - Fork 0
/
bimodal_dist_plot.py
49 lines (38 loc) · 1.37 KB
/
bimodal_dist_plot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
# Bimodal distribution (mixture of two 1d Gaussians)
# Based on https://github.com/probml/pmtk3/blob/master/demos/bimodalDemo.m
import numpy as np
import matplotlib.pyplot as plt
import os
figdir = os.path.join(os.environ["PYPROBML"], "figures")
def save_fig(fname): plt.savefig(os.path.join(figdir, fname))
from scipy.stats import norm
# Define two normal distrubutions and their corresponding weights.
mu = [0, 2]
sigma = [1, 0.05]
n = [norm(loc=mu[i], scale=sigma[i]) for i in range(2)]
w = [0.5, 0.5]
# Define a set of x points for graphing.
xs = np.linspace(-2, 2*mu[1], 600)
# Combine the two distributions by their weights, evaluated at the x points.
p = sum(w[i] * n[i].pdf(xs) for i in range(2))
# Calculate the mean of the final distribution.
mean_p = np.mean(xs * p)
# Plot the final distribution and its mean.
linewidth = 3
plt.figure()
plt.plot(xs, p, 'black', linewidth=linewidth)
plt.vlines(mean_p, ymin=0, ymax=max(p), color='red', linewidth=linewidth)
save_fig('bimodalSpike.pdf')
plt.show()
# Another example, with two modes
mu = [0, 2]
sigma = [0.5, 0.5]
n = [norm(loc=mu[i], scale=sigma[i]) for i in range(2)]
w = [0.5, 0.5]
xs = np.linspace(-2, 2*mu[1], 600)
p = sum(w[i] * n[i].pdf(xs) for i in range(2))
plt.figure()
linewidth = 3
plt.plot(xs, p, 'black', linewidth=linewidth)
save_fig('bimodalDistribution.pdf')
plt.show()