Merge pull request #48 from ChristopherMayes/add_bunching

Add bunching

mkdocs.yml

@@ -12,6 +12,7 @@ nav:
             - examples/read_examples.ipynb
             - examples/write_examples.ipynb
             - examples/plot_examples.ipynb
+            - examples/bunching.ipynb
         - Fields:
             - examples/fields/field_examples.ipynb
             - examples/fields/field_expansion.ipynb

pmd_beamphysics/labels.py

@@ -1,4 +1,4 @@
-from pmd_beamphysics.units import pg_units
+from pmd_beamphysics.units import pg_units, parse_bunching_str, nice_array
 
 
 TEXLABEL = {
@@ -119,6 +119,11 @@ def texlabel(key: str):
         tex1 = texlabel(subkeys[1])
         return fr'\left<{tex0}, {tex1}\right>'
 
+    if key.startswith('bunching'):
+        wavelength = parse_bunching_str(key)
+        x, _, prefix = nice_array(wavelength)
+        return f'\mathrm{{bunching~at}}~{x:.1f}~\mathrm{{ {prefix}m }}'
+
     return None
 
 

pmd_beamphysics/particles.py

@@ -1,4 +1,4 @@
-from pmd_beamphysics.units import dimension, dimension_name, SI_symbol, pg_units, c_light
+from pmd_beamphysics.units import dimension, dimension_name, SI_symbol, pg_units, c_light, parse_bunching_str
 
 from pmd_beamphysics.interfaces.astra import write_astra
 from pmd_beamphysics.interfaces.bmad import write_bmad
@@ -17,6 +17,8 @@
 from pmd_beamphysics.species import charge_of, mass_of
 
 from pmd_beamphysics.statistics import norm_emit_calc, normalized_particle_coordinate, particle_amplitude, particle_twiss_dispersion, matched_particles, resample_particles, slice_statistics
+import pmd_beamphysics.statistics as statistics 
+
 from pmd_beamphysics.writers import write_pmd_bunch, pmd_init
 
 from h5py import File
@@ -675,6 +677,48 @@ def average_current(self):
             dt = self.z.ptp() / (self.avg('beta_z')*c_light)
         return self.charge / dt
 
+    def bunching(self, wavelength):
+        """
+        Calculate the normalized bunching parameter, which is the magnitude of the 
+        complex sum of weighted exponentials at a given point.
+    
+        The formula for bunching is given by:
+    
+        $$
+        B(z, \lambda) = \frac{\left|\sum w_i e^{i k z_i}\right|}{\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.
+    
+        Parameters
+        ----------
+        wavelength : float
+            Wavelength of the wave.
+
+    
+        Returns
+        -------
+        float
+            The normalized bunching parameter.
+    
+        Raises
+        ------
+        ValueError
+            If `wavelength` is not a positive number.
+        """        
+
+        if self.in_z_coordinates:
+            # Approximate z
+            z = self.t * self.avg('beta_z')*c_light
+        else:
+            z = self.z
+
+        return statistics.bunching(z, wavelength, weight=self.weight)
+
     def __getitem__(self, key):
         """
         Returns a property or statistical quantity that can be computed:
@@ -705,7 +749,10 @@ def __getitem__(self, key):
         elif key.startswith('max_'):
             return self.max(key[4:])     
         elif key.startswith('ptp_'):
-            return self.ptp(key[4:])         
+            return self.ptp(key[4:])   
+        elif key.startswith('bunching_'):
+            wavelength = parse_bunching_str(key)
+            return self.bunching(wavelength)
 
         else:
             return getattr(self, key) 

pmd_beamphysics/statistics.py

@@ -521,4 +521,54 @@ 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.
+
+    The formula for bunching is given by:
+
+    $$
+    B(z, \lambda) = \frac{\left|\sum w_i e^{i k z_i}\right|}{\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.
+
+    Parameters
+    ----------
+    z : np.ndarray
+        Array of positions where the bunching parameter is calculated.
+    wavelength : float
+        Wavelength of the wave.
+    weight : np.ndarray, optional
+        Weights for each exponential term. Default is 1 for all terms.
+
+    Returns
+    -------
+    float
+        The normalized bunching parameter.
+
+    Raises
+    ------
+    ValueError
+        If `wavelength` is not a positive number.
+    """
+    if wavelength <= 0:
+        raise ValueError("Wavelength must be a positive number.")
+
+    if weight is None:
+        weight = np.ones(len(z))
+    if len(weight) != len(z):
+        raise ValueError(f"Weight array has length {len(weight)} != length of the z array, {len(z)}")         
+
+    k = 2 * np.pi / wavelength
+    f = np.exp(1j * k * z)
+    return np.abs(np.sum(weight * f)) / np.sum(weight)
+
+
+
 

pmd_beamphysics/units.py

@@ -442,7 +442,7 @@ def plottable_array(x, nice=True, lim=None):
     PARTICLEGROUP_UNITS[k] = unit('eV/c')
 for k in ['x', 'y', 'z', 'r', 'Jx', 'Jy']:
     PARTICLEGROUP_UNITS[k] = unit('m')
-for k in ['beta', 'beta_x', 'beta_y', 'beta_z', 'gamma']:    
+for k in ['beta', 'beta_x', 'beta_y', 'beta_z', 'gamma', 'bunching']:    
     PARTICLEGROUP_UNITS[k] = unit('1')
 for k in ['theta']:    
     PARTICLEGROUP_UNITS[k] = unit('rad')
@@ -504,10 +504,53 @@ def pg_units(key):
         return unit('V/m')  
     if key.startswith('magneticField'):
         return unit('T')
+    if key.startswith('bunching_'):
+        return unit('1')
 
 
     raise ValueError(f'No known unit for: {key}')
 
+
+# -------------------------
+# Special parsers
+
+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' 
+        
+    Returns
+    -------
+    wavelength: float
+    
+    """
+    assert s.startswith('bunching_')
+
+    x = s.split('_')
+
+    wavelength = float(x[1])
+
+    if len(x) == 2:
+        factor = 1
+    elif len(x) == 3:
+        unit = x[2]
+        if unit == 'm':
+            factor = 1
+        elif unit == 'mm':
+            factor = 1e-3
+        elif unit == 'µm' or unit == 'um':
+            factor = 1e-6
+        elif unit == 'nm':
+            factor = 1e-9
+        else:
+            raise ValueError(f'Unparsable unit: {unit}')
+    else:
+        raise ValueError(f'Cannot parse {s}')
+
+
+    return wavelength * factor
+
 
 
 # -------------------------
@@ -591,6 +634,9 @@ def write_dataset_and_unit_h5(h5, name, data, unit=None):
 
     if unit:
         write_unit_h5(h5[name], unit)
+
+
+