From 377f4c0056557d4675adc2ecd5184edd4f794fe8 Mon Sep 17 00:00:00 2001
From: Jan Kaiser <jan.kaiser.email@googlemail.com>
Date: Sun, 3 Sep 2023 15:29:04 +0200
Subject: [PATCH] Comparison to Ocelot test with ARES EA

---
 cheetah/utils.py            |  3 ++-
 test/test_compare_ocelot.py | 45 +++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/cheetah/utils.py b/cheetah/utils.py
index a95c6178..a6df01f9 100644
--- a/cheetah/utils.py
+++ b/cheetah/utils.py
@@ -113,7 +113,8 @@ def ocelot2cheetah(element, warnings: bool = True) -> "acc.Element":
     else:
         if warnings:
             print(
-                f"WARNING: Unknown element {element.id}, replacing with drift section."
+                f"WARNING: Unknown element {element.id} of type {type(element)},"
+                " replacing with drift section."
             )
         return acc.Drift(element.l, name=element.id)
 
diff --git a/test/test_compare_ocelot.py b/test/test_compare_ocelot.py
index 0dbc062b..922e48c4 100644
--- a/test/test_compare_ocelot.py
+++ b/test/test_compare_ocelot.py
@@ -1,3 +1,4 @@
+import test.ARESlatticeStage3v1_9 as ares
 from copy import deepcopy
 
 import numpy as np
@@ -157,3 +158,47 @@ def test_solenoid():
     assert np.allclose(
         outgoing_beam.particles[:, :6], outgoing_p_array.rparticles.transpose()
     )
+
+
+def test_ares_ea():
+    """
+    Test that the tracking results through a Experimental Area (EA) lattice of the ARES
+    accelerator at DESY match those using Ocelot.
+    """
+    cell = cheetah.utils.subcell_of_ocelot(ares.cell, "AREASOLA1", "AREABSCR1")
+    ares.areamqzm1.k1 = 5.0
+    ares.areamqzm2.k1 = -5.0
+    ares.areamcvm1.k1 = 1e-3
+    ares.areamqzm3.k1 = 5.0
+    ares.areamchm1.k1 = -2e-3
+
+    # Cheetah
+    incoming_beam = cheetah.ParticleBeam.from_astra(
+        "benchmark/cheetah/ACHIP_EA1_2021.1351.001"
+    )
+    cheetah_segment = cheetah.Segment.from_ocelot(cell)
+    outgoing_beam = cheetah_segment.track(incoming_beam)
+
+    # Ocelot
+    incoming_p_array = ocelot.astraBeam2particleArray(
+        "benchmark/cheetah/ACHIP_EA1_2021.1351.001", print_params=False
+    )
+    lattice = ocelot.MagneticLattice(cell)
+    navigator = ocelot.Navigator(lattice)
+    _, outgoing_p_array = ocelot.track(lattice, deepcopy(incoming_p_array), navigator)
+
+    assert np.isclose(outgoing_beam.mu_x, outgoing_p_array.x().mean())
+    assert np.isclose(outgoing_beam.mu_xp, outgoing_p_array.px().mean())
+    assert np.isclose(outgoing_beam.mu_y, outgoing_p_array.y().mean())
+    assert np.isclose(outgoing_beam.mu_yp, outgoing_p_array.py().mean())
+    assert np.isclose(outgoing_beam.mu_s, outgoing_p_array.tau().mean(), atol=1e-7)
+    assert np.isclose(outgoing_beam.mu_p, outgoing_p_array.p().mean())
+
+    assert np.allclose(outgoing_beam.xs, outgoing_p_array.x())
+    assert np.allclose(outgoing_beam.xps, outgoing_p_array.px())
+    assert np.allclose(outgoing_beam.ys, outgoing_p_array.y())
+    assert np.allclose(outgoing_beam.yps, outgoing_p_array.py())
+    assert np.allclose(
+        outgoing_beam.ss, outgoing_p_array.tau(), atol=1e-7, rtol=1e-1
+    )  # TODO: Why do we need such large tolerances?
+    assert np.allclose(outgoing_beam.ps, outgoing_p_array.p())