diff --git a/.gitignore b/.gitignore index 47650c0e..1583fd4f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +.DS_Store +.idea build PyORBIT.egg-info .vscode diff --git a/examples/Envelope/inputs/sns_ring.lat b/examples/Envelope/inputs/sns_ring.lat new file mode 100644 index 00000000..cbfb8df1 --- /dev/null +++ b/examples/Envelope/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/Envelope/plot.py b/examples/Envelope/plot.py new file mode 100644 index 00000000..99a77f80 --- /dev/null +++ b/examples/Envelope/plot.py @@ -0,0 +1,133 @@ +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.patches as patches + + +def calc_rms_ellipse_params(cov_matrix: np.ndarray) -> tuple[float, float, float]: + """Return rms ellipse dimensions and orientation.""" + i, j = (0, 1) + + sii = cov_matrix[i, i] + sjj = cov_matrix[j, j] + sij = cov_matrix[i, j] + + angle = -0.5 * np.arctan2(2.0 * sij, sii - sjj) + + _sin = np.sin(angle) + _cos = np.cos(angle) + _sin2 = _sin**2 + _cos2 = _cos**2 + + c1 = np.sqrt(abs(sii * _cos2 + sjj * _sin2 - 2 * sij * _sin * _cos)) + c2 = np.sqrt(abs(sii * _sin2 + sjj * _cos2 + 2 * sij * _sin * _cos)) + + return (c1, c2, angle) + + +def plot_ellipse( + r1: float = 1.0, + r2: float = 1.0, + angle: float = 0.0, + center: tuple[float, float] = None, + ax=None, + **kws, +): + kws.setdefault("fill", False) + kws.setdefault("color", "black") + kws.setdefault("lw", 1.25) + + if center is None: + center = (0.0, 0.0) + + d1 = r1 * 2.0 + d2 = r2 * 2.0 + angle = -np.degrees(angle) + + ax.add_patch(patches.Ellipse(center, d1, d2, angle=angle, **kws)) + return ax + + +def plot_rms_ellipse( + cov_matrix: np.ndarray, + level: float = 1.0, + ax=None, + **ellipse_kws, +): + """Plot rms ellipse from 2 x 2 covariance matrix.""" + r1, r2, angle = calc_rms_ellipse_params(cov_matrix) + plot_ellipse(r1 * level, r2 * level, angle=angle, ax=ax, **ellipse_kws) + return ax + + +def plot_corner( + particles: np.ndarray, + limits: list[tuple[float, float]] = None, + bins: int = 64, + labels: list[str] = None, + blur: float = None, +) -> tuple: + """Generate corner plot.""" + ndim = particles.shape[1] + + if limits is None: + xmax = np.max(particles, axis=0) + xmin = np.min(particles, axis=0) + limits = list(zip(xmin, xmax)) + + if labels is None: + labels = ndim * [""] + + fig, axs = plt.subplots( + ncols=ndim, nrows=ndim, sharex=None, sharey=None, figsize=(8, 8) + ) + for i in range(ndim): + for j in range(ndim): + axis = (j, i) + ax = axs[i, j] + if i > j: + values, edges = np.histogramdd( + particles[:, axis], bins=bins, range=[limits[k] for k in axis] + ) + if blur: + values = scipy.ndimage.gaussian_filter(values, sigma=blur) + ax.pcolormesh( + edges[0], + edges[1], + values.T, + linewidth=0.0, + rasterized=True, + shading="auto", + ) + elif i == j: + values, edges = np.histogram( + particles[:, i], bins=bins, range=limits[i] + ) + if blur: + values = scipy.ndimage.gaussian_filter(values, sigma=blur) + ax.stairs(values, edges, lw=1.5, color="black") + else: + ax.axis("off") + + for i in range(0, ndim - 1): + for j in range(0, ndim): + axs[i, j].set_xticklabels([]) + for i in range(0, ndim): + for j in range(1, ndim): + axs[i, j].set_yticklabels([]) + + for ax in axs.flat: + for loc in ["top", "right"]: + ax.spines[loc].set_visible(False) + + for i, label in enumerate(labels): + axs[-1, i].set_xlabel(label) + for i, label in enumerate(labels[1:], start=1): + axs[i, 0].set_ylabel(label) + + axs[0, 0].set_yticklabels([]) + axs[0, 0].set_ylabel(None) + + fig.align_ylabels() + fig.align_xlabels() + + return fig, axs diff --git a/examples/Envelope/style.mplstyle b/examples/Envelope/style.mplstyle new file mode 100644 index 00000000..71f36605 --- /dev/null +++ b/examples/Envelope/style.mplstyle @@ -0,0 +1,8 @@ +axes.linewidth: 1.25 +axes.titlesize: "medium" +image.cmap: "Greys" +figure.constrained_layout.use: True +savefig.dpi: 300 +savefig.format: "png" +xtick.minor.visible: True +ytick.minor.visible: True diff --git a/examples/Envelope/test_env.py b/examples/Envelope/test_env.py new file mode 100644 index 00000000..3cb46a2d --- /dev/null +++ b/examples/Envelope/test_env.py @@ -0,0 +1,188 @@ +import numpy as np + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +from orbit.bunch_utils import collect_bunch +from orbit.lattice import AccNode +from orbit.lattice import AccLattice +from orbit.teapot import QuadTEAPOT +from orbit.teapot import BendTEAPOT +from orbit.teapot import DriftTEAPOT +from orbit.teapot import TEAPOT_Lattice +from orbit.utils.consts import mass_proton +from orbit.envelope import Envelope +from orbit.envelope import EnvelopeTracker + + +def calc_bunch_cov(bunch: Bunch) -> np.ndarray: + coords = collect_bunch(bunch)["coords"] + return np.cov(coords.T) + + # twiss_calc = BunchTwissAnalysis() + # twiss_calc.computeBunchMoments(bunch, 2, 0, 0) + # + # cov_matrix = np.zeros((6, 6)) + # for i in range(6): + # for j in range(i + 1): + # cov_matrix[i, j] = twiss_calc.getCorrelation(j, i) + # cov_matrix[j, i] = cov_matrix[i, j] + # return cov_matrix + + +def make_lattice(nodes: list[AccNode]) -> AccLattice: + lattice = TEAPOT_Lattice() + for node in nodes: + lattice.addNode(node) + lattice.initialize() + for node in lattice.getNodes(): + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + return lattice + + +def track_and_compare( + nodes: list[AccNode], + kin_energy: float, + cov_matrix: np.ndarray, + nparts: int = 100_000, +) -> dict: + + cov_scale = 1e6 + + data = {} + for k1 in ["env", "bunch"]: + data[k1] = {} + for k2 in ["rms", "cov"]: + data[k1][k2] = {} + for k3 in ["env", "bunch"]: + data[k1][k2][k3] = {} + + lattice = make_lattice(nodes) + + bunch = Bunch() + bunch.mass(mass_proton) + + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(kin_energy) + + envelope = Envelope(sync_part=sync_part, cov_matrix=cov_matrix) + + envelope_tracker = EnvelopeTracker(lattice=lattice) + + data["env"]["cov"]["in"] = cov_scale * envelope.cov() + envelope_tracker.track(envelope) + data["env"]["cov"]["out"] = cov_scale * envelope.cov() + + particles = np.random.multivariate_normal(np.zeros(6), cov_matrix, size=nparts) + for x in particles: + bunch.addParticle(*x) + + data["bunch"]["cov"]["in"] = cov_scale * calc_bunch_cov(bunch) + lattice.trackBunch(bunch) + data["bunch"]["cov"]["out"] = cov_scale * calc_bunch_cov(bunch) + + for mode in ["env", "bunch"]: + for loc in ["in", "out"]: + data[mode]["rms"][loc] = np.sqrt(np.diag(data[mode]["cov"][loc])) + + dims = ["x", "xp", "y", "yp", "z", "dE"] + for key in ["in", "out"]: + print(key.upper()) + for i in range(6): + print(" rms {}:".format(dims[i])) + print(" env: {}".format(data["env"]["rms"][key][i])) + print(" bunch: {}".format(data["bunch"]["rms"][key][i])) + + for key in ["in", "out"]: + assert np.all(np.isclose(data["env"]["cov"][key], data["bunch"]["cov"][key])) + + +def make_default_cov_matrix(scale: float = 0.001) -> np.ndarray: + cov_matrix = np.zeros((6, 6)) + cov_matrix[0, 0] = scale ** 2 + cov_matrix[2, 2] = scale ** 2 + return cov_matrix + + +def test_drift(kin_energy: float = 0.0025, length: float = 1.0, cov_matrix: np.ndarray = None): + nodes = [ + DriftTEAPOT(length=1.0), + ] + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare(nodes, kin_energy, cov_matrix) + +def test_quad(kin_energy: float = 0.0025, length: float = 1.0, kq: float = 1.0, cov_matrix: np.ndarray = None): + nodes = [ + QuadTEAPOT(length=length, kq=kq), + ] + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare(nodes, kin_energy, cov_matrix) + + +def test_dipole(kin_energy: float = 0.0025, length: float = 1.0, theta: float = 20.0, cov_matrix: np.ndarray = None): + nodes = [ + BendTEAPOT(length=length, theta=np.radians(theta)) + ] + if cov_matrix is None: + cov_matrix = make_default_cov_matrix() + track_and_compare(nodes, kin_energy, cov_matrix) + + +def test_dipole_matrix(): + from orbit.envelope.matrix import MatrixFactory + + bunch = Bunch() + bunch.mass(mass_proton) + + sync_part = bunch.getSyncParticle() + sync_part.kinEnergy(0.0025) + + node = BendTEAPOT(length=1.0, theta=np.radians(20.0)) + + matrix_factory = MatrixFactory() + matrix = matrix_factory.bend(length=node.getLength(), theta=node.getParam("theta"), sync_part=sync_part) + + print(np.round(matrix, 2)) + + x = np.zeros(6) + x[1] = 0.001 + x[1] = 0.001 + + bunch.addParticle(*x) + + node.trackBunch(bunch) + + x_out = [bunch.x(0), bunch.xp(0), bunch.y(0), bunch.yp(0), bunch.z(0), bunch.dE(0)] + x_out = np.array(x_out) + + print(x) + print(x_out) + print(np.matmul(matrix[:6, :6], x)) + + cov_matrix = make_default_cov_matrix() + envelope = Envelope(sync_part=bunch.getSyncParticle(), cov_matrix=cov_matrix) + envelope.apply_transfer_matrix(matrix) + print(np.round(1e6 * envelope.cov(), 2)) + + particles_in = np.random.multivariate_normal(np.zeros(6), cov_matrix, size=100_000) + particles_out = np.matmul(particles_in, matrix[:6, :6].T) + print(np.round(1e6 * np.cov(particles_out.T), 2)) + + bunch.deleteAllParticles() + for x in particles_in: + bunch.addParticle(*x) + node.trackBunch(bunch) + particles_out = collect_bunch(bunch)["coords"] + print(np.round(1e6 * np.cov(particles_out.T), 2)) + + +if __name__ == "__main__": + # test_drift() + # test_quad() + test_dipole() + # test_dipole_matrix() + + + diff --git a/examples/Envelope/test_env_2d_fodo.py b/examples/Envelope/test_env_2d_fodo.py new file mode 100644 index 00000000..8184f5bf --- /dev/null +++ b/examples/Envelope/test_env_2d_fodo.py @@ -0,0 +1,327 @@ +"""Test 2D envelope tracker in FODO lattice.""" + +import argparse +import copy +import math +import os +import pathlib + +import numpy as np +import matplotlib.pyplot as plt + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +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 +from orbit.space_charge.sc2p5d import setSC2p5DAccNodes +from orbit.teapot import DriftTEAPOT +from orbit.teapot import QuadTEAPOT +from orbit.teapot import TEAPOT_Lattice +from orbit.teapot import TEAPOT_MATRIX_Lattice +from orbit.utils.consts import mass_proton + +from plot import plot_rms_ellipse +from plot import plot_corner +from utils import gen_dist +from utils import build_rotation_matrix_xy +from utils import project_cov_matrix + +plt.style.use("style.mplstyle") + + +# Parse arguments +# ------------------------------------------------------------------------------ + +parser = argparse.ArgumentParser() +parser.add_argument("--zrms", type=float, default=5.0) +parser.add_argument("--kin-energy", type=float, default=0.0025) +parser.add_argument("--intensity", type=float, default=5e9) + +parser.add_argument( + "--dist", type=str, default="kv", choices=["kv", "waterbag", "gauss"] +) +parser.add_argument("--mismatch-x", type=float, default=0.0) +parser.add_argument("--mismatch-y", type=float, default=0.0) +parser.add_argument("--offset-x", type=float, default=0.0) +parser.add_argument("--offset-y", type=float, default=0.0) +parser.add_argument("--tilt", type=float, default=0) + +parser.add_argument("--nslice", type=int, default=10) +parser.add_argument("--kq", type=float, default=0.25) + +parser.add_argument("--nparts", type=int, default=100_000) +parser.add_argument("--turns", type=int, default=25) +parser.add_argument("--sc", type=int, default=0) +args = parser.parse_args() + + +# Setup +# ------------------------------------------------------------------------------ + +path = pathlib.Path(__file__) +output_dir = os.path.join("outputs", path.stem) +os.makedirs(output_dir, exist_ok=True) + + +# Create lattice +# ------------------------------------------------------------------------------ + +nodes = [ + QuadTEAPOT(length=0.5, kq=+args.kq), + DriftTEAPOT(length=1.0), + QuadTEAPOT(length=1.0, kq=-args.kq), + DriftTEAPOT(length=1.0), + QuadTEAPOT(length=0.5, kq=+args.kq), +] + +lattice = TEAPOT_Lattice() +for node in nodes: + node.setnParts(args.nslice) + node.setUsageFringeFieldIN(False) + node.setUsageFringeFieldOUT(False) + lattice.addNode(node) + +lattice.initialize() + + +# Create envelope +# ------------------------------------------------------------------------------ + +# Create bunch +bunch = Bunch() +bunch.mass(mass_proton) +sync_part = bunch.getSyncParticle() +sync_part.kinEnergy(args.kin_energy) + +# Find periodic lattice parameters +matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) +matrix_lattice_params = matrix_lattice.getRingParametersDict() +alpha_x = matrix_lattice_params["alpha x"] +alpha_y = matrix_lattice_params["alpha y"] +beta_x = matrix_lattice_params["beta x [m]"] +beta_y = matrix_lattice_params["beta y [m]"] +eps_x = 0.25e-06 +eps_y = eps_x + +# Generate covariance matrix +cov_matrix = np.zeros((6, 6)) +cov_matrix[0, 0] = eps_x * beta_x +cov_matrix[2, 2] = eps_y * beta_y +cov_matrix[0, 1] = cov_matrix[1, 0] = -eps_x * alpha_x +cov_matrix[2, 3] = cov_matrix[3, 2] = -eps_y * alpha_y +cov_matrix[1, 1] = eps_x * (1.0 + alpha_x**2) / beta_x +cov_matrix[3, 3] = eps_y * (1.0 + alpha_y**2) / beta_y +cov_matrix[4, 4] = args.zrms**2 +cov_matrix[5, 5] = 0.0 + +# Tilt +if args.tilt: + rot_matrix = np.identity(6) + rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi)) + cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T]) + +# Mismatch +cov_matrix[0, 0] *= (1.0 + args.mismatch_x) ** 2 +cov_matrix[2, 2] *= (1.0 + args.mismatch_y) ** 2 +cov_matrix_init = np.copy(cov_matrix) + +# Offset +centroid_init = np.zeros(6) +centroid_init[0] += args.offset_x +centroid_init[2] += args.offset_y + +# Create envelope +envelope = Envelope( + sync_part=sync_part, + cov_matrix=cov_matrix_init, + centroid=centroid_init, + intensity=args.intensity, +) + + +# Track envelope +# ------------------------------------------------------------------------------ + +print("TRACK ENVELOPE") + +tracker = EnvelopeTracker(lattice, space_charge=("2d" if args.sc else None)) + +history = {"xrms": [], "yrms": [], "xavg": [], "yavg": []} +for turn in range(args.turns): + if turn > 0: + tracker.track(envelope) + + cov_matrix = envelope.cov() + centroid = envelope.centroid() + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + xavg = 1000.0 * centroid[0] + yavg = 1000.0 * centroid[2] + + print( + f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} xavg={xavg:0.3f} yavg={yavg:0.3f}" + ) + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["xavg"].append(xavg) + history["yavg"].append(yavg) + +histories = {} +histories["envelope"] = copy.deepcopy(history) + + +# Track bunch +# ------------------------------------------------------------------------------ + +print("TRACK BUNCH") + +rng = np.random.default_rng() + +bunch_coords = np.zeros((args.nparts, 6)) +bunch_coords[:, :4] = gen_dist( + n=args.nparts, cov_matrix=cov_matrix_init[0:4, 0:4], name=args.dist +) +bunch_coords[:, 4] = 2.0 * rng.uniform(-args.zrms, args.zrms, size=args.nparts) +bunch_coords += centroid_init[None, :6] + +for i in range(bunch_coords.shape[0]): + bunch.addParticle(*bunch_coords[i]) + +if args.sc: + sc_calc = SpaceChargeCalc2p5D(128, 128, 1) + sc_path_length_min = 1.00e-06 + sc_nodes = setSC2p5DAccNodes(lattice, sc_path_length_min, sc_calc) + + bunch_size = bunch.getSizeGlobal() + bunch.macroSize(args.intensity / bunch_size) + +history = {"xrms": [], "yrms": [], "xavg": [], "yavg": []} +for turn in range(args.turns): + if turn > 0: + lattice.trackBunch(bunch) + + twiss_calc = BunchTwissAnalysis() + twiss_calc.computeBunchMoments(bunch, 2, 0, 0) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(i + 1): + cov_matrix[i, j] = twiss_calc.getCorrelation(j, i) + cov_matrix[j, i] = cov_matrix[i, j] + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + xavg = 1000.0 * twiss_calc.getAverage(0) + yavg = 1000.0 * twiss_calc.getAverage(2) + + print( + f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} xavg={xavg:0.3f} yavg={yavg:0.3f}" + ) + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["xavg"].append(xavg) + history["yavg"].append(yavg) + +histories["bunch"] = copy.deepcopy(history) + + +# Analysis +# ------------------------------------------------------------------------------ + +for history in histories.values(): + for key in history: + history[key] = np.array(history[key]) + +# Print errors +for key in histories["envelope"]: + deltas = histories["bunch"][key] - histories["envelope"][key] + print("key:", key) + print("max_abs_delta:", np.max(np.abs(deltas))) + print("avg_abs_delta:", np.mean(np.abs(deltas))) + +# Plot rms bunch sizes +for key in ["xrms", "yrms"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(0.0, ax.get_ylim()[1] * 2.0) + ax.set_xlabel("Turn") + ax.set_ylabel("RMS [mm]") + ax.legend(loc="upper right") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + +# Plot centroids +for key in ["xavg", "yavg"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(-5.0, 5.0) + ax.set_xlabel("Turn") + ax.set_ylabel("AVG [mm]") + ax.legend(loc="upper right") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + + +# Collect bunch/envelope data on final turn. +particles = collect_bunch(bunch)["coords"] +particles[:, :4] *= 1000.0 + +env_cov_matrix = envelope.cov() +env_cov_matrix[:4, :4] *= 1000.0**2 + +env_centroid = envelope.centroid() +env_centroid[:4] *= 1000.0 + +xmax = 4.0 * np.std(particles, axis=0) +limits = list(zip(-xmax, xmax)) +labels = ["x [mm]", "xp [mrad]", "y [mm]", "yp [mrad]", "z [m]", "dE [GeV]"] + + +# Plot x-x' +fig, ax = plt.subplots(figsize=(4, 4)) +ax.hist2d(particles[:, 0], particles[:, 1], bins=100, range=[limits[0], limits[1]]) +plot_rms_ellipse( + env_cov_matrix[0:2, 0:2], + center=(env_centroid[0], env_centroid[1]), + level=2.0, + color="red", + ax=ax, +) +ax.set_xlabel(labels[0]) +ax.set_ylabel(labels[1]) +plt.savefig(os.path.join(output_dir, "fig_dist_x_xp")) +plt.close() + +# Plot corner +fig, axs = plot_corner( + particles, + limits=limits, + bins=100, + labels=labels, +) +for i in range(6): + for j in range(i): + env_cov_matrix_proj = project_cov_matrix(env_cov_matrix, axis=(j, i)) + plot_rms_ellipse( + env_cov_matrix_proj, + center=(env_centroid[j], env_centroid[i]), + level=2.0, + color="red", + ax=axs[i, j], + ) +plt.savefig(os.path.join(output_dir, "fig_dist_corner")) +plt.close() diff --git a/examples/Envelope/test_env_3d_drift.py b/examples/Envelope/test_env_3d_drift.py new file mode 100644 index 00000000..5e344249 --- /dev/null +++ b/examples/Envelope/test_env_3d_drift.py @@ -0,0 +1,249 @@ +"""Test 3D envelope tracker in drift.""" + +import argparse +import copy +import math +import os +import pathlib + +import numpy as np +import matplotlib.pyplot as plt + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +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 +from orbit.utils.consts import mass_proton + +from plot import plot_rms_ellipse +from plot import plot_corner +from utils import gen_dist +from utils import project_cov_matrix + +plt.style.use("style.mplstyle") + + +# Parse arguments +# ------------------------------------------------------------------------------ + +parser = argparse.ArgumentParser() +parser.add_argument("--kin-energy", type=float, default=0.0025) +parser.add_argument("--intensity", type=float, default=5e10) + +parser.add_argument("--xrms", type=float, default=0.010) +parser.add_argument("--yrms", type=float, default=0.010) +parser.add_argument("--zrms", type=float, default=0.010) + +parser.add_argument("--nslice", type=int, default=10) +parser.add_argument("--length", type=float, default=0.1) +parser.add_argument("--turns", type=int, default=20) +parser.add_argument("--sc-grid", type=int, default=64) + +parser.add_argument("--nparts", type=int, default=100_000) +parser.add_argument("--sc", type=int, default=0) +args = parser.parse_args() + + +# Setup +# ------------------------------------------------------------------------------ + +path = pathlib.Path(__file__) +output_dir = os.path.join("outputs", path.stem) +os.makedirs(output_dir, exist_ok=True) + + +# Create lattice +# ------------------------------------------------------------------------------ + +node = DriftTEAPOT(length=args.length) +node.setLength(args.length) +node.setnParts(args.nslice) + +lattice = TEAPOT_Lattice() +lattice.addNode(node) +lattice.initialize() + + +# Create envelope +# ------------------------------------------------------------------------------ + +bunch = Bunch() +bunch.mass(mass_proton) +sync_part = bunch.getSyncParticle() +sync_part.kinEnergy(args.kin_energy) + +cov_matrix_init = np.zeros((6, 6)) +cov_matrix_init[0, 0] = args.xrms**2 +cov_matrix_init[2, 2] = args.yrms**2 +cov_matrix_init[4, 4] = (args.zrms / sync_part.gamma()) ** 2 + +centroid_init = np.zeros(6) + +envelope = Envelope( + sync_part=sync_part, + cov_matrix=cov_matrix_init, + centroid=centroid_init, + intensity=args.intensity, +) + +# Track envelope +# ------------------------------------------------------------------------------ + +print("TRACK ENVELOPE") + +tracker = EnvelopeTracker(lattice, space_charge=("3d" if args.sc else None)) + +history = {"xrms": [], "yrms": [], "zrms": []} +for turn in range(args.turns): + if turn > 0: + tracker.track(envelope) + + cov_matrix = envelope.cov() + centroid = envelope.centroid() + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + zrms = 1000.0 * math.sqrt(cov_matrix[4, 4]) * envelope.gamma() + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["zrms"].append(zrms) + + print(f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} zrms={zrms:0.3f}") + +histories = {} +histories["envelope"] = copy.deepcopy(history) + + +# Track bunch +# ------------------------------------------------------------------------------ + +print("TRACK BUNCH") + +bunch_coords = np.zeros((args.nparts, 6)) +bunch_coords[:, (0, 2, 4)] = gen_dist( + args.nparts, cov_matrix=np.eye(3), name="waterbag" +) +bunch_coords[:, 0] *= args.xrms +bunch_coords[:, 2] *= args.yrms +bunch_coords[:, 4] *= args.zrms / sync_part.gamma() + +for x, xp, y, yp, z, dE in bunch_coords: + bunch.addParticle(x, xp, y, yp, z, dE) + +size_global = bunch.getSizeGlobal() +bunch.macroSize(args.intensity / size_global) + +if args.sc: + sc_calc = SpaceChargeCalc3D(args.sc_grid, args.sc_grid, args.sc_grid) + sc_path_length_min = 0.01 + sc_nodes = setSC3DAccNodes(lattice, sc_path_length_min, sc_calc) + +history = {"xrms": [], "yrms": [], "zrms": []} +for turn in range(args.turns): + if turn > 0: + lattice.trackBunch(bunch) + + twiss_calc = BunchTwissAnalysis() + twiss_calc.computeBunchMoments(bunch, 2, 0, 0) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(i + 1): + cov_matrix[i, j] = twiss_calc.getCorrelation(j, i) + cov_matrix[j, i] = cov_matrix[i, j] + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + zrms = 1000.0 * math.sqrt(cov_matrix[4, 4]) * bunch.getSyncParticle().gamma() + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["zrms"].append(zrms) + + print(f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} zrms={zrms:0.3f}") + +histories["bunch"] = copy.deepcopy(history) + + +# Analysis +# ------------------------------------------------------------------------------ + +for history in histories.values(): + for key in history: + history[key] = np.array(history[key]) + +# Print errors +for key in histories["envelope"]: + deltas = histories["bunch"][key] - histories["envelope"][key] + print("key:", key) + print("max_abs_delta:", np.max(np.abs(deltas))) + print("avg_abs_delta:", np.mean(np.abs(deltas))) + +# Plot rms bunch sizes +for key in ["xrms", "yrms", "zrms"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(0.0, ax.get_ylim()[1]) + ax.set_xlabel("Turn") + ax.set_ylabel("RMS [mm]") + ax.legend(loc="upper left") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + +# Collect bunch/envelope data on final turn. +particles = collect_bunch(bunch)["coords"] +particles *= 1e3 + +env_cov_matrix = envelope.cov() +env_cov_matrix *= 1e6 + +env_centroid = envelope.centroid() +env_centroid *= 1e3 + +xmax = 4.0 * np.std(particles, axis=0) +limits = list(zip(-xmax, xmax)) +labels = ["x [mm]", "xp [mrad]", "y [mm]", "yp [mrad]", "z [mm]", "dE [MeV]"] + +# Plot x-x' +fig, ax = plt.subplots(figsize=(4, 4)) +ax.hist2d(particles[:, 0], particles[:, 1], bins=100, range=[limits[0], limits[1]]) +plot_rms_ellipse( + env_cov_matrix[0:2, 0:2], + center=(env_centroid[0], env_centroid[1]), + level=2.0, + color="red", + ax=ax, +) +ax.set_xlabel(labels[0]) +ax.set_ylabel(labels[1]) +plt.savefig(os.path.join(output_dir, "fig_dist_x_xp")) +plt.close() + +# Plot corner +fig, axs = plot_corner( + particles, + limits=limits, + bins=100, + labels=labels, +) +for i in range(6): + for j in range(i): + env_cov_matrix_proj = project_cov_matrix(env_cov_matrix, axis=(j, i)) + plot_rms_ellipse( + env_cov_matrix_proj, + center=(env_centroid[j], env_centroid[i]), + level=2.0, + color="red", + ax=axs[i, j], + ) +plt.savefig(os.path.join(output_dir, "fig_dist_corner")) +plt.close() diff --git a/examples/Envelope/test_env_sns_ring.py b/examples/Envelope/test_env_sns_ring.py new file mode 100644 index 00000000..f5ef19ef --- /dev/null +++ b/examples/Envelope/test_env_sns_ring.py @@ -0,0 +1,325 @@ +"""Test envelope tracker in SNS ring.""" + +import argparse +import copy +import math +import os +import pathlib + +import numpy as np +import matplotlib.pyplot as plt + +from orbit.core.bunch import Bunch +from orbit.core.bunch import BunchTwissAnalysis +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.utils.consts import mass_proton + +from plot import plot_rms_ellipse +from plot import plot_corner +from utils import gen_dist +from utils import build_rotation_matrix_xy +from utils import project_cov_matrix + +plt.style.use("style.mplstyle") + + +# Parse arguments +# ------------------------------------------------------------------------------ + +parser = argparse.ArgumentParser() +parser.add_argument("--bunch-length", type=float, default=120.0) +parser.add_argument("--kin-energy", type=float, default=1.300) +parser.add_argument("--intensity", type=float, default=5e14) + +parser.add_argument("--dist", type=str, default="kv", choices=["kv", "waterbag", "gauss"]) +parser.add_argument("--mismatch-x", type=float, default=0.0) +parser.add_argument("--mismatch-y", type=float, default=0.0) +parser.add_argument("--offset-x", type=float, default=0.0) +parser.add_argument("--offset-y", type=float, default=0.0) +parser.add_argument("--tilt", type=float, default=0) + +parser.add_argument("--nparts", type=int, default=100_000) +parser.add_argument("--turns", type=int, default=25) +parser.add_argument("--sol", type=int, default=0) +parser.add_argument("--sc", type=int, default=0) +parser.add_argument("--sc-grid", type=int, default=64) + +parser.add_argument("--handle-unknown", type=str, default=None, choices=["drift", "fit"]) +args = parser.parse_args() + + +# Setup +# ------------------------------------------------------------------------------ + +path = pathlib.Path(__file__) +output_dir = os.path.join("outputs", path.stem) +os.makedirs(output_dir, exist_ok=True) + + +# Create lattice +# ------------------------------------------------------------------------------ + +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 + +if args.sol: + for name in ["scbdsol_c13a", "scbdsol_c13b"]: + node = lattice.getNodeForName(name) + node.setParam("B", 0.15) + + +# Create envelope +# ------------------------------------------------------------------------------ + +# Create bunch +bunch = Bunch() +bunch.mass(mass_proton) +sync_part = bunch.getSyncParticle() +sync_part.kinEnergy(args.kin_energy) + +# Find periodic lattice parameters +matrix_lattice = TEAPOT_MATRIX_Lattice(lattice, bunch) +matrix_lattice_params = matrix_lattice.getRingParametersDict() +alpha_x = matrix_lattice_params["alpha x"] +alpha_y = matrix_lattice_params["alpha y"] +beta_x = matrix_lattice_params["beta x [m]"] +beta_y = matrix_lattice_params["beta y [m]"] +eps_x = 25.0e-06 +eps_y = eps_x + +# Generate covariance matrix +cov_matrix = np.zeros((6, 6)) +cov_matrix[0, 0] = eps_x * beta_x +cov_matrix[2, 2] = eps_y * beta_y +cov_matrix[0, 1] = cov_matrix[1, 0] = -eps_x * alpha_x +cov_matrix[2, 3] = cov_matrix[3, 2] = -eps_y * alpha_y +cov_matrix[1, 1] = eps_x * (1.0 + alpha_x**2) / beta_x +cov_matrix[3, 3] = eps_y * (1.0 + alpha_y**2) / beta_y +cov_matrix[4, 4] = (args.bunch_length / 4.0) ** 2 +cov_matrix[5, 5] = 0.0 + +# Tilt +if args.tilt: + rot_matrix = np.identity(6) + rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi)) + cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T]) + +# Mismatch +cov_matrix[0, 0] *= (1.0 + args.mismatch_x) ** 2 +cov_matrix[2, 2] *= (1.0 + args.mismatch_y) ** 2 +cov_matrix_init = np.copy(cov_matrix) + +# Offset +centroid_init = np.zeros(6) +centroid_init[0] += args.offset_x +centroid_init[2] += args.offset_y + +# Create envelope +envelope = Envelope( + sync_part=sync_part, + cov_matrix=cov_matrix_init, + centroid=centroid_init, + intensity=args.intensity, +) + + +# Track envelope +# ------------------------------------------------------------------------------ + +print("TRACK ENVELOPE") + +tracker = EnvelopeTracker( + lattice, + handle_unknown=args.handle_unknown, + space_charge=("2d" if args.sc else None), +) + +history = {"xrms": [], "yrms": [], "xavg": [], "yavg": []} +for turn in range(args.turns): + if turn > 0: + tracker.track(envelope) + + cov_matrix = envelope.cov() + centroid = envelope.centroid() + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + xavg = 1000.0 * centroid[0] + yavg = 1000.0 * centroid[2] + + print( + f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} xavg={xavg:0.3f} yavg={yavg:0.3f}" + ) + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["xavg"].append(xavg) + history["yavg"].append(yavg) + +histories = {} +histories["envelope"] = copy.deepcopy(history) + + +# Track bunch +# ------------------------------------------------------------------------------ + +print("TRACK BUNCH") + +rng = np.random.default_rng() + +bunch_coords = np.zeros((args.nparts, 6)) +bunch_coords[:, :4] = gen_dist( + n=args.nparts, cov_matrix=cov_matrix_init[0:4, 0:4], name=args.dist +) +bunch_coords[:, 4] = args.bunch_length * rng.uniform(-0.5, 0.5, size=args.nparts) +bunch_coords += centroid_init[None, :6] + +for i in range(bunch_coords.shape[0]): + bunch.addParticle(*bunch_coords[i]) + +if args.sc: + sc_calc = SpaceChargeCalc2p5D(128, 128, 1) + sc_path_length_min = 1.00e-06 + sc_nodes = setSC2p5DAccNodes(lattice, sc_path_length_min, sc_calc) + + bunch_size = bunch.getSizeGlobal() + bunch.macroSize(args.intensity / bunch_size) + +history = {"xrms": [], "yrms": [], "xavg": [], "yavg": []} +for turn in range(args.turns): + if turn > 0: + lattice.trackBunch(bunch) + + twiss_calc = BunchTwissAnalysis() + twiss_calc.computeBunchMoments(bunch, 2, 0, 0) + + cov_matrix = np.zeros((6, 6)) + for i in range(6): + for j in range(i + 1): + cov_matrix[i, j] = twiss_calc.getCorrelation(j, i) + cov_matrix[j, i] = cov_matrix[i, j] + + xrms = 1000.0 * math.sqrt(cov_matrix[0, 0]) + yrms = 1000.0 * math.sqrt(cov_matrix[2, 2]) + xavg = 1000.0 * twiss_calc.getAverage(0) + yavg = 1000.0 * twiss_calc.getAverage(2) + + print( + f"turn={turn} xrms={xrms:0.3f} yrms={yrms:0.3f} xavg={xavg:0.3f} yavg={yavg:0.3f}" + ) + + history["xrms"].append(xrms) + history["yrms"].append(yrms) + history["xavg"].append(xavg) + history["yavg"].append(yavg) + +histories["bunch"] = copy.deepcopy(history) + + +# Analysis +# ------------------------------------------------------------------------------ + +for history in histories.values(): + for key in history: + history[key] = np.array(history[key]) + +# Print errors +for key in histories["envelope"]: + deltas = histories["bunch"][key] - histories["envelope"][key] + print("key:", key) + print("max_abs_delta:", np.max(np.abs(deltas))) + print("avg_abs_delta:", np.mean(np.abs(deltas))) + +# Plot rms bunch sizes +for key in ["xrms", "yrms"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(0.0, ax.get_ylim()[1] * 2.0) + ax.set_xlabel("Turn") + ax.set_ylabel("RMS [mm]") + ax.legend(loc="upper right") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + +# Plot centroids +for key in ["xavg", "yavg"]: + fig, ax = plt.subplots(figsize=(5, 3)) + for i, model in enumerate(["envelope", "bunch"]): + color = ["black", "red"][i] + lw = [None, 0][i] + ax.plot(histories[model][key], marker=".", lw=lw, color=color, label=model) + ax.set_ylim(-5.0, 5.0) + ax.set_xlabel("Turn") + ax.set_ylabel("AVG [mm]") + ax.legend(loc="upper right") + plt.savefig(os.path.join(output_dir, f"fig_{key}")) + plt.close() + + +# Collect bunch/envelope data on final turn. +particles = collect_bunch(bunch)["coords"] +particles[:, :4] *= 1000.0 + +env_cov_matrix = envelope.cov() +env_cov_matrix[:4, :4] *= 1000.0**2 + +env_centroid = envelope.centroid() +env_centroid[:4] *= 1000.0 + +xmax = 4.0 * np.std(particles, axis=0) +limits = list(zip(-xmax, xmax)) +labels = ["x [mm]", "xp [mrad]", "y [mm]", "yp [mrad]", "z [m]", "dE [GeV]"] + + +# Plot x-x' +fig, ax = plt.subplots(figsize=(4, 4)) +ax.hist2d(particles[:, 0], particles[:, 1], bins=100, range=[limits[0], limits[1]]) +plot_rms_ellipse( + env_cov_matrix[0:2, 0:2], + center=(env_centroid[0], env_centroid[1]), + level=2.0, + color="red", + ax=ax, +) +ax.set_xlabel(labels[0]) +ax.set_ylabel(labels[1]) +plt.savefig(os.path.join(output_dir, "fig_dist_x_xp")) +plt.close() + +# Plot corner +fig, axs = plot_corner( + particles, + limits=limits, + bins=100, + labels=labels, +) +for i in range(6): + for j in range(i): + env_cov_matrix_proj = project_cov_matrix(env_cov_matrix, axis=(j, i)) + plot_rms_ellipse( + env_cov_matrix_proj, + center=(env_centroid[j], env_centroid[i]), + level=2.0, + color="red", + ax=axs[i, j], + ) +plt.savefig(os.path.join(output_dir, "fig_dist_corner")) +plt.close() diff --git a/examples/Envelope/utils.py b/examples/Envelope/utils.py new file mode 100644 index 00000000..432e4945 --- /dev/null +++ b/examples/Envelope/utils.py @@ -0,0 +1,59 @@ +import numpy as np + + +def gen_dist_gauss(n: int, cov_matrix: np.ndarray) -> np.ndarray: + return np.random.multivariate_normal( + mean=np.zeros(cov_matrix.shape[0]), + cov=cov_matrix, + size=n, + ) + + +def gen_dist_kv(n: int, cov_matrix: np.ndarray) -> np.ndarray: + X = np.random.normal(size=(n, cov_matrix.shape[0])) + X /= np.linalg.norm(X, axis=1)[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist_waterbag(n: int, cov_matrix: np.ndarray) -> np.ndarray: + X = gen_dist_kv(n, cov_matrix) + dim = X.shape[1] + r = np.random.uniform(size=n) ** (1.0 / dim) + X *= r[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist(n: int, cov_matrix: np.ndarray, name: str) -> np.ndarray: + if name == "kv": + X = gen_dist_kv(n, cov_matrix) + elif name == "waterbag": + X = gen_dist_waterbag(n, cov_matrix) + elif name == "gauss": + X = gen_dist_gauss(n, cov_matrix) + else: + raise ValueError(f"Invalid distribution name: {name}") + + L = np.linalg.cholesky(cov_matrix) + return np.matmul(X, L.T) + + +def build_rotation_matrix_xy(angle: float) -> np.ndarray: + cs = np.cos(angle) + sn = np.sin(angle) + + matrix = np.identity(4) + matrix[0, 0] = matrix[1, 1] = +cs + matrix[0, 2] = matrix[1, 3] = +sn + matrix[2, 0] = matrix[3, 1] = -sn + matrix[2, 2] = matrix[3, 3] = +cs + return matrix + + +def project_cov_matrix(cov_matrix: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: + cov_matrix_proj = np.zeros((len(axis), len(axis))) + for i in range(len(axis)): + for j in range(len(axis)): + cov_matrix_proj[i, j] = cov_matrix[axis[i], axis[j]] + return cov_matrix_proj diff --git a/py/orbit/envelope/__init__.py b/py/orbit/envelope/__init__.py new file mode 100644 index 00000000..70a5fcbc --- /dev/null +++ b/py/orbit/envelope/__init__.py @@ -0,0 +1,2 @@ +from .envelope import Envelope +from .envelope import EnvelopeTracker diff --git a/py/orbit/envelope/envelope.py b/py/orbit/envelope/envelope.py new file mode 100644 index 00000000..cce9aa0c --- /dev/null +++ b/py/orbit/envelope/envelope.py @@ -0,0 +1,268 @@ +import math + +import numpy as np +import scipy.special + +from ..core.bunch import SyncParticle +from ..lattice import AccNode +from ..lattice import AccLattice +from ..utils.consts import speed_of_light +from ..utils.consts import charge_electron + +from .matrix import MatrixFactory +from .utils import gen_dist +from .utils import proj_cov_matrix + +ENTRANCE = AccNode.ENTRANCE +BODY = AccNode.BODY +EXIT = AccNode.EXIT + +BEFORE = AccNode.BEFORE +AFTER = AccNode.AFTER + +CLASSICAL_PROTON_RADIUS = 1.534697049469832e-18 # [m] + + +def build_diag_matrix_from_xyz_eig(eigenvectors: np.ndarray) -> np.ndarray: + A = np.eye(7) + for i in range(eigenvectors.shape[0]): + for j in range(eigenvectors.shape[1]): + row = i * 2 + col = j * 2 + A[row, col] = A[row + 1, col + 1] = eigenvectors[i, j] + return A + + +class Envelope: + """Represents beam envelope and centroid. + + Attributes: + matrix: 7 x 7 covariance matrix for augmented phase space vector. + + Define the phase space vector X = [x, x', y, y', z, dE]^T and + augmented vector Y = [x, x', y, y', z, dE, 1]. + + Let X evolve according to X -> MX + U, where M is a 6 x 6 transfer matrix + and U is 6 x 1 "driving" vector. The augmented vector Y evolves according + to Y -> NY, where N = [[M, U], [0, 1]] is a 7 x 7 matrix. + + We track the 7 x 7 covariance matrix of Y: + + R = = [[, ], [, 1]], + + which contains both the phase space covariance matrix and centroid vector. + R evolves according to R -> N R N^T. + """ + + def __init__( + self, + sync_part: SyncParticle, + cov_matrix: np.ndarray = None, + centroid: np.ndarray = None, + intensity: float = 0.0, + ) -> None: + self.sync_part = sync_part + + if centroid is None: + centroid = np.zeros(6) + + if cov_matrix is None: + cov_matrix = np.eye(6) + + self.matrix = np.zeros((7, 7)) + self.matrix[0:6, 0:6] = cov_matrix + self.matrix[0:6, 6] = centroid + self.matrix[6, 0:6] = centroid + self.matrix[6, 6] = 1.0 + + self.intensity = 0.0 + self.perveance_2d = 0.0 + self.perveance_3d = 0.0 + self.set_intensity(intensity) + + def set_intensity(self, intensity: float) -> None: + self.intensity = intensity + self.perveance = ( + 2.0 + * intensity + * CLASSICAL_PROTON_RADIUS + / (self.beta() ** 2 * self.gamma() ** 3) + ) + + def gamma(self) -> float: + return self.sync_part.gamma() + + def beta(self) -> float: + return self.sync_part.beta() + + def mass(self) -> float: + return self.sync_part.mass() + + def centroid(self) -> np.ndarray: + return np.copy(self.matrix[0:6, 6]) + + def cov(self) -> np.ndarray: + return np.copy(self.matrix[0:6, 0:6]) + + def rms(self, axis: int = None) -> float | np.ndarray: + rms_arr = np.sqrt(np.diag(self.cov())) + return rms_arr[axis] + + def apply_transfer_matrix(self, transfer_matrix: np.ndarray) -> None: + self.matrix = np.linalg.multi_dot( + [transfer_matrix, self.matrix, transfer_matrix.T] + ) + + def sample(self, n: int, dist: str = "kv") -> np.ndarray: + # Issue: covariance matrix is becoming non semi-positive definite, + # giving error in cholesky decomposition. + particles = gen_dist(n=n, cov_matrix=self.cov(), name=dist) + particles = particles + self.centroid() + return particles + + def sc_transfer_matrix_2d(self, length: float) -> np.ndarray: + # Extract beam centroid and covariance matrix. + centroid = self.centroid() + cov_matrix = self.cov() + + # Project covariance matrix onto x-y plane. + cov_matrix_proj = proj_cov_matrix(cov_matrix, axis=(0, 2)) + + # Compute eigenvalues and eigenvectors of x-y covariance matrix. + cov_eig_res = np.linalg.eig(cov_matrix_proj) + cov_eig_vals = cov_eig_res.eigenvalues + cov_eig_vecs = cov_eig_res.eigenvectors + + # Compute rms beam sizes in upright frame. + rx = 2.0 * math.sqrt(cov_eig_vals[0]) + ry = 2.0 * math.sqrt(cov_eig_vals[1]) + + # Build transfer matrix in upright frame. + bunch_length = 4.0 * self.rms(axis=4) + perveance = self.perveance / bunch_length + factor = 2.0 * perveance / (rx + ry) + kappa_x = factor / rx + kappa_y = factor / ry + + M = np.identity(7) + M[1, 0] = kappa_x * length + M[3, 2] = kappa_y * length + + # Build matrix to undo x-y diagonalization. + A = build_diag_matrix_from_xyz_eig(cov_eig_vecs) + + # Build matrix to translate to centroid. + T = np.identity(7) + T[0, -1] = centroid[0] + T[2, -1] = centroid[2] + + # Compute matrix in lab frame. + # x = V u = T A u. + # u -> M u + # x -> V M V^-1 x + V = np.matmul(T, A) + V_inv = np.linalg.inv(V) + return np.linalg.multi_dot([V, M, V_inv]) + + def sc_transfer_matrix_3d(self, length: float) -> np.ndarray: + centroid = self.centroid() + centroid[4] *= self.gamma() + + cov_matrix = self.cov() + cov_matrix[4, 4] *= self.gamma() ** 2 + cov_matrix_proj = proj_cov_matrix(cov_matrix, axis=(0, 2, 4)) + + # Compute eigenvalues and eigenvectors of x-y covariance matrix. + cov_eig_res = np.linalg.eig(cov_matrix_proj) + cov_eig_vals = cov_eig_res.eigenvalues + cov_eig_vecs = cov_eig_res.eigenvectors + + # Build transfer matrix in upright frame. + cov_xx, cov_yy, cov_zz = cov_eig_vals + RDx = scipy.special.elliprd(cov_yy, cov_zz, cov_xx) + RDy = scipy.special.elliprd(cov_xx, cov_zz, cov_yy) + RDz = scipy.special.elliprd(cov_xx, cov_yy, cov_zz) + + factor = 0.5 * self.perveance * ((1.0 / 5.0) ** 1.5) + kappa_x = factor * RDx # [1 / m] + kappa_y = factor * RDy # [1 / m] + kappa_z = factor * RDz # [1 / m] + kappa_z *= self.gamma() ** 3 * self.beta() ** 2 * self.mass() # [GeV / m] + + M = np.identity(7) + M[1, 0] = kappa_x * length + M[3, 2] = kappa_y * length + M[5, 4] = kappa_z * length + + # Build matrix to undo x-y-z diagonalization. + A = build_diag_matrix_from_xyz_eig(cov_eig_vecs) + + # Build matrix for translation to centroid. + T = np.identity(7) + for i in (0, 2, 4): + T[i, -1] = centroid[i] + + # Build matrix for Lorentz boost (length contraction). + L = np.identity(7) + L[4, 4] = 1.0 / self.gamma() + + # Compute matrix in lab frame. + # x = L V u = L T A u. + # u -> M u + # x -> V M V^-1 x + V = np.matmul(T, A) + V_inv = np.linalg.inv(V) + return np.linalg.multi_dot([V, M, V_inv]) + + +class EnvelopeTracker: + """Tracks envelope through linear lattice with optional linear space charge kicks.""" + + def __init__( + self, + lattice: AccLattice, + space_charge: str | None = None, + handle_unknown: str | None = None, + ) -> None: + self.lattice = lattice + self.matrix_factory = MatrixFactory(handle_unknown=handle_unknown) + self.space_charge = space_charge + + def track(self, envelope: Envelope) -> None: + for node in self.lattice.getNodes(): + for child_node in node.getChildNodes(ENTRANCE): + envelope.apply_transfer_matrix( + self.matrix_factory(child_node, envelope.sync_part) + ) + + for part_index in range(node.getnParts()): + for child_node in node.getChildNodes(BODY, part_index, place_in_part=BEFORE): + envelope.apply_transfer_matrix( + self.matrix_factory(child_node, envelope.sync_part) + ) + + if self.space_charge: + length = node.getLength(part_index) + if self.space_charge == "2d": + matrix = envelope.sc_transfer_matrix_2d(length) + elif self.space_charge == "3d": + matrix = envelope.sc_transfer_matrix_3d(length) + else: + raise ValueError( + f"Invalid space charge model: {self.space_charge}" + ) + envelope.apply_transfer_matrix(matrix) + + envelope.apply_transfer_matrix( + self.matrix_factory(node, envelope.sync_part, part_index) + ) + + for child_node in node.getChildNodes(BODY, part_index, place_in_part=AFTER): + envelope.apply_transfer_matrix( + self.matrix_factory(child_node, envelope.sync_part) + ) + + for child_node in node.getChildNodes(EXIT): + envelope.apply_transfer_matrix( + self.matrix_factory(child_node, envelope.sync_part) + ) diff --git a/py/orbit/envelope/matrix.py b/py/orbit/envelope/matrix.py new file mode 100644 index 00000000..6196262a --- /dev/null +++ b/py/orbit/envelope/matrix.py @@ -0,0 +1,186 @@ +import math + +import numpy as np + +from ..core.bunch import Bunch +from ..core.bunch import SyncParticle +from ..lattice import AccNode +from ..teapot import DriftTEAPOT +from ..teapot import QuadTEAPOT +from ..teapot import BendTEAPOT +from ..teapot import TiltTEAPOT +from ..teapot import KickTEAPOT +from ..teapot import ApertureTEAPOT +from ..teapot import BunchWrapTEAPOT +from ..teapot import FringeFieldTEAPOT +from ..teapot import MonitorTEAPOT +from ..teapot import TurnCounterTEAPOT + +from ..utils import speed_of_light + + +def get_dp_p_coeff(sync_part: SyncParticle) -> float: + return 1.0 / (sync_part.momentum() * sync_part.beta()) + + +class MatrixFactory: + """Factory for 7 x 7 transfer matrices. + + Units: x [m], x' [rad], y [m], y' [rad], z [m], dE [GeV] + """ + def __init__(self, handle_unknown: str | None = None) -> None: + self.ignore_node_types = [ + ApertureTEAPOT, + BunchWrapTEAPOT, + FringeFieldTEAPOT, + MonitorTEAPOT, + TurnCounterTEAPOT, + ] + self.handle_unknown = handle_unknown + + def drift(self, length: float, sync_part: SyncParticle) -> np.ndarray: + matrix = np.identity(7) + matrix[0, 1] = length + matrix[2, 3] = length + matrix[4, 5] = length / sync_part.gamma() ** 2 + + # Matrix above is for dp_p; switch to dE. + matrix[4, 5] *= get_dp_p_coeff(sync_part) + matrix[5, 4] /= get_dp_p_coeff(sync_part) + return matrix + + def quad(self, length: float, kq: float, sync_part: SyncParticle) -> np.ndarray: + sqrt_abs_kq = math.sqrt(abs(kq)) + + matrix = 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) + matrix[0, 0] = cx + matrix[0, 1] = +sx / sqrt_abs_kq + matrix[1, 0] = -sx * sqrt_abs_kq + matrix[1, 1] = cx + matrix[2, 2] = cy + matrix[2, 3] = sy / sqrt_abs_kq + matrix[3, 2] = sy * sqrt_abs_kq + matrix[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) + matrix[0, 0] = cx + matrix[0, 1] = sx / sqrt_abs_kq + matrix[1, 0] = sx * sqrt_abs_kq + matrix[1, 1] = cx + matrix[2, 2] = cy + matrix[2, 3] = +sy / sqrt_abs_kq + matrix[3, 2] = -sy * sqrt_abs_kq + matrix[3, 3] = cy + + matrix[4, 5] = length / sync_part.gamma() ** 2 + matrix[4, 5] *= get_dp_p_coeff(sync_part) + matrix[5, 4] /= get_dp_p_coeff(sync_part) + return matrix + + def bend(self, length: float, theta: float, sync_part: SyncParticle) -> np.ndarray: + if length <= 0: + return np.identity(7) + + v = speed_of_light * sync_part.beta() + sync_part.time(sync_part.time() + length / v) + + rho = length / theta + cx = math.cos(theta) + sx = math.sin(theta) + + betasq = sync_part.beta() ** 2 + + rho = length / theta + cx = math.cos(theta) + sx = math.sin(theta) + + matrix = np.identity(7) + matrix[0, 0] = cx + matrix[0, 1] = rho * sx + matrix[0, 5] = rho * (1.0 - cx) + matrix[1, 0] = -sx / rho + matrix[1, 1] = cx + matrix[1, 5] = sx + matrix[2, 3] = length + matrix[4, 0] = -sx + matrix[4, 1] = -rho * (1.0 - cx) + matrix[4, 5] = -betasq * length + rho * sx + + matrix[4, 5] *= get_dp_p_coeff(sync_part) + matrix[5, 4] /= get_dp_p_coeff(sync_part) + return matrix + + def tilt(self, angle: float) -> np.ndarray: + matrix = np.identity(7) + matrix[0, 0] = matrix[1, 1] = +math.cos(angle) + matrix[0, 2] = matrix[1, 3] = +math.sin(angle) + matrix[2, 0] = matrix[3, 1] = -math.sin(angle) + matrix[2, 2] = matrix[3, 3] = +math.cos(angle) + return matrix + + def translation(self, x: float = 0.0, y: float = 0.0, z: float = 0.0) -> np.ndarray: + matrix = np.identity(7) + matrix[0, 6] = x + matrix[2, 6] = y + matrix[4, 6] = z + return matrix + + def kick(self, kx: float, ky: float, dE: float) -> np.ndarray: + matrix = np.identity(7) + matrix[1, 6] = kx + matrix[3, 6] = ky + matrix[5, 6] = dE + return matrix + + def __call__(self, node: AccNode, sync_part: SyncParticle, part_index: int = 0) -> np.ndarray: + if type(node) is DriftTEAPOT: + length = node.getLength(part_index) + return self.drift(length=length, sync_part=sync_part) + elif type(node) is QuadTEAPOT: + length = node.getLength(part_index) + + scale = 1.0 + if node.waveform: + scale = node.waveform.getStrength() + + kq = scale * node.getParam("kq") + return self.quad(length=length, kq=kq, sync_part=sync_part) + + elif type(node) is BendTEAPOT: + length = node.getLength(part_index) + theta = node.getParam("theta") / node.getnParts() + return self.bend(length=length, theta=theta, sync_part=sync_part) + + elif type(node) is KickTEAPOT: + nparts = node.getnParts() + + scale = 1.0 + if node.waveform: + scale = node.waveform.getStrength() + + kx = scale * node.getParam("kx") / nparts + ky = scale * node.getParam("ky") / nparts + dE = node.getParam("dE") / nparts + return self.kick(kx, ky, dE) + + elif type(node) is TiltTEAPOT: + angle = node.getTiltAngle() + return self.tilt(angle) + + elif type(node) in self.ignore_node_types: + return np.identity(7) + + else: + if self.handle_unknown == "drift": + return self.drift(length=node.getLength(), sync_part=sync_part) + elif self.handle_unknown == "fit": + raise NotImplementedError() + raise NotImplementedError("Unsupported node type: {}. See `handle_unknown` attribute.".format(type(node))) diff --git a/py/orbit/envelope/meson.build b/py/orbit/envelope/meson.build new file mode 100644 index 00000000..31c9033b --- /dev/null +++ b/py/orbit/envelope/meson.build @@ -0,0 +1,13 @@ +py_sources = files([ + '__init__.py', + 'envelope.py', + 'matrix.py', + 'utils.py' +]) + +python.install_sources( + py_sources, + subdir: 'orbit/envelope', + # pure: true, +) + diff --git a/py/orbit/envelope/utils.py b/py/orbit/envelope/utils.py new file mode 100644 index 00000000..83d8e847 --- /dev/null +++ b/py/orbit/envelope/utils.py @@ -0,0 +1,47 @@ +import numpy as np + + +def gen_dist_gauss(n: int, cov_matrix: np.ndarray) -> np.ndarray: + return np.random.multivariate_normal( + mean=np.zeros(cov_matrix.shape[0]), + cov=cov_matrix, + size=n, + ) + + +def gen_dist_kv(n: int, cov_matrix: np.ndarray) -> np.ndarray: + X = np.random.normal(size=(n, cov_matrix.shape[0])) + X /= np.linalg.norm(X, axis=1)[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist_waterbag(n: int, cov_matrix: np.ndarray) -> np.ndarray: + X = gen_dist_kv(n, cov_matrix) + dim = X.shape[1] + r = np.random.uniform(size=n) ** (1.0 / dim) + X *= r[:, None] + X /= np.std(X, axis=0) + return X + + +def gen_dist(n: int, cov_matrix: np.ndarray, name: str) -> np.ndarray: + if name == "kv": + X = gen_dist_kv(n, cov_matrix) + elif name == "waterbag": + X = gen_dist_waterbag(n, cov_matrix) + elif name == "gauss": + X = gen_dist_gauss(n, cov_matrix) + else: + raise ValueError(f"Invalid distribution name: {name}") + + L = np.linalg.cholesky(cov_matrix) + return np.matmul(X, L.T) + + +def proj_cov_matrix(cov_matrix: np.ndarray, axis: tuple[int, ...]) -> np.ndarray: + cov_matrix_proj = np.zeros((len(axis), len(axis))) + for i in range(len(axis)): + for j in range(len(axis)): + cov_matrix_proj[i, j] = cov_matrix[axis[i], axis[j]] + return cov_matrix_proj diff --git a/py/orbit/meson.build b/py/orbit/meson.build index 2ba16e73..651b14d5 100644 --- a/py/orbit/meson.build +++ b/py/orbit/meson.build @@ -24,6 +24,7 @@ subdir('space_charge') subdir('errors') subdir('matrix_lattice') subdir('teapot') +subdir('envelope') py_sources = files([ diff --git a/py/orbit/space_charge/sc1d/.ipynb_checkpoints/sc1DNode-checkpoint.py b/py/orbit/space_charge/sc1d/.ipynb_checkpoints/sc1DNode-checkpoint.py new file mode 100644 index 00000000..0bf968e1 --- /dev/null +++ b/py/orbit/space_charge/sc1d/.ipynb_checkpoints/sc1DNode-checkpoint.py @@ -0,0 +1,304 @@ +""" +Module. Includes classes for 1D longidutinal space charge accelerator nodes. +""" + +import sys +import os +import math + +from orbit.core.bunch import Bunch +from orbit.core.spacecharge import LSpaceChargeCalc +from orbit.lattice import AccLattice +from orbit.lattice import AccNode +from orbit.lattice import AccActionsContainer +from orbit.lattice import AccNodeBunchTracker +from orbit.utils import consts +from orbit.utils import orbitFinalize +from orbit.teapot import DriftTEAPOT + + +class SC1D_AccNode(DriftTEAPOT): + """Longitudinal space charge node.""" + + def __init__( + self, + b_a: float, + phase_length: float, + nmacros_min: float, + use_sc: float, + nbins: float, + nmodes: int = None, + use_grad: bool = False, + name="long sc node", + ) -> None: + """ + Constructor. Creates the SC1D-teapot element. + """ + DriftTEAPOT.__init__(self, name) + self.lspacecharge = LSpaceChargeCalc(b_a, phase_length, nmacros_min, use_sc, nbins) + self.setNumModes(nmodes) + # self.setUseGrad(use_grad) + self.setType("long sc node") + self.setLength(0.0) + + def setUseGrad(self, use_grad: bool) -> None: + """Sets whether to use gradient-based solver instead of impedance solver.""" + self.lspacecharge.setUseGrad(int(use_grad)) + + def setNumModes(self, n: int) -> None: + """Sets number of FFT modes used to calculate energy kick.""" + self.lspacecharge.setNumModes(n) + + def trackBunch(self, bunch: Bunch) -> None: + """ + The SC1D-teapot class implementation of the + AccNodeBunchTracker class trackBunch(probe) method. + """ + length = self.getLength(self.getActivePartIndex()) + self.lspacecharge.trackBunch(bunch) # track method goes here + + def track(self, params_dict: dict) -> None: + """ + The SC1D-teapot class implementation of the + AccNodeBunchTracker class track(probe) method. + """ + length = self.getLength(self.getActivePartIndex()) + bunch = params_dict["bunch"] + self.lspacecharge.trackBunch(bunch) + + def assignImpedance(self, py_cmplx_arr: list[float]) -> None: + self.lspacecharge.assignImpedance(py_cmplx_arr) + + +class FreqDep_SC1D_AccNode(DriftTEAPOT): + """Longitudinal space charge node (frequency-dependent).""" + + def __init__( + self, + b_a: float, + phase_length: float, + nmacros_min: int, + use_sc: int, + nbins: int, + bunch: Bunch, + imp_dict: dict, + name: str = "freq. dep. long sc node", + ) -> None: + """ + Constructor. Creates the FreqDep_SC1D-teapot element. + """ + DriftTEAPOT.__init__(self, name) + self.lspacecharge = LSpaceChargeCalc( + b_a, phase_length, nmacros_min, use_sc, nbins + ) + self.setType("freq. dep. long sc node") + self.setLength(0.0) + self.phase_length = phase_length + self.nbins = nbins + self.localDict = imp_dict + self.freq_tuple = self.localDict["freqs"] + self.freq_range = len(self.freq_tuple) - 1 + self.z_tuple = self.localDict["z_imp"] + self.c = consts.speed_of_light + BetaRel = bunch.getSyncParticle().beta() + Freq0 = (BetaRel * self.c) / self.phase_length + Z = [] + for n in range(self.nbins // 2 - 1): + freq_mode = Freq0 * (n + 1) + z_mode = interp(freq_mode, self.freq_range, self.freq_tuple, self.z_tuple) + Z.append(z_mode) + self.lspacecharge.assignImpedance(Z) + + def trackBunch(self, bunch: Bunch) -> None: + """ + The FreqDep_SC1D-teapot class implementation of + the AccNodeBunchTracker class track(probe) method. + """ + length = self.getLength(self.getActivePartIndex()) + BetaRel = bunch.getSyncParticle().beta() + Freq0 = (BetaRel * self.c) / self.phase_length + Z = [] + for n in range(self.nbins // 2 - 1): + freq_mode = Freq0 * (n + 1) + z_mode = interp(freq_mode, self.freq_range, self.freq_tuple, self.z_tuple) + Z.append(z_mode) + self.lspacecharge.assignImpedance(Z) + self.lspacecharge.trackBunch(bunch) + + def track(self, params_dict: dict) -> None: + """ + The FreqDep_SC1D-teapot class implementation of + the AccNodeBunchTracker class track(probe) method. + """ + length = self.getLength(self.getActivePartIndex()) + bunch = params_dict["bunch"] + BetaRel = bunch.getSyncParticle().beta() + Freq0 = (BetaRel * self.c) / self.phase_length + Z = [] + for n in range(self.nbins // 2 - 1): + freq_mode = Freq0 * (n + 1) + z_mode = interp(freq_mode, self.freq_range, self.freq_tuple, self.z_tuple) + Z.append(z_mode) + self.lspacecharge.assignImpedance(Z) + self.lspacecharge.trackBunch(bunch) + + +class BetFreqDep_SC1D_AccNode(DriftTEAPOT): + """Longitudinal space charge node (frequency- and velocity-dependent).""" + + def __init__( + self, + b_a: float, + phase_length: float, + nmacros_min: float, + use_sc: int, + nbins: int, + bunch: Bunch, + imp_dict: dict, + name: str = "freq. dep. long sc node", + ) -> None: + """ + Constructor. Creates the BetFreqDep_SC1D-teapot element. + """ + DriftTEAPOT.__init__(self, name) + self.lspacecharge = LSpaceChargeCalc( + b_a, phase_length, nmacros_min, use_sc, nbins + ) + self.setType("beta-freq. dep. long sc node") + self.setLength(0.0) + self.phase_length = phase_length + self.nbins = nbins + self.localDict = imp_dict + self.bet_tuple = self.localDict["betas"] + self.bet_range = len(self.bet_tuple) - 1 + self.freq_tuple = self.localDict["freqs"] + self.freq_range = len(self.freq_tuple) - 1 + self.z_bf = self.localDict["z_imp"] + self.c = consts.speed_of_light + BetaRel = bunch.getSyncParticle().beta() + Freq0 = (BetaRel * self.c) / self.phase_length + Z = [] + for n in range(self.nbins / 2 - 1): + freq_mode = Freq0 * (n + 1) + z_mode = bilinterp( + BetaRel, + freq_mode, + self.bet_range, + self.freq_range, + self.bet_tuple, + self.freq_tuple, + self.z_bf, + ) + Z.append(z_mode) + self.lspacecharge.assignImpedance(Z) + + def trackBunch(self, bunch: Bunch) -> None: + """ + The BetFreqDep_SC1D-teapot class implementation of + the AccNodeBunchTracker class track(probe) method. + """ + length = self.getLength(self.getActivePartIndex()) + BetaRel = bunch.getSyncParticle().beta() + Freq0 = (BetaRel * self.c) / self.phase_length + Z = [] + for n in range(self.nbins / 2 - 1): + freq_mode = Freq0 * (n + 1) + z_mode = bilinterp( + BetaRel, + freq_mode, + self.bet_range, + self.freq_range, + self.bet_tuple, + self.freq_tuple, + self.z_bf, + ) + Z.append(z_mode) + self.lspacecharge.assignImpedance(Z) + self.lspacecharge.trackBunch(bunch) + + def track(self, params_dict: dict) -> None: + """ + The BetFreqDep_SC1D-teapot class implementation of + the AccNodeBunchTracker class track(probe) method. + """ + length = self.getLength(self.getActivePartIndex()) + bunch = params_dict["bunch"] + BetaRel = bunch.getSyncParticle().beta() + Freq0 = (BetaRel * self.c) / self.phase_length + Z = [] + for n in range(self.nbins / 2 - 1): + freq_mode = Freq0 * (n + 1) + z_mode = bilinterp( + BetaRel, + freq_mode, + self.bet_range, + self.freq_range, + self.bet_tuple, + self.freq_tuple, + self.z_bf, + ) + Z.append(z_mode) + self.lspacecharge.assignImpedance(Z) + self.lspacecharge.trackBunch(bunch) + + +def interp(x: float, n_tuple: int, x_tuple: list[float], y_tuple: list[float]) -> float: + """ + Linear interpolation: Given n-tuple + 1 points, + x_tuple and y_tuple, routine finds y = y_tuple + at x in x_tuple. Assumes x_tuple is increasing array. + """ + if x < x_tuple[0]: + y = y_tuple[0] + return y + if x > x_tuple[n_tuple]: + y = y_tuple[n_tuple] + return y + dxp = x - x_tuple[0] + for n in range(n_tuple): + dxm = dxp + dxp = x - x_tuple[n + 1] + dxmp = dxm * dxp + if dxmp <= 0: + break + y = (-dxp * y_tuple[n] + dxm * y_tuple[n + 1]) / (dxm - dxp) + return y + + +def bilinterp( + x: float, + y: float, + nx_tuple: int, + ny_tuple: int, + x_tuple: list[float], + y_tuple: list[float], + fxy: list[list[float]], +) -> float: + """ + Bilinear interpolation: Given nx-tuple + 1 x-points, + ny-tuple + 1 y-points, x_tuple and y_tuple, + routine finds f(x, y) = fxy at (x, y) in (x_tuple, y_tuple). + Assumes x_tuple and y_tuple are increasing arrays. + """ + f_tuple = [] + if x < x_tuple[0]: + for ny in range(ny_tuple + 1): + vf = fxy[0][ny] + f_tuple.append(vf) + elif x > x_tuple[nx_tuple]: + for ny in range(ny_tuple + 1): + vf = fxy[x_tuple][ny] + f_tuple.append(vf) + else: + dxp = x - x_tuple[0] + for nx in range(nx_tuple): + dxm = dxp + dxp = x - x_tuple[nx + 1] + dxmp = dxm * dxp + if dxmp <= 0: + break + for ny in range(ny_tuple + 1): + vf = (-dxp * fxy[nx][ny] + dxm * fxy[nx + 1][ny]) / (dxm - dxp) + f_tuple.append(vf) + f = interp(y, ny_tuple, y_tuple, f_tuple) + return f diff --git a/py/orbit/teapot/__init__.py b/py/orbit/teapot/__init__.py index 945450d4..9886b7fe 100644 --- a/py/orbit/teapot/__init__.py +++ b/py/orbit/teapot/__init__.py @@ -16,6 +16,10 @@ from .teapot import SolenoidTEAPOT from .teapot import TiltTEAPOT from .teapot import NodeTEAPOT +from .teapot import TurnCounterTEAPOT +from .teapot import ApertureTEAPOT +from .teapot import MonitorTEAPOT +from .teapot import BunchWrapTEAPOT from .teapot import TPB @@ -38,3 +42,6 @@ __all__.append("NodeTEAPOT") __all__.append("TPB") __all__.append("TEAPOT_MATRIX_Lattice") +__all__.append("TurnCounterTEAPOT") +__all__.append("ApertureTEAPOT") +__all__.append("MonitorTEAPOT")