diff --git a/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py index 033dcec37..94c205c80 100644 --- a/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py +++ b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py @@ -170,7 +170,6 @@ def __init__(self, stage_i, surrogate_model, surrogate_length, stage_start): self.ds = surrogate_length def surrogate_push(self, pc, step): - ref_part = pc.ref_particle() ref_z_i = ref_part.z ref_z_i_LPA = ref_z_i - self.stage_start diff --git a/src/python/impactx/MADXParser.pyi b/src/python/impactx/MADXParser.pyi index 315a53563..768dfdf70 100644 --- a/src/python/impactx/MADXParser.pyi +++ b/src/python/impactx/MADXParser.pyi @@ -2,63 +2,69 @@ from __future__ import annotations import os as os import re as re import warnings as warnings -__all__ = ['MADXInputError', 'MADXInputWarning', 'MADXParser', 'MADXParserError', 'os', 're', 'warnings'] + +__all__ = [ + "MADXInputError", + "MADXInputWarning", + "MADXParser", + "MADXParserError", + "os", + "re", + "warnings", +] + class MADXInputError(MADXParserError): - def __init__(self, args, with_traceback): - ... + def __init__(self, args, with_traceback): ... + class MADXInputWarning(UserWarning): pass + class MADXParser: """ - Simple MADX parser. - It expects a single line per element. + Simple MADX parser. + It expects a single line per element. """ - def __init__(self): - ... - def __str__(self): - ... + def __init__(self): ... + def __str__(self): ... def _combine(self, lattice): """ - Combine to one list of all basic - elements. + Combine to one list of all basic + elements. - return a list of of element dictionaries + return a list of of element dictionaries """ def _flatten(self, line): """ - Find sublines. + Find sublines. """ def _noWhitespace(self, string): """ - Remove white space from a string. + Remove white space from a string. - 14. Oct. 2017, - https://stackoverflow.com/questions/3739909/how-to-strip-all-whitespace-from-string + 14. Oct. 2017, + https://stackoverflow.com/questions/3739909/how-to-strip-all-whitespace-from-string """ - def getBeamline(self): - ... - def getEtot(self): - ... - def getParticle(self): - ... - def nonblank_lines_to_lowercase(self, f): - ... + def getBeamline(self): ... + def getEtot(self): ... + def getParticle(self): ... + def nonblank_lines_to_lowercase(self, f): ... def parse(self, fn): """ - fn (str) filename + fn (str) filename """ + class MADXParserError(Exception): pass diff --git a/src/python/impactx/__init__.pyi b/src/python/impactx/__init__.pyi index 96de2c48b..e76f3ec5c 100644 --- a/src/python/impactx/__init__.pyi +++ b/src/python/impactx/__init__.pyi @@ -1,20 +1,23 @@ """ - impactx_pybind - -------------- - .. currentmodule:: impactx_pybind +impactx_pybind +-------------- +.. currentmodule:: impactx_pybind - .. autosummary:: - :toctree: _generate - ImpactX - distribution - elements +.. autosummary:: + :toctree: _generate + ImpactX + distribution + elements """ + from __future__ import annotations from amrex import space3d as amr from impactx.extensions.ImpactXParIter import register_ImpactXParIter_extension -from impactx.extensions.ImpactXParticleContainer import register_ImpactXParticleContainer_extension +from impactx.extensions.ImpactXParticleContainer import ( + register_ImpactXParticleContainer_extension, +) from impactx.impactx_pybind import Config from impactx.impactx_pybind import CoordSystem from impactx.impactx_pybind import ImpactX @@ -33,10 +36,38 @@ from . import MADXParser from . import extensions from . import impactx_pybind from . import madx_to_impactx -__all__ = ['Config', 'CoordSystem', 'ImpactX', 'ImpactXParConstIter', 'ImpactXParIter', 'ImpactXParticleContainer', 'MADXParser', 'RefPart', 'amr', 'coordinate_transformation', 'cxx', 'distribution', 'elements', 'extensions', 'impactx_pybind', 'madx_to_impactx', 'os', 'push', 'read_beam', 'read_lattice', 'register_ImpactXParIter_extension', 'register_ImpactXParticleContainer_extension', 's', 't'] -__author__: str = 'Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al.' -__license__: str = 'BSD-3-Clause-LBNL' -__version__: str = '24.08' + +__all__ = [ + "Config", + "CoordSystem", + "ImpactX", + "ImpactXParConstIter", + "ImpactXParIter", + "ImpactXParticleContainer", + "MADXParser", + "RefPart", + "amr", + "coordinate_transformation", + "cxx", + "distribution", + "elements", + "extensions", + "impactx_pybind", + "madx_to_impactx", + "os", + "push", + "read_beam", + "read_lattice", + "register_ImpactXParIter_extension", + "register_ImpactXParticleContainer_extension", + "s", + "t", +] +__author__: str = ( + "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al." +) +__license__: str = "BSD-3-Clause-LBNL" +__version__: str = "24.08" s: impactx_pybind.CoordSystem # value = t: impactx_pybind.CoordSystem # value = cxx = impactx_pybind diff --git a/src/python/impactx/dashboard/Input/generalFunctions.py b/src/python/impactx/dashboard/Input/generalFunctions.py index 60164bb62..dcb4717f0 100644 --- a/src/python/impactx/dashboard/Input/generalFunctions.py +++ b/src/python/impactx/dashboard/Input/generalFunctions.py @@ -22,7 +22,6 @@ class generalFunctions: - @staticmethod def documentation(section_name): """ diff --git a/src/python/impactx/dashboard/Toolbar/toolbarMain.py b/src/python/impactx/dashboard/Toolbar/toolbarMain.py index ff5894770..cdf57ccb5 100644 --- a/src/python/impactx/dashboard/Toolbar/toolbarMain.py +++ b/src/python/impactx/dashboard/Toolbar/toolbarMain.py @@ -61,8 +61,8 @@ def run_toolbar(): Builds toolbar for the 'Run' page. """ - vuetify.VSpacer(), - ToolbarElements.run_simulation_button(), + (vuetify.VSpacer(),) + (ToolbarElements.run_simulation_button(),) @staticmethod def analyze_toolbar(): diff --git a/src/python/impactx/extensions/ImpactXParIter.pyi b/src/python/impactx/extensions/ImpactXParIter.pyi index e18789e82..5757d68dd 100644 --- a/src/python/impactx/extensions/ImpactXParIter.pyi +++ b/src/python/impactx/extensions/ImpactXParIter.pyi @@ -6,12 +6,21 @@ Copyright 2023 ImpactX contributors Authors: Axel Huebl License: BSD-3-Clause-LBNL """ + from __future__ import annotations -__all__ = ['register_ImpactXParIter_extension', 'soa', 'soa_int_comps', 'soa_real_comps'] + +__all__ = [ + "register_ImpactXParIter_extension", + "soa", + "soa_int_comps", + "soa_real_comps", +] + def register_ImpactXParIter_extension(impactx_pybind): """ ImpactXParIter helper methods """ + def soa(self): """ Get the StructOfArrays on the current tile @@ -22,37 +31,39 @@ def soa(self): used to query particle container component names """ + def soa_int_comps(pti, num_comps): """ - Name the ImpactX int components in SoA. + Name the ImpactX int components in SoA. - Parameters - ---------- - pti : ImpactXParIter or ImpactXParConstIter - used to query particle container component names - num_comps : int - number of components to generate names for. + Parameters + ---------- + pti : ImpactXParIter or ImpactXParConstIter + used to query particle container component names + num_comps : int + number of components to generate names for. - Returns - ------- - A list of length num_comps with values. + Returns + ------- + A list of length num_comps with values. """ + def soa_real_comps(pti, num_comps): """ - Name the ImpactX ParticleReal components in SoA. + Name the ImpactX ParticleReal components in SoA. - Parameters - ---------- - pti : ImpactXParIter or ImpactXParConstIter - used to query particle container component names - num_comps : int - number of components to generate names for. + Parameters + ---------- + pti : ImpactXParIter or ImpactXParConstIter + used to query particle container component names + num_comps : int + number of components to generate names for. - Returns - ------- - A list of length num_comps with values. + Returns + ------- + A list of length num_comps with values. """ diff --git a/src/python/impactx/extensions/ImpactXParticleContainer.pyi b/src/python/impactx/extensions/ImpactXParticleContainer.pyi index f7da3d329..dfb5ebefa 100644 --- a/src/python/impactx/extensions/ImpactXParticleContainer.pyi +++ b/src/python/impactx/extensions/ImpactXParticleContainer.pyi @@ -6,28 +6,32 @@ Copyright 2023 ImpactX contributors Authors: Axel Huebl License: BSD-3-Clause-LBNL """ + from __future__ import annotations -__all__ = ['ix_pc_plot_mpl_phasespace', 'register_ImpactXParticleContainer_extension'] -def ix_pc_plot_mpl_phasespace(self, num_bins = 50, root_rank = 0): + +__all__ = ["ix_pc_plot_mpl_phasespace", "register_ImpactXParticleContainer_extension"] + +def ix_pc_plot_mpl_phasespace(self, num_bins=50, root_rank=0): """ - Plot the longitudinal and transverse phase space projections with matplotlib. + Plot the longitudinal and transverse phase space projections with matplotlib. - Parameters - ---------- - self : ImpactXParticleContainer_* - The particle container class in ImpactX - num_bins : int, default=50 - The number of bins for spatial and momentum directions per plot axis. - root_rank : int, default=0 - MPI root rank to reduce to in parallel runs. + Parameters + ---------- + self : ImpactXParticleContainer_* + The particle container class in ImpactX + num_bins : int, default=50 + The number of bins for spatial and momentum directions per plot axis. + root_rank : int, default=0 + MPI root rank to reduce to in parallel runs. - Returns - ------- - A matplotlib figure with containing the plot. - For MPI-parallel ranks, the figure is only created on the root_rank. + Returns + ------- + A matplotlib figure with containing the plot. + For MPI-parallel ranks, the figure is only created on the root_rank. """ + def register_ImpactXParticleContainer_extension(ixpc): """ ImpactXParticleContainer helper methods diff --git a/src/python/impactx/extensions/__init__.pyi b/src/python/impactx/extensions/__init__.pyi index 20b5567f8..0ce7fb449 100644 --- a/src/python/impactx/extensions/__init__.pyi +++ b/src/python/impactx/extensions/__init__.pyi @@ -1,4 +1,5 @@ from __future__ import annotations from . import ImpactXParIter from . import ImpactXParticleContainer -__all__ = ['ImpactXParIter', 'ImpactXParticleContainer'] + +__all__ = ["ImpactXParIter", "ImpactXParticleContainer"] diff --git a/src/python/impactx/impactx_pybind/__init__.pyi b/src/python/impactx/impactx_pybind/__init__.pyi index 9f0781d16..95962ecf9 100644 --- a/src/python/impactx/impactx_pybind/__init__.pyi +++ b/src/python/impactx/impactx_pybind/__init__.pyi @@ -1,16 +1,17 @@ """ - impactx_pybind - -------------- - .. currentmodule:: impactx_pybind +impactx_pybind +-------------- +.. currentmodule:: impactx_pybind - .. autosummary:: - :toctree: _generate - ImpactX - distribution - elements +.. autosummary:: + :toctree: _generate + ImpactX + distribution + elements """ + from __future__ import annotations from amrex import space3d as amr import amrex.space3d.amrex_3d_pybind @@ -18,12 +19,30 @@ import pybind11_stubgen.typing_ext import typing from . import distribution from . import elements -__all__ = ['Config', 'CoordSystem', 'ImpactX', 'ImpactXParConstIter', 'ImpactXParIter', 'ImpactXParticleContainer', 'RefPart', 'amr', 'coordinate_transformation', 'distribution', 'elements', 'push', 's', 't'] + +__all__ = [ + "Config", + "CoordSystem", + "ImpactX", + "ImpactXParConstIter", + "ImpactXParIter", + "ImpactXParticleContainer", + "RefPart", + "amr", + "coordinate_transformation", + "distribution", + "elements", + "push", + "s", + "t", +] + class Config: gpu_backend = None have_gpu: typing.ClassVar[bool] = False have_mpi: typing.ClassVar[bool] = True have_omp: typing.ClassVar[bool] = True + class CoordSystem: """ Members: @@ -32,43 +51,47 @@ class CoordSystem: t """ - __members__: typing.ClassVar[dict[str, CoordSystem]] # value = {'s': , 't': } + + __members__: typing.ClassVar[ + dict[str, CoordSystem] + ] # value = {'s': , 't': } s: typing.ClassVar[CoordSystem] # value = t: typing.ClassVar[CoordSystem] # value = - def __eq__(self, other: typing.Any) -> bool: - ... - def __getstate__(self) -> int: - ... - def __hash__(self) -> int: - ... - def __index__(self) -> int: - ... - def __init__(self, value: int) -> None: - ... - def __int__(self) -> int: - ... - def __ne__(self, other: typing.Any) -> bool: - ... - def __repr__(self) -> str: - ... - def __setstate__(self, state: int) -> None: - ... - def __str__(self) -> str: - ... - @property - def name(self) -> str: - ... - @property - def value(self) -> int: - ... + def __eq__(self, other: typing.Any) -> bool: ... + def __getstate__(self) -> int: ... + def __hash__(self) -> int: ... + def __index__(self) -> int: ... + def __init__(self, value: int) -> None: ... + def __int__(self) -> int: ... + def __ne__(self, other: typing.Any) -> bool: ... + def __repr__(self) -> str: ... + def __setstate__(self, state: int) -> None: ... + def __str__(self) -> str: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... + class ImpactX: - def DistributionMap(self, lev: int) -> amrex.space3d.amrex_3d_pybind.DistributionMapping: - ... - def Geom(self, lev: int) -> amrex.space3d.amrex_3d_pybind.Geometry: - ... - def __init__(self) -> None: - ... - def add_particles(self, bunch_charge: float, distr: distribution.Empty | distribution.Gaussian | distribution.Kurth4D | distribution.Kurth6D | distribution.KVdist | distribution.Thermal | distribution.Triangle | distribution.Semigaussian | distribution.Waterbag, npart: int) -> None: + def DistributionMap( + self, lev: int + ) -> amrex.space3d.amrex_3d_pybind.DistributionMapping: ... + def Geom(self, lev: int) -> amrex.space3d.amrex_3d_pybind.Geometry: ... + def __init__(self) -> None: ... + def add_particles( + self, + bunch_charge: float, + distr: distribution.Empty + | distribution.Gaussian + | distribution.Kurth4D + | distribution.Kurth6D + | distribution.KVdist + | distribution.Thermal + | distribution.Triangle + | distribution.Semigaussian + | distribution.Waterbag, + npart: int, + ) -> None: """ Generate and add n particles to the particle container. @@ -76,8 +99,7 @@ class ImpactX: distribution's extent and then redistribute particles in according AMReX grid boxes. """ - def boxArray(self, lev: int) -> amrex.space3d.amrex_3d_pybind.BoxArray: - ... + def boxArray(self, lev: int) -> amrex.space3d.amrex_3d_pybind.BoxArray: ... def deposit_charge(self) -> None: """ Deposit charge in x,y,z. @@ -90,18 +112,15 @@ class ImpactX: """ Deallocate all contexts and data. """ - def init_beam_distribution_from_inputs(self) -> None: - ... + def init_beam_distribution_from_inputs(self) -> None: ... def init_grids(self) -> None: """ Initialize AMReX blocks/grids for domain decomposition & space charge mesh. This must come first, before particle beams and lattice elements are initialized. """ - def init_lattice_elements_from_inputs(self) -> None: - ... - def load_inputs_file(self, arg0: str) -> None: - ... + def init_lattice_elements_from_inputs(self) -> None: ... + def load_inputs_file(self, arg0: str) -> None: ... def particle_container(self) -> ImpactXParticleContainer: """ Access the beam particle container. @@ -118,7 +137,9 @@ class ImpactX: """ charge density per level """ - def space_charge_field(self, lev: int, comp: str) -> amrex.space3d.amrex_3d_pybind.MultiFab: + def space_charge_field( + self, lev: int, comp: str + ) -> amrex.space3d.amrex_3d_pybind.MultiFab: """ space charge force (vector: x,y,z) per level """ @@ -129,8 +150,7 @@ class ImpactX: if there are unused parameters in the input. """ @abort_on_unused_inputs.setter - def abort_on_unused_inputs(self, arg1: int) -> None: - ... + def abort_on_unused_inputs(self, arg1: int) -> None: ... @property def abort_on_warning_threshold(self) -> str: """ @@ -139,8 +159,7 @@ class ImpactX: Valid choices are: ['low', 'medium', 'high']. """ @abort_on_warning_threshold.setter - def abort_on_warning_threshold(self, arg1: str) -> None: - ... + def abort_on_warning_threshold(self, arg1: str) -> None: ... @property def always_warn_immediately(self) -> int: """ @@ -148,32 +167,28 @@ class ImpactX: as soon as it is generated. """ @always_warn_immediately.setter - def always_warn_immediately(self, arg1: int) -> None: - ... + def always_warn_immediately(self, arg1: int) -> None: ... @property def blocking_factor_x(self) -> list[int]: """ AMReX blocking factor for a direction, per MR level. """ @blocking_factor_x.setter - def blocking_factor_x(self, arg1: list[int]) -> None: - ... + def blocking_factor_x(self, arg1: list[int]) -> None: ... @property def blocking_factor_y(self) -> list[int]: """ AMReX blocking factor for a direction, per MR level. """ @blocking_factor_y.setter - def blocking_factor_y(self, arg1: list[int]) -> None: - ... + def blocking_factor_y(self, arg1: list[int]) -> None: ... @property def blocking_factor_z(self) -> list[int]: """ AMReX blocking factor for a direction, per MR level. """ @blocking_factor_z.setter - def blocking_factor_z(self, arg1: list[int]) -> None: - ... + def blocking_factor_z(self, arg1: list[int]) -> None: ... @property def diag_file_min_digits(self) -> int: """ @@ -181,8 +196,7 @@ class ImpactX: number appended to the diagnostic file names. """ @diag_file_min_digits.setter - def diag_file_min_digits(self, arg1: int) -> None: - ... + def diag_file_min_digits(self, arg1: int) -> None: ... @property def diagnostics(self) -> bool: """ @@ -190,24 +204,21 @@ class ImpactX: Disabling this is mostly used for benchmarking. """ @diagnostics.setter - def diagnostics(self, arg1: bool) -> None: - ... + def diagnostics(self, arg1: bool) -> None: ... @property def domain(self) -> amrex.space3d.amrex_3d_pybind.RealBox: """ The physical extent of the full simulation domain, relative to the reference particle position, in meters. """ @domain.setter - def domain(self, arg1: amrex.space3d.amrex_3d_pybind.RealBox) -> None: - ... + def domain(self, arg1: amrex.space3d.amrex_3d_pybind.RealBox) -> None: ... @property def dynamic_size(self) -> bool: """ Use dynamic (``true``) resizing of the field mesh or static sizing (``false``). """ @dynamic_size.setter - def dynamic_size(self, arg1: bool) -> None: - ... + def dynamic_size(self, arg1: bool) -> None: ... @property def finest_level(self) -> int: """ @@ -219,56 +230,52 @@ class ImpactX: Access the accelerator element lattice. """ @lattice.setter - def lattice(self, arg0: elements.KnownElementsList) -> None: - ... + def lattice(self, arg0: elements.KnownElementsList) -> None: ... @property def max_level(self) -> int: """ The maximum mesh-refinement level for the simulation. """ @max_level.setter - def max_level(self, arg1: int) -> None: - ... + def max_level(self, arg1: int) -> None: ... @property def mlmg_absolute_tolerance(self) -> bool: """ The absolute tolerance with which the space-charge fields should be calculated in units of V/m^2. More specifically, the acceptable residual with which the solution can be considered converged. In general this should be left as the default, but in cases where the simulation state changes very little between steps it can occur that the initial guess for the MLMG solver is so close to the converged value that it fails to improve that solution sufficiently to reach the mlmg_relative_tolerance value. """ @mlmg_absolute_tolerance.setter - def mlmg_absolute_tolerance(self, arg1: float) -> None: - ... + def mlmg_absolute_tolerance(self, arg1: float) -> None: ... @property def mlmg_max_iters(self) -> bool: """ Maximum number of iterations used for MLMG solver for space-charge fields calculation. In case if MLMG converges but fails to reach the desired self_fields_required_precision, this parameter may be increased. """ @mlmg_max_iters.setter - def mlmg_max_iters(self, arg1: int) -> None: - ... + def mlmg_max_iters(self, arg1: int) -> None: ... @property def mlmg_relative_tolerance(self) -> bool: """ The relative precision with which the electrostatic space-charge fields should be calculated. More specifically, the space-charge fields are computed with an iterative Multi-Level Multi-Grid (MLMG) solver. This solver can fail to reach the default precision within a reasonable time. """ @mlmg_relative_tolerance.setter - def mlmg_relative_tolerance(self, arg1: float) -> None: - ... + def mlmg_relative_tolerance(self, arg1: float) -> None: ... @property def mlmg_verbosity(self) -> bool: """ The verbosity used for MLMG solver for space-charge fields calculation. Currently MLMG solver looks for verbosity levels from 0-5. A higher number results in more verbose output. """ @mlmg_verbosity.setter - def mlmg_verbosity(self, arg1: int) -> None: - ... + def mlmg_verbosity(self, arg1: int) -> None: ... @property def n_cell(self) -> list[int]: """ The number of grid points along each direction on the coarsest level. """ @n_cell.setter - def n_cell(self, arg1: typing.Annotated[list[int], pybind11_stubgen.typing_ext.FixedSize(3)]) -> None: - ... + def n_cell( + self, + arg1: typing.Annotated[list[int], pybind11_stubgen.typing_ext.FixedSize(3)], + ) -> None: ... @property def particle_lost_diagnostics_backend(self) -> str: """ @@ -277,40 +284,35 @@ class ImpactX: See the ``BeamMonitor`` element for backend values. """ @particle_lost_diagnostics_backend.setter - def particle_lost_diagnostics_backend(self, arg1: str) -> None: - ... + def particle_lost_diagnostics_backend(self, arg1: str) -> None: ... @property def particle_shape(self) -> int: """ Whether to calculate space charge effects. """ @particle_shape.setter - def particle_shape(self, arg1: int) -> None: - ... + def particle_shape(self, arg1: int) -> None: ... @property def periods(self) -> int: """ The number of periods to repeat the lattice. """ @periods.setter - def periods(self, arg1: int) -> None: - ... + def periods(self, arg1: int) -> None: ... @property def poisson_solver(self) -> str: """ The numerical solver to solve the Poisson equation when calculating space charge effects. Either multigrid (default) or fft. """ @poisson_solver.setter - def poisson_solver(self, arg1: str) -> None: - ... + def poisson_solver(self, arg1: str) -> None: ... @property def prob_relative(self) -> float: """ The field mesh spans, per direction, multiple times the maximum physical extent of beam particles, as given by this factor. """ @prob_relative.setter - def prob_relative(self, arg1: list[float]) -> None: - ... + def prob_relative(self, arg1: list[float]) -> None: ... @property def slice_step_diagnostics(self) -> bool: """ @@ -320,16 +322,14 @@ class ImpactX: Enabling this flag will write diagnostics every step and slice step. """ @slice_step_diagnostics.setter - def slice_step_diagnostics(self, arg1: bool) -> None: - ... + def slice_step_diagnostics(self, arg1: bool) -> None: ... @property def space_charge(self) -> bool: """ Enable or disable space charge calculations (default: enabled). """ @space_charge.setter - def space_charge(self, arg1: bool) -> None: - ... + def space_charge(self, arg1: bool) -> None: ... @property def verbose(self) -> int: """ @@ -337,17 +337,27 @@ class ImpactX: ``0`` for silent, higher is more verbose. Default is ``1``. """ @verbose.setter - def verbose(self, arg1: int) -> None: - ... -class ImpactXParConstIter(amrex.space3d.amrex_3d_pybind.ParConstIter_pureSoA_8_0_default): + def verbose(self, arg1: int) -> None: ... + +class ImpactXParConstIter( + amrex.space3d.amrex_3d_pybind.ParConstIter_pureSoA_8_0_default +): @typing.overload - def __init__(self, particle_container: amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default, level: int) -> None: - ... + def __init__( + self, + particle_container: amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default, + level: int, + ) -> None: ... @typing.overload - def __init__(self, particle_container: amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default, level: int, info: amrex.space3d.amrex_3d_pybind.MFItInfo) -> None: - ... - def pc(self) -> amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default: - ... + def __init__( + self, + particle_container: amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default, + level: int, + info: amrex.space3d.amrex_3d_pybind.MFItInfo, + ) -> None: ... + def pc( + self, + ) -> amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default: ... def soa(self): """ Get the StructOfArrays on the current tile @@ -358,15 +368,24 @@ class ImpactXParConstIter(amrex.space3d.amrex_3d_pybind.ParConstIter_pureSoA_8_0 used to query particle container component names """ + class ImpactXParIter(amrex.space3d.amrex_3d_pybind.ParIter_pureSoA_8_0_default): @typing.overload - def __init__(self, particle_container: amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default, level: int) -> None: - ... + def __init__( + self, + particle_container: amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default, + level: int, + ) -> None: ... @typing.overload - def __init__(self, particle_container: amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default, level: int, info: amrex.space3d.amrex_3d_pybind.MFItInfo) -> None: - ... - def pc(self) -> amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default: - ... + def __init__( + self, + particle_container: amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default, + level: int, + info: amrex.space3d.amrex_3d_pybind.MFItInfo, + ) -> None: ... + def pc( + self, + ) -> amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default: ... def soa(self): """ Get the StructOfArrays on the current tile @@ -377,10 +396,23 @@ class ImpactXParIter(amrex.space3d.amrex_3d_pybind.ParIter_pureSoA_8_0_default): used to query particle container component names """ -class ImpactXParticleContainer(amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default): + +class ImpactXParticleContainer( + amrex.space3d.amrex_3d_pybind.ParticleContainer_pureSoA_8_0_default +): const_iterator = ImpactXParConstIter iterator = ImpactXParIter - def add_n_particles(self, x: amrex.space3d.amrex_3d_pybind.PODVector_real_std, y: amrex.space3d.amrex_3d_pybind.PODVector_real_std, t: amrex.space3d.amrex_3d_pybind.PODVector_real_std, px: amrex.space3d.amrex_3d_pybind.PODVector_real_std, py: amrex.space3d.amrex_3d_pybind.PODVector_real_std, pt: amrex.space3d.amrex_3d_pybind.PODVector_real_std, qm: float, bchchg: float) -> None: + def add_n_particles( + self, + x: amrex.space3d.amrex_3d_pybind.PODVector_real_std, + y: amrex.space3d.amrex_3d_pybind.PODVector_real_std, + t: amrex.space3d.amrex_3d_pybind.PODVector_real_std, + px: amrex.space3d.amrex_3d_pybind.PODVector_real_std, + py: amrex.space3d.amrex_3d_pybind.PODVector_real_std, + pt: amrex.space3d.amrex_3d_pybind.PODVector_real_std, + qm: float, + bchchg: float, + ) -> None: """ Add new particles to the container for fixed s. @@ -409,27 +441,29 @@ class ImpactXParticleContainer(amrex.space3d.amrex_3d_pybind.ParticleContainer_p :return: x_min, y_min, z_min, x_max, y_max, z_max """ - def plot_phasespace(self, num_bins = 50, root_rank = 0): + def plot_phasespace(self, num_bins=50, root_rank=0): """ - Plot the longitudinal and transverse phase space projections with matplotlib. + Plot the longitudinal and transverse phase space projections with matplotlib. - Parameters - ---------- - self : ImpactXParticleContainer_* - The particle container class in ImpactX - num_bins : int, default=50 - The number of bins for spatial and momentum directions per plot axis. - root_rank : int, default=0 - MPI root rank to reduce to in parallel runs. + Parameters + ---------- + self : ImpactXParticleContainer_* + The particle container class in ImpactX + num_bins : int, default=50 + The number of bins for spatial and momentum directions per plot axis. + root_rank : int, default=0 + MPI root rank to reduce to in parallel runs. - Returns - ------- - A matplotlib figure with containing the plot. - For MPI-parallel ranks, the figure is only created on the root_rank. + Returns + ------- + A matplotlib figure with containing the plot. + For MPI-parallel ranks, the figure is only created on the root_rank. """ - def redistribute(self, arg0: int, arg1: int, arg2: int, arg3: int, arg4: bool) -> None: + def redistribute( + self, arg0: int, arg1: int, arg2: int, arg3: int, arg4: bool + ) -> None: """ Redistribute particles in the current mesh in x, y, z """ @@ -460,15 +494,16 @@ class ImpactXParticleContainer(amrex.space3d.amrex_3d_pybind.ParticleContainer_p """ Get the name of each int SoA component """ + class RefPart: @staticmethod def load_file(ref: RefPart, madx_file): """ - Function that reads elements from a MAD-X file into a list of ImpactX.KnownElements - :param RefPart ref: ImpactX reference particle (passed by reference) - :param madx_file: file name to MAD-X file with beamline elements - :return: list of ImpactX.KnownElements + Function that reads elements from a MAD-X file into a list of ImpactX.KnownElements + :param RefPart ref: ImpactX reference particle (passed by reference) + :param madx_file: file name to MAD-X file with beamline elements + :return: list of ImpactX.KnownElements """ def __init__(self) -> None: @@ -504,8 +539,7 @@ class RefPart: reference charge, in C """ @charge.setter - def charge(self, arg0: float) -> None: - ... + def charge(self, arg0: float) -> None: ... @property def charge_qe(self) -> float: """ @@ -527,8 +561,7 @@ class RefPart: reference rest mass, in kg """ @mass.setter - def mass(self, arg0: float) -> None: - ... + def mass(self, arg0: float) -> None: ... @property def mass_MeV(self) -> float: """ @@ -540,32 +573,28 @@ class RefPart: energy deviation, normalized by rest energy """ @pt.setter - def pt(self, arg0: float) -> None: - ... + def pt(self, arg0: float) -> None: ... @property def px(self) -> float: """ momentum in x, normalized to proper velocity """ @px.setter - def px(self, arg0: float) -> None: - ... + def px(self, arg0: float) -> None: ... @property def py(self) -> float: """ momentum in y, normalized to proper velocity """ @py.setter - def py(self, arg0: float) -> None: - ... + def py(self, arg0: float) -> None: ... @property def pz(self) -> float: """ momentum in z, normalized to proper velocity """ @pz.setter - def pz(self, arg0: float) -> None: - ... + def pz(self, arg0: float) -> None: ... @property def qm_qeeV(self) -> float: """ @@ -582,50 +611,83 @@ class RefPart: integrated orbit path length, in meters """ @s.setter - def s(self, arg0: float) -> None: - ... + def s(self, arg0: float) -> None: ... @property def t(self) -> float: """ clock time * c in meters """ @t.setter - def t(self, arg0: float) -> None: - ... + def t(self, arg0: float) -> None: ... @property def x(self) -> float: """ horizontal position x, in meters """ @x.setter - def x(self, arg0: float) -> None: - ... + def x(self, arg0: float) -> None: ... @property def y(self) -> float: """ vertical position y, in meters """ @y.setter - def y(self, arg0: float) -> None: - ... + def y(self, arg0: float) -> None: ... @property def z(self) -> float: """ longitudinal position y, in meters """ @z.setter - def z(self, arg0: float) -> None: - ... -def coordinate_transformation(pc: ImpactXParticleContainer, direction: CoordSystem) -> None: + def z(self, arg0: float) -> None: ... + +def coordinate_transformation( + pc: ImpactXParticleContainer, direction: CoordSystem +) -> None: """ Transform coordinates from fixed s to fixed to or vice versa. """ -def push(pc: ImpactXParticleContainer, element: elements.Empty | elements.Aperture | elements.Buncher | elements.CFbend | elements.ChrAcc | elements.ChrDrift | elements.ChrPlasmaLens | elements.ChrQuad | elements.ConstF | elements.BeamMonitor | elements.DipEdge | elements.Drift | elements.ExactDrift | elements.ExactSbend | elements.Kicker | elements.Multipole | elements.NonlinearLens | elements.Programmable | elements.PRot | elements.Quad | elements.RFCavity | elements.Sbend | elements.ShortRF | elements.SoftSolenoid | elements.SoftQuadrupole | elements.Sol | elements.TaperedPL | elements.ThinDipole, step: int = 0) -> None: + +def push( + pc: ImpactXParticleContainer, + element: elements.Empty + | elements.Aperture + | elements.Buncher + | elements.CFbend + | elements.ChrAcc + | elements.ChrDrift + | elements.ChrPlasmaLens + | elements.ChrQuad + | elements.ConstF + | elements.BeamMonitor + | elements.DipEdge + | elements.Drift + | elements.ExactDrift + | elements.ExactSbend + | elements.Kicker + | elements.Multipole + | elements.NonlinearLens + | elements.Programmable + | elements.PRot + | elements.Quad + | elements.RFCavity + | elements.Sbend + | elements.ShortRF + | elements.SoftSolenoid + | elements.SoftQuadrupole + | elements.Sol + | elements.TaperedPL + | elements.ThinDipole, + step: int = 0, +) -> None: """ Push particles through an element """ -__author__: str = 'Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al.' -__license__: str = 'BSD-3-Clause-LBNL' -__version__: str = '24.08' + +__author__: str = ( + "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al." +) +__license__: str = "BSD-3-Clause-LBNL" +__version__: str = "24.08" s: CoordSystem # value = t: CoordSystem # value = diff --git a/src/python/impactx/impactx_pybind/distribution.pyi b/src/python/impactx/impactx_pybind/distribution.pyi index f477b7a3b..105a1c40a 100644 --- a/src/python/impactx/impactx_pybind/distribution.pyi +++ b/src/python/impactx/impactx_pybind/distribution.pyi @@ -1,60 +1,166 @@ """ Particle beam distributions in ImpactX """ + from __future__ import annotations -__all__ = ['Empty', 'Gaussian', 'KVdist', 'Kurth4D', 'Kurth6D', 'Semigaussian', 'Thermal', 'Triangle', 'Waterbag'] + +__all__ = [ + "Empty", + "Gaussian", + "KVdist", + "Kurth4D", + "Kurth6D", + "Semigaussian", + "Thermal", + "Triangle", + "Waterbag", +] + class Empty: def __init__(self) -> None: """ Sets all values to zero. """ + class Gaussian: - def __init__(self, lambdaX: float, lambdaY: float, lambdaT: float, lambdaPx: float, lambdaPy: float, lambdaPt: float, muxpx: float = 0.0, muypy: float = 0.0, mutpt: float = 0.0) -> None: + def __init__( + self, + lambdaX: float, + lambdaY: float, + lambdaT: float, + lambdaPx: float, + lambdaPy: float, + lambdaPt: float, + muxpx: float = 0.0, + muypy: float = 0.0, + mutpt: float = 0.0, + ) -> None: """ A 6D Gaussian distribution """ + class KVdist: - def __init__(self, lambdaX: float, lambdaY: float, lambdaT: float, lambdaPx: float, lambdaPy: float, lambdaPt: float, muxpx: float = 0.0, muypy: float = 0.0, mutpt: float = 0.0) -> None: + def __init__( + self, + lambdaX: float, + lambdaY: float, + lambdaT: float, + lambdaPx: float, + lambdaPy: float, + lambdaPt: float, + muxpx: float = 0.0, + muypy: float = 0.0, + mutpt: float = 0.0, + ) -> None: """ A K-V distribution transversely + a uniform distribution in t + a Gaussian distribution in pt """ + class Kurth4D: - def __init__(self, lambdaX: float, lambdaY: float, lambdaT: float, lambdaPx: float, lambdaPy: float, lambdaPt: float, muxpx: float = 0.0, muypy: float = 0.0, mutpt: float = 0.0) -> None: + def __init__( + self, + lambdaX: float, + lambdaY: float, + lambdaT: float, + lambdaPx: float, + lambdaPy: float, + lambdaPt: float, + muxpx: float = 0.0, + muypy: float = 0.0, + mutpt: float = 0.0, + ) -> None: """ A 4D Kurth distribution transversely + a uniform distribution in t + a Gaussian distribution in pt """ + class Kurth6D: - def __init__(self, lambdaX: float, lambdaY: float, lambdaT: float, lambdaPx: float, lambdaPy: float, lambdaPt: float, muxpx: float = 0.0, muypy: float = 0.0, mutpt: float = 0.0) -> None: + def __init__( + self, + lambdaX: float, + lambdaY: float, + lambdaT: float, + lambdaPx: float, + lambdaPy: float, + lambdaPt: float, + muxpx: float = 0.0, + muypy: float = 0.0, + mutpt: float = 0.0, + ) -> None: """ A 6D Kurth distribution R. Kurth, Quarterly of Applied Mathematics vol. 32, pp. 325-329 (1978) C. Mitchell, K. Hwang and R. D. Ryne, IPAC2021, WEPAB248 (2021) """ + class Semigaussian: - def __init__(self, lambdaX: float, lambdaY: float, lambdaT: float, lambdaPx: float, lambdaPy: float, lambdaPt: float, muxpx: float = 0.0, muypy: float = 0.0, mutpt: float = 0.0) -> None: + def __init__( + self, + lambdaX: float, + lambdaY: float, + lambdaT: float, + lambdaPx: float, + lambdaPy: float, + lambdaPt: float, + muxpx: float = 0.0, + muypy: float = 0.0, + mutpt: float = 0.0, + ) -> None: """ A 6D Semi-Gaussian distribution (uniform in position, Gaussian in momentum). """ + class Thermal: - def __init__(self, k: float, kT: float, kT_halo: float, normalize: float, normalize_halo: float, halo: float = 0.0) -> None: + def __init__( + self, + k: float, + kT: float, + kT_halo: float, + normalize: float, + normalize_halo: float, + halo: float = 0.0, + ) -> None: """ A stationary thermal or bithermal distribution R. D. Ryne, J. Qiang, and A. Adelmann, in Proc. EPAC2004, pp. 1942-1944 (2004) """ + class Triangle: - def __init__(self, lambdaX: float, lambdaY: float, lambdaT: float, lambdaPx: float, lambdaPy: float, lambdaPt: float, muxpx: float = 0.0, muypy: float = 0.0, mutpt: float = 0.0) -> None: + def __init__( + self, + lambdaX: float, + lambdaY: float, + lambdaT: float, + lambdaPx: float, + lambdaPy: float, + lambdaPt: float, + muxpx: float = 0.0, + muypy: float = 0.0, + mutpt: float = 0.0, + ) -> None: """ A triangle distribution for laser-plasma acceleration related applications. A ramped, triangular current profile with a Gaussian energy spread (possibly correlated). The transverse distribution is a 4D waterbag. """ + class Waterbag: - def __init__(self, lambdaX: float, lambdaY: float, lambdaT: float, lambdaPx: float, lambdaPy: float, lambdaPt: float, muxpx: float = 0.0, muypy: float = 0.0, mutpt: float = 0.0) -> None: + def __init__( + self, + lambdaX: float, + lambdaY: float, + lambdaT: float, + lambdaPx: float, + lambdaPy: float, + lambdaPt: float, + muxpx: float = 0.0, + muypy: float = 0.0, + mutpt: float = 0.0, + ) -> None: """ A 6D Waterbag distribution """ diff --git a/src/python/impactx/impactx_pybind/elements.pyi b/src/python/impactx/impactx_pybind/elements.pyi index 16af10de2..2be7826e0 100644 --- a/src/python/impactx/impactx_pybind/elements.pyi +++ b/src/python/impactx/impactx_pybind/elements.pyi @@ -1,10 +1,46 @@ """ Accelerator lattice elements in ImpactX """ + from __future__ import annotations import impactx.impactx_pybind import typing -__all__ = ['Alignment', 'Aperture', 'BeamMonitor', 'Buncher', 'CFbend', 'ChrAcc', 'ChrDrift', 'ChrPlasmaLens', 'ChrQuad', 'ConstF', 'DipEdge', 'Drift', 'Empty', 'ExactDrift', 'ExactSbend', 'Kicker', 'KnownElementsList', 'Multipole', 'NonlinearLens', 'PRot', 'Programmable', 'Quad', 'RFCavity', 'Sbend', 'ShortRF', 'SoftQuadrupole', 'SoftSolenoid', 'Sol', 'TaperedPL', 'Thick', 'Thin', 'ThinDipole'] + +__all__ = [ + "Alignment", + "Aperture", + "BeamMonitor", + "Buncher", + "CFbend", + "ChrAcc", + "ChrDrift", + "ChrPlasmaLens", + "ChrQuad", + "ConstF", + "DipEdge", + "Drift", + "Empty", + "ExactDrift", + "ExactSbend", + "Kicker", + "KnownElementsList", + "Multipole", + "NonlinearLens", + "PRot", + "Programmable", + "Quad", + "RFCavity", + "Sbend", + "ShortRF", + "SoftQuadrupole", + "SoftSolenoid", + "Sol", + "TaperedPL", + "Thick", + "Thin", + "ThinDipole", +] + class Alignment: def __init__(self) -> None: """ @@ -16,32 +52,39 @@ class Alignment: horizontal translation error in m """ @dx.setter - def dx(self, arg1: float) -> None: - ... + def dx(self, arg1: float) -> None: ... @property def dy(self) -> float: """ vertical translation error in m """ @dy.setter - def dy(self, arg1: float) -> None: - ... + def dy(self, arg1: float) -> None: ... @property def rotation(self) -> float: """ rotation error in the transverse plane in degree """ @rotation.setter - def rotation(self, arg1: float) -> None: - ... + def rotation(self, arg1: float) -> None: ... + class Aperture(Thin, Alignment): - def __init__(self, xmax: float, ymax: float, shape: str = 'rectangular', dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, + xmax: float, + ymax: float, + shape: str = "rectangular", + dx: float = 0, + dy: float = 0, + rotation: float = 0, + ) -> None: """ A short collimator element applying a transverse aperture boundary. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -51,30 +94,32 @@ class Aperture(Thin, Alignment): aperture type (rectangular, elliptical) """ @shape.setter - def shape(self, arg1: str) -> None: - ... + def shape(self, arg1: str) -> None: ... @property def xmax(self) -> float: """ maximum horizontal coordinate """ @xmax.setter - def xmax(self, arg1: float) -> None: - ... + def xmax(self, arg1: float) -> None: ... @property def ymax(self) -> float: """ maximum vertical coordinate """ @ymax.setter - def ymax(self, arg1: float) -> None: - ... + def ymax(self, arg1: float) -> None: ... + class BeamMonitor(Thin): - def __init__(self, name: str, backend: str = 'default', encoding: str = 'g') -> None: + def __init__( + self, name: str, backend: str = "default", encoding: str = "g" + ) -> None: """ This element writes the particle beam out to openPMD data. """ - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -85,8 +130,7 @@ class BeamMonitor(Thin): Horizontal and vertical values must be equal. """ @alpha.setter - def alpha(self, arg1: float) -> None: - ... + def alpha(self, arg1: float) -> None: ... @property def beta(self) -> float: """ @@ -94,16 +138,14 @@ class BeamMonitor(Thin): Horizontal and vertical values must be equal. """ @beta.setter - def beta(self, arg1: float) -> None: - ... + def beta(self, arg1: float) -> None: ... @property def cn(self) -> float: """ Scale factor (in meters^(1/2)) of the IOTA nonlinear magnetic insert element used for computing H and I. """ @cn.setter - def cn(self, arg1: float) -> None: - ... + def cn(self, arg1: float) -> None: ... @property def name(self) -> str: """ @@ -115,24 +157,26 @@ class BeamMonitor(Thin): Compute and output the invariants H and I within the nonlinear magnetic insert element """ @nonlinear_lens_invariants.setter - def nonlinear_lens_invariants(self, arg1: bool) -> None: - ... + def nonlinear_lens_invariants(self, arg1: bool) -> None: ... @property def tn(self) -> float: """ Dimensionless strength of the IOTA nonlinear magnetic insert element used for computing H and I. """ @tn.setter - def tn(self, arg1: float) -> None: - ... + def tn(self, arg1: float) -> None: ... + class Buncher(Thin, Alignment): - def __init__(self, V: float, k: float, dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, V: float, k: float, dx: float = 0, dy: float = 0, rotation: float = 0 + ) -> None: """ A short linear RF cavity element at zero-crossing for bunching. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -142,24 +186,33 @@ class Buncher(Thin, Alignment): Normalized RF voltage drop V = Emax*L/(c*Brho) """ @V.setter - def V(self, arg1: float) -> None: - ... + def V(self, arg1: float) -> None: ... @property def k(self) -> float: """ Wavenumber of RF in 1/m """ @k.setter - def k(self, arg1: float) -> None: - ... + def k(self, arg1: float) -> None: ... + class CFbend(Thick, Alignment): - def __init__(self, ds: float, rc: float, k: float, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + rc: float, + k: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ An ideal combined function bend (sector bend with quadrupole component). """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -169,24 +222,33 @@ class CFbend(Thick, Alignment): Quadrupole strength in m^(-2) (MADX convention) = (gradient in T/m) / (rigidity in T-m) k > 0 horizontal focusing k < 0 horizontal defocusing """ @k.setter - def k(self, arg1: float) -> None: - ... + def k(self, arg1: float) -> None: ... @property def rc(self) -> float: """ Radius of curvature in m """ @rc.setter - def rc(self, arg1: float) -> None: - ... + def rc(self, arg1: float) -> None: ... + class ChrAcc(Thick, Alignment): - def __init__(self, ds: float, ez: float, bz: float, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + ez: float, + bz: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ A region of Uniform Acceleration, with chromatic effects included. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -196,35 +258,53 @@ class ChrAcc(Thick, Alignment): magnetic field strength in 1/m """ @bz.setter - def bz(self, arg1: float) -> None: - ... + def bz(self, arg1: float) -> None: ... @property def ez(self) -> float: """ electric field strength in 1/m """ @ez.setter - def ez(self, arg1: float) -> None: - ... + def ez(self, arg1: float) -> None: ... + class ChrDrift(Thick, Alignment): - def __init__(self, ds: float, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ A Drift with chromatic effects included. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ + class ChrPlasmaLens(Thick, Alignment): - def __init__(self, ds: float, k: float, unit: int = 0, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + k: float, + unit: int = 0, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ An active Plasma Lens with chromatic effects included. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -234,24 +314,33 @@ class ChrPlasmaLens(Thick, Alignment): focusing strength in 1/m^2 (or T/m) """ @k.setter - def k(self, arg1: float) -> None: - ... + def k(self, arg1: float) -> None: ... @property def unit(self) -> int: """ unit specification for focusing strength """ @unit.setter - def unit(self, arg1: int) -> None: - ... + def unit(self, arg1: int) -> None: ... + class ChrQuad(Thick, Alignment): - def __init__(self, ds: float, k: float, unit: int = 0, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + k: float, + unit: int = 0, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ A Quadrupole magnet with chromatic effects included. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -261,24 +350,34 @@ class ChrQuad(Thick, Alignment): quadrupole strength in 1/m^2 (or T/m) """ @k.setter - def k(self, arg1: float) -> None: - ... + def k(self, arg1: float) -> None: ... @property def unit(self) -> int: """ unit specification for quad strength """ @unit.setter - def unit(self, arg1: int) -> None: - ... + def unit(self, arg1: int) -> None: ... + class ConstF(Thick, Alignment): - def __init__(self, ds: float, kx: float, ky: float, kt: float, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + kx: float, + ky: float, + kt: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ A linear Constant Focusing element. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -288,32 +387,40 @@ class ConstF(Thick, Alignment): focusing t strength in 1/m """ @kt.setter - def kt(self, arg1: float) -> None: - ... + def kt(self, arg1: float) -> None: ... @property def kx(self) -> float: """ focusing x strength in 1/m """ @kx.setter - def kx(self, arg1: float) -> None: - ... + def kx(self, arg1: float) -> None: ... @property def ky(self) -> float: """ focusing y strength in 1/m """ @ky.setter - def ky(self, arg1: float) -> None: - ... + def ky(self, arg1: float) -> None: ... + class DipEdge(Thin, Alignment): - def __init__(self, psi: float, rc: float, g: float, K2: float, dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, + psi: float, + rc: float, + g: float, + K2: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + ) -> None: """ Edge focusing associated with bend entry or exit. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -323,73 +430,100 @@ class DipEdge(Thin, Alignment): Fringe field integral (unitless) """ @K2.setter - def K2(self, arg1: float) -> None: - ... + def K2(self, arg1: float) -> None: ... @property def g(self) -> float: """ Gap parameter in m """ @g.setter - def g(self, arg1: float) -> None: - ... + def g(self, arg1: float) -> None: ... @property def psi(self) -> float: """ Pole face angle in rad """ @psi.setter - def psi(self, arg1: float) -> None: - ... + def psi(self, arg1: float) -> None: ... @property def rc(self) -> float: """ Radius of curvature in m """ @rc.setter - def rc(self, arg1: float) -> None: - ... + def rc(self, arg1: float) -> None: ... + class Drift(Thick, Alignment): - def __init__(self, ds: float, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ A drift. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ + class Empty(Thin): def __init__(self) -> None: """ This element does nothing. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ + class ExactDrift(Thick, Alignment): - def __init__(self, ds: float, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ A Drift using the exact nonlinear map. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ + class ExactSbend(Thick, Alignment): - def __init__(self, ds: float, phi: float, B: float = 0.0, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + phi: float, + B: float = 0.0, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ An ideal sector bend using the exact nonlinear map. When B = 0, the reference bending radius is defined by r0 = length / (angle in rad), corresponding to a magnetic field of B = rigidity / r0; otherwise the reference bending radius is defined by r0 = rigidity / B. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -399,24 +533,32 @@ class ExactSbend(Thick, Alignment): Magnetic field in Tesla; when B = 0 (default), the reference bending radius is defined by r0 = length / (angle in rad), corresponding to a magnetic field of B = rigidity / r0; otherwise the reference bending radius is defined by r0 = rigidity / B """ @B.setter - def B(self, arg1: float) -> None: - ... + def B(self, arg1: float) -> None: ... @property def phi(self) -> float: """ Bend angle in degrees """ @phi.setter - def phi(self, arg1: float) -> None: - ... + def phi(self, arg1: float) -> None: ... + class Kicker(Thin, Alignment): - def __init__(self, xkick: float, ykick: float, unit: str = 'dimensionless', dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, + xkick: float, + ykick: float, + unit: str = "dimensionless", + dx: float = 0, + dy: float = 0, + rotation: float = 0, + ) -> None: """ A thin transverse kicker element. Kicks are for unit "dimensionless" or in "T-m". """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -426,33 +568,119 @@ class Kicker(Thin, Alignment): horizontal kick strength (dimensionless OR T-m) """ @xkick.setter - def xkick(self, arg1: float) -> None: - ... + def xkick(self, arg1: float) -> None: ... @property def ykick(self) -> float: """ vertical kick strength (dimensionless OR T-m) """ @ykick.setter - def ykick(self, arg1: float) -> None: - ... + def ykick(self, arg1: float) -> None: ... + class KnownElementsList: @typing.overload - def __init__(self) -> None: - ... + def __init__(self) -> None: ... @typing.overload - def __init__(self, arg0: Empty | Aperture | Buncher | CFbend | ChrAcc | ChrDrift | ChrPlasmaLens | ChrQuad | ConstF | BeamMonitor | DipEdge | Drift | ExactDrift | ExactSbend | Kicker | Multipole | NonlinearLens | Programmable | PRot | Quad | RFCavity | Sbend | ShortRF | SoftSolenoid | SoftQuadrupole | Sol | TaperedPL | ThinDipole) -> None: - ... + def __init__( + self, + arg0: Empty + | Aperture + | Buncher + | CFbend + | ChrAcc + | ChrDrift + | ChrPlasmaLens + | ChrQuad + | ConstF + | BeamMonitor + | DipEdge + | Drift + | ExactDrift + | ExactSbend + | Kicker + | Multipole + | NonlinearLens + | Programmable + | PRot + | Quad + | RFCavity + | Sbend + | ShortRF + | SoftSolenoid + | SoftQuadrupole + | Sol + | TaperedPL + | ThinDipole, + ) -> None: ... @typing.overload - def __init__(self, arg0: list) -> None: - ... - def __iter__(self) -> typing.Iterator[Empty | Aperture | Buncher | CFbend | ChrAcc | ChrDrift | ChrPlasmaLens | ChrQuad | ConstF | BeamMonitor | DipEdge | Drift | ExactDrift | ExactSbend | Kicker | Multipole | NonlinearLens | Programmable | PRot | Quad | RFCavity | Sbend | ShortRF | SoftSolenoid | SoftQuadrupole | Sol | TaperedPL | ThinDipole]: - ... + def __init__(self, arg0: list) -> None: ... + def __iter__( + self, + ) -> typing.Iterator[ + Empty + | Aperture + | Buncher + | CFbend + | ChrAcc + | ChrDrift + | ChrPlasmaLens + | ChrQuad + | ConstF + | BeamMonitor + | DipEdge + | Drift + | ExactDrift + | ExactSbend + | Kicker + | Multipole + | NonlinearLens + | Programmable + | PRot + | Quad + | RFCavity + | Sbend + | ShortRF + | SoftSolenoid + | SoftQuadrupole + | Sol + | TaperedPL + | ThinDipole + ]: ... def __len__(self) -> int: """ The length of the list. """ - def append(self, arg0: Empty | Aperture | Buncher | CFbend | ChrAcc | ChrDrift | ChrPlasmaLens | ChrQuad | ConstF | BeamMonitor | DipEdge | Drift | ExactDrift | ExactSbend | Kicker | Multipole | NonlinearLens | Programmable | PRot | Quad | RFCavity | Sbend | ShortRF | SoftSolenoid | SoftQuadrupole | Sol | TaperedPL | ThinDipole) -> None: + def append( + self, + arg0: Empty + | Aperture + | Buncher + | CFbend + | ChrAcc + | ChrDrift + | ChrPlasmaLens + | ChrQuad + | ConstF + | BeamMonitor + | DipEdge + | Drift + | ExactDrift + | ExactSbend + | Kicker + | Multipole + | NonlinearLens + | Programmable + | PRot + | Quad + | RFCavity + | Sbend + | ShortRF + | SoftSolenoid + | SoftQuadrupole + | Sol + | TaperedPL + | ThinDipole, + ) -> None: """ Add a single element to the list. """ @@ -470,20 +698,29 @@ class KnownElementsList: """ Add a list of elements to the list. """ - def load_file(self, madx_file, nslice = 1): - ... + def load_file(self, madx_file, nslice=1): ... def pop_back(self) -> None: """ Return and remove the last element of the list. """ + class Multipole(Thin, Alignment): - def __init__(self, multipole: int, K_normal: float, K_skew: float, dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, + multipole: int, + K_normal: float, + K_skew: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + ) -> None: """ A general thin multipole element. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -493,32 +730,38 @@ class Multipole(Thin, Alignment): Integrated normal multipole coefficient (1/meter^m) """ @K_normal.setter - def K_normal(self, arg1: float) -> None: - ... + def K_normal(self, arg1: float) -> None: ... @property def K_skew(self) -> float: """ Integrated skew multipole coefficient (1/meter^m) """ @K_skew.setter - def K_skew(self, arg1: float) -> None: - ... + def K_skew(self, arg1: float) -> None: ... @property def multipole(self) -> int: """ index m (m=1 dipole, m=2 quadrupole, m=3 sextupole etc.) """ @multipole.setter - def multipole(self, arg1: float) -> None: - ... + def multipole(self, arg1: float) -> None: ... + class NonlinearLens(Thin, Alignment): - def __init__(self, knll: float, cnll: float, dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, + knll: float, + cnll: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + ) -> None: """ Single short segment of the nonlinear magnetic insert element. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -528,24 +771,24 @@ class NonlinearLens(Thin, Alignment): distance of singularities from the origin (m) """ @cnll.setter - def cnll(self, arg1: float) -> None: - ... + def cnll(self, arg1: float) -> None: ... @property def knll(self) -> float: """ integrated strength of the nonlinear lens (m) """ @knll.setter - def knll(self, arg1: float) -> None: - ... + def knll(self, arg1: float) -> None: ... + class PRot(Thin): def __init__(self, phi_in: float, phi_out: float) -> None: """ An exact pole-face rotation in the x-z plane. Both angles are in degrees. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -555,16 +798,15 @@ class PRot(Thin): angle of the reference particle with respect to the longitudinal (z) axis in the original frame in degrees """ @phi_in.setter - def phi_in(self, arg1: float) -> None: - ... + def phi_in(self, arg1: float) -> None: ... @property def phi_out(self) -> float: """ angle of the reference particle with respect to the longitudinal (z) axis in the rotated frame in degrees """ @phi_out.setter - def phi_out(self, arg1: float) -> None: - ... + def phi_out(self, arg1: float) -> None: ... + class Programmable: ds: float nslice: int @@ -572,48 +814,72 @@ class Programmable: """ A programmable beam optics element. """ - def __repr__(self) -> str: - ... + def __repr__(self) -> str: ... @property - def beam_particles(self) -> typing.Callable[[impactx.impactx_pybind.ImpactXParIter, impactx.impactx_pybind.RefPart], None]: + def beam_particles( + self, + ) -> typing.Callable[ + [impactx.impactx_pybind.ImpactXParIter, impactx.impactx_pybind.RefPart], None + ]: """ hook for beam particles (pti, RefPart) """ @beam_particles.setter - def beam_particles(self, arg1: typing.Callable[[impactx.impactx_pybind.ImpactXParIter, impactx.impactx_pybind.RefPart], None]) -> None: - ... - @property - def push(self) -> typing.Callable[[impactx.impactx_pybind.ImpactXParticleContainer, int], None]: + def beam_particles( + self, + arg1: typing.Callable[ + [impactx.impactx_pybind.ImpactXParIter, impactx.impactx_pybind.RefPart], + None, + ], + ) -> None: ... + @property + def push( + self, + ) -> typing.Callable[[impactx.impactx_pybind.ImpactXParticleContainer, int], None]: """ hook for push of whole container (pc, step) """ @push.setter - def push(self, arg1: typing.Callable[[impactx.impactx_pybind.ImpactXParticleContainer, int], None]) -> None: - ... + def push( + self, + arg1: typing.Callable[ + [impactx.impactx_pybind.ImpactXParticleContainer, int], None + ], + ) -> None: ... @property def ref_particle(self) -> typing.Callable[[impactx.impactx_pybind.RefPart], None]: """ hook for reference particle (RefPart) """ @ref_particle.setter - def ref_particle(self, arg1: typing.Callable[[impactx.impactx_pybind.RefPart], None]) -> None: - ... + def ref_particle( + self, arg1: typing.Callable[[impactx.impactx_pybind.RefPart], None] + ) -> None: ... @property def threadsafe(self) -> bool: """ allow threading via OpenMP for the particle iterator loop, default=False (note: if OMP backend is active) """ @threadsafe.setter - def threadsafe(self, arg1: bool) -> None: - ... + def threadsafe(self, arg1: bool) -> None: ... + class Quad(Thick, Alignment): - def __init__(self, ds: float, k: float, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + k: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ A Quadrupole magnet. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -623,16 +889,30 @@ class Quad(Thick, Alignment): Quadrupole strength in m^(-2) (MADX convention) = (gradient in T/m) / (rigidity in T-m) k > 0 horizontal focusing k < 0 horizontal defocusing """ @k.setter - def k(self, arg1: float) -> None: - ... + def k(self, arg1: float) -> None: ... + class RFCavity(Thick, Alignment): - def __init__(self, ds: float, escale: float, freq: float, phase: float, cos_coefficients: list[float], sin_coefficients: list[float], dx: float = 0, dy: float = 0, rotation: float = 0, mapsteps: int = 1, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + escale: float, + freq: float, + phase: float, + cos_coefficients: list[float], + sin_coefficients: list[float], + dx: float = 0, + dy: float = 0, + rotation: float = 0, + mapsteps: int = 1, + nslice: int = 1, + ) -> None: """ An RF cavity. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -642,40 +922,46 @@ class RFCavity(Thick, Alignment): scaling factor for on-axis RF electric field in 1/m = (peak on-axis electric field Ez in MV/m) / (particle rest energy in MeV) """ @escale.setter - def escale(self, arg1: float) -> None: - ... + def escale(self, arg1: float) -> None: ... @property def freq(self) -> float: """ RF frequency in Hz """ @freq.setter - def freq(self, arg1: float) -> None: - ... + def freq(self, arg1: float) -> None: ... @property def mapsteps(self) -> int: """ number of integration steps per slice used for map and reference particle push in applied fields """ @mapsteps.setter - def mapsteps(self, arg1: int) -> None: - ... + def mapsteps(self, arg1: int) -> None: ... @property def phase(self) -> float: """ RF driven phase in degrees """ @phase.setter - def phase(self, arg1: float) -> None: - ... + def phase(self, arg1: float) -> None: ... + class Sbend(Thick, Alignment): - def __init__(self, ds: float, rc: float, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + rc: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ An ideal sector bend. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -685,16 +971,25 @@ class Sbend(Thick, Alignment): Radius of curvature in m """ @rc.setter - def rc(self, arg1: float) -> None: - ... + def rc(self, arg1: float) -> None: ... + class ShortRF(Thin, Alignment): - def __init__(self, V: float, freq: float, phase: float = -90.0, dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, + V: float, + freq: float, + phase: float = -90.0, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + ) -> None: """ A short RF cavity element. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -704,32 +999,42 @@ class ShortRF(Thin, Alignment): Normalized RF voltage V = maximum energy gain/(m*c^2) """ @V.setter - def V(self, arg1: float) -> None: - ... + def V(self, arg1: float) -> None: ... @property def freq(self) -> float: """ RF frequency in Hz """ @freq.setter - def freq(self, arg1: float) -> None: - ... + def freq(self, arg1: float) -> None: ... @property def phase(self) -> float: """ RF synchronous phase in degrees (phase = 0 corresponds to maximum energy gain, phase = -90 corresponds go zero energy gain for bunching) """ @phase.setter - def phase(self, arg1: float) -> None: - ... + def phase(self, arg1: float) -> None: ... + class SoftQuadrupole(Thick, Alignment): - def __init__(self, ds: float, gscale: float, cos_coefficients: list[float], sin_coefficients: list[float], dx: float = 0, dy: float = 0, rotation: float = 0, mapsteps: int = 1, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + gscale: float, + cos_coefficients: list[float], + sin_coefficients: list[float], + dx: float = 0, + dy: float = 0, + rotation: float = 0, + mapsteps: int = 1, + nslice: int = 1, + ) -> None: """ A soft-edge quadrupole. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -739,24 +1044,36 @@ class SoftQuadrupole(Thick, Alignment): Scaling factor for on-axis field gradient in inverse meters """ @gscale.setter - def gscale(self, arg1: float) -> None: - ... + def gscale(self, arg1: float) -> None: ... @property def mapsteps(self) -> int: """ number of integration steps per slice used for map and reference particle push in applied fields """ @mapsteps.setter - def mapsteps(self, arg1: int) -> None: - ... + def mapsteps(self, arg1: int) -> None: ... + class SoftSolenoid(Thick, Alignment): - def __init__(self, ds: float, bscale: float, cos_coefficients: list[float], sin_coefficients: list[float], unit: float = 0, dx: float = 0, dy: float = 0, rotation: float = 0, mapsteps: int = 1, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + bscale: float, + cos_coefficients: list[float], + sin_coefficients: list[float], + unit: float = 0, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + mapsteps: int = 1, + nslice: int = 1, + ) -> None: """ A soft-edge solenoid. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -766,32 +1083,39 @@ class SoftSolenoid(Thick, Alignment): Scaling factor for on-axis magnetic field Bz in inverse meters (if unit = 0) or magnetic field Bz in T (SI units, if unit = 1) """ @bscale.setter - def bscale(self, arg1: float) -> None: - ... + def bscale(self, arg1: float) -> None: ... @property def mapsteps(self) -> int: """ number of integration steps per slice used for map and reference particle push in applied fields """ @mapsteps.setter - def mapsteps(self, arg1: int) -> None: - ... + def mapsteps(self, arg1: int) -> None: ... @property def unit(self) -> int: """ specification of units for scaling of the on-axis longitudinal magnetic field """ @unit.setter - def unit(self, arg1: float) -> None: - ... + def unit(self, arg1: float) -> None: ... + class Sol(Thick, Alignment): - def __init__(self, ds: float, ks: float, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + ks: float, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + nslice: int = 1, + ) -> None: """ An ideal hard-edge Solenoid magnet. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -801,10 +1125,18 @@ class Sol(Thick, Alignment): Solenoid strength in m^(-1) (MADX convention) in (magnetic field Bz in T) / (rigidity in T-m) """ @ks.setter - def ks(self, arg1: float) -> None: - ... + def ks(self, arg1: float) -> None: ... + class TaperedPL(Thin, Alignment): - def __init__(self, k: float, taper: float, unit: int = 0, dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, + k: float, + taper: float, + unit: int = 0, + dx: float = 0, + dy: float = 0, + rotation: float = 0, + ) -> None: """ A thin nonlinear plasma lens with transverse (horizontal) taper @@ -814,9 +1146,10 @@ class TaperedPL(Thin, Alignment): where :math:`g` is the (linear) field gradient in T/m and :math:`D_x` is the targeted horizontal dispersion in m. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -826,24 +1159,22 @@ class TaperedPL(Thin, Alignment): integrated focusing strength in m^(-1) (if unit = 0) or integrated focusing strength in T (if unit = 1) """ @k.setter - def k(self, arg1: float) -> None: - ... + def k(self, arg1: float) -> None: ... @property def taper(self) -> float: """ horizontal taper parameter in m^(-1) = 1 / (target horizontal dispersion in m) """ @taper.setter - def taper(self, arg1: float) -> None: - ... + def taper(self, arg1: float) -> None: ... @property def unit(self) -> int: """ specification of units for plasma lens focusing strength """ @unit.setter - def unit(self, arg1: int) -> None: - ... + def unit(self, arg1: int) -> None: ... + class Thick: def __init__(self, ds: float, nslice: float = 1) -> None: """ @@ -855,16 +1186,15 @@ class Thick: segment length in m """ @ds.setter - def ds(self, arg1: float) -> None: - ... + def ds(self, arg1: float) -> None: ... @property def nslice(self) -> int: """ number of slices used for the application of space charge """ @nslice.setter - def nslice(self, arg1: int) -> None: - ... + def nslice(self, arg1: int) -> None: ... + class Thin: def __init__(self) -> None: """ @@ -880,14 +1210,18 @@ class Thin: """ number of slices used for the application of space charge """ + class ThinDipole(Thin, Alignment): - def __init__(self, theta: float, rc: float, dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, theta: float, rc: float, dx: float = 0, dy: float = 0, rotation: float = 0 + ) -> None: """ A thin kick model of a dipole bend. """ - def __repr__(self) -> str: - ... - def push(self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0) -> None: + def __repr__(self) -> str: ... + def push( + self, pc: impactx.impactx_pybind.ImpactXParticleContainer, step: int = 0 + ) -> None: """ Push first the reference particle, then all other particles. """ @@ -897,13 +1231,11 @@ class ThinDipole(Thin, Alignment): Effective curvature radius (meters) """ @rc.setter - def rc(self, arg1: float) -> None: - ... + def rc(self, arg1: float) -> None: ... @property def theta(self) -> float: """ Bend angle (degrees) """ @theta.setter - def theta(self, arg1: float) -> None: - ... + def theta(self, arg1: float) -> None: ... diff --git a/src/python/impactx/madx_to_impactx.pyi b/src/python/impactx/madx_to_impactx.pyi index 004d42e70..c5d68f3e6 100644 --- a/src/python/impactx/madx_to_impactx.pyi +++ b/src/python/impactx/madx_to_impactx.pyi @@ -5,62 +5,81 @@ from impactx.impactx_pybind import RefPart from impactx.impactx_pybind import elements import typing import warnings as warnings -__all__ = ['MADXParser', 'RefPart', 'beam', 'elements', 'lattice', 'read_beam', 'read_lattice', 'sc', 'warnings'] + +__all__ = [ + "MADXParser", + "RefPart", + "beam", + "elements", + "lattice", + "read_beam", + "read_lattice", + "sc", + "warnings", +] + class sc: """ - This class is used in lieu of scipy.constants - to avoid a direct dependency on it. - At the time of writing, this file was the only - one requiring scipy in the ImpactX source outside - of examples. + This class is used in lieu of scipy.constants + to avoid a direct dependency on it. + At the time of writing, this file was the only + one requiring scipy in the ImpactX source outside + of examples. """ + c: typing.ClassVar[float] = 299792458.0 electron_volt: typing.ClassVar[float] = 1.602176634e-19 m_e: typing.ClassVar[float] = 9.1093837015e-31 m_p: typing.ClassVar[float] = 1.67262192369e-27 m_u: typing.ClassVar[float] = 1.6605390666e-27 - physical_constants: typing.ClassVar[dict] = {'electron-muon mass ratio': (0.00483633169, '', 1.1e-10)} -def beam(particle, charge = None, mass = None, energy = None): + physical_constants: typing.ClassVar[dict] = { + "electron-muon mass ratio": (0.00483633169, "", 1.1e-10) + } + +def beam(particle, charge=None, mass=None, energy=None): """ - Function that converts a list of beam parameter dictionaries in the MADXParser format into ImpactX format + Function that converts a list of beam parameter dictionaries in the MADXParser format into ImpactX format - Rules following https://mad.web.cern.ch/mad/releases/5.02.08/madxuguide.pdf pages 55f. + Rules following https://mad.web.cern.ch/mad/releases/5.02.08/madxuguide.pdf pages 55f. - :param str particle: reference particle name - :param float charge: particle charge (proton charge units) - :param float mass: particle mass (electron masses) - :param float energy: total particle energy (GeV) - - MAD-X default: 1 GeV - :return dict: dictionary containing particle and beam attributes in ImpactX units + :param str particle: reference particle name + :param float charge: particle charge (proton charge units) + :param float mass: particle mass (electron masses) + :param float energy: total particle energy (GeV) + - MAD-X default: 1 GeV + :return dict: dictionary containing particle and beam attributes in ImpactX units """ -def lattice(parsed_beamline, nslice = 1): + +def lattice(parsed_beamline, nslice=1): """ - Function that converts a list of elements in the MADXParser format into ImpactX format - :param parsed_beamline: list of dictionaries - :param nslice: number of ds slices per element - :return: list of translated dictionaries + Function that converts a list of elements in the MADXParser format into ImpactX format + :param parsed_beamline: list of dictionaries + :param nslice: number of ds slices per element + :return: list of translated dictionaries """ + def read_beam(ref: impactx.impactx_pybind.RefPart, madx_file): """ - Function that reads elements from a MAD-X file into a list of ImpactX.KnownElements - :param RefPart ref: ImpactX reference particle (passed by reference) - :param madx_file: file name to MAD-X file with beamline elements - :return: list of ImpactX.KnownElements + Function that reads elements from a MAD-X file into a list of ImpactX.KnownElements + :param RefPart ref: ImpactX reference particle (passed by reference) + :param madx_file: file name to MAD-X file with beamline elements + :return: list of ImpactX.KnownElements """ -def read_lattice(madx_file, nslice = 1): + +def read_lattice(madx_file, nslice=1): """ - Function that reads elements from a MAD-X file into a list of ImpactX.KnownElements - :param madx_file: file name to MAD-X file with beamline elements - :param nslice: number of ds slices per element - :return: list of ImpactX.KnownElements + Function that reads elements from a MAD-X file into a list of ImpactX.KnownElements + :param madx_file: file name to MAD-X file with beamline elements + :param nslice: number of ds slices per element + :return: list of ImpactX.KnownElements """