From d03e4934d4dffac1c09be18877ff7da4bd43dea5 Mon Sep 17 00:00:00 2001
From: Anton Golovanov <agolovanov256@gmail.com>
Date: Thu, 9 May 2024 16:01:21 +0300
Subject: [PATCH] Add test comparing Quadrupole and FieldQuadrupole

---
 .gitignore                     |  3 ++
 tests/test_field_quadrupole.py | 50 ++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 tests/test_field_quadrupole.py

diff --git a/.gitignore b/.gitignore
index 49d2f212..9110d945 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,6 +10,9 @@ __pycache__
 /examples/tests.py
 /old_stuff
 
+# Test output
+/tests_output
+
 # IDEs
 /.vscode
 
diff --git a/tests/test_field_quadrupole.py b/tests/test_field_quadrupole.py
new file mode 100644
index 00000000..2eb2543a
--- /dev/null
+++ b/tests/test_field_quadrupole.py
@@ -0,0 +1,50 @@
+import copy
+
+import numpy as np
+import scipy.constants as ct
+
+from wake_t.beamline_elements import FieldQuadrupole, Quadrupole
+from wake_t.utilities.bunch_generation import get_gaussian_bunch_from_twiss
+
+
+def test_field_vs_tm_quadrupole():
+    """
+    This test checks that the FieldElement-based quadrupole (FieldQuadrupole)
+    and the TM quadrupole (Quadrupole) produce similar results.
+    """
+
+    emitt_nx = emitt_ny = 1e-6  # m
+    beta_x = beta_y = 1.  # m
+    s_t = 100.  # fs
+    gamma_avg = 1000
+    ene_spread = 0.1  # %
+    q_bunch = 30  # pC
+    xi_avg = 0.  # m
+    n_part = 1e4
+    bunch_1 = get_gaussian_bunch_from_twiss(
+        en_x=emitt_nx, en_y=emitt_ny, a_x=0, a_y=0, b_x=beta_x, b_y=beta_y,
+        ene=gamma_avg, ene_sp=ene_spread, s_t=s_t, xi_c=xi_avg,
+        q_tot=q_bunch, n_part=n_part, name='elec_bunch')
+
+    bunch_2 = copy.deepcopy(bunch_1)
+
+    foc_strength = 100  # T/m
+    quadrupole_length = 0.05  # m
+    k1 = foc_strength * ct.e / ct.m_e / ct.c / gamma_avg
+
+    field_quadrupole = FieldQuadrupole(quadrupole_length, foc_strength)
+    tm_quadrupole = Quadrupole(quadrupole_length, k1)
+
+    field_quadrupole.track(bunch_1)
+    tm_quadrupole.track(bunch_2)
+
+    np.testing.assert_allclose(bunch_1.x, bunch_2.x, rtol=1e-3, atol=1e-7)
+    np.testing.assert_allclose(bunch_1.y, bunch_2.y, rtol=1e-3, atol=1e-7)
+    np.testing.assert_allclose(bunch_1.xi, bunch_2.xi, rtol=1e-3, atol=1e-7)
+    np.testing.assert_allclose(bunch_1.px, bunch_2.px, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(bunch_1.py, bunch_2.py, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(bunch_1.pz, bunch_2.pz, rtol=1e-3, atol=1e-3)
+
+
+if __name__ == '__main__':
+    test_field_vs_tm_quadrupole()
\ No newline at end of file