FieldQuadrupole element based on FieldElement

.gitignore

@@ -10,6 +10,9 @@ __pycache__
 /examples/tests.py
 /old_stuff
 
+# Test output
+/tests_output
+
 # IDEs
 /.vscode
 

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()

wake_t/beamline_elements/__init__.py

@@ -4,8 +4,10 @@
 from .active_plasma_lens import ActivePlasmaLens
 from .beamline import Beamline
 from .field_element import FieldElement
+from .field_quadrupole import FieldQuadrupole
 
 
 __all__ = [
     'Drift', 'Dipole', 'Quadrupole', 'Sextupole', 'PlasmaStage', 'PlasmaRamp',
-    'ActivePlasmaLens', 'Beamline', 'FieldElement', 'TMElement']
+    'ActivePlasmaLens', 'Beamline', 'FieldElement', 'TMElement',
+    'FieldQuadrupole']

wake_t/beamline_elements/field_quadrupole.py

@@ -0,0 +1,71 @@
+from typing import Literal, Optional
+
+import numpy as np
+import scipy.constants as ct
+
+from .plasma_stage import DtBunchType
+from .field_element import FieldElement
+from wake_t.physics_models.em_fields.quadrupole import QuadrupoleField
+
+
+class FieldQuadrupole(FieldElement):
+    """
+    Class defining a quadrupole as a field element.
+
+    Parameters
+    ----------
+    length : float
+        Length of the quadrupole lens in :math:`m`.
+    foc_strength : float
+        Focusing strength of the quadrupole in :math:`T/m`. Defined so
+        that a positive value is focusing for electrons in the :math:`x`
+        plane and defocusing in the :math:`y` plane.
+    dt_bunch : float, str, or list of float and str
+        The time step for evolving the particle bunches. If ``'auto'``, it will
+        be automatically set to :math:`dt = T/(10*2*pi)`, where T is the
+        betatron period of the particle with the lowest energy in the bunch.
+        A list of values can also be provided. In this case, the list
+        should have the same order as the list of bunches given to the
+        ``track`` method.
+    bunch_pusher : str
+        The pusher used to evolve the particle bunches in time within
+        the specified fields. Possible values are ``'rk4'`` (Runge-Kutta
+        method of 4th order) or ``'boris'`` (Boris method).
+    n_out : int, optional
+        Number of times along the lens in which the particle distribution
+        should be returned (A list with all output bunches is returned
+        after tracking).
+    name : str, optional
+        Name of the quadrupole. This is only used for displaying the
+        progress bar during tracking. By default, 'quadrupole'
+    """
+
+    def __init__(
+        self,
+        length: float,
+        foc_strength: float,
+        dt_bunch: Optional[DtBunchType] = 'auto',
+        bunch_pusher: Literal['boris', 'rk4'] = 'boris',
+        n_out: Optional[int] = 1,
+        name: Optional[str] = 'quadrupole',
+    ) -> None:
+        self.foc_strength = foc_strength
+        super().__init__(
+            length=length,
+            dt_bunch=dt_bunch,
+            bunch_pusher=bunch_pusher,
+            n_out=n_out,
+            name=name,
+            fields=[QuadrupoleField(foc_strength)],
+            auto_dt_bunch=self._get_optimized_dt,
+        )
+
+    def _get_optimized_dt(self, beam):
+        """ Get tracking time step. """
+        # Get minimum gamma in the bunch (assumes px,py << pz).
+        q_over_m = beam.q_species / beam.m_species
+        min_gamma = np.sqrt(np.min(beam.pz)**2 + 1)
+        w_x = np.sqrt(np.abs(q_over_m*ct.c * self.foc_strength/min_gamma))
+        T_x = 1/w_x
+        dt = 0.1*T_x
+        return dt

wake_t/physics_models/em_fields/quadrupole.py

@@ -0,0 +1,41 @@
+""" Defines a magnetic field of a quadrupole """
+
+from wake_t.fields.analytical_field import AnalyticalField
+from wake_t.utilities.numba import prange
+
+
+def b_x(x, y, z, t, bx, constants):
+    """B_x component."""
+    k = - constants[0]
+    for i in prange(x.shape[0]):
+        bx[i] += k * y[i]
+
+
+def b_y(x, y, z, t, by, constants):
+    """B_y component."""
+    k = - constants[0]
+    for i in prange(x.shape[0]):
+        by[i] += k * x[i]
+
+
+class QuadrupoleField(AnalyticalField):
+    """Defines a field of a magnetic quadrupole of constant focusing gradient
+    `k`.
+
+    In Cartesian coordinates, the field is given by:
+    ```
+        b_x = - k * y
+        b_y = - k * x
+    ```
+
+    When `k > 0`, it corresponds to focussing electrons in the `x` direction
+    and defocussing in `y`. When `k < 0`, the result is the opposite.
+
+    Parameters
+    ----------
+    foc_gradient : float
+        Uniform focusing gradient in T/m.
+    """
+
+    def __init__(self, foc_gradient):
+        super().__init__(b_x=b_x, b_y=b_y, constants=[foc_gradient])