From 53b215cbeafe7a3d4595b3a4d1a50b5164f7ae66 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <gamelin@synchrotron-soleil.fr>
Date: Fri, 14 Feb 2020 20:06:41 +0100
Subject: [PATCH] Add bunch_X properties to Beam class
Add bunch_X properties to Beam class
Add a generator over "not empty" bunches to Beam class
Modify Bunch class in order to have mean, std and emit as property and not method
---
Tracking/beam.py | 115 ++++++++++++++++++++++++++++++++---------------
1 file changed, 79 insertions(+), 36 deletions(-)
diff --git a/Tracking/beam.py b/Tracking/beam.py
index 1ffede3..2b0e039 100644
--- a/Tracking/beam.py
+++ b/Tracking/beam.py
@@ -8,7 +8,6 @@ Beam and bunch elements
import numpy as np
import pandas as pd
-import warnings
class Bunch:
"""
@@ -38,10 +37,7 @@ class Bunch:
Bunch current in [A]
Methods
- ------
- stats()
- Compute the mean and the standard deviation of alive particles for each
- coordinates.
+ -------
init_gaussian(cov=None, mean=None, **kwargs)
Initialize bunch particles with 6D gaussian phase space.
"""
@@ -132,7 +128,8 @@ class Bunch:
@current.setter
def current(self, value):
self.charge_per_mp = value * self.ring.T0 / self.__len__()
-
+
+ @property
def mean(self):
"""
Compute the mean position of alive particles for each
@@ -146,6 +143,7 @@ class Bunch:
mean = [[self[name].mean()] for name in self]
return np.array(mean)
+ @property
def std(self):
"""
Compute the standard deviation of the position of alive
@@ -158,7 +156,8 @@ class Bunch:
"""
std = [[self[name].std()] for name in self]
return np.array(std)
-
+
+ @property
def emit(self):
"""
Compute the bunch emittance for each plane [1].
@@ -245,7 +244,7 @@ class Beam:
Total bunch charge in [C]
particle_number : int
Total number of particle in the beam
- filling_pattern : list of bool
+ filling_pattern : array of bool
Filling pattern of the beam
Methods
@@ -270,9 +269,8 @@ class Beam:
def __len__(self):
"""Return the number of (not empty) bunches"""
length = 0
- for bunch in self:
- if bunch.current != 0:
- length += 1
+ for bunch in self.not_empty:
+ length += 1
return length
def __getitem__(self, i):
@@ -284,8 +282,17 @@ class Beam:
self.bunch_list.__setitem__(i, value)
def __iter__(self):
- """Iterate over the bunches"""
+ """Iterate over all bunches"""
return self.bunch_list.__iter__()
+
+ @property
+ def not_empty(self):
+ """Return a generator to iterate over not empty bunches"""
+ for bunch in self:
+ if bunch.current == 0:
+ pass
+ else:
+ yield bunch
def init_beam(self, filling_pattern, current_per_bunch=1e-3,
mp_per_bunch=1e3):
@@ -333,41 +340,77 @@ class Beam:
self.bunch_list = bunch_list
- for bunch in self:
+ for bunch in self.not_empty:
bunch.init_gaussian()
+
+ @property
+ def filling_pattern(self):
+ """Return an array with the filling pattern of the beam as bool"""
+ filling_pattern = []
+ for bunch in self:
+ if filling_pattern != 0:
+ filling_pattern.append(True)
+ else:
+ filling_pattern.append(False)
+ return np.array(filling_pattern)
+ @property
+ def bunch_current(self):
+ """Return an array with the current in each bunch in [A]"""
+ bunch_current = [bunch.current for bunch in self]
+ return np.array(bunch_current)
+
+ @property
+ def bunch_charge(self):
+ """Return an array with the charge in each bunch in [C]"""
+ bunch_charge = [bunch.charge for bunch in self]
+ return np.array(bunch_charge)
+
+ @property
+ def bunch_particle(self):
+ """Return an array with the particle number in each bunch"""
+ bunch_particle = [bunch.particle_number for bunch in self]
+ return np.array(bunch_particle)
+
@property
def current(self):
"""Total beam current in [A]"""
- current = 0
- for bunch in self:
- current += bunch.current
- return current
+ return np.sum(self.bunch_current)
@property
def charge(self):
"""Total beam charge in [C]"""
- charge = 0
- for bunch in self:
- charge += bunch.charge
- return charge
+ return np.sum(self.bunch_charge)
@property
def particle_number(self):
- """Total number of particle in the beam"""
- particle_number = 0
- for bunch in self:
- particle_number += bunch.particle_number
- return particle_number
+ """Total number of particles in the beam"""
+ return np.sum(self.bunch_particle)
@property
- def filling_pattern(self):
- """Return an array with the filling pattern of the beam as bool"""
- filling_pattern = []
- for bunch in self:
- if filling_pattern != 0:
- filling_pattern.append(True)
- else:
- filling_pattern.append(False)
- return np.array(filling_pattern)
-
+ def bunch_mean(self):
+ """Return an array with the mean position of alive particles for each
+ bunches"""
+ bunch_mean = np.zeros((6,self.ring.h))
+ for index, bunch in enumerate(self):
+ bunch_mean[:,index] = np.squeeze(bunch.mean)
+ return bunch_mean
+
+ @property
+ def bunch_std(self):
+ """Return an array with the standard deviation of the position of alive
+ particles for each bunches"""
+ bunch_std = np.zeros((6,self.ring.h))
+ for index, bunch in enumerate(self):
+ bunch_std[:,index] = np.squeeze(bunch.std)
+ return bunch_std
+
+ @property
+ def bunch_emit(self):
+ """Return an array with the bunch emittance of alive particles for each
+ bunches and each plane"""
+ bunch_emit = np.zeros((3,self.ring.h))
+ for index, bunch in enumerate(self):
+ bunch_emit[:,index] = np.squeeze(bunch.emit)
+ return bunch_emit
+
\ No newline at end of file
--
GitLab