Skip to content

Commit

Permalink
Merge pull request #54 from ken-lauer/fix_lint
Browse files Browse the repository at this point in the history
MAINT: Fix lint
  • Loading branch information
ChristopherMayes authored May 31, 2024
2 parents 670f9ac + 44e30d8 commit 62fe9fe
Show file tree
Hide file tree
Showing 17 changed files with 69 additions and 61 deletions.
11 changes: 10 additions & 1 deletion pmd_beamphysics/__init__.py
Original file line number Diff line number Diff line change
@@ -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",
]
3 changes: 1 addition & 2 deletions pmd_beamphysics/fields/analysis.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
3 changes: 1 addition & 2 deletions pmd_beamphysics/fields/expansion.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -368,4 +367,4 @@ def expand_fieldmesh_from_onaxis(fieldmesh, *,
attrs['gridSize'] = (nr, 1, nz)
attrs['gridSpacing'] = (dr, 0, dz)

return fieldmesh
return fieldmesh
3 changes: 1 addition & 2 deletions pmd_beamphysics/fields/fieldmesh.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
import numpy as np
from copy import deepcopy
import os
import functools


c_light = 299792458.
Expand Down Expand Up @@ -446,7 +445,7 @@ def from_ansys_ascii_3d(cls, *,
"""

if frequency is None:
raise ValueError(f"Please provide a frequency")
raise ValueError("Please provide a frequency")

data = read_ansys_ascii_3d_fields(efile, hfile, frequency=frequency)
return cls(data=data)
Expand Down
1 change: 0 additions & 1 deletion pmd_beamphysics/interfaces/astra.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 0 additions & 2 deletions pmd_beamphysics/interfaces/bmad.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion pmd_beamphysics/interfaces/gpt.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 4 additions & 3 deletions pmd_beamphysics/interfaces/impact.py
Original file line number Diff line number Diff line change
@@ -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.




Expand Down Expand Up @@ -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 = {}
Expand Down
4 changes: 2 additions & 2 deletions pmd_beamphysics/interfaces/lucretia.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ def lucretia_to_data(filename, ele_name='BEGINNING', t_ref=0, exclude_dead_parti
elif (len(ele_list) == 1):
ele_name = ele_list[0]

mdat = sio.loadmat(filename);
mdat = sio.loadmat(filename)

coords = mdat['bstore'][ele_name][0,0]['Bunch'][0,0]['x'][0,0]
charges = mdat['bstore'][ele_name][0,0]['Bunch'][0,0]['Q'][0,0][0]
Expand Down Expand Up @@ -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)
23 changes: 14 additions & 9 deletions pmd_beamphysics/interfaces/simion.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)) )
Expand Down
4 changes: 2 additions & 2 deletions pmd_beamphysics/interfaces/superfish.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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

Expand Down
30 changes: 15 additions & 15 deletions pmd_beamphysics/labels.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,20 +36,20 @@
'gamma': r'\gamma',
'theta': r'\theta',
'charge': 'Q',
'twiss_alpha_x' : 'Twiss\ '+r'\alpha_x',
'twiss_beta_x' : 'Twiss\ '+r'\beta_x',
'twiss_gamma_x' : 'Twiss\ '+r'\gamma_x',
'twiss_eta_x' : 'Twiss\ '+r'\eta_x',
'twiss_etap_x' : 'Twiss\ '+r"\eta'_x",
'twiss_emit_x' : 'Twiss\ '+r'\epsilon_{x}',
'twiss_norm_emit_x' : 'Twiss\ '+r'\epsilon_{n, x}',
'twiss_alpha_y' : 'Twiss\ '+r'\alpha_y',
'twiss_beta_y' : 'Twiss\ '+r'\beta_y',
'twiss_gamma_y' : 'Twiss\ '+r'\gamma_y',
'twiss_eta_y' : 'Twiss\ '+r'\eta_y',
'twiss_etap_y' : 'Twiss\ '+r"\eta'_y",
'twiss_emit_y' : 'Twiss\ '+r'\epsilon_{y}',
'twiss_norm_emit_y' : '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}',
Expand Down Expand Up @@ -179,4 +179,4 @@ def mathlabel(*keys, units=None, tex=True):
if units:
label = fr'{label} ({units})'

return label
return label
2 changes: 0 additions & 2 deletions pmd_beamphysics/particles.py
Original file line number Diff line number Diff line change
Expand Up @@ -1158,8 +1158,6 @@ def load_bunch_data(h5):
"""
Load particles into structured numpy array.
"""
n = len(h5['position/x'])

attrs = dict(h5.attrs)
data = {}
data['species'] = attrs['speciesType'].decode('utf-8') # String
Expand Down
24 changes: 13 additions & 11 deletions pmd_beamphysics/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -574,4 +576,4 @@ def plot_fieldmesh_cylindrical_1d(fm,
if return_figure:
return fig



3 changes: 1 addition & 2 deletions pmd_beamphysics/units.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
6 changes: 3 additions & 3 deletions pmd_beamphysics/writers.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,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.
Expand Down Expand Up @@ -118,8 +118,8 @@ def write_pmd_field(h5, data, name=None):
val = val.astype(complex)

# Write
g2 = write_component_data(g, key, val, unit=u)
write_component_data(g, key, val, unit=u)


def write_component_data(h5, name, data, unit=None):
"""
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import versioneer
from setuptools import setup, find_packages
from os import path, environ
from os import path

cur_dir = path.abspath(path.dirname(__file__))

Expand Down

0 comments on commit 62fe9fe

Please sign in to comment.