Merge branch 'master' into allow_species_in_paths

pmd_beamphysics/__init__.py

@@ -1,8 +1,17 @@
-from .particles import ParticleGroup,  single_particle
+from .particles import ParticleGroup, single_particle
 from .status import ParticleStatus
 from .fields.fieldmesh import FieldMesh
 from .readers import particle_paths
 from .writers import pmd_init
 
 from . import _version
 __version__ = _version.get_versions()['version']
+
+__all__ = [
+    "FieldMesh",
+    "ParticleGroup",
+    "ParticleStatus",
+    "particle_paths",
+    "pmd_init",
+    "single_particle",
+]

pmd_beamphysics/fields/analysis.py

@@ -122,7 +122,6 @@ def track_field_1d(z,
 
     # Make interpolating function
     field = interpolate.interp1d(z, Ez * q0 * c_light, fill_value='extrapolate')
-    zmin = z.min()
     zmax = z.max()
     tmax = 100/frequency
     omega = 2*np.pi*frequency
@@ -470,4 +469,4 @@ def phase_scale_f(phase_deg, scale):
     if verbose:
         print(f'    Pass 2 delta energy: {phase_scale_f(phase2, scale2)} at phase  {phase2} deg')
 
-    return phase2, scale2
+    return phase2, scale2

pmd_beamphysics/fields/expansion.py

@@ -51,7 +51,6 @@ def expand_1d_dynamic_fieldmap(z, Ez0, frequency=0, spline_s=0):
 
     omega = 2*np.pi*frequency
     c_light=299792458.
-    ooc2 = (omega/c_light)**2 
 
     # Make spline and derivatives
     S = UnivariateSpline(z, Ez0, k=5, s=spline_s)
@@ -368,4 +367,4 @@ def expand_fieldmesh_from_onaxis(fieldmesh, *,
     attrs['gridSize'] = (nr, 1, nz)
     attrs['gridSpacing'] = (dr, 0, dz)
 
-    return fieldmesh
+    return fieldmesh

pmd_beamphysics/interfaces/astra.py

@@ -258,7 +258,6 @@ def astra_1d_fieldmap_data(fm):
 
     assert fm.shape[1] == 1, 'Cylindrical symmetry required'
 
-    dat = {}
     z = fm.coord_vec('z')
 
     if fm.is_static:

pmd_beamphysics/interfaces/bmad.py

@@ -150,8 +150,6 @@ def write_bmad(particle_group,
     status  = particle_group.status 
     weight  = particle_group.weight
 
-    zeros = np.full(n, 0)
-
     if p0c:
         # s-based coordinates
 

pmd_beamphysics/interfaces/gpt.py

@@ -75,7 +75,7 @@ def run_asci2gdf(outfile, asci2gdf_bin, verbose=False):
     cmd = [asci2gdf_bin, '-o', outfile, tempfile]
     if verbose:
         print(' '.join(cmd))
-    result = subprocess.run(cmd)
+    subprocess.run(cmd)
 
     # Cleanup
     os.remove(tempfile)    

pmd_beamphysics/interfaces/impact.py

@@ -1,11 +1,12 @@
+import os
 import numpy as np
 from numpy import pi
 from scipy import fft
 
-c_light = 299792458.
-
 from pmd_beamphysics.fields.expansion import spline_derivative_array
 
+c_light = 299792458.
+
 
 
 
@@ -627,7 +628,7 @@ def create_impact_solrf_fieldmap_derivatives(field_mesh,
     rfdata = []
 
     z = field_mesh.coord_vec('z')
-    zmin, zmax, L = z.min(), z.max(), z.ptp()
+    L = z.ptp()
 
     info = {'format':'solrf'}
     field = {}    

pmd_beamphysics/interfaces/lucretia.py

@@ -217,4 +217,4 @@ def list_element_names(filePath):
 
     """
     dat = sio.loadmat(filePath)
-    return list(dat['bstore'].dtype.fields)
+    return list(dat['bstore'].dtype.fields)

pmd_beamphysics/interfaces/simion.py

@@ -73,15 +73,20 @@ def write_simion(particle_group, outfile, verbose=0, color=0, flip_z_to_x=True):
     header=';0'
 
     simion_params= ['TOB', 'MASS', 'CHARGE', 'X', 'Y', 'Z', 'AZ', 'EL', 'KE', 'CWF', 'COLOR']
-
-    simion_units = {'TOB':'usec', 
-                    'MASS':'amu', 
-                    'CHARGE':'e', 
-                    'X':'mm', 'Y':'mm', 'Z':'mm', 
-                    'AZ':'deg', 'EL':'deg', 
-                    'CWF':'', 
-                    'COLOR':''}
-
+
+    # simion_units = {
+    #     "TOB": "usec",
+    #     "MASS": "amu",
+    #     "CHARGE": "e",
+    #     "X": "mm",
+    #     "Y": "mm",
+    #     "Z": "mm",
+    #     "AZ": "deg",
+    #     "EL": "deg",
+    #     "CWF": "",
+    #     "COLOR": "",
+    # }
+
     N = len(particle_group)
 
     data = np.zeros( (N, len(simion_params)) )

pmd_beamphysics/interfaces/superfish.py

@@ -277,7 +277,7 @@ def read_superfish_t7(filename,
         # FISH problem   
         zmin, zmax, nz =  float(line1[0])*1e-2, float(line1[1])*1e-2, int(line1[2])+1
         frequency = float(line2[0])*1e6 # MHz -> Hz
-        rmin, rmax, nr =  float(line3[0])*1e-2, float(line3[1])*1e-2, int(line3[2])+1  
+        _rmin, rmax, nr =  float(line3[0])*1e-2, float(line3[1])*1e-2, int(line3[2])+1  
 
         # Read and reshape
         # dat = np.loadtxt(filename, skiprows=3)
@@ -293,7 +293,7 @@ def read_superfish_t7(filename,
         # zmin(cm), zmax(cm), ny-1    # z in cylindrical geometry
 
         # POISSON problem
-        rmin, rmax, nr =  float(line1[0])*1e-2, float(line1[1])*1e-2, int(line1[2])+1
+        _rmin, rmax, nr =  float(line1[0])*1e-2, float(line1[1])*1e-2, int(line1[2])+1
         zmin, zmax, nz =  float(line2[0])*1e-2, float(line2[1])*1e-2, int(line2[2])+1        
         frequency=0
 

pmd_beamphysics/labels.py

@@ -36,20 +36,20 @@
  'gamma': r'\gamma',
  'theta': r'\theta',
  'charge': 'Q',
- 'twiss_alpha_x'     : r'Twiss\ '+r'\alpha_x',
- 'twiss_beta_x'      : r'Twiss\ '+r'\beta_x',   
- 'twiss_gamma_x'     : r'Twiss\ '+r'\gamma_x',   
- 'twiss_eta_x'       : r'Twiss\ '+r'\eta_x',   
- 'twiss_etap_x'      : r'Twiss\ '+r"\eta'_x",     
- 'twiss_emit_x'      : r'Twiss\ '+r'\epsilon_{x}',
- 'twiss_norm_emit_x' : r'Twiss\ '+r'\epsilon_{n, x}',
- 'twiss_alpha_y'     : r'Twiss\ '+r'\alpha_y',
- 'twiss_beta_y'      : r'Twiss\ '+r'\beta_y',   
- 'twiss_gamma_y'     : r'Twiss\ '+r'\gamma_y',   
- 'twiss_eta_y'       : r'Twiss\ '+r'\eta_y',   
- 'twiss_etap_y'      : r'Twiss\ '+r"\eta'_y",     
- 'twiss_emit_y'      : r'Twiss\ '+r'\epsilon_{y}',
- 'twiss_norm_emit_y' : r'Twiss\ '+r'\epsilon_{n, y}',    
+ 'twiss_alpha_x'     : r"Twiss\ \alpha_x",
+ 'twiss_beta_x'      : r"Twiss\ \beta_x",
+ 'twiss_gamma_x'     : r"Twiss\ \gamma_x",
+ 'twiss_eta_x'       : r"Twiss\ \eta_x",
+ 'twiss_etap_x'      : r"Twiss\ \eta'_x",     
+ 'twiss_emit_x'      : r"Twiss\ \epsilon_{x}",
+ 'twiss_norm_emit_x' : r"Twiss\ \epsilon_{n, x}",
+ 'twiss_alpha_y'     : r"Twiss\ \alpha_y",
+ 'twiss_beta_y'      : r"Twiss\ \beta_y",
+ 'twiss_gamma_y'     : r"Twiss\ \gamma_y",
+ 'twiss_eta_y'       : r"Twiss\ \eta_y",
+ 'twiss_etap_y'      : r"Twiss\ \eta'_y",     
+ 'twiss_emit_y'      : r"Twiss\ \epsilon_{y}",
+ 'twiss_norm_emit_y' : r"Twiss\ \epsilon_{n, y}",
 # 'species_charge',
 # 'weight',
  'average_current': r'I_{av}',
@@ -179,4 +179,4 @@ def mathlabel(*keys, units=None, tex=True):
         if units:
             label = fr'{label} ({units})'
 
-        return label
+        return label

pmd_beamphysics/plot.py

@@ -3,24 +3,26 @@
 
 """
 
-from pmd_beamphysics.units import nice_array, plottable_array, nice_scale_prefix
-from pmd_beamphysics.labels import mathlabel
-
+from copy import copy
 
-from .statistics import slice_statistics
-import matplotlib.pyplot as plt
 import matplotlib
-from matplotlib.gridspec import GridSpec
+import matplotlib.pyplot as plt
 import numpy as np
-from copy import copy
+from matplotlib.gridspec import GridSpec
+# For field legends
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from pmd_beamphysics.labels import mathlabel
+from pmd_beamphysics.units import (nice_array, nice_scale_prefix,
+                                   plottable_array)
+
+from .statistics import slice_statistics
+
 CMAP0 = copy(plt.get_cmap('viridis'))
 CMAP0.set_under('white')
 
 CMAP1 = copy(plt.get_cmap('plasma'))
 
-# For field legends
-from mpl_toolkits.axes_grid1 import make_axes_locatable
-
 
 def plt_histogram(a, weights=None, bins=40):
     """
@@ -574,4 +576,4 @@ def plot_fieldmesh_cylindrical_1d(fm,
     if return_figure:
         return fig        
 
-
+

pmd_beamphysics/units.py

@@ -7,15 +7,14 @@
 
 
 """
+import numpy as np
 import scipy.constants
 
 mec2 = scipy.constants.value('electron mass energy equivalent in MeV')*1e6
 mpc2 = scipy.constants.value('proton mass energy equivalent in MeV')*1e6
 c_light = scipy.constants.c
 e_charge = scipy.constants.e
 mu_0 = scipy.constants.mu_0 # Note that this is no longer 4pi*10^-7 !
-
-import numpy as np
 
 
 class pmd_unit:

pmd_beamphysics/writers.py

@@ -77,7 +77,7 @@ def write_pmd_bunch(h5, data, name=None):
         u = pg_units(key)        
 
         # Write
-        g2 = write_component_data(g, g2_name, data[key], unit=u)
+        write_component_data(g, g2_name, data[key], unit=u)
 
 
     # Optional id. This does not have any units.