11
11
from matplotlib .collections import PolyCollection
12
12
import matplotlib .pyplot as plt
13
13
import numpy as np
14
+ from scipy .stats import poisson
14
15
15
16
# Fixing random state for reproducibility
16
17
np .random .seed (19680801 )
17
18
18
19
19
- def polygon_under_graph (xlist , ylist ):
20
+ def polygon_under_graph (x , y ):
20
21
"""
21
22
Construct the vertex list which defines the polygon filling the space under
22
- the (xlist, ylist ) line graph. Assumes the xs are in ascending order.
23
+ the (x, y ) line graph. This assumes x is in ascending order.
23
24
"""
24
- return [(xlist [0 ], 0. ), * zip (xlist , ylist ), (xlist [- 1 ], 0. )]
25
+ return [(x [0 ], 0. ), * zip (x , y ), (x [- 1 ], 0. )]
25
26
26
27
27
28
ax = plt .figure ().add_subplot (projection = '3d' )
28
29
29
- # Make verts a list such that verts[i] is a list of (x, y) pairs defining
30
- # polygon i.
31
- verts = []
30
+ x = np .linspace (0. , 10. , 31 )
31
+ lambdas = range (1 , 9 )
32
32
33
- # Set up the x sequence
34
- xs = np .linspace (0. , 10. , 26 )
33
+ # verts[i] is a list of (x, y) pairs defining polygon i.
34
+ verts = [polygon_under_graph (x , poisson .pmf (l , x )) for l in lambdas ]
35
+ facecolors = plt .get_cmap ('viridis_r' )(np .linspace (0 , 1 , len (verts )))
35
36
36
- # The ith polygon will appear on the plane y = zs[i]
37
- zs = range ( 4 )
37
+ poly = PolyCollection ( verts , facecolors = facecolors , alpha = .7 )
38
+ ax . add_collection3d ( poly , zs = lambdas , zdir = 'y' )
38
39
39
- for i in zs :
40
- ys = np .random .rand (len (xs ))
41
- verts .append (polygon_under_graph (xs , ys ))
40
+ ax .set (xlim = (0 , 10 ), ylim = (1 , 9 ), zlim = (0 , 0.35 ),
41
+ xlabel = 'x' , ylabel = '$\lambda$' , zlabel = 'probability' )
42
42
43
- poly = PolyCollection (verts , facecolors = ['r' , 'g' , 'b' , 'y' ], alpha = .6 )
44
- ax .add_collection3d (poly , zs = zs , zdir = 'y' )
45
-
46
- ax .set_xlabel ('X' )
47
- ax .set_ylabel ('Y' )
48
- ax .set_zlabel ('Z' )
49
- ax .set_xlim (0 , 10 )
50
- ax .set_ylim (- 1 , 4 )
51
- ax .set_zlim (0 , 1 )
52
-
53
- plt .show ()
43
+ plt .show ()
0 commit comments