diff --git a/pmd_beamphysics/__init__.py b/pmd_beamphysics/__init__.py index 5d8e595..66288fa 100644 --- a/pmd_beamphysics/__init__.py +++ b/pmd_beamphysics/__init__.py @@ -1,7 +1,7 @@ -from .particles import ParticleGroup, single_particle -from .status import ParticleStatus from .fields.fieldmesh import FieldMesh +from .particles import ParticleGroup, single_particle from .readers import particle_paths +from .status import ParticleStatus from .writers import pmd_init try: diff --git a/pmd_beamphysics/fields/analysis.py b/pmd_beamphysics/fields/analysis.py index 344294a..a4857b7 100644 --- a/pmd_beamphysics/fields/analysis.py +++ b/pmd_beamphysics/fields/analysis.py @@ -1,11 +1,10 @@ -from pmd_beamphysics.units import mec2, c_light -from pmd_beamphysics.species import charge_state, mass_of - +import numpy as np from scipy import interpolate from scipy.integrate import solve_ivp from scipy.optimize import brent, brentq -import numpy as np +from pmd_beamphysics.species import charge_state, mass_of +from pmd_beamphysics.units import c_light, mec2 # Numpy migration per https://numpy.org/doc/stable/numpy_2_0_migration_guide.html if np.lib.NumpyVersion(np.__version__) >= "2.0.0": diff --git a/pmd_beamphysics/fields/corrector_modeling.py b/pmd_beamphysics/fields/corrector_modeling.py index 541a9ba..d24c5d6 100644 --- a/pmd_beamphysics/fields/corrector_modeling.py +++ b/pmd_beamphysics/fields/corrector_modeling.py @@ -1,9 +1,8 @@ +import numpy as np +from matplotlib import pyplot as plt from scipy.constants import mu_0 as u0 from scipy.constants import pi -from matplotlib import pyplot as plt -import numpy as np - from pmd_beamphysics import FieldMesh diff --git a/pmd_beamphysics/fields/expansion.py b/pmd_beamphysics/fields/expansion.py index 5433a4c..e0b4abe 100644 --- a/pmd_beamphysics/fields/expansion.py +++ b/pmd_beamphysics/fields/expansion.py @@ -1,6 +1,7 @@ -from scipy.interpolate import UnivariateSpline -from scipy import fft import numpy as np +from scipy import fft +from scipy.interpolate import UnivariateSpline + from pmd_beamphysics.units import c_light diff --git a/pmd_beamphysics/fields/fieldmesh.py b/pmd_beamphysics/fields/fieldmesh.py index 353a88e..0bd1fd2 100644 --- a/pmd_beamphysics/fields/fieldmesh.py +++ b/pmd_beamphysics/fields/fieldmesh.py @@ -1,56 +1,48 @@ -from pmd_beamphysics.units import pg_units - -from pmd_beamphysics.readers import ( - component_data, - expected_record_unit_dimension, - field_record_components, - field_paths, - component_from_alias, - load_field_attrs, - component_alias, -) +import functools +import os +from copy import deepcopy -from pmd_beamphysics.writers import write_pmd_field, pmd_field_init +import numpy as np +from h5py import File +from scipy.interpolate import RegularGridInterpolator from pmd_beamphysics import tools - -from pmd_beamphysics.plot import ( - plot_fieldmesh_cylindrical_2d, - plot_fieldmesh_cylindrical_1d, - plot_fieldmesh_rectangular_1d, - plot_fieldmesh_rectangular_2d, +from pmd_beamphysics.fields.conversion import ( + fieldmesh_rectangular_to_cylindrically_symmetric_data, ) - +from pmd_beamphysics.fields.expansion import expand_fieldmesh_from_onaxis from pmd_beamphysics.interfaces.ansys import read_ansys_ascii_3d_fields from pmd_beamphysics.interfaces.astra import ( - write_astra_1d_fieldmap, + astra_1d_fieldmap_data, read_astra_3d_fieldmaps, + write_astra_1d_fieldmap, write_astra_3d_fieldmaps, - astra_1d_fieldmap_data, ) from pmd_beamphysics.interfaces.gpt import write_gpt_fieldmesh from pmd_beamphysics.interfaces.impact import ( + create_impact_emfield_cartesian_ele, create_impact_solrf_ele, parse_impact_emfield_cartesian, write_impact_emfield_cartesian, - create_impact_emfield_cartesian_ele, ) -from pmd_beamphysics.interfaces.superfish import write_superfish_t7, read_superfish_t7 - -from pmd_beamphysics.fields.expansion import expand_fieldmesh_from_onaxis -from pmd_beamphysics.fields.conversion import ( - fieldmesh_rectangular_to_cylindrically_symmetric_data, +from pmd_beamphysics.interfaces.superfish import read_superfish_t7, write_superfish_t7 +from pmd_beamphysics.plot import ( + plot_fieldmesh_cylindrical_1d, + plot_fieldmesh_cylindrical_2d, + plot_fieldmesh_rectangular_1d, + plot_fieldmesh_rectangular_2d, ) - -import functools - -from scipy.interpolate import RegularGridInterpolator - -from h5py import File -import numpy as np -from copy import deepcopy -import os - +from pmd_beamphysics.readers import ( + component_alias, + component_data, + component_from_alias, + expected_record_unit_dimension, + field_paths, + field_record_components, + load_field_attrs, +) +from pmd_beamphysics.units import pg_units +from pmd_beamphysics.writers import pmd_field_init, write_pmd_field # ----------------------------------------- # Classes diff --git a/pmd_beamphysics/interfaces/ansys.py b/pmd_beamphysics/interfaces/ansys.py index daaf59d..8426142 100644 --- a/pmd_beamphysics/interfaces/ansys.py +++ b/pmd_beamphysics/interfaces/ansys.py @@ -1,4 +1,5 @@ import numpy as np + from pmd_beamphysics.units import mu_0 diff --git a/pmd_beamphysics/interfaces/astra.py b/pmd_beamphysics/interfaces/astra.py index 7c973f1..992ad1d 100644 --- a/pmd_beamphysics/interfaces/astra.py +++ b/pmd_beamphysics/interfaces/astra.py @@ -1,8 +1,9 @@ -import numpy as np -from pmd_beamphysics.status import ParticleStatus -from pmd_beamphysics.readers import component_alias import os +import numpy as np + +from pmd_beamphysics.readers import component_alias +from pmd_beamphysics.status import ParticleStatus astra_species_name = {1: "electron", 2: "positron", 3: "proton", 4: "hydrogen"} astra_species_index = {v: k for k, v in astra_species_name.items()} # Inverse mapping diff --git a/pmd_beamphysics/interfaces/bmad.py b/pmd_beamphysics/interfaces/bmad.py index 9ff9511..be14b8f 100644 --- a/pmd_beamphysics/interfaces/bmad.py +++ b/pmd_beamphysics/interfaces/bmad.py @@ -1,6 +1,7 @@ import numpy as np -from pmd_beamphysics.units import c_light + from pmd_beamphysics.species import mass_of +from pmd_beamphysics.units import c_light # Remove from below, because this docstring is used directly in ParticleGroup diff --git a/pmd_beamphysics/interfaces/elegant.py b/pmd_beamphysics/interfaces/elegant.py index dd0288e..7bb2747 100644 --- a/pmd_beamphysics/interfaces/elegant.py +++ b/pmd_beamphysics/interfaces/elegant.py @@ -1,11 +1,11 @@ -from pmd_beamphysics.units import mec2 - -import numpy as np -import subprocess import os +import subprocess +import numpy as np from h5py import File +from pmd_beamphysics.units import mec2 + def write_elegant(particle_group, outfile, verbose=False): """ diff --git a/pmd_beamphysics/interfaces/genesis.py b/pmd_beamphysics/interfaces/genesis.py index 7fc9351..9104542 100644 --- a/pmd_beamphysics/interfaces/genesis.py +++ b/pmd_beamphysics/interfaces/genesis.py @@ -1,10 +1,10 @@ import os + import numpy as np from h5py import File from pmd_beamphysics.statistics import twiss_calc -from pmd_beamphysics.units import mec2, c_light, write_unit_h5, unit - +from pmd_beamphysics.units import c_light, mec2, unit, write_unit_h5 # Genesis 1.3 # ------------- diff --git a/pmd_beamphysics/interfaces/gpt.py b/pmd_beamphysics/interfaces/gpt.py index d6eb4da..1ef4892 100644 --- a/pmd_beamphysics/interfaces/gpt.py +++ b/pmd_beamphysics/interfaces/gpt.py @@ -1,9 +1,10 @@ -from pmd_beamphysics.units import e_charge, c_light -from pmd_beamphysics.interfaces.superfish import fish_complex_to_real_fields +import os +import subprocess import numpy as np -import subprocess -import os + +from pmd_beamphysics.interfaces.superfish import fish_complex_to_real_fields +from pmd_beamphysics.units import c_light, e_charge def write_gpt(particle_group, outfile, asci2gdf_bin=None, verbose=False): diff --git a/pmd_beamphysics/interfaces/impact.py b/pmd_beamphysics/interfaces/impact.py index 40fc52e..1bc554b 100644 --- a/pmd_beamphysics/interfaces/impact.py +++ b/pmd_beamphysics/interfaces/impact.py @@ -1,5 +1,6 @@ import os import re + import numpy as np from numpy import pi from scipy import fft diff --git a/pmd_beamphysics/interfaces/litrack.py b/pmd_beamphysics/interfaces/litrack.py index 1eafe2a..03843a5 100644 --- a/pmd_beamphysics/interfaces/litrack.py +++ b/pmd_beamphysics/interfaces/litrack.py @@ -1,6 +1,7 @@ -from pmd_beamphysics.units import c_light import numpy as np +from pmd_beamphysics.units import c_light + def write_litrack(particle_group, outfile="litrack.zd", p0c=None, verbose=False): """ diff --git a/pmd_beamphysics/interfaces/lucretia.py b/pmd_beamphysics/interfaces/lucretia.py index 6713d3f..1f6e7d9 100644 --- a/pmd_beamphysics/interfaces/lucretia.py +++ b/pmd_beamphysics/interfaces/lucretia.py @@ -1,5 +1,5 @@ -import scipy.io as sio import numpy as np +import scipy.io as sio def lucretia_to_data( diff --git a/pmd_beamphysics/interfaces/simion.py b/pmd_beamphysics/interfaces/simion.py index b2e6d48..eb2d67a 100644 --- a/pmd_beamphysics/interfaces/simion.py +++ b/pmd_beamphysics/interfaces/simion.py @@ -1,9 +1,8 @@ import numpy as np - -from pmd_beamphysics.species import mec2, e_charge - from scipy.constants import physical_constants +from pmd_beamphysics.species import e_charge, mec2 + amu_to_rest_mass_energy = ( physical_constants["atomic mass constant energy equivalent in MeV"][0] * 1e6 ) diff --git a/pmd_beamphysics/interfaces/superfish.py b/pmd_beamphysics/interfaces/superfish.py index 28eaa5d..944ce1a 100644 --- a/pmd_beamphysics/interfaces/superfish.py +++ b/pmd_beamphysics/interfaces/superfish.py @@ -1,5 +1,5 @@ -import scipy.constants import numpy as np +import scipy.constants mu_0 = scipy.constants.mu_0 diff --git a/pmd_beamphysics/labels.py b/pmd_beamphysics/labels.py index e24a3cf..53d5b42 100644 --- a/pmd_beamphysics/labels.py +++ b/pmd_beamphysics/labels.py @@ -1,5 +1,4 @@ -from pmd_beamphysics.units import parse_bunching_str, nice_array - +from pmd_beamphysics.units import nice_array, parse_bunching_str TEXLABEL = { # 'status' diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py index 35e013d..2bef90f 100644 --- a/pmd_beamphysics/particles.py +++ b/pmd_beamphysics/particles.py @@ -1,12 +1,19 @@ -from pmd_beamphysics.units import pg_units, c_light, parse_bunching_str +import functools +import os +from copy import deepcopy + +import numpy as np +from h5py import File -from pmd_beamphysics.interfaces.astra import write_astra import pmd_beamphysics.interfaces.bmad as bmad +import pmd_beamphysics.statistics as statistics +from pmd_beamphysics.interfaces.astra import write_astra +from pmd_beamphysics.interfaces.elegant import write_elegant from pmd_beamphysics.interfaces.genesis import ( - write_genesis4_distribution, genesis2_beam_data, write_genesis2_beam_file, write_genesis4_beam, + write_genesis4_distribution, ) from pmd_beamphysics.interfaces.gpt import write_gpt from pmd_beamphysics.interfaces.impact import write_impact @@ -14,32 +21,20 @@ from pmd_beamphysics.interfaces.lucretia import write_lucretia from pmd_beamphysics.interfaces.opal import write_opal from pmd_beamphysics.interfaces.simion import write_simion -from pmd_beamphysics.interfaces.elegant import write_elegant - from pmd_beamphysics.plot import density_plot, marginal_plot, slice_plot - from pmd_beamphysics.readers import particle_array, particle_paths from pmd_beamphysics.species import charge_of, mass_of - from pmd_beamphysics.statistics import ( + matched_particles, 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 -import numpy as np -from copy import deepcopy -import functools -import os - +from pmd_beamphysics.units import c_light, parse_bunching_str, pg_units +from pmd_beamphysics.writers import pmd_init, write_pmd_bunch # ----------------------------------------- # Classes diff --git a/pmd_beamphysics/plot.py b/pmd_beamphysics/plot.py index a73b939..4886ae0 100644 --- a/pmd_beamphysics/plot.py +++ b/pmd_beamphysics/plot.py @@ -10,7 +10,6 @@ # 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 diff --git a/pmd_beamphysics/readers.py b/pmd_beamphysics/readers.py index d5c45e7..af4503e 100644 --- a/pmd_beamphysics/readers.py +++ b/pmd_beamphysics/readers.py @@ -1,9 +1,8 @@ -from .units import dimension, dimension_name, SI_symbol, c_light, e_charge -from .tools import decode_attrs, decode_attr - - import numpy as np +from .tools import decode_attr, decode_attrs +from .units import SI_symbol, c_light, dimension, dimension_name, e_charge + # ----------------------------------------- # General Utilities diff --git a/pmd_beamphysics/tools.py b/pmd_beamphysics/tools.py index 70554b0..1fd2a3f 100644 --- a/pmd_beamphysics/tools.py +++ b/pmd_beamphysics/tools.py @@ -1,4 +1,5 @@ import datetime + import numpy as np diff --git a/pmd_beamphysics/writers.py b/pmd_beamphysics/writers.py index ddbc98e..c5d4e69 100644 --- a/pmd_beamphysics/writers.py +++ b/pmd_beamphysics/writers.py @@ -1,8 +1,8 @@ import numpy as np -from .units import pg_units from .readers import component_from_alias, load_field_attrs -from .tools import fstr, encode_attrs +from .tools import encode_attrs, fstr +from .units import pg_units def pmd_init(h5, basePath="/data/%T/", particlesPath="./"):