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 = (xleft)/(rightleft)*(efficiency_rightefficiency_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 (nonuniform)

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

