diff --git a/tracking/particles.py b/tracking/particles.py index 94d525a47503c5bcb10e3276491bf01952b1d16c..e16847f6a7975deb5feebe21df463005651b7e79 100644 --- a/tracking/particles.py +++ b/tracking/particles.py @@ -10,6 +10,7 @@ import numpy as np import pandas as pd from tracking.parallel import Mpi from scipy.constants import c, m_e, m_p, e +import matplotlib.pyplot as plt class Particle: """ @@ -271,6 +272,53 @@ class Bunch: self.particles["tau"] = values[:,4] self.particles["delta"] = values[:,5] + def binning(self, dimension="tau", n_bin=75): + """ + Bin macro-particles. + + Parameters + ---------- + dimension : str, optional + Dimension in which the binning is done. The default is "tau". + n_bin : int, optional + Number of bins. The default is 75. + + Returns + ------- + bins : IntervalIndex of shape (n_bin,) + Bins where the particles are sorted. + sorted_index : Series of shape (self.mp_number,) + Bin number of each macro-particles. + profile : Series of shape (n_bin,) + Number of marco-particles in each bin. + + """ + bins = pd.interval_range(start=min(self[dimension]), + end=max(self[dimension]), + periods=n_bin, + name=dimension) + sorted_index = pd.cut(self[dimension], bins) + sorted_index = sorted_index.reindex(self.alive.index) + profile = sorted_index.value_counts().sort_index() + return (bins, sorted_index, profile) + + def plot_profile(self, dimension="tau", n_bin=75): + """ + Plot bunch profile. + + Parameters + ---------- + dimension : str, optional + Dimension to plot. The default is "tau". + n_bin : int, optional + Number of bins. The default is 75. + + """ + bins, sorted_index, profile = self.binning(dimension, n_bin) + fig = plt.figure() + ax = fig.gca() + ax.plot(bins.mid, profile) + class Beam: """ Define a Beam object composed of several Bunch objects. @@ -355,12 +403,29 @@ class Beam: @property def not_empty(self): - """Return a generator to iterate over not empty bunches""" - for bunch in self: - if bunch.current == 0: - pass + """Return a generator to iterate over not empty bunches.""" + for index, value in enumerate(self.filling_pattern): + if value == True: + yield self[index] else: - yield bunch + pass + + @property + def distance_between_bunches(self): + filling_pattern = self.filling_pattern + distance = np.zeros(filling_pattern.shape) + for index, value in enumerate(filling_pattern): + if value == False: + pass + elif value == True: + count = 1 + for value2 in filling_pattern[index+1:]: + if value2 == False: + count += 1 + elif value2 == True: + break + distance[index] = count + return distance def init_beam(self, filling_pattern, current_per_bunch=1e-3, mp_per_bunch=1e3): @@ -410,17 +475,23 @@ class Beam: for bunch in self.not_empty: bunch.init_gaussian() + + self.update_filling_pattern() - @property - def filling_pattern(self): - """Return an array with the filling pattern of the beam as bool""" + def update_filling_pattern(self): + """Update the beam filling pattern.""" filling_pattern = [] for bunch in self: if bunch.current != 0: filling_pattern.append(True) else: filling_pattern.append(False) - return np.array(filling_pattern) + self._filling_pattern = np.array(filling_pattern) + + @property + def filling_pattern(self): + """Return an array with the filling pattern of the beam as bool""" + return self._filling_pattern @property def bunch_current(self): diff --git a/tracking/synchrotron.py b/tracking/synchrotron.py index 038248540226a1ad46854473f32f2a10665de7c0..a622c0458f5ba5efc77ac2492c3f0dc4f245223b 100644 --- a/tracking/synchrotron.py +++ b/tracking/synchrotron.py @@ -53,6 +53,8 @@ class Synchrotron: Revolution frequency in [Hz]. omega0 : float Angular revolution frequency in [Hz.rad] + T1 : flaot + Fundamental RF period in [s]. f1 : float Fundamental RF frequency in [Hz]. omega1 : float @@ -112,6 +114,7 @@ class Synchrotron: def L(self,value): self._L = value self._T0 = self.L/c + self._T1 = self.T0/self.h self._f0 = 1/self.T0 self._omega0 = 2*np.pi*self.f0 self._f1 = self.h*self.f0 @@ -127,6 +130,15 @@ class Synchrotron: def T0(self, value): self.L = c*value + @property + def T1(self): + """"Fundamental RF period [s]""" + return self._T1 + + @T1.setter + def T1(self, value): + self.L = c*value*self.h + @property def f0(self): """Revolution frequency [Hz]"""