From da44ae8d5b26a4e32f230c37fb3a98946131061e Mon Sep 17 00:00:00 2001 From: delaossa Date: Fri, 14 Apr 2023 14:08:53 +0200 Subject: [PATCH 1/2] Enables passing a function for `auto_dt_bunch` in `PlasmaStage` --- wake_t/beamline_elements/field_element.py | 4 ++-- wake_t/beamline_elements/plasma_stage.py | 26 +++++++++++++++++------ 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/wake_t/beamline_elements/field_element.py b/wake_t/beamline_elements/field_element.py index cecf66db..41a155a0 100644 --- a/wake_t/beamline_elements/field_element.py +++ b/wake_t/beamline_elements/field_element.py @@ -1,4 +1,4 @@ -from typing import Optional, Union, List +from typing import Optional, Union, List, Callable import scipy.constants as ct @@ -48,7 +48,7 @@ def __init__( n_out: Optional[int] = 1, name: Optional[str] = 'field element', fields: Optional[List[Field]] = [], - auto_dt_bunch: Optional[str] = None + auto_dt_bunch: Optional[Callable[[ParticleBunch], float]] = None ) -> None: self.length = length self.bunch_pusher = bunch_pusher diff --git a/wake_t/beamline_elements/plasma_stage.py b/wake_t/beamline_elements/plasma_stage.py index b0280f5b..05821118 100644 --- a/wake_t/beamline_elements/plasma_stage.py +++ b/wake_t/beamline_elements/plasma_stage.py @@ -8,6 +8,7 @@ import wake_t.physics_models.plasma_wakefields as wf from wake_t.fields.base import Field from .field_element import FieldElement +from wake_t.particles.particle_bunch import ParticleBunch wakefield_models = { @@ -39,10 +40,13 @@ class PlasmaStage(FieldElement): the specified fields. Possible values are ``'rk4'`` (Runge-Kutta method of 4th order) or ``'boris'`` (Boris method). dt_bunch : float - The time step for evolving the particle bunches. If ``None``, it - will be automatically set to ``dt = T/(10*2*pi)``, where T is the - smallest expected betatron period of the bunch along the plasma - stage. + The time step for evolving the particle bunches. An adaptive time + step can be used if this parameter is set to ``'auto'`` and a + ``auto_dt_bunch`` function is provided. + By default the automatic function is set to ``dt = T/(10*2*pi)``, + where T is the smallest expected betatron period of the bunch + along the plasma stage. + n_out : int Number of times along the stage in which the particle distribution should be returned (A list with all output bunches is returned @@ -50,6 +54,10 @@ class PlasmaStage(FieldElement): name : str Name of the plasma stage. This is only used for displaying the progress bar during tracking. By default, ``'Plasma stage'``. + auto_dt_bunch : callable, optional + Function used to determine the adaptive time step for bunches in + which the time step is set to ``'auto'``. The function should take + solely a ``ParticleBunch`` as argument. **model_params Keyword arguments which will be given to the wakefield model. Each model requires a different set of parameters. See the documentation @@ -68,10 +76,11 @@ def __init__( density: Union[float, Callable[[float], float]], wakefield_model: Optional[str] = 'simple_blowout', bunch_pusher: Optional[str] = 'rk4', - dt_bunch: Optional[Union[float, int]] = 'auto', + dt_bunch: Optional[Union[float, str]] = 'auto', n_out: Optional[int] = 1, name: Optional[str] = 'Plasma stage', external_fields: Optional[List[Field]] = [], + auto_dt_bunch: Optional[Callable[[ParticleBunch], float]] = None, **model_params ) -> None: self.density = self._get_density_profile(density) @@ -81,6 +90,10 @@ def __init__( if self.wakefield is not None: fields.append(self.wakefield) fields.extend(self.external_fields) + if auto_dt_bunch is not None: + self.auto_dt_bunch = auto_dt_bunch + else: + self.auto_dt_bunch = self._get_optimized_dt super().__init__( length=length, dt_bunch=dt_bunch, @@ -88,7 +101,8 @@ def __init__( n_out=n_out, name=name, fields=fields, - auto_dt_bunch=self._get_optimized_dt + # auto_dt_bunch=self._get_optimized_dt + auto_dt_bunch=self.auto_dt_bunch ) def _get_density_profile(self, density): From 7c54ac9cc7ca6f13c5c6cafbd56723c2c490cdff Mon Sep 17 00:00:00 2001 From: delaossa Date: Fri, 5 Jul 2024 11:20:32 +0200 Subject: [PATCH 2/2] Update changes. --- wake_t/beamline_elements/field_element.py | 6 ++- wake_t/beamline_elements/plasma_stage.py | 53 +++++++++++++++-------- 2 files changed, 39 insertions(+), 20 deletions(-) diff --git a/wake_t/beamline_elements/field_element.py b/wake_t/beamline_elements/field_element.py index e6c8615d..891bae68 100644 --- a/wake_t/beamline_elements/field_element.py +++ b/wake_t/beamline_elements/field_element.py @@ -1,4 +1,4 @@ -from typing import Optional, Union, List, Literal, Callable +from typing import Optional, Union, Callable, List, Literal import scipy.constants as ct @@ -64,7 +64,8 @@ def __init__( n_out: Optional[int] = 1, name: Optional[str] = 'field element', fields: Optional[List[Field]] = [], - auto_dt_bunch: Optional[Callable[[ParticleBunch], float]] = None + auto_dt_bunch: Optional[Callable[[ParticleBunch], float]] = None, + push_bunches_before_diags: Optional[bool] = True, ) -> None: self.length = length self.bunch_pusher = bunch_pusher @@ -73,6 +74,7 @@ def __init__( self.name = name self.fields = fields self.auto_dt_bunch = auto_dt_bunch + self.push_bunches_before_diags = push_bunches_before_diags def track( self, diff --git a/wake_t/beamline_elements/plasma_stage.py b/wake_t/beamline_elements/plasma_stage.py index 10bdaf10..bdbee4b1 100644 --- a/wake_t/beamline_elements/plasma_stage.py +++ b/wake_t/beamline_elements/plasma_stage.py @@ -42,25 +42,41 @@ class PlasmaStage(FieldElement): 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). - dt_bunch : float - The time step for evolving the particle bunches. An adaptive time - step can be used if this parameter is set to ``'auto'`` and a - ``auto_dt_bunch`` function is provided. - By default the automatic function is set to ``dt = T/(10*2*pi)``, - where T is the smallest expected betatron period of the bunch - along the plasma stage. - - n_out : int + dt_bunch : float, str or list of float or str, optional + The time step for evolving the particle bunches. If ``'auto'``, it + will be automatically set to ``dt = T/(10*2*pi)``, where T is the + smallest expected betatron period of the bunch along the plasma + stage. 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. + auto_dt_bunch : callable, optional + Function used to determine the adaptive time step for bunches in + which the time step is set to ``'auto'``. The function should take + solely a ``ParticleBunch`` as argument. + push_bunches_before_diags : bool, optional + Whether to push the bunches before saving them to the diagnostics. + Since the time step of the diagnostics can be different from that + of the bunches, it could happen that the bunches appear in the + diagnostics as they were at the last push, but not at the actual + time of the diagnostics. Setting this parameter to ``True`` + (default) ensures that an additional push is given to all bunches + to evolve them to the diagnostics time before saving. + This additional push will always have a time step smaller than + the the time step of the bunch, so it has no detrimental impact + on the accuracy of the simulation. However, it could make + convergence studies more difficult to interpret, + since the number of pushes will depend on `n_diags`. Therefore, + it is exposed as an option so that it can be disabled if needed. + n_out : int, optional Number of times along the stage in which the particle distribution should be returned (A list with all output bunches is returned after tracking). name : str Name of the plasma stage. This is only used for displaying the progress bar during tracking. By default, ``'Plasma stage'``. - auto_dt_bunch : callable, optional - Function used to determine the adaptive time step for bunches in - which the time step is set to ``'auto'``. The function should take - solely a ``ParticleBunch`` as argument. + external_fields : list of Field, optional + A list of fields to apply to the particle bunches in + addition to the plasma wakefields. **model_params Keyword arguments which will be given to the wakefield model. Each model requires a different set of parameters. See the documentation @@ -78,12 +94,13 @@ def __init__( length: float, density: Union[float, Callable[[float], float]], wakefield_model: Optional[str] = 'simple_blowout', - bunch_pusher: Optional[str] = 'rk4', - dt_bunch: Optional[Union[float, str]] = 'auto', + bunch_pusher: Optional[Literal['boris', 'rk4']] = 'boris', + dt_bunch: Optional[DtBunchType] = 'auto', + auto_dt_bunch: Optional[Callable[[ParticleBunch], float]] = None, + push_bunches_before_diags: Optional[bool] = True, n_out: Optional[int] = 1, name: Optional[str] = 'Plasma stage', external_fields: Optional[List[Field]] = [], - auto_dt_bunch: Optional[Callable[[ParticleBunch], float]] = None, **model_params ) -> None: self.density = self._get_density_profile(density) @@ -104,8 +121,8 @@ def __init__( n_out=n_out, name=name, fields=fields, - # auto_dt_bunch=self._get_optimized_dt - auto_dt_bunch=self.auto_dt_bunch + auto_dt_bunch=self.auto_dt_bunch, + push_bunches_before_diags=push_bunches_before_diags, ) def _get_density_profile(self, density):