diff --git a/tracking/particles.py b/tracking/particles.py index 7a97fa35b64ca960442bcb856103b0e4c3ab8722..b327da5c7fb24e70ceb7446e772eb6bd6cf6a859 100644 --- a/tracking/particles.py +++ b/tracking/particles.py @@ -97,7 +97,8 @@ class Bunch: Springer, Eq.(8.39) of p224. """ - def __init__(self, ring, mp_number=1e3, current=1e-3, alive=True): + def __init__(self, ring, mp_number=1e3, current=1e-3, alive=True, + track_alive=False): self.ring = ring if not alive: @@ -105,20 +106,22 @@ class Bunch: current = 0 self._mp_number = int(mp_number) - particles = {"x":np.zeros((self.mp_number,)), - "xp":np.zeros((self.mp_number,)), - "y":np.zeros((self.mp_number,)), - "yp":np.zeros((self.mp_number,)), - "tau":np.zeros((self.mp_number,)), - "delta":np.zeros((self.mp_number,)), - } - self.particles = pd.DataFrame(particles) - self.alive = pd.Series(np.ones((self.mp_number,),dtype=bool)) + self.dtype = np.dtype([('x',np.float), + ('xp',np.float), + ('y',np.float), + ('yp',np.float), + ('tau',np.float), + ('delta',np.float)]) + + self.particles = np.zeros(self.mp_number, self.dtype) + self.track_alive = track_alive + + self.alive = np.ones((self.mp_number,),dtype=bool) self.current = current if not alive: - self.alive = pd.Series(np.zeros((self.mp_number,),dtype=bool)) + self.alive = np.zeros((self.mp_number,),dtype=bool) - self._energy_change = pd.Series(np.zeros((self.mp_number,),dtype=float)) + self._energy_change = np.zeros((self.mp_number,),dtype=float) def __len__(self): """Return the number of alive particles""" @@ -126,15 +129,21 @@ class Bunch: def __getitem__(self, label): """Return the columns label for alive particles""" - return self.particles.loc[self.alive, label] + if self.track_alive is True: + return self.particles[label][self.alive] + else: + return self.particles[label] def __setitem__(self, label, value): """Set value to the columns label for alive particles""" - self.particles.loc[self.alive, label] = value + if self.track_alive is True: + self.particles[label][self.alive] = value + else: + self.particles[label] = value def __iter__(self): """Iterate over labels""" - return self[:].__iter__() + return self.dtype.names.__iter__() def __repr__(self): """Return representation of alive particles""" @@ -222,11 +231,17 @@ class Bunch: the SynchrotronRadiation class to compute radiation damping. To be changed by Element objects which change the energy, for example RF cavities.""" - return self._energy_change.loc[self.alive] + if self.track_alive is True: + return self._energy_change[self.alive] + else: + return self._energy_change @energy_change.setter def energy_change(self, value): - self._energy_change.loc[self.alive] = value + if self.track_alive is True: + self._energy_change[self.alive] = value + else: + self._energy_change = value def init_gaussian(self, cov=None, mean=None, **kwargs): """ @@ -297,14 +312,18 @@ class Bunch: 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) + bin_min = self[dimension].min() + bin_min = min(bin_min*0.99, bin_min*1.01) + bin_max = self[dimension].max() + bin_max = max(bin_max*0.99, bin_max*1.01) + + bins = np.linspace(bin_min, bin_max, n_bin) + center = (bins[1:] + bins[:-1])/2 + sorted_index = np.searchsorted(bins, self[dimension], side='left') + sorted_index -= 1 + profile = np.bincount(sorted_index, minlength=n_bin-1) + + return (bins, sorted_index, profile, center) def plot_profile(self, dimension="tau", n_bin=75): """ @@ -318,10 +337,10 @@ class Bunch: Number of bins. The default is 75. """ - bins, sorted_index, profile = self.binning(dimension, n_bin) + bins, sorted_index, profile, center = self.binning(dimension, n_bin) fig = plt.figure() ax = fig.gca() - ax.plot(bins.mid, profile) + ax.plot(center, profile) def plot_phasespace(self, x_var="tau", y_var="delta", plot_type="j"): """ diff --git a/tracking/wakepotential.py b/tracking/wakepotential.py index f9197bca6bf62f3d9633bb683e2a723bf7f469f7..5e31e275b34c1f1e0c6f42aa0b9f99051166880c 100644 --- a/tracking/wakepotential.py +++ b/tracking/wakepotential.py @@ -72,18 +72,19 @@ class WakePotential(Element): """ # Get binning data - a, b, c = bunch.binning(n_bin=self.n_bin) + a, b, c, d = bunch.binning(n_bin=self.n_bin) self.bins = a self.sorted_index = b self.profile = c - self.bin_size = self.bins.length + self.center = d + self.bin_size = self.bins[1] - self.bins[0] # Compute charge density self.rho = bunch.charge_per_mp*self.profile/(self.bin_size*bunch.charge) self.rho = np.array(self.rho) # Compute time array - self.tau = np.array(self.bins.mid) + self.tau = np.array(self.center) self.dtau = self.tau[1] - self.tau[0] # Add N values before and after rho and tau