## 2021-05-24 craig » density_toy.py

 1 ```import numpy ``` ```import matplotlib ``` ```import matplotlib.pyplot ``` ```class Density(object): ``` ``` def __init__(self, sample_size): ``` ``` """ ``` ``` Set empty data ``` ``` """ ``` ``` self.initial_sample_size = sample_size ``` ``` self.distribution = numpy.matrix([]) ``` ``` self.fig = matplotlib.pyplot.figure(figsize=(20,10)) ``` ``` self.axes = [ ``` ``` self.fig.add_subplot(1, 2, 1), ``` ``` self.fig.add_subplot(1, 2, 2), ``` ``` ] ``` ``` self.bin_size = 0.1 ``` ``` self.n_bins = 20 ``` ``` def gen_distribution(self, left, right): ``` ``` """ ``` ``` Generate a triangular sawtooth distribution size initial_sample_size ``` ``` """ ``` ``` self.distribution = numpy.random.triangular(left, left, right, self.initial_sample_size) ``` ``` def survival_probability(self, x, left, right, efficiency_left, efficiency_right): ``` ``` """ ``` ``` Calculate survival probability for sawtooth having efficiency at min efficiency_left and max efficiency_right ``` ``` """ ``` ``` survival_probability = (x-left)/(right-left)*(efficiency_right-efficiency_left)+efficiency_left ``` ``` if survival_probability < 0: ``` ``` return 0.0 ``` ``` if survival_probability > 1: ``` ``` return 1.0 ``` ``` return survival_probability ``` ``` def inefficiency(self, left, right, efficiency_left, efficiency_right): ``` ``` """ ``` ``` Generate a mock inefficiency (non-uniform) ``` ``` """ ``` ``` print("Inefficiency", len(self.distribution), "...", end = " ") ``` ``` ran_1 = numpy.random.uniform(0.0, 1.0, len(self.distribution)) ``` ``` distribution_copy = [] ``` ``` for ran, x in zip(ran_1, self.distribution): ``` ``` survival_probability = self.survival_probability(x, left, right, efficiency_left, efficiency_right) ``` ``` if ran < survival_probability: ``` ``` distribution_copy.append(x) ``` ``` self.distribution = numpy.array(distribution_copy) ``` ``` print(len(self.distribution)) ``` ``` def plot_distribution(self): ``` ``` """ ``` ``` Plot a histogram of the distribution. Bin size and number of bins user controlled ``` ``` """ ``` ``` if self.fig == None: ``` ``` self.fig = matplotlib.pyplot.figure() ``` ``` self.axes[0].hist(self.distribution, self.n_bins, [0.0, self.n_bins*self.bin_size]) ``` ``` self.axes[0].set_xlabel("x [m]") ``` ``` self.axes[0].set_ylabel("Number") ``` ``` def plot_density(self): ``` ``` """ ``` ``` Plot a graph of the number of events per unit length i.e. 1D number density ``` ``` """ ``` ``` self.bin_count = [0 for i in range(self.n_bins)] ``` ``` self.bin_centre = [(i+0.5)*self.bin_size for i in range(self.n_bins)] ``` ``` for i in range(self.n_bins): ``` ``` low = i*self.bin_size ``` ``` high = low+self.bin_size ``` ``` for x in self.distribution: ``` ``` if x >= low and x < high: ``` ``` self.bin_count[i] += 1 ``` ``` self.density = [n/self.bin_size for n in self.bin_count] ``` ``` self.axes[1].plot(self.bin_centre, self.density) ``` ``` self.axes[1].set_xlabel("x [m]") ``` ``` self.axes[1].set_ylabel("density [number m\$^{-1}\$]") ``` ```unique_id = 0 ``` ```def do_plots(bin_size_1, n_bins_1, bin_size_2, n_bins_2): ``` ``` global unique_id ``` ``` my_density = Density(1000000) ``` ``` my_density.bin_size = bin_size_1 ``` ``` my_density.n_bins = n_bins_1 ``` ``` my_density.gen_distribution(0.0, 2.0) ``` ``` my_density.plot_distribution() ``` ``` my_density.plot_density() ``` ``` for i in range(0, 21): ``` ``` x = i/10.0 ``` ``` p = my_density.survival_probability(x, 0.0, 1.0, 1.0, 0.0) ``` ``` print(format(x, "4.2g"), format(p, "4.2g")) ``` ``` my_density.bin_size = bin_size_2 ``` ``` my_density.n_bins = n_bins_2 ``` ``` my_density.inefficiency(0.0, 1.0, 1.0, 0.0) ``` ``` my_density.plot_distribution() ``` ``` my_density.plot_density() ``` ``` my_density.fig.suptitle("bin size 1: "+str(bin_size_1)+" number of bins 1: "+str(n_bins_1)+"\n"+\ ``` ``` "bin size 2: "+str(bin_size_2)+" number of bins 2: "+str(n_bins_2)) ``` ``` unique_id += 1 ``` ``` my_density.fig.savefig("density_"+str(unique_id)+".png") ``` ```def main(): ``` ``` do_plots(0.01, 200, 0.01, 200) ``` ``` do_plots(0.1, 20, 0.1, 20) ``` ``` do_plots(1.0, 2, 1.0, 2) ``` ``` do_plots(2.0, 2, 1.0, 2) ``` ```if __name__ == "__main__": ``` ``` main() ``` ``` matplotlib.pyplot.show(block=False) ``` ``` input("Press to finish") ```
(1-1/4)