Add math label for bunching

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
@@ -750,7 +750,7 @@ def __getitem__(self, key):
         elif key.startswith('ptp_'):
             return self.ptp(key[4:])   
         elif key.startswith('bunching_'):
-            wavelength = statistics.parse_bunching_str(key)
+            wavelength = parse_bunching_str(key)
             return self.bunching(wavelength)
 
         else:

pmd_beamphysics/statistics.py

@@ -569,42 +569,6 @@ def bunching(z: np.ndarray, wavelength: float, weight: np.ndarray = None) -> flo
     f = np.exp(1j * k * z)
     return np.abs(np.sum(weight * f)) / np.sum(weight)
 
-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
 
 
 

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)
+
+
+