From 3bbd0ac9c36665b0017cdc9045264eb1100d5ab7 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 13 Jul 2022 15:14:22 +0200
Subject: [PATCH 1/3] implementing a simple strategy to track ballistically a
 bunch for prop_distance < z_injection

---
 wake_t/particles/particle_bunch.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/wake_t/particles/particle_bunch.py b/wake_t/particles/particle_bunch.py
index a2cac027..e59930fc 100644
--- a/wake_t/particles/particle_bunch.py
+++ b/wake_t/particles/particle_bunch.py
@@ -18,7 +18,8 @@ class ParticleBunch():
 
     def __init__(self, q, x=None, y=None, xi=None, px=None, py=None, pz=None,
                  bunch_matrix=None, matrix_type='standard', gamma_ref=None,
-                 tags=None, prop_distance=0, t_flight=0, name=None):
+                 tags=None, prop_distance=0, t_flight=0, z_injection=None,
+                 name=None):
         """
         Initialize particle bunch.
 
@@ -73,6 +74,9 @@ def __init__(self, q, x=None, y=None, xi=None, px=None, py=None, pz=None,
         t_flight : float
             Time of flight of the bunch along the beamline.
 
+        z_injection: float (in meters)
+            Particles have a ballistic motion for z<z_injection.
+
         name : str
             Name of the particle bunch. Used for species identification
             in openPMD diagnostics.
@@ -96,6 +100,7 @@ def __init__(self, q, x=None, y=None, xi=None, px=None, py=None, pz=None,
         self.tags = tags
         self.prop_distance = prop_distance
         self.t_flight = t_flight
+        self.z_injection = z_injection
         self.x_ref = 0
         self.theta_ref = 0
         self.set_name(name)
@@ -315,6 +320,10 @@ def evolve(self, fields, t, dt, pusher='rk4'):
             The particle pusher to use. Either 'rk4' or 'boris'. By
             default 'rk4'.
         """
+
+        if self.prop_distance < self.z_injection:
+            fields = []
+
         if pusher == 'rk4':
             apply_rk4_pusher(self, fields, t, dt)
         elif pusher == 'boris':

From c5f69bce41faa85b94e86f098c78ea6e0510361d Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 13 Jul 2022 15:35:52 +0200
Subject: [PATCH 2/3] change on the condition to account for the \xi position
 of the bunch

---
 wake_t/particles/particle_bunch.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/wake_t/particles/particle_bunch.py b/wake_t/particles/particle_bunch.py
index e59930fc..8e407164 100644
--- a/wake_t/particles/particle_bunch.py
+++ b/wake_t/particles/particle_bunch.py
@@ -321,7 +321,7 @@ def evolve(self, fields, t, dt, pusher='rk4'):
             default 'rk4'.
         """
 
-        if self.prop_distance < self.z_injection:
+        if (np.amax(self.xi) + self.prop_distance) < self.z_injection:
             fields = []
 
         if pusher == 'rk4':

From a164bf4502415d8f9a0b9e392407626028233232 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 13 Jul 2022 16:47:59 +0200
Subject: [PATCH 3/3] corrects a bug that made the tests fail

---
 wake_t/particles/particle_bunch.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/wake_t/particles/particle_bunch.py b/wake_t/particles/particle_bunch.py
index 8e407164..b1722be2 100644
--- a/wake_t/particles/particle_bunch.py
+++ b/wake_t/particles/particle_bunch.py
@@ -321,8 +321,9 @@ def evolve(self, fields, t, dt, pusher='rk4'):
             default 'rk4'.
         """
 
-        if (np.amax(self.xi) + self.prop_distance) < self.z_injection:
-            fields = []
+        if self.z_injection is not None:
+            if (np.amax(self.xi) + self.prop_distance) < self.z_injection:
+                fields = []
 
         if pusher == 'rk4':
             apply_rk4_pusher(self, fields, t, dt)