diff --git a/.gitignore b/.gitignore index 49d2f21..9110d94 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 0000000..65e93bd --- /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() diff --git a/wake_t/beamline_elements/__init__.py b/wake_t/beamline_elements/__init__.py index 87d5931..92c081e 100644 --- a/wake_t/beamline_elements/__init__.py +++ b/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'] diff --git a/wake_t/beamline_elements/field_quadrupole.py b/wake_t/beamline_elements/field_quadrupole.py new file mode 100755 index 0000000..345f828 --- /dev/null +++ b/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 diff --git a/wake_t/physics_models/em_fields/quadrupole.py b/wake_t/physics_models/em_fields/quadrupole.py new file mode 100644 index 0000000..e4e3109 --- /dev/null +++ b/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])