1
|
import numpy
|
2
|
import matplotlib
|
3
|
import matplotlib.pyplot
|
4
|
|
5
|
class Density(object):
|
6
|
def __init__(self, sample_size):
|
7
|
"""
|
8
|
Set empty data
|
9
|
"""
|
10
|
self.initial_sample_size = sample_size
|
11
|
self.distribution = numpy.matrix([])
|
12
|
self.fig = matplotlib.pyplot.figure(figsize=(20,10))
|
13
|
self.axes = [
|
14
|
self.fig.add_subplot(1, 2, 1),
|
15
|
self.fig.add_subplot(1, 2, 2),
|
16
|
]
|
17
|
self.bin_size = 0.1
|
18
|
self.n_bins = 20
|
19
|
|
20
|
def gen_distribution(self, left, right):
|
21
|
"""
|
22
|
Generate a triangular sawtooth distribution size initial_sample_size
|
23
|
"""
|
24
|
self.distribution = numpy.random.triangular(left, left, right, self.initial_sample_size)
|
25
|
|
26
|
def survival_probability(self, x, left, right, efficiency_left, efficiency_right):
|
27
|
"""
|
28
|
Calculate survival probability for sawtooth having efficiency at min efficiency_left and max efficiency_right
|
29
|
"""
|
30
|
survival_probability = (x-left)/(right-left)*(efficiency_right-efficiency_left)+efficiency_left
|
31
|
if survival_probability < 0:
|
32
|
return 0.0
|
33
|
if survival_probability > 1:
|
34
|
return 1.0
|
35
|
return survival_probability
|
36
|
|
37
|
def inefficiency(self, left, right, efficiency_left, efficiency_right):
|
38
|
"""
|
39
|
Generate a mock inefficiency (non-uniform)
|
40
|
"""
|
41
|
print("Inefficiency", len(self.distribution), "...", end = " ")
|
42
|
ran_1 = numpy.random.uniform(0.0, 1.0, len(self.distribution))
|
43
|
distribution_copy = []
|
44
|
for ran, x in zip(ran_1, self.distribution):
|
45
|
survival_probability = self.survival_probability(x, left, right, efficiency_left, efficiency_right)
|
46
|
if ran < survival_probability:
|
47
|
distribution_copy.append(x)
|
48
|
self.distribution = numpy.array(distribution_copy)
|
49
|
print(len(self.distribution))
|
50
|
|
51
|
def plot_distribution(self):
|
52
|
"""
|
53
|
Plot a histogram of the distribution. Bin size and number of bins user controlled
|
54
|
"""
|
55
|
if self.fig == None:
|
56
|
self.fig = matplotlib.pyplot.figure()
|
57
|
self.axes[0].hist(self.distribution, self.n_bins, [0.0, self.n_bins*self.bin_size])
|
58
|
self.axes[0].set_xlabel("x [m]")
|
59
|
self.axes[0].set_ylabel("Number")
|
60
|
|
61
|
def plot_density(self):
|
62
|
"""
|
63
|
Plot a graph of the number of events per unit length i.e. 1D number density
|
64
|
"""
|
65
|
self.bin_count = [0 for i in range(self.n_bins)]
|
66
|
self.bin_centre = [(i+0.5)*self.bin_size for i in range(self.n_bins)]
|
67
|
for i in range(self.n_bins):
|
68
|
low = i*self.bin_size
|
69
|
high = low+self.bin_size
|
70
|
for x in self.distribution:
|
71
|
if x >= low and x < high:
|
72
|
self.bin_count[i] += 1
|
73
|
self.density = [n/self.bin_size for n in self.bin_count]
|
74
|
self.axes[1].plot(self.bin_centre, self.density)
|
75
|
self.axes[1].set_xlabel("x [m]")
|
76
|
self.axes[1].set_ylabel("density [number m$^{-1}$]")
|
77
|
|
78
|
|
79
|
unique_id = 0
|
80
|
def do_plots(bin_size_1, n_bins_1, bin_size_2, n_bins_2):
|
81
|
global unique_id
|
82
|
my_density = Density(1000000)
|
83
|
my_density.bin_size = bin_size_1
|
84
|
my_density.n_bins = n_bins_1
|
85
|
my_density.gen_distribution(0.0, 2.0)
|
86
|
my_density.plot_distribution()
|
87
|
my_density.plot_density()
|
88
|
for i in range(0, 21):
|
89
|
x = i/10.0
|
90
|
p = my_density.survival_probability(x, 0.0, 1.0, 1.0, 0.0)
|
91
|
print(format(x, "4.2g"), format(p, "4.2g"))
|
92
|
my_density.bin_size = bin_size_2
|
93
|
my_density.n_bins = n_bins_2
|
94
|
my_density.inefficiency(0.0, 1.0, 1.0, 0.0)
|
95
|
my_density.plot_distribution()
|
96
|
my_density.plot_density()
|
97
|
my_density.fig.suptitle("bin size 1: "+str(bin_size_1)+" number of bins 1: "+str(n_bins_1)+"\n"+\
|
98
|
"bin size 2: "+str(bin_size_2)+" number of bins 2: "+str(n_bins_2))
|
99
|
unique_id += 1
|
100
|
my_density.fig.savefig("density_"+str(unique_id)+".png")
|
101
|
|
102
|
def main():
|
103
|
do_plots(0.01, 200, 0.01, 200)
|
104
|
do_plots(0.1, 20, 0.1, 20)
|
105
|
do_plots(1.0, 2, 1.0, 2)
|
106
|
do_plots(2.0, 2, 1.0, 2)
|
107
|
|
108
|
if __name__ == "__main__":
|
109
|
main()
|
110
|
matplotlib.pyplot.show(block=False)
|
111
|
input("Press <CR> to finish")
|
112
|
|