chromaticty calc example

chroma.py

+import at
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.special import factorial
+
+#simple_ring = at.simple_ring(3e9, 528, 176, 0.1, 0.1, 1e6, 0.000306, U0=0.0, tauz=1000, espread=0.000769,Qpy=1.0, chromx_arr=[0,1e3,1e5])
+simple_ring = at.simple_ring(3e9, 528, 176, 0.1, 0.1, 1e6, 0.000306, U0=0.0, tauz=1000, espread=0.000769,Qpy=1.0)
+simple_ring.enable_6d()
+
+fitOrder = 4
+(fitH, fitV), dpa, qz = at.chromaticity(simple_ring, method='linopt', dpm=3e-3, npoints=20, order=fitOrder, dp=0)
+
+fitH_pv = fitH/factorial(np.arange(fitOrder+1))
+fitH_pv = fitH_pv[::-1]
+
+fitV_pv = fitV/factorial(np.arange(fitOrder+1))
+fitV_pv = fitV_pv[::-1]
+
+labelH = r'$Q={:.1f}$'.format(fitH[0]) + '\n' +  \
+         r"$Q'={:.2f}$".format(fitH[1]) + '\n' +   \
+         r"$Q''={:.4f}$".format(fitH[2])
+
+labelV = r'$Q={:.1f}$'.format(fitV[0]) + '\n' +  \
+         r"$Q'={:.2f}$".format(fitV[1]) + '\n' +   \
+         r"$Q''={:.4f}$".format(fitV[2])
+
+fig, ax = plt.subplots(2,1)
+ax[0].plot(dpa, np.polyval(fitH_pv, dpa), label=labelH)
+ax[0].plot(dpa, qz[:,0], linestyle='None', marker='x')
+ax[1].plot(dpa, np.polyval(fitV_pv, dpa), label=labelV)
+ax[1].plot(dpa, qz[:,1], linestyle='None', marker='x')
+ax[0].legend()
+ax[1].legend()
+
+plt.show()