Project

General

Profile

2021-05-24 craig » density_toy.py

Rogers, Chris, 26 May 2021 13:22

 
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

    
(1-1/4)