From 32d6c79e4e3fc2543ca96cb2a86f13c16fa609e1 Mon Sep 17 00:00:00 2001
From: Claudio Scafuri <claudio.scafuri@elettra.eu>
Date: Tue, 18 Mar 2025 10:44:20 +0100
Subject: [PATCH] chromaticty calc example

---
 chroma.py | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)
 create mode 100644 chroma.py

diff --git a/chroma.py b/chroma.py
new file mode 100644
index 0000000..e27cb1a
--- /dev/null
+++ b/chroma.py
@@ -0,0 +1,35 @@
+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()
-- 
GitLab