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

 1 ```import copy ``` ```import numpy ``` ```import matplotlib ``` ```import matplotlib.pyplot ``` ```import scipy.stats ``` ```class Lattice(object): ``` ``` def __init__(self, step_size, n_steps): ``` ``` """ ``` ``` Generalised tracking element. Overload "acceleration" to generate ``` ``` arbitrary symplectic tracking. Overload "transport_one" to do ``` ``` non-symplectic things. ``` ``` """ ``` ``` self.step_size = step_size ``` ``` self.n_steps = n_steps ``` ``` def transport(self, distribution): ``` ``` """ ``` ``` Transport a distribution. This just iterates over the psv_list to do the ``` ``` transport. ``` ``` """ ``` ``` for i, psv in enumerate(distribution.psv_list): ``` ``` distribution.psv_list[i] = self.transport_one(psv) ``` ``` distribution.psv_list = [psv for psv in distribution.psv_list if psv is not None] ``` ``` return distribution ``` ``` def transport_one(self, psv): ``` ``` """ ``` ``` Following "leap frog integration" which is a symplectic 2nd order integrator ``` ``` http://www.artcompsci.org/vol_1/v1_web/node34.html ``` ``` It is symplectic, because at each step we perform 2 shears in phase ``` ``` space - a shear in x followed by a shear in v. ``` ``` """ ``` ``` dt = self.step_size ``` ``` x_im1 = psv[0] ``` ``` v_imh = psv[1]+self.acceleration(x_im1, 0)*dt/2.0 ``` ``` for i in range(self.n_steps): ``` ``` # calculate the values for the next time step ``` ``` x_i = x_im1+v_imh*dt ``` ``` a_i = self.acceleration(x_i, 0) ``` ``` v_iph = v_imh + a_i*dt ``` ``` # update the values ready for the next step. We can go directly but ``` ``` # I include this step for clarity ``` ``` x_im1 = x_i ``` ``` v_imh = v_iph ``` ``` psv[0] = x_im1 ``` ``` psv[1] = v_imh+self.acceleration(x_im1, 0)*dt/2.0 ``` ``` return psv ``` ``` def acceleration(self, x, z): ``` ``` """ ``` ``` Force is some function of phase space vector. As long as force is a ``` ``` function only of position then area conserving. ``` ``` """ ``` ``` raise NotImplementedError("Abstract base class of Transport") ``` ```class ConstantNonLinearFocusing(Lattice): ``` ``` """ ``` ``` Class to do focussing that is constant polynomial ``` ``` i.e. force(x, z) = a_0 x + a_1 x^2 + a_2 x^3 +... ``` ``` """ ``` ``` def __init__(self, step_size, n_steps, focusing_strength_vector): ``` ``` """ ``` ``` Instantiate with step size, number of steps and focusing strength ``` ``` polynomial coefficients ``` ``` """ ``` ``` super().__init__(step_size, n_steps) ``` ``` self.f = focusing_strength_vector ``` ``` def acceleration(self, x, z): ``` ``` """ ``` ``` Acceleration as a function of x. ``` ``` """ ``` ``` xpow = 1 # x, x^2, x^3, etc ``` ``` acc = 0.0 ``` ``` for f_i in self.f: ``` ``` xpow *= x ``` ``` acc += f_i*xpow ``` ``` return acc ``` ```class Scraper(Lattice): ``` ``` def __init__(self, aperture): ``` ``` """ ``` ``` Scraper kills particles having |x| > aperture ``` ``` """ ``` ``` super().__init__(1, 1) ``` ``` self.aperture = aperture ``` ``` def transport_one(self, psv): ``` ``` """ ``` ``` Kills the particles. Returns None if position is greater than the ``` ``` aperture or psv. ``` ``` """ ``` ``` if abs(psv[0]) > self.aperture: ``` ``` return None ``` ``` return psv ``` ```class Damping(Lattice): ``` ``` def __init__(self, x_damping, v_damping): ``` ``` """Damps the particle.""" ``` ``` super().__init__(1, 1) ``` ``` self.x_damping = x_damping ``` ``` self.v_damping = v_damping ``` ``` def transport_one(self, psv): ``` ``` """Damps by x_damping, v_damping for each particle.""" ``` ``` psv[0] *= self.x_damping ``` ``` psv[1] *= self.v_damping ``` ``` return psv ``` ```class Distribution(object): ``` ``` def __init__(self): ``` ``` """Distribution just holds a collection of particles.""" ``` ``` self.psv_list = [] ``` ``` def __deepcopy__(self, memo): ``` ``` """Deepcopy - allocate new memory for the psv_list""" ``` ``` my_copy = Distribution() ``` ``` my_copy.psv_list = copy.deepcopy(self.psv_list, memo) ``` ``` return my_copy ``` ``` def __str__(self): ``` ``` """Return a string representation of the distribution - x v one psv per line""" ``` ``` my_str = "" ``` ``` for psv in self.psv_list: ``` ``` my_str += format(psv[0], "+8.4f")+format(psv[1], "+9.4f")+"\n" ``` ``` return my_str ``` ``` def plot(self, axes, vmin, vmax): ``` ``` """ ``` ``` Make a 2d histogram of the particles. axes is the matplotlib axes to draw on ``` ``` vmin, vmax are the minimum/maximum vertical axes (set to None to automatically calculate) ``` ``` """ ``` ``` x_y = list(zip(*self.psv_list)) ``` ``` clr = axes.hist2d(x_y[0], x_y[1], 50, range=[[-2, 2],[-2, 2]], vmin=vmin, vmax=vmax)[3] ``` ``` axes.get_figure().colorbar(clr) ``` ``` @classmethod ``` ``` def gen_gaussian_distribution(cls, sample_size, covariance): ``` ``` """ ``` ``` Generate a multivariatae gaussian distirbution ``` ``` """ ``` ``` dist = Distribution() ``` ``` dist.psv_list = numpy.random.multivariate_normal([0.0, 0.0], ``` ``` covariance, ``` ``` sample_size) ``` ``` return dist ``` ``` def get_amplitude_list(self, psv_list, cov = None): ``` ``` """ ``` ``` Calculate a set of amplitudes from the psv_list. If cov is defined, ``` ``` use this covariance matrix to calculate the amplitudes. ``` ``` Returns a tuple of the (list of amplitudes, covariance matrix). ``` ``` """ ``` ``` if cov is None: ``` ``` cov = numpy.cov(psv_list, rowvar=False) ``` ``` cov /= numpy.linalg.det(cov)**(1.0/cov.shape[0]) ``` ``` cov_inv = numpy.linalg.inv(cov) ``` ``` amp_list = [] ``` ``` for psv in psv_list: ``` ``` amplitude = numpy.dot(psv, cov_inv) ``` ``` amplitude = numpy.dot(amplitude, numpy.transpose(psv)) ``` ``` amp_list.append(amplitude) ``` ``` return amp_list, cov ``` ``` def amplitude(self, amp_cut): ``` ``` """ ``` ``` Calculate amplitudes. amp_cut is a float that sets the maximum amplitude ``` ``` for applying the cut. We iterate on amp_cut until no events are included ``` ``` in covariance matrix calculation having amplitude > amp_cut. ``` ``` """ ``` ``` test_list = numpy.array(self.psv_list) ``` ``` sample_size = len(test_list) ``` ``` cut_list = [0] ``` ``` cov = None ``` ``` while amp_cut and len(cut_list) > 0: ``` ``` amp_list, cov = self.get_amplitude_list(test_list) ``` ``` cut_list = [i for i, amp in enumerate(amp_list) if amp > amp_cut] ``` ``` test_list = numpy.delete(test_list, cut_list, 0) ``` ``` amp_list, cov = self.get_amplitude_list(numpy.array(self.psv_list), cov) ``` ``` return amp_list, cov ``` ```def plot_phase_space_trajectory(psv, lattice, n_steps, axes): ``` ``` """ ``` ``` Plot the trajectory in phase space for a particle traversing *lattice*. Plot ``` ``` n_steps number of iterations of lattice transport. Axes is the matplotlib axes ``` ``` used for tracking. ``` ``` """ ``` ``` x_list = [psv[0]]+[None]*n_steps ``` ``` y_list = [psv[1]]+[None]*n_steps ``` ``` for i in range(n_steps): ``` ``` lattice.transport_one(psv) ``` ``` x_list[i+1] = psv[0] ``` ``` y_list[i+1] = psv[1] ``` ``` axes.scatter(x_list, y_list, s=0.1) ``` ```def plot_distribution(lattice_list_1, contour_element, dist_in, n_turns, title, amplitude_cut): ``` ``` """ ``` ``` Transport a distribution a number of turns and then plot it. Make a 2D ``` ``` histogram of x v on the left and an amplitude distribution on the right. ``` ``` - lattice_list is used for tracking the distribution ``` ``` - contour_element is a single element that is used for plotting contours ``` ``` - dist_in is the distribution ``` ``` - n_turns is the number of turns to track distribution before plotting ``` ``` - title is a string title for the plot ``` ``` - amplitude_cut is the amplitude cut to use when calculating amplitudes ``` ``` """ ``` ``` figure = matplotlib.pyplot.figure(figsize=(20,10)) ``` ``` axes1 = figure.add_subplot(1, 2, 1) ``` ``` n_events = len(dist_in.psv_list) ``` ``` dist_out = copy.deepcopy(dist_in) ``` ``` for i in range(n_turns): ``` ``` for element in lattice_list_1: ``` ``` element.transport(dist_out) ``` ``` dist_out.plot(axes1, None, None) ``` ``` plot_phase_space_trajectory([0.1, 0], contour_element, 1000, axes1) ``` ``` plot_phase_space_trajectory([0.75, 0], contour_element, 1000, axes1) ``` ``` plot_phase_space_trajectory([1.249, 0], contour_element, 1000, axes1) ``` ``` axes1.get_figure().suptitle(title+"\n"+str(n_turns)+" turns") ``` ``` axes1.set_xlabel("Position [au]") ``` ``` axes1.set_ylabel("Momentum [au]") ``` ``` axes2 = axes1.get_figure().add_subplot(1, 2, 2) ``` ``` amplitude_list, cov = dist_out.amplitude(amplitude_cut) ``` ``` axes2.hist(amplitude_list, bins=50, range=[0.0, 2.0]) ``` ``` axes2.set_title(str(len(amplitude_list))+" events") ``` ``` axes2.set_xlabel("Amplitude [au]") ``` ``` x_list = [i/100. for i in range(401)] ``` ``` chi2_list = [2*chi2*n_events*2.0/50 for chi2 in scipy.stats.chi2.pdf(x_list, 2)] ``` ``` x_list = [x/2.0 for x in x_list] ``` ``` axes2.plot(x_list, chi2_list) ``` ``` return figure ``` ```def savefig(figure, title, n_turns): ``` ``` """ ``` ``` Save figure, using title and n_turns with special characters removed ``` ``` """ ``` ``` fname = title.replace(",", "") ``` ``` fname = fname.replace("\n", "") ``` ``` fname = fname.replace(" ", "_") ``` ``` fname += "_"+str(n_turns)+"_turns" ``` ``` figure.savefig(fname+".png") ``` ```def distributions(title, focusing, damping, aperture, amplitude_cut): ``` ``` """ ``` ``` Generate a distribution and lattice. Then track the distribution and plot it ``` ``` using plot_distribution. ``` ``` """ ``` ``` suptitle = title+"Focusing: "+str(focusing)+" Damping: "+str(damping)+" Aperture: "+str(aperture)+" Amplitude cut: "+str(amplitude_cut) ``` ``` lattice = ConstantNonLinearFocusing(1, 1, focusing)#, 0.0, 0.06]) ``` ``` damping = Damping(*damping) ``` ``` scraper = Scraper(aperture) ``` ``` distribution = Distribution.gen_gaussian_distribution(100000, [[0.5, 0.0], [0.0, 0.5]]) ``` ``` for n_turns in [0, 1, 10]: ``` ``` figure = plot_distribution([lattice, scraper, damping], lattice, distribution, n_turns, suptitle, amplitude_cut) ``` ``` savefig(figure, title, n_turns) ``` ```def main(): ``` ``` """ ``` ``` Main function ``` ``` """ ``` ``` distributions("Basic lattice\n", [-0.1], (1.0, 1.0), 100.0, 0.5) ``` ``` distributions("Scraping\n", [-0.1], (1.0, 1.0), 1.8, 0.5) ``` ``` distributions("Damping\n", [-0.1], (1.0, 0.9), 100.0, 0.5) ``` ``` distributions("Nonlinear\n", [-0.1, 0.0, 0.03], (1.0, 1.0), 100.0, 0.05) ``` ``` distributions("Scraping, damping, nonlinear\n", [-0.1, 0.0, 0.01], (1.0, 0.9), 1.0, 0.5) ``` ``` distributions("Scraping, nonlinear\n", [-0.1, 0.0, 0.01], (1.0, 1.0), 1.0, 0.5) ``` ```if __name__ == "__main__": ``` ``` main() ``` ``` matplotlib.pyplot.show(block=False) ``` ``` input("Press to finish") ```
(4-4/4)