From 9529ba3450a85fe23cb7b5858dae4402ba40f5d9 Mon Sep 17 00:00:00 2001
From: gubaidulinvadim <gubaidulinvadim@gmail.com>
Date: Tue, 30 Jul 2024 12:05:09 +0200
Subject: [PATCH] added get_chroma() method to compute chromaticity up to
 arbitrary order from the lattice file.

---
 mbtrack2/tracking/synchrotron.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/mbtrack2/tracking/synchrotron.py b/mbtrack2/tracking/synchrotron.py
index 15235ec..ee55af1 100644
--- a/mbtrack2/tracking/synchrotron.py
+++ b/mbtrack2/tracking/synchrotron.py
@@ -387,6 +387,16 @@ class Synchrotron:
         coef_yy = np.array([det.A3 / 2, 0])
         self.adts = [coef_xx, coef_yx, coef_xy, coef_yy]
 
+    def get_chroma(self, order=4, dpm=0.02, n_points=100):
+        fit, dpa, tune = at.physics.nonlinear.chromaticity(
+            self.ring.optics.lattice,
+            method='linopt',
+            dpm=0.02,
+            n_points=100,
+            order=order)
+        chrox, chroy = fit
+        return chrox, chroy
+
     def get_mcf_order(self, add=True, show_fit=False):
         """
         Compute momentum compaction factor up to 3rd order from AT lattice.
-- 
GitLab