From 8821687f2cab77a00b8260cd75ca8cc8d798cd26 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 20 Oct 2025 13:01:43 -0400 Subject: [PATCH 1/3] Format imports --- py/orbit/teapot/teapot.py | 46 +++++++++++++++------------------------ 1 file changed, 18 insertions(+), 28 deletions(-) diff --git a/py/orbit/teapot/teapot.py b/py/orbit/teapot/teapot.py index 12cd1c34..c3c868a7 100644 --- a/py/orbit/teapot/teapot.py +++ b/py/orbit/teapot/teapot.py @@ -1,46 +1,36 @@ """ Module. Includes classes for all TEAPOT elements. The approach is based on the original ORBIT approach developed by J. Holmes. + +Elements: +- Drift +- Bend +- Quad +- Multipole +- Solenoid +- Kicker +- RingRF +- Monitor """ import sys import os import math -# import teapot base functions from wrapper around C++ functions +from ..lattice import AccLattice +from ..lattice import AccNode +from ..lattice import AccActionsContainer +from ..lattice import AccNodeBunchTracker from ..teapot_base import TPB - -# import the function that creates multidimensional arrays from ..utils import orbitFinalize - -# import general accelerator elements and lattice -from ..lattice import AccLattice, AccNode, AccActionsContainer, AccNodeBunchTracker - -# import the MAD parser to construct lattices of TEAPOT elements. -from ..parsers.mad_parser import MAD_Parser, MAD_LattElement - -# import the MADX parser to construct lattices of TEAPOT elements. -from ..parsers.madx_parser import MADX_Parser, MADX_LattElement - -# import aperture +from ..parsers.mad_parser import MAD_Parser +from ..parsers.mad_parser import MAD_LattElement +from ..parsers.madx_parser import MADX_Parser +from ..parsers.madx_parser import MADX_LattElement from orbit.core.aperture import Aperture - -# monitor from orbit.core.bunch import BunchTwissAnalysis -""" -Drift -Bend -Quad -Multipole -Solenoid -Kicker -RingRF -monitor -""" - - class TEAPOT_Lattice(AccLattice): """ The subclass of the AccLattice class. Shell class for the TEAPOT nodes collection. From f9e57ef433d48709086212f340add5aeb2681fcd Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 20 Oct 2025 13:18:40 -0400 Subject: [PATCH 2/3] Add madx parser example --- .../TEAPOT/inputs/sns_ring_dual_solenoid.lat | 769 ++++++++++++++++++ examples/TEAPOT/test_sns_ring.py | 11 + 2 files changed, 780 insertions(+) create mode 100644 examples/TEAPOT/inputs/sns_ring_dual_solenoid.lat create mode 100644 examples/TEAPOT/test_sns_ring.py diff --git a/examples/TEAPOT/inputs/sns_ring_dual_solenoid.lat b/examples/TEAPOT/inputs/sns_ring_dual_solenoid.lat new file mode 100644 index 00000000..39b53151 --- /dev/null +++ b/examples/TEAPOT/inputs/sns_ring_dual_solenoid.lat @@ -0,0 +1,769 @@ +none = 0; +kvx12 := kdc; +kdc = -3.410059477; +kta11c12 := 0; +brho := 1e9*(pc/c); +pc := sqrt(ek*(ek+2*e0)); +ek := 0.8; +e0 := 0.93827231; +c := 299792458; +khx13 := kfc; +kfc = 3.160004902; +kta10b13 := 0; +kdhta13 := 0; +kdvta13 := 0; +vkck12 := 0; +hkck12 := 0; +vkck13 := 0; +hkck13 := 0; +kdvtb1 := 0; +ksc_b01 := 0; +kssb1_9 = 0; +kvx1 := kdee; +kdee = -2.579585121; +ktb1_9 := 0; +ksv1 := 0; +ksc_b02 := 0; +kssb2_8 = 0; +khx2 := kf; +kf = 3.388663868; +ktb2d8 := 0; +ksh2 := 0; +kdvtb3 := 0; +ksc_b03 := 0; +kvx3 := kd; +kd = -3.733058392; +ktb357 := 0; +ksv3 := chrm3; +chrm3 := 0; +kdhtb4 := 0; +khx4 := kf26; +kf26 := kf*(lf/lf26); +lf := 0.25; +lf26 := 0.2705; +kta4b6 := 0; +ksh4 := chrm4; +chrm4 := 0; +ksv5 := chrm5; +chrm5 := 0; +kvx5 := kd; +kdvtb5 := 0; +ksc_b05 := 0; +ksh6 := chrm6; +chrm6 := 0; +khx6 := kf26; +kdhtb6 := 0; +ksv7 := chrm7; +chrm7 := chrm3; +kvx7 := kd; +kdvtb7 := 0; +ksc_b07 := 0; +ko_b08 := 0; +khx8 := kf; +kdhtb8 := 0; +ksc_b08 := 0; +ko_b09 := 0; +kvx9 := kdee; +kdvtb9 := 0; +ksc_b09 := 0; +kdhtb10 := 0; +kdvtb10 := 0; +khx10 := kfc; +kvx11 := kdc; +ktb11d12 := 0; +kdhtb13 := 0; +kdvtb13 := 0; +kdvtc1 := 0; +ksc_c01 := 0; +kssc1_9 = 0; +ktc1_9 := 0; +kdhtc2 := 0; +kssc2_8 = 0; +ksc_c02 := 0; +kta2c8 := 0; +kdvtc3 := 0; +ksc_c03 := 0; +ktc357 := 0; +kdhtc4 := 0; +ktc4d6 := 0; +kdvtc5 := 0; +ksc_c05 := 0; +kdhtc6 := 0; +kdvtc7 := 0; +ksc_c07 := 0; +ko_c08 := 0; +kdhtc8 := 0; +ksc_c08 := 0; +ko_c09 := 0; +kdvtc9 := 0; +ksc_c09 := 0; +kdhtc10 := 0; +kdvtc10 := 0; +ktc10d13 := 0; +kdhtc13 := 0; +kdvtc13 := 0; +kssol := solfield/brho; +solfield := 0.297599; +ksisol := 1.22174*kssol; +kdvtd1 := 0; +ksc_d01 := 0; +kssd1_9 = 0; +ktd1_9 := 0; +kdhtd2 := 0; +ksc_d02 := 0; +kssd2_8 = 0; +kdvtd3 := 0; +ksc_d03 := 0; +ktd357 := 0; +lsv3 := lsxt; +lsxt := 0.317; +kdhtd4 := 0; +kdvtd5 := 0; +ksc_d05 := 0; +kdhtd6 := 0; +kdvtd7 := 0; +ksc_d07 := 0; +ko_d08 := 0; +kdhtd8 := 0; +ksc_d08 := 0; +ko_d09 := 0; +kdvtd9 := 0; +ksc_d09 := 0; +kdhtd10 := 0; +kdvtd10 := 0; +kdhtd13 := 0; +kdvtd13 := 0; +kdvta1 := 0; +ksc_a01 := 0; +kssa1_9 = 0; +kta1_9 := 0; +kdhta2 := 0; +ksc_a02 := 0; +kssa2_8 = 0; +kdvta3 := 0; +ksc_a03 := 0; +kta357 := 0; +kdhta4 := 0; +kdvta5 := 0; +kdhta6 := 0; +kdvta7 := 0; +ksc_a07 := 0; +ko_a08 := 0; +kdhta8 := 0; +ksc_a08 := 0; +ko_a09 := 0; +kdvta9 := 0; +ksc_a09 := 0; +hkck10 := 0; +vkck10 := 0; +hkck11 := 0; +vkck11 := 0; +kdhta10 := 0; +kdvta10 := 0; +injm1: marker; +dh_a12: sbend,l:= 0.9903829659,angle:= 0.0436,e1:= -0.003,e2:= 0.0466; +dh_a13: sbend,l:= 0.8903221964,angle:= -0.0466,e1:= -0.0466,e2:= -6.938893904e-18; +qtv_a12: quadrupole,l:= 0.533,k1:=( kvx12 + kta11c12 ) / brho ; +injm4: marker; +qth_a13: quadrupole,l:= 0.673,k1:=( khx13 + kta10b13 ) / brho ; +bpm_a13: monitor; +dchv_a13: kicker,l:= 0,hkick:=kdhta13 ,vkick:=kdvta13 ; +ikickv_a12: vkicker,l:= 0.428,kick:=vkck12 ; +ikickh_a12: hkicker,l:= 0.428,kick:=hkck12 ; +ikickv_a13: vkicker,l:= 0.839,kick:=vkck13 ; +ikickh_a13: hkicker,l:= 0.839,kick:=hkck13 ; +injm2: marker; +dmcv_b01: vkicker,l:= 0,kick:=kdvtb1 ; +qsc_b01: multipole,knl:={ 0,ksc_b01 }; +ssxc_b01: multipole,ksl:={ 0, 0,kssb1_9 }; +bpm_b01: monitor; +qtv_b01: quadrupole,l:= 0.5,k1:=( kvx1 + ktb1_9 ) / brho ; +scv_b01: sextupole,l:= 0.354,k2:=ksv1 ; +dh_b01: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_b02: hkicker,l:= 0,kick:= 0; +qsc_b02: multipole,knl:={ 0,ksc_b02 }; +ssxc_b02: multipole,ksl:={ 0, 0,kssb2_8 }; +bpm_b02: monitor; +qth_b02: quadrupole,l:= 0.5,k1:=( khx2 + ktb2d8 ) / brho ; +sch_b02: sextupole,l:= 0.354,k2:=ksh2 ; +dh_b02: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmcv_b03: vkicker,l:= 0,kick:=kdvtb3 ; +qsc_b03: multipole,knl:={ 0,ksc_b03 }; +bpm_b03: monitor; +qtv_b03: quadrupole,l:= 0.5,k1:=( kvx3 + ktb357 ) / brho ; +sv_b03: sextupole,l:= 0.317,k2:=ksv3 ; +dh_b03: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_b04: hkicker,l:= 0,kick:=kdhtb4 ; +bpm_b04: monitor; +qth_b04: quadrupole,l:= 0.541,k1:=( khx4 - kta4b6 ) / brho ; +sh_b04: sextupole,l:= 0.33,k2:=ksh4 ; +dh_b04: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_b05: sextupole,l:= 0.317,k2:=ksv5 ; +qtv_b05: quadrupole,l:= 0.5,k1:=( kvx5 + ktb357 ) / brho ; +bpm_b05: monitor; +dmcv_b05: vkicker,l:= 0,kick:=kdvtb5 ; +qsc_b05: multipole,knl:={ 0,ksc_b05 }; +dh_b06: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sh_b06: sextupole,l:= 0.33,k2:=ksh6 ; +qth_b06: quadrupole,l:= 0.541,k1:=( khx6 - kta4b6 ) / brho ; +bpm_b06: monitor; +dmch_b06: hkicker,l:= 0,kick:=kdhtb6 ; +dh_b07: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_b07: sextupole,l:= 0.317,k2:=ksv7 ; +qtv_b07: quadrupole,l:= 0.5,k1:=( kvx7 + ktb357 ) / brho ; +bpm_b07: monitor; +dmcv_b07: vkicker,l:= 0,kick:=kdvtb7 ; +qsc_b07: multipole,knl:={ 0,ksc_b07 }; +dh_b08: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_b08: octupole,l:= 0.33,k3:=ko_b08 ; +qth_b08: quadrupole,l:= 0.5,k1:=( khx8 + ktb2d8 ) / brho ; +bpm_b08: monitor; +dmch_b08: hkicker,l:= 0,kick:=kdhtb8 ; +qsc_b08: multipole,knl:={ 0,ksc_b08 }; +ssxc_b08: multipole,ksl:={ 0, 0,kssb2_8 }; +dh_b09: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_b09: octupole,l:= 0.33,k3:=ko_b09 ; +qtv_b09: quadrupole,l:= 0.5,k1:=( kvx9 + ktb1_9 ) / brho ; +bpm_b09: monitor; +dmcv_b09: vkicker,l:= 0,kick:=kdvtb9 ; +qsc_b09: multipole,knl:={ 0,ksc_b09 }; +ssxc_b09: multipole,ksl:={ 0, 0,kssb1_9 }; +dchv_b10: kicker,l:= 0,hkick:=kdhtb10 ,vkick:=kdvtb10 ; +bpm_b10: monitor; +qth_b10: quadrupole,l:= 0.673,k1:=( khx10 - kta10b13 ) / brho ; +qtv_b11: quadrupole,l:= 0.533,k1:=( kvx11 + ktb11d12 ) / brho ; +qtv_b12: quadrupole,l:= 0.533,k1:=( kvx12 + ktb11d12 ) / brho ; +qth_b13: quadrupole,l:= 0.673,k1:=( khx13 - kta10b13 ) / brho ; +bpm_b13: monitor; +dchv_b13: kicker,l:= 0,hkick:=kdhtb13 ,vkick:=kdvtb13 ; +dampkicker1: marker; +dampkicker2: marker; +qmmkicker: marker; +tunekicker: marker; +dmcv_c01: vkicker,l:= 0,kick:=kdvtc1 ; +qsc_c01: multipole,knl:={ 0,ksc_c01 }; +ssxc_c01: multipole,ksl:={ 0, 0,kssc1_9 }; +bpm_c01: monitor; +qtv_c01: quadrupole,l:= 0.5,k1:=( kvx1 + ktc1_9 ) / brho ; +scv_c01: sextupole,l:= 0.354,k2:=ksv1 ; +dh_c01: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_c02: hkicker,l:= 0,kick:=kdhtc2 ; +ssxc_c02: multipole,ksl:={ 0, 0,kssc2_8 }; +qsc_c02: multipole,knl:={ 0,ksc_c02 }; +bpm_c02: monitor; +qth_c02: quadrupole,l:= 0.5,k1:=( khx2 + kta2c8 ) / brho ; +sch_c02: sextupole,l:= 0.354,k2:=ksh2 ; +dh_c02: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmcv_c03: vkicker,l:= 0,kick:=kdvtc3 ; +qsc_c03: multipole,knl:={ 0,ksc_c03 }; +bpm_c03: monitor; +qtv_c03: quadrupole,l:= 0.5,k1:=( kvx3 + ktc357 ) / brho ; +sv_c03: sextupole,l:= 0.317,k2:=ksv3 ; +dh_c03: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_c04: hkicker,l:= 0,kick:=kdhtc4 ; +bpm_c04: monitor; +qth_c04: quadrupole,l:= 0.541,k1:=( khx4 + ktc4d6 ) / brho ; +sh_c04: sextupole,l:= 0.33,k2:=ksh4 ; +dh_c04: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_c05: sextupole,l:= 0.317,k2:=ksv5 ; +qtv_c05: quadrupole,l:= 0.5,k1:=( kvx5 + ktc357 ) / brho ; +bpm_c05: monitor; +dmcv_c05: vkicker,l:= 0,kick:=kdvtc5 ; +qsc_c05: multipole,knl:={ 0,ksc_c05 }; +dh_c06: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sh_c06: sextupole,l:= 0.33,k2:=ksh6 ; +qth_c06: quadrupole,l:= 0.541,k1:=( khx6 + ktc4d6 ) / brho ; +bpm_c06: monitor; +dmch_c06: hkicker,l:= 0,kick:=kdhtc6 ; +dh_c07: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_c07: sextupole,l:= 0.317,k2:=ksv7 ; +qtv_c07: quadrupole,l:= 0.5,k1:=( kvx7 + ktc357 ) / brho ; +bpm_c07: monitor; +dmcv_c07: vkicker,l:= 0,kick:=kdvtc7 ; +qsc_c07: multipole,knl:={ 0,ksc_c07 }; +dh_c08: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_c08: octupole,l:= 0.33,k3:=ko_c08 ; +qth_c08: quadrupole,l:= 0.5,k1:=( khx8 + kta2c8 ) / brho ; +bpm_c08: monitor; +dmch_c08: hkicker,l:= 0,kick:=kdhtc8 ; +qsc_c08: multipole,knl:={ 0,ksc_c08 }; +ssxc_c08: multipole,ksl:={ 0,kssc2_8 }; +dh_c09: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_c09: octupole,l:= 0.33,k3:=ko_c09 ; +qtv_c09: quadrupole,l:= 0.5,k1:=( kvx9 + ktc1_9 ) / brho ; +bpm_c09: monitor; +dmcv_c09: vkicker,l:= 0,kick:=kdvtc9 ; +qsc_c09: multipole,knl:={ 0,ksc_c09 }; +ssxc_c09: multipole,ksl:={ 0, 0,kssc1_9 }; +ekick01: vkicker,l:= 0.4; +ekick02: vkicker,l:= 0.4; +ekick03: vkicker,l:= 0.4; +ekick04: vkicker,l:= 0.505; +ekick05: vkicker,l:= 0.505; +ekick06: vkicker,l:= 0.505; +ekick07: vkicker,l:= 0.505; +dchv_c10: kicker,l:= 0,hkick:=kdhtc10 ,vkick:=kdvtc10 ; +bpm_c10: monitor; +qth_c10: quadrupole,l:= 0.673,k1:=( khx10 + ktc10d13 ) / brho ; +qtv_c11: quadrupole,l:= 0.533,k1:=( kvx11 - kta11c12 ) / brho ; +ekick08: vkicker,l:= 0.4275; +ekick09: vkicker,l:= 0.4275; +ekick10: vkicker,l:= 0.4275; +ekick11: vkicker,l:= 0.4275; +ekick12: vkicker,l:= 0.39; +ekick13: vkicker,l:= 0.39; +ekick14: vkicker,l:= 0.39; +qtv_c12: quadrupole,l:= 0.533,k1:=( kvx12 - kta11c12 ) / brho ; +qth_c13: quadrupole,l:= 0.673,k1:=( khx13 + ktc10d13 ) / brho ; +bpm_c13: monitor; +dchv_c13: kicker,l:= 0,hkick:=kdhtc13 ,vkick:=kdvtc13 ; +scbdsol_c13a: solenoid,l:= 1.22174,ks:=kssol ,ksi:=ksisol ; +scbdsol_c13b: solenoid,l:= 1.22174,ks:=kssol ,ksi:=ksisol ; +dmcv_d01: vkicker,l:= 0,kick:=kdvtd1 ; +qsc_d01: multipole,knl:={ 0,ksc_d01 }; +ssxc_d01: multipole,ksl:={ 0,kssd1_9 }; +bpm_d01: monitor; +qtv_d01: quadrupole,l:= 0.5,k1:=( kvx1 + ktd1_9 ) / brho ; +scv_d01: sextupole,l:= 0.354,k2:=ksv1 ; +dh_d01: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_d02: hkicker,l:= 0,kick:=kdhtd2 ; +qsc_d02: multipole,knl:={ 0,ksc_d02 }; +ssxc_d02: multipole,ksl:={ 0, 0,kssd2_8 }; +bpm_d02: monitor; +qth_d02: quadrupole,l:= 0.5,k1:=( khx2 + ktb2d8 ) / brho ; +sch_d02: sextupole,l:= 0.354,k2:=ksh2 ; +dh_d02: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmcv_d03: vkicker,l:= 0,kick:=kdvtd3 ; +qsc_d03: multipole,knl:={ 0,ksc_d03 }; +bpm_d03: monitor; +qtv_d03: quadrupole,l:= 0.5,k1:=( kvx3 + ktd357 ) / brho ; +sv_d03: sextupole,l:=lsv3 ,k2:=ksv3 ; +dh_d03: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_d04: hkicker,l:= 0,kick:=kdhtd4 ; +bpm_d04: monitor; +qth_d04: quadrupole,l:= 0.541,k1:=( khx4 - ktc4d6 ) / brho ; +sh_d04: sextupole,l:= 0.33,k2:=ksh4 ; +dh_d04: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_d05: sextupole,l:= 0.317,k2:=ksv5 ; +qtv_d05: quadrupole,l:= 0.5,k1:=( kvx5 + ktd357 ) / brho ; +bpm_d05: monitor; +dmcv_d05: vkicker,l:= 0,kick:=kdvtd5 ; +qsc_d05: multipole,knl:={ 0,ksc_d05 }; +dh_d06: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sh_d06: sextupole,l:= 0.33,k2:=ksh6 ; +qth_d06: quadrupole,l:= 0.541,k1:=( khx6 - ktc4d6 ) / brho ; +bpm_d06: monitor; +dmch_d06: hkicker,l:= 0,kick:=kdhtd6 ; +dh_d07: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_d07: sextupole,l:= 0.317,k2:=ksv7 ; +qtv_d07: quadrupole,l:= 0.5,k1:=( kvx7 + ktd357 ) / brho ; +bpm_d07: monitor; +dmcv_d07: vkicker,l:= 0,kick:=kdvtd7 ; +qsc_d07: multipole,knl:={ 0,ksc_d07 }; +dh_d08: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_d08: octupole,l:= 0.33,k3:=ko_d08 ; +qth_d08: quadrupole,l:= 0.5,k1:=( khx8 + ktb2d8 ) / brho ; +bpm_d08: monitor; +dmch_d08: hkicker,l:= 0,kick:=kdhtd8 ; +qsc_d08: multipole,knl:={ 0,ksc_d08 }; +ssxc_d08: multipole,ksl:={ 0, 0,kssd2_8 }; +dh_d09: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_d09: octupole,l:= 0.33,k3:=ko_d09 ; +qtv_d09: quadrupole,l:= 0.5,k1:=( kvx9 + ktd1_9 ) / brho ; +bpm_d09: monitor; +dmcv_d09: vkicker,l:= 0,kick:=kdvtd9 ; +qsc_d09: multipole,knl:={ 0,ksc_d09 }; +ssxc_d09: multipole,ksl:={ 0, 0,kssd1_9 }; +tunepickup: marker; +qmmpickup: marker; +wcm: marker; +bcm: marker; +dchv_d10: kicker,l:= 0,hkick:=kdhtd10 ,vkick:=kdvtd10 ; +bpm_d10: monitor; +qth_d10: quadrupole,l:= 0.673,k1:=( khx10 - ktc10d13 ) / brho ; +qtv_d11: quadrupole,l:= 0.533,k1:=( kvx11 - ktb11d12 ) / brho ; +cav_01: rfcavity,l:= 2.1466,volt:= 0.0133,harmon:= 1; +cav_02: rfcavity,l:= 2.1466,volt:= 0.0133,harmon:= 1; +cav_03: rfcavity,l:= 2.1466,volt:= 0.0133,harmon:= 1; +cav_04: rfcavity,l:= 2.1466,volt:= -0.02,harmon:= 2; +qtv_d12: quadrupole,l:= 0.533,k1:=( kvx12 - ktb11d12 ) / brho ; +qth_d13: quadrupole,l:= 0.673,k1:=( khx13 - ktc10d13 ) / brho ; +bpm_d13: monitor; +dchv_d13: kicker,l:= 0,hkick:=kdhtd13 ,vkick:=kdvtd13 ; +haloscanner1: marker; +wirescanner1: marker; +ipm1: marker; +ipm2: marker; +wirescanner2: marker; +haloscanner2: marker; +dmcv_a01: vkicker,l:= 0,kick:=kdvta1 ; +qsc_a01: multipole,knl:={ 0,ksc_a01 }; +ssxc_a01: multipole,ksl:={ 0, 0,kssa1_9 }; +bpm_a01: monitor; +qtv_a01: quadrupole,l:= 0.5,k1:=( kvx1 + kta1_9 ) / brho ; +scv_a01: sextupole,l:= 0.354,k2:=ksv1 ; +dh_a01: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_a02: hkicker,l:= 0,kick:=kdhta2 ; +qsc_a02: multipole,knl:={ 0,ksc_a02 }; +ssxc_a02: multipole,ksl:={ 0, 0,kssa2_8 }; +bpm_a02: monitor; +qth_a02: quadrupole,l:= 0.5,k1:=( khx2 + kta2c8 ) / brho ; +sch_a02: sextupole,l:= 0.354,k2:=ksh2 ; +dh_a02: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmcv_a03: vkicker,l:= 0,kick:=kdvta3 ; +qsc_a03: multipole,knl:={ 0,ksc_a03 }; +bpm_a03: monitor; +qtv_a03: quadrupole,l:= 0.5,k1:=( kvx3 + kta357 ) / brho ; +sv_a03: sextupole,l:= 0.317,k2:=ksv3 ; +dh_a03: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +dmch_a04: hkicker,l:= 0,kick:=kdhta4 ; +bpm_a04: monitor; +qth_a04: quadrupole,l:= 0.541,k1:=( khx4 + kta4b6 ) / brho ; +sh_a04: sextupole,l:= 0.33,k2:=ksh4 ; +dh_a04: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_a05: sextupole,l:= 0.317,k2:=ksv5 ; +qtv_a05: quadrupole,l:= 0.5,k1:=( kvx5 + kta357 ) / brho ; +bpm_a05: monitor; +dmcv_a05: vkicker,l:= 0,kick:=kdvta5 ; +qsc_a05: multipole,knl:={ 0}; +dh_a06: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sh_a06: sextupole,l:= 0.33,k2:=ksh6 ; +qth_a06: quadrupole,l:= 0.541,k1:=( khx6 + kta4b6 ) / brho ; +bpm_a06: monitor; +dmch_a06: hkicker,l:= 0,kick:=kdhta6 ; +dh_a07: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +sv_a07: sextupole,l:= 0.317,k2:=ksv7 ; +qtv_a07: quadrupole,l:= 0.5,k1:=( kvx7 + kta357 ) / brho ; +bpm_a07: monitor; +dmcv_a07: vkicker,l:= 0,kick:=kdvta7 ; +qsc_a07: multipole,knl:={ 0,ksc_a07 }; +dh_a08: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_a08: octupole,l:= 0.33,k3:=ko_a08 ; +qth_a08: quadrupole,l:= 0.5,k1:=( khx8 + kta2c8 ) / brho ; +bpm_a08: monitor; +dmch_a08: hkicker,l:= 0,kick:=kdhta8 ; +qsc_a08: multipole,knl:={ 0,ksc_a08 }; +ssxc_a08: multipole,ksl:={ 0, 0,kssa2_8 }; +dh_a09: sbend,l:= 1.4407,angle:= 0.1963495408,e1:= 0,e2:= 0; +oct_a09: octupole,l:= 0.33,k3:=ko_a09 ; +qtv_a09: quadrupole,l:= 0.5,k1:=( kvx9 + kta1_9 ) / brho ; +bpm_a09: monitor; +dmcv_a09: vkicker,l:= 0,kick:=kdvta9 ; +qsc_a09: multipole,knl:={ 0,ksc_a09 }; +ssxc_a09: multipole,ksl:={ 0, 0,kssa1_9 }; +ikickh_a10: hkicker,l:= 0.839,kick:=hkck10 ; +ikickv_a10: vkicker,l:= 0.839,kick:=vkck10 ; +ikickh_a11: hkicker,l:= 0.428,kick:=hkck11 ; +ikickv_a11: vkicker,l:= 0.428,kick:=vkck11 ; +dchv_a10: kicker,l:= 0,hkick:=kdhta10 ,vkick:=kdvta10 ; +bpm_a10: monitor; +qth_a10: quadrupole,l:= 0.673,k1:=( khx10 + kta10b13 ) / brho ; +injm3: marker; +qtv_a11: quadrupole,l:= 0.533,k1:=( kvx11 + kta11c12 ) / brho ; +dh_a10: sbend,l:= 0.8632537742,angle:= -0.042,e1:= 0,e2:= -0.042; +dh_a11: sbend,l:= 0.8722394086,angle:= 0.045,e1:= 0.042,e2:= 0.003; +rnginjsol: sequence, l = 248.0098418; +injm1, at = 0; +dh_a12, at = 1.378195456; +dh_a13, at = 3.408731523; +qtv_a12, at = 7.004892621; +injm4, at = 7.693392621; +qth_a13, at = 8.029892621; +bpm_a13, at = 8.547265121; +dchv_a13, at = 8.683688621; +ikickv_a12, at = 10.59989262; +ikickh_a12, at = 11.13989262; +ikickv_a13, at = 12.66989262; +ikickh_a13, at = 13.82989262; +injm2, at = 14.24939262; +dmcv_b01, at = 14.92668062; +qsc_b01, at = 14.92668062; +ssxc_b01, at = 14.92668062; +bpm_b01, at = 15.12176002; +qtv_b01, at = 15.47989262; +scv_b01, at = 16.03310462; +dh_b01, at = 17.47998825; +dmch_b02, at = 18.92687188; +qsc_b02, at = 18.92687188; +ssxc_b02, at = 18.92687188; +bpm_b02, at = 19.11879438; +bpm_b02, at = 19.23008388; +qth_b02, at = 19.48008388; +sch_b02, at = 20.03329588; +dh_b02, at = 21.4801795; +dmcv_b03, at = 22.92706313; +qsc_b03, at = 22.92706313; +bpm_b03, at = 23.11546513; +qtv_b03, at = 23.48027513; +sv_b03, at = 24.01443713; +dh_b03, at = 25.48037076; +dmch_b04, at = 26.89474238; +bpm_b04, at = 27.11860888; +qth_b04, at = 27.48046638; +sh_b04, at = 28.08524038; +dh_b04, at = 29.48056201; +sv_b05, at = 30.94649564; +qtv_b05, at = 31.48065764; +bpm_b05, at = 31.84546764; +dmcv_b05, at = 32.03386964; +qsc_b05, at = 32.03386964; +dh_b06, at = 33.48075326; +sh_b06, at = 34.87607489; +qth_b06, at = 35.48084889; +bpm_b06, at = 35.84270639; +dmch_b06, at = 36.06657289; +dh_b07, at = 37.48094452; +sv_b07, at = 38.94687815; +qtv_b07, at = 39.48104015; +bpm_b07, at = 39.84585015; +dmcv_b07, at = 40.03425215; +qsc_b07, at = 40.03425215; +dh_b08, at = 41.48113577; +oct_b08, at = 42.9280194; +qth_b08, at = 43.4812314; +bpm_b08, at = 43.8425209; +dmch_b08, at = 44.0344434; +qsc_b08, at = 44.0344434; +ssxc_b08, at = 44.0344434; +dh_b09, at = 45.48132703; +oct_b09, at = 46.92821065; +qtv_b09, at = 47.48142265; +bpm_b09, at = 47.83955525; +dmcv_b09, at = 48.03463465; +qsc_b09, at = 48.03463465; +ssxc_b09, at = 48.03463465; +dchv_b10, at = 54.27762665; +bpm_b10, at = 54.41405015; +qth_b10, at = 54.93142265; +qtv_b11, at = 55.95642265; +qtv_b12, at = 69.00642265; +qth_b13, at = 70.03142265; +bpm_b13, at = 70.54879515; +dchv_b13, at = 70.68521865; +dampkicker1, at = 73.19024865; +dampkicker2, at = 73.69024765; +qmmkicker, at = 74.31523965; +tunekicker, at = 75.44023065; +dmcv_c01, at = 76.92821065; +qsc_c01, at = 76.92821065; +ssxc_c01, at = 76.92821065; +bpm_c01, at = 77.12329005; +qtv_c01, at = 77.48142265; +scv_c01, at = 78.03463465; +dh_c01, at = 79.48151828; +dmch_c02, at = 80.92840191; +ssxc_c02, at = 80.92840191; +qsc_c02, at = 80.92840191; +bpm_c02, at = 81.12032441; +qth_c02, at = 81.48161391; +sch_c02, at = 82.03482591; +dh_c02, at = 83.48170954; +dmcv_c03, at = 84.92859316; +qsc_c03, at = 84.92859316; +bpm_c03, at = 85.11699516; +qtv_c03, at = 85.48180516; +sv_c03, at = 86.01596716; +dh_c03, at = 87.48190079; +dmch_c04, at = 88.89627242; +bpm_c04, at = 89.12013892; +qth_c04, at = 89.48199642; +sh_c04, at = 90.08677042; +dh_c04, at = 91.48209204; +sv_c05, at = 92.94802567; +qtv_c05, at = 93.48218767; +bpm_c05, at = 93.84699767; +dmcv_c05, at = 94.03539967; +qsc_c05, at = 94.03539967; +dh_c06, at = 95.4822833; +sh_c06, at = 96.87760493; +qth_c06, at = 97.48237893; +bpm_c06, at = 97.84423643; +dmch_c06, at = 98.06810293; +dh_c07, at = 99.48247455; +sv_c07, at = 100.9484082; +qtv_c07, at = 101.4825702; +bpm_c07, at = 101.8473802; +dmcv_c07, at = 102.0357822; +qsc_c07, at = 102.0357822; +dh_c08, at = 103.4826658; +oct_c08, at = 104.9295494; +qth_c08, at = 105.4827614; +bpm_c08, at = 105.8440509; +dmch_c08, at = 106.0359734; +qsc_c08, at = 106.0359734; +ssxc_c08, at = 106.0359734; +dh_c09, at = 107.4828571; +oct_c09, at = 108.9297407; +qtv_c09, at = 109.4829527; +bpm_c09, at = 109.8410853; +dmcv_c09, at = 110.0361647; +qsc_c09, at = 110.0361647; +ssxc_c09, at = 110.0361647; +ekick01, at = 112.3299527; +ekick02, at = 112.8099527; +ekick03, at = 113.2899527; +ekick04, at = 113.8229527; +ekick05, at = 114.4079527; +ekick06, at = 114.9929527; +ekick07, at = 115.5779527; +dchv_c10, at = 116.2791567; +bpm_c10, at = 116.4155802; +qth_c10, at = 116.9329527; +qtv_c11, at = 117.9579527; +ekick08, at = 118.9059527; +ekick09, at = 119.4149527; +ekick10, at = 119.9249527; +ekick11, at = 120.4309527; +ekick12, at = 120.9202527; +ekick13, at = 121.3895527; +ekick14, at = 121.8606527; +qtv_c12, at = 131.0079527; +qth_c13, at = 132.0329527; +bpm_c13, at = 132.5503252; +dchv_c13, at = 132.6867487; +scbdsol_c13a, at = 134.3897627; +scbdsol_c13b, at = 138.0255227; +dmcv_d01, at = 138.9297407; +qsc_d01, at = 138.9297407; +ssxc_d01, at = 138.9297407; +bpm_d01, at = 139.1248201; +qtv_d01, at = 139.4829527; +scv_d01, at = 140.0361647; +dh_d01, at = 141.4830483; +dmch_d02, at = 142.9299319; +qsc_d02, at = 142.9299319; +ssxc_d02, at = 142.9299319; +bpm_d02, at = 143.1218544; +qth_d02, at = 143.4831439; +sch_d02, at = 144.0363559; +dh_d02, at = 145.4832396; +dmcv_d03, at = 146.9301232; +qsc_d03, at = 146.9301232; +bpm_d03, at = 147.1185252; +qtv_d03, at = 147.4833352; +sv_d03, at = 148.0174972; +dh_d03, at = 149.4834308; +dmch_d04, at = 150.8978024; +bpm_d04, at = 151.1216689; +qth_d04, at = 151.4835264; +sh_d04, at = 152.0883004; +dh_d04, at = 153.4836221; +sv_d05, at = 154.9495557; +qtv_d05, at = 155.4837177; +bpm_d05, at = 155.8485277; +dmcv_d05, at = 156.0369297; +qsc_d05, at = 156.0369297; +dh_d06, at = 157.4838133; +sh_d06, at = 158.879135; +qth_d06, at = 159.483909; +bpm_d06, at = 159.8457665; +dmch_d06, at = 160.069633; +dh_d07, at = 161.4840046; +sv_d07, at = 162.9499382; +qtv_d07, at = 163.4841002; +bpm_d07, at = 163.8489102; +dmcv_d07, at = 164.0373122; +qsc_d07, at = 164.0373122; +dh_d08, at = 165.4841958; +oct_d08, at = 166.9310795; +qth_d08, at = 167.4842915; +bpm_d08, at = 167.845581; +dmch_d08, at = 168.0375035; +qsc_d08, at = 168.0375035; +ssxc_d08, at = 168.0375035; +dh_d09, at = 169.4843871; +oct_d09, at = 170.9312707; +qtv_d09, at = 171.4844827; +bpm_d09, at = 171.8426153; +dmcv_d09, at = 172.0376947; +qsc_d09, at = 172.0376947; +ssxc_d09, at = 172.0376947; +tunepickup, at = 173.5472277; +qmmpickup, at = 174.6722197; +wcm, at = 175.7778927; +bcm, at = 176.7642687; +dchv_d10, at = 178.2806867; +bpm_d10, at = 178.4171102; +qth_d10, at = 178.9344827; +qtv_d11, at = 179.9594827; +cav_01, at = 183.0386827; +cav_02, at = 185.3358827; +cav_03, at = 187.6330827; +cav_04, at = 189.9302827; +qtv_d12, at = 193.0094827; +qth_d13, at = 194.0344827; +bpm_d13, at = 194.5518552; +dchv_d13, at = 194.6882787; +haloscanner1, at = 196.4528147; +wirescanner1, at = 196.6309197; +ipm1, at = 197.0093417; +ipm2, at = 199.3134267; +wirescanner2, at = 199.6918487; +haloscanner2, at = 199.8699527; +dmcv_a01, at = 200.9312707; +qsc_a01, at = 200.9312707; +ssxc_a01, at = 200.9312707; +bpm_a01, at = 201.1263501; +qtv_a01, at = 201.4844827; +scv_a01, at = 202.0376947; +dh_a01, at = 203.4845783; +dmch_a02, at = 204.931462; +qsc_a02, at = 204.931462; +ssxc_a02, at = 204.931462; +bpm_a02, at = 205.1233845; +qth_a02, at = 205.484674; +sch_a02, at = 206.037886; +dh_a02, at = 207.4847696; +dmcv_a03, at = 208.9316532; +qsc_a03, at = 208.9316532; +bpm_a03, at = 209.1200552; +qtv_a03, at = 209.4848652; +sv_a03, at = 210.0190272; +dh_a03, at = 211.4849609; +dmch_a04, at = 212.8993325; +bpm_a04, at = 213.123199; +qth_a04, at = 213.4850565; +sh_a04, at = 214.0898305; +dh_a04, at = 215.4851521; +sv_a05, at = 216.9510857; +qtv_a05, at = 217.4852477; +bpm_a05, at = 217.8500577; +dmcv_a05, at = 218.0384597; +qsc_a05, at = 218.0384597; +dh_a06, at = 219.4853434; +sh_a06, at = 220.880665; +qth_a06, at = 221.485439; +bpm_a06, at = 221.8472965; +dmch_a06, at = 222.071163; +dh_a07, at = 223.4855346; +sv_a07, at = 224.9514682; +qtv_a07, at = 225.4856302; +bpm_a07, at = 225.8504402; +dmcv_a07, at = 226.0388422; +qsc_a07, at = 226.0388422; +dh_a08, at = 227.4857259; +oct_a08, at = 228.9326095; +qth_a08, at = 229.4858215; +bpm_a08, at = 229.847111; +dmch_a08, at = 230.0390335; +qsc_a08, at = 230.0390335; +ssxc_a08, at = 230.0390335; +dh_a09, at = 231.4859171; +oct_a09, at = 232.9328008; +qtv_a09, at = 233.4860128; +bpm_a09, at = 233.8441454; +dmcv_a09, at = 234.0392248; +qsc_a09, at = 234.0392248; +ssxc_a09, at = 234.0392248; +ikickh_a10, at = 235.1360128; +ikickv_a10, at = 236.2960128; +ikickh_a11, at = 237.8260128; +ikickv_a11, at = 238.3660128; +dchv_a10, at = 240.2822168; +bpm_a10, at = 240.4186403; +qth_a10, at = 240.9360128; +injm3, at = 241.2725128; +qtv_a11, at = 241.9610128; +dh_a10, at = 245.1911396; +dh_a11, at = 247.5737221; +endsequence; diff --git a/examples/TEAPOT/test_sns_ring.py b/examples/TEAPOT/test_sns_ring.py new file mode 100644 index 00000000..bd509a63 --- /dev/null +++ b/examples/TEAPOT/test_sns_ring.py @@ -0,0 +1,11 @@ +from orbit.core.bunch import Bunch +from orbit.lattice import AccLattice +from orbit.lattice import AccNode +from orbit.teapot import TEAPOT_Lattice + + +lattice = TEAPOT_Lattice() +lattice.readMADX("inputs/sns_ring_dual_solenoid.lat", "rnginjsol") + +for node in lattice.getNodes(): + print(node) From 5dac9b7dc4a8e1d4e1f40da91a171f4762664644 Mon Sep 17 00:00:00 2001 From: austin-hoover Date: Mon, 20 Oct 2025 13:39:53 -0400 Subject: [PATCH 3/3] Add type hints to teapot.py --- py/orbit/teapot/teapot.py | 195 +++++++++++++++++++++++--------------- 1 file changed, 116 insertions(+), 79 deletions(-) diff --git a/py/orbit/teapot/teapot.py b/py/orbit/teapot/teapot.py index c3c868a7..9ad5a70a 100644 --- a/py/orbit/teapot/teapot.py +++ b/py/orbit/teapot/teapot.py @@ -13,9 +13,14 @@ - Monitor """ +from __future__ import annotations + import sys import os import math +from typing import Any +from typing import Callable +from typing import Union from ..lattice import AccLattice from ..lattice import AccNode @@ -27,7 +32,9 @@ from ..parsers.mad_parser import MAD_LattElement from ..parsers.madx_parser import MADX_Parser from ..parsers.madx_parser import MADX_LattElement + from orbit.core.aperture import Aperture +from orbit.core.bunch import Bunch from orbit.core.bunch import BunchTwissAnalysis @@ -37,11 +44,11 @@ class TEAPOT_Lattice(AccLattice): TEAPOT has the ability to read MAD files. """ - def __init__(self, name="no name"): + def __init__(self, name: str = "no name") -> None: AccLattice.__init__(self, name) self.useCharge = 1 - def readMAD(self, mad_file_name, lineName): + def readMAD(self, mad_file_name: str, lineName: str) -> None: """ It creates the teapot lattice from MAD file. """ @@ -67,7 +74,7 @@ def readMAD(self, mad_file_name, lineName): self._addChildren() self.initialize() - def readMADX(self, madx_file_name, seqName): + def readMADX(self, madx_file_name: str, seqName: str) -> None: """ It creates the teapot lattice from MAD file. """ @@ -93,13 +100,13 @@ def readMADX(self, madx_file_name, seqName): self._addChildren() self.initialize() - def _addChildren(self): + def _addChildren(self) -> None: """ It does nothing here. It is a place holder for TEAPOT_Ring class. """ pass - def initialize(self): + def initialize(self) -> None: AccLattice.initialize(self) # set up ring length for RF nodes ringRF_Node = RingRFTEAPOT() @@ -107,11 +114,7 @@ def initialize(self): if node.getType() == ringRF_Node.getType(): node.getParamsDict()["ring_length"] = self.getLength() - def getSubLattice( - self, - index_start=-1, - index_stop=-1, - ): + def getSubLattice(self, index_start: int = -1, index_stop: int = -1): # -> Self (python 3.11+) """ It returns the new TEAPOT_Lattice with children with indexes between index_start and index_stop inclusive @@ -120,7 +123,14 @@ def getSubLattice( new_teapot_lattice.setUseRealCharge(self.getUseRealCharge()) return new_teapot_lattice - def trackBunch(self, bunch, paramsDict={}, actionContainer=None, index_start=-1, index_stop=-1): + def trackBunch( + self, + bunch: Bunch, + paramsDict: dict = {}, + actionContainer: AccActionsContainer = None, + index_start: str = -1, + index_stop: str = -1, + ) -> None: """ It tracks the bunch through the lattice. Indexes index_start and index_stop are inclusive. """ @@ -137,11 +147,11 @@ def track(paramsDict): self.trackActions(actionContainer, paramsDict, index_start, index_stop) actionContainer.removeAction(track, AccActionsContainer.BODY) - def setUseRealCharge(self, useCharge=1): + def setUseRealCharge(self, useCharge: int = 1) -> None: """If useCharge != 1 the trackBunch(...) method will assume the charge = +1""" self.useCharge = useCharge - def getUseRealCharge(self): + def getUseRealCharge(self) -> int: """If useCharge != 1 the trackBunch(...) method will assume the charge = +1""" return self.useCharge @@ -152,10 +162,10 @@ class TEAPOT_Ring(TEAPOT_Lattice): collection for rings. TEAPOT has the ability to read MAD files. """ - def __init__(self, name="no name"): + def __init__(self, name: str = "no name") -> None: TEAPOT_Lattice.__init__(self, name) - def _addChildren(self): + def _addChildren(self) -> None: """ Adds Bunch wrapping nodes to all lattice nodes of 1st level. These wrapping nodes will move particles from head ( tail ) to tail ( head) @@ -178,10 +188,10 @@ class _teapotFactory: Class. Factory to produce TEAPOT accelerator elements. """ - def __init__(self): + def __init__(self) -> None: pass - def getElements(self, madElem): + def getElements(self, madElem: Union[MAD_LattElement, MADX_LattElement]) -> list[NodeTEAPOT]: """ Method produces TEAPOT accelerator elements. It returns the array of TEAPOT nodes. Usually there is only one such element, but sometimes (RF with a non zero length) @@ -340,8 +350,8 @@ def getElements(self, madElem): return [drft_1, elem, drft_2] # ==========Others elements such as markers,monitor,rcollimator if ( - madElem.getType().lower() == "marker" - or # madElem.getType().lower() == "monitor" or \ + madElem.getType().lower() == "marker" # madElem.getType().lower() == "monitor" or \ + or # madElem.getType().lower() == "hmonitor" or \ # madElem.getType().lower() == "vmonitor" or \ madElem.getType().lower() == "rcolimator" @@ -404,7 +414,7 @@ def getElements(self, madElem): class BaseTEAPOT(AccNodeBunchTracker): """The base abstract class of the TEAPOT accelerator elements hierarchy.""" - def __init__(self, name="no name"): + def __init__(self, name: str = "no name") -> None: """ Constructor. Creates the base TEAPOT element. This is a superclass for all TEAPOT elements. """ @@ -413,14 +423,14 @@ def __init__(self, name="no name"): class TurnCounterTEAPOT(BaseTEAPOT): - def __init__(self, name="TurnCounter"): + def __init__(self, name: str = "TurnCounter") -> None: """ Constructor. Creates the TEAPOT for turn count in the Ring lattice. """ BaseTEAPOT.__init__(self, name) self.setType("turn counter") - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The Turn Counter class implementation of the AccNodeBunchTracker class track(probe) method. """ @@ -431,7 +441,7 @@ def track(self, paramsDict): class NodeTEAPOT(BaseTEAPOT): - def __init__(self, name="no name"): + def __init__(self, name: str = "no name") -> None: """ Constructor. Creates the real TEAPOT element. This is a superclass for all real TEAPOT elements. The term real means that element is included in TEAPOT list of elements. @@ -453,7 +463,7 @@ def __init__(self, name="no name"): self.addParam("tilt", self.__tiltNodeIN.getTiltAngle()) self.setType("node teapot") - def setTiltAngle(self, angle=0.0): + def setTiltAngle(self, angle: float = 0.0) -> None: """ Sets the tilt angle for the tilt operation. """ @@ -461,86 +471,86 @@ def setTiltAngle(self, angle=0.0): self.__tiltNodeIN.setTiltAngle(angle) self.__tiltNodeOUT.setTiltAngle((-1.0) * angle) - def getTiltAngle(self): + def getTiltAngle(self) -> float: """ Returns the tilt angle for the tilt operation. """ return self.__tiltNodeIN.getTiltAngle() - def getNodeFringeFieldIN(self): + def getNodeFringeFieldIN(self) -> FringeFieldTEAPOT: """ Returns the FringeFieldTEAPOT instance before this TEAPOT node """ return self.__fringeFieldIN - def getNodeFringeFieldOUT(self): + def getNodeFringeFieldOUT(self) -> FringeFieldTEAPOT: """ Returns the FringeFieldTEAPOT instance after this TEAPOT node """ return self.__fringeFieldOUT - def getNodeTiltIN(self): + def getNodeTiltIN(self) -> TiltTEAPOT: """ Returns the TiltTEAPOT instance before this TEAPOT node """ return self.__tiltNodeIN - def getNodeTiltOUT(self): + def getNodeTiltOUT(self) -> TiltTEAPOT: """ Returns the TiltTEAPOT instance after this TEAPOT node """ return self.__tiltNodeOUT - def setFringeFieldFunctionIN(self, trackFunction=None): + def setFringeFieldFunctionIN(self, trackFunction: Callable = None) -> None: """ Sets the fringe field function that will track the bunch through the fringe at the entrance of the node. """ self.__fringeFieldIN.setFringeFieldFunction(trackFunction) - def setFringeFieldFunctionOUT(self, trackFunction=None): + def setFringeFieldFunctionOUT(self, trackFunction: Callable = None) -> None: """ Sets the fringe field function that will track the bunch through the fringe at the exit of the element. """ self.__fringeFieldOUT.setFringeFieldFunction(trackFunction) - def getFringeFieldFunctionIN(self, trackFunction=None): + def getFringeFieldFunctionIN(self) -> Callable: """ Returns the fringe field function that will track the bunch through the fringe at the entrance of the node. """ return self.__fringeFieldIN.getFringeFieldFunction() - def getFringeFieldFunctionOUT(self, trackFunction=None): + def getFringeFieldFunctionOUT(self) -> Callable: """ Returns the fringe field function that will track the bunch through the fringe at the exit of the element. """ return self.__fringeFieldOUT.getFringeFieldFunction() - def setUsageFringeFieldIN(self, usage=True): + def setUsageFringeFieldIN(self, usage: bool = True) -> None: """ Sets the property describing if the IN fringe field will be used in calculation. """ self.__fringeFieldIN.setUsage(usage) - def getUsageFringeFieldIN(self): + def getUsageFringeFieldIN(self) -> bool: """ Returns the property describing if the IN fringe field will be used in calculation. """ return self.__fringeFieldIN.getUsage() - def setUsageFringeFieldOUT(self, usage=True): + def setUsageFringeFieldOUT(self, usage: bool = True) -> None: """ Sets the property describing if the OUT fringe field will be used in calculation. """ self.__fringeFieldOUT.setUsage(usage) - def getUsageFringeFieldOUT(self): + def getUsageFringeFieldOUT(self) -> bool: """ Returns the property describing if the OUT fringe field will be used in calculation. @@ -553,14 +563,14 @@ class DriftTEAPOT(NodeTEAPOT): Drift TEAPOT element. """ - def __init__(self, name="drift no name"): + def __init__(self, name: str = "drift no name") -> None: """ Constructor. Creates the Drift TEAPOT element. """ NodeTEAPOT.__init__(self, name) self.setType("drift teapot") - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The drift class implementation of the AccNodeBunchTracker class track(probe) method. """ @@ -574,7 +584,7 @@ class ApertureTEAPOT(NodeTEAPOT): Aperture TEAPOT element. """ - def __init__(self, name="aperture no name"): + def __init__(self, name: str = "aperture no name") -> None: """ Constructor. Creates the aperutre element. """ @@ -583,7 +593,7 @@ def __init__(self, name="aperture no name"): self.addParam("aperture", []) self.addParam("apertype", 0.0) - def initialize(self): + def initialize(self) -> None: shape = self.getParam("apertype") dim = self.getParam("aperture") if len(dim) > 0: @@ -594,7 +604,7 @@ def initialize(self): if shape == 3: self.aperture = Aperture(shape, dim[0], dim[1], 0.0, 0.0) - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The aperture class implementation of the ApertueNode. """ @@ -608,7 +618,7 @@ class MonitorTEAPOT(NodeTEAPOT): Monitor TEAPOT element. """ - def __init__(self, name="Monitor no name"): + def __init__(self, name: str = "Monitor no name") -> None: """ Constructor. Creates the aperutre element. """ @@ -616,7 +626,7 @@ def __init__(self, name="Monitor no name"): self.setType("monitor teapot") self.twiss = BunchTwissAnalysis() - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The bunchtuneanalysis-teapot class implementation of the AccNodeBunchTracker class track(probe) method. """ @@ -634,7 +644,7 @@ class BunchWrapTEAPOT(NodeTEAPOT): Drift TEAPOT element. """ - def __init__(self, name="drift no name"): + def __init__(self, name: str = "drift no name") -> None: """ Constructor. Creates the Bunch wrapper TEAPOT element used in Ring lattices. """ @@ -642,7 +652,7 @@ def __init__(self, name="drift no name"): self.setType("bunch_wrap_teapot") self.addParam("ring_length", 0.0) - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The bunch wrap class implementation of the AccNodeBunchTracker class track(probe) method. """ @@ -657,7 +667,7 @@ class SolenoidTEAPOT(NodeTEAPOT): Solenoid TEAPOT element. """ - def __init__(self, name="solenoid no name"): + def __init__(self, name: str = "solenoid no name") -> None: """ Constructor. Creates the Solenoid TEAPOT element. """ @@ -666,7 +676,7 @@ def __init__(self, name="solenoid no name"): self.addParam("B", 0.0) self.waveform = None - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The Solenoid TEAPOT class implementation of the AccNodeBunchTracker class track(probe) method. @@ -683,7 +693,7 @@ def track(self, paramsDict): useCharge = paramsDict["useCharge"] TPB.soln(bunch, length, B, useCharge) - def setWaveform(self, waveform): + def setWaveform(self, waveform: Any) -> None: """ Sets the time dependent waveform function """ @@ -695,7 +705,7 @@ class MultipoleTEAPOT(NodeTEAPOT): Multipole Combined Function TEAPOT element. """ - def __init__(self, name="multipole no name"): + def __init__(self, name: str = "multipole no name") -> None: """ Constructor. Creates the Multipole Combined Function TEAPOT element. @@ -758,7 +768,7 @@ def fringeOUT(node, paramsDict): self.setType("multipole teapot") - def initialize(self): + def initialize(self) -> None: """ The Multipole Combined Function TEAPOT class implementation of the AccNode class initialize() method. @@ -783,7 +793,7 @@ def initialize(self): for i in range(nParts - 2): self.setLength(lengthStep, i + 1) - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The Multipole Combined Function TEAPOT class implementation of the AccNodeBunchTracker class @@ -822,7 +832,7 @@ def track(self, paramsDict): TPB.drift(bunch, length) return - def setWaveform(self, waveform): + def setWaveform(self, waveform: Any) -> None: """ Sets the time dependent waveform function """ @@ -834,7 +844,7 @@ class QuadTEAPOT(NodeTEAPOT): Quad Combined Function TEAPOT element. """ - def __init__(self, name="quad no name"): + def __init__(self, name: str = "quad no name") -> None: """ Constructor. Creates the Quad Combined Function TEAPOT element. @@ -907,7 +917,7 @@ def fringeOUT(node, paramsDict): self.setType("quad teapot") - def initialize(self): + def initialize(self) -> None: """ The Quad Combined Function TEAPOT class implementation of the AccNode class initialize() method. @@ -932,7 +942,7 @@ def initialize(self): for i in range(nParts - 2): self.setLength(lengthStep, i + 1) - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The Quad Combined Function TEAPOT class implementation of the AccNodeBunchTracker class track(probe) method. @@ -975,7 +985,7 @@ def track(self, paramsDict): TPB.quad1(bunch, length, kq, useCharge) return - def setWaveform(self, waveform): + def setWaveform(self, waveform: Any) -> None: """ Sets the time dependent waveform function """ @@ -987,7 +997,7 @@ class BendTEAPOT(NodeTEAPOT): Bend Combined Functions TEAPOT element. """ - def __init__(self, name="bend no name"): + def __init__(self, name: str = "bend no name") -> None: """ Constructor. Creates the Bend Combined Functions TEAPOT element . """ @@ -1031,7 +1041,18 @@ def fringeIN(node, paramsDict): TPB.multpfringeIN(bunch, pole, kl, skew, useCharge) frinout = 1 TPB.wedgerotate(bunch, e, frinout) - TPB.wedgebendCF(bunch, e, inout, rho, len(poleArr), poleArr, klArr, skewArr, nParts - 1, useCharge) + TPB.wedgebendCF( + bunch, + e, + inout, + rho, + len(poleArr), + poleArr, + klArr, + skewArr, + nParts - 1, + useCharge, + ) else: if usageIN: TPB.bendfringeIN(bunch, rho) @@ -1057,7 +1078,18 @@ def fringeOUT(node, paramsDict): nParts = paramsDict["parentNode"].getnParts() if e != 0.0: inout = 1 - TPB.wedgebendCF(bunch, e, inout, rho, len(poleArr), poleArr, klArr, skewArr, nParts - 1, useCharge) + TPB.wedgebendCF( + bunch, + e, + inout, + rho, + len(poleArr), + poleArr, + klArr, + skewArr, + nParts - 1, + useCharge, + ) if usageOUT: frinout = 0 TPB.wedgerotate(bunch, -e, frinout) @@ -1086,7 +1118,7 @@ def fringeOUT(node, paramsDict): self.setType("bend teapot") - def initialize(self): + def initialize(self) -> None: """ The Bend Combined Functions TEAPOT class implementation of the AccNode class initialize() method. @@ -1113,7 +1145,7 @@ def initialize(self): for i in range(nParts - 2): self.setLength(lengthStep, i + 1) - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The Bend Combined Functions TEAPOT class implementation of the AccNodeBunchTracker class track(probe) method. @@ -1167,7 +1199,7 @@ class RingRFTEAPOT(NodeTEAPOT): Ring RF TEAPOT element. """ - def __init__(self, name="RingRF no name"): + def __init__(self, name: str = "RingRF no name") -> None: """ Constructor. Creates the Ring RF TEAPOT element. Harmonics numbers are 1,2,3, ... @@ -1183,7 +1215,7 @@ def __init__(self, name="RingRF no name"): self.setnParts(1) - def addRF(self, harmonic, voltage, phase): + def addRF(self, harmonic: float, voltage: float, phase: float) -> None: """ Method. It adds the RF component with sertain harmonic, voltage, and phase. @@ -1192,7 +1224,7 @@ def addRF(self, harmonic, voltage, phase): self.getParam("voltages").append(voltage) self.getParam("phases").append(phase) - def initialize(self): + def initialize(self) -> None: """ The Ring RF TEAPOT class implementation of the AccNode class initialize() method. @@ -1212,7 +1244,7 @@ def initialize(self): msg = msg + "length =" + str(self.getLength()) orbitFinalize(msg) - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The Ring RF TEAPOT class implementation of the AccNodeBunchTracker class track(probe) method. @@ -1237,7 +1269,7 @@ class KickTEAPOT(NodeTEAPOT): Kick TEAPOT element. """ - def __init__(self, name="kick no name"): + def __init__(self, name: str = "kick no name") -> None: """ Constructor. Creates the Kick TEAPOT element . """ @@ -1249,7 +1281,7 @@ def __init__(self, name="kick no name"): self.setnParts(2) self.waveform = None - def initialize(self): + def initialize(self) -> None: """ The Kicker TEAPOT class implementation of the AccNode class initialize() method. @@ -1274,7 +1306,7 @@ def initialize(self): for i in range(nParts - 2): self.setLength(lengthStep, i + 1) - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ The Kick TEAPOT class implementation of the AccNodeBunchTracker class track(probe) method. @@ -1316,7 +1348,7 @@ class TiltTEAPOT(BaseTEAPOT): The class to do tilt at the entrance of an TEAPOT element. """ - def __init__(self, name="tilt no name", angle=0.0): + def __init__(self, name: str = "tilt no name", angle: float = 0.0) -> None: """ Constructor. Creates the Tilt TEAPOT element. """ @@ -1324,19 +1356,19 @@ def __init__(self, name="tilt no name", angle=0.0): self.__angle = angle self.setType("tilt teapot") - def setTiltAngle(self, angle=0.0): + def setTiltAngle(self, angle: float = 0.0) -> None: """ Sets the tilt angle for the tilt operation. """ self.__angle = angle - def getTiltAngle(self): + def getTiltAngle(self) -> float: """ Returns the tilt angle for the tilt operation. """ return self.__angle - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ It is tracking the dictionary with parameters through the titlt node. @@ -1351,7 +1383,12 @@ class FringeFieldTEAPOT(BaseTEAPOT): The class is a base class for the fringe field classes for others TEAPOT elements. """ - def __init__(self, parentNode, trackFunction=None, name="fringe field no name"): + def __init__( + self, + parentNode: AccNode, + trackFunction: Callable = None, + name: str = "fringe field no name", + ) -> None: """ Constructor. Creates the Fringe Field TEAPOT element. """ @@ -1361,7 +1398,7 @@ def __init__(self, parentNode, trackFunction=None, name="fringe field no name"): self.__usage = True self.setType("fringeField teapot") - def track(self, paramsDict): + def track(self, paramsDict: dict) -> None: """ It is tracking the dictionary with parameters through the fringe field node. @@ -1369,26 +1406,26 @@ def track(self, paramsDict): if self.__trackFunc != None: self.__trackFunc(self, paramsDict) - def setFringeFieldFunction(self, trackFunction=None): + def setFringeFieldFunction(self, trackFunction: Callable) -> None: """ Sets the fringe field function that will track the bunch through the fringe. """ self.__trackFunc = trackFunction - def getFringeFieldFunction(self): + def getFringeFieldFunction(self) -> Callable: """ Returns the fringe field function. """ return self.__trackFunc - def setUsage(self, usage=True): + def setUsage(self, usage: bool = True) -> None: """ Sets the boolean flag describing if the fringe field will be used in calculation. """ self.__usage = usage - def getUsage(self): + def getUsage(self) -> bool: """ Returns the boolean flag describing if the fringe field will be used in calculation.