From 7d6ebafbfdbf1e97e525d34d0fd025f4a8f851cd Mon Sep 17 00:00:00 2001 From: Anton Golovanov Date: Wed, 24 Apr 2024 14:14:17 +0300 Subject: [PATCH 01/10] Add a FieldQuadrupole element Quadrupole element based on the FieldElement class for tracking a bunch through a quadrupole without approximations --- wake_t/beamline_elements/__init__.py | 4 +- wake_t/beamline_elements/field_quadrupole.py | 50 +++++++++++++++++++ wake_t/physics_models/em_fields/quadrupole.py | 41 +++++++++++++++ 3 files changed, 94 insertions(+), 1 deletion(-) create mode 100755 wake_t/beamline_elements/field_quadrupole.py create mode 100644 wake_t/physics_models/em_fields/quadrupole.py diff --git a/wake_t/beamline_elements/__init__.py b/wake_t/beamline_elements/__init__.py index 87d59310..92c081eb 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 00000000..548a7b6a --- /dev/null +++ b/wake_t/beamline_elements/field_quadrupole.py @@ -0,0 +1,50 @@ +from typing import List, Literal +from .field_element import FieldElement +from wake_t.physics_models.em_fields.quadrupole import QuadrupoleField + + +class FieldQuadrupole(FieldElement): + def __init__( + self, + length: float, + foc_strength: float, + dt_bunch: float | str | List[float | str], + bunch_pusher: Literal['boris', 'rk4'] = 'boris', + n_out: int | None = 1, + name: str | None = 'quadrupole', + ) -> None: + """ + 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 plasma lens in :math:`T/m`. Defined so + that a positive value is focusing for electrons in the :math:`y` + plane and defocusing in the :math:`x` plane. + dt_bunch : float | str | List[float | str] + The time step for evolving the particle bunches. + 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 plasma lens. This is only used for displaying the + progress bar during tracking. By default, 'quadrupole' + """ + 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=None, + ) 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 00000000..341b9f69 --- /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 `y` direction + and defocussing in `x`. 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]) From 0f7d008f29d63e0ed839bc974b71fb3516dca216 Mon Sep 17 00:00:00 2001 From: Anton Golovanov Date: Tue, 30 Apr 2024 11:31:33 +0300 Subject: [PATCH 02/10] FieldQuadrupole - improve docstring - move to the class definition for consistency - rename "plasma lens" to "quadrupole" in the description. --- wake_t/beamline_elements/field_quadrupole.py | 51 ++++++++++---------- 1 file changed, 26 insertions(+), 25 deletions(-) diff --git a/wake_t/beamline_elements/field_quadrupole.py b/wake_t/beamline_elements/field_quadrupole.py index 548a7b6a..9e76c009 100755 --- a/wake_t/beamline_elements/field_quadrupole.py +++ b/wake_t/beamline_elements/field_quadrupole.py @@ -4,6 +4,32 @@ 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:`y` + plane and defocusing in the :math:`x` plane. + dt_bunch : float | str | List[float | str] + The time step for evolving the particle bunches. + 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, @@ -13,31 +39,6 @@ def __init__( n_out: int | None = 1, name: str | None = 'quadrupole', ) -> None: - """ - 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 plasma lens in :math:`T/m`. Defined so - that a positive value is focusing for electrons in the :math:`y` - plane and defocusing in the :math:`x` plane. - dt_bunch : float | str | List[float | str] - The time step for evolving the particle bunches. - 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 plasma lens. This is only used for displaying the - progress bar during tracking. By default, 'quadrupole' - """ self.foc_strength = foc_strength super().__init__( length=length, From 1aad191b53ee92077ce3912b6175fdeb90dacaa7 Mon Sep 17 00:00:00 2001 From: Anton Golovanov Date: Tue, 30 Apr 2024 14:00:53 +0300 Subject: [PATCH 03/10] Fix type annotations for Python<3.10 The | operator for types was introduced only in Python 3.10. --- wake_t/beamline_elements/field_quadrupole.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/wake_t/beamline_elements/field_quadrupole.py b/wake_t/beamline_elements/field_quadrupole.py index 9e76c009..8409ee3d 100755 --- a/wake_t/beamline_elements/field_quadrupole.py +++ b/wake_t/beamline_elements/field_quadrupole.py @@ -1,4 +1,4 @@ -from typing import List, Literal +from typing import List, Literal, Union, Optional from .field_element import FieldElement from wake_t.physics_models.em_fields.quadrupole import QuadrupoleField @@ -34,10 +34,10 @@ def __init__( self, length: float, foc_strength: float, - dt_bunch: float | str | List[float | str], + dt_bunch: Union[float, str, List[Union[float, str]]], bunch_pusher: Literal['boris', 'rk4'] = 'boris', - n_out: int | None = 1, - name: str | None = 'quadrupole', + n_out: Optional[int] = 1, + name: Optional[str] = 'quadrupole', ) -> None: self.foc_strength = foc_strength super().__init__( From 1f2d7380a39c17511675a78c686d8e7ca0ae833e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= Date: Fri, 3 May 2024 10:32:17 +0200 Subject: [PATCH 04/10] Enable automatic time step in `FieldQuadrupole` --- wake_t/beamline_elements/field_quadrupole.py | 30 ++++++++++++++++---- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/wake_t/beamline_elements/field_quadrupole.py b/wake_t/beamline_elements/field_quadrupole.py index 8409ee3d..416a68cf 100755 --- a/wake_t/beamline_elements/field_quadrupole.py +++ b/wake_t/beamline_elements/field_quadrupole.py @@ -1,4 +1,9 @@ -from typing import List, Literal, Union, Optional +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 @@ -15,8 +20,13 @@ class FieldQuadrupole(FieldElement): Focusing strength of the quadrupole in :math:`T/m`. Defined so that a positive value is focusing for electrons in the :math:`y` plane and defocusing in the :math:`x` plane. - dt_bunch : float | str | List[float | str] - The time step for evolving the particle bunches. + 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 @@ -34,7 +44,7 @@ def __init__( self, length: float, foc_strength: float, - dt_bunch: Union[float, str, List[Union[float, str]]], + dt_bunch: Optional[DtBunchType] = 'auto', bunch_pusher: Literal['boris', 'rk4'] = 'boris', n_out: Optional[int] = 1, name: Optional[str] = 'quadrupole', @@ -47,5 +57,15 @@ def __init__( n_out=n_out, name=name, fields=[QuadrupoleField(foc_strength)], - auto_dt_bunch=None, + 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 \ No newline at end of file From 45345782ad9347c18cec8ae4899732899ce31fe9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= Date: Fri, 3 May 2024 10:32:50 +0200 Subject: [PATCH 05/10] Formatting --- wake_t/beamline_elements/field_quadrupole.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wake_t/beamline_elements/field_quadrupole.py b/wake_t/beamline_elements/field_quadrupole.py index 416a68cf..b07ff0f1 100755 --- a/wake_t/beamline_elements/field_quadrupole.py +++ b/wake_t/beamline_elements/field_quadrupole.py @@ -68,4 +68,4 @@ def _get_optimized_dt(self, beam): 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 \ No newline at end of file + return dt From b41c128fda77940ff1f561d82dd517f5a466d475 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= Date: Fri, 3 May 2024 11:14:08 +0200 Subject: [PATCH 06/10] Formatting --- wake_t/beamline_elements/field_quadrupole.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wake_t/beamline_elements/field_quadrupole.py b/wake_t/beamline_elements/field_quadrupole.py index b07ff0f1..8363e3f7 100755 --- a/wake_t/beamline_elements/field_quadrupole.py +++ b/wake_t/beamline_elements/field_quadrupole.py @@ -1,4 +1,4 @@ -from typing import Literal, Optional +from typing import Literal, Optional import numpy as np import scipy.constants as ct From 57f3f557c0b4e2391237c8279f84cd21febf8bb0 Mon Sep 17 00:00:00 2001 From: Anton Golovanov Date: Thu, 9 May 2024 16:00:51 +0300 Subject: [PATCH 07/10] Make foc_gradient sign in FieldQuadrupole consistent with Quadrupole --- wake_t/physics_models/em_fields/quadrupole.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wake_t/physics_models/em_fields/quadrupole.py b/wake_t/physics_models/em_fields/quadrupole.py index 341b9f69..9f4960aa 100644 --- a/wake_t/physics_models/em_fields/quadrupole.py +++ b/wake_t/physics_models/em_fields/quadrupole.py @@ -6,14 +6,14 @@ def b_x(x, y, z, t, bx, constants): """B_x component.""" - k = constants[0] + 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] + k = - constants[0] for i in prange(x.shape[0]): by[i] += k * x[i] From d03e4934d4dffac1c09be18877ff7da4bd43dea5 Mon Sep 17 00:00:00 2001 From: Anton Golovanov Date: Thu, 9 May 2024 16:01:21 +0300 Subject: [PATCH 08/10] 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 From 07346a42304ff647ff40564cd20a8b139e851668 Mon Sep 17 00:00:00 2001 From: Anton Golovanov Date: Thu, 9 May 2024 16:13:54 +0300 Subject: [PATCH 09/10] Fix docstrings after the change in sign --- wake_t/beamline_elements/field_quadrupole.py | 4 ++-- wake_t/physics_models/em_fields/quadrupole.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/wake_t/beamline_elements/field_quadrupole.py b/wake_t/beamline_elements/field_quadrupole.py index 8363e3f7..345f8288 100755 --- a/wake_t/beamline_elements/field_quadrupole.py +++ b/wake_t/beamline_elements/field_quadrupole.py @@ -18,8 +18,8 @@ class FieldQuadrupole(FieldElement): 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:`y` - plane and defocusing in the :math:`x` plane. + 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 diff --git a/wake_t/physics_models/em_fields/quadrupole.py b/wake_t/physics_models/em_fields/quadrupole.py index 9f4960aa..e4e31096 100644 --- a/wake_t/physics_models/em_fields/quadrupole.py +++ b/wake_t/physics_models/em_fields/quadrupole.py @@ -24,12 +24,12 @@ class QuadrupoleField(AnalyticalField): In Cartesian coordinates, the field is given by: ``` - b_x = k * y - b_y = k * x + b_x = - k * y + b_y = - k * x ``` - When `k > 0`, it corresponds to focussing electrons in the `y` direction - and defocussing in `x`. When `k < 0`, the result is the opposite. + 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 ---------- From 6c77ead37cea56b6208317e0056fd0c87cf0245d Mon Sep 17 00:00:00 2001 From: Angel Ferran Pousa Date: Wed, 15 May 2024 17:12:10 +0200 Subject: [PATCH 10/10] Add final newline --- tests/test_field_quadrupole.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_field_quadrupole.py b/tests/test_field_quadrupole.py index 2eb2543a..65e93bdd 100644 --- a/tests/test_field_quadrupole.py +++ b/tests/test_field_quadrupole.py @@ -47,4 +47,4 @@ def test_field_vs_tm_quadrupole(): if __name__ == '__main__': - test_field_vs_tm_quadrupole() \ No newline at end of file + test_field_vs_tm_quadrupole()