diff --git a/examples/Diagnostics/Tunes/inputs/sns_ring.lat b/examples/Diagnostics/Tunes/inputs/sns_ring.lat new file mode 100644 index 00000000..cbfb8df1 --- /dev/null +++ b/examples/Diagnostics/Tunes/inputs/sns_ring.lat @@ -0,0 +1,768 @@ +kvx12 := kdc; +kdc = -3.412580596; +kta11c12 := 0; +brho := 1e9*(pc/c); +pc := sqrt(ek*(ek+2*e0)); +ek := 0.8; +e0 := 0.93827231; +c := 299792458; +khx13 := kfc; +kfc = 3.161218767; +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.579970923; +ktb1_9 := 0; +ksv1 := 0; +ksc_b02 := 0; +kssb2_8 = 0; +khx2 := kf; +kf = 3.388342139; +ktb2d8 := 0; +ksh2 := 0; +kdvtb3 := 0; +ksc_b03 := 0; +kvx3 := kd; +kd = -3.731394588; +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; +ksisol := 1.22174*kssol; +kssol := solfield/brho; +solfield = 0; +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:= 0,ksi:=ksisol ; +scbdsol_c13b: solenoid,l:= 1.22174,ks:= 0,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/Diagnostics/Tunes/test_sns_2d.py b/examples/Diagnostics/Tunes/test_sns_2d.py new file mode 100644 index 00000000..cdd04265 --- /dev/null +++ b/examples/Diagnostics/Tunes/test_sns_2d.py @@ -0,0 +1,155 @@ +"""Test one-turn tune estimation in uncoupled lattice. + +This example tracks a Gaussian distribution through a FODO lattice. The tunes +are estimated from the phase space coordinates before/after tracking using the +`BunchTuneAnalysis` class. +""" +import argparse +import math +import os +import pathlib +import random +from pprint import pprint + +import numpy as np +import pandas as pd + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.bunch_generators import TwissContainer +from orbit.bunch_generators import GaussDist2D +from orbit.diagnostics import TeapotTuneAnalysisNode +from orbit.diagnostics.matrix import TransferMatrixAnalysis +from orbit.teapot import TEAPOT_Lattice +from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.utils.consts import mass_proton + +from utils import make_lattice_sns +from utils import get_tmat + + +# Arguments +parser = argparse.ArgumentParser() +parser.add_argument("--norm-from", type=str, default="tmat", choices=["tmat", "cov", "twiss"]) +args = parser.parse_args() + +# Setup +path = pathlib.Path(__file__) +output_dir = os.path.join("outputs", path.stem) +os.makedirs(output_dir, exist_ok=True) + +# Initialize lattice and bunch +lattice = make_lattice_sns() + +bunch = Bunch() +bunch.mass(mass_proton) +bunch.getSyncParticle().kinEnergy(1.000) + +# Calculate transfer matrix +matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) +lattice_params = matrix_lattice.getRingParametersDict() +pprint(lattice_params) + +# Store some parameters as variables +lattice_alpha_x = lattice_params["alpha x"] +lattice_alpha_y = lattice_params["alpha y"] +lattice_beta_x = lattice_params["beta x [m]"] +lattice_beta_y = lattice_params["beta y [m]"] +lattice_eta_x = lattice_params["dispersion x [m]"] +lattice_etap_x = lattice_params["dispersion prime x"] + +# Add tune diagnostic node +tune_node = TeapotTuneAnalysisNode() + +if args.norm_from == "twiss": + tune_node.setNormMatrixFromTwiss( + betax=lattice_beta_x, + alphax=lattice_alpha_x, + etax=lattice_eta_x, + etapx=lattice_etap_x, + betay=lattice_beta_y, + alphay=lattice_alpha_y, + ) +elif args.norm_from == "tmat": + tmat = get_tmat(lattice, bunch) + tune_node.setNormMatrixFromTransferMatrix(tmat) +elif args.norm_from == "cov": + tmat = get_tmat(lattice, bunch) + tmat_analysis = TransferMatrixAnalysis(tmat) + cov_matrix = tmat_analysis.cov_matrix(1e-7, 1e-7) + tune_node.setNormMatrixFromCovMatrix(cov_matrix) +else: + raise ValueError("Invalid norm_from argument") + +lattice.getNodes()[0].addChildNode(tune_node, 0) + +# Generate particles +emittance_x = 0.1e-06 +emittance_y = 0.1e-06 +twiss_x = TwissContainer(lattice_alpha_x, lattice_beta_x, emittance_x) +twiss_y = TwissContainer(lattice_alpha_y, lattice_beta_y, emittance_y) +dist = GaussDist2D(twiss_x, twiss_y) + +for index in range(1000): + x, xp, y, yp = dist.getCoordinates() + z = random.uniform(-25.0, 25.0) + bunch.addParticle(x, xp, y, yp, z, 0.0) + +# Track particles +for turn in range(10): + lattice.trackBunch(bunch) + + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + xrms = 1000.0 * math.sqrt(twiss_calc.getCorrelation(0, 0)) + yrms = 1000.0 * math.sqrt(twiss_calc.getCorrelation(2, 2)) + print("turn={} xrms={:0.3f} yrms={:0.3f}".format(turn + 1, xrms, yrms)) + +# Test writing to file +filename = "bunch.dat" +filename = os.path.join(output_dir, filename) +bunch.dumpBunch(filename) + +# Collect phase data from bunch +phase_data = tune_node.getData(bunch) +phase_data = pd.DataFrame(phase_data) + +# Read phase data from file +particles = np.loadtxt(filename, comments="%") +particles = pd.DataFrame( + particles, + columns=[ # https://github.com/PyORBIT-Collaboration/PyORBIT3/issues/78 + "x", + "xp", + "y", + "yp", + "z", + "dE", + "phase_x", + "phase_y", + "tune_x", + "tune_y", + "action_x", + "action_y", + ], +) +print(particles.iloc[:, 6:]) + +# Check against tune from transfer matrix +tune_x_true = lattice_params["fractional tune x"] +tune_y_true = lattice_params["fractional tune y"] +tune_x_calc = np.mean(phase_data["tune_1"]) +tune_y_calc = np.mean(phase_data["tune_2"]) + +tune_x_err = tune_x_calc - tune_x_true +tune_y_err = tune_y_calc - tune_y_true + +print("tune_x_true", tune_x_true) +print("tune_x_calc", tune_x_calc) +print("tune_y_true", tune_y_true) +print("tune_y_calc", tune_y_calc) +print("tune_x_err", tune_x_err) +print("tune_y_err", tune_y_err) + +assert np.abs(tune_x_err) < 1.00e-06 +assert np.abs(tune_y_err) < 1.00e-06 diff --git a/examples/Diagnostics/Tunes/test_sns_4d.py b/examples/Diagnostics/Tunes/test_sns_4d.py new file mode 100644 index 00000000..be53b774 --- /dev/null +++ b/examples/Diagnostics/Tunes/test_sns_4d.py @@ -0,0 +1,107 @@ +"""Test one-turn tune estimation in coupled lattice.""" + +import argparse +import math +import os +import pathlib + +import numpy as np +import pandas as pd + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.diagnostics import TeapotTuneAnalysisNode +from orbit.diagnostics.matrix import TransferMatrixAnalysis +from orbit.teapot import TEAPOT_Ring +from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.utils.consts import mass_proton + +from utils import make_lattice_sns +from utils import get_tmat + + +# Arguments +parser = argparse.ArgumentParser() +parser.add_argument("--coupled", type=int, default=1) +parser.add_argument("--norm-from", type=str, default="tmat", choices=["tmat", "cov"]) +parser.add_argument("--eps1", type=float, default=1.0) +parser.add_argument("--eps2", type=float, default=1.01) +args = parser.parse_args() + +# Setup +path = pathlib.Path(__file__) +output_dir = os.path.join("outputs", path.stem) +os.makedirs(output_dir, exist_ok=True) + +# Initialize lattice and bunch +lattice = make_lattice_sns(sol=args.coupled) +lattice.initialize() + +bunch = Bunch() +bunch.mass(mass_proton) +bunch.getSyncParticle().kinEnergy(1.000) + +# Calculate transfer matrix +M = get_tmat(lattice, bunch) +tmat_analysis = TransferMatrixAnalysis(M) +tune_1_true = tmat_analysis.eigtunes[0] +tune_2_true = tmat_analysis.eigtunes[1] + +# Calculate matched covariance matrix +eps_1 = args.eps1 * 1e-6 +eps_2 = args.eps2 * 1e-6 +S_matched = tmat_analysis.cov_matrix(eps_1, eps_2) + +# Add tune diagnostic node +tune_node = TeapotTuneAnalysisNode() +if args.norm_from == "tmat": + tune_node.setNormMatrixFromTransferMatrix(M) +else: + tune_node.setNormMatrixFromCovMatrix(S_matched) +lattice.getNodes()[0].addChildNode(tune_node, 0) + +# Generate particles +rng = np.random.default_rng(123) +particles = np.zeros((1000, 6)) +particles[:, :4] = rng.multivariate_normal( + mean=np.zeros(4), cov=S_matched, size=particles.shape[0] +) +particles[:, 4] = rng.uniform(-25.0, 25.0, size=particles.shape[0]) +particles[:, 5] = 0.0 + +for index in range(particles.shape[0]): + bunch.addParticle(*particles[index]) + +# Track particles +for turn in range(10): + lattice.trackBunch(bunch) + + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + xrms = 1000.0 * math.sqrt(twiss_calc.getCorrelation(0, 0)) + yrms = 1000.0 * math.sqrt(twiss_calc.getCorrelation(2, 2)) + print("turn={} xrms={:0.3f} yrms={:0.3f}".format(turn + 1, xrms, yrms)) + +# Analysis +phase_data = tune_node.getData(bunch) # phase_data = pd.DataFrame(phase_data) + +tune_1_calc = np.mean(phase_data["tune_1"]) +tune_2_calc = np.mean(phase_data["tune_2"]) + +# Order depends on eps1/eps2. To check values against transfer matrix, +# sort tunes by magnitude. +if tune_1_calc < tune_2_calc: + (tune_1_calc, tune_2_calc) = (tune_2_calc, tune_1_calc) + +tune_1_err = tune_1_calc - tune_1_true +tune_2_err = tune_2_calc - tune_2_true + +print("tune_1_true", tune_1_true) +print("tune_1_calc", tune_1_calc) +print("tune_2_true", tune_2_true) +print("tune_2_calc", tune_2_calc) +print("tune_1_err", tune_1_err) +print("tune_2_err", tune_2_err) + +assert np.abs(tune_1_err) < 1.00e-05 +assert np.abs(tune_2_err) < 1.00e-05 diff --git a/examples/Diagnostics/Tunes/test_update_norm_mat.py b/examples/Diagnostics/Tunes/test_update_norm_mat.py new file mode 100644 index 00000000..4ae4c196 --- /dev/null +++ b/examples/Diagnostics/Tunes/test_update_norm_mat.py @@ -0,0 +1,99 @@ +"""Test updating normalization matrix during tracking. + +Average tunes are printed out after each turn. The normalization matrix +is updated twice during tracking. A warning message should print that +indicates the normalization matrix has changed, and that the tunes +will be inaccurate until the next tracked turn. +""" +import os +import pathlib +import random +from pprint import pprint + +import numpy as np + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.bunch_generators import TwissContainer +from orbit.bunch_generators import GaussDist2D +from orbit.diagnostics import TeapotTuneAnalysisNode +from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.utils.consts import mass_proton + +from utils import make_lattice_sns + +# Setup +path = pathlib.Path(__file__) +output_dir = os.path.join("outputs", path.stem) +os.makedirs(output_dir, exist_ok=True) + +# Initialize lattice and bunch +lattice = make_lattice_sns() + +bunch = Bunch() +bunch.mass(mass_proton) +bunch.getSyncParticle().kinEnergy(1.000) + +# Calculate transfer matrix +matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) +lattice_params = matrix_lattice.getRingParametersDict() +pprint(lattice_params) + +# Store some parameters as variables +lattice_alpha_x = lattice_params["alpha x"] +lattice_alpha_y = lattice_params["alpha y"] +lattice_beta_x = lattice_params["beta x [m]"] +lattice_beta_y = lattice_params["beta y [m]"] +lattice_eta_x = lattice_params["dispersion x [m]"] +lattice_etap_x = lattice_params["dispersion prime x"] + +# Add tune diagnostic node +tune_node = TeapotTuneAnalysisNode() + +tune_node.setNormMatrixFromTwiss( + betax=lattice_beta_x, + alphax=lattice_alpha_x, + etax=lattice_eta_x, + etapx=lattice_etap_x, + betay=lattice_beta_y, + alphay=lattice_alpha_y, +) + +lattice.getNodes()[0].addChildNode(tune_node, 0) + +# Generate particles +emittance_x = 0.1e-06 +emittance_y = 0.1e-06 +twiss_x = TwissContainer(lattice_alpha_x, lattice_beta_x, emittance_x) +twiss_y = TwissContainer(lattice_alpha_y, lattice_beta_y, emittance_y) +dist = GaussDist2D(twiss_x, twiss_y) + +for index in range(10_000): + x, xp, y, yp = dist.getCoordinates() + z = random.uniform(-25.0, 25.0) + bunch.addParticle(x, xp, y, yp, z, 0.0) + +# Track particles +for turn in range(20): + lattice.trackBunch(bunch) + + + if turn >= 0: + phase_data = tune_node.getData(bunch) + nu_x = np.mean(phase_data["tune_1"]) + nu_y = np.mean(phase_data["tune_2"]) + + print("turn{} nux={:0.5f} nuy={:0.5f}".format(turn, nu_x, nu_y)) + + if turn == 5: + tune_node.setNormMatrixFromBunch(bunch) + + if turn == 10: + tune_node.setNormMatrixFromTwiss( + betax=lattice_beta_x, + alphax=lattice_alpha_x, + etax=lattice_eta_x, + etapx=lattice_etap_x, + betay=lattice_beta_y, + alphay=lattice_alpha_y, + ) diff --git a/examples/Diagnostics/Tunes/utils.py b/examples/Diagnostics/Tunes/utils.py new file mode 100644 index 00000000..dc25c6c3 --- /dev/null +++ b/examples/Diagnostics/Tunes/utils.py @@ -0,0 +1,38 @@ +import numpy as np + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.lattice import AccLattice +from orbit.teapot import TEAPOT_Ring +from orbit.teapot import TEAPOT_MATRIX_Lattice + + +def get_tmat(lattice: AccLattice, bunch: Bunch) -> np.ndarray: + matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) + + M = np.zeros((4, 4)) + for i in range(4): + for j in range(4): + M[i, j] = matrix_lattice.getOneTurnMatrix().get(i, j) + return M + + +def make_lattice_sns(sol: bool = False) -> TEAPOT_Ring: + lattice = TEAPOT_Ring() + lattice.readMADX("inputs/sns_ring.lat", "rnginjsol") + lattice.initialize() + + for node in lattice.getNodes(): + try: + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + except: + pass + + for name in ["scbdsol_c13a", "scbdsol_c13b"]: + node = lattice.getNodeForName(name) + B = 0.0 + if sol: + B = 0.15 / 2.0 + node.setParam("B", B) + return lattice \ No newline at end of file diff --git a/examples/Envelope/run.sh b/examples/Envelope/run.sh index 98cf808c..078af457 100755 --- a/examples/Envelope/run.sh +++ b/examples/Envelope/run.sh @@ -2,15 +2,28 @@ set -x -python test_env_2d_fodo.py -python test_env_2d_fodo_speed.py -python test_env_3d_drift.py +python test_env_2d_fodo.py --sc 0 +python test_env_2d_fodo.py --sc 1 +python test_env_2d_fodo.py --sc 1 --offset-x 0.001 +python test_env_2d_fodo.py --sc 1 --tilt 45.0 +python test_env_2d_fodo_speed.py --sc 0 +python test_env_2d_fodo_speed.py --sc 1 +python test_env_3d_drift.py --sc 0 +python test_env_3d_drift.py --sc 1 +python test_env_3d_drift.py --sc 1 --rms-y 0.002 --tilt-z 45.0 +python test_env_3d_drift.py --sc 1 --rms-z 0.002 --tilt-x 45.0 cd sns_linac -python test_sns_linac.py +python test_sns_linac.py --sc 0 +python test_sns_linac.py --sc 1 --dist kv +python test_sns_linac.py --sc 1 --dist waterbag +python test_sns_linac.py --sc 1 --dist gauss cd .. cd sns_ring -python test_sns_ring.py -python test_sns_ring_speed.py +python test_sns_ring.py --sc 0 +python test_sns_ring.py --sc 1 +python test_sns_ring.py --sc 1 --tilt 45.0 +python test_sns_ring_speed.py --sc 0 +python test_sns_ring_speed.py --sc 1 cd .. diff --git a/examples/Envelope/sns_linac/test_sns_linac.py b/examples/Envelope/sns_linac/test_sns_linac.py index 7614d54f..11cccdd9 100755 --- a/examples/Envelope/sns_linac/test_sns_linac.py +++ b/examples/Envelope/sns_linac/test_sns_linac.py @@ -11,8 +11,10 @@ import argparse import math import os +import pathlib import random import sys +import time import numpy as np import matplotlib.pyplot as plt @@ -30,7 +32,6 @@ from orbit.bunch_generators import KVDist3D from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.lattice import AccLattice from orbit.lattice import AccNode from orbit.lattice import AccActionsContainer @@ -66,16 +67,47 @@ def parse_args() -> argparse.Namespace: return parser.parse_args() -def main(args: argparse.Namespace) -> None: +def make_lattice(args: argparse.Namespace) -> LinacAccLattice: + seq_names = [ + "MEBT", + "DTL1", + "DTL2", + "DTL3", + "DTL4", + "DTL5", + "DTL6", + "CCL1", + "CCL2", + "CCL3", + "CCL4", + "SCLMed", + "SCLHigh", + "HEBT1", + "HEBT2", + ] + if args.seq_stop: + index = seq_names.index(args.seq_stop) + 1 + seq_names = seq_names[:index] - output_dir = "outputs" - os.makedirs(output_dir, exist_ok=True) + sns_linac_factory = SNS_LinacLatticeFactory() + sns_linac_factory.setMaxDriftLength(args.sc_path_length_min) + lattice = sns_linac_factory.getLinacAccLattice(seq_names, "inputs/sns_linac.xml") - random.seed(23) + for node in lattice.getNodes(): + try: + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + except: + pass - # Bunch - # -------------------------------------------------------------------------------- + rf_gaps = lattice.getRF_Gaps() + for rf_gap in rf_gaps: + rf_gap.setCppGapModel(MatrixRfGap()) + return lattice + + +def make_bunch(args: argparse.Namespace) -> Bunch: kin_energy = 0.0025 # [GeV] mass = mass_proton + 2.0 * mass_electron frequency = 402.5e06 @@ -110,72 +142,34 @@ def main(args: argparse.Namespace) -> None: for _ in range(args.nparts): bunch.addParticle(*dist.getCoordinates()) + return bunch - # Lattice - # -------------------------------------------------------------------------------- - seq_names = [ - "MEBT", - "DTL1", - "DTL2", - "DTL3", - "DTL4", - "DTL5", - "DTL6", - "CCL1", - "CCL2", - "CCL3", - "CCL4", - "SCLMed", - "SCLHigh", - "HEBT1", - "HEBT2", - ] - if args.seq_stop: - index = seq_names.index(args.seq_stop) + 1 - seq_names = seq_names[:index] - - sns_linac_factory = SNS_LinacLatticeFactory() - sns_linac_factory.setMaxDriftLength(args.sc_path_length_min) - lattice = sns_linac_factory.getLinacAccLattice(seq_names, "inputs/sns_linac.xml") - - for node in lattice.getNodes(): - try: - node.setUsageFringeFieldIN(False) - node.setUsageFringeFieldOUT(False) - except: - pass - - rf_gaps = lattice.getRF_Gaps() - for rf_gap in rf_gaps: - rf_gap.setCppGapModel(MatrixRfGap()) +def main(args: argparse.Namespace) -> None: - for index, node in enumerate(lattice.getNodes()): - print(index, type(node), node.getName()) + path = pathlib.Path(__file__) + output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S")) + os.makedirs(output_dir, exist_ok=True) - lattice.trackDesignBunch(bunch) + random.seed(23) # Track envelope - # -------------------------------------------------------------------------------- - - twiss_calc = BunchTwissAnalysis() - twiss_calc.analyzeBunch(bunch) - - cov_matrix = np.zeros((6, 6)) - for i in range(6): - for j in range(6): - cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j) + bunch = make_bunch(args) + envelope = Envelope(bunch=bunch) - envelope = Envelope(bunch=bunch, cov_matrix=cov_matrix, intensity=intensity) - - tracker = EnvelopeTracker(lattice, sc=("3d" if args.sc else None)) + lattice = make_lattice(args) + lattice.trackDesignBunch(bunch) histories = {} - histories["envelope"] = tracker.track_history(envelope) + histories["envelope"] = lattice.trackEnvelope( + envelope, + history=True, + sc=("3d" if args.sc else None) + ) # Track bunch - # -------------------------------------------------------------------------------- - + bunch = make_bunch(args) + lattice = make_lattice(args) lattice.trackDesignBunch(bunch) if args.sc: diff --git a/examples/Envelope/sns_ring/test_sns_ring.py b/examples/Envelope/sns_ring/test_sns_ring.py index 4efe4289..3fafa636 100644 --- a/examples/Envelope/sns_ring/test_sns_ring.py +++ b/examples/Envelope/sns_ring/test_sns_ring.py @@ -5,8 +5,8 @@ import math import os import pathlib -import time import sys +import time import numpy as np import matplotlib.pyplot as plt @@ -16,12 +16,12 @@ from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.space_charge.sc2p5d import setSC2p5DAccNodes from orbit.teapot import TEAPOT_Ring from orbit.teapot import TEAPOT_MATRIX_Lattice from orbit.teapot import teapot +from orbit.teapot import BendTEAPOT from orbit.utils.consts import mass_proton sys.path.append("..") @@ -65,7 +65,7 @@ def parse_args() -> argparse.Namespace: def main(args: argparse.Namespace) -> None: path = pathlib.Path(__file__) - output_dir = os.path.join("outputs", path.stem) + output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S")) os.makedirs(output_dir, exist_ok=True) # Lattice @@ -122,7 +122,7 @@ def main(args: argparse.Namespace) -> None: if args.tilt: rot_matrix = np.identity(6) - rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi)) + rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=np.radians(args.tilt)) cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T]) if args.mismatch_x or args.mismatch_y: @@ -157,12 +157,11 @@ def main(args: argparse.Namespace) -> None: print("TRACK ENVELOPE") envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, centroid=centroid_init, intensity=args.intensity, ) - tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) history_keys = [ "rms_x", @@ -178,7 +177,7 @@ def main(args: argparse.Namespace) -> None: for turn in range(args.turns + 1): if turn > 0: - tracker.track_ring(envelope) + lattice.trackEnvelopeRing(envelope, sc=("2d" if args.sc else None)) cov_matrix = envelope.cov_matrix centroid = envelope.centroid diff --git a/examples/Envelope/sns_ring/test_sns_ring_speed.py b/examples/Envelope/sns_ring/test_sns_ring_speed.py index 773febc6..ea119c71 100644 --- a/examples/Envelope/sns_ring/test_sns_ring_speed.py +++ b/examples/Envelope/sns_ring/test_sns_ring_speed.py @@ -12,16 +12,17 @@ from orbit.core.bunch import Bunch from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.space_charge.sc2p5d import setSC2p5DAccNodes from orbit.teapot import TEAPOT_Ring from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.teapot import BendTEAPOT from orbit.utils.consts import mass_proton sys.path.append("..") from utils import gen_dist + parser = argparse.ArgumentParser() parser.add_argument("--bunch-length", type=float, default=120.0) parser.add_argument("--kin-energy", type=float, default=1.300) @@ -45,6 +46,11 @@ except: pass +for node in lattice.getNodes(): + if type(node) is BendTEAPOT: + node.setParam("ea1", 0.0) + node.setParam("ea2", 0.0) + for node in lattice.getNodes(): max_length = 1.0 if node.getLength() > max_length: @@ -80,11 +86,11 @@ print("ENVELOPE") envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, intensity=args.intensity, ) -tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) +envelope_sc = "2d" if args.sc else None start_time = time.time() @@ -92,7 +98,7 @@ profiler.enable() for turn in trange(args.turns): - tracker.track_ring(envelope) + lattice.trackEnvelopeRing(envelope, sc=envelope_sc) time_per_turn = (time.time() - start_time) / args.turns diff --git a/examples/Envelope/test_env_2d_fodo.py b/examples/Envelope/test_env_2d_fodo.py index 824cc3b7..969e99ff 100644 --- a/examples/Envelope/test_env_2d_fodo.py +++ b/examples/Envelope/test_env_2d_fodo.py @@ -5,6 +5,7 @@ import math import os import pathlib +import time import numpy as np import matplotlib.pyplot as plt @@ -14,7 +15,6 @@ from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.lattice import AccLattice from orbit.lattice import AccNode from orbit.core.spacecharge import SpaceChargeCalc2p5D @@ -59,12 +59,8 @@ def parse_args() -> argparse.Namespace: def main(args: argparse.Namespace) -> None: - - # Setup - # ------------------------------------------------------------------------------ - path = pathlib.Path(__file__) - output_dir = os.path.join("outputs", path.stem) + output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S")) os.makedirs(output_dir, exist_ok=True) # Create lattice @@ -120,7 +116,7 @@ def main(args: argparse.Namespace) -> None: # Tilt if args.tilt: rot_matrix = np.identity(6) - rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi)) + rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=np.radians(args.tilt)) cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T]) # Mismatch @@ -135,7 +131,7 @@ def main(args: argparse.Namespace) -> None: # Create envelope envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, centroid=centroid_init, intensity=args.intensity, @@ -146,12 +142,12 @@ def main(args: argparse.Namespace) -> None: print("TRACK ENVELOPE") - tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) + envelope_sc = "2d" if args.sc else None history = {"xrms": [], "yrms": [], "xavg": [], "yavg": []} for turn in range(args.turns): if turn > 0: - tracker.track_ring(envelope) + lattice.trackEnvelopeRing(envelope, sc=envelope_sc) cov_matrix = envelope.cov_matrix centroid = envelope.centroid diff --git a/examples/Envelope/test_env_2d_fodo_speed.py b/examples/Envelope/test_env_2d_fodo_speed.py index 063ddfaf..7462d289 100644 --- a/examples/Envelope/test_env_2d_fodo_speed.py +++ b/examples/Envelope/test_env_2d_fodo_speed.py @@ -11,7 +11,6 @@ from orbit.core.bunch import Bunch from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.core.spacecharge import SpaceChargeCalc2p5D from orbit.space_charge.sc2p5d import setSC2p5DAccNodes from orbit.teapot import QuadTEAPOT @@ -31,7 +30,7 @@ parser.add_argument("--kq", type=float, default=0.25) parser.add_argument("--nparts", type=int, default=10_000) -parser.add_argument("--turns", type=int, default=5000) +parser.add_argument("--turns", type=int, default=500) parser.add_argument("--sc", type=int, default=0) parser.add_argument("--sc-grid", type=int, default=64) args = parser.parse_args() @@ -82,11 +81,11 @@ print("ENVELOPE") envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, intensity=args.intensity, ) -tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None)) +envelope_sc = "2d" if args.sc else None start_time = time.time() @@ -94,7 +93,7 @@ profiler.enable() for turn in trange(args.turns): - tracker.track_ring(envelope) + lattice.trackEnvelopeRing(envelope, sc=envelope_sc) time_per_turn = (time.time() - start_time) / args.turns diff --git a/examples/Envelope/test_env_3d_drift.py b/examples/Envelope/test_env_3d_drift.py index 798d4448..992ec27d 100644 --- a/examples/Envelope/test_env_3d_drift.py +++ b/examples/Envelope/test_env_3d_drift.py @@ -9,6 +9,7 @@ import math import os import pathlib +import time import numpy as np import matplotlib.pyplot as plt @@ -19,7 +20,6 @@ from orbit.core.spacecharge import SpaceChargeCalc3D from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.space_charge.sc3d import setSC3DAccNodes from orbit.teapot import DriftTEAPOT from orbit.teapot import TEAPOT_Lattice @@ -42,9 +42,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--rms-y", type=float, default=0.010) parser.add_argument("--rms-z", type=float, default=0.010) - parser.add_argument("--rot-x", type=float, default=0.0) - parser.add_argument("--rot-y", type=float, default=0.0) - parser.add_argument("--rot-z", type=float, default=0.0) + parser.add_argument("--tilt-x", type=float, default=0.0) + parser.add_argument("--tilt-y", type=float, default=0.0) + parser.add_argument("--tilt-z", type=float, default=0.0) parser.add_argument("--nslice", type=int, default=10) parser.add_argument("--length", type=float, default=0.1) @@ -74,7 +74,7 @@ def build_cov_matrix_xyz( def main(args: argparse.Namespace) -> None: path = pathlib.Path(__file__) - output_dir = os.path.join("outputs", path.stem) + output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S")) os.makedirs(output_dir, exist_ok=True) # Create lattice @@ -99,7 +99,7 @@ def main(args: argparse.Namespace) -> None: cov_matrix_init = np.zeros((6, 6)) rotation_matrix = rotation_matrix_3d( - math.radians(args.rot_x), math.radians(args.rot_y), math.radians(args.rot_z) + math.radians(args.tilt_x), math.radians(args.tilt_y), math.radians(args.tilt_z) ) print(rotation_matrix) @@ -120,7 +120,7 @@ def main(args: argparse.Namespace) -> None: centroid_init = np.zeros(6) envelope = Envelope( - bunch=bunch, + sync_part=sync_part, cov_matrix=cov_matrix_init, centroid=centroid_init, intensity=args.intensity, @@ -131,12 +131,12 @@ def main(args: argparse.Namespace) -> None: print("TRACK ENVELOPE") - tracker = EnvelopeTracker(lattice, sc=("3d" if args.sc else None)) + envelope_sc = "3d" if args.sc else None history = {"xrms": [], "yrms": [], "zrms": []} for turn in range(args.turns): if turn > 0: - tracker.track(envelope) + lattice.trackEnvelope(envelope, sc=envelope_sc) cov_matrix = envelope.cov_matrix diff --git a/meson.build b/meson.build index b2bb9c79..d37025ea 100644 --- a/meson.build +++ b/meson.build @@ -1,4 +1,13 @@ -project('PyORBIT', ['cpp'], version: run_command('python', '-m', 'setuptools_scm').stdout().strip()) +project('PyORBIT', + ['cpp'], + version: run_command('python', '-m', 'setuptools_scm').stdout().strip(), + default_options: [ + 'buildtype=release', + 'b_lto=true', + 'cpp_std=c++11', + 'warning_level=3', + ], +) subdir('src') subdir('py/orbit') diff --git a/py/orbit/diagnostics/TeapotDiagnosticsNode.py b/py/orbit/diagnostics/TeapotDiagnosticsNode.py index 4d7b2b42..ab26978f 100644 --- a/py/orbit/diagnostics/TeapotDiagnosticsNode.py +++ b/py/orbit/diagnostics/TeapotDiagnosticsNode.py @@ -1,247 +1,347 @@ -""" -This module is a collimator node class for TEAPOT lattice -""" +"""TEAPOT-style bunch diagnostic nodes.""" +from typing import IO -import os -import math +import numpy as np -# import the auxiliary classes -from ..utils import orbitFinalize, NamedObject, ParamsDictObject +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.core.bunch import BunchTuneAnalysis +from orbit.teapot import DriftTEAPOT -# import general accelerator elements and lattice -from ..lattice import AccNode, AccActionsContainer, AccNodeBunchTracker - - -# import Diagnostics classes -from .diagnostics import StatLats, StatLatsSetMember -from .diagnostics import Moments, MomentsSetMember, BPMSignal - -# import teapot drift class -from ..teapot import DriftTEAPOT - - -# import Bunch diagnostics -from orbit.core import bunch - -BunchTuneAnalysis = bunch.BunchTuneAnalysis +from .diagnostics import StatLats +from .diagnostics import StatLatsSetMember +from .diagnostics import Moments +from .diagnostics import MomentsSetMember +from .diagnostics import BPMSignal +from .matrix import build_norm_matrix_from_cov +from .matrix import build_norm_matrix_from_tmat +from .matrix import TransferMatrixAnalysis class TeapotStatLatsNode(DriftTEAPOT): - """ - The statlats node class for TEAPOT lattice - """ - - def __init__(self, filename, name="statlats no name"): - """ - Constructor. Creates the StatLats TEAPOT element. - """ + def __init__(self, filename: str, name: str = "statlats no name") -> None: DriftTEAPOT.__init__(self, name) self.statlats = StatLats(filename) self.setType("statlats teapot") self.setLength(0.0) self.position = 0.0 - self.lattlength = 0.0 + self.lattice_length = 0.0 self.file_out = open(filename, "w") - def track(self, paramsDict): - """ - The statlats-teapot class implementation of the AccNodeBunchTracker class track(probe) method. - """ - length = self.getLength(self.getActivePartIndex()) - bunch = paramsDict["bunch"] - self.statlats.writeStatLats(self.position, bunch, self.lattlength) + def track(self, params_dict: dict) -> None: + bunch = params_dict["bunch"] + self.statlats.writeStatLats(self.position, bunch, self.lattice_length) - def setPosition(self, pos): - self.position = pos + def setPosition(self, position: float) -> None: + self.position = position - def closeStatLats(self): + def closeStatLats(self) -> None: self.file_out.close() - def setLatticeLength(self, lattlength): - self.lattlength = lattlength + def setLatticeLength(self, length: float) -> None: + self.lattice_length = length class TeapotStatLatsNodeSetMember(DriftTEAPOT): - """ - The statlats node class for TEAPOT lattice - """ - - def __init__(self, file, name="statlats no name"): - """ - Constructor. Creates the StatLats TEAPOT element. - """ + def __init__(self, file: IO[str], name: str = "statlats no name") -> None: DriftTEAPOT.__init__(self, name) self.statlats = StatLatsSetMember(file) self.setType("statlats teapot") self.setLength(0.0) self.position = 0.0 - self.lattlength = 0.0 + self.lattice_length = 0.0 self.active = True self.file = file - def track(self, paramsDict): - """ - The statlats-teapot class implementation of the AccNodeBunchTracker class track(probe) method. - """ + def track(self, params_dict: dict) -> None: if self.active: - length = self.getLength(self.getActivePartIndex()) - bunch = paramsDict["bunch"] - self.statlats.writeStatLats(self.position, bunch, self.lattlength) + bunch = params_dict["bunch"] + self.statlats.writeStatLats(self.position, bunch, self.lattice_length) - def setPosition(self, pos): - self.position = pos + def setPosition(self, position: float) -> None: + self.position = position - def setLatticeLength(self, lattlength): - self.lattlength = lattlength + def setLatticeLength(self, length: float) -> None: + self.lattice_length = length - def activate(self): + def activate(self) -> None: self.active = True - def deactivate(self): + def deactivate(self) -> None: self.active = False - def resetFile(self, file): + def resetFile(self, file: IO[str]) -> None: self.file = file self.statlats.resetFile(self.file) class TeapotMomentsNode(DriftTEAPOT): - """ - The moments node class for TEAPOT lattice - """ - - def __init__(self, filename, order, nodispersion=True, emitnorm=False, name="moments no name"): - """ - Constructor. Creates the StatLats TEAPOT element. - """ + def __init__(self, filename: str, order: int, no_dispersion: bool = True, emit_norm: bool = False, name: str = "moments no name") -> None: DriftTEAPOT.__init__(self, name) - self.moments = Moments(filename, order, nodispersion, emitnorm) + self.moments = Moments(filename, order, no_dispersion, emit_norm) self.setType("moments teapot") self.setLength(0.0) self.position = 0.0 - self.lattlength = 0.0 + self.lattice_length = 0.0 self.file_out = open(filename, "w") - def track(self, paramsDict): - """ - The moments-teapot class implementation of the AccNodeBunchTracker class track(probe) method. - """ - length = self.getLength(self.getActivePartIndex()) - bunch = paramsDict["bunch"] - self.moments.writeMoments(self.position, bunch, self.lattlength) + def track(self, params_dict: dict) -> None: + bunch = params_dict["bunch"] + self.moments.writeMoments(self.position, bunch, self.lattice_length) - def setPosition(self, pos): - self.position = pos + def setPosition(self, position: float) -> None: + self.position = position - def closeMoments(self): + def closeMoments(self) -> None: self.file_out.close() - def setLatticeLength(self, lattlength): - self.lattlength = lattlength + def setLatticeLength(self, length: float) -> None: + self.lattice_length = length class TeapotMomentsNodeSetMember(DriftTEAPOT): - """ - The moments node class for TEAPOT lattice - """ - - def __init__(self, file, order, nodispersion=True, emitnorm=False, name="moments no name"): - """ - Constructor. Creates the Moments TEAPOT element. - """ + def __init__(self, file: IO[str], order: int, no_dispersion: bool = True, emit_norm: bool = False, name: str = "moments no name") -> None: DriftTEAPOT.__init__(self, str(name)) self.file = file - self.moments = MomentsSetMember(self.file, order, nodispersion, emitnorm) + self.moments = MomentsSetMember(self.file, order, no_dispersion, emit_norm) self.setType("moments teapot") self.setLength(0.0) self.position = 0.0 - self.lattlength = 0.0 + self.lattice_length = 0.0 self.active = True - def track(self, paramsDict): - """ - The moments-teapot class implementation of the AccNodeBunchTracker class track(probe) method. - """ + def track(self, params_dict: dict) -> None: if self.active: - length = self.getLength(self.getActivePartIndex()) - bunch = paramsDict["bunch"] - self.moments.writeMoments(self.position, bunch, self.lattlength) + bunch = params_dict["bunch"] + self.moments.writeMoments(self.position, bunch, self.lattice_length) - def setPosition(self, pos): - self.position = pos + def setPosition(self, position: float) -> None: + self.position = position - def setLatticeLength(self, lattlength): - self.lattlength = lattlength + def setLatticeLength(self, length: float) -> None: + self.lattice_length = length - def activate(self): + def activate(self) -> None: self.active = True - def deactivate(self): + def deactivate(self) -> None: self.active = False - def resetFile(self, file): + def resetFile(self, file: IO[str]) -> None: self.file = file self.moments.resetFile(self.file) class TeapotTuneAnalysisNode(DriftTEAPOT): - def __init__(self, name="tuneanalysis no name"): - """ - Constructor. Creates the StatLats TEAPOT element. - """ + """Estimates tunes from coordinates on neighboring turns. + + This node computes the tunes and actions of each particle in the bunch. + We use the Average Phase Advance (APA) method to estimate the tunes [1]. + We use only a single turn rather than the average of multiple turns. + + [1] https://cds.cern.ch/record/292773/files/p147.pdf + [2] https://arxiv.org/pdf/1207.5526 + [3] S. Y. Lee, *Accelerator Physics* + """ + def __init__(self, name: str = "TeapotTuneAnalysis no name") -> None: DriftTEAPOT.__init__(self, name) - self.bunchtune = BunchTuneAnalysis() + self.tune_calc = BunchTuneAnalysis() self.setType("tune calculator teapot") - self.lattlength = 0.0 self.setLength(0.0) self.position = 0.0 + self.active = True + self.keys = ["phase_1", "phase_2", "tune_1", "tune_2", "action_1", "action_2"] - def track(self, paramsDict): - """ - The bunchtuneanalysis-teapot class implementation of the AccNodeBunchTracker class track(probe) method. + def track(self, params_dict: dict) -> None: + if self.active: + self.tune_calc.analyzeBunch(params_dict["bunch"]) + + def activate(self) -> None: + self.active = True + + def deactivate(self) -> None: + self.active = False + + def setPosition(self, position: float) -> None: + self.position = position + + def setNormMatrix(self, norm_matrix: np.ndarray) -> None: + ndim = norm_matrix.shape[0] + for i in range(ndim): + for j in range(ndim): + self.tune_calc.setNormMatrixElement(i, j, norm_matrix[i, j]) + + def getNormMatrix(self) -> np.ndarray: + norm_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(6): + norm_matrix[i][j] = self.tune_calc.getNormMatrixElement(i, j) + return norm_matrix + + def setNormMatrixFromTwiss( + self, + betax: float, + alphax: float, + etax: float, + etapx: float, + betay: float, + alphay: float, + ) -> None: + """Set normalization matrix from Twiss parameters (x, y) and dispersion. + + betax{y}: Beta parameter in x{y} plane. + alphax{y}: Alpha parameter in x{y} plane. + etax: Dispersion in x plane. + etapx: Disperion prime in x plane. """ - length = self.getLength(self.getActivePartIndex()) - bunch = paramsDict["bunch"] - self.bunchtune.analyzeBunch(bunch) + self.tune_calc.setNormMatrixFromTwiss(betax, alphax, etax, etapx, betay, alphay) - def setPosition(self, pos): - self.position = pos + def setNormMatrixFromTransferMatrix(self, transfer_matrix: np.ndarray) -> None: + """Set normalization matrix from transfer matrix. - def setLatticeLength(self, lattlength): - self.lattlength = lattlength + Assumes transfer matrix is periodic and stable. - def assignTwiss(self, betax, alphax, etax, etapx, betay, alphay): - self.bunchtune.assignTwiss(betax, alphax, etax, etapx, betay, alphay) + Args: + transfer_matrix: 4x4 or 6x6 transfer matrix. + """ + assert transfer_matrix.shape[0] == transfer_matrix.shape[1] + assert transfer_matrix.shape[0] in (4, 6) + norm_matrix = build_norm_matrix_from_tmat(transfer_matrix) + self.setNormMatrix(norm_matrix) + + def setNormMatrixFromCovMatrix(self, cov_matrix: np.ndarray) -> None: + """Set normalization matrix from covariance matrix. + + Assume that S = M S M^T, where S is the covariance matrix + and M is the transfer matrix. Then M and SU (U is the Poisson matrix) + have different eigenvalues but the same eigenvectors So we can compute + the normalization matrix directly from SU, without knowing M. + + I'm not sure how to order the eigenvectors of SU as there is no + guaranteed ordering from np.linalg.eig. By default, we sort the + eigenvectors of SU by their eigenvalues (eigenemittances), so the + smallest eigenemittance is mode 1, the next is mode 2, and so on. + So if you compare this method to `setNormMatrixFromTransferMatrix`, + you may get {nu1, nu2} -> {nu2, nu1}. + + This is only a problem in coupled lattices with 4D normalization. + With 2D normalization there is no ambiguity. This function will + check if there are off-block-diagonal terms in the covariance + matrix to determine whether to use 2D or 4D normalization. + norm_matrix = build_norm_matrix_from_cov(cov_matrix) + self.setNormMatrix(norm_matrix) + + Args: + cov_matrix: 4x4 or 6x6 covariance matrix. + """ + assert cov_matrix.shape[0] == cov_matrix.shape[1] + assert cov_matrix.shape[0] in (4, 6) + norm_matrix = build_norm_matrix_from_cov(cov_matrix) + self.setNormMatrix(norm_matrix) + def setNormMatrixFromBunch(self, bunch: Bunch, dim: int = 4) -> None: + """Set normalization matrix from bunch covariance matrix. -class TeapotBPMSignalNode(DriftTEAPOT): - def __init__(self, name="BPMSignal no name"): + The bunch covariance matrix is calculated from the macroparticles. + + Args: + bunch: Bunch object with at least one macroparticle. + """ + assert dim in (4, 6) + assert bunch.getSizeGlobal() > 0 + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + cov_matrix = np.zeros((dim, dim)) + for i in range(dim): + for j in range(dim): + cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j) + self.setNormMatrixFromCovMatrix(cov_matrix) + + def getData(self, bunch: Bunch, index: int = None) -> dict[str, float] | dict[str, np.ndarray]: + """Return tune and action data. + + Args: + bunch: A Bunch object. + index: Particle index. If None, return data for all particles. + + Returns: + data: Dictionary with the following keys: + - "phase_1" + - "phase_2" + - "tune_1" + - "tune_2" + - "action_1" + - "action_2" + - "action_3" + + If `index` is provided, each value is a float. Otherwise each value + is a list of floats. If the lattice is uncoupled, 1->x and 2->y. """ - Constructor. Creates the StatLats TEAPOT element. + data = {} + if index is None: + for j, key in enumerate(self.keys): + data[key] = [] + for index in range(bunch.getSize()): + value = bunch.partAttrValue("ParticlePhaseAttributes", index, j) + data[key].append(value) + data[key] = np.array(data[key]) + else: + index = int(index) + bunch_size = bunch.getSize() + if (index < bunch_size): + raise ValueError("particle index < 0") + if (index > bunch_size - 1): + raise ValueError("particle index > bunch.getSize() - 1") + for j, key in enumerate(self.keys): + data[key] = bunch.partAttrValue("ParticlePhaseAttributes", index, j) + return data + + def getTunes(self, bunch: Bunch, index: int = None) -> tuple[float, float] | tuple[np.ndarray, np.ndarray]: + """Return fractional tunes (nu_1, nu_2). + + Args: + bunch: A Bunch object. + index: Particle index (not ID). + + Returns: + tune_1: Fractional tune (mode 1). + tune_2: Fractional tune (mode 2). """ + data = self.getData(bunch, index) + return tuple([data[key] for key in ["tune_1", "tune_2"]]) + + def getActions(self, bunch: Bunch, index: int) -> tuple[float, float] | tuple[np.ndarray, np.ndarray]: + """Return actions (J_1, J_2). + + Args: + bunch: Bunch object. + index: Particle index (not ID). + + Returns: + J_1: Action (mode 1). + J_2: Action (mode 2). + """ + data = self.getData(bunch, index) + return tuple([data[key] for key in ["action_1", "action_2"]]) + + +class TeapotBPMSignalNode(DriftTEAPOT): + def __init__(self, name: str = "BPMSignal no name") -> None: DriftTEAPOT.__init__(self, name) self.bpm = BPMSignal() self.setType("BPMSignal") - self.lattlength = 0.0 self.setLength(0.0) self.position = 0.0 - def track(self, paramsDict): - """ - The bunchtuneanalysis-teapot class implementation of the AccNodeBunchTracker class track(probe) method. - """ - length = self.getLength(self.getActivePartIndex()) - bunch = paramsDict["bunch"] - self.bpm.analyzeSignal(bunch) - - def setPosition(self, pos): - self.position = pos + def track(self, params_dict: dict) -> None: + bunch = params_dict["bunch"] + self.bpm.analyzeSignal(params_dict["bunch"]) - def setLatticeLength(self, lattlength): - self.lattlength = lattlength + def setPosition(self, position: float) -> None: + self.position = position - def getSignal(self): - xAvg = self.bpm.getSignalX() - yAvg = self.bpm.getSignalY() - return xAvg, yAvg + def getSignal(self) -> tuple[float, float]: + x_avg = self.bpm.getSignalX() + y_avg = self.bpm.getSignalY() + return x_avg, y_avg diff --git a/py/orbit/diagnostics/__init__.py b/py/orbit/diagnostics/__init__.py index c1d197a3..51c3da15 100644 --- a/py/orbit/diagnostics/__init__.py +++ b/py/orbit/diagnostics/__init__.py @@ -13,6 +13,7 @@ from .TeapotDiagnosticsNode import TeapotStatLatsNode, TeapotStatLatsNodeSetMember from .TeapotDiagnosticsNode import TeapotMomentsNode, TeapotMomentsNodeSetMember from .TeapotDiagnosticsNode import TeapotTuneAnalysisNode +from . import matrix __all__ = [] @@ -30,3 +31,4 @@ __all__.append("addTeapotMomentsNodeSet") __all__.append("TeapotTuneAnalysisNode") __all__.append("profiles") +__all__.append("matrix") diff --git a/py/orbit/diagnostics/matrix.py b/py/orbit/diagnostics/matrix.py new file mode 100644 index 00000000..36f51e66 --- /dev/null +++ b/py/orbit/diagnostics/matrix.py @@ -0,0 +1,124 @@ +import numpy as np + + +def calc_eigtune(eigval: float) -> float: + return np.arccos(np.real(eigval)) / (2.0 * np.pi) + + +def build_poisson_matrix(ndim: int) -> np.ndarray: + U = np.zeros((ndim, ndim)) + for i in range(0, ndim, 2): + U[i : i + 2, i : i + 2] = [[0.0, 1.0], [-1.0, 0.0]] + return U + + +def normalize_eigvec(v: np.ndarray) -> np.ndarray: + v = np.copy(v) + U = build_poisson_matrix(len(v)) + + def complex_amplitude(v): + return np.conj(v).T @ U @ v + + if np.imag(complex_amplitude(v)) > 0.0: + v = np.conj(v) + + v *= np.sqrt(2.0 / np.abs(complex_amplitude(v))) + + assert np.isclose(np.imag(complex_amplitude(v)), -2.0) + assert np.isclose(np.real(complex_amplitude(v)), +0.0) + return v + + +def build_norm_matrix_from_eigvecs(*eigvecs: list[np.ndarray]) -> np.ndarray: + dim = len(eigvecs[0]) + + V = np.zeros((dim, dim)) + for i, v in enumerate(eigvecs): + V[:, 2 * i + 0] = +np.real(v) + V[:, 2 * i + 1] = -np.imag(v) + return np.linalg.inv(V) + + +def symplectic_eig(matrix: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + eigvals, eigvecs = np.linalg.eig(matrix) + eigvecs = eigvecs.T + + eigvals = eigvals[::2] + eigvecs = eigvecs[::2] + + for i in range(eigvecs.shape[0]): + eigvecs[i] = normalize_eigvec(eigvecs[i]) + + return eigvals, eigvecs + + +def build_norm_matrix_from_tmat(matrix: np.ndarray) -> np.ndarray: + eigvals, eigvecs = symplectic_eig(matrix) + return build_norm_matrix_from_eigvecs(*eigvecs) + + +def build_norm_matrix_from_cov_uncoupled(cov_matrix: np.ndarray) -> np.ndarray: + S = cov_matrix + U = build_poisson_matrix(S.shape[0]) + + V_inv = np.zeros_like(S) + for i in range(0, S.shape[0], 2): + eigvals, eigvecs = symplectic_eig(S[i:i+2, i:i+2] @ U[i:i+2, i:i+2]) + V_inv[i:i+2, i:i+2] = build_norm_matrix_from_eigvecs(*eigvecs) + return V_inv + + +def build_norm_matrix_from_cov(cov_matrix: np.ndarray) -> np.ndarray: + if not is_coupled(cov_matrix): + return build_norm_matrix_from_cov_uncoupled(cov_matrix) + + S = cov_matrix + U = build_poisson_matrix(S.shape[0]) + + eigvals, eigvecs = symplectic_eig(S @ U) + + emittances = np.abs(np.imag(eigvals)) + if np.all(np.abs(emittances - emittances[0]) < 1e-15): + raise ValueError("Eigenemittances are equal eigenvectors are degenerate and V will not be correct.") + + # Order by emittances + order = np.argsort(emittances) + eigvals = eigvals[order] + eigvecs = eigvecs[order] + return build_norm_matrix_from_eigvecs(*eigvecs) + + +def is_coupled(matrix: np.ndarray) -> bool: + assert matrix.ndim == 2 + assert matrix.shape[0] == matrix.shape[1] + assert matrix.shape[0] % 2 == 0 + matrix = np.copy(matrix) + for i in range(0, matrix.shape[0], 2): + matrix[i:i+2, i:i+2] = 0.0 + return not np.all(np.isclose(matrix, 0.0)) + + +class TransferMatrixAnalysis: + def __init__(self, M: np.ndarray) -> None: + self.M = M + self.ndim = M.shape[0] + + self.eigvals, self.eigvecs = np.linalg.eig(M) + self.eigvecs = self.eigvecs.T + self.eigvecs = self.eigvecs[::2] + for i in range(self.eigvecs.shape[0]): + self.eigvecs[i] = normalize_eigvec(self.eigvecs[i]) + + self.eigvals = self.eigvals[::2] + self.eigtunes = [calc_eigtune(eigval) for eigval in self.eigvals] + + self.V_inv = build_norm_matrix_from_eigvecs(*self.eigvecs) + self.V = np.linalg.inv(self.V_inv) + + self.is_coupled = is_coupled(M) + + def cov_matrix(self, *emittances: float) -> np.ndarray: + S = np.eye(self.ndim) + if len(emittances) > 0: + S = np.diag(np.repeat(emittances, 2)) + return self.V @ S @ self.V.T diff --git a/py/orbit/diagnostics/meson.build b/py/orbit/diagnostics/meson.build index 26e40c8e..a91a9016 100644 --- a/py/orbit/diagnostics/meson.build +++ b/py/orbit/diagnostics/meson.build @@ -6,7 +6,8 @@ py_sources = files([ 'TeapotDiagnosticsNode.py', 'diagnosticsLatticeModifications.py', '__init__.py', - 'profiles.py' + 'profiles.py', + 'matrix.py' ]) python.install_sources( diff --git a/py/orbit/envelope/__init__.py b/py/orbit/envelope/__init__.py index 7a4586ae..e2a3e379 100644 --- a/py/orbit/envelope/__init__.py +++ b/py/orbit/envelope/__init__.py @@ -1,2 +1 @@ from .envelope import Envelope -from .track import EnvelopeTracker \ No newline at end of file diff --git a/py/orbit/envelope/envelope.py b/py/orbit/envelope/envelope.py index 003c24ab..3aeab4a9 100644 --- a/py/orbit/envelope/envelope.py +++ b/py/orbit/envelope/envelope.py @@ -5,13 +5,34 @@ from orbit.core.bunch import Bunch from orbit.core.bunch import BunchTwissAnalysis from orbit.core.bunch import SyncParticle +from orbit.utils.matrix import convert_matrix_zp_to_dE -from .utils import convert_matrix_zp_to_dE from .utils import gen_dist from .utils import get_classical_radius from .utils import proj_cov_matrix +def get_bunch_cov_matrix(bunch: Bunch) -> np.ndarray: + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(6): + cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j) + return cov_matrix + + +def get_bunch_centroid(bunch: Bunch) -> np.ndarray: + twiss_calc = BunchTwissAnalysis() + twiss_calc.analyzeBunch(bunch) + + centroid = np.zeros(6) + for i in range(6): + centroid[i] = twiss_calc.getAverage(i) + return centroid + + def build_diag_matrix_from_xyz_eig(eigenvectors: np.ndarray) -> np.ndarray: A = np.eye(7) for i in range(eigenvectors.shape[0]): @@ -26,7 +47,7 @@ class Envelope: """Represents beam envelope/centroid. Attributes: - bunch: Bunch containing synchronous particle and (optionally) test particles. + sync_part: Synchronous particle. cov_matrix: 6 x 6 covariance matrix centroid: 6 x 1 centroid vector. intensity: Total number of particles. @@ -34,45 +55,40 @@ class Envelope: def __init__( self, - bunch: Bunch, + sync_part: SyncParticle = None, cov_matrix: np.ndarray = None, centroid: np.ndarray = None, intensity: float = 0.0, + bunch: Bunch = None, ) -> None: + """Constructor. - empty_bunch = Bunch() - bunch.copyEmptyBunchTo(empty_bunch) + Args: + sync_part: Synchronous particle. + cov_matrix: 6 x 6 covariance matrix. + centroid: 6 x 1 centroid vector. + intensity: Total number of particles. + bunch: If provided, the parameters above are calculated from the bunch particles. + """ - self.bunch = empty_bunch - self.sync_part = self.bunch.getSyncParticle() + if bunch is not None: + sync_part = bunch.getSyncParticle() + cov_matrix = get_bunch_cov_matrix(bunch) + centroid = get_bunch_centroid(bunch) + intensity = bunch.getSize() * bunch.macroSize() + + self.sync_part = sync_part self.centroid = centroid if self.centroid is None: - if bunch.getSize(): - twiss_calc = BunchTwissAnalysis() - twiss_calc.analyzeBunch(bunch) - self.centroid = np.zeros(6) - for i in range(6): - self.centroid[i] = twiss_calc.getAverage(i) - else: - self.centroid = np.zeros(6) + self.centroid = np.zeros(6) self.cov_matrix = cov_matrix if self.cov_matrix is None: - if bunch.getSize(): - twiss_calc = BunchTwissAnalysis() - twiss_calc.analyzeBunch(bunch) - self.cov_matrix = np.zeros((6, 6)) - for i in range(6): - for j in range(6): - self.cov_matrix[i, j] = twiss_calc.getCorrelation(i, j) - self.cov_matrix[j, i] = self.cov_matrix[i, j] - else: - self.cov_matrix = np.eye(6) + self.cov_matrix = np.eye(6) self.intensity = intensity self.classical_radius = get_classical_radius(self.charge, self.mass) - self.charge_sign = self.charge / abs(self.charge) # For a uniform one-dimensional distribution over length L, the standard # deviation is L * sqrt(12). This quantity is used to calculate the line @@ -81,9 +97,9 @@ def __init__( def copy(self): return Envelope( - bunch=self.bunch, - cov_matrix=self.cov_matrix, - centroid=self.centroid, + sync_part=self.sync_part, + cov_matrix=self.cov_matrix.copy(), + centroid=self.centroid.copy(), intensity=self.intensity ) @@ -105,7 +121,7 @@ def mass(self) -> float: @property def charge(self) -> float: - return self.bunch.charge() + return self.sync_part.charge() @property def momentum(self) -> float: @@ -130,6 +146,18 @@ def sample(self, size: int, dist: str = "kv") -> np.ndarray: particles = particles + self.centroid return particles + def to_bunch(self, size: int, dist: str = "gauss") -> Bunch: + bunch = Bunch() + bunch.mass(self.mass) + bunch.charge(self.charge) + bunch.getSyncParticle().kinEnergy(self.kin_energy) + bunch.macroSize(self.intensity / size) + + particles = self.sample(size=size, dist=dist) + for i in range(particles.shape[0]): + bunch.addParticle(*particles[i]) + return bunch + def sc_matrix_2d(self, length: float) -> np.ndarray: centroid = self.centroid cov_matrix = self.cov_matrix diff --git a/py/orbit/envelope/matrix.py b/py/orbit/envelope/matrix.py deleted file mode 100644 index ee0cebdb..00000000 --- a/py/orbit/envelope/matrix.py +++ /dev/null @@ -1,443 +0,0 @@ -import math - -import numpy as np - -from orbit.core.bunch import Bunch -from orbit.core.bunch import SyncParticle -from orbit.lattice import AccNode -from orbit.teapot import ApertureTEAPOT -from orbit.teapot import DriftTEAPOT -from orbit.teapot import BendTEAPOT -from orbit.teapot import KickTEAPOT -from orbit.teapot import MonitorTEAPOT -from orbit.teapot import MultipoleTEAPOT -from orbit.teapot import NodeTEAPOT -from orbit.teapot import QuadTEAPOT -from orbit.teapot import SolenoidTEAPOT -from orbit.teapot import FringeFieldTEAPOT -from orbit.teapot import BunchWrapTEAPOT -from orbit.teapot import TiltTEAPOT -from orbit.teapot import ContinuousLinearFocusingTEAPOT -from orbit.teapot import TurnCounterTEAPOT -from orbit.py_linac.lattice import MarkerLinacNode as MarkerLINAC -from orbit.py_linac.lattice import Drift as DriftLINAC -from orbit.py_linac.lattice import Quad as QuadLINAC -from orbit.py_linac.lattice import Bend as BendLINAC -from orbit.py_linac.lattice import DCorrectorH as DCorrectorHLINAC -from orbit.py_linac.lattice import DCorrectorV as DCorrectorVLINAC -from orbit.py_linac.lattice import Solenoid as SolenoidLINAC -from orbit.py_linac.lattice import TiltElement as TiltLINAC -from orbit.py_linac.lattice import FringeField as FringeFieldLINAC -from orbit.py_linac.lattice import BaseRF_Gap as BaseRF_Gap -from orbit.py_linac.lattice import LinacApertureNode as ApertureLINAC -from orbit.utils.consts import speed_of_light - -from .envelope import Envelope -from .utils import get_dp_p_coeff - - -IGNORE_NODE_TYPES = [ - NodeTEAPOT, - MonitorTEAPOT, - FringeFieldTEAPOT, - ApertureTEAPOT, - BunchWrapTEAPOT, - TurnCounterTEAPOT, - MarkerLINAC, - FringeFieldLINAC, -] - - -def get_matrix_tilt(angle: float) -> np.ndarray: - cos_phi = math.cos(angle) - sin_phi = math.sin(angle) - - M = np.identity(7) - M[0, 0] = M[1, 1] = +cos_phi - M[0, 2] = M[1, 3] = -sin_phi - M[2, 0] = M[3, 1] = +sin_phi - M[2, 2] = M[3, 3] = +cos_phi - return M - - -def get_matrix_kick(kx: float = 0.0, ky: float = 0.0, kE: float = 0.0) -> np.ndarray: - M = np.identity(7) - M[1, -1] = kx - M[3, -1] = ky - M[5, -1] = kE - return M - - -def get_matrix_drift(envelope: Envelope, length: float) -> np.ndarray: - sync_part = envelope.sync_part - - M = np.identity(7) - M[0, 1] = length - M[2, 3] = length - M[4, 5] = length / (sync_part.gamma() ** 2) - M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_quad(envelope: Envelope, length: float, kq: float) -> np.ndarray: - if abs(kq) == 0: - return get_matrix_drift(envelope=envelope, length=length) - - sync_part = envelope.sync_part - - sqrt_abs_kq = math.sqrt(abs(kq)) - - M = np.identity(7) - if kq > 0: - cx = np.cos(sqrt_abs_kq * length) - sx = np.sin(sqrt_abs_kq * length) - cy = np.cosh(sqrt_abs_kq * length) - sy = np.sinh(sqrt_abs_kq * length) - M[0, 0] = cx - M[0, 1] = +sx / sqrt_abs_kq - M[1, 0] = -sx * sqrt_abs_kq - M[1, 1] = cx - M[2, 2] = cy - M[2, 3] = sy / sqrt_abs_kq - M[3, 2] = sy * sqrt_abs_kq - M[3, 3] = cy - elif kq < 0: - cx = np.cosh(sqrt_abs_kq * length) - sx = np.sinh(sqrt_abs_kq * length) - cy = np.cos(sqrt_abs_kq * length) - sy = np.sin(sqrt_abs_kq * length) - M[0, 0] = cx - M[0, 1] = sx / sqrt_abs_kq - M[1, 0] = sx * sqrt_abs_kq - M[1, 1] = cx - M[2, 2] = cy - M[2, 3] = +sy / sqrt_abs_kq - M[3, 2] = -sy * sqrt_abs_kq - M[3, 3] = cy - - M[4, 5] = length / (sync_part.gamma()**2) - M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_bend(envelope: Envelope, length: float, theta: float) -> np.ndarray: - sync_part = envelope.sync_part - - rho = length / theta - cx = math.cos(theta) - sx = math.sin(theta) - - M = np.identity(7) - M[0, 0] = cx - M[0, 1] = rho * sx - M[0, 5] = rho * (1.0 - cx) - M[1, 0] = -sx / rho - M[1, 1] = cx - M[1, 5] = sx - M[2, 3] = length - M[4, 0] = -sx - M[4, 1] = -rho * (1.0 - cx) - M[4, 5] = -(sync_part.beta() ** 2) * length + rho * sx - M[:5, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_solenoid(envelope: Envelope, length: float, B: float) -> np.ndarray: - if B == 0: - return get_matrix_drift(envelope=envelope, length=length) - - sync_part = envelope.sync_part - - phase = B * length - - V = np.identity(7) - V[:4, :4] = 0.0 - V[0, 1] = -1.0 / B - V[0, 2] = 0.5 - V[1, 0] = 0.5 * B - V[1, 3] = 1.0 - V[2, 1] = 1.0 / B - V[2, 2] = 0.5 - V[3, 0] = -0.5 * B - V[3, 3] = 1.0 - - M = np.identity(7) - M[0, 0] = +1.0 - M[1, 1] = -1.0 - M[2, 2] = math.cos(phase) - M[2, 3] = math.sin(phase) / B - M[3, 2] = math.sin(phase) * B * -1.0 - M[3, 3] = math.cos(phase) - M[4, 5] = length / (sync_part.gamma()**2) - - M = np.linalg.inv(V) @ M @ V - M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_cf(envelope: Envelope, length: float, kq: float) -> np.ndarray: - if kq == 0: - return get_matrix_drift(envelope=envelope, length=length) - - sync_part = envelope.sync_part - - sqrt_abs_kq = math.sqrt(abs(kq)) - - cx = math.cos(sqrt_abs_kq * length) - sx = math.sin(sqrt_abs_kq * length) - - M = np.identity(7) - M[0, 0] = M[2, 2] = cx - M[0, 1] = M[2, 3] = +sx / sqrt_abs_kq - M[1, 0] = M[3, 2] = -sx * sqrt_abs_kq - M[1, 1] = M[3, 3] = cx - M[4, 5] = length / (sync_part.gamma()**2) - M[4, 5] *= get_dp_p_coeff(sync_part) - - sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) - return M - - -def get_matrix_rf_gap(envelope: Envelope, frequency: float, E0TL: float, phase: float) -> np.ndarray: - sync_part = envelope.sync_part - - gamma = sync_part.gamma() - beta = sync_part.beta() - mass = sync_part.mass() - charge = envelope.charge - - kin_energy_in = sync_part.kinEnergy() - charge_E0TL_sin = charge * E0TL * math.sin(phase) - kin_energy_delta = charge * E0TL * math.cos(phase) - - # Calculate parameters in the center of the gap. - sync_part.momentum(sync_part.energyToMomentum(kin_energy_in + kin_energy_delta / 2.0)) - gamma_gap = sync_part.gamma() - beta_gap = sync_part.beta() - - # Move to the end of the gap. - kin_energy_out = kin_energy_in + kin_energy_delta - sync_part.momentum(sync_part.energyToMomentum(kin_energy_out)) - - # The base RF gap is simple - no phase correction. - gamma_out = sync_part.gamma() - beta_out = sync_part.beta() - prime_coeff = (beta * gamma) / (beta_out * gamma_out) - - # Wave momentum - k = 2.0 * math.pi * frequency / speed_of_light - phase_time_coeff = k / beta - - # Transverse focusing coefficient - kappa = -charge * E0TL * k / (2.0 * mass * beta_gap**2 * beta_out * gamma_gap**2 * gamma_out) - d_rp = kappa * math.sin(phase) - - M = np.eye(7) - M[5, 4] = charge_E0TL_sin * phase_time_coeff - M[4, 4] = beta_out / beta - M[1, 1] = prime_coeff - M[3, 3] = prime_coeff - M[1, 0] = d_rp - M[3, 2] = d_rp - return M - - -def get_matrix(node: AccNode, envelope: Envelope, part_index: int = -1) -> np.ndarray | None: - """Calculate transfer matrix and update synchronous particle. - - This function maps various accelerator nodes to 7 x 7 transfer matrices - for envelope tracking. For non-accelerating, finite-length nodes, the - synchronous particle time is updated as in a drift. Accelerating nodes - such as RF gaps will update the synchronous particle energy. - - Args: - node: The accelerator node. - envelope: The beam envelope. - part_index: Index of the part within the node. An index of -1 returns - the transfer matrix for the entire node. - Returns: - 7 x 7 transfer matrix or None. If None, the node can be ignored during - envelope tracking. - """ - - node_type = type(node) - if node_type in IGNORE_NODE_TYPES: - return None - - length = node.getLength(part_index) - nparts = node.getnParts() - - if node_type is DriftTEAPOT: - if length <= 0: - return None - return get_matrix_drift(envelope=envelope, length=length) - - elif node_type is SolenoidTEAPOT: - if length <= 0: - return None - - B = node.getParam("B") - if node.waveform: - B *= node.waveform.getStrength() - B *= envelope.charge_sign - - return get_matrix_solenoid(envelope=envelope, length=length, B=B) - - elif node_type is MultipoleTEAPOT: - if length <= 0: - return None - - if np.all(np.abs(node.getParam("kls")) == 0): - return get_matrix_drift(envelope=envelope, length=length) - - elif node_type is QuadTEAPOT: - if length <= 0: - return None - - kq = node.getParam("kq") - if node.waveform: - kq *= node.waveform.getStrength() - kq *= envelope.charge_sign - - return get_matrix_quad(envelope=envelope, length=length, kq=kq) - - elif node_type is BendTEAPOT: - if length <= 0: - return None - - theta = node.getParam("theta") / (nparts - 1) - if part_index == 0 or part_index == nparts - 1: - theta *= 0.5 - theta *= envelope.charge_sign - - return get_matrix_bend(envelope=envelope, length=length, theta=theta) - - elif node_type is KickTEAPOT: - scale = 1.0 - if node.waveform is not None: - scale = node.waveform.getStrength() - - scale /= (nparts - 1) - kx = scale * node.getParam("kx") - ky = scale * node.getParam("ky") - kE = node.getParam("dE") - - if abs(kx) > 0 or abs(ky) > 0 or abs(kE) > 0: - return np.matmul( - get_matrix_kick(kx=kx, ky=ky, kE=kE), - get_matrix_drift(envelope=envelope, length=length), - ) - else: - return get_matrix_drift(envelope=envelope, length=length) - - elif node_type is TiltTEAPOT: - angle = node.getTiltAngle() - if angle == 0: - return None - return get_matrix_tilt(angle) - - elif node_type is ContinuousLinearFocusingTEAPOT: - if length <= 0: - return None - - kq = node.getParam("kq") - kq *= envelope.charge_sign - if node.waveform: - kq *= node.waveform.getStrength() - - return get_matrix_cf(envelope=envelope, length=length, kq=kq) - - elif node_type is DriftLINAC: - if length <= 0: - return None - return get_matrix_drift(envelope=envelope, length=length) - - elif node_type is QuadLINAC: - if length <= 0: - return None - - brho = 3.335640952 * envelope.momentum / envelope.charge - kq = node.getParam("dB/dr") / brho - return get_matrix_quad(envelope=envelope, length=length, kq=kq) - - elif node_type is BendLINAC: - if length <= 0: - return None - - theta = node.getParam("theta") / (nparts - 1) - if part_index == 0 or part_index == nparts - 1: - theta *= 0.5 - theta *= envelope.charge_sign - - return get_matrix_bend(envelope=envelope, length=length, theta=theta) - - elif node_type is DCorrectorHLINAC: - length = node.getParam("effLength") / nparts - field = node.getParam("B") - delta_xp = -field * envelope.charge * length * 0.299792 / envelope.momentum - if delta_xp == 0: - return None - return get_matrix_kick(kx=delta_xp, ky=0.0, kE=0.0) - - elif node_type is DCorrectorVLINAC: - length = node.getParam("effLength") / nparts - field = node.getParam("B") - delta_yp = -field * envelope.charge * length * 0.299792 / envelope.momentum - if delta_yp == 0: - return None - return get_matrix_kick(kx=0.0, ky=delta_yp, kE=0.0) - - elif node_type is SolenoidLINAC: - if length <= 0: - return None - B = node.getParam("B") * envelope.charge_sign - return get_matrix_solenoid(envelope=envelope, length=length, B=B) - - elif node_type is TiltLINAC: - angle = node.getTiltAngle() - if angle == 0: - return None - return get_matrix_tilt(angle=angle) - - elif node_type is BaseRF_Gap: - E0TL = node.getParam("E0TL") - mode_phase = node.getParam("mode") * math.pi - - cavity = node.getRF_Cavity() - frequency = cavity.getFrequency() - phase = cavity.getPhase() + mode_phase - amplitude = cavity.getAmp() - - sync_part = envelope.sync_part - arrival_time = sync_part.time() - arrival_time_design = cavity.getDesignArrivalTime() - - if node.isFirstRFGap(): - if cavity.isDesignSetUp(): - phase = math.fmod(frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase, 2.0 * math.pi) - else: - orbitFinalize("Run `trackDesign` first to initialize cavity phases.") - else: - phase = math.fmod(frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase,2.0 * math.pi) - - node.setGapPhase(phase) - - if amplitude == 0.0: - return None - - return get_matrix_rf_gap( - envelope=envelope, - frequency=frequency, - E0TL=(E0TL * amplitude), - phase=phase, - ) - - raise NotImplementedError(str(node)) \ No newline at end of file diff --git a/py/orbit/envelope/meson.build b/py/orbit/envelope/meson.build index 90eab844..1e9fe35c 100644 --- a/py/orbit/envelope/meson.build +++ b/py/orbit/envelope/meson.build @@ -1,8 +1,6 @@ py_sources = files([ '__init__.py', 'envelope.py', - 'matrix.py', - 'track.py', 'utils.py' ]) diff --git a/py/orbit/envelope/track.py b/py/orbit/envelope/track.py deleted file mode 100644 index 16b62b9d..00000000 --- a/py/orbit/envelope/track.py +++ /dev/null @@ -1,290 +0,0 @@ -import numpy as np -import warnings - -from orbit.core.bunch import Bunch -from orbit.core.bunch import SyncParticle - -from orbit.lattice import AccNode -from orbit.lattice import AccLattice -from orbit.teapot import BendTEAPOT -from orbit.py_linac.lattice import Bend as BendLINAC - -from .matrix import get_matrix -from .envelope import Envelope - - -ENTRANCE = AccNode.ENTRANCE -BODY = AccNode.BODY -EXIT = AccNode.EXIT - -BEFORE = AccNode.BEFORE -AFTER = AccNode.AFTER - - -class EnvelopeTracker: - def __init__(self, lattice: AccLattice, sc: str | None = None) -> None: - """Constructor. - - Args: - lattice: The accelerator lattice. - sc: Envelope space charge model {"2d", "3d", None}. - """ - self.lattice = lattice - self.sc = sc - - # For pre-computing elements - self.elements = [] - self.one_turn_matrix = None - - for node in self.lattice.getNodes(): - if type(node) in (BendTEAPOT, BendLINAC): - if node.getParam("ea1") != 0.0 or node.getParam("ea2") != 0.0: - message = f"Found bend ea1 or ea2 != 0.0 ({node.getName()}.)" - message += " Nonzero edge angles are not yet supported in envelope tracking." - message += " Setting ea1 and ea2 to 0.0." - warnings.warn(message) - - node.setParam("ea1", 0.0) - node.setParam("ea2", 0.0) - - def track(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> None: - """Track envelope through lattice. - - This is not recursive, so grandchild nodes are not tracked. - """ - nodes = self.lattice.getNodes() - nodes = nodes[index_start : index_stop] - - for node_index, node in enumerate(nodes): - for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - envelope.transform(matrix) - - for part_index in range(node.getnParts()): - for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - envelope.transform(matrix) - - matrix_sc = None - if self.sc: - length = node.getLength(part_index) - if length > 0: - if self.sc == "2d": - matrix_sc = envelope.sc_matrix_2d(length) - elif self.sc == "3d": - matrix_sc = envelope.sc_matrix_3d(length) - else: - raise ValueError - - matrix = get_matrix(node, envelope=envelope, part_index=part_index) - if matrix is not None: - if matrix_sc is not None: - matrix = matrix @ matrix_sc - envelope.transform(matrix) - - for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - envelope.transform(matrix) - - for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - envelope.transform(matrix) - - def track_history(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> dict[str, list]: - """Track and return envelope parameters vs. position in lattice.""" - history_keys = [ - "s", - "kin_energy", - "gamma", - "beta", - "mean", - "cov", - "rms_x", - "rms_y", - "rms_z", - "eps_x", - "eps_y", - ] - history = {key: [] for key in history_keys} - - def observe(envelope: Envelope) -> None: - parameters = {} - parameters["gamma"] = envelope.gamma - parameters["beta"] = envelope.beta - parameters["kin_energy"] = envelope.kin_energy - parameters["mean"] = envelope.centroid.copy() - parameters["cov"] = envelope.cov_matrix.copy() - parameters["rms_x"] = np.sqrt(parameters["cov"][0, 0]) - parameters["rms_y"] = np.sqrt(parameters["cov"][2, 2]) - parameters["rms_z"] = np.sqrt(parameters["cov"][4, 4]) - return parameters - - def update_history(envelope: Envelope, position: float) -> None: - history["s"].append(position) - parameters = observe(envelope) - for key in parameters: - history[key].append(parameters[key]) - - path_length = 0.0 - update_history(envelope, path_length) - - nodes = self.lattice.getNodes() - nodes = nodes[index_start : index_stop] - - for node_index, node in enumerate(nodes): - for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - envelope.transform(matrix) - - for part_index in range(node.getnParts()): - for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - envelope.transform(matrix) - - matrix_sc = None - if self.sc: - length = node.getLength(part_index) - if length > 0: - if self.sc == "2d": - matrix_sc = envelope.sc_matrix_2d(length) - elif self.sc == "3d": - matrix_sc = envelope.sc_matrix_3d(length) - else: - raise ValueError - - matrix = get_matrix(node, envelope=envelope, part_index=part_index) - if matrix is not None: - if matrix_sc is not None: - matrix = matrix @ matrix_sc - envelope.transform(matrix) - - path_length += node.getLength(part_index) - update_history(envelope, path_length) - - for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - envelope.transform(matrix) - - for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - envelope.transform(matrix) - - return history - - def precompute_matrices(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> None: - """Pre-compute transfer matrices for each node. - - For each node, return tuple (node, matrix). Mark space charge kicks as ("sc", length). - """ - nodes = self.lattice.getNodes() - nodes = nodes[index_start : index_stop] - - self.elements = [] - for node_index, node in enumerate(nodes): - for child_node in node.getChildNodes(ENTRANCE): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - self.elements.append((child_node, matrix)) - - for part_index in range(node.getnParts()): - for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - self.elements.append((child_node, matrix)) - - if self.sc: - length = node.getLength(part_index) - if length > 0: - self.elements.append(("sc", length)) - - matrix = get_matrix(node, envelope=envelope, part_index=part_index) - if matrix is not None: - self.elements.append((node, matrix)) - - for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - self.elements.append((node, matrix)) - - for child_node in node.getChildNodes(EXIT): - matrix = get_matrix(child_node, envelope=envelope) - if matrix is not None: - self.elements.append((node, matrix)) - - def track_ring(self, envelope: Envelope) -> None: - """Track using pre-computed transfer matrices. - - The method assumes that all nodes are static and that there is no - change in the synchronous particle energy. In this case the matrices - can be computed once and reused on each turn. If there is no space charge, - we track using the one-turn matrix. - """ - - # Pre-compute transfer matrices on the first turn. - if not self.elements: - self.precompute_matrices(envelope) - self.one_turn_matrix = None - - # If there is no space charge, apply the one-turn transfer matrix. - if not self.sc: - if self.one_turn_matrix is None: - self.one_turn_matrix = np.identity(7) - for (node, matrix) in self.elements: - self.one_turn_matrix = matrix @ self.one_turn_matrix - return envelope.transform(self.one_turn_matrix) - - # If there is space charge, apply the matrices one-by-one. - for element in self.elements: - if element[0] == "sc": - length = element[1] - if self.sc == "2d": - envelope.transform(envelope.sc_matrix_2d(length)) - elif self.sc == "3d": - envelope.transform(envelope.sc_matrix_3d(length)) - else: - raise ValueError - else: - node, matrix = element - envelope.transform(matrix) - - def get_transfer_matrix(self, envelope: Envelope, index_start: int = 0, index_stop: int = None) -> np.ndarray: - """Return total transfer matrix (including linear space charge).""" - self.precompute_matrices(envelope, index_start, index_stop) - - if index_stop is None: - index_stop = len(self.elements) - - elements = self.elements[index_start : index_stop] - - if not self.sc: - total_matrix = np.identity(7) - for (node, matrix) in self.elements: - envelope.transform(matrix) - total_matrix = matrix @ total_matrix - return total_matrix - - total_matrix = np.identity(7) - for element in self.elements: - if element[0] == "sc": - length = element[1] - if self.sc == "2d": - matrix = envelope.sc_matrix_2d(length) - elif self.sc == "3d": - matrix = envelope.sc_matrix_3d(length) - else: - raise ValueError - envelope.transform(matrix) - total_matrix = matrix @ total_matrix - else: - node, matrix = element - envelope.transform(matrix) - total_matrix = matrix @ total_matrix - return total_matrix \ No newline at end of file diff --git a/py/orbit/envelope/utils.py b/py/orbit/envelope/utils.py index 71499e31..9102794e 100644 --- a/py/orbit/envelope/utils.py +++ b/py/orbit/envelope/utils.py @@ -2,7 +2,6 @@ import numpy as np from scipy.constants import epsilon_0 -from orbit.core.bunch import SyncParticle from orbit.utils.consts import charge_electron @@ -12,47 +11,6 @@ def get_classical_radius(charge: float, mass: float) -> float: return q**2 / (4.0 * math.pi * epsilon_0 * rest_energy) -def get_dp_p_coeff(sync_part: SyncParticle) -> float: - # dE/E = (beta^2) * dp/p - # dE = (beta^2 * E) * dp/p - # dE = (beta^2 * gamma * m * c^2) * dp/p - beta = sync_part.beta() - gamma = sync_part.gamma() - rest_energy = sync_part.mass() # GeV - return 1.0 / (beta**2 * gamma * rest_energy) - - -def get_zp_coeff(sync_part: SyncParticle) -> float: - # dE/E = (beta^2) * dp/p = (beta^2) * (gamma^2) z' - # dE = (beta^2 * gamma^2 * E) * z' - # dE = (beta^2 * gamma^3 * m * c^2) * z' - beta = sync_part.beta() - gamma = sync_part.gamma() - rest_energy = sync_part.mass() - return 1.0 / (beta**2 * gamma**3 * rest_energy) - - -def convert_matrix_dp_p_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: - # v = [x, x', y, y', z, dp/p] - # w = [x, x', y, y', z, dE] - # v = A w - # v -> M v - # w -> A M A^-1 - dp_p_coeff = get_dp_p_coeff(sync_part) - matrix[:5, 5] *= dp_p_coeff - matrix[5, :5] /= dp_p_coeff - matrix[5, 6] /= dp_p_coeff # driving term - return matrix - - -def convert_matrix_zp_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: - zp_coeff = get_zp_coeff(sync_part) - matrix[:5, 5] *= zp_coeff - matrix[5, :5] /= zp_coeff - matrix[5, 6] /= zp_coeff # driving term - return matrix - - def gen_dist_gauss(size: int, cov_matrix: np.ndarray) -> np.ndarray: return np.random.multivariate_normal( mean=np.zeros(cov_matrix.shape[0]), diff --git a/py/orbit/lattice/AccLattice.py b/py/orbit/lattice/AccLattice.py index 84a466f8..b16da010 100644 --- a/py/orbit/lattice/AccLattice.py +++ b/py/orbit/lattice/AccLattice.py @@ -1,12 +1,19 @@ -import sys +from __future__ import annotations + import os +from typing import TYPE_CHECKING + +import numpy as np from ..utils import orbitFinalize from ..utils import NamedObject from ..utils import TypedObject -from ..lattice import AccActionsContainer -from ..lattice import AccNode +from .AccActionsContainer import AccActionsContainer +from .AccNode import AccNode + +if TYPE_CHECKING: + from orbit.envelope.envelope import Envelope class AccLattice(NamedObject, TypedObject): @@ -31,6 +38,9 @@ def __init__(self, name="no name"): self.__isInitialized = False self.__children = [] self.__childPositions = {} + self.__envelopeElements = [] + self.__envelopeOneTurnMatrix = None + self.__envelopeSpaceCharge = None def initialize(self): """ @@ -261,3 +271,274 @@ def trackActions(self, actionsContainer, paramsDict={}, index_start=-1, index_st paramsDict["node"] = node paramsDict["parentNode"] = self node.trackActions(actionsContainer, paramsDict) + + def _getNodesInRange(self, index_start: int = 0, index_stop: int = None): + if index_stop is None: + index_stop = len(self.__children) - 1 + return self.__children[index_start : index_stop + 1] + + def _prepareEnvelopeTracking(self) -> None: + for node in self.__children: + node_type = type(node) + is_teapot_bend = node_type.__name__ == "BendTEAPOT" and node_type.__module__ == "orbit.teapot.teapot" + is_linac_bend = node_type.__name__ == "Bend" and node_type.__module__ == "orbit.py_linac.lattice.LinacAccNodes" + if is_teapot_bend or is_linac_bend: + if node.getParam("ea1") != 0.0 or node.getParam("ea2") != 0.0: + message = f"Found bend ea1 or ea2 != 0.0 ({node.getName()}.)" + message += " Nonzero edge angles are not yet supported in envelope tracking." + message += " Please set them to zero:" + message += " `node.setParam('ea1', 0.0)`" + message += " `node.setParam('ea2', 0.0)`" + raise RuntimeError(message) + + def _getEnvelopeSpaceChargeMatrix(self, envelope: Envelope, length: float, sc: str | None) -> np.ndarray | None: + if not sc or length <= 0: + return None + if sc == "2d": + return envelope.sc_matrix_2d(length) + if sc == "3d": + return envelope.sc_matrix_3d(length) + raise ValueError(f"Invalid envelope space charge option `{sc}`") + + def setEnvelopeSpaceCharge(self, sc: str | None) -> None: + self.__envelopeSpaceCharge = sc + + def trackEnvelope( + self, + envelope: Envelope, + index_start: int = 0, + index_stop: int = None, + sc: str | None = None, + history: bool = False, + ) -> None | dict[str, list]: + """ + Track envelope through lattice. + """ + if history: + return self.trackEnvelopeHistory( + envelope, + index_start=index_start, + index_stop=index_stop, + sc=sc + ) + + self._prepareEnvelopeTracking() + self.setEnvelopeSpaceCharge(sc) + sync_part = envelope.sync_part + + for node in self._getNodesInRange(index_start, index_stop): + for child_node in node.getChildNodes(AccNode.ENTRANCE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.BEFORE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + matrix_sc = self._getEnvelopeSpaceChargeMatrix(envelope, node.getLength(part_index), sc) + matrix = node.getMatrix(sync_part, part_index=part_index) + if matrix is not None: + if matrix_sc is not None: + matrix = matrix @ matrix_sc + envelope.transform(matrix) + + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.AFTER): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + for child_node in node.getChildNodes(AccNode.EXIT): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + def trackEnvelopeHistory( + self, + envelope: Envelope, + index_start: int = 0, + index_stop: int = None, + sc: str | None = None, + ) -> dict[str, list]: + """ + Track envelope and return parameters vs. position in lattice. + """ + self._prepareEnvelopeTracking() + self.setEnvelopeSpaceCharge(sc) + sync_part = envelope.sync_part + + history_keys = [ + "s", + "kin_energy", + "gamma", + "beta", + "mean", + "cov", + "rms_x", + "rms_y", + "rms_z", + "eps_x", + "eps_y", + ] + history = {key: [] for key in history_keys} + + def observe(envelope: Envelope) -> dict: + parameters = {} + parameters["gamma"] = envelope.gamma + parameters["beta"] = envelope.beta + parameters["kin_energy"] = envelope.kin_energy + parameters["mean"] = envelope.centroid.copy() + parameters["cov"] = envelope.cov_matrix.copy() + parameters["rms_x"] = np.sqrt(parameters["cov"][0, 0]) + parameters["rms_y"] = np.sqrt(parameters["cov"][2, 2]) + parameters["rms_z"] = np.sqrt(parameters["cov"][4, 4]) + return parameters + + def update_history(envelope: Envelope, position: float) -> None: + history["s"].append(position) + parameters = observe(envelope) + for key in parameters: + history[key].append(parameters[key]) + + path_length = 0.0 + update_history(envelope, path_length) + + for node in self._getNodesInRange(index_start, index_stop): + for child_node in node.getChildNodes(AccNode.ENTRANCE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.BEFORE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + matrix_sc = self._getEnvelopeSpaceChargeMatrix(envelope, node.getLength(part_index), sc) + matrix = node.getMatrix(sync_part, part_index=part_index) + if matrix is not None: + if matrix_sc is not None: + matrix = matrix @ matrix_sc + envelope.transform(matrix) + + path_length += node.getLength(part_index) + update_history(envelope, path_length) + + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.AFTER): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + + for child_node in node.getChildNodes(AccNode.EXIT): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + envelope.transform(matrix) + return history + + def precomputeEnvelopeMatrices( + self, + envelope: Envelope, + index_start: int = 0, + index_stop: int = None, + sc: str | None = None, + ) -> list: + """ + Pre-compute transfer matrices for each node. + + For each node, store tuple (node, matrix). Space charge kicks are + stored as ("sc", length). + """ + self._prepareEnvelopeTracking() + sync_part = envelope.sync_part + + self.__envelopeElements = [] + self.__envelopeOneTurnMatrix = None + self.__envelopeSpaceCharge = sc + + for node in self._getNodesInRange(index_start, index_stop): + for child_node in node.getChildNodes(AccNode.ENTRANCE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + self.__envelopeElements.append((child_node, matrix)) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.BEFORE): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + self.__envelopeElements.append((child_node, matrix)) + + if sc: + length = node.getLength(part_index) + if length > 0: + self.__envelopeElements.append(("sc", length)) + + matrix = node.getMatrix(sync_part, part_index=part_index) + if matrix is not None: + self.__envelopeElements.append((node, matrix)) + + for child_node in node.getChildNodes(AccNode.BODY, part_index, place_in_part=AccNode.AFTER): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + self.__envelopeElements.append((child_node, matrix)) + + for child_node in node.getChildNodes(AccNode.EXIT): + matrix = child_node.getMatrix(sync_part) + if matrix is not None: + self.__envelopeElements.append((child_node, matrix)) + + return self.__envelopeElements + + def trackEnvelopeRing(self, envelope: Envelope, sc: str | None = None) -> None: + """ + Track using pre-computed transfer matrices. + + The method assumes that all nodes are static and that there is no + change in the synchronous particle energy. In this case the matrices + can be computed once and reused on each turn. If there is no space charge, + we track using the one-turn matrix. + """ + if not self.__envelopeElements or self.__envelopeSpaceCharge != sc: + self.precomputeEnvelopeMatrices(envelope, sc=sc) + + if not sc: + if self.__envelopeOneTurnMatrix is None: + self.__envelopeOneTurnMatrix = np.identity(7) + for node, matrix in self.__envelopeElements: + self.__envelopeOneTurnMatrix = matrix @ self.__envelopeOneTurnMatrix + envelope.transform(self.__envelopeOneTurnMatrix) + return + + for element in self.__envelopeElements: + if element[0] == "sc": + length = element[1] + matrix = self._getEnvelopeSpaceChargeMatrix(envelope, length, sc) + envelope.transform(matrix) + else: + node, matrix = element + envelope.transform(matrix) + + def getEnvelopeTransferMatrix( + self, + envelope: Envelope, + index_start: int = 0, + index_stop: int = None, + sc: str | None = None, + ) -> np.ndarray: + """ + Return total transfer matrix, including linear space charge when requested. + """ + elements = self.precomputeEnvelopeMatrices(envelope, index_start, index_stop, sc=sc) + + total_matrix = np.identity(7) + for element in elements: + if element[0] == "sc": + length = element[1] + matrix = self._getEnvelopeSpaceChargeMatrix(envelope, length, sc) + else: + node, matrix = element + envelope.transform(matrix) + total_matrix = matrix @ total_matrix + return total_matrix diff --git a/py/orbit/lattice/AccNode.py b/py/orbit/lattice/AccNode.py index fdc145dc..1926ff9c 100644 --- a/py/orbit/lattice/AccNode.py +++ b/py/orbit/lattice/AccNode.py @@ -1,13 +1,20 @@ +from __future__ import annotations + import sys import os import math +from typing import TYPE_CHECKING from ..utils import orbitFinalize from ..utils import NamedObject from ..utils import TypedObject from ..utils import ParamsDictObject -from ..lattice import AccActionsContainer +from .AccActionsContainer import AccActionsContainer + +if TYPE_CHECKING: + import numpy as np + from orbit.core.bunch import SyncParticle class AccNode(NamedObject, TypedObject, ParamsDictObject): @@ -145,6 +152,12 @@ def initialize(self): """ pass + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + """ + Return the transfer matrix for this node and update the synchronous particle. + """ + raise NotImplementedError(str(self)) + def getNumberOfChildren(self): """ Returns the total number of direct children diff --git a/py/orbit/py_linac/lattice/LinacAccNodes.py b/py/orbit/py_linac/lattice/LinacAccNodes.py index ba72fef5..1d4e2a7e 100755 --- a/py/orbit/py_linac/lattice/LinacAccNodes.py +++ b/py/orbit/py_linac/lattice/LinacAccNodes.py @@ -8,16 +8,26 @@ import os import math +import numpy as np # import the finalization function from orbit.utils import orbitFinalize # import general accelerator elements and lattice from orbit.lattice import AccNode, AccActionsContainer, AccNodeBunchTracker +from orbit.core.bunch import SyncParticle # import teapot base functions from wrapper around C++ functions from orbit.teapot_base import TPB +from orbit.utils.matrix import get_matrix_bend +from orbit.utils.matrix import get_matrix_cf +from orbit.utils.matrix import get_matrix_drift +from orbit.utils.matrix import get_matrix_kick +from orbit.utils.matrix import get_matrix_quad +from orbit.utils.matrix import get_matrix_solenoid +from orbit.utils.matrix import get_matrix_tilt + # Import the linac specific tracking from linac_tracking. This module has # the following functions duplicated the original TEAPOT functions # drift - linac drift tracking @@ -137,6 +147,9 @@ def __init__(self, name="none"): BaseLinacNode.__init__(self, name) self.setType("markerLinacNode") + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class LinacNode(BaseLinacNode): """ @@ -304,6 +317,12 @@ def track(self, paramsDict): bunch = paramsDict["bunch"] self.tracking_module.drift(bunch, length) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + return get_matrix_drift(sync_part, length) + class Quad(LinacMagnetNode): """ @@ -513,6 +532,15 @@ def getTotalField(self, z): G = self.getParam("dB/dr") return G + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + brho = 3.335640952 * sync_part.momentum() / sync_part.charge() + kq = self.getParam("dB/dr") / brho + return get_matrix_quad(sync_part, length=length, kq=kq) + class Bend(LinacMagnetNode): """ @@ -705,6 +733,18 @@ def track(self, paramsDict): TPB.bend1(bunch, length, theta / 2.0) return + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + nparts = self.getnParts() + theta = self.getParam("theta") / (nparts - 1) + if part_index == 0 or part_index == nparts - 1: + theta *= 0.5 + + return get_matrix_bend(sync_part, length=length, theta=theta) + class DCorrectorH(LinacMagnetNode): """ @@ -750,6 +790,14 @@ def track(self, paramsDict): kick = -field * charge * length * 0.299792 / momentum self.tracking_module.kick(bunch, kick, 0.0, 0.0) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getParam("effLength") / self.getnParts() + field = self.getParam("B") + delta_xp = -field * sync_part.charge() * length * 0.299792 / sync_part.momentum() + if delta_xp == 0: + return None + return get_matrix_kick(kx=delta_xp, ky=0.0, kE=0.0) + class DCorrectorV(LinacMagnetNode): """ @@ -795,6 +843,14 @@ def track(self, paramsDict): kick = field * charge * length * 0.299792 / momentum self.tracking_module.kick(bunch, 0, kick, 0.0) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getParam("effLength") / self.getnParts() + field = self.getParam("B") + delta_yp = -field * sync_part.charge() * length * 0.299792 / sync_part.momentum() + if delta_yp == 0: + return None + return get_matrix_kick(kx=0.0, ky=delta_yp, kE=0.0) + class ThickKick(LinacMagnetNode): """ @@ -892,6 +948,13 @@ def track(self, paramsDict): useCharge = paramsDict["useCharge"] TPB.soln(bunch, length, B, useCharge) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + B = self.getParam("B") * np.sign(sync_part.charge()) + return get_matrix_solenoid(sync_part, length=length, B=B) + class AbstractRF_Gap(BaseLinacNode): """ This is an abstarct class for all RF Gap classes. @@ -1010,6 +1073,12 @@ def track(self, paramsDict): bunch = paramsDict["bunch"] TPB.rotatexy(bunch, self.__angle) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + angle = self.getTiltAngle() + if angle == 0: + return None + return get_matrix_tilt(angle=angle) + class FringeField(BaseLinacNode): """ @@ -1034,6 +1103,9 @@ def track(self, paramsDict): if self.__trackFunc != None and self.__usage == True: self.__trackFunc(self, paramsDict) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + def setFringeFieldFunction(self, trackFunction=None): """ Sets the fringe field function that will track the bunch through the fringe. diff --git a/py/orbit/py_linac/lattice/LinacRfGapNodes.py b/py/orbit/py_linac/lattice/LinacRfGapNodes.py index 30b1fbe8..6f482f39 100644 --- a/py/orbit/py_linac/lattice/LinacRfGapNodes.py +++ b/py/orbit/py_linac/lattice/LinacRfGapNodes.py @@ -5,6 +5,7 @@ import os import math +import numpy as np # ---- MPI module function and classes from orbit.core.orbit_mpi import mpi_comm, mpi_datatype, MPI_Comm_rank, MPI_Bcast @@ -26,6 +27,8 @@ # The abstract RF gap import from orbit.py_linac.lattice.LinacAccNodes import AbstractRF_Gap +from orbit.utils.matrix import get_matrix_rf_gap + # import teapot base functions from wrapper around C++ functions # Import the linac specific tracking from linac_tracking. This module has @@ -35,6 +38,7 @@ # quad2 - linac quad non-linear part of tracking from orbit.core.bunch import Bunch +from orbit.core.bunch import SyncParticle class BaseRF_Gap(AbstractRF_Gap): @@ -251,6 +255,44 @@ def track(self, paramsDict): # print "debug delta_time in deg=",frequency*(arrival_time - designArrivalTime)*380. # print "debug RF =",self.getName()," E0TL=",E0TL," phase=",(phase*180./math.pi - 180.)," eKin[MeV]=",bunch.getSyncParticle().kinEnergy()*1.0e+3 + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + E0TL = self.getParam("E0TL") + mode_phase = self.getParam("mode") * math.pi + + cavity = self.getRF_Cavity() + frequency = cavity.getFrequency() + phase = cavity.getPhase() + mode_phase + amplitude = cavity.getAmp() + + arrival_time = sync_part.time() + arrival_time_design = cavity.getDesignArrivalTime() + + if self.isFirstRFGap(): + if cavity.isDesignSetUp(): + phase = math.fmod( + frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase, + 2.0 * math.pi, + ) + else: + raise ValueError("Run `trackDesign` first to initialize cavity phases.") + else: + phase = math.fmod( + frequency * (arrival_time - arrival_time_design) * 2.0 * math.pi + phase, + 2.0 * math.pi, + ) + + self.setGapPhase(phase) + + if amplitude == 0.0: + return None + + return get_matrix_rf_gap( + sync_part=sync_part, + frequency=frequency, + E0TL=(E0TL * amplitude), + phase=phase, + ) + def trackDesign(self, paramsDict): """ The RF First Gap node setups the design time of passage diff --git a/py/orbit/teapot/teapot.py b/py/orbit/teapot/teapot.py index 63a09b27..cd216b8d 100644 --- a/py/orbit/teapot/teapot.py +++ b/py/orbit/teapot/teapot.py @@ -19,23 +19,32 @@ import sys import os import math +import numpy as np from typing import Any from typing import Callable from typing import Union -from ..lattice import AccLattice -from ..lattice import AccNode -from ..lattice import AccActionsContainer -from ..lattice import AccNodeBunchTracker -from ..teapot_base import TPB -from ..utils import orbitFinalize -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.lattice import AccLattice +from orbit.lattice import AccNode +from orbit.lattice import AccActionsContainer +from orbit.lattice import AccNodeBunchTracker +from orbit.teapot_base import TPB +from orbit.utils import orbitFinalize +from orbit.parsers.mad_parser import MAD_Parser +from orbit.parsers.mad_parser import MAD_LattElement +from orbit.parsers.madx_parser import MADX_Parser +from orbit.parsers.madx_parser import MADX_LattElement +from orbit.utils.matrix import get_matrix_bend +from orbit.utils.matrix import get_matrix_cf +from orbit.utils.matrix import get_matrix_drift +from orbit.utils.matrix import get_matrix_kick +from orbit.utils.matrix import get_matrix_quad +from orbit.utils.matrix import get_matrix_solenoid +from orbit.utils.matrix import get_matrix_tilt from orbit.core.aperture import Aperture from orbit.core.bunch import Bunch +from orbit.core.bunch import SyncParticle from orbit.core.bunch import BunchTwissAnalysis @@ -440,6 +449,9 @@ def track(self, paramsDict: dict) -> None: turn = bunch.bunchAttrInt("TurnNumber") bunch.bunchAttrInt("TurnNumber", turn + 1) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class NodeTEAPOT(BaseTEAPOT): def __init__(self, name: str = "no name") -> None: @@ -464,6 +476,11 @@ def __init__(self, name: str = "no name") -> None: self.addParam("tilt", self.__tiltNodeIN.getTiltAngle()) self.setType("node teapot") + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + if type(self) is NodeTEAPOT: + return None + return super().getMatrix(sync_part, part_index=part_index) + def setTiltAngle(self, angle: float = 0.0) -> None: """ Sets the tilt angle for the tilt operation. @@ -582,6 +599,12 @@ def track(self, paramsDict: dict) -> None: bunch = paramsDict["bunch"] TPB.drift(bunch, length) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + return get_matrix_drift(sync_part, length) + class ApertureTEAPOT(NodeTEAPOT): """ @@ -622,6 +645,9 @@ def track(self, paramsDict: dict) -> None: lostbunch = paramsDict["lostbunch"] self.aperture.checkBunch(bunch, lostbunch) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class MonitorTEAPOT(NodeTEAPOT): """ @@ -649,6 +675,9 @@ def track(self, paramsDict: dict) -> None: self.addParam("yAvg", self.twiss.getAverage(2)) self.addParam("ypAvg", self.twiss.getAverage(3)) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class BunchWrapTEAPOT(NodeTEAPOT): """ @@ -673,6 +702,9 @@ def track(self, paramsDict: dict) -> None: length = self.getParam("ring_length") TPB.wrapbunch(bunch, length) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + class SolenoidTEAPOT(NodeTEAPOT): """ @@ -721,6 +753,18 @@ def setWaveform(self, waveform: Any) -> None: """ self.waveform = waveform + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + B = self.getParam("B") + if self.waveform: + B *= self.waveform.getStrength() + B *= np.sign(sync_part.charge()) + + return get_matrix_solenoid(sync_part, length=length, B=B) + class MultipoleTEAPOT(NodeTEAPOT): """ @@ -871,6 +915,16 @@ def setWaveform(self, waveform: Any) -> None: """ self.waveform = waveform + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + if np.all(np.abs(self.getParam("kls")) == 0): + return get_matrix_drift(sync_part, length) + + raise NotImplementedError(str(self)) + class QuadTEAPOT(NodeTEAPOT): """ @@ -1034,6 +1088,17 @@ def setWaveform(self, waveform: Any) -> None: """ self.waveform = waveform + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + kq = self.getParam("kq") + if self.waveform: + kq *= self.waveform.getStrength() + + return get_matrix_quad(sync_part, length=length, kq=kq) + class BendTEAPOT(NodeTEAPOT): """ @@ -1248,6 +1313,18 @@ def track(self, paramsDict: dict) -> None: TPB.bend1(bunch, length, theta / 2.0) return + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + nparts = self.getnParts() + theta = self.getParam("theta") / (nparts - 1) + if part_index == 0 or part_index == nparts - 1: + theta *= 0.5 + + return get_matrix_bend(sync_part, length=length, theta=theta) + class RingRFTEAPOT(NodeTEAPOT): """ @@ -1414,6 +1491,26 @@ def setWaveform(self, waveform): """ self.waveform = waveform + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + nparts = self.getnParts() + + scale = 1.0 + if self.waveform is not None: + scale = self.waveform.getStrength() + + scale /= nparts - 1 + kx = scale * self.getParam("kx") + ky = scale * self.getParam("ky") + kE = self.getParam("dE") + + if abs(kx) > 0 or abs(ky) > 0 or abs(kE) > 0: + return np.matmul( + get_matrix_kick(kx=kx, ky=ky, kE=kE), + get_matrix_drift(sync_part, length), + ) + return get_matrix_drift(sync_part, length) + class TiltTEAPOT(BaseTEAPOT): """ @@ -1449,6 +1546,12 @@ def track(self, paramsDict: dict) -> None: bunch = paramsDict["bunch"] TPB.rotatexy(bunch, self.__angle) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + angle = self.getTiltAngle() + if angle == 0: + return None + return get_matrix_tilt(angle) + class FringeFieldTEAPOT(BaseTEAPOT): """ @@ -1478,6 +1581,9 @@ def track(self, paramsDict: dict) -> None: if self.__trackFunc != None: self.__trackFunc(self, paramsDict) + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + return None + def setFringeFieldFunction(self, trackFunction: Callable) -> None: """ Sets the fringe field function that will track the bunch through the fringe. @@ -1595,3 +1701,14 @@ def track(self, paramsDict): def setWaveform(self, waveform): self.waveform = waveform + + def getMatrix(self, sync_part: SyncParticle, part_index: int = -1) -> np.ndarray | None: + length = self.getLength(part_index) + if length <= 0: + return None + + kq = self.getParam("kq") + if self.waveform: + kq *= self.waveform.getStrength() + + return get_matrix_cf(sync_part, length=length, kq=kq) diff --git a/py/orbit/utils/matrix.py b/py/orbit/utils/matrix.py new file mode 100644 index 00000000..0f884f2f --- /dev/null +++ b/py/orbit/utils/matrix.py @@ -0,0 +1,247 @@ +"""Transfer matrix definitions.""" +import math + +import numpy as np + +from orbit.core.bunch import SyncParticle +from orbit.core.orbit_utils import Matrix +from orbit.utils.consts import speed_of_light + + +def get_dp_p_coeff(sync_part: SyncParticle) -> float: + # dE/E = (beta^2) * dp/p + # dE = (beta^2 * E) * dp/p + # dE = (beta^2 * gamma * m * c^2) * dp/p + beta = sync_part.beta() + gamma = sync_part.gamma() + rest_energy = sync_part.mass() # GeV + return 1.0 / (beta**2 * gamma * rest_energy) + + +def get_zp_coeff(sync_part: SyncParticle) -> float: + # dE/E = (beta^2) * dp/p = (beta^2) * (gamma^2) z' + # dE = (beta^2 * gamma^2 * E) * z' + # dE = (beta^2 * gamma^3 * m * c^2) * z' + beta = sync_part.beta() + gamma = sync_part.gamma() + rest_energy = sync_part.mass() + return 1.0 / (beta**2 * gamma**3 * rest_energy) + + +def convert_matrix_dp_p_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: + # v = [x, x', y, y', z, dp/p] + # w = [x, x', y, y', z, dE] + # v = A w + # v -> M v + # w -> A M A^-1 + dp_p_coeff = get_dp_p_coeff(sync_part) + matrix[:5, 5] *= dp_p_coeff + matrix[5, :5] /= dp_p_coeff + matrix[5, 6] /= dp_p_coeff # driving term + return matrix + + +def convert_matrix_zp_to_dE(matrix: np.ndarray, sync_part: SyncParticle) -> np.ndarray: + zp_coeff = get_zp_coeff(sync_part) + matrix[:5, 5] *= zp_coeff + matrix[5, :5] /= zp_coeff + matrix[5, 6] /= zp_coeff # driving term + return matrix + + +def get_matrix_tilt(angle: float) -> np.ndarray: + cos_phi = math.cos(angle) + sin_phi = math.sin(angle) + + M = np.identity(7) + M[0, 0] = M[1, 1] = +cos_phi + M[0, 2] = M[1, 3] = -sin_phi + M[2, 0] = M[3, 1] = +sin_phi + M[2, 2] = M[3, 3] = +cos_phi + return M + + +def get_matrix_kick(kx: float = 0.0, ky: float = 0.0, kE: float = 0.0) -> np.ndarray: + M = np.identity(7) + M[1, -1] = kx + M[3, -1] = ky + M[5, -1] = kE + return M + + +def get_matrix_drift(sync_part: SyncParticle, length: float) -> np.ndarray: + M = np.identity(7) + M[0, 1] = length + M[2, 3] = length + M[4, 5] = length / (sync_part.gamma() ** 2) + M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_quad(sync_part: SyncParticle, length: float, kq: float) -> np.ndarray: + if abs(kq) == 0: + return get_matrix_drift(sync_part, length) + + sqrt_abs_kq = math.sqrt(abs(kq)) + + M = np.identity(7) + if kq > 0: + cx = np.cos(sqrt_abs_kq * length) + sx = np.sin(sqrt_abs_kq * length) + cy = np.cosh(sqrt_abs_kq * length) + sy = np.sinh(sqrt_abs_kq * length) + M[0, 0] = cx + M[0, 1] = +sx / sqrt_abs_kq + M[1, 0] = -sx * sqrt_abs_kq + M[1, 1] = cx + M[2, 2] = cy + M[2, 3] = sy / sqrt_abs_kq + M[3, 2] = sy * sqrt_abs_kq + M[3, 3] = cy + elif kq < 0: + cx = np.cosh(sqrt_abs_kq * length) + sx = np.sinh(sqrt_abs_kq * length) + cy = np.cos(sqrt_abs_kq * length) + sy = np.sin(sqrt_abs_kq * length) + M[0, 0] = cx + M[0, 1] = sx / sqrt_abs_kq + M[1, 0] = sx * sqrt_abs_kq + M[1, 1] = cx + M[2, 2] = cy + M[2, 3] = +sy / sqrt_abs_kq + M[3, 2] = -sy * sqrt_abs_kq + M[3, 3] = cy + + M[4, 5] = length / (sync_part.gamma()**2) + M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_bend(sync_part: SyncParticle, length: float, theta: float) -> np.ndarray: + rho = length / theta + cx = math.cos(theta) + sx = math.sin(theta) + + M = np.identity(7) + M[0, 0] = cx + M[0, 1] = rho * sx + M[0, 5] = rho * (1.0 - cx) + M[1, 0] = -sx / rho + M[1, 1] = cx + M[1, 5] = sx + M[2, 3] = length + M[4, 0] = -sx + M[4, 1] = -rho * (1.0 - cx) + M[4, 5] = -(sync_part.beta() ** 2) * length + rho * sx + M[:5, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_solenoid(sync_part: SyncParticle, length: float, B: float) -> np.ndarray: + if B == 0: + return get_matrix_drift(sync_part, length) + + phase = B * length + + V = np.identity(7) + V[:4, :4] = 0.0 + V[0, 1] = -1.0 / B + V[0, 2] = 0.5 + V[1, 0] = 0.5 * B + V[1, 3] = 1.0 + V[2, 1] = 1.0 / B + V[2, 2] = 0.5 + V[3, 0] = -0.5 * B + V[3, 3] = 1.0 + + M = np.identity(7) + M[0, 0] = +1.0 + M[1, 1] = -1.0 + M[2, 2] = math.cos(phase) + M[2, 3] = math.sin(phase) / B + M[3, 2] = math.sin(phase) * B * -1.0 + M[3, 3] = math.cos(phase) + M[4, 5] = length / (sync_part.gamma()**2) + + M = np.linalg.inv(V) @ M @ V + M[4, 5] *= get_dp_p_coeff(sync_part) # convert_matrix_dp_p_to_dE(M, sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_cf(sync_part: SyncParticle, length: float, kq: float) -> np.ndarray: + if kq == 0: + return get_matrix_drift(sync_part, length) + + sqrt_abs_kq = math.sqrt(abs(kq)) + + cx = math.cos(sqrt_abs_kq * length) + sx = math.sin(sqrt_abs_kq * length) + + M = np.identity(7) + M[0, 0] = M[2, 2] = cx + M[0, 1] = M[2, 3] = +sx / sqrt_abs_kq + M[1, 0] = M[3, 2] = -sx * sqrt_abs_kq + M[1, 1] = M[3, 3] = cx + M[4, 5] = length / (sync_part.gamma()**2) + M[4, 5] *= get_dp_p_coeff(sync_part) + + sync_part.time(sync_part.time() + length / (sync_part.beta() * speed_of_light)) + return M + + +def get_matrix_rf_gap(sync_part: SyncParticle, frequency: float, E0TL: float, phase: float) -> np.ndarray: + gamma = sync_part.gamma() + beta = sync_part.beta() + mass = sync_part.mass() + charge = sync_part.charge() + + kin_energy_in = sync_part.kinEnergy() + charge_E0TL_sin = charge * E0TL * math.sin(phase) + kin_energy_delta = charge * E0TL * math.cos(phase) + + # Calculate parameters in the center of the gap. + sync_part.momentum(sync_part.energyToMomentum(kin_energy_in + kin_energy_delta / 2.0)) + gamma_gap = sync_part.gamma() + beta_gap = sync_part.beta() + + # Move to the end of the gap. + kin_energy_out = kin_energy_in + kin_energy_delta + sync_part.momentum(sync_part.energyToMomentum(kin_energy_out)) + + # The base RF gap is simple - no phase correction. + gamma_out = sync_part.gamma() + beta_out = sync_part.beta() + prime_coeff = (beta * gamma) / (beta_out * gamma_out) + + # Wave momentum + k = 2.0 * math.pi * frequency / speed_of_light + phase_time_coeff = k / beta + + # Transverse focusing coefficient + kappa = -charge * E0TL * k / (2.0 * mass * beta_gap**2 * beta_out * gamma_gap**2 * gamma_out) + d_rp = kappa * math.sin(phase) + + M = np.eye(7) + M[5, 4] = charge_E0TL_sin * phase_time_coeff + M[4, 4] = beta_out / beta + M[1, 1] = prime_coeff + M[3, 3] = prime_coeff + M[1, 0] = d_rp + M[3, 2] = d_rp + return M + + +def orbit_matrix_to_numpy(matrix: Matrix) -> np.ndarray: + matrix_out = np.zeros(matrix.size()) + for i in range(matrix_out.shape[0]): + for j in range(matrix_out.shape[1]): + matrix_out[i, j] = matrix.get(i, j) + return matrix_out diff --git a/py/orbit/utils/meson.build b/py/orbit/utils/meson.build index 6d969897..7ada54d2 100644 --- a/py/orbit/utils/meson.build +++ b/py/orbit/utils/meson.build @@ -12,7 +12,8 @@ py_sources = files([ 'consts.py', 'multiDimArray.py', 'NamedObject.py', - 'orbitFinalize.py' + 'orbitFinalize.py', + 'matrix.py' ]) python.install_sources( diff --git a/src/meson.build b/src/meson.build index c223553d..c7474e58 100644 --- a/src/meson.build +++ b/src/meson.build @@ -1,14 +1,13 @@ - -#project('orbit', ['cpp'], version: '0.0.1') - - # Add Python installation details python = import('python').find_installation('python3', pure: false) # Add C++ compiler details cpp = meson.get_compiler('cpp') -base = meson.current_source_dir() + '/core' +# Add -march=native if the compiler supports it +if cpp.has_argument('-march=native') + add_project_arguments('-march=native', language: 'cpp') +endif dependencies = [] @@ -17,412 +16,303 @@ dependencies += dependency('fftw3', version: '>= 3.0.0', required: true) # Detecting if MPICH or OPENMPI are installed and enabling support if present - mpi_use = get_option('USE_MPI') -# message('MPI_USE is set to', mpi_use) if mpi_use == 'mpich' - message('Requested to use MPICH as the MPI implementation.') - dependencies += dependency('mpich', version: '>= 4.0.0', required: true) - cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1'] + message('Requested to use MPICH as the MPI implementation.') + dependencies += dependency('mpich', version: '>= 4.0.0', required: true) + add_project_arguments('-DUSE_MPI=1', language: 'cpp') - # Configure dependencies or settings specific to MPICH + # Configure dependencies or settings specific to MPICH elif mpi_use == 'ompi' - message('Requested to use OpenMPI as the MPI implementation.') - dependencies += dependency('ompi', version: '>= 4.0.0', required: true) - cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1'] + message('Requested to use OpenMPI as the MPI implementation.') + dependencies += dependency('ompi', version: '>= 4.0.0', required: true) + add_project_arguments('-DUSE_MPI=1', language: 'cpp') elif mpi_use == 'none' - message('Requested to not use MPI.') - cpp_args = ['-fPIC', '-std=c++11'] + message('Requested to not use MPI.') else mpich_dependency = dependency('mpich', version: '>= 4.0.0', required: false) openmpi_dependency = dependency('ompi', version: '>= 4.0.0', required: false) if mpich_dependency.found() - cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1'] - dependencies += mpich_dependency - message('Using MPICH as the MPI implementation.') + add_project_arguments('-DUSE_MPI=1', language: 'cpp') + dependencies += mpich_dependency + message('Using MPICH as the MPI implementation.') elif openmpi_dependency.found() - cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1'] - dependencies += openmpi_dependency - message('Using OpenMPI as the MPI implementation.') + add_project_arguments('-DUSE_MPI=1', language: 'cpp') + dependencies += openmpi_dependency + message('Using OpenMPI as the MPI implementation.') else - cpp_args = ['-fPIC', '-std=c++11'] - message('MPI will not be used.') + message('MPI will not be used.') endif endif sources = files([ - 'linac/wrap_linacmodule.cc', - 'linac/tracking/linac_tracking.cc', - 'linac/tracking/wrap_linac_tracking.cc', - 'linac/rfgap/wrap_BaseRfGap.cc', - 'linac/rfgap/BaseRfGap_slow.cc', - 'linac/rfgap/wrap_RfGapTTF_slow.cc', - 'linac/rfgap/wrap_RfGapThreePointTTF.cc', - 'linac/rfgap/BaseRfGap.cc', - 'linac/rfgap/RfGapThreePointTTF.cc', - 'linac/rfgap/MatrixRfGap.cc', - 'linac/rfgap/wrap_BaseRfGap_slow.cc', - 'linac/rfgap/wrap_RfGapThreePointTTF_slow.cc', - 'linac/rfgap/wrap_SuperFishFieldSource.cc', - 'linac/rfgap/RfGapTTF.cc', - 'linac/rfgap/SuperFishFieldSource.cc', - 'linac/rfgap/RfGapTTF_slow.cc', - 'linac/rfgap/RfGapThreePointTTF_slow.cc', - 'linac/rfgap/wrap_RfGapTTF.cc', - 'linac/rfgap/wrap_MatrixRfGap.cc', - 'orbit/OrbitConst.cc', - 'orbit/wrap_bunch.cc', - 'orbit/wrap_syncpart.cc', - 'orbit/Bunch.cc', - 'orbit/SyncPart.cc', - 'orbit/BunchDiagnostics/wrap_bunch_tune_analysis.cc', - 'orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.cc', - 'orbit/BunchDiagnostics/BunchTwissAnalysis.cc', - 'orbit/BunchDiagnostics/BunchTuneAnalysis.cc', - 'orbit/FieldTracker/wrap_fieldtracker.cc', - 'orbit/FieldTracker/FieldTracker.cc', - 'orbit/MaterialInteractions/wrap_foil.cc', - 'orbit/MaterialInteractions/Foil.cc', - 'orbit/MaterialInteractions/Collimator.cc', - 'orbit/MaterialInteractions/numrecipes.cc', - 'orbit/MaterialInteractions/cross_sections.cc', - 'orbit/MaterialInteractions/wrap_collimator.cc', - 'orbit/MaterialInteractions/MaterialInteractions.cc', - 'orbit/RFCavities/Dual_Harmonic_Cav.cc', - 'orbit/RFCavities/Frequency_Cav.cc', - 'orbit/RFCavities/Barrier_Cav.cc', - 'orbit/RFCavities/wrap_Harmonic_Cav.cc', - 'orbit/RFCavities/Harmonic_Cav.cc', - 'orbit/RFCavities/wrap_Dual_Harmonic_Cav.cc', - 'orbit/RFCavities/wrap_Barrier_Cav.cc', - 'orbit/RFCavities/wrap_Frequency_Cav.cc', - 'orbit/RFCavities/wrap_rfcavities.cc', - 'orbit/Impedances/wrap_TImpedance.cc', - 'orbit/Impedances/TImpedance.cc', - 'orbit/Impedances/wrap_impedances.cc', - 'orbit/Impedances/LImpedance.cc', - 'orbit/Impedances/wrap_LImpedance.cc', - 'orbit/SynchPartRedefinition/SynchPartRedefinitionZdE.cc', - 'orbit/SynchPartRedefinition/wrap_synch_part_redefinition_z_de.cc', - 'orbit/Errors/errorbase.cc', - 'orbit/Errors/wrap_errorbase.cc', - 'orbit/ParticlesAttributes/ParticleInitialCoordinates.cc', - 'orbit/ParticlesAttributes/ParticleIdNumber.cc', - 'orbit/ParticlesAttributes/ParticleMacroSize.cc', - 'orbit/ParticlesAttributes/ParticleAttributesFactory.cc', - 'orbit/ParticlesAttributes/WaveFunctionAmplitudes.cc', - 'orbit/ParticlesAttributes/pq_coordinates.cc', - 'orbit/ParticlesAttributes/Evolution.cc', - 'orbit/ParticlesAttributes/ParticleAttributes.cc', - 'orbit/ParticlesAttributes/TurnNumberAttributes.cc', - 'orbit/ParticlesAttributes/LostParticleAttributes.cc', - 'orbit/ParticlesAttributes/ParticlePhaseAttributes.cc', - 'orbit/ParticlesAttributes/part_time.cc', - 'orbit/ParticlesAttributes/AtomPopulations.cc', - 'orbit/Apertures/CircleApertureShape.cc', - 'orbit/Apertures/PhaseAperture.cc', - 'orbit/Apertures/Aperture.cc', - 'orbit/Apertures/wrap_aperture.cc', - 'orbit/Apertures/wrap_CompositeApertureShape.cc', - 'orbit/Apertures/RectangularApertureShape.cc', - 'orbit/Apertures/PyBaseApertureShape.cc', - 'orbit/Apertures/ConvexApertureShape.cc', - 'orbit/Apertures/wrap_ConvexApertureShape.cc', - 'orbit/Apertures/wrap_PyBaseApertureShape.cc', - 'orbit/Apertures/EnergyAperture.cc', - 'orbit/Apertures/BaseAperture.cc', - 'orbit/Apertures/EllipseApertureShape.cc', - 'orbit/Apertures/wrap_BaseAperture.cc', - 'orbit/Apertures/wrap_PrimitiveApertureShape.cc', - 'orbit/Apertures/CompositeApertureShape.cc', - 'orbit/Apertures/wrap_EnergyAperture.cc', - 'orbit/Apertures/wrap_TAperture.cc', - 'orbit/Apertures/wrap_PhaseAperture.cc', - 'orbit/Apertures/BaseApertureShape.cc', - 'trackerrk4/wrap_runge_kutta_tracker.cc', - 'trackerrk4/wrap_py_external_effects.cc', - 'trackerrk4/ExternalEffects.cc', - 'trackerrk4/ExtEffectsContainer.cc', - 'trackerrk4/wrap_trackerrk4.cc', - 'trackerrk4/PyExternalEffects.cc', - 'trackerrk4/RungeKuttaTracker.cc', - 'trackerrk4/wrap_ext_effects_container.cc', - 'mpi/wrap_orbit_mpi.cc', - 'mpi/wrap_mpi_comm.cc', - 'mpi/orbit_mpi.cc', - 'mpi/wrap_mpi_datatype.cc', - 'mpi/wrap_mpi_request.cc', - 'mpi/wrap_mpi_op.cc', - 'mpi/wrap_mpi_status.cc', - 'mpi/wrap_mpi_group.cc', - 'spacecharge/PoissonSolverFFT3D.cc', - 'spacecharge/SpaceChargeCalcUnifEllipse.cc', - 'spacecharge/ForceSolver2D.cc', - 'spacecharge/PoissonSolverFFT2D.cc', - 'spacecharge/wrap_poissonsolverfft2d.cc', - 'spacecharge/wrap_poissonsolverfft3d.cc', - 'spacecharge/PoissonSolver2D.cc', - 'spacecharge/SpaceChargeForceCalc2p5D.cc', - 'spacecharge/PoissonSolver3D.cc', - 'spacecharge/SpaceChargeCalc3D.cc', - 'spacecharge/wrap_boundary2d.cc', - 'spacecharge/wrap_forcesolverfft2d.cc', - 'spacecharge/ForceSolverFFT2D.cc', - 'spacecharge/SpaceChargeCalcSliceBySlice2D.cc', - 'spacecharge/Grid3D.cc', - 'spacecharge/wrap_spacecharge.cc', - 'spacecharge/wrap_spacechargecalc2p5d_rb.cc', - 'spacecharge/Grid2D.cc', - 'spacecharge/SpaceChargeCalc2p5Drb.cc', - 'spacecharge/Grid1D.cc', - 'spacecharge/LSpaceChargeCalc.cc', - 'spacecharge/wrap_grid1D.cc', - 'spacecharge/wrap_uniform_ellipsoid_field_calculator.cc', - 'spacecharge/ShapedBoundary2D.cc', - 'spacecharge/wrap_grid3D.cc', - 'spacecharge/BaseBoundary2D.cc', - 'spacecharge/wrap_grid2D.cc', - 'spacecharge/wrap_spacechargecalc_uniform_ellipse.cc', - 'spacecharge/wrap_spacechargecalc3d.cc', - 'spacecharge/UniformEllipsoidFieldCalculator.cc', - 'spacecharge/wrap_spacechargeforcecalc2p5d.cc', - 'spacecharge/SpaceChargeCalc2p5D.cc', - 'spacecharge/wrap_spacechargecalc_slicebyslice_2D.cc', - 'spacecharge/wrap_spacechargecalc2p5d.cc', - 'spacecharge/wrap_lspacechargecalc.cc', - 'utils/wrap_function.cc', - 'utils/FieldSourceContainer.cc', - 'utils/AttributesBucket.cc', - 'utils/bessel.cc', - 'utils/PyBaseFieldSource.cc', - 'utils/wrap_field_source_container.cc', - 'utils/OU_SplineCH.cc', - 'utils/CppPyWrapper.cc', - 'utils/wrap_random.cc', - 'utils/wrap_utils.cc', - 'utils/StringUtils.cc', - 'utils/wrap_py_base_field_source.cc', - 'utils/wrap_splinech.cc', - 'utils/Random.cc', - 'utils/wrap_numrecipes.cc', - 'utils/BaseFieldSource.cc', - 'utils/OU_Function.cc', - 'utils/MathPolynomial.cc', - 'utils/BufferStore.cc', - 'utils/statistics/wrap_statmoments2d.cc', - 'utils/statistics/StatMoments2D.cc', - 'utils/field_sources/wrap_loop_field_source.cc', - 'utils/field_sources/LoopFieldSource.cc', - 'utils/field_sources/QuadFieldSource.cc', - 'utils/field_sources/ShiftedFieldSource.cc', - 'utils/field_sources/wrap_dipole_field_source.cc', - 'utils/field_sources/wrap_quad_field_source.cc', - 'utils/field_sources/wrap_magnetic_field_source_grid3d.cc', - 'utils/field_sources/DipoleFieldSource.cc', - 'utils/field_sources/wrap_field_sources_module.cc', - 'utils/field_sources/MagnetFieldSourceGrid3D.cc', - 'utils/integration/gauss_legendre_points.cc', - 'utils/integration/wrap_gauss_legendre_integrator.cc', - 'utils/integration/GaussLegendreIntegrator.cc', - 'utils/bunch/ParticlesWithIdFunctions.cc', - 'utils/bunch/wrap_bunch_extrema_calculator.cc', - 'utils/bunch/TwissFilteringFunctions.cc', - 'utils/bunch/InitialCoordsAttrFunctions.cc', - 'utils/bunch/BunchExtremaCalculator.cc', - 'utils/bunch/wrap_bunch_utils_functions.cc', - 'utils/harmonic_analysis/HarmonicData.cc', - 'utils/harmonic_analysis/wrap_harmonic_data.cc', - 'utils/polynomial/wrap_polynomial.cc', - 'utils/polynomial/OU_Polynomial.cc', - 'utils/matrix/wrap_matrix.cc', - 'utils/matrix/PhaseVector.cc', - 'utils/matrix/MatrixOperations.cc', - 'utils/matrix/wrap_phase_vector.cc', - 'utils/matrix/Matrix.cc', - 'teapot/wrap_teapotbase.cc', - 'teapot/wrap_matrix_generator.cc', - 'teapot/teapotbase.cc', - 'teapot/MatrixGenerator.cc' + 'linac/wrap_linacmodule.cc', + 'linac/tracking/linac_tracking.cc', + 'linac/tracking/wrap_linac_tracking.cc', + 'linac/rfgap/wrap_BaseRfGap.cc', + 'linac/rfgap/BaseRfGap_slow.cc', + 'linac/rfgap/wrap_RfGapTTF_slow.cc', + 'linac/rfgap/wrap_RfGapThreePointTTF.cc', + 'linac/rfgap/BaseRfGap.cc', + 'linac/rfgap/RfGapThreePointTTF.cc', + 'linac/rfgap/MatrixRfGap.cc', + 'linac/rfgap/wrap_BaseRfGap_slow.cc', + 'linac/rfgap/wrap_RfGapThreePointTTF_slow.cc', + 'linac/rfgap/wrap_SuperFishFieldSource.cc', + 'linac/rfgap/RfGapTTF.cc', + 'linac/rfgap/SuperFishFieldSource.cc', + 'linac/rfgap/RfGapTTF_slow.cc', + 'linac/rfgap/RfGapThreePointTTF_slow.cc', + 'linac/rfgap/wrap_RfGapTTF.cc', + 'linac/rfgap/wrap_MatrixRfGap.cc', + 'orbit/OrbitConst.cc', + 'orbit/wrap_bunch.cc', + 'orbit/wrap_syncpart.cc', + 'orbit/Bunch.cc', + 'orbit/SyncPart.cc', + 'orbit/BunchDiagnostics/wrap_bunch_tune_analysis.cc', + 'orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.cc', + 'orbit/BunchDiagnostics/BunchTwissAnalysis.cc', + 'orbit/BunchDiagnostics/BunchTuneAnalysis.cc', + 'orbit/FieldTracker/wrap_fieldtracker.cc', + 'orbit/FieldTracker/FieldTracker.cc', + 'orbit/MaterialInteractions/wrap_foil.cc', + 'orbit/MaterialInteractions/Foil.cc', + 'orbit/MaterialInteractions/Collimator.cc', + 'orbit/MaterialInteractions/numrecipes.cc', + 'orbit/MaterialInteractions/cross_sections.cc', + 'orbit/MaterialInteractions/wrap_collimator.cc', + 'orbit/MaterialInteractions/MaterialInteractions.cc', + 'orbit/RFCavities/Dual_Harmonic_Cav.cc', + 'orbit/RFCavities/Frequency_Cav.cc', + 'orbit/RFCavities/Barrier_Cav.cc', + 'orbit/RFCavities/wrap_Harmonic_Cav.cc', + 'orbit/RFCavities/Harmonic_Cav.cc', + 'orbit/RFCavities/wrap_Dual_Harmonic_Cav.cc', + 'orbit/RFCavities/wrap_Barrier_Cav.cc', + 'orbit/RFCavities/wrap_Frequency_Cav.cc', + 'orbit/RFCavities/wrap_rfcavities.cc', + 'orbit/Impedances/wrap_TImpedance.cc', + 'orbit/Impedances/TImpedance.cc', + 'orbit/Impedances/wrap_impedances.cc', + 'orbit/Impedances/LImpedance.cc', + 'orbit/Impedances/wrap_LImpedance.cc', + 'orbit/SynchPartRedefinition/SynchPartRedefinitionZdE.cc', + 'orbit/SynchPartRedefinition/wrap_synch_part_redefinition_z_de.cc', + 'orbit/Errors/errorbase.cc', + 'orbit/Errors/wrap_errorbase.cc', + 'orbit/ParticlesAttributes/ParticleInitialCoordinates.cc', + 'orbit/ParticlesAttributes/ParticleIdNumber.cc', + 'orbit/ParticlesAttributes/ParticleMacroSize.cc', + 'orbit/ParticlesAttributes/ParticleAttributesFactory.cc', + 'orbit/ParticlesAttributes/WaveFunctionAmplitudes.cc', + 'orbit/ParticlesAttributes/pq_coordinates.cc', + 'orbit/ParticlesAttributes/Evolution.cc', + 'orbit/ParticlesAttributes/ParticleAttributes.cc', + 'orbit/ParticlesAttributes/TurnNumberAttributes.cc', + 'orbit/ParticlesAttributes/LostParticleAttributes.cc', + 'orbit/ParticlesAttributes/ParticlePhaseAttributes.cc', + 'orbit/ParticlesAttributes/part_time.cc', + 'orbit/ParticlesAttributes/AtomPopulations.cc', + 'orbit/Apertures/CircleApertureShape.cc', + 'orbit/Apertures/PhaseAperture.cc', + 'orbit/Apertures/Aperture.cc', + 'orbit/Apertures/wrap_aperture.cc', + 'orbit/Apertures/wrap_CompositeApertureShape.cc', + 'orbit/Apertures/RectangularApertureShape.cc', + 'orbit/Apertures/PyBaseApertureShape.cc', + 'orbit/Apertures/ConvexApertureShape.cc', + 'orbit/Apertures/wrap_ConvexApertureShape.cc', + 'orbit/Apertures/wrap_PyBaseApertureShape.cc', + 'orbit/Apertures/EnergyAperture.cc', + 'orbit/Apertures/BaseAperture.cc', + 'orbit/Apertures/EllipseApertureShape.cc', + 'orbit/Apertures/wrap_BaseAperture.cc', + 'orbit/Apertures/wrap_PrimitiveApertureShape.cc', + 'orbit/Apertures/CompositeApertureShape.cc', + 'orbit/Apertures/wrap_EnergyAperture.cc', + 'orbit/Apertures/wrap_TAperture.cc', + 'orbit/Apertures/wrap_PhaseAperture.cc', + 'orbit/Apertures/BaseApertureShape.cc', + 'trackerrk4/wrap_runge_kutta_tracker.cc', + 'trackerrk4/wrap_py_external_effects.cc', + 'trackerrk4/ExternalEffects.cc', + 'trackerrk4/ExtEffectsContainer.cc', + 'trackerrk4/wrap_trackerrk4.cc', + 'trackerrk4/PyExternalEffects.cc', + 'trackerrk4/RungeKuttaTracker.cc', + 'trackerrk4/wrap_ext_effects_container.cc', + 'mpi/wrap_orbit_mpi.cc', + 'mpi/wrap_mpi_comm.cc', + 'mpi/orbit_mpi.cc', + 'mpi/wrap_mpi_datatype.cc', + 'mpi/wrap_mpi_request.cc', + 'mpi/wrap_mpi_op.cc', + 'mpi/wrap_mpi_status.cc', + 'mpi/wrap_mpi_group.cc', + 'spacecharge/PoissonSolverFFT3D.cc', + 'spacecharge/SpaceChargeCalcUnifEllipse.cc', + 'spacecharge/ForceSolver2D.cc', + 'spacecharge/PoissonSolverFFT2D.cc', + 'spacecharge/wrap_poissonsolverfft2d.cc', + 'spacecharge/wrap_poissonsolverfft3d.cc', + 'spacecharge/PoissonSolver2D.cc', + 'spacecharge/SpaceChargeForceCalc2p5D.cc', + 'spacecharge/PoissonSolver3D.cc', + 'spacecharge/SpaceChargeCalc3D.cc', + 'spacecharge/wrap_boundary2d.cc', + 'spacecharge/wrap_forcesolverfft2d.cc', + 'spacecharge/ForceSolverFFT2D.cc', + 'spacecharge/SpaceChargeCalcSliceBySlice2D.cc', + 'spacecharge/Grid3D.cc', + 'spacecharge/wrap_spacecharge.cc', + 'spacecharge/wrap_spacechargecalc2p5d_rb.cc', + 'spacecharge/Grid2D.cc', + 'spacecharge/SpaceChargeCalc2p5Drb.cc', + 'spacecharge/Grid1D.cc', + 'spacecharge/LSpaceChargeCalc.cc', + 'spacecharge/wrap_grid1D.cc', + 'spacecharge/wrap_uniform_ellipsoid_field_calculator.cc', + 'spacecharge/ShapedBoundary2D.cc', + 'spacecharge/wrap_grid3D.cc', + 'spacecharge/BaseBoundary2D.cc', + 'spacecharge/wrap_grid2D.cc', + 'spacecharge/wrap_spacechargecalc_uniform_ellipse.cc', + 'spacecharge/wrap_spacechargecalc3d.cc', + 'spacecharge/UniformEllipsoidFieldCalculator.cc', + 'spacecharge/wrap_spacechargeforcecalc2p5d.cc', + 'spacecharge/SpaceChargeCalc2p5D.cc', + 'spacecharge/wrap_spacechargecalc_slicebyslice_2D.cc', + 'spacecharge/wrap_spacechargecalc2p5d.cc', + 'spacecharge/wrap_lspacechargecalc.cc', + 'utils/wrap_function.cc', + 'utils/FieldSourceContainer.cc', + 'utils/AttributesBucket.cc', + 'utils/bessel.cc', + 'utils/PyBaseFieldSource.cc', + 'utils/wrap_field_source_container.cc', + 'utils/OU_SplineCH.cc', + 'utils/CppPyWrapper.cc', + 'utils/wrap_random.cc', + 'utils/wrap_utils.cc', + 'utils/StringUtils.cc', + 'utils/wrap_py_base_field_source.cc', + 'utils/wrap_splinech.cc', + 'utils/Random.cc', + 'utils/wrap_numrecipes.cc', + 'utils/BaseFieldSource.cc', + 'utils/OU_Function.cc', + 'utils/MathPolynomial.cc', + 'utils/BufferStore.cc', + 'utils/statistics/wrap_statmoments2d.cc', + 'utils/statistics/StatMoments2D.cc', + 'utils/field_sources/wrap_loop_field_source.cc', + 'utils/field_sources/LoopFieldSource.cc', + 'utils/field_sources/QuadFieldSource.cc', + 'utils/field_sources/ShiftedFieldSource.cc', + 'utils/field_sources/wrap_dipole_field_source.cc', + 'utils/field_sources/wrap_quad_field_source.cc', + 'utils/field_sources/wrap_magnetic_field_source_grid3d.cc', + 'utils/field_sources/DipoleFieldSource.cc', + 'utils/field_sources/wrap_field_sources_module.cc', + 'utils/field_sources/MagnetFieldSourceGrid3D.cc', + 'utils/integration/gauss_legendre_points.cc', + 'utils/integration/wrap_gauss_legendre_integrator.cc', + 'utils/integration/GaussLegendreIntegrator.cc', + 'utils/bunch/ParticlesWithIdFunctions.cc', + 'utils/bunch/wrap_bunch_extrema_calculator.cc', + 'utils/bunch/TwissFilteringFunctions.cc', + 'utils/bunch/InitialCoordsAttrFunctions.cc', + 'utils/bunch/BunchExtremaCalculator.cc', + 'utils/bunch/wrap_bunch_utils_functions.cc', + 'utils/harmonic_analysis/HarmonicData.cc', + 'utils/harmonic_analysis/wrap_harmonic_data.cc', + 'utils/polynomial/wrap_polynomial.cc', + 'utils/polynomial/OU_Polynomial.cc', + 'utils/matrix/wrap_matrix.cc', + 'utils/matrix/PhaseVector.cc', + 'utils/matrix/MatrixOperations.cc', + 'utils/matrix/wrap_phase_vector.cc', + 'utils/matrix/Matrix.cc', + 'teapot/wrap_teapotbase.cc', + 'teapot/wrap_matrix_generator.cc', + 'teapot/teapotbase.cc', + 'teapot/MatrixGenerator.cc' ]) + inc = include_directories([ - 'main', - 'orbit/SynchPartRedefinition', - 'trackerrk4', - 'utils/ellipticalint', - 'teapot', - 'mpi', - 'orbit/ParticlesAttributes', - 'orbit/Errors', - 'orbit/FieldTracker', - 'utils/statistics', - 'utils', - 'linac/tracking', - 'spacecharge', - 'orbit/MaterialInteractions', - 'utils/field_sources', - 'utils/bunch', - 'orbit/Impedances', - 'utils/harmonic_analysis', - 'linac', - 'utils/polynomial', - 'linac/rfgap', - 'orbit/RFCavities', - 'utils/matrix', - 'orbit/BunchDiagnostics', - 'orbit', - 'utils/integration', - 'orbit/Apertures' + 'main', + 'orbit/SynchPartRedefinition', + 'trackerrk4', + 'utils/ellipticalint', + 'teapot', + 'mpi', + 'orbit/ParticlesAttributes', + 'orbit/Errors', + 'orbit/FieldTracker', + 'utils/statistics', + 'utils', + 'linac/tracking', + 'spacecharge', + 'orbit/MaterialInteractions', + 'utils/field_sources', + 'utils/bunch', + 'orbit/Impedances', + 'utils/harmonic_analysis', + 'linac', + 'utils/polynomial', + 'linac/rfgap', + 'orbit/RFCavities', + 'utils/matrix', + 'orbit/BunchDiagnostics', + 'orbit', + 'utils/integration', + 'orbit/Apertures' ]) - -core_lib = library('core', - sources: sources, - include_directories: inc, - cpp_args: cpp_args, - override_options: ['b_lundef=false'], - dependencies: dependencies, - install: true, +core_lib = library( + 'core', + sources: sources, + include_directories: inc, + override_options: ['b_lundef=false'], + dependencies: dependencies, + install: true, ) core_dep = declare_dependency(link_with : core_lib) -python.extension_module('orbit_mpi', - sources: [base + '/mpi_init.cc'], - include_directories: inc, - cpp_args: cpp_args, - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - +ext_modules = { + 'orbit_mpi': 'mpi_init.cc', + 'bunch': 'bunch_init.cc', + 'spacecharge': 'spacecharge_init.cc', + 'trackerrk4': 'trackerrk4_init.cc', + 'teapot_base': 'teapot_base_init.cc', + 'linac': 'linac_init.cc', + 'orbit_utils': 'utils_init.cc', + 'aperture': 'aperture_init.cc', + 'foil': 'foil_init.cc', + 'field_sources': 'field_sources_init.cc', + 'rfcavities': 'rfcavities_init.cc', + 'impedances': 'impedances_init.cc', + 'fieldtracker': 'fieldtracker_init.cc', + 'collimator': 'collimator_init.cc', + 'error_base': 'error_base_init.cc', +} -python.extension_module('bunch', - sources: [base + '/bunch_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('spacecharge', - sources: [base + '/spacecharge_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('trackerrk4', - sources: [base + '/trackerrk4_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('teapot_base', - sources: [base + '/teapot_base_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('linac', - sources: [base + '/linac_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('orbit_utils', - sources: [base + '/utils_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('aperture', - sources: [base + '/aperture_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('foil', - sources: [base + '/foil_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('field_sources', - sources: [base + '/field_sources_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('rfcavities', - sources: [base + '/rfcavities_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('impedances', - sources: [base + '/impedances_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('fieldtracker', - sources: [base + '/fieldtracker_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('collimator', - sources: [base + '/collimator_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) +base = meson.current_source_dir() + '/core' -python.extension_module('error_base', - sources: [base + '/error_base_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) +foreach name, src : ext_modules + python.extension_module( + name, + sources: base / src, + include_directories: inc, + dependencies: [core_dep], + install: true, + subdir: 'orbit/core', + ) +endforeach diff --git a/src/orbit/BunchDiagnostics/BunchTuneAnalysis.cc b/src/orbit/BunchDiagnostics/BunchTuneAnalysis.cc index 12d11937..659b58a6 100644 --- a/src/orbit/BunchDiagnostics/BunchTuneAnalysis.cc +++ b/src/orbit/BunchDiagnostics/BunchTuneAnalysis.cc @@ -7,97 +7,155 @@ #include #include -/** Constructor */ -BunchTuneAnalysis::BunchTuneAnalysis(): CppPyWrapper(NULL) -{ - betax = 0; - alphax = 0; - etax = 0; - etapx = 0; - betay = 0; - alphay = 0; + +BunchTuneAnalysis::BunchTuneAnalysis(): CppPyWrapper(NULL) { + double matrix[6][6] = { + {1.0, 0.0, 0.0, 0.0, 0.0, 0.0}, + {0.0, 1.0, 0.0, 0.0, 0.0, 0.0}, + {0.0, 0.0, 1.0, 0.0, 0.0, 0.0}, + {0.0, 0.0, 0.0, 1.0, 0.0, 0.0}, + {0.0, 0.0, 0.0, 0.0, 1.0, 0.0}, + {0.0, 0.0, 0.0, 0.0, 0.0, 1.0} + }; + int erase = 1; } -/** Destructor */ -BunchTuneAnalysis::~BunchTuneAnalysis() -{ + +BunchTuneAnalysis::~BunchTuneAnalysis() {} + +double BunchTuneAnalysis::getNormMatrixElement(int i, int j) { + return matrix[i][j]; } -void BunchTuneAnalysis::assignTwiss(double bx, double ax, double dx, double dpx, double by, double ay){ - betax = bx; - alphax = ax; - etax = dx; - etapx = dpx; - betay = by; - alphay = ay; +void BunchTuneAnalysis::setNormMatrixElement(int i, int j, double value) { + matrix[i][j] = value; + erase = 1; +} + +void BunchTuneAnalysis::setNormMatrixFromTwiss(double betax, double alphax, double etax, double etapx, double betay, double alphay) { + erase = 1; + + // Set V_{-1} = I + for (int i = 0; i < 6; i++) { + for (int j = 0; j < 6; j++) { + matrix[i][j] = 0.0; + } + } + for (int i = 0; i < 6; i++) { + matrix[i][i] = 1.0; + } + + // 2D normalization (x-x') + matrix[0][0] = 1.0 / sqrt(betax); + matrix[0][1] = 0.0; + matrix[1][0] = alphax / sqrt(betax); + matrix[1][1] = sqrt(betax); + + // 2D normalization (y-y') + matrix[2][2] = 1.0 / sqrt(betay); + matrix[2][3] = 0.0; + matrix[3][2] = alphay / sqrt(betay); + matrix[3][3] = sqrt(betay); + + // Dispersion (x-x') + matrix[0][5] = -etax / sqrt(betax); + matrix[1][5] = -etax * (alphax / sqrt(betax)) - etapx * sqrt(betax); + + // Dispersion (y-y') + double etay = 0.0; + double etapy = 0.0; + matrix[2][5] = -etay / sqrt(betay); + matrix[3][5] = -etay * (alphay / sqrt(betay)) - etapy * sqrt(betay); + } -/** Performs the Tune analysis of the bunch */ void BunchTuneAnalysis::analyzeBunch(Bunch* bunch){ - //initialization bunch->compress(); SyncPart* syncPart = bunch->getSyncPart(); double beta = syncPart->getBeta(); double** part_coord_arr = bunch->coordArr(); if(!bunch->hasParticleAttributes("ParticlePhaseAttributes")){ - cerr<<"adding particle phase information attribute\n"; + cerr<<"BunchTuneAnalysis: Adding particle phase information attribute.\n"; std::map tunemap; - tunemap.insert(std::make_pair("xLastPhase", 0)); - tunemap.insert(std::make_pair("yLastPhase", 0)); - tunemap.insert(std::make_pair("xLastTune", 0)); - tunemap.insert(std::make_pair("yLastTune", 0)); - tunemap.insert(std::make_pair("xAction", 0)); - tunemap.insert(std::make_pair("yAction", 0)); + tunemap.insert(std::make_pair("phase_1", 0)); + tunemap.insert(std::make_pair("phase_2", 0)); + tunemap.insert(std::make_pair("tune_1", 0)); + tunemap.insert(std::make_pair("tune_2", 0)); + tunemap.insert(std::make_pair("action_1", 0)); + tunemap.insert(std::make_pair("action_2", 0)); bunch->addParticleAttributes("ParticlePhaseAttributes", tunemap); } - if(bunch->hasParticleAttributes("ParticlePhaseAttributes")){ - for (int i=0; i < bunch->getSize(); i++) - { - double x = part_coord_arr[i][0]; + if (bunch->hasParticleAttributes("ParticlePhaseAttributes") && erase == 1) { + cerr<<"BunchTuneAnalysis: Normalization matrix has been updated. Setting particle phases to zero. Tunes will be accurate after the next `analyzeBunch` call.\n"; + for (int i=0; i < bunch->getSize(); i++) { + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 0) = 0.0; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 1) = 0.0; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 2) = 0.0; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 3) = 0.0; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 4) = 0.0; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 5) = 0.0; + } + } + + if (bunch->hasParticleAttributes("ParticlePhaseAttributes")) { + for (int i=0; i < bunch->getSize(); i++) { + // Extract phase space coordinates + double x = part_coord_arr[i][0]; double xp = part_coord_arr[i][1]; - double y = part_coord_arr[i][2]; + double y = part_coord_arr[i][2]; double yp = part_coord_arr[i][3]; + double z = part_coord_arr[i][4]; double Etot = syncPart->getEnergy() + syncPart->getMass(); - double dpp = 1/(beta*beta)*part_coord_arr[i][5]/Etot; - - double xval = (x - etax * dpp)/sqrt(betax); - double xpval = (xp - etapx * dpp) * sqrt(betax) + xval * alphax; - double yval = y / sqrt(betay); - double ypval = (yp + y * alphay/betay) * sqrt(betay); - - double angle = atan2(xpval, xval); - if(angle < 0.) angle += (2.0*OrbitConst::PI); - double xPhase = angle; - double xPhaseOld = bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i,0); - double xTune = (xPhaseOld - xPhase) / (2.0*OrbitConst::PI); - if(xTune < 0.) xTune += 1.; - bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 0) = xPhase; - bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 2) = xTune; - - angle = atan2(ypval, yval); - if(angle < 0.) angle += (2.0*OrbitConst::PI); - double yPhase = angle; - double yPhaseOld = bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i,1); - double yTune = (yPhaseOld - yPhase) / (2.0*OrbitConst::PI); - if(yTune < 0.) yTune += 1.; - bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 1) = yPhase; - bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 3) = yTune; - - double xcanonical = x - etax * dpp; - double ycanonical = y; - double xpfac = xp - etapx * dpp; - double ypfac = yp; - double pxcanonical = xpfac + xcanonical * (alphax/betax); - double pycanonical = ypfac + ycanonical * (alphay/betay); - double xAction = xcanonical * xcanonical / betax + pxcanonical * pxcanonical * betax; - double yAction = ycanonical * ycanonical / betay + pycanonical * pycanonical * betay; - - bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 4) = xAction; - bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 5) = yAction; + double dpp = 1.0 / (beta * beta) * part_coord_arr[i][5] / Etot; + + // Normalize phase space coordinates + double u1 = matrix[0][0] * x + matrix[0][1] * xp + matrix[0][2] * y + matrix[0][3] * yp + matrix[0][4] * z + matrix[0][5] * dpp; + double u1p = matrix[1][0] * x + matrix[1][1] * xp + matrix[1][2] * y + matrix[1][3] * yp + matrix[1][4] * z + matrix[1][5] * dpp; + double u2 = matrix[2][0] * x + matrix[2][1] * xp + matrix[2][2] * y + matrix[2][3] * yp + matrix[2][4] * z + matrix[2][5] * dpp; + double u2p = matrix[3][0] * x + matrix[3][1] * xp + matrix[3][2] * y + matrix[3][3] * yp + matrix[3][4] * z + matrix[3][5] * dpp; + double u3 = matrix[4][0] * x + matrix[4][1] * xp + matrix[4][2] * y + matrix[4][3] * yp + matrix[4][4] * z + matrix[4][5] * dpp; + double u3p = matrix[5][0] * x + matrix[5][1] * xp + matrix[5][2] * y + matrix[5][3] * yp + matrix[5][4] * z + matrix[5][5] * dpp; + + // Compute phase advance (mode 1) + double angle = atan2(u1p, u1); + if (angle < 0.0) { + angle += (2.0 * OrbitConst::PI); + } + double phase1 = angle; + double phase1Old = bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 0); + double tune1 = (phase1Old - phase1) / (2.0 * OrbitConst::PI); + if (tune1 < 0.0) { + tune1 += 1.0; } + + // Compute phase advance (mode 2) + angle = atan2(u2p, u2); + if (angle < 0.0) { + angle += (2.0 * OrbitConst::PI); + } + double phase2 = angle; + double phase2Old = bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 1); + double tune2 = (phase2Old - phase2) / (2.0 * OrbitConst::PI); + if (tune2 < 0.0) { + tune2 += 1.0; + } + + // Compute actions + double action1 = (u1 * u1 + u1p * u1p) / 2.0; + double action2 = (u2 * u2 + u2p * u2p) / 2.0; + + // Update ParticlePhaseAttributes + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 0) = phase1; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 1) = phase2; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 2) = tune1; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 3) = tune2; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 4) = action1; + bunch->getParticleAttributes("ParticlePhaseAttributes")->attValue(i, 5) = action2; + } + erase = 0; } -} +} \ No newline at end of file diff --git a/src/orbit/BunchDiagnostics/BunchTuneAnalysis.hh b/src/orbit/BunchDiagnostics/BunchTuneAnalysis.hh index b3701788..84b137b0 100644 --- a/src/orbit/BunchDiagnostics/BunchTuneAnalysis.hh +++ b/src/orbit/BunchDiagnostics/BunchTuneAnalysis.hh @@ -1,7 +1,6 @@ #ifndef BUNCH_TUNE_ANALYSIS_H #define BUNCH_TUNE_ANALYSIS_H -//pyORBIT utils #include "CppPyWrapper.hh" #include "Bunch.hh" @@ -9,40 +8,37 @@ using namespace std; -/** - The BunchTuneAnalysis class calculates the particle tunes -*/ +/** Estimates particle tunes using average phase advance (APA) over one turn. */ class BunchTuneAnalysis: public OrbitUtils::CppPyWrapper { public: - /** Constructor*/ BunchTuneAnalysis(); /** Destructor */ virtual ~BunchTuneAnalysis(); - /** Performs the Twiss analysis of the bunch */ + /** Estimates tunes. */ void analyzeBunch(Bunch* bunch); - //** Assigns Twiss values at location of calculator */ - void assignTwiss(double bx, double ax, double dx, double dpx, double by, double ay); + /** Sets element of normalization matrix. */ + void setNormMatrixElement(int i, int j, double value); - /** Returns the average value for coordinate with index ic */ - double getTune(int ic); + /** Returns element of normalization matrix. */ + double getNormMatrixElement(int i, int j); + /** Sets normalization matrix based on uncoupled Twiss parameters. */ + void setNormMatrixFromTwiss(double betax, double alphax, double etax, double etapx, double betay, double alphay); private: - //** Twiss */ - double betax; - double alphax; - double etax; - double etapx; - double betay; - double alphay; + // Normalization matrix V^{-1} + double matrix[6][6]; + // Flag to erase stored phase/amplitude info. + int erase; }; + #endif //endif for BUNCH_TUNE_ANALYSIS_H diff --git a/src/orbit/BunchDiagnostics/wrap_bunch_tune_analysis.cc b/src/orbit/BunchDiagnostics/wrap_bunch_tune_analysis.cc index cce58059..3adeff0a 100644 --- a/src/orbit/BunchDiagnostics/wrap_bunch_tune_analysis.cc +++ b/src/orbit/BunchDiagnostics/wrap_bunch_tune_analysis.cc @@ -16,134 +16,162 @@ namespace wrap_bunch_tune_analysis{ extern "C" { #endif - /** - Constructor for python class wrapping c++ BunchTuneAnalysis instance. - It never will be called directly. - */ - static PyObject* BunchTuneAnalysis_new(PyTypeObject *type, PyObject *args, PyObject *kwds){ - pyORBIT_Object* self; - self = (pyORBIT_Object *) type->tp_alloc(type, 0); - self->cpp_obj = NULL; - return (PyObject *) self; + +/** +Constructor for Python class wrapping C++ BunchTuneAnalysis instance. +It never will be called directly. +*/ +static PyObject* BunchTuneAnalysis_new(PyTypeObject *type, PyObject *args, PyObject *kwds){ + pyORBIT_Object* self; + self = (pyORBIT_Object *) type->tp_alloc(type, 0); + self->cpp_obj = NULL; + return (PyObject *) self; +} + +/** Implementation of the __init__ method */ +static int BunchTuneAnalysis_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds){ + self->cpp_obj = new BunchTuneAnalysis(); + ((BunchTuneAnalysis*) self->cpp_obj)->setPyWrapper((PyObject*) self); + return 0; +} + +/** Analyzes the bunch. */ +static PyObject* BunchTuneAnalysis_analyzeBunch(PyObject *self, PyObject *args){ + BunchTuneAnalysis* cpp_BunchTuneAnalysis = (BunchTuneAnalysis*)((pyORBIT_Object*) self)->cpp_obj; + PyObject* pyBunch; + if(!PyArg_ParseTuple(args,"O:analyzeBunch", &pyBunch)){ + ORBIT_MPI_Finalize("BunchTuneAnalysis - analyzeBunch(Bunch* bunch) - parameter are needed."); + } + PyObject* pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); + if(!PyObject_IsInstance(pyBunch,pyORBIT_Bunch_Type)){ + ORBIT_MPI_Finalize("BunchTuneAnalysis - analyzeBunch(Bunch* bunch) - method needs a Bunch."); } + Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object*)pyBunch)->cpp_obj; + cpp_BunchTuneAnalysis->analyzeBunch(cpp_bunch); + Py_INCREF(Py_None); + return Py_None; +} - /** This is implementation of the __init__ method */ - static int BunchTuneAnalysis_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds){ - self->cpp_obj = new BunchTuneAnalysis(); - ((BunchTuneAnalysis*) self->cpp_obj)->setPyWrapper((PyObject*) self); - return 0; - } - - /** Performs the Tune analysis of the bunch */ - static PyObject* BunchTuneAnalysis_analyzeBunch(PyObject *self, PyObject *args){ - BunchTuneAnalysis* cpp_BunchTuneAnalysis = (BunchTuneAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - PyObject* pyBunch; - if(!PyArg_ParseTuple(args,"O:analyzeBunch",&pyBunch)){ - ORBIT_MPI_Finalize("BunchTuneAnalysis - analyzeBunch(Bunch* bunch) - parameter are needed."); - } - PyObject* pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); - if(!PyObject_IsInstance(pyBunch,pyORBIT_Bunch_Type)){ - ORBIT_MPI_Finalize("BunchTuneAnalysis - analyzeBunch(Bunch* bunch) - method needs a Bunch."); - } - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object*)pyBunch)->cpp_obj; - cpp_BunchTuneAnalysis->analyzeBunch(cpp_bunch); - Py_INCREF(Py_None); - return Py_None; - } - - /** Performs the Tune analysis of the bunch */ - static PyObject* BunchTuneAnalysis_assignTwiss(PyObject *self, PyObject *args){ - BunchTuneAnalysis* cpp_BunchTuneAnalysis = (BunchTuneAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - double betax; - double alphax; - double etax; - double etapx; - double betay; - double alphay; - if(!PyArg_ParseTuple(args,"dddddd:assignTwiss",&betax, &alphax,&etax,&etapx,&betay,&alphay)){ - ORBIT_MPI_Finalize("BunchTuneAnalysis - getTwiss(double betax, double alphax, double etax, double etapx, double betay, double alphay) - parameter are needed."); - } - cpp_BunchTuneAnalysis->assignTwiss(betax, alphax, etax, etapx, betay, alphay); - Py_INCREF(Py_None); - return Py_None; +/** Sets normalization matrix element. */ +static PyObject* BunchTuneAnalysis_setNormMatrixElement(PyObject *self, PyObject *args){ + BunchTuneAnalysis* cpp_BunchTuneAnalysis = (BunchTuneAnalysis*)((pyORBIT_Object*) self)->cpp_obj; + double value; + int i; + int j; + if(!PyArg_ParseTuple(args,"iid:setNormMatrixElement", &i, &j, &value)){ + ORBIT_MPI_Finalize("BunchTuneAnalysis - setNormMatrixElement(int i, int j, double value) - parameters are needed."); } + cpp_BunchTuneAnalysis->setNormMatrixElement(i, j, value); + Py_INCREF(Py_None); + return Py_None; +} +/** Gets normalization matrix element. */ +static PyObject* BunchTuneAnalysis_getNormMatrixElement(PyObject *self, PyObject *args){ + BunchTuneAnalysis* cpp_BunchTuneAnalysis = (BunchTuneAnalysis*)((pyORBIT_Object*) self)->cpp_obj; + int i; + int j; + if(!PyArg_ParseTuple(args,"ii:getNormMatrixElement", &i, &j)){ + ORBIT_MPI_Finalize("BunchTuneAnalysis - getNormMatrixElement(int i, int j) - parameters are needed."); + } + double value = cpp_BunchTuneAnalysis->getNormMatrixElement(i, j); + return Py_BuildValue("d", value); + Py_INCREF(Py_None); + return Py_None; +} - //-------------------------------------------------------------- - //destructor for python BunchTuneAnalysis class (__del__ method). - //--------------------------------------------------------------- - static void BunchTuneAnalysis_del(pyORBIT_Object* self){ - //std::cerr<<"The BunchTuneAnalysis __del__ has been called!"<cpp_obj); - self->ob_base.ob_type->tp_free((PyObject*)self); - } - - // defenition of the methods of the python BunchTuneAnalysis wrapper class - // they will be vailable from python level - static PyMethodDef BunchTuneAnalysisClassMethods[] = { - { "analyzeBunch", BunchTuneAnalysis_analyzeBunch, METH_VARARGS,"Performs the Tune analysis of the bunch."}, - { "assignTwiss", BunchTuneAnalysis_assignTwiss, METH_VARARGS,"Assigns Twiss at location of tune calculator."}, - {NULL} - }; - - // defenition of the memebers of the python BunchTwissAnalysis wrapper class - // they will be vailable from python level - static PyMemberDef BunchTuneAnalysisClassMembers [] = { - {NULL} - }; - - //new python BunchTwissAnalysis wrapper type definition - static PyTypeObject pyORBIT_BunchTuneAnalysis_Type = { - PyVarObject_HEAD_INIT(NULL, 0) - "BunchTuneAnalysis", /*tp_name*/ - sizeof(pyORBIT_Object), /*tp_basicsize*/ - 0, /*tp_itemsize*/ - (destructor) BunchTuneAnalysis_del , /*tp_dealloc*/ - 0, /*tp_print*/ - 0, /*tp_getattr*/ - 0, /*tp_setattr*/ - 0, /*tp_compare*/ - 0, /*tp_repr*/ - 0, /*tp_as_number*/ - 0, /*tp_as_sequence*/ - 0, /*tp_as_mapping*/ - 0, /*tp_hash */ - 0, /*tp_call*/ - 0, /*tp_str*/ - 0, /*tp_getattro*/ - 0, /*tp_setattro*/ - 0, /*tp_as_buffer*/ - Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ - "The BunchTuneAnalysis python wrapper", /* tp_doc */ - 0, /* tp_traverse */ - 0, /* tp_clear */ - 0, /* tp_richcompare */ - 0, /* tp_weaklistoffset */ - 0, /* tp_iter */ - 0, /* tp_iternext */ - BunchTuneAnalysisClassMethods, /* tp_methods */ - BunchTuneAnalysisClassMembers, /* tp_members */ - 0, /* tp_getset */ - 0, /* tp_base */ - 0, /* tp_dict */ - 0, /* tp_descr_get */ - 0, /* tp_descr_set */ - 0, /* tp_dictoffset */ - (initproc) BunchTuneAnalysis_init, /* tp_init */ - 0, /* tp_alloc */ - BunchTuneAnalysis_new, /* tp_new */ - }; - - - //-------------------------------------------------- - //Initialization BunchTwissAnalysis of the pyBunchTwissAnalysis class - //-------------------------------------------------- - void initbunchtuneanalysis(PyObject* module){ - if (PyType_Ready(&pyORBIT_BunchTuneAnalysis_Type) < 0) return; - Py_INCREF(&pyORBIT_BunchTuneAnalysis_Type); - PyModule_AddObject(module, "BunchTuneAnalysis", (PyObject *)&pyORBIT_BunchTuneAnalysis_Type); +/** Sets normalization matrix based on uncoupled Twiss parameters. */ +static PyObject* BunchTuneAnalysis_setNormMatrixFromTwiss(PyObject *self, PyObject *args){ + BunchTuneAnalysis* cpp_BunchTuneAnalysis = (BunchTuneAnalysis*)((pyORBIT_Object*) self)->cpp_obj; + double betax; + double alphax; + double etax; + double etapx; + double betay; + double alphay; + double etay; + double etapy; + if(!PyArg_ParseTuple(args,"dddddd:setNormMatrixFromTwiss",&betax, &alphax, &etax, &etapx, &betay, &alphay)){ + ORBIT_MPI_Finalize("BunchTuneAnalysis - getTwiss(double betax, double alphax, double etax, double etapx, double betay, double alphay) - parameter are needed."); } + cpp_BunchTuneAnalysis->setNormMatrixFromTwiss(betax, alphax, etax, etapx, betay, alphay); + Py_INCREF(Py_None); + return Py_None; +} + +/** Destructor for python BunchTuneAnalysis class (__del__ method). */ +static void BunchTuneAnalysis_del(pyORBIT_Object* self){ + //std::cerr<<"The BunchTuneAnalysis __del__ has been called!"<cpp_obj); + self->ob_base.ob_type->tp_free((PyObject*)self); +} + +// Definition of the methods of the python BunchTuneAnalysis wrapper class. +// They will be available from python level. +static PyMethodDef BunchTuneAnalysisClassMethods[] = { + { "analyzeBunch", BunchTuneAnalysis_analyzeBunch, METH_VARARGS,"Analyzes the bunch."}, + { "setNormMatrixElement", BunchTuneAnalysis_setNormMatrixElement, METH_VARARGS,"Sets normalization matrix element."}, + { "getNormMatrixElement", BunchTuneAnalysis_getNormMatrixElement, METH_VARARGS,"Gets normalization matrix element."}, + { "setNormMatrixFromTwiss", BunchTuneAnalysis_setNormMatrixFromTwiss, METH_VARARGS,"Sets normalization matrix based on uncoupled Twiss parameters."}, + {NULL} +}; + +// Definition of the memebers of the python BunchTuneAnalysis wrapper class. +// They will be vailable from python level. +static PyMemberDef BunchTuneAnalysisClassMembers [] = { + {NULL} +}; + +// New python BunchTuneAnalysis wrapper type definition. +static PyTypeObject pyORBIT_BunchTuneAnalysis_Type = { + PyVarObject_HEAD_INIT(NULL, 0) + "BunchTuneAnalysis", /*tp_name*/ + sizeof(pyORBIT_Object), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + (destructor) BunchTuneAnalysis_del , /*tp_dealloc*/ + 0, /*tp_print*/ + 0, /*tp_getattr*/ + 0, /*tp_setattr*/ + 0, /*tp_compare*/ + 0, /*tp_repr*/ + 0, /*tp_as_number*/ + 0, /*tp_as_sequence*/ + 0, /*tp_as_mapping*/ + 0, /*tp_hash */ + 0, /*tp_call*/ + 0, /*tp_str*/ + 0, /*tp_getattro*/ + 0, /*tp_setattro*/ + 0, /*tp_as_buffer*/ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ + "The BunchTuneAnalysis python wrapper", /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + BunchTuneAnalysisClassMethods, /* tp_methods */ + BunchTuneAnalysisClassMembers, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + (initproc) BunchTuneAnalysis_init, /* tp_init */ + 0, /* tp_alloc */ + BunchTuneAnalysis_new, /* tp_new */ +}; + +// Initialization of the pyBunchTuneAnalysis class. +void initbunchtuneanalysis(PyObject* module){ + if (PyType_Ready(&pyORBIT_BunchTuneAnalysis_Type) < 0) return; + Py_INCREF(&pyORBIT_BunchTuneAnalysis_Type); + PyModule_AddObject(module, "BunchTuneAnalysis", (PyObject *)&pyORBIT_BunchTuneAnalysis_Type); +} + #ifdef __cplusplus } diff --git a/src/orbit/BunchDiagnostics/wrap_bunch_tune_analysis.hh b/src/orbit/BunchDiagnostics/wrap_bunch_tune_analysis.hh index af0f0c0a..6511745e 100644 --- a/src/orbit/BunchDiagnostics/wrap_bunch_tune_analysis.hh +++ b/src/orbit/BunchDiagnostics/wrap_bunch_tune_analysis.hh @@ -7,7 +7,7 @@ extern "C" { #endif - namespace wrap_bunch_tune_analysis{ + namespace wrap_bunch_tune_analysis { void initbunchtuneanalysis(PyObject* module); } diff --git a/src/orbit/SyncPart.cc b/src/orbit/SyncPart.cc index 6494ecbf..320f9850 100644 --- a/src/orbit/SyncPart.cc +++ b/src/orbit/SyncPart.cc @@ -320,6 +320,10 @@ double SyncPart::getMass(){ return bunch->getMass(); } +double SyncPart::getCharge(){ + return bunch->getCharge(); +} + void SyncPart::readSyncPart(const char* fileName){ //for MPI diff --git a/src/orbit/SyncPart.hh b/src/orbit/SyncPart.hh index d14aa9ba..2081e4bf 100644 --- a/src/orbit/SyncPart.hh +++ b/src/orbit/SyncPart.hh @@ -68,6 +68,11 @@ class SyncPart */ double getMass(); + /** + Charge in elementary charge units + */ + double getCharge(); + /** time in seconds */ diff --git a/src/orbit/wrap_syncpart.cc b/src/orbit/wrap_syncpart.cc index 787c00a7..419ae556 100644 --- a/src/orbit/wrap_syncpart.cc +++ b/src/orbit/wrap_syncpart.cc @@ -129,6 +129,20 @@ extern "C" { return Py_BuildValue("d",val); } + // charge() - returns charge in elementary charge units + static PyObject* SyncPart_charge(PyObject *self, PyObject *args){ + pyORBIT_Object* pySyncPart = (pyORBIT_Object*) self; + int nVars = PyTuple_Size(args); + if(nVars > 0){ + error("PySyncPart - charge() - you should use bunch.charge(charge_value) instead!"); + } + double val = 0.; + SyncPart* cpp_SyncPart = (SyncPart*) pySyncPart->cpp_obj; + val = cpp_SyncPart->getCharge(); + return Py_BuildValue("d",val); + } + + //Sets or returns the momentum for the SyncPart object // the action is depended on the number of arguments // momentum() - returns momentum @@ -543,37 +557,38 @@ extern "C" { return Py_BuildValue("d",val); } - // defenition of the methods of the python SyncPart wrapper class - // they will be vailable from python level + // Definition of the methods of the python SyncPart wrapper class. + // They will be available from the Python level. static PyMethodDef SyncPartClassMethods[] = { - { "mass", SyncPart_mass ,METH_VARARGS,"Returns mass in GeV"}, - { "momentum", SyncPart_momentum ,METH_VARARGS,"Returns or sets momentum in GeV/c."}, - { "beta", SyncPart_beta ,METH_VARARGS,"Returns beta=v/c"}, - { "gamma", SyncPart_gamma ,METH_VARARGS,"Returns gamma=1/sqrt(1-(v/c)**2)"}, - { "kinEnergy", SyncPart_kinEnergy ,METH_VARARGS,"Returns or sets kinetic energy of the synchronous particle in MeV"}, - { "time", SyncPart_time ,METH_VARARGS,"Sets or returns time in sec"}, - { "x", SyncPart_x ,METH_VARARGS,"Sets or returns the x-coordinate"}, - { "y", SyncPart_y ,METH_VARARGS,"Sets or returns the y-coordinate"}, - { "z", SyncPart_z ,METH_VARARGS,"Sets or returns the z-coordinate"}, - { "px", SyncPart_px ,METH_VARARGS,"Sets or returns the x-momentum"}, - { "py", SyncPart_py ,METH_VARARGS,"Sets or returns the y-momentum"}, - { "pz", SyncPart_pz ,METH_VARARGS,"Sets or returns the z-momentum"}, - { "pVector", SyncPart_pVector ,METH_VARARGS,"Sets or returns the momentum vector as a tuple"}, - { "rVector", SyncPart_rVector ,METH_VARARGS,"Sets or returns the position vector as a tuple"}, - { "nxVector", SyncPart_nxVector ,METH_VARARGS,"Sets or returns the x-axis vector as a tuple"}, - { "nyVector", SyncPart_nyVector ,METH_VARARGS,"Returns the y-axis vector as a tuple"}, - { "energyToMomentum", SyncPart_eToP ,METH_VARARGS,"Transforms the kinetic energy to momentum"}, - { "momentumToEnergy", SyncPart_pToE ,METH_VARARGS,"Transforms the momentum to kinetic energy"}, + { "mass", SyncPart_mass, METH_VARARGS,"Returns mass in GeV"}, + { "charge", SyncPart_charge, METH_VARARGS,"Returns charge in elementary charge units"}, + { "momentum", SyncPart_momentum, METH_VARARGS,"Returns or sets momentum in GeV/c."}, + { "beta", SyncPart_beta, METH_VARARGS,"Returns beta=v/c"}, + { "gamma", SyncPart_gamma, METH_VARARGS,"Returns gamma=1/sqrt(1-(v/c)**2)"}, + { "kinEnergy", SyncPart_kinEnergy, METH_VARARGS,"Returns or sets kinetic energy of the synchronous particle in MeV"}, + { "time", SyncPart_time, METH_VARARGS,"Sets or returns time in sec"}, + { "x", SyncPart_x, METH_VARARGS,"Sets or returns the x-coordinate"}, + { "y", SyncPart_y, METH_VARARGS,"Sets or returns the y-coordinate"}, + { "z", SyncPart_z, METH_VARARGS,"Sets or returns the z-coordinate"}, + { "px", SyncPart_px, METH_VARARGS,"Sets or returns the x-momentum"}, + { "py", SyncPart_py, METH_VARARGS,"Sets or returns the y-momentum"}, + { "pz", SyncPart_pz, METH_VARARGS,"Sets or returns the z-momentum"}, + { "pVector", SyncPart_pVector, METH_VARARGS,"Sets or returns the momentum vector as a tuple"}, + { "rVector", SyncPart_rVector, METH_VARARGS,"Sets or returns the position vector as a tuple"}, + { "nxVector", SyncPart_nxVector, METH_VARARGS,"Sets or returns the x-axis vector as a tuple"}, + { "nyVector", SyncPart_nyVector, METH_VARARGS,"Returns the y-axis vector as a tuple"}, + { "energyToMomentum", SyncPart_eToP, METH_VARARGS,"Transforms the kinetic energy to momentum"}, + { "momentumToEnergy", SyncPart_pToE, METH_VARARGS,"Transforms the momentum to kinetic energy"}, {NULL} }; - // defenition of the memebers of the python SyncPart wrapper class - // they will be vailable from python level + // Definition of the members of the python SyncPart wrapper class. + // They will be available from python level. static PyMemberDef SyncPartClassMembers [] = { {NULL} }; - //new python SyncPart wrapper type definition + // New Python SyncPart wrapper type definition. static PyTypeObject pyORBIT_SyncPart_Type = { PyVarObject_HEAD_INIT(NULL, 0) "SyncParticle", /*tp_name*/ diff --git a/tests/py/orbit/test_env.py b/tests/py/orbit/test_env.py index 5275ab6c..337c18db 100644 --- a/tests/py/orbit/test_env.py +++ b/tests/py/orbit/test_env.py @@ -1,11 +1,11 @@ import numpy as np +import pytest from orbit.core.bunch import Bunch from orbit.core.bunch import BunchTwissAnalysis from orbit.core.linac import MatrixRfGap from orbit.bunch_utils import collect_bunch from orbit.envelope import Envelope -from orbit.envelope import EnvelopeTracker from orbit.lattice import AccNode from orbit.lattice import AccLattice from orbit.py_linac.lattice import Drift @@ -64,6 +64,7 @@ def track_and_compare_rms( kin_energy: float, cov_matrix: np.ndarray, nparts: int = 100_000, + charge: float = 1.0, verbose: int = 1, ) -> dict: """Track bunch/envelope and compare rms beam sizes. @@ -82,13 +83,16 @@ def track_and_compare_rms( data[k1] = {} for k2 in ["rms", "cov"]: data[k1][k2] = {} - for k3 in ["env", "bunch"]: + for k3 in ["in", "out"]: data[k1][k2][k3] = {} # Initialize bunch bunch = Bunch() bunch.mass(mass_proton) - bunch.getSyncParticle().kinEnergy(kin_energy) + bunch.charge(charge) + + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(kin_energy) # Track bunch particles = np.random.multivariate_normal(np.zeros(6), cov_matrix, size=nparts) @@ -103,11 +107,10 @@ def track_and_compare_rms( data["bunch"]["cov"]["out"] = cov_scale * calc_bunch_cov(bunch) # Track envelope - envelope = Envelope(bunch=bunch, cov_matrix=cov_matrix) - envelope_tracker = EnvelopeTracker(lattice=lattice) + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix) data["env"]["cov"]["in"] = cov_scale * envelope.cov_matrix - envelope_tracker.track(envelope) + lattice.trackEnvelope(envelope) data["env"]["cov"]["out"] = cov_scale * envelope.cov_matrix # Compare @@ -148,196 +151,125 @@ def make_default_cov_matrix( return np.diag(np.square([rms_x, rms_xp, rms_y, rms_yp, rms_z, rms_dE])) -def test_drift_teapot( - kin_energy: float = 0.0025, - length: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 6, -) -> None: - node = DriftTEAPOT(length=length, nparts=nparts) +def test_drift_teapot(): + node = DriftTEAPOT(length=1.0, nparts=6) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix) -def test_drift_linac( - kin_energy: float = 0.0025, - length: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 6, -) -> None: +def test_drift_linac(): node = Drift() - node.setLength(length) - node.setnParts(nparts) + node.setLength(1.0) + node.setnParts(6) nodes = [node] lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_quad_teapot( - kin_energy: float = 0.0025, - length: float = 1.0, - kq: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 10, -) -> None: - node = QuadTEAPOT(length=length, kq=kq, nparts=nparts) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_quad_teapot(charge: float): + node = QuadTEAPOT(length=1.0, kq=1.0, nparts=10) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_cf_teapot( - kin_energy: float = 0.0025, - length: float = 10.0, - kq: float = 1.0, - nparts: int = 10, -) -> None: - node = ContinuousLinearFocusingTEAPOT(length=length, kq=kq, nparts=nparts) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_cf_teapot(charge: float): + node = ContinuousLinearFocusingTEAPOT(length=10.0, kq=1.0, nparts=10) lattice = make_lattice([node]) - cov_matrix = np.diag(np.square([1e-3, 0, 1e-3, 0.0, 0.0, 0.0])) - track_and_compare_rms(lattice, kin_energy, cov_matrix) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) -def test_quad_linac( - kin_energy: float = 0.0025, - length: float = 1.0, - field_grad: float = 0.23, - cov_matrix: np.ndarray = None, - nparts: int = 10, -) -> None: +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_quad_linac(charge: float): node = Quad() - node.setLength(length) - node.setnParts(nparts) - node.setParam("dB/dr", field_grad) + node.setLength(1.0) + node.setnParts(10) + node.setParam("dB/dr", 0.23) nodes = [node] lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_bend_teapot( - kin_energy: float = 0.0025, - length: float = 1.0, - theta: float = 20.0, - cov_matrix: np.ndarray = None, - nparts: int = 5, -) -> None: - node = BendTEAPOT(length=length, theta=np.radians(theta), nparts=nparts) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_bend_teapot(charge: float): + node = BendTEAPOT(length=1.0, theta=np.radians(20.0), nparts=5) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_bend_linac( - kin_energy: float = 0.0025, - length: float = 1.0, - theta: float = 20.0, - cov_matrix: np.ndarray = None, - nparts: int = 5, -) -> None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_bend_linac(charge: float): node = Bend() - node.setLength(length) - node.setnParts(nparts) - node.setParam("theta", np.radians(theta)) + node.setLength(1.0) + node.setnParts(5) + node.setParam("theta", np.radians(20.0)) nodes = [node] lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_kick_teapot( - kin_energy: float = 0.0025, - length: float = 0.1, - kx: float = 0.001, - ky: float = 0.001, - dE: float = 0.00001, - cov_matrix: np.ndarray = None, - nparts: int = 4, -) -> None: - node = KickTEAPOT(kx=kx, ky=ky, dE=dE, length=length, nparts=nparts) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_kick_teapot(charge: float): + node = KickTEAPOT(kx=0.001, ky=0.001, dE=0.00001, length=0.1, nparts=4) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) -def test_tilt_teapot( - kin_energy: float = 0.0025, - angle: float = 0.25 * np.pi, - cov_matrix: np.ndarray = None, -) -> None: - node = TiltTEAPOT(angle=angle) +def test_tilt_teapot(): + node = TiltTEAPOT(angle=(0.25 * np.pi)) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix) -def test_tilt_linac( - kin_energy: float = 0.0025, - angle: float = 0.25 * np.pi, - cov_matrix: np.ndarray = None, -) -> None: +def test_tilt_linac(): node = TiltElement() - node.setTiltAngle(angle) - nodes = [node] - lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_solenoid_teapot( - kin_energy: float = 0.0025, - length: float = 2.0, - B: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 10, -) -> None: - node = SolenoidTEAPOT(length=length, B=B, nparts=nparts) + node.setTiltAngle(0.25 * np.pi) lattice = make_lattice([node]) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_solenoid_linac( - kin_energy: float = 0.0025, - length: float = 2.0, - B: float = 1.0, - cov_matrix: np.ndarray = None, - nparts: int = 10, -) -> None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_solenoid_teapot(charge: float): + node = SolenoidTEAPOT(length=2.0, B=1.0, nparts=10) + lattice = make_lattice([node]) + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_solenoid_linac(charge: float): node = Solenoid() - node.setLength(length) - node.setnParts(nparts) - node.setParam("B", B) + node.setLength(2.0) + node.setnParts(10) + node.setParam("B", 1.0) nodes = [node] lattice = make_lattice(nodes) - if cov_matrix is None: - cov_matrix = make_default_cov_matrix() - track_and_compare_rms(lattice, kin_energy, cov_matrix) - - -def test_rf_gap_matrix( - kin_energy: float = 0.0025, - frequency: float = 402.5e06, - E0TL: float = 0.001, - phase: float = 0.0, - charge: float = -1.0, -) -> None: + cov_matrix = make_default_cov_matrix() + track_and_compare_rms(lattice, kin_energy=0.0025, cov_matrix=cov_matrix, charge=charge) + + +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_rf_gap_matrix(charge: float): + kin_energy = 0.0025 + frequency = 402.5e06 + E0TL = 0.001 + phase = 0.0 + cov_matrix = make_default_cov_matrix() bunch_in = Bunch() @@ -357,14 +289,14 @@ def test_rf_gap_matrix( coords_out_1 = collect_bunch(bunch_out_1)["coords"] - from orbit.envelope.matrix import get_matrix_rf_gap + from orbit.utils.matrix import get_matrix_rf_gap bunch_out_2 = Bunch() bunch_in.copyBunchTo(bunch_out_2) - envelope = Envelope(bunch=bunch_in) + sync_part = bunch_in.getSyncParticle() matrix = get_matrix_rf_gap( - envelope=envelope, + sync_part, frequency=frequency, E0TL=E0TL, phase=phase, @@ -386,10 +318,12 @@ def test_sc_3d_cold_expansion(): def test_track_sublattice_no_error(): bunch = Bunch() bunch.mass(mass_proton) - bunch.getSyncParticle().kinEnergy(0.001) + + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(0.001) cov_matrix = np.diag(np.square([1e-3, 0, 1e-3, 0.0, 1e-3, 0.0])) - envelope = Envelope(bunch, cov_matrix=cov_matrix) + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix) lattice = TEAPOT_Lattice() @@ -397,28 +331,30 @@ def test_track_sublattice_no_error(): for _ in range(n): lattice.addNode(DriftTEAPOT(length=0.1)) - tracker = EnvelopeTracker(lattice) for i in range(n): - tracker.track(envelope, index_start=i) - tracker.track(envelope, index_stop=-i) + lattice.trackEnvelope(envelope, index_start=i) + lattice.trackEnvelope(envelope, index_stop=-i) + -def test_get_total_matrix() -> None: +@pytest.mark.parametrize("charge", [1.0, -1.0]) +def test_get_total_matrix(charge: float) -> None: node = DriftTEAPOT(length=2.0, nparts=50) lattice = make_lattice([node]) bunch = Bunch() bunch.mass(mass_proton) - bunch.getSyncParticle().kinEnergy(0.001) + bunch.charge(charge) - cov_matrix = make_default_cov_matrix() - envelope = Envelope(bunch, cov_matrix=cov_matrix, intensity=1e7) + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(0.001) - tracker = EnvelopeTracker(lattice, sc="2d") + cov_matrix = make_default_cov_matrix() + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix, intensity=1e7) envelope_out_a = envelope.copy() - tracker.track(envelope_out_a) + lattice.trackEnvelope(envelope_out_a, sc="2d") - matrix = tracker.get_transfer_matrix(envelope.copy()) + matrix = lattice.getEnvelopeTransferMatrix(envelope.copy(), sc="2d") envelope_out_b = envelope.copy() envelope_out_b.transform(matrix) assert np.all(np.isclose(envelope_out_a.cov_matrix, envelope_out_b.cov_matrix)) diff --git a/tests/py/orbit/test_sync_part.py b/tests/py/orbit/test_sync_part.py new file mode 100644 index 00000000..f0b0ee7b --- /dev/null +++ b/tests/py/orbit/test_sync_part.py @@ -0,0 +1,18 @@ +from orbit.core.bunch import Bunch +from orbit.core.bunch import SyncParticle + + +def make_bunch(): + bunch = Bunch() + bunch.mass(0.938272) + bunch.charge(1.0) + bunch.getSyncParticle().kinEnergy(1.0) + return bunch + + +def test_get_mass_charge(): + bunch = make_bunch() + sync_part = bunch.getSyncParticle() + + assert sync_part.mass() == bunch.mass() + assert sync_part.charge() == bunch.charge()