Merge pull request #50 from ChristopherMayes/bunching_phase

Add bunching_phase

pmd_beamphysics/particles.py

@@ -702,7 +702,7 @@ def bunching(self, wavelength):
     
         Returns
         -------
-        float
+        complex
             The normalized bunching parameter.
     
         Raises
@@ -750,9 +750,15 @@ def __getitem__(self, key):
             return self.max(key[4:])     
         elif key.startswith('ptp_'):
             return self.ptp(key[4:])   
-        elif key.startswith('bunching_'):
+        elif 'bunching' in key:
             wavelength = parse_bunching_str(key)
-            return self.bunching(wavelength)
+            bunching = self.bunching(wavelength) # complex
+
+            # abs or arg (angle):
+            if 'phase_' in key:
+                return np.angle(bunching)
+            else:
+                return np.abs(bunching)
 
         else:
             return getattr(self, key) 

pmd_beamphysics/statistics.py

@@ -521,22 +521,22 @@ def resample_particles(particle_group, n=0):
     return data
 
 
-def bunching(z: np.ndarray, wavelength: float, weight: np.ndarray = None) -> float:
-    """
-    Calculate the normalized bunching parameter, which is the magnitude of the 
-    complex sum of weighted exponentials at a given point.
+def bunching(z: np.ndarray, wavelength: float, weight: np.ndarray = None) -> complex:
+    r"""
+    Calculate the normalized bunching parameter, which is the
+    complex sum of weighted exponentials.
 
     The formula for bunching is given by:
 
     $$
-    B(z, \lambda) = \frac{\left|\sum w_i e^{i k z_i}\right|}{\sum w_i}
+    B(z, \lambda) = \frac{\sum w_i e^{i k z_i}}{\sum w_i}
     $$
 
     where:
-    - \( z \) is the position array,
-    - \( \lambda \) is the wavelength,
-    - \( k = \frac{2\pi}{\lambda} \) is the wave number,
-    - \( w_i \) are the weights.
+    - $z$ is the position array,
+    - $\lambda$ is the wavelength,
+    - $k = \frac{2\pi}{\lambda}$  is the wave number,
+    - $w_i$ are the weights.
 
     Parameters
     ----------
@@ -549,8 +549,8 @@ def bunching(z: np.ndarray, wavelength: float, weight: np.ndarray = None) -> flo
 
     Returns
     -------
-    float
-        The normalized bunching parameter.
+    complex
+        The bunching parameter
 
     Raises
     ------
@@ -567,7 +567,7 @@ def bunching(z: np.ndarray, wavelength: float, weight: np.ndarray = None) -> flo
 
     k = 2 * np.pi / wavelength
     f = np.exp(1j * k * z)
-    return np.abs(np.sum(weight * f)) / np.sum(weight)
+    return np.sum(weight * f) / np.sum(weight)
 
 
 

pmd_beamphysics/units.py

@@ -444,7 +444,7 @@ def plottable_array(x, nice=True, lim=None):
     PARTICLEGROUP_UNITS[k] = unit('m')
 for k in ['beta', 'beta_x', 'beta_y', 'beta_z', 'gamma', 'bunching']:    
     PARTICLEGROUP_UNITS[k] = unit('1')
-for k in ['theta']:    
+for k in ['theta', 'bunching_phase']:    
     PARTICLEGROUP_UNITS[k] = unit('rad')
 for k in ['charge', 'species_charge', 'weight']:
     PARTICLEGROUP_UNITS[k] = unit('C')
@@ -504,7 +504,9 @@ def pg_units(key):
         return unit('V/m')  
     if key.startswith('magneticField'):
         return unit('T')
-    if key.startswith('bunching_'):
+    if key.startswith('bunching_phase'):
+        return unit('rad')    
+    if key.startswith('bunching'):
         return unit('1')
 
 
@@ -519,6 +521,7 @@ def parse_bunching_str(s):
     Parse a string of the on of the forms to extract the wavelength:
         'bunching_1.23e-4' 
         'bunching_1.23e-4_nm' 
+        'bunching_phase_1.23e-4' 
         
     Returns
     -------
@@ -527,14 +530,18 @@ def parse_bunching_str(s):
     """
     assert s.startswith('bunching_')
 
+    # Remove bunching and phase prefixes
+    s = s.replace('bunching_', '')
+    s = s.replace('phase_', '')
+
     x = s.split('_')
 
-    wavelength = float(x[1])
+    wavelength = float(x[0])
 
-    if len(x) == 2:
+    if len(x) == 1:
         factor = 1
-    elif len(x) == 3:
-        unit = x[2]
+    elif len(x) == 2:
+        unit = x[1]
         if unit == 'm':
             factor = 1
         elif unit == 'mm':