diff --git a/src/python/impactx/impactx_pybind/__init__.pyi b/src/python/impactx/impactx_pybind/__init__.pyi index 508b1ae6d..c990aec77 100644 --- a/src/python/impactx/impactx_pybind/__init__.pyi +++ b/src/python/impactx/impactx_pybind/__init__.pyi @@ -11,6 +11,7 @@ elements """ + from __future__ import annotations import typing @@ -21,12 +22,29 @@ from amrex import space3d as amr from . import distribution, 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: @@ -35,43 +53,49 @@ 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.None | 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. @@ -79,57 +103,64 @@ 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 evolve(self) -> None: """ Run the main simulation loop for a number of steps. """ + def finalize(self) -> None: """ 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. """ + def phi(self, lev: int) -> amrex.space3d.amrex_3d_pybind.MultiFab: """ scalar potential per level """ + def resize_mesh(self) -> None: """ Resize the mesh :py:attr:`~domain` based on the :py:attr:`~dynamic_size` and related parameters. """ + def rho(self, lev: int) -> amrex.space3d.amrex_3d_pybind.MultiFab: """ 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 """ + @property def abort_on_unused_inputs(self) -> int: """ Configure simulation to abort AFTER it has run 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: """ @@ -137,137 +168,141 @@ class ImpactX: has to abort when a warning is recorded. 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: """ If set to 1, immediately prints every warning message 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: """ The minimum number of digits (default: 6) used for the step 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: """ Enable or disable diagnostics generally (default: enabled). 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: """ The currently finest level of mesh-refinement used. This is always less or equal to max_level. """ + @property def lattice(self) -> elements.KnownElementsList: """ 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: """ @@ -275,33 +310,33 @@ 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 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: """ @@ -310,35 +345,46 @@ class ImpactX: By default, diagnostics is performed at the beginning and end of the simulation. 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: """ Controls how much information is printed to the terminal, when running 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 @@ -349,15 +395,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 @@ -368,10 +423,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. @@ -388,235 +456,300 @@ class ImpactXParticleContainer(amrex.space3d.amrex_3d_pybind.ParticleContainer_p :param qm: charge over mass in 1/eV :param bchchg: total charge within a bunch in C """ + def mean_and_std_positions(self) -> tuple[float, float, float, float, float, float]: """ Compute the mean and std of the particle position in each dimension. :return: x_mean, x_std, y_mean, y_std, z_mean, z_std """ + def min_and_max_positions(self) -> tuple[float, float, float, float, float, float]: """ Compute the min and max of the particle position in each dimension. :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 """ + def reduced_beam_characteristics(self) -> dict[str, float]: """ Compute reduced beam characteristics like the position and momentum moments of the particle distribution, as well as emittance and Twiss parameters. """ + def ref_particle(self) -> RefPart: """ Access the reference particle. """ + def set_ref_particle(self, refpart: RefPart) -> None: """ Set reference particle attributes. """ + @property def RealSoA_names(self) -> list[str]: """ Get the name of each ParticleReal SoA component """ + @property def coord_system(self) -> CoordSystem: """ Get the current coordinate system of particles in this container """ + @property def intSoA_names(self) -> list[str]: """ 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: """ This struct stores the reference particle attributes stored in ImpactXParticleContainer. """ + def set_charge_qe(self, charge_qe: float) -> RefPart: """ Set reference particle charge (positive elementary charge) """ + def set_kin_energy_MeV(self, kin_energy_MeV: float) -> RefPart: """ Set reference particle kinetic energy (MeV) """ + def set_mass_MeV(self, mass_MeV: float) -> RefPart: """ Set reference particle rest mass (MeV/c^2) """ + @property def beta(self) -> float: """ Get reference particle relativistic beta """ + @property def beta_gamma(self) -> float: """ Get reference particle beta*gamma """ + @property def charge(self) -> float: """ reference charge, in C """ + @charge.setter - def charge(self, arg0: float) -> None: - ... + def charge(self, arg0: float) -> None: ... @property def charge_qe(self) -> float: """ Get reference particle charge (positive elementary charge) """ + @property def gamma(self) -> float: """ Get reference particle relativistic gamma """ + @property def kin_energy_MeV(self) -> float: """ Get reference particle energy (MeV) """ + @property def mass(self) -> float: """ 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: """ Get reference particle rest mass (MeV/c^2) """ + @property def pt(self) -> float: """ 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: """ Get reference particle charge to mass ratio (charge/eV) """ + @property def rigidity_Tm(self) -> float: """ Get reference particle magnetic rigidity Brho (T*m) """ + @property def s(self) -> float: """ 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.None | 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.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.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.04' + +__author__: str = ( + "Axel Huebl, Chad Mitchell, Ryan Sandberg, Marco Garten, Ji Qiang, et al." +) +__license__: str = "BSD-3-Clause-LBNL" +__version__: str = "24.04" 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 cb7d84779..105a1c40a 100644 --- a/src/python/impactx/impactx_pybind/distribution.pyi +++ b/src/python/impactx/impactx_pybind/distribution.pyi @@ -1,61 +1,166 @@ """ Particle beam distributions in ImpactX """ + from __future__ import annotations -__all__ = ['Gaussian', 'KVdist', 'Kurth4D', 'Kurth6D', 'None', '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 None: - def __init__(self) -> None: - """ - Sets all values to zero. - """ + 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 6d874c471..6a2d57eb1 100644 --- a/src/python/impactx/impactx_pybind/elements.pyi +++ b/src/python/impactx/impactx_pybind/elements.pyi @@ -1,421 +1,702 @@ """ Accelerator lattice elements in ImpactX """ + from __future__ import annotations import typing import impactx.impactx_pybind -__all__ = ['Alignment', 'Aperture', 'BeamMonitor', 'Buncher', 'CFbend', 'ChrAcc', 'ChrDrift', 'ChrPlasmaLens', 'ChrQuad', 'ConstF', 'DipEdge', 'Drift', 'ExactDrift', 'ExactSbend', 'Kicker', 'KnownElementsList', 'Multipole', 'None', 'NonlinearLens', 'PRot', 'Programmable', 'Quad', 'RFCavity', 'Sbend', 'ShortRF', 'SoftQuadrupole', 'SoftSolenoid', 'Sol', '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", + "Thick", + "Thin", + "ThinDipole", +] + class Alignment: def __init__(self) -> None: """ Mixin class for lattice elements with horizontal/vertical alignment errors. """ + @property def dx(self) -> float: """ 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. """ + @property def shape(self) -> str: """ 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. """ + @property def alpha(self) -> float: """ Twiss alpha of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I. 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: """ Twiss beta (in meters) of the bare linear lattice at the location of output for the nonlinear IOTA invariants H and I. 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: """ name of the series """ + @property def nonlinear_lens_invariants(self) -> bool: """ 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. """ + 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. """ + 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. """ + @property def bz(self) -> float: """ 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, units: int = 0, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + k: float, + units: 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. """ + @property def k(self) -> float: """ 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 units(self) -> int: """ unit specification for focusing strength """ + @units.setter - def units(self, arg1: int) -> None: - ... + def units(self, arg1: int) -> None: ... + class ChrQuad(Thick, Alignment): - def __init__(self, ds: float, k: float, units: int = 0, dx: float = 0, dy: float = 0, rotation: float = 0, nslice: int = 1) -> None: + def __init__( + self, + ds: float, + k: float, + units: 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. """ + @property def k(self) -> float: """ 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 units(self) -> int: """ unit specification for quad strength """ + @units.setter - def units(self, arg1: int) -> None: - ... + def units(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. """ + @property def kt(self) -> float: """ 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. """ + 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: """ 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. """ + class Kicker(Thin, Alignment): - def __init__(self, xkick: float, ykick: float, units: str = 'dimensionless', dx: float = 0, dy: float = 0, rotation: float = 0) -> None: + def __init__( + self, + xkick: float, + ykick: float, + units: str = "dimensionless", + dx: float = 0, + dy: float = 0, + rotation: float = 0, + ) -> None: """ A thin transverse kicker element. Kicks are for units "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. """ + class KnownElementsList: @typing.overload - def __init__(self) -> None: - ... + def __init__(self) -> None: ... @typing.overload - def __init__(self, arg0: None | 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 | 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 + | ThinDipole + ), + ) -> None: ... @typing.overload - def __init__(self, arg0: list) -> None: - ... - def __iter__(self) -> typing.Iterator: - ... + def __init__(self, arg0: list) -> None: ... + def __iter__(self) -> typing.Iterator: ... def __len__(self) -> int: """ The length of the list. """ - def append(self, arg0: None | 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 | 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 + | ThinDipole + ), + ) -> None: """ Add a single element to the list. """ + def clear(self) -> None: """ Clear the list to become empty. """ + @typing.overload def extend(self, arg0: KnownElementsList) -> KnownElementsList: """ Add a list of elements to the list. """ + @typing.overload def extend(self, arg0: list) -> 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: - """ - Push first the reference particle, then all other particles. - """ -class None(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 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. """ + 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. """ + class Programmable: ds: float nslice: int @@ -423,161 +704,278 @@ 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: - ... + 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 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. """ + 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. """ + 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. """ + 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. """ + 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. """ + class SoftSolenoid(Thick, Alignment): - def __init__(self, ds: float, bscale: 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, + bscale: 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 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. """ + 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. """ + class Thick: def __init__(self, ds: float, nslice: float = 1) -> None: """ Mixin class for lattice elements with finite length. """ + @property def ds(self) -> float: """ 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: """ Mixin class for lattice elements with zero length. """ + @property def ds(self) -> float: """ segment length in m """ + @property def nslice(self) -> int: """ 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. """