From 290c774d1820a24eb861de6deb1df2095c6caad6 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Sun, 23 Nov 2025 23:14:35 +0100
Subject: [PATCH 01/84] Added Stage.length_num_beta_osc() to calculate the
 stage length that gives a given number of main beam betatron oscillations.

---
 abel/classes/stage/stage.py | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 7455bdac..78e62f1f 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1323,6 +1323,40 @@ def matched_beta_function_flattop(self, energy):
         return beta_matched(self.plasma_density, energy)
     
 
+    # ==================================================
+    def length_num_beta_osc(self, num_beta_osc, initial_energy):
+        """
+        Calculate the stage length required for the main beam to undergo 
+        ``num_beta_osc`` betatron oscillations.
+
+        Parameters
+        ----------
+        num_beta_osc : int
+            ...
+
+        initial_energy : [eV] float
+            ...
+
+            
+        Returns
+        -------
+        length : [m] float
+            ...
+        """
+
+        from abel.utilities.plasma_physics import k_p
+
+        if not isinstance(num_beta_osc, int) or num_beta_osc < 0:
+            raise ValueError('Number of input betatron oscillations must be a positive integer.')
+
+        phase_advance = num_beta_osc * 2*np.pi
+        phase_advance_factor = phase_advance / k_p(self.plasma_density) * np.sqrt(2/(SI.m_e*SI.c**2))
+
+        length = (phase_advance_factor/2)**2 * SI.e*self.nom_accel_gradient_flattop + np.sqrt(initial_energy*SI.e) * phase_advance_factor
+
+        return length
+    
+
     # ==================================================
     def energy_usage(self):
         return self.driver_source.energy_usage()

From 3535df02e8d8240306ad9c25da660fa2fb87c099 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Sun, 23 Nov 2025 23:40:13 +0100
Subject: [PATCH 02/84] Edited Stage.length_num_beta_osc() and added
 Stage.length2num_beta_osc() to calculate the number of betatron oscillations
 a particle can undergo in the stage.

---
 abel/classes/stage/stage.py | 52 +++++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 8 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 78e62f1f..8f0de744 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1324,19 +1324,22 @@ def matched_beta_function_flattop(self, energy):
     
 
     # ==================================================
-    def length_num_beta_osc(self, num_beta_osc, initial_energy):
+    def length_num_beta_osc(self, num_beta_osc, initial_energy, nom_accel_gradient, q=SI.e, m=SI.m_e):
         """
-        Calculate the stage length required for the main beam to undergo 
-        ``num_beta_osc`` betatron oscillations.
+        Calculate the stage length that gives ``num_beta_osc`` betatron 
+        oscillations for a particle with given energy.
 
         Parameters
         ----------
-        num_beta_osc : int
+        num_beta_osc : float
             ...
 
         initial_energy : [eV] float
             ...
 
+        nom_accel_gradient : [V/m] float
+            ...
+
             
         Returns
         -------
@@ -1346,17 +1349,50 @@ def length_num_beta_osc(self, num_beta_osc, initial_energy):
 
         from abel.utilities.plasma_physics import k_p
 
-        if not isinstance(num_beta_osc, int) or num_beta_osc < 0:
-            raise ValueError('Number of input betatron oscillations must be a positive integer.')
+        if num_beta_osc < 0:
+            raise ValueError('Number of input betatron oscillations must be positive.')
 
         phase_advance = num_beta_osc * 2*np.pi
-        phase_advance_factor = phase_advance / k_p(self.plasma_density) * np.sqrt(2/(SI.m_e*SI.c**2))
+        phase_advance_factor = phase_advance / k_p(self.plasma_density) * np.sqrt(2/(m*SI.c**2))
 
-        length = (phase_advance_factor/2)**2 * SI.e*self.nom_accel_gradient_flattop + np.sqrt(initial_energy*SI.e) * phase_advance_factor
+        length = (phase_advance_factor/2)**2 * q*nom_accel_gradient + np.sqrt(initial_energy*q) * phase_advance_factor
 
         return length
     
 
+    # ==================================================
+    def length2num_beta_osc(self, length, initial_energy, nom_accel_gradient, q=SI.e, m=SI.m_e):
+        """
+        Calculate the number of betatron oscillations a particle can undergo in 
+        the stage.
+
+        Parameters
+        ----------
+        length : [m] float
+            ...
+
+        initial_energy : [eV] float
+            ...
+
+        nom_accel_gradient : [V/m] float
+            ...
+
+            
+        Returns
+        -------
+        length : [m] float
+            ...
+        """
+
+        from abel.utilities.plasma_physics import k_p
+
+        integral = 2*np.sqrt(initial_energy*q + q*nom_accel_gradient*length)/(q*nom_accel_gradient) - 2*np.sqrt(initial_energy*q)/(q*nom_accel_gradient)
+
+        num_beta_osc = k_p(self.plasma_density)*np.sqrt(m*SI.c**2/2) * integral/(2*np.pi)
+
+        return num_beta_osc
+    
+
     # ==================================================
     def energy_usage(self):
         return self.driver_source.energy_usage()

From f02f1545f67724ee6656d8b919077a49092dc80c Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 24 Nov 2025 15:43:14 +0100
Subject: [PATCH 03/84] Added Stage.calc_flattop_num_beta_osc() for calculating
 the number of betatron oscillations that should be performed in the main
 flattop by subtracting the contributions from the ramps.

---
 abel/classes/stage/stage.py | 99 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 98 insertions(+), 1 deletion(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 8f0de744..7488fcab 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1324,7 +1324,16 @@ def matched_beta_function_flattop(self, energy):
     
 
     # ==================================================
-    def length_num_beta_osc(self, num_beta_osc, initial_energy, nom_accel_gradient, q=SI.e, m=SI.m_e):
+    # def set_flattop_length(self, num_beta_osc, initial_energy):
+    #     """
+        
+    #     """
+        
+    #     flattop_num_beta_osc = self.calc_flattop_num_beta_osc(num_beta_osc, initial_energy)
+
+
+    # ==================================================
+    def match_length2num_beta_osc(self, num_beta_osc, initial_energy, nom_accel_gradient, q=SI.e, m=SI.m_e):
         """
         Calculate the stage length that gives ``num_beta_osc`` betatron 
         oscillations for a particle with given energy.
@@ -1360,6 +1369,94 @@ def length_num_beta_osc(self, num_beta_osc, initial_energy, nom_accel_gradient,
         return length
     
 
+    # ==================================================
+    def calc_flattop_num_beta_osc(self, num_beta_osc):
+        """
+        If the stage has uniform ramps, calculates the number of betatron 
+        oscillations that the are performed in the main flattop by subtracting 
+        the contributions from the ramps.
+
+        The contributions from the ramps are calculated from the phase advances 
+        in the ramps as L_ramp/beta_matched_ramp.
+
+        - If the stage does have ramps that have not been fully set up, a 
+        deepcopy of the stage is created to set up its ramps using 
+        :func:`Stage._prepare_ramps() <abel.Stage._prepare_ramps>`. Subse
+
+        - If the stage does not have ramps, will simply return the input total 
+        number of betatron oscillations ``num_beta_osc``.
+
+        
+        Parameters
+        ----------
+        num_beta_osc : float
+            Total number of desired betatron oscillations that the main beam 
+            should perform through the whole plasma stage including ramps.
+
+            
+        Returns
+        -------
+        float
+            The number of betatron oscillations the main beam should undergo in 
+            the flattop stage after the contributions from the ramps (if 
+            applicable) have been subtracted from ``num_beta_osc``.
+        """
+
+        if num_beta_osc < 0:
+            raise ValueError('Number of input betatron oscillations must be positive.')
+
+        if not self.has_ramp():
+            return num_beta_osc
+        
+        # Make a copy of the stage and set up its ramps if they are not set yp
+        ramps_not_set_up = (
+            (self.upramp is not None and self.upramp.length is None) or
+            (self.downramp is not None and self.downramp.length is None)
+        )
+        if ramps_not_set_up:
+            stage_copy = copy.deepcopy(self)
+            stage_copy._prepare_ramps()
+        else: 
+            stage_copy = self
+
+        # Calculate the upramp length and matched beta function
+        if stage_copy.upramp is not None:
+            # if self.upramp.length is None:
+            #     upramp_length = beta_matched(self.plasma_density, initial_energy)*np.pi/(2*np.sqrt(1/self.upramp.ramp_beta_mag))
+            # else:
+            #     upramp_length = self.upramp.length
+            
+            upramp_length = stage_copy.upramp.length
+            upramp_beta = beta_matched(stage_copy.plasma_density, self.nom_energy)*stage_copy.upramp.ramp_beta_mag
+            upramp_phase_advance = upramp_length/upramp_beta
+        else:
+            upramp_phase_advance = 0.0
+        
+        # Calculate the downramp length and matched beta function
+        if stage_copy.nom_energy_gain_flattop is None:
+            raise ValueError('Stage.nom_energy_gain_flattop not set.')
+        else:
+            downramp_input_energy = self.nom_energy+stage_copy.nom_energy_gain_flattop
+
+        if stage_copy.downramp is not None:
+            # if self.downramp.length is None:
+            #     downramp_length = beta_matched(self.plasma_density, downramp_input_energy)*np.pi/(2*np.sqrt(1/self.downramp.ramp_beta_mag))
+            # else:
+            #     downramp_length = self.downramp.length
+
+            downramp_length = stage_copy.downramp.length
+            downramp_beta = beta_matched(stage_copy.plasma_density, downramp_input_energy)*stage_copy.downramp.ramp_beta_mag
+            downramp_phase_advance = downramp_length/downramp_beta
+        else:
+            downramp_phase_advance = 0.0
+
+        # Calculate the phase advance in the flattop stage
+        tot_phase_advance = num_beta_osc * 2*np.pi
+        flattop_phase_advance = tot_phase_advance - upramp_phase_advance - downramp_phase_advance
+
+        return flattop_phase_advance/(2*np.pi)
+        
+    
     # ==================================================
     def length2num_beta_osc(self, length, initial_energy, nom_accel_gradient, q=SI.e, m=SI.m_e):
         """

From 6e86c7a10a4e8b05ef7036ada669c37324868084 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 24 Nov 2025 18:13:04 +0100
Subject: [PATCH 04/84] Edited the docstrings of
 Stage.calc_flattop_num_beta_osc(), and edited
 Stage.flattop_length2num_beta_osc() for calculating the number of betatron
 oscillations when the gradient is small (usually the case for ramps).

---
 abel/classes/stage/stage.py | 84 ++++++++++++++++++++++++-------------
 1 file changed, 56 insertions(+), 28 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 7488fcab..94826cd6 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1372,16 +1372,18 @@ def match_length2num_beta_osc(self, num_beta_osc, initial_energy, nom_accel_grad
     # ==================================================
     def calc_flattop_num_beta_osc(self, num_beta_osc):
         """
-        If the stage has uniform ramps, calculates the number of betatron 
-        oscillations that the are performed in the main flattop by subtracting 
-        the contributions from the ramps.
+        For a given total number of betatron oscillations ``num_beta_osc`` that 
+        an electron with energy ``self.nom_energy`` will undergo across the 
+        whole plasma stage inclusing its uniform ramps, this function calculates 
+        the number of betatron oscillations that should be performed in the main 
+        flattop plasma stage by subtracting the contributions from the ramps.
 
         The contributions from the ramps are calculated from the phase advances 
         in the ramps as L_ramp/beta_matched_ramp.
 
         - If the stage does have ramps that have not been fully set up, a 
         deepcopy of the stage is created to set up its ramps using 
-        :func:`Stage._prepare_ramps() <abel.Stage._prepare_ramps>`. Subse
+        :func:`Stage._prepare_ramps() <abel.Stage._prepare_ramps>`. 
 
         - If the stage does not have ramps, will simply return the input total 
         number of betatron oscillations ``num_beta_osc``.
@@ -1390,14 +1392,14 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         Parameters
         ----------
         num_beta_osc : float
-            Total number of desired betatron oscillations that the main beam 
+            Total number of design betatron oscillations that the electron 
             should perform through the whole plasma stage including ramps.
 
             
         Returns
         -------
         float
-            The number of betatron oscillations the main beam should undergo in 
+            The number of betatron oscillations the electron should undergo in 
             the flattop stage after the contributions from the ramps (if 
             applicable) have been subtracted from ``num_beta_osc``.
         """
@@ -1421,11 +1423,6 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
 
         # Calculate the upramp length and matched beta function
         if stage_copy.upramp is not None:
-            # if self.upramp.length is None:
-            #     upramp_length = beta_matched(self.plasma_density, initial_energy)*np.pi/(2*np.sqrt(1/self.upramp.ramp_beta_mag))
-            # else:
-            #     upramp_length = self.upramp.length
-            
             upramp_length = stage_copy.upramp.length
             upramp_beta = beta_matched(stage_copy.plasma_density, self.nom_energy)*stage_copy.upramp.ramp_beta_mag
             upramp_phase_advance = upramp_length/upramp_beta
@@ -1439,11 +1436,6 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
             downramp_input_energy = self.nom_energy+stage_copy.nom_energy_gain_flattop
 
         if stage_copy.downramp is not None:
-            # if self.downramp.length is None:
-            #     downramp_length = beta_matched(self.plasma_density, downramp_input_energy)*np.pi/(2*np.sqrt(1/self.downramp.ramp_beta_mag))
-            # else:
-            #     downramp_length = self.downramp.length
-
             downramp_length = stage_copy.downramp.length
             downramp_beta = beta_matched(stage_copy.plasma_density, downramp_input_energy)*stage_copy.downramp.ramp_beta_mag
             downramp_phase_advance = downramp_length/downramp_beta
@@ -1458,34 +1450,70 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         
     
     # ==================================================
-    def length2num_beta_osc(self, length, initial_energy, nom_accel_gradient, q=SI.e, m=SI.m_e):
+    def flattop_length2num_beta_osc(self, length_flattop=None, initial_energy=None, nom_accel_gradient_flattop=None, plasma_density=None, q=SI.e, m=SI.m_e):
         """
         Calculate the number of betatron oscillations a particle can undergo in 
-        the stage.
+        the stage (excluding ramps).
 
         Parameters
         ----------
-        length : [m] float
-            ...
+        length_flattop : [m] float, optional
+            Length of a plasma stage excluding ramps that the particle can 
+            perform betatron 
+            oscillations in. Defaults to 
+            ``self.length_flattop``.
 
-        initial_energy : [eV] float
-            ...
+        initial_energy : [eV] float, optional
+            The initial energy of the particle at the start of the plasma stage. 
+            Defaults to 
+            ``self.plasma_density``.
 
-        nom_accel_gradient : [V/m] float
-            ...
+        nom_accel_gradient_flattop : [V/m] float, optional
+            Nominal accelerating gradient of the plasma stage exclusing ramps. Defaults to 
+            ``self.plasma_density``.
+
+        plasma_density : [m^-3] float, optional
+            The plasma density of the plasma stage. Defaults to 
+            ``self.plasma_density``.
+
+        q : [C] float, optional
+            Particle charge. Defaults to elementary charge.
+
+        m : [kg] float, optional
+            Particle mass. Defaults to electron mass.
 
             
         Returns
         -------
-        length : [m] float
-            ...
+        num_beta_osc : float
+            Total number of betatron oscillations that the particle will perform
+            across the plasma stage.
         """
 
         from abel.utilities.plasma_physics import k_p
 
-        integral = 2*np.sqrt(initial_energy*q + q*nom_accel_gradient*length)/(q*nom_accel_gradient) - 2*np.sqrt(initial_energy*q)/(q*nom_accel_gradient)
+        if length_flattop is None:
+            length_flattop = self.length_flattop
+
+        if initial_energy is None:
+            initial_energy = self.nom_energy
+
+        if nom_accel_gradient_flattop is None:
+            nom_accel_gradient_flattop = self.nom_accel_gradient_flattop
+
+        if plasma_density is None:
+            plasma_density = self.plasma_density
+
+        if nom_accel_gradient_flattop < 1e-15: # Need to treat very small gradients separately. Often the case for ramps.
+            if self.parent is not None:
+                beta = beta_matched(self.parent.plasma_density, initial_energy)*self.ramp_beta_mag
+            else:
+                beta = self.matched_beta_function(initial_energy)
+            num_beta_osc = length_flattop/beta/(2*np.pi)
+        else:
+            integral = 2*np.sqrt(initial_energy*q + q*nom_accel_gradient_flattop*length_flattop)/(q*nom_accel_gradient_flattop) - 2*np.sqrt(initial_energy*q)/(q*nom_accel_gradient_flattop)
 
-        num_beta_osc = k_p(self.plasma_density)*np.sqrt(m*SI.c**2/2) * integral/(2*np.pi)
+            num_beta_osc = k_p(plasma_density)*np.sqrt(m*SI.c**2/2) * integral/(2*np.pi)
 
         return num_beta_osc
     

From c72de2b3640cc5c7edec9d67e6645db38474d432 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 24 Nov 2025 18:26:51 +0100
Subject: [PATCH 05/84] Added default parameters to
 Stage.match_length2num_beta_osc() and edited various docstrings.

---
 abel/classes/stage/stage.py | 53 ++++++++++++++++++++++++++-----------
 1 file changed, 38 insertions(+), 15 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 94826cd6..7d51cbb0 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1333,36 +1333,61 @@ def matched_beta_function_flattop(self, energy):
 
 
     # ==================================================
-    def match_length2num_beta_osc(self, num_beta_osc, initial_energy, nom_accel_gradient, q=SI.e, m=SI.m_e):
+    def match_length2num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, q=SI.e, m=SI.m_e):
         """
         Calculate the stage length that gives ``num_beta_osc`` betatron 
-        oscillations for a particle with given energy.
+        oscillations for a particle with given initial energy ``initial_energy`` 
+        in a uniform plasma stage (excluding ramps) with nominal acceleration 
+        gradient ``nom_accel_gradient`` and plasma density ``plasma_density``.
 
         Parameters
         ----------
         num_beta_osc : float
-            ...
+            Total number of design betatron oscillations that the electron 
+            should perform through the plasma stage excluding ramps. 
+
+        initial_energy : [eV] float, optional
+            The initial energy of the particle at the start of the plasma stage. 
+            Defaults to ``self.nom_energy``.
 
-        initial_energy : [eV] float
-            ...
+        nom_accel_gradient : [V/m] float, optional
+            Nominal accelerating gradient of the plasma stage exclusing ramps. 
+            Defaults to ``self.nom_accel_gradient_flattop``.
 
-        nom_accel_gradient : [V/m] float
-            ...
+        plasma_density : [m^-3] float, optional
+            The plasma density of the plasma stage. Defaults to 
+            ``self.plasma_density``.
+
+        q : [C] float, optional
+            Particle charge. Defaults to elementary charge.
+
+        m : [kg] float, optional
+            Particle mass. Defaults to electron mass.
 
             
         Returns
         -------
         length : [m] float
-            ...
+            Length of the plasma stage excluding ramps matched to the given 
+            number of betatron oscillations.
         """
 
         from abel.utilities.plasma_physics import k_p
 
+        if initial_energy is None:
+            initial_energy = self.nom_energy
+
+        if nom_accel_gradient is None:
+            nom_accel_gradient = self.nom_accel_gradient_flattop
+
+        if plasma_density is None:
+            plasma_density = self.plasma_density
+
         if num_beta_osc < 0:
             raise ValueError('Number of input betatron oscillations must be positive.')
 
         phase_advance = num_beta_osc * 2*np.pi
-        phase_advance_factor = phase_advance / k_p(self.plasma_density) * np.sqrt(2/(m*SI.c**2))
+        phase_advance_factor = phase_advance / k_p(plasma_density) * np.sqrt(2/(m*SI.c**2))
 
         length = (phase_advance_factor/2)**2 * q*nom_accel_gradient + np.sqrt(initial_energy*q) * phase_advance_factor
 
@@ -1459,18 +1484,16 @@ def flattop_length2num_beta_osc(self, length_flattop=None, initial_energy=None,
         ----------
         length_flattop : [m] float, optional
             Length of a plasma stage excluding ramps that the particle can 
-            perform betatron 
-            oscillations in. Defaults to 
+            perform betatron scillations in. Defaults to 
             ``self.length_flattop``.
 
         initial_energy : [eV] float, optional
             The initial energy of the particle at the start of the plasma stage. 
-            Defaults to 
-            ``self.plasma_density``.
+            Defaults to ``self.nom_energy``.
 
         nom_accel_gradient_flattop : [V/m] float, optional
-            Nominal accelerating gradient of the plasma stage exclusing ramps. Defaults to 
-            ``self.plasma_density``.
+            Nominal accelerating gradient of the plasma stage exclusing ramps. 
+            Defaults to ``self.nom_accel_gradient_flattop``.
 
         plasma_density : [m^-3] float, optional
             The plasma density of the plasma stage. Defaults to 

From 85b48c3b35e48ea07c4cd3f73b21d9d0faa14e25 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 25 Nov 2025 16:48:38 +0100
Subject: [PATCH 06/84] Added the function phase_advance() to
 abel/utilities/beam_physics.py.

---
 abel/utilities/beam_physics.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/abel/utilities/beam_physics.py b/abel/utilities/beam_physics.py
index e73375e3..9a6dba80 100644
--- a/abel/utilities/beam_physics.py
+++ b/abel/utilities/beam_physics.py
@@ -1091,5 +1091,16 @@ def evolve_chromatic_amplitude(ls, inv_rhos, ks, ms, taus, beta0, alpha0=0, Dx0=
         ax.set_ylabel('W_x')
 
     return W, evolution
-    
 
+
+# =============================================
+def phase_advance(ss, betas):
+    """
+    Calculate the phase advance in one dimesion by using the composite Simpson’s 
+    rule (:func:`scipy.integrate.simpson() <scipy.integrate.simpson>`) to 
+    integrate two arrays containing the location and the beta function.
+    """
+    
+    from scipy import integrate
+    inv_betas = 1/betas
+    return integrate.simpson(y=inv_betas, x=ss)
\ No newline at end of file

From 5171f30ae9db7b707d57085927295537edb1a2a2 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 25 Nov 2025 16:49:50 +0100
Subject: [PATCH 07/84] Added controls for stage parameters that may not be set
 in Stage.flattop_length2num_beta_osc().

---
 abel/classes/stage/stage.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 298c8462..6850fb88 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1516,12 +1516,18 @@ def flattop_length2num_beta_osc(self, length_flattop=None, initial_energy=None,
         from abel.utilities.plasma_physics import k_p
 
         if length_flattop is None:
+            if self.length_flattop is None:
+                raise ValueError('Stage.length_flattop not set.')
             length_flattop = self.length_flattop
 
         if initial_energy is None:
+            if self.nom_energy is None:
+                raise ValueError('Stage.nom_energy not set.')
             initial_energy = self.nom_energy
 
         if nom_accel_gradient_flattop is None:
+            if self.nom_accel_gradient_flattop is None:
+                raise ValueError('Stage.nom_accel_gradient_flattop not set.')
             nom_accel_gradient_flattop = self.nom_accel_gradient_flattop
 
         if plasma_density is None:

From 712059362ec64acdccdfe497ef6abb0cc01eed11 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 26 Nov 2025 17:57:43 +0100
Subject: [PATCH 08/84] Added control to Stage.calc_flattop_num_beta_osc() to
 ensure that the ramps are uniform type.

---
 abel/classes/stage/stage.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 6850fb88..cd3352b4 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1435,6 +1435,9 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         if not self.has_ramp():
             return num_beta_osc
         
+        if self.upramp.ramp_shape != 'uniform' or self.downramp.ramp_shape != 'uniform':
+            raise ValueError('This method assumes uniform ramps.')
+
         # Make a copy of the stage and set up its ramps if they are not set yp
         ramps_not_set_up = (
             (self.upramp is not None and self.upramp.length is None) or

From e0fff13e2f029498b1b10cbe1b13bebd9037b8c3 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 27 Nov 2025 15:54:00 +0100
Subject: [PATCH 09/84] Added Stage.phase_advance_beta_evolution() and in the
 process of modifying Stage.flattop_length2num_beta_osc() and
 Stage.calc_flattop_num_beta_osc().

---
 abel/classes/stage/stage.py | 127 ++++++++++++++++++++++++++++++++----
 1 file changed, 116 insertions(+), 11 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index cd3352b4..33b7fbc0 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1437,6 +1437,9 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         
         if self.upramp.ramp_shape != 'uniform' or self.downramp.ramp_shape != 'uniform':
             raise ValueError('This method assumes uniform ramps.')
+        
+        from abel.utilities.beam_physics import evolve_beta_function
+        from abel.utilities.beam_physics import phase_advance
 
         # Make a copy of the stage and set up its ramps if they are not set yp
         ramps_not_set_up = (
@@ -1449,24 +1452,43 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         else: 
             stage_copy = self
 
-        # Calculate the upramp length and matched beta function
+        # Calculate the upramp length and phase advance
         if stage_copy.upramp is not None:
             upramp_length = stage_copy.upramp.length
-            upramp_beta = beta_matched(stage_copy.plasma_density, self.nom_energy)*stage_copy.upramp.ramp_beta_mag
-            upramp_phase_advance = upramp_length/upramp_beta
+            ramp_length = beta_matched(stage_copy.plasma_density, stage_copy.upramp.nom_energy)*np.pi/(2*np.sqrt(1/stage_copy.upramp.ramp_beta_mag))
+            if np.isclose(upramp_length, ramp_length, rtol=1e-3):
+                upramp_phase_advance = np.pi/2
+            else:
+                g_ion_upramp = SI.e*stage_copy.upramp.plasma_density/(2*SI.epsilon_0)
+                p0 = np.sqrt((stage_copy.upramp.nom_energy*SI.e)**2-(SI.m_e*SI.c**2)**2)/SI.c
+                upramp_beta = beta_matched(stage_copy.plasma_density, self.nom_energy)*stage_copy.upramp.ramp_beta_mag
+                ls = np.array([stage_copy.upramp.length_flattop])
+                ks = np.array([g_ion_upramp*SI.e/SI.c/p0])
+                _, _, beta_evolution = evolve_beta_function(ls=ls, ks=ks, beta0=upramp_beta, fast=False, plot=False)
+                upramp_phase_advance = phase_advance(beta_evolution[0,:], beta_evolution[1,:])
         else:
             upramp_phase_advance = 0.0
         
         # Calculate the downramp length and matched beta function
         if stage_copy.nom_energy_gain_flattop is None:
             raise ValueError('Stage.nom_energy_gain_flattop not set.')
-        else:
-            downramp_input_energy = self.nom_energy+stage_copy.nom_energy_gain_flattop
+        # else:
+        #     downramp_input_energy = self.nom_energy+stage_copy.nom_energy_gain_flattop
 
+        # Calculate the downramp length and phase advance
         if stage_copy.downramp is not None:
             downramp_length = stage_copy.downramp.length
-            downramp_beta = beta_matched(stage_copy.plasma_density, downramp_input_energy)*stage_copy.downramp.ramp_beta_mag
-            downramp_phase_advance = downramp_length/downramp_beta
+            ramp_length = beta_matched(stage_copy.plasma_density, stage_copy.downramp.nom_energy)*np.pi/(2*np.sqrt(1/stage_copy.downramp.ramp_beta_mag))
+            if np.isclose(downramp_length, ramp_length, rtol=1e-3):
+                downramp_phase_advance = np.pi/2
+            else:
+                g_ion_downramp = SI.e*stage_copy.downramp.plasma_density/(2*SI.epsilon_0)
+                p0 = np.sqrt((stage_copy.downramp.nom_energy*SI.e)**2-(SI.m_e*SI.c**2)**2)/SI.c
+                downramp_beta = beta_matched(stage_copy.plasma_density, self.nom_energy)*stage_copy.downramp.ramp_beta_mag
+                ls = np.array([stage_copy.downramp.length_flattop])
+                ks = np.array([g_ion_downramp*SI.e/SI.c/p0])
+                _, _, beta_evolution = evolve_beta_function(ls=ls, ks=ks, beta0=downramp_beta, fast=False, plot=False)
+                downramp_phase_advance = phase_advance(beta_evolution[0,:], beta_evolution[1,:])
         else:
             downramp_phase_advance = 0.0
 
@@ -1537,11 +1559,25 @@ def flattop_length2num_beta_osc(self, length_flattop=None, initial_energy=None,
             plasma_density = self.plasma_density
 
         if nom_accel_gradient_flattop < 1e-15: # Need to treat very small gradients separately. Often the case for ramps.
-            if self.parent is not None:
-                beta = beta_matched(self.parent.plasma_density, initial_energy)*self.ramp_beta_mag
+            if self.parent is not None: # self is a ramp
+
+                return self.phase_advance_beta_evolution()/(2*np.pi)
+            
+                # # Extract correct ramp_beta_mag
+                # if self.ramp_beta_mag is not None:
+                #     ramp_beta_mag = self.ramp_beta_mag
+                # elif self.parent.ramp_beta_mag is not None:
+                #     ramp_beta_mag = self.parent.ramp_beta_mag
+                # else:
+                #     raise ValueError('No ramp_beta_mag defined.')
+                
+                # # Check that the ramp length is set as to give pi/2 phase advance
+                # ramp_length = beta_matched(self.parent.plasma_density, self.nom_energy)*np.pi/(2*np.sqrt(1/ramp_beta_mag))
+                # if np.isclose(self.length_flattop, ramp_length, rtol=1e-5):
+                #     return 1/4
+
             else:
-                beta = self.matched_beta_function(initial_energy)
-            num_beta_osc = length_flattop/beta/(2*np.pi)
+                return self.phase_advance_beta_evolution()/(2*np.pi)
         else:
             integral = 2*np.sqrt(initial_energy*q + q*nom_accel_gradient_flattop*length_flattop)/(q*nom_accel_gradient_flattop) - 2*np.sqrt(initial_energy*q)/(q*nom_accel_gradient_flattop)
 
@@ -1549,6 +1585,75 @@ def flattop_length2num_beta_osc(self, length_flattop=None, initial_energy=None,
 
         return num_beta_osc
     
+    
+    # ==================================================
+    def phase_advance_beta_evolution(self, beta0=None):
+        """
+        Calculate the phase advance in a stage by evolving the beta function 
+        through a single element lattice set up using the stage's focusing 
+        strength and length. The evolved beta function is then integrated along 
+        the stage.
+
+        Parameters
+        ----------
+        beta0 : [m] float
+            The initial beta function at the start of the stage. If ``None``, 
+            will calculate the matched beta function for the flattop stage and 
+            scale it according the the ramp's ramp_beta_mag if ``self`` is a 
+            ramp.
+
+            
+        Returns
+        -------
+        float
+            The phase advance calculated by integrating the beta function 
+            evolution through the stage.
+        """
+
+        from abel.utilities.beam_physics import evolve_beta_function
+        from abel.utilities.beam_physics import phase_advance
+        
+        g_ion = SI.e*self.plasma_density/(2*SI.epsilon_0)
+        p0 = np.sqrt((self.nom_energy*SI.e)**2-(SI.m_e*SI.c**2)**2)/SI.c
+        if beta0 is None:
+            if self.is_upramp():
+                beta0 = self.matched_beta_function(self.nom_energy, match_entrance=True)
+            elif self.is_downramp():
+                beta0 = beta_matched(self.parent.plasma_density, self.nom_energy)
+            else:
+                beta0 = beta_matched(self.plasma_density, self.nom_energy)
+
+        ls = np.array([self.length_flattop])
+        ks = np.array([g_ion*SI.e/SI.c/p0])
+        _, _, beta_evolution = evolve_beta_function(ls=ls, ks=ks, beta0=beta0, fast=False, plot=False)
+        return phase_advance(beta_evolution[0,:], beta_evolution[1,:])
+
+
+        # if self.upramp is not None and self.upramp.ramp_beta_mag is not None:
+        #         return beta_matched(self.plasma_density, energy_incoming)*self.upramp.ramp_beta_mag
+        #     else:
+        #         return beta_matched(self.plasma_density, energy_incoming)
+
+
+        # g_ion_downramp = SI.e*stage_copy.downramp.plasma_density/(2*SI.epsilon_0)
+        # p0 = np.sqrt((stage_copy.downramp.nom_energy*SI.e)**2-(SI.m_e*SI.c**2)**2)/SI.c
+        # downramp_beta = beta_matched(stage_copy.plasma_density, self.nom_energy)*stage_copy.downramp.ramp_beta_mag
+        # ls = np.array([stage_copy.downramp.length_flattop])
+        # ks = np.array([g_ion_downramp*SI.e/SI.c/p0])
+        # _, _, beta_evolution = evolve_beta_function(ls=ls, ks=ks, beta0=downramp_beta, fast=False, plot=False)
+        # downramp_phase_advance = phase_advance(beta_evolution[0,:], beta_evolution[1,:])
+    
+
+    # if ramp.ramp_beta_mag is not None:
+    #         ramp_beta_mag = ramp.ramp_beta_mag
+    #     elif self.ramp_beta_mag is not None:
+    #         ramp_beta_mag = self.ramp_beta_mag
+    #     else:
+    #         raise ValueError('No ramp_beta_mag defined.')
+        
+    #     ramp_length = beta_matched(self.plasma_density, ramp.nom_energy)*np.pi/(2*np.sqrt(1/ramp_beta_mag))
+    #     if ramp_length < 0.0:
+    
 
     # ==================================================
     def energy_usage(self):

From 938add41d1e41a26e2f014b4862e632b4383860b Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 27 Nov 2025 23:04:35 +0100
Subject: [PATCH 10/84] Cleaned up in various betatron phase matching functions
 in abel/classes/stage/stage.py.

---
 abel/classes/stage/stage.py | 74 ++-----------------------------------
 1 file changed, 4 insertions(+), 70 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 33b7fbc0..e8bc370a 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1454,18 +1454,7 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
 
         # Calculate the upramp length and phase advance
         if stage_copy.upramp is not None:
-            upramp_length = stage_copy.upramp.length
-            ramp_length = beta_matched(stage_copy.plasma_density, stage_copy.upramp.nom_energy)*np.pi/(2*np.sqrt(1/stage_copy.upramp.ramp_beta_mag))
-            if np.isclose(upramp_length, ramp_length, rtol=1e-3):
-                upramp_phase_advance = np.pi/2
-            else:
-                g_ion_upramp = SI.e*stage_copy.upramp.plasma_density/(2*SI.epsilon_0)
-                p0 = np.sqrt((stage_copy.upramp.nom_energy*SI.e)**2-(SI.m_e*SI.c**2)**2)/SI.c
-                upramp_beta = beta_matched(stage_copy.plasma_density, self.nom_energy)*stage_copy.upramp.ramp_beta_mag
-                ls = np.array([stage_copy.upramp.length_flattop])
-                ks = np.array([g_ion_upramp*SI.e/SI.c/p0])
-                _, _, beta_evolution = evolve_beta_function(ls=ls, ks=ks, beta0=upramp_beta, fast=False, plot=False)
-                upramp_phase_advance = phase_advance(beta_evolution[0,:], beta_evolution[1,:])
+            upramp_phase_advance = stage_copy.upramp.phase_advance_beta_evolution()
         else:
             upramp_phase_advance = 0.0
         
@@ -1477,18 +1466,7 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
 
         # Calculate the downramp length and phase advance
         if stage_copy.downramp is not None:
-            downramp_length = stage_copy.downramp.length
-            ramp_length = beta_matched(stage_copy.plasma_density, stage_copy.downramp.nom_energy)*np.pi/(2*np.sqrt(1/stage_copy.downramp.ramp_beta_mag))
-            if np.isclose(downramp_length, ramp_length, rtol=1e-3):
-                downramp_phase_advance = np.pi/2
-            else:
-                g_ion_downramp = SI.e*stage_copy.downramp.plasma_density/(2*SI.epsilon_0)
-                p0 = np.sqrt((stage_copy.downramp.nom_energy*SI.e)**2-(SI.m_e*SI.c**2)**2)/SI.c
-                downramp_beta = beta_matched(stage_copy.plasma_density, self.nom_energy)*stage_copy.downramp.ramp_beta_mag
-                ls = np.array([stage_copy.downramp.length_flattop])
-                ks = np.array([g_ion_downramp*SI.e/SI.c/p0])
-                _, _, beta_evolution = evolve_beta_function(ls=ls, ks=ks, beta0=downramp_beta, fast=False, plot=False)
-                downramp_phase_advance = phase_advance(beta_evolution[0,:], beta_evolution[1,:])
+            downramp_phase_advance = stage_copy.downramp.phase_advance_beta_evolution()
         else:
             downramp_phase_advance = 0.0
 
@@ -1500,7 +1478,7 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         
     
     # ==================================================
-    def flattop_length2num_beta_osc(self, length_flattop=None, initial_energy=None, nom_accel_gradient_flattop=None, plasma_density=None, q=SI.e, m=SI.m_e):
+    def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None, nom_accel_gradient_flattop=None, plasma_density=None, q=SI.e, m=SI.m_e):
         """
         Calculate the number of betatron oscillations a particle can undergo in 
         the stage (excluding ramps).
@@ -1559,25 +1537,7 @@ def flattop_length2num_beta_osc(self, length_flattop=None, initial_energy=None,
             plasma_density = self.plasma_density
 
         if nom_accel_gradient_flattop < 1e-15: # Need to treat very small gradients separately. Often the case for ramps.
-            if self.parent is not None: # self is a ramp
-
-                return self.phase_advance_beta_evolution()/(2*np.pi)
-            
-                # # Extract correct ramp_beta_mag
-                # if self.ramp_beta_mag is not None:
-                #     ramp_beta_mag = self.ramp_beta_mag
-                # elif self.parent.ramp_beta_mag is not None:
-                #     ramp_beta_mag = self.parent.ramp_beta_mag
-                # else:
-                #     raise ValueError('No ramp_beta_mag defined.')
-                
-                # # Check that the ramp length is set as to give pi/2 phase advance
-                # ramp_length = beta_matched(self.parent.plasma_density, self.nom_energy)*np.pi/(2*np.sqrt(1/ramp_beta_mag))
-                # if np.isclose(self.length_flattop, ramp_length, rtol=1e-5):
-                #     return 1/4
-
-            else:
-                return self.phase_advance_beta_evolution()/(2*np.pi)
+            return self.phase_advance_beta_evolution()/(2*np.pi)
         else:
             integral = 2*np.sqrt(initial_energy*q + q*nom_accel_gradient_flattop*length_flattop)/(q*nom_accel_gradient_flattop) - 2*np.sqrt(initial_energy*q)/(q*nom_accel_gradient_flattop)
 
@@ -1629,32 +1589,6 @@ def phase_advance_beta_evolution(self, beta0=None):
         return phase_advance(beta_evolution[0,:], beta_evolution[1,:])
 
 
-        # if self.upramp is not None and self.upramp.ramp_beta_mag is not None:
-        #         return beta_matched(self.plasma_density, energy_incoming)*self.upramp.ramp_beta_mag
-        #     else:
-        #         return beta_matched(self.plasma_density, energy_incoming)
-
-
-        # g_ion_downramp = SI.e*stage_copy.downramp.plasma_density/(2*SI.epsilon_0)
-        # p0 = np.sqrt((stage_copy.downramp.nom_energy*SI.e)**2-(SI.m_e*SI.c**2)**2)/SI.c
-        # downramp_beta = beta_matched(stage_copy.plasma_density, self.nom_energy)*stage_copy.downramp.ramp_beta_mag
-        # ls = np.array([stage_copy.downramp.length_flattop])
-        # ks = np.array([g_ion_downramp*SI.e/SI.c/p0])
-        # _, _, beta_evolution = evolve_beta_function(ls=ls, ks=ks, beta0=downramp_beta, fast=False, plot=False)
-        # downramp_phase_advance = phase_advance(beta_evolution[0,:], beta_evolution[1,:])
-    
-
-    # if ramp.ramp_beta_mag is not None:
-    #         ramp_beta_mag = ramp.ramp_beta_mag
-    #     elif self.ramp_beta_mag is not None:
-    #         ramp_beta_mag = self.ramp_beta_mag
-    #     else:
-    #         raise ValueError('No ramp_beta_mag defined.')
-        
-    #     ramp_length = beta_matched(self.plasma_density, ramp.nom_energy)*np.pi/(2*np.sqrt(1/ramp_beta_mag))
-    #     if ramp_length < 0.0:
-    
-
     # ==================================================
     def energy_usage(self):
         return self.driver_source.energy_usage()

From e22e7421ce3615094665e7f1da080661983baddd Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 27 Nov 2025 23:07:41 +0100
Subject: [PATCH 11/84] Reduced the number of time steps used in ramp tracking
 in StageReducedModels.

---
 abel/classes/stage/impl/stage_reduced_models.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/abel/classes/stage/impl/stage_reduced_models.py b/abel/classes/stage/impl/stage_reduced_models.py
index fd285037..1a2a11f9 100644
--- a/abel/classes/stage/impl/stage_reduced_models.py
+++ b/abel/classes/stage/impl/stage_reduced_models.py
@@ -674,9 +674,9 @@ def track_upramp(self, beam0, driver0, shot_path=None):
                 raise TypeError('upramp is not a StageReducedModels.')
             
             # Set a new time step for the ramp
-            n_steps = 25  # Do n_steps time steps in the tracking of the ramp
+            n_steps = 15  # Do n_steps time steps in the tracking of the ramp
             lambda_beta = upramp.matched_beta_function_flattop(beam0_energy) * 2*np.pi  # [m], betatron wavelength
-            upramp.time_step_mod = min(upramp.length_flattop / (lambda_beta*n_steps), 0.02)  # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c.
+            upramp.time_step_mod = min(upramp.length_flattop / (lambda_beta*n_steps), 0.04)  # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c.
 
         elif type(self.upramp) is Stage:
             upramp = self.upramp  # Allow for other types of ramps
@@ -774,9 +774,9 @@ def track_downramp(self, beam0, driver0, shot_path=None):
                 raise TypeError('downramp is not a StageReducedModels.')
             
             # Set a new time step for the ramp
-            n_steps = 25  # Do n_steps time steps in the tracking of the ramp
+            n_steps = 15  # Do n_steps time steps in the tracking of the ramp
             lambda_beta = downramp.matched_beta_function_flattop(beam0_energy) * 2*np.pi  # [m], betatron wavelength
-            downramp.time_step_mod = min(downramp.length_flattop / (lambda_beta*n_steps), 0.02)  # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c.
+            downramp.time_step_mod = min(downramp.length_flattop / (lambda_beta*n_steps), 0.04)  # Step size in in units of betatron wavelength, equivalent to time step size in units of betatron wavelength/c.
             
         elif type(self.downramp) is Stage:
             downramp = self.downramp  # Allow for other types of ramps

From ee98ccb304b6da5523cb359aaad194169500dff7 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 28 Nov 2025 02:48:59 +0100
Subject: [PATCH 12/84] Correction in Stage.match_length2num_beta_osc().

---
 abel/classes/stage/stage.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index e8bc370a..da3a9071 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1389,7 +1389,7 @@ def match_length2num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel
         phase_advance = num_beta_osc * 2*np.pi
         phase_advance_factor = phase_advance / k_p(plasma_density) * np.sqrt(2/(m*SI.c**2))
 
-        length = (phase_advance_factor/2)**2 * q*nom_accel_gradient + np.sqrt(initial_energy*q) * phase_advance_factor
+        length = (phase_advance_factor/2)**2 * q*nom_accel_gradient + np.sqrt(initial_energy*SI.e) * phase_advance_factor
 
         return length
     

From 0d194cd1f17c68af846ffbdc74458890dbff68e5 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 28 Nov 2025 14:49:18 +0100
Subject: [PATCH 13/84] Changed the name of Stage.match_length2num_beta_osc()
 to Stage.calc_length_num_beta_osc() and removed unnecessary comments and
 imports in some of the beatron oscillation length and phase advance class
 methods.

---
 abel/classes/stage/stage.py | 16 +---------------
 1 file changed, 1 insertion(+), 15 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index da3a9071..5b350069 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1321,19 +1321,10 @@ def matched_beta_function_flattop(self, energy):
         '''
 
         return beta_matched(self.plasma_density, energy)
-    
-
-    # ==================================================
-    # def set_flattop_length(self, num_beta_osc, initial_energy):
-    #     """
-        
-    #     """
-        
-    #     flattop_num_beta_osc = self.calc_flattop_num_beta_osc(num_beta_osc, initial_energy)
 
 
     # ==================================================
-    def match_length2num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, q=SI.e, m=SI.m_e):
+    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, q=SI.e, m=SI.m_e):
         """
         Calculate the stage length that gives ``num_beta_osc`` betatron 
         oscillations for a particle with given initial energy ``initial_energy`` 
@@ -1437,9 +1428,6 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         
         if self.upramp.ramp_shape != 'uniform' or self.downramp.ramp_shape != 'uniform':
             raise ValueError('This method assumes uniform ramps.')
-        
-        from abel.utilities.beam_physics import evolve_beta_function
-        from abel.utilities.beam_physics import phase_advance
 
         # Make a copy of the stage and set up its ramps if they are not set yp
         ramps_not_set_up = (
@@ -1461,8 +1449,6 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         # Calculate the downramp length and matched beta function
         if stage_copy.nom_energy_gain_flattop is None:
             raise ValueError('Stage.nom_energy_gain_flattop not set.')
-        # else:
-        #     downramp_input_energy = self.nom_energy+stage_copy.nom_energy_gain_flattop
 
         # Calculate the downramp length and phase advance
         if stage_copy.downramp is not None:

From 138f90b39d3447055fc2ab9cd02f0370e15f63a2 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 1 Dec 2025 16:27:57 +0100
Subject: [PATCH 14/84] Added tests for checking the ramp phase advance and
 setting the stage flattop lengths to match the number of betatron
 oscillations.

---
 tests/test_StageReducedModels_beamline.py | 109 +++++++++++++++-------
 1 file changed, 75 insertions(+), 34 deletions(-)

diff --git a/tests/test_StageReducedModels_beamline.py b/tests/test_StageReducedModels_beamline.py
index f17c6461..4e0e048f 100644
--- a/tests/test_StageReducedModels_beamline.py
+++ b/tests/test_StageReducedModels_beamline.py
@@ -12,7 +12,7 @@
 import pytest
 from abel import *
 import matplotlib
-import shutil
+import shutil, copy
 matplotlib.use('Agg')  # Use a backend that does not display figure to suppress plots.
 
 def setup_trapezoid_driver_source(enable_xy_jitter=False, enable_xpyp_jitter=False):
@@ -72,12 +72,12 @@ def setup_basic_main_source(plasma_density, ramp_beta_mag=1.0, energy=361.8e9):
     return main
 
 
-def setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability=True, enable_radiation_reaction=True, enable_ion_motion=False, use_ramps=False, drive_beam_update_period=0, save_final_step=False):
+def setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, length_flattop=1.56, enable_tr_instability=True, enable_radiation_reaction=True, enable_ion_motion=False, use_ramps=False, drive_beam_update_period=0, save_final_step=False):
     
     stage = StageReducedModels()
-    stage.time_step_mod = 0.03*2                                                      # In units of betatron wavelengths/c.
-    stage.nom_energy_gain = 7.8e9/5                                                # [eV]
-    stage.length_flattop = 7.8/5                                                      # [m]
+    stage.time_step_mod = 0.03*2                                                    # In units of betatron wavelengths/c.
+    stage.length_flattop = length_flattop                                           # [m]
+    stage.nom_energy_gain = stage.length_flattop*1e9                                # [eV]
     stage.plasma_density = plasma_density                                           # [m^-3]
     stage.driver_source = driver_source
     stage.main_source = main_source
@@ -136,11 +136,10 @@ def test_baseline_linac():
     enable_radiation_reaction = False
     enable_ion_motion = False
     use_ramps = False
-    drive_beam_update_period = 0
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag=1.0)
-    stage = setup_StageReducedModels(plasma_density, driver_source, main_source, 1.0, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=1.0, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -208,6 +207,7 @@ def test_ramped_linac():
     """
 
     np.random.seed(42)
+    from abel.utilities.plasma_physics import beta_matched
 
     num_stages = 2
     plasma_density = 6.0e+20                                                        # [m^-3]
@@ -218,14 +218,25 @@ def test_ramped_linac():
     enable_radiation_reaction = False
     enable_ion_motion = False
     use_ramps = True
-    drive_beam_update_period = 0
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
-    main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-    stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+    driver_source.z_offset = 1602e-6                                                # [m]
+    main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=3.0e9)
+
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag,  enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+
+    # Adjust the lengths of the two stages in the linac to match the number of betatron oscillation 
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc=9.5, initial_energy=main_source.energy, nom_accel_gradient=1e9)
+    stage.nom_energy_gain = stage.length_flattop*1e9
+    last_stage = copy.deepcopy(stage)
+    last_stage.length_flattop = last_stage.calc_length_num_beta_osc(num_beta_osc=9.5, initial_energy=main_source.energy+stage.nom_energy_gain, nom_accel_gradient=1e9)
+    last_stage.nom_energy_gain = last_stage.length_flattop*1e9
+
+    # Set up the interstage
     interstage = setup_InterstageImpactX(stage)
 
-    linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
+    # Assemble the linac
+    linac = PlasmaLinac(source=main_source, stage=stage, last_stage=last_stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
 
     # Perform tracking
     linac.run('test_ramped_linac', overwrite=True, verbose=False)
@@ -233,15 +244,16 @@ def test_ramped_linac():
     # Check the machine parameters
     stages = linac.stages
     assert len(stages) == num_stages
-    assert np.isclose(linac.nom_energy, 364920000000.0, rtol=1e-5, atol=0.0)
+    assert np.isclose(linac.nom_energy, 6462752491.36345, rtol=1e-5, atol=0.0)
     assert np.isclose(linac.nom_energy, stages[-1].nom_energy + stages[-1].nom_energy_gain, rtol=1e-5, atol=0.0)
-    assert np.isclose(linac.get_length(), 41.586517149595515, rtol=1e-5, atol=0.0)
+    assert np.isclose(linac.nom_energy, main_source.energy + stages[0].nom_energy_gain + stages[-1].nom_energy_gain, rtol=1e-5, atol=0.0)
+    assert np.isclose(linac.get_length(), 7.775398161868525, rtol=1e-5, atol=0.0)
     
-    assert np.isclose(stages[0].nom_energy, 361.8e9)
-    assert np.isclose(stages[0].nom_energy_gain, 1560000000.0, rtol=1e-15, atol=0.0)
-    assert np.isclose(stages[1].nom_energy, stages[0].nom_energy + stages[0].nom_energy_gain)
-    assert np.isclose(stages[1].nom_energy_gain, 1560000000.0, rtol=1e-15, atol=0.0)
-    assert np.isclose(stages[-1].upramp.length, 0.5747274948097144, rtol=1e-2, atol=0.0)
+    assert np.isclose(stages[0].nom_energy, 3.0e9)
+    assert np.isclose(stages[0].nom_energy_gain, 1567293075.7156029, rtol=1e-15, atol=0.0)
+    assert np.isclose(stages[1].nom_energy, stages[0].nom_energy + stages[0].nom_energy_gain, rtol=1e-15, atol=0.0)
+    assert np.isclose(stages[1].nom_energy_gain, 1895459415.6478472, rtol=1e-15, atol=0.0)
+    assert np.isclose(stages[-1].upramp.length, 0.06443515179509757, rtol=1e-2, atol=0.0)
     assert np.isclose(stages[-1].upramp.length_flattop, stages[-1].upramp.length, rtol=1e-15, atol=0.0)
     assert np.isclose(stages[-1].upramp.nom_energy, stages[-1].nom_energy, rtol=1e-15, atol=0.0)
     assert np.isclose(stages[-1].upramp.nom_energy_flattop, stages[-1].upramp.nom_energy, rtol=1e-15, atol=0.0)
@@ -249,7 +261,7 @@ def test_ramped_linac():
     assert np.isclose(stages[-1].upramp.nom_energy_gain_flattop, stages[-1].upramp.nom_energy_gain, rtol=1e-15, atol=0.0)
     assert np.isclose(stages[-1].upramp.nom_accel_gradient, 0.0, rtol=1e-15, atol=0.0)
     assert np.isclose(stages[-1].upramp.nom_accel_gradient_flattop, stages[-1].upramp.nom_accel_gradient, rtol=1e-15, atol=0.0)
-    assert np.isclose(stages[-1].downramp.length, 0.5759599015745918, rtol=1e-2, atol=0.0)
+    assert np.isclose(stages[-1].downramp.length, 0.07664823833842702, rtol=1e-2, atol=0.0)
     assert np.isclose(stages[-1].downramp.length_flattop, stages[-1].downramp.length, rtol=1e-15, atol=0.0)
     assert np.isclose(stages[-1].downramp.nom_energy, stages[-1].nom_energy + stages[-1].nom_energy_gain, rtol=1e-15, atol=0.0)
     assert np.isclose(stages[-1].downramp.nom_energy_flattop, stages[-1].downramp.nom_energy, rtol=1e-15, atol=0.0)
@@ -278,7 +290,7 @@ def test_ramped_linac():
     assert np.isclose(final_beam.energy(), stages[-1].nom_energy + stages[-1].nom_energy_gain, rtol=1e-2, atol=0.0)
     assert np.isclose(final_beam.bunch_length(), main_source.bunch_length, rtol=1e-1, atol=0.0)
     assert np.isclose(final_beam.charge(), main_source.charge, rtol=1e-5, atol=0.0)
-    assert np.isclose(final_beam.rel_energy_spread(), 0.02071572458689477, rtol=1e-1, atol=0.0)
+    assert np.isclose(final_beam.rel_energy_spread(), 0.01815229983882398, rtol=1e-1, atol=0.0)
 
     nom_beam_size_x = (stages[0].nom_energy/stages[-1].nom_energy)**(1/4)*initial_beam.beam_size_x()
     nom_beam_size_y = (stages[0].nom_energy/stages[-1].nom_energy)**(1/4)*initial_beam.beam_size_y()
@@ -291,6 +303,40 @@ def test_ramped_linac():
     assert np.isclose(final_beam.norm_emittance_x(), main_source.emit_nx, rtol=1e-1, atol=0.0)
     assert np.isclose(final_beam.norm_emittance_y(), main_source.emit_ny, rtol=1e-1, atol=0.0)
 
+    # Check ramp parameters
+    upramp1 = stages[0].upramp
+    downramp1 = stages[0].downramp
+    assert np.isclose(upramp1.plasma_density, stage.plasma_density/upramp1.ramp_beta_mag, rtol=1e-10, atol=0.0)
+    assert np.isclose(downramp1.plasma_density, stage.plasma_density/downramp1.ramp_beta_mag, rtol=1e-10, atol=0.0)
+    upramp1_length = beta_matched(stage.plasma_density, upramp1.nom_energy) * np.pi/2 * np.sqrt(upramp1.ramp_beta_mag)
+    assert np.isclose(upramp1.length_flattop, upramp1_length, rtol=1e-10, atol=0.0)
+    downramp1_length = beta_matched(stage.plasma_density, downramp1.nom_energy) * np.pi/2 * np.sqrt(downramp1.ramp_beta_mag)
+    assert np.isclose(downramp1.length_flattop, downramp1_length, rtol=1e-10, atol=0.0)
+    upramp2 = stages[1].upramp
+    downramp2 = stages[1].downramp
+    assert np.isclose(upramp2.plasma_density, stage.plasma_density/upramp2.ramp_beta_mag, rtol=1e-10, atol=0.0)
+    assert np.isclose(downramp2.plasma_density, stage.plasma_density/downramp2.ramp_beta_mag, rtol=1e-10, atol=0.0)
+    upramp2_length = beta_matched(stage.plasma_density, upramp2.nom_energy) * np.pi/2 * np.sqrt(upramp2.ramp_beta_mag)
+    assert np.isclose(upramp2.length_flattop, upramp2_length, rtol=1e-10, atol=0.0)
+    downramp2_length = beta_matched(stage.plasma_density, downramp2.nom_energy) * np.pi/2 * np.sqrt(downramp2.ramp_beta_mag)
+    assert np.isclose(downramp2.length_flattop, downramp2_length, rtol=1e-10, atol=0.0)
+
+    # Check phase advances in ramps and flattop stage
+    assert np.isclose(upramp1.phase_advance_beta_evolution(), np.pi/2, rtol=1e-3, atol=0.0)  # Calculate the phase advance by integrating the beta function.
+    assert np.isclose(upramp1.length_flattop2num_beta_osc(), 1/4, rtol=1e-3, atol=0.0)  # Calculate the number of betatron oscillations a particle can undergo along the ramp.
+    assert np.isclose(downramp1.phase_advance_beta_evolution(), np.pi/2, rtol=1e-3, atol=0.0)
+    assert np.isclose(downramp1.length_flattop2num_beta_osc(), 1/4, rtol=1e-3, atol=0.0)
+
+    assert np.isclose(upramp2.phase_advance_beta_evolution(), np.pi/2, rtol=1e-3, atol=0.0)  # Calculate the phase advance by integrating the beta function.
+    assert np.isclose(upramp2.length_flattop2num_beta_osc(), 1/4, rtol=1e-3, atol=0.0)  # Calculate the number of betatron oscillations a particle can undergo along the ramp.
+    assert np.isclose(downramp2.phase_advance_beta_evolution(), np.pi/2, rtol=1e-3, atol=0.0)
+    assert np.isclose(downramp2.length_flattop2num_beta_osc(), 1/4, rtol=1e-3, atol=0.0)
+
+    assert np.isclose(stages[0].length_flattop2num_beta_osc(), 9.5, rtol=1e-3, atol=0.0)  # Calculate the number of betatron oscillations a particle can undergo along the stage.
+    assert np.isclose(stages[0].length_flattop2num_beta_osc() + upramp1.length_flattop2num_beta_osc() + downramp1.length_flattop2num_beta_osc(), 10.0, rtol=1e-3, atol=0.0)  # Calculate the total phase adcance.
+    assert np.isclose(stages[1].length_flattop2num_beta_osc(), 9.5, rtol=1e-3, atol=0.0)  # Calculate the number of betatron oscillations a particle can undergo along the stage.
+    assert np.isclose(stages[1].length_flattop2num_beta_osc() + upramp2.length_flattop2num_beta_osc() + downramp2.length_flattop2num_beta_osc(), 10.0, rtol=1e-3, atol=0.0)  # Calculate the total phase adcance.
+
     # Test plotting
     linac.stages[-1].plot_evolution(bunch='beam')
     linac.plot_survey()
@@ -327,7 +373,7 @@ def test_ramped_linac():
 
 #     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
 #     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-#     stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+#     stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, length_flattop=length_flattop, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, drive_beam_update_period)
 #     interstage = setup_InterstageImpactX(stage)
 
 #     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -361,7 +407,7 @@ def test_ramped_linac():
 
 #     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
 #     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-#     stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+#     stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, length_flattop=length_flattop, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, drive_beam_update_period)
 #     interstage = setup_InterstageImpactX(stage)
 
 #     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -399,7 +445,7 @@ def test_angular_jitter_linac():
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-    stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -429,11 +475,10 @@ def test_angular_jitter_ramped_linac():
     enable_radiation_reaction = False
     enable_ion_motion = False
     use_ramps = True
-    drive_beam_update_period = 0
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-    stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -467,12 +512,11 @@ def test_trInstability_linac():
     enable_radiation_reaction = True
     enable_ion_motion = False
     use_ramps = False
-    drive_beam_update_period = 0
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
     main_source.energy = 3.0e9                                                      # [eV], HALHF v2 start energy
-    stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -506,7 +550,7 @@ def test_jitter_trInstability_ramped_linac():
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=3.0e9)
-    stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -631,11 +675,10 @@ def test_ionMotion_linac():
     enable_radiation_reaction = True
     enable_ion_motion = True
     use_ramps = False
-    drive_beam_update_period = 0
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=361.8e9)
-    stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -666,11 +709,10 @@ def test_jitter_trInstability_ionMotion_linac():
     enable_radiation_reaction = True
     enable_ion_motion = True
     use_ramps = False
-    drive_beam_update_period = 0
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=361.8e9)
-    stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -701,11 +743,10 @@ def test_jitter_trInstability_ionMotion_ramped_linac():
     enable_radiation_reaction = True
     enable_ion_motion = True
     use_ramps = True
-    drive_beam_update_period = 0
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=81.0e9)  # Choosing an energy that gives a sensible number of time steps.
-    stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period, save_final_step=True)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, save_final_step=True)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)

From 37249a46b14aa19bfac3a19d46c7dc0134e7ec85 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 4 Dec 2025 01:27:43 +0100
Subject: [PATCH 15/84] Added a class method to StageHipace for estimating the
 orbit that a drive beam with an angular offset will follow when driver
 guiding is applied. Also added getter and setter for
 StageHipace.external_focusing. The setter will also calculate the gradient
 for the external focusing field.

---
 abel/classes/stage/impl/stage_hipace.py | 151 +++++++++++++++++++++++-
 1 file changed, 150 insertions(+), 1 deletion(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index c62046ff..80445142 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -87,7 +87,7 @@ class StageHipace(Stage):
         external magnetic field across the beams. If ``True``, the field 
         gradient of the external field is set to enforce the drive beam to 
         undergo an half-interger number of betatron oscillations along the 
-            stage. Defaults to ``False``.
+        stage. Defaults to ``False``.
 
     mesh_refinement : bool, optional
         Enable HiPACE++ mesh refinement. See the 
@@ -679,6 +679,155 @@ def get_length(self):
             #ss = density_table[:,0]
             return ss.max()-ss.min()
         return super().get_length()
+    
+
+    # # =============================================
+    # @property
+    # def external_focusing_gradient(self):
+    #     return self._external_focusing_gradient
+    # @external_focusing_gradient.setter
+    # def external_focusing_gradient(self, gradient : float | None, num_half_oscillations=1):
+    #     if self.external_focusing and gradient is None:
+    #         #if self.get_length() is None:
+
+    #         # Make a copy of the stage and set up its ramps if they are not set yp
+    #         ramps_not_set_up = (
+    #             (self.upramp is not None and self.upramp.length is None) or
+    #             (self.downramp is not None and self.downramp.length is None)
+    #         )
+    #         if ramps_not_set_up:
+    #             stage_copy = copy.deepcopy(self)
+    #             stage_copy._prepare_ramps()
+    #         else: 
+    #             stage_copy = self
+
+    #         self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=num_half_oscillations)  # [T/m]
+
+    #     elif self.external_focusing and gradient is not None:
+    #         self._external_focusing_gradient = gradient
+    # _external_focusing_gradient = None
+
+
+    # =============================================
+    @property
+    def external_focusing(self) -> bool:
+        return self._external_focusing
+    @external_focusing.setter
+    def external_focusing(self, enable_external_focusing=False):
+        self._external_focusing = enable_external_focusing
+
+        if enable_external_focusing is False:
+            self._external_focusing_gradient = 0.0
+        elif enable_external_focusing and self._external_focusing_gradient is None:
+            #if self.get_length() is None:
+
+            # Make a copy of the stage and set up its ramps if they are not set yp
+            ramps_not_set_up = (
+                (self.upramp is not None and self.upramp.length is None) or
+                (self.downramp is not None and self.downramp.length is None)
+            )
+            if ramps_not_set_up:
+                stage_copy = copy.deepcopy(self)
+                stage_copy._prepare_ramps()
+            else: 
+                stage_copy = self
+
+            self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=1)  # [T/m]
+
+    _external_focusing = False
+
+
+    def calc_external_focusing_gradient(self, num_half_oscillations=1):
+        """
+        Calculate the external focusing gradient g for an azimuthal magnetic 
+        field B=[gy,-gx,0] that gives ``num_half_oscillations`` half 
+        oscillations for the drive beam over the length of the stage.
+        """
+        driver_source = self.get_driver_source()
+        if driver_source is None:
+            raise ValueError('The driver of the stage is not set.')
+        elif driver_source.energy is None:
+            raise ValueError('The energy of the driver source of the stage is not set.')
+        if self.get_length() is None:
+            raise ValueError('Stage length is not set.')
+        #return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.get_length())**2  # [T/m]
+        return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.length_flattop)**2 
+
+
+    # =============================================
+    def driver_guiding_orbit(self, driver, dacc_gradient=0.0, num_steps_per_half_osc=100):
+        """
+        Estimate the orbit in one dimension that the drive beam will follow 
+        when driver guiding with an external linear azimuthal magnetic field is 
+        applied to a drive beam with an initial angular offset. The calculations 
+        are done by integrating simplified equations of motion.
+        """
+
+        from abel.utilities.relativity import energy2momentum
+        from abel.utilities.statistics import weighted_mean
+
+        energy_thres = 10*driver.particle_mass*SI.c**2/SI.e  # [eV], 10 * particle rest energy. Gives beta=0.995.
+        pz_thres = energy2momentum(energy_thres, unit='eV', m=driver.particle_mass)
+        pz0 = energy2momentum(driver.energy(), unit='eV', m=driver.particle_mass)
+
+        if pz0 < pz_thres:
+            raise ValueError('This estimate is only valid for a relativistic beam.')
+        
+        # Make a copy of the stage and set up its ramps if they are not set yp
+        ramps_not_set_up = (
+            (self.upramp is not None and self.upramp.length is None) or
+            (self.downramp is not None and self.downramp.length is None)
+        )
+        if ramps_not_set_up:
+            stage_copy = copy.deepcopy(self)
+            stage_copy._prepare_ramps()
+        else: 
+            stage_copy = self
+        
+        q = driver.particle_charge()  # [C], particle charge including charge sign.
+        g = self._external_focusing_gradient  # [T/m]
+        #num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.get_length()
+        num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.length_flattop
+        ds = self.length_flattop/num_half_oscillations/num_steps_per_half_osc  # [m], step size
+
+        prop_length = 0
+        s_orbit = np.array([0.0])
+        #x = driver.x_offset()
+        y0 = driver.y_offset()
+        y_orbit = np.array([y0])  # [m], records the orbit
+        y = y0
+        py = weighted_mean(driver.pys(), driver.weightings(), clean=False)
+        pz = pz0 # Can add option for deceleration using a gradient
+
+        while prop_length < self.length_flattop:
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            y = y + py/pz*1/2*ds
+
+            # Kick
+            dpy = q*g*y*ds
+            py = py + dpy
+            pz = pz0 + q * dacc_gradient * prop_length/SI.c # dacc_gradient>0
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            y = y + py/pz*1/2*ds
+            s_orbit = np.append(s_orbit, prop_length)
+            y_orbit = np.append(y_orbit, y)
+
+        s_orbit = s_orbit + driver.z_offset()
+
+        return s_orbit, y_orbit
+
+
+
+
+        
+
+
+
+
 
     
     # ==================================================

From 285df8b32f59cccaaf7668582700f5cddb1daf1b Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 4 Dec 2025 01:30:01 +0100
Subject: [PATCH 16/84] Removed commented out code in
 abel/classes/stage/impl/stage_hipace.py.

---
 abel/classes/stage/impl/stage_hipace.py | 30 +------------------------
 1 file changed, 1 insertion(+), 29 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 80445142..818881eb 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -679,33 +679,6 @@ def get_length(self):
             #ss = density_table[:,0]
             return ss.max()-ss.min()
         return super().get_length()
-    
-
-    # # =============================================
-    # @property
-    # def external_focusing_gradient(self):
-    #     return self._external_focusing_gradient
-    # @external_focusing_gradient.setter
-    # def external_focusing_gradient(self, gradient : float | None, num_half_oscillations=1):
-    #     if self.external_focusing and gradient is None:
-    #         #if self.get_length() is None:
-
-    #         # Make a copy of the stage and set up its ramps if they are not set yp
-    #         ramps_not_set_up = (
-    #             (self.upramp is not None and self.upramp.length is None) or
-    #             (self.downramp is not None and self.downramp.length is None)
-    #         )
-    #         if ramps_not_set_up:
-    #             stage_copy = copy.deepcopy(self)
-    #             stage_copy._prepare_ramps()
-    #         else: 
-    #             stage_copy = self
-
-    #         self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=num_half_oscillations)  # [T/m]
-
-    #     elif self.external_focusing and gradient is not None:
-    #         self._external_focusing_gradient = gradient
-    # _external_focusing_gradient = None
 
 
     # =============================================
@@ -731,9 +704,8 @@ def external_focusing(self, enable_external_focusing=False):
                 stage_copy._prepare_ramps()
             else: 
                 stage_copy = self
-
+                
             self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=1)  # [T/m]
-
     _external_focusing = False
 
 

From d3718d514cf11edc58b917dfb7e1ab6acb6ff75e Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 4 Dec 2025 01:48:31 +0100
Subject: [PATCH 17/84] Added another dimension to the driver orbit tracking in
 StageHipace.driver_guiding_orbit().

---
 abel/classes/stage/impl/stage_hipace.py | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 818881eb..7697a4e8 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -704,7 +704,7 @@ def external_focusing(self, enable_external_focusing=False):
                 stage_copy._prepare_ramps()
             else: 
                 stage_copy = self
-                
+
             self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=1)  # [T/m]
     _external_focusing = False
 
@@ -729,10 +729,10 @@ def calc_external_focusing_gradient(self, num_half_oscillations=1):
     # =============================================
     def driver_guiding_orbit(self, driver, dacc_gradient=0.0, num_steps_per_half_osc=100):
         """
-        Estimate the orbit in one dimension that the drive beam will follow 
-        when driver guiding with an external linear azimuthal magnetic field is 
-        applied to a drive beam with an initial angular offset. The calculations 
-        are done by integrating simplified equations of motion.
+        Estimate the orbit that the drive beam will follow when driver guiding 
+        with an external linear azimuthal magnetic field is applied to a drive 
+        beam with an initial angular offset. The calculations  are done by 
+        integrating simplified equations of motion.
         """
 
         from abel.utilities.relativity import energy2momentum
@@ -764,10 +764,13 @@ def driver_guiding_orbit(self, driver, dacc_gradient=0.0, num_steps_per_half_osc
 
         prop_length = 0
         s_orbit = np.array([0.0])
-        #x = driver.x_offset()
+        x0 = driver.x_offset()
+        x_orbit = np.array([x0])  # [m], records the orbit
+        x = x0
         y0 = driver.y_offset()
         y_orbit = np.array([y0])  # [m], records the orbit
         y = y0
+        px = weighted_mean(driver.pxs(), driver.weightings(), clean=False)
         py = weighted_mean(driver.pys(), driver.weightings(), clean=False)
         pz = pz0 # Can add option for deceleration using a gradient
 
@@ -775,22 +778,27 @@ def driver_guiding_orbit(self, driver, dacc_gradient=0.0, num_steps_per_half_osc
 
             # Drift
             prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
             y = y + py/pz*1/2*ds
 
             # Kick
+            dpx = q*g*x*ds
+            px = px + dpx
             dpy = q*g*y*ds
             py = py + dpy
             pz = pz0 + q * dacc_gradient * prop_length/SI.c # dacc_gradient>0
 
             # Drift
             prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
             y = y + py/pz*1/2*ds
             s_orbit = np.append(s_orbit, prop_length)
+            x_orbit = np.append(x_orbit, x)
             y_orbit = np.append(y_orbit, y)
 
         s_orbit = s_orbit + driver.z_offset()
 
-        return s_orbit, y_orbit
+        return s_orbit, x_orbit, y_orbit
 
 
 

From 7c425133defc30a1f54658a930300dc4092ab890 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 8 Dec 2025 12:08:02 +0100
Subject: [PATCH 18/84] Minor edit in abel/classes/stage/impl/stage_hipace.py.

---
 abel/classes/stage/impl/stage_hipace.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 7697a4e8..7f3cebfa 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -689,9 +689,9 @@ def external_focusing(self) -> bool:
     def external_focusing(self, enable_external_focusing=False):
         self._external_focusing = enable_external_focusing
 
-        if enable_external_focusing is False:
+        if self._external_focusing is False or self._external_focusing is None:
             self._external_focusing_gradient = 0.0
-        elif enable_external_focusing and self._external_focusing_gradient is None:
+        elif self._external_focusing and self._external_focusing_gradient is None:
             #if self.get_length() is None:
 
             # Make a copy of the stage and set up its ramps if they are not set yp

From e854afb701ccbc04a3b6fdc8f247b26517822fa2 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 8 Dec 2025 12:38:21 +0100
Subject: [PATCH 19/84] Changed the name StageHipace.driver_guiding_orbit() to
 StageHipace.driver_guiding_trajectory() and edited docstrings.

---
 abel/classes/stage/impl/stage_hipace.py | 72 ++++++++++++++++---------
 1 file changed, 47 insertions(+), 25 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 7f3cebfa..e756e07d 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -727,12 +727,40 @@ def calc_external_focusing_gradient(self, num_half_oscillations=1):
 
 
     # =============================================
-    def driver_guiding_orbit(self, driver, dacc_gradient=0.0, num_steps_per_half_osc=100):
+    def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_half_osc=100):
         """
-        Estimate the orbit that the drive beam will follow when driver guiding 
-        with an external linear azimuthal magnetic field is applied to a drive 
-        beam with an initial angular offset. The calculations  are done by 
+        Estimate the trajectory that the drive beam will follow when driver 
+        guiding with an external linear azimuthal magnetic field is applied to a 
+        drive beam with an initial angular offset. The calculations  are done by 
         integrating simplified equations of motion.
+
+        Parameters
+        ----------
+        driver : ``Beam``
+            The drive beam.
+        
+        dacc_gradient : [V/m] float, optional
+            The decceleration gradient. Drive beam charge * decceleration 
+            gradient must be negative. Defaults to 0.0.
+
+        num_steps_per_half_osc : int, optional
+            Number of calcualtion steps per half-oscillation of the drive beam. 
+            The number of half-oscillations is set in 
+            :func:`StageHipace.calc_external_focusing_gradient() <abel.classes.stage.impl.stage_hipace.calc_external_focusing_gradient>`. 
+            Defaults to 100.
+        
+
+        Returns
+        -------
+        s_trajectory : [m] float
+            Longitudinal coordinate of the drive beam trajectory. Reference is 
+            set at the start of the plasma stage.
+
+        x_trajectory : [m] float
+            x-coordinate of the drive beam trajectory.
+        
+        y_trajectory : [m] float
+            y-coordinate of the drive beam trajectory.
         """
 
         from abel.utilities.relativity import energy2momentum
@@ -745,6 +773,10 @@ def driver_guiding_orbit(self, driver, dacc_gradient=0.0, num_steps_per_half_osc
         if pz0 < pz_thres:
             raise ValueError('This estimate is only valid for a relativistic beam.')
         
+        q = driver.particle_charge()  # [C], particle charge including charge sign.
+        if q * dacc_gradient > 0.0:
+            raise ValueError('Drive beam charge * decceleration gradient must be negative.')
+        
         # Make a copy of the stage and set up its ramps if they are not set yp
         ramps_not_set_up = (
             (self.upramp is not None and self.upramp.length is None) or
@@ -756,19 +788,18 @@ def driver_guiding_orbit(self, driver, dacc_gradient=0.0, num_steps_per_half_osc
         else: 
             stage_copy = self
         
-        q = driver.particle_charge()  # [C], particle charge including charge sign.
         g = self._external_focusing_gradient  # [T/m]
         #num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.get_length()
         num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.length_flattop
         ds = self.length_flattop/num_half_oscillations/num_steps_per_half_osc  # [m], step size
 
         prop_length = 0
-        s_orbit = np.array([0.0])
+        s_trajectory = np.array([0.0])
         x0 = driver.x_offset()
-        x_orbit = np.array([x0])  # [m], records the orbit
+        x_trajectory = np.array([x0])  # [m], records the trajectory
         x = x0
         y0 = driver.y_offset()
-        y_orbit = np.array([y0])  # [m], records the orbit
+        y_trajectory = np.array([y0])  # [m], records the trajectory
         y = y0
         px = weighted_mean(driver.pxs(), driver.weightings(), clean=False)
         py = weighted_mean(driver.pys(), driver.weightings(), clean=False)
@@ -792,22 +823,13 @@ def driver_guiding_orbit(self, driver, dacc_gradient=0.0, num_steps_per_half_osc
             prop_length = prop_length + 1/2*ds
             x = x + px/pz*1/2*ds
             y = y + py/pz*1/2*ds
-            s_orbit = np.append(s_orbit, prop_length)
-            x_orbit = np.append(x_orbit, x)
-            y_orbit = np.append(y_orbit, y)
-
-        s_orbit = s_orbit + driver.z_offset()
-
-        return s_orbit, x_orbit, y_orbit
-
-
-
-
-        
-
-
+            s_trajectory = np.append(s_trajectory, prop_length)
+            x_trajectory = np.append(x_trajectory, x)
+            y_trajectory = np.append(y_trajectory, y)
 
+        s_trajectory = s_trajectory + driver.z_offset()
 
+        return s_trajectory, x_trajectory, y_trajectory
 
     
     # ==================================================
@@ -817,7 +839,7 @@ def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsi
         Applies waterfall function to all HiPACE++ HDF5 output files in 
         ``data_dir``.
 
-         Parameters
+        Parameters
         ----------
         fcns : A list of ``Beam`` class methods
             Beam class profile methods such as ``Beam.current_profile``, 
@@ -846,7 +868,7 @@ def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsi
             
             
         Returns
-        ----------
+        -------
         waterfalls : list of 2D float ndarrays
             Each element in ``waterfalls`` corresponds to the output of one 
             function in ``fcns`` applied across all files (i.e., simulation 
@@ -922,7 +944,7 @@ def plot_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=20, save_
 
 
         Returns
-        ----------
+        -------
         ``None``
         '''
 

From ee8dcb0c770c5a86af75909249b30aa77424b17b Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 8 Dec 2025 12:46:43 +0100
Subject: [PATCH 20/84] Edited a hyperlink in the docstring of
 StageHipace.driver_guiding_trajectory().

---
 abel/classes/stage/impl/stage_hipace.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index e756e07d..6c9d7a2b 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -746,7 +746,7 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
         num_steps_per_half_osc : int, optional
             Number of calcualtion steps per half-oscillation of the drive beam. 
             The number of half-oscillations is set in 
-            :func:`StageHipace.calc_external_focusing_gradient() <abel.classes.stage.impl.stage_hipace.calc_external_focusing_gradient>`. 
+            :meth:`StageHipace.calc_external_focusing_gradient() <abel.StageHipace.calc_external_focusing_gradient>`. 
             Defaults to 100.
         
 

From ca9d97a5c280b11c382f620ae3b83ad4d8846c37 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 9 Dec 2025 00:30:27 +0100
Subject: [PATCH 21/84] Added arc_lengths() for calculating trajectory length
 to abel/utilities/beam_physics.py.

---
 abel/utilities/beam_physics.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/abel/utilities/beam_physics.py b/abel/utilities/beam_physics.py
index 9a6dba80..63fbb0a1 100644
--- a/abel/utilities/beam_physics.py
+++ b/abel/utilities/beam_physics.py
@@ -1103,4 +1103,23 @@ def phase_advance(ss, betas):
     
     from scipy import integrate
     inv_betas = 1/betas
-    return integrate.simpson(y=inv_betas, x=ss)
\ No newline at end of file
+    return integrate.simpson(y=inv_betas, x=ss)
+
+
+# =============================================
+def arc_lengths(s_trajectory, x_trajectory):
+    """
+    Docstring for arc_length
+    
+    :param s_trajectory: Description
+    :param x_trajectory: Description
+    """
+
+    ds = np.diff(s_trajectory)
+    dx = np.diff(x_trajectory)
+
+    length = np.cumsum(np.sqrt(ds**2 + dx**2))
+
+    length = np.insert(length, 0, 0.0)
+
+    return length

From f60e0e5526d6cc09163818707f5acc5f7ed9be55 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 9 Dec 2025 15:45:51 +0100
Subject: [PATCH 22/84] Edited StageHipace.driver_guiding_trajectory() to raise
 an error if the energy depletion is too severe.

---
 abel/classes/stage/impl/stage_hipace.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 6c9d7a2b..2ad6a7fc 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -787,11 +787,16 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
             stage_copy._prepare_ramps()
         else: 
             stage_copy = self
+
+        #L = stage_copy.get_length()  # [m]
+        L = stage_copy.length_flattop  # [m]
+        if pz0 + q * dacc_gradient * L/SI.c < pz_thres:
+            raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')
         
         g = self._external_focusing_gradient  # [T/m]
-        #num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.get_length()
-        num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.length_flattop
-        ds = self.length_flattop/num_half_oscillations/num_steps_per_half_osc  # [m], step size
+        #num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*L
+        num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*L
+        ds = L/num_half_oscillations/num_steps_per_half_osc  # [m], step size
 
         prop_length = 0
         s_trajectory = np.array([0.0])
@@ -805,7 +810,7 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
         py = weighted_mean(driver.pys(), driver.weightings(), clean=False)
         pz = pz0 # Can add option for deceleration using a gradient
 
-        while prop_length < self.length_flattop:
+        while prop_length < L:
 
             # Drift
             prop_length = prop_length + 1/2*ds

From 9f20dbe09da0adbe560fd2dba31edfa7913af5b9 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 12 Dec 2025 01:37:46 +0100
Subject: [PATCH 23/84] Correction in StageHipace.driver_guiding_trajectory()
 for the cases without driver guiding.

---
 abel/classes/stage/impl/stage_hipace.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 2ad6a7fc..1d3b0f49 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -787,16 +787,23 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
             stage_copy._prepare_ramps()
         else: 
             stage_copy = self
-
+        
         #L = stage_copy.get_length()  # [m]
         L = stage_copy.length_flattop  # [m]
+        
         if pz0 + q * dacc_gradient * L/SI.c < pz_thres:
             raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')
         
         g = self._external_focusing_gradient  # [T/m]
-        #num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*L
-        num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*L
-        ds = L/num_half_oscillations/num_steps_per_half_osc  # [m], step size
+        if g is None:
+            g = 0.0
+            num_half_oscillations = 1
+        elif g < 1e-15:
+            num_half_oscillations = 1
+        else:
+            #num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.get_length()
+            num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.length_flattop
+        ds = self.length_flattop/num_half_oscillations/num_steps_per_half_osc  # [m], step size
 
         prop_length = 0
         s_trajectory = np.array([0.0])

From 704e77a9032c67e199408f5e02d8609b19df8aa0 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 15 Dec 2025 20:09:04 +0100
Subject: [PATCH 24/84] Bugfix for setting a transverse gradient for the
 azimuthal magnetic field used for driver guiding in
 abel/classes/stage/impl/stage_hipace.py.

---
 abel/classes/stage/impl/stage_hipace.py | 27 +++++++++----------------
 1 file changed, 9 insertions(+), 18 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 1d3b0f49..a2a2dd5b 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -214,8 +214,8 @@ def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, drive
         self.no_plasma = no_plasma
 
         # external focusing (APL-like) [T/m]
-        self.external_focusing = external_focusing
         self._external_focusing_gradient = None
+        self.external_focusing = external_focusing
 
         # plasma profile
         self.plasma_profile = SimpleNamespace()
@@ -262,13 +262,6 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
         self._prepare_ramps()
         self._make_ramp_profile(tmpfolder)
         
-        # set external focusing
-        if self.external_focusing == False:
-            self._external_focusing_gradient = 0
-        if self.external_focusing == True and self._external_focusing_gradient is None:
-            num_half_oscillations = 1
-            self._external_focusing_gradient = self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.get_length())**2  # [T/m]
-        
         beam0 = beam_incoming
         driver0 = driver_incoming
         
@@ -686,12 +679,12 @@ def get_length(self):
     def external_focusing(self) -> bool:
         return self._external_focusing
     @external_focusing.setter
-    def external_focusing(self, enable_external_focusing=False):
-        self._external_focusing = enable_external_focusing
+    def external_focusing(self, enable_external_focusing : bool | None):
+        self._external_focusing = bool(enable_external_focusing)
 
-        if self._external_focusing is False or self._external_focusing is None:
+        if self._external_focusing is False:
             self._external_focusing_gradient = 0.0
-        elif self._external_focusing and self._external_focusing_gradient is None:
+        elif self._external_focusing_gradient is None or self._external_focusing_gradient < 1e-15:
             #if self.get_length() is None:
 
             # Make a copy of the stage and set up its ramps if they are not set yp
@@ -704,7 +697,7 @@ def external_focusing(self, enable_external_focusing=False):
                 stage_copy._prepare_ramps()
             else: 
                 stage_copy = self
-
+                
             self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=1)  # [T/m]
     _external_focusing = False
 
@@ -716,14 +709,12 @@ def calc_external_focusing_gradient(self, num_half_oscillations=1):
         oscillations for the drive beam over the length of the stage.
         """
         driver_source = self.get_driver_source()
-        if driver_source is None:
-            raise ValueError('The driver of the stage is not set.')
-        elif driver_source.energy is None:
-            raise ValueError('The energy of the driver source of the stage is not set.')
+        #if driver_source.energy is None:
+        #    raise ValueError('The energy of the driver source of the stage is not set.')
         if self.get_length() is None:
             raise ValueError('Stage length is not set.')
         #return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.get_length())**2  # [T/m]
-        return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.length_flattop)**2 
+        return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.length_flattop)**2  # [T/m]
 
 
     # =============================================

From adc578ad76704b16caead4a33d9ce110318c06fd Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 3 Feb 2026 11:08:35 +0100
Subject: [PATCH 25/84] Added controls to Stage.calc_length_num_beta_osc() and
 Stage.length_flattop2num_beta_osc() for ensuring valid inputs.

---
 abel/classes/stage/stage.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 37d81ab3..ab7c31b6 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1350,7 +1350,8 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             ``self.plasma_density``.
 
         q : [C] float, optional
-            Particle charge. Defaults to elementary charge.
+            Particle charge. q * nom_accel_gradient must be positive. Defaults 
+            to elementary charge.
 
         m : [kg] float, optional
             Particle mass. Defaults to electron mass.
@@ -1366,11 +1367,18 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
         from abel.utilities.plasma_physics import k_p
 
         if initial_energy is None:
+            if self.nom_energy is None:
+                raise ValueError('Stage.nom_energy not set.')
             initial_energy = self.nom_energy
 
         if nom_accel_gradient is None:
+            if self.nom_accel_gradient_flattop is None:
+                raise ValueError('Stage.nom_accel_gradient_flattop not set.')
             nom_accel_gradient = self.nom_accel_gradient_flattop
 
+        if q * nom_accel_gradient < 0:
+            raise ValueError('q * nom_accel_gradient must be positive.')
+
         if plasma_density is None:
             plasma_density = self.plasma_density
 
@@ -1489,7 +1497,8 @@ def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None,
             ``self.plasma_density``.
 
         q : [C] float, optional
-            Particle charge. Defaults to elementary charge.
+            Particle charge. q * nom_accel_gradient_flattop must be positive. 
+            Defaults to elementary charge.
 
         m : [kg] float, optional
             Particle mass. Defaults to electron mass.
@@ -1519,12 +1528,16 @@ def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None,
                 raise ValueError('Stage.nom_accel_gradient_flattop not set.')
             nom_accel_gradient_flattop = self.nom_accel_gradient_flattop
 
+        if q * nom_accel_gradient_flattop < 0:
+            raise ValueError('q * nom_accel_gradient_flattop must be positive.')
+
         if plasma_density is None:
             plasma_density = self.plasma_density
 
         if nom_accel_gradient_flattop < 1e-15: # Need to treat very small gradients separately. Often the case for ramps.
             return self.phase_advance_beta_evolution()/(2*np.pi)
         else:
+            
             integral = 2*np.sqrt(initial_energy*q + q*nom_accel_gradient_flattop*length_flattop)/(q*nom_accel_gradient_flattop) - 2*np.sqrt(initial_energy*q)/(q*nom_accel_gradient_flattop)
 
             num_beta_osc = k_p(plasma_density)*np.sqrt(m*SI.c**2/2) * integral/(2*np.pi)

From 0e79de1d031acea95997dd93390f6631fbb56f97 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 3 Feb 2026 11:36:38 +0100
Subject: [PATCH 26/84] Added _external_focusing_gradient as a Stage class
 attribute and modified Stage.length_flattop2num_beta_osc() to take into
 account of the effects of any external guiding fields.

---
 abel/classes/stage/stage.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index ab7c31b6..48584e22 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -126,6 +126,8 @@ def __init__(self, nom_accel_gradient, nom_energy_gain, plasma_density, driver_s
         self.ramp_beta_mag = ramp_beta_mag
         
         self.stage_number = None
+
+        self._external_focusing_gradient = None
         
         # nominal initial energy
         self.nom_energy = None 
@@ -1537,12 +1539,16 @@ def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None,
         if nom_accel_gradient_flattop < 1e-15: # Need to treat very small gradients separately. Often the case for ramps.
             return self.phase_advance_beta_evolution()/(2*np.pi)
         else:
+            g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
+            if self._external_focusing_gradient is not None:
+                g = g + self._external_focusing_gradient
             
-            integral = 2*np.sqrt(initial_energy*q + q*nom_accel_gradient_flattop*length_flattop)/(q*nom_accel_gradient_flattop) - 2*np.sqrt(initial_energy*q)/(q*nom_accel_gradient_flattop)
+            prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient_flattop)
+            energy_scaling = np.sqrt(initial_energy*SI.e + q*nom_accel_gradient_flattop*length_flattop) - np.sqrt(initial_energy*SI.e)
 
-            num_beta_osc = k_p(plasma_density)*np.sqrt(m*SI.c**2/2) * integral/(2*np.pi)
+            num_beta_osc = prefactor * energy_scaling / (2*np.pi)
 
-        return num_beta_osc
+            return num_beta_osc
     
     
     # ==================================================

From 78ffa697a9728a4af27fd32e9b8df43f774393d1 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 5 Feb 2026 15:20:28 +0100
Subject: [PATCH 27/84] Edited Stage.calc_length_num_beta_osc() in such a way
 that a child class can call this method while doing some slight changes with
 its own version of calc_length_num_beta_osc().

---
 abel/classes/stage/stage.py | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 48584e22..4d89329b 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1326,7 +1326,7 @@ def matched_beta_function_flattop(self, energy):
 
 
     # ==================================================
-    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, q=SI.e, m=SI.m_e):
+    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, rhs=None, q=SI.e, m=SI.m_e):
         """
         Calculate the stage length that gives ``num_beta_osc`` betatron 
         oscillations for a particle with given initial energy ``initial_energy`` 
@@ -1351,6 +1351,9 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             The plasma density of the plasma stage. Defaults to 
             ``self.plasma_density``.
 
+        rhs : ...
+            ...
+
         q : [C] float, optional
             Particle charge. q * nom_accel_gradient must be positive. Defaults 
             to elementary charge.
@@ -1366,12 +1369,13 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             number of betatron oscillations.
         """
 
-        from abel.utilities.plasma_physics import k_p
+        from scipy.optimize import fsolve
 
         if initial_energy is None:
             if self.nom_energy is None:
                 raise ValueError('Stage.nom_energy not set.')
             initial_energy = self.nom_energy
+        initial_energy = initial_energy*SI.e  # [J]
 
         if nom_accel_gradient is None:
             if self.nom_accel_gradient_flattop is None:
@@ -1382,15 +1386,27 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             raise ValueError('q * nom_accel_gradient must be positive.')
 
         if plasma_density is None:
+            if self.plasma_density is None:
+                raise ValueError('Stage.plasma_density not set.')
             plasma_density = self.plasma_density
 
         if num_beta_osc < 0:
             raise ValueError('Number of input betatron oscillations must be positive.')
 
-        phase_advance = num_beta_osc * 2*np.pi
-        phase_advance_factor = phase_advance / k_p(plasma_density) * np.sqrt(2/(m*SI.c**2))
+        if rhs is None:
+            def rhs(L):
+                g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
+
+                prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient)
+                energy_scaling = np.sqrt(initial_energy + q*nom_accel_gradient*L) - np.sqrt(initial_energy)
+                return prefactor * energy_scaling
+
+        # Solve 2*np.pi*num_beta_osc = rhs(L)
+        solution = fsolve(lambda L: rhs(L) - 2*np.pi*num_beta_osc , x0=1)
+        length = solution[0]
+
 
-        length = (phase_advance_factor/2)**2 * q*nom_accel_gradient + np.sqrt(initial_energy*SI.e) * phase_advance_factor
+        #self._external_focusing_gradient = self.driver_source.energy/SI.c*(2.0*np.pi/length)**2 ############TODO: delete
 
         return length
     

From da5644ad2cf154c4c3e3772987e801aab60df759 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 5 Feb 2026 15:25:37 +0100
Subject: [PATCH 28/84] Edited StageHipace.calc_length_num_beta_osc() to
 include the contribution from the external fields used for driver guiding
 when calculating the stage length to match a given number of main beam
 betatron oscillations. This class method partially calls the parent class'
 method. Also needed to edit StageHipace.calc_external_focusing_gradient() to
 take a length as an extra optional argument. Also needed to edit
 StageHipace.external_focusing setter to bypass the case for when the stage
 length is not set (in this case, the gradient for the external driver guiding
 field is not defined).

---
 abel/classes/stage/impl/stage_hipace.py | 105 ++++++++++++++++++++++--
 1 file changed, 97 insertions(+), 8 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index a2a2dd5b..b7477bb8 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -697,24 +697,27 @@ def external_focusing(self, enable_external_focusing : bool | None):
                 stage_copy._prepare_ramps()
             else: 
                 stage_copy = self
+            
+            if self.get_length() is not None:
+                self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=1)  # [T/m]
+            else:
+                self._external_focusing_gradient = None
                 
-            self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=1)  # [T/m]
     _external_focusing = False
 
 
-    def calc_external_focusing_gradient(self, num_half_oscillations=1):
+    def calc_external_focusing_gradient(self, num_half_oscillations=1, L=None):
         """
         Calculate the external focusing gradient g for an azimuthal magnetic 
         field B=[gy,-gx,0] that gives ``num_half_oscillations`` half 
         oscillations for the drive beam over the length of the stage.
         """
-        driver_source = self.get_driver_source()
-        #if driver_source.energy is None:
-        #    raise ValueError('The energy of the driver source of the stage is not set.')
-        if self.get_length() is None:
-            raise ValueError('Stage length is not set.')
+        if L is None:
+            if self.get_length() is None:
+                raise ValueError('Stage length is not set.')
+            L = self.length_flattop
         #return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.get_length())**2  # [T/m]
-        return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.length_flattop)**2  # [T/m]
+        return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/L)**2  # [T/m]
 
 
     # =============================================
@@ -833,8 +836,94 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
         s_trajectory = s_trajectory + driver.z_offset()
 
         return s_trajectory, x_trajectory, y_trajectory
+    
+
+    # ==================================================
+    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, num_half_oscillations=1, q=SI.e, m=SI.m_e):
+        """
+        Calculate the stage length that gives ``num_beta_osc`` betatron 
+        oscillations for a particle with given initial energy ``initial_energy`` 
+        in a uniform plasma stage (excluding ramps) with nominal acceleration 
+        gradient ``nom_accel_gradient`` and plasma density ``plasma_density``.
+
+        Parameters
+        ----------
+        num_beta_osc : float
+            Total number of design betatron oscillations that the electron 
+            should perform through the plasma stage excluding ramps. 
+
+        initial_energy : [eV] float, optional
+            The initial energy of the particle at the start of the plasma stage. 
+            Defaults to ``self.nom_energy``.
+
+        nom_accel_gradient : [V/m] float, optional
+            Nominal accelerating gradient of the plasma stage exclusing ramps. 
+            Defaults to ``self.nom_accel_gradient_flattop``.
+
+        plasma_density : [m^-3] float, optional
+            The plasma density of the plasma stage. Defaults to 
+            ``self.plasma_density``.
+
+        num_half_oscillations : ...
+            ...
+
+        q : [C] float, optional
+            Particle charge. q * nom_accel_gradient must be positive. Defaults 
+            to elementary charge.
+
+        m : [kg] float, optional
+            Particle mass. Defaults to electron mass.
+
+            
+        Returns
+        -------
+        length : [m] float
+            Length of the plasma stage excluding ramps matched to the given 
+            number of betatron oscillations.
+        """
+
+        if initial_energy is None:
+            if self.nom_energy is None:
+                raise ValueError('Stage.nom_energy not set.')
+            initial_energy = self.nom_energy
+        initial_energy = initial_energy*SI.e  # [J]
+
+        if nom_accel_gradient is None:
+            if self.nom_accel_gradient_flattop is None:
+                raise ValueError('Stage.nom_accel_gradient_flattop not set.')
+            nom_accel_gradient = self.nom_accel_gradient_flattop
+
+        if q * nom_accel_gradient < 0:
+            raise ValueError('q * nom_accel_gradient must be positive.')
 
+        if plasma_density is None:
+            if self.plasma_density is None:
+                raise ValueError('Stage.plasma_density not set.')
+            plasma_density = self.plasma_density
+
+        if num_beta_osc < 0:
+            raise ValueError('Number of input betatron oscillations must be positive.')
+
+        def rhs(L):
+            g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
+
+            if self.external_focusing:  # Add contribution from external field used for driver guiding
+                g = g + self.calc_external_focusing_gradient(num_half_oscillations=num_half_oscillations, L=L)
+
+            prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient)
+            energy_scaling = np.sqrt(initial_energy + q*nom_accel_gradient*L) - np.sqrt(initial_energy)
+            return prefactor * energy_scaling
+        
+        # Solve 2*np.pi*num_beta_osc = rhs(L)
+        length = super().calc_length_num_beta_osc(num_beta_osc, initial_energy, nom_accel_gradient, plasma_density, rhs=rhs, q=q, m=m)
+
+        # Set the focusing gradient if not already set
+        if self.external_focusing and self._external_focusing_gradient is None:
+            self._external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=num_half_oscillations, L=length)
+
+        return length
     
+
     # ==================================================
     # Apply waterfall function to all beam dump files
     def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsigma=None, args=None):

From e2538f67dd5b3af00be1da8e14a6966612b64afc Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 5 Feb 2026 15:49:13 +0100
Subject: [PATCH 29/84] Decoupled Stage.calc_length_num_beta_osc() and
 StageHipace.calc_length_num_beta_osc() from each other, so that
 StageHipace.calc_length_num_beta_osc() no longer calls the parent method.
 Also changed the way the stage length is calculated inside
 Stage.calc_length_num_beta_osc().

---
 abel/classes/stage/impl/stage_hipace.py | 25 +++++++++++-----------
 abel/classes/stage/stage.py             | 28 +++++--------------------
 2 files changed, 18 insertions(+), 35 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index b7477bb8..e5288bae 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -683,7 +683,7 @@ def external_focusing(self, enable_external_focusing : bool | None):
         self._external_focusing = bool(enable_external_focusing)
 
         if self._external_focusing is False:
-            self._external_focusing_gradient = 0.0
+            self._external_focusing_gradient = 0.0  # TODO: set to None instead?
         elif self._external_focusing_gradient is None or self._external_focusing_gradient < 1e-15:
             #if self.get_length() is None:
 
@@ -702,7 +702,7 @@ def external_focusing(self, enable_external_focusing : bool | None):
                 self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=1)  # [T/m]
             else:
                 self._external_focusing_gradient = None
-                
+
     _external_focusing = False
 
 
@@ -839,7 +839,7 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
     
 
     # ==================================================
-    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, num_half_oscillations=1, q=SI.e, m=SI.m_e):
+    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, driver_half_oscillations=1.0, q=SI.e):
         """
         Calculate the stage length that gives ``num_beta_osc`` betatron 
         oscillations for a particle with given initial energy ``initial_energy`` 
@@ -864,16 +864,14 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             The plasma density of the plasma stage. Defaults to 
             ``self.plasma_density``.
 
-        num_half_oscillations : ...
-            ...
+        driver_half_oscillations : float, optional
+            Number of half betatron oscillations that the drive beam is 
+            expected to perform. Defaults to 1.0.
 
         q : [C] float, optional
             Particle charge. q * nom_accel_gradient must be positive. Defaults 
             to elementary charge.
 
-        m : [kg] float, optional
-            Particle mass. Defaults to electron mass.
-
             
         Returns
         -------
@@ -882,6 +880,8 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             number of betatron oscillations.
         """
 
+        from scipy.optimize import fsolve
+
         if initial_energy is None:
             if self.nom_energy is None:
                 raise ValueError('Stage.nom_energy not set.')
@@ -908,18 +908,19 @@ def rhs(L):
             g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
 
             if self.external_focusing:  # Add contribution from external field used for driver guiding
-                g = g + self.calc_external_focusing_gradient(num_half_oscillations=num_half_oscillations, L=L)
+                g = g + self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=L)
 
             prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient)
             energy_scaling = np.sqrt(initial_energy + q*nom_accel_gradient*L) - np.sqrt(initial_energy)
             return prefactor * energy_scaling
         
         # Solve 2*np.pi*num_beta_osc = rhs(L)
-        length = super().calc_length_num_beta_osc(num_beta_osc, initial_energy, nom_accel_gradient, plasma_density, rhs=rhs, q=q, m=m)
+        solution = fsolve(lambda L: rhs(L) - 2*np.pi*num_beta_osc , x0=1)
+        length = solution[0]
 
-        # Set the focusing gradient if not already set
+        # Set the external focusing gradient for the driver guiding field if not already set
         if self.external_focusing and self._external_focusing_gradient is None:
-            self._external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=num_half_oscillations, L=length)
+            self._external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=length)
 
         return length
     
diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 4d89329b..9e45fd6d 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1326,7 +1326,7 @@ def matched_beta_function_flattop(self, energy):
 
 
     # ==================================================
-    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, rhs=None, q=SI.e, m=SI.m_e):
+    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, q=SI.e):
         """
         Calculate the stage length that gives ``num_beta_osc`` betatron 
         oscillations for a particle with given initial energy ``initial_energy`` 
@@ -1351,16 +1351,10 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             The plasma density of the plasma stage. Defaults to 
             ``self.plasma_density``.
 
-        rhs : ...
-            ...
-
         q : [C] float, optional
             Particle charge. q * nom_accel_gradient must be positive. Defaults 
             to elementary charge.
 
-        m : [kg] float, optional
-            Particle mass. Defaults to electron mass.
-
             
         Returns
         -------
@@ -1369,8 +1363,6 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             number of betatron oscillations.
         """
 
-        from scipy.optimize import fsolve
-
         if initial_energy is None:
             if self.nom_energy is None:
                 raise ValueError('Stage.nom_energy not set.')
@@ -1393,20 +1385,10 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
         if num_beta_osc < 0:
             raise ValueError('Number of input betatron oscillations must be positive.')
 
-        if rhs is None:
-            def rhs(L):
-                g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
+        g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
+        gradient_prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient)
 
-                prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient)
-                energy_scaling = np.sqrt(initial_energy + q*nom_accel_gradient*L) - np.sqrt(initial_energy)
-                return prefactor * energy_scaling
-
-        # Solve 2*np.pi*num_beta_osc = rhs(L)
-        solution = fsolve(lambda L: rhs(L) - 2*np.pi*num_beta_osc , x0=1)
-        length = solution[0]
-
-
-        #self._external_focusing_gradient = self.driver_source.energy/SI.c*(2.0*np.pi/length)**2 ############TODO: delete
+        length = ((2*np.pi*num_beta_osc/gradient_prefactor + np.sqrt(initial_energy))**2 - initial_energy) / (q*nom_accel_gradient)
 
         return length
     
@@ -1558,7 +1540,7 @@ def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None,
             g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
             if self._external_focusing_gradient is not None:
                 g = g + self._external_focusing_gradient
-            
+
             prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient_flattop)
             energy_scaling = np.sqrt(initial_energy*SI.e + q*nom_accel_gradient_flattop*length_flattop) - np.sqrt(initial_energy*SI.e)
 

From d5663c89d04ec52810a8f476f8fbd9323c019890 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 5 Feb 2026 16:15:53 +0100
Subject: [PATCH 30/84] Added more docstrings to
 StageHipace.calc_length_num_beta_osc() and edited StageHipace.track() to
 ensure that StageHipace._external_focusing_gradient is set before running
 HiPACE++.

---
 abel/classes/stage/impl/stage_hipace.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index e5288bae..d8973acd 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -337,6 +337,10 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
         # input file
         filename_input = 'input_file'
         path_input = tmpfolder + filename_input
+
+        if self.external_focusing and self._external_focusing_gradient is None:
+            self._external_focusing_gradient = self.calc_external_focusing_gradient() # Set the gradient for external focusing fields if not already set.
+
         hipace_write_inputs(path_input, filename_beam, filename_driver, self.plasma_density, self.num_steps, time_step, box_range_z, box_size_xy, ion_motion=self.ion_motion, ion_species=self.ion_species, beam_ionization=self.beam_ionization, radiation_reaction=self.radiation_reaction, output_period=output_period, num_cell_xy=self.num_cell_xy, num_cell_z=num_cell_z, driver_only=self.driver_only, density_table_file=density_table_file, no_plasma=self.no_plasma, external_focusing_gradient=self._external_focusing_gradient, mesh_refinement=self.mesh_refinement, do_spin_tracking=self.do_spin_tracking)
         
         
@@ -846,6 +850,14 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
         in a uniform plasma stage (excluding ramps) with nominal acceleration 
         gradient ``nom_accel_gradient`` and plasma density ``plasma_density``.
 
+        Will take into account the contribution from an external linear magnetic 
+        field B=[gy,-gx,0] if :attr:`self._external_focusing_gradient <abel.Stage.StageHipace.external_focusing>` 
+        is set to ``True`` before calling this method.
+
+        Also set :attr:`self._external_focusing_gradient <abel.Stage.StageHipace._external_focusing_gradient>`
+        if it is not already set, and :attr:`self._external_focusing_gradient <abel.Stage.StageHipace.external_focusing>` 
+        is ``True``. 
+
         Parameters
         ----------
         num_beta_osc : float

From 3cd19ff5f1713d78ce5ba4a30ef5a1e11751faf3 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 5 Feb 2026 17:16:31 +0100
Subject: [PATCH 31/84] Created the file tests/test_StageHipace.py and added
 tests for StageHipace.calc_length_num_beta_osc() in
 tests/test_StageHipace.py::test_external_focusing.

---
 tests/test_StageHipace.py | 165 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 165 insertions(+)
 create mode 100644 tests/test_StageHipace.py

diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py
new file mode 100644
index 00000000..96b3b6e9
--- /dev/null
+++ b/tests/test_StageHipace.py
@@ -0,0 +1,165 @@
+# -*- coding: utf-8 -*-
+# This file is part of ABEL
+# Copyright 2025, The ABEL Authors
+# Authors: C.A.Lindstrøm(1), J.B.B.Chen(1), O.G.Finnerud(1), D.Kalvik(1), E.Hørlyk(1), A.Huebl(2), K.N.Sjobak(1), E.Adli(1)
+# Affiliations: 1) University of Oslo, 2) LBNL
+# License: GPL-3.0-or-later
+
+
+"""
+ABEL : StageHipace unit tests
+"""
+
+import pytest
+from abel import *
+#import shutil
+import numpy as np
+
+
+def setup_trapezoid_driver_source(enable_xy_jitter=False, enable_xpyp_jitter=False, x_angle=0.0, y_angle=0.0):
+    driver = SourceTrapezoid()
+    driver.current_head = 0.1e3                                                   # [A]
+    driver.bunch_length = 1050e-6                                                 # [m]
+    driver.z_offset = 1602e-6                                                     # [m]
+    driver.x_angle = x_angle                                                      # [rad]
+    driver.y_angle = y_angle                                                      # [rad]
+
+    driver.num_particles = 1000000                                                 
+    driver.charge = 5.0e10 * -SI.e                                                # [C]
+    driver.energy = 50e9                                                          # [eV] 
+    driver.gaussian_blur = 50e-6                                                  # [m]
+    driver.rel_energy_spread = 0.01                                              
+
+    driver.emit_nx = 50e-6 * driver.energy/4.5e9                                  # [m rad]
+    driver.emit_ny = 100e-6 * driver.energy/4.5e9                                 # [m rad]
+    driver.beta_x, driver.beta_y = 0.5, 0.5                                       # [m]
+
+    if enable_xy_jitter:
+        driver.jitter.x = 100e-9                                                  # [m], std
+        driver.jitter.y = 100e-9                                                  # [m], std
+
+    if enable_xpyp_jitter:
+        driver.jitter.xp = 1.0e-6                                                 # [rad], std
+        driver.jitter.yp = 1.0e-6                                                 # [rad], std
+
+    driver.symmetrize = True
+
+    return driver
+
+
+
+
+
+@pytest.mark.StageHipace
+def test_external_focusing():
+    """
+    Tests for StageHipace.calc_length_num_beta_osc() for accurately matching the 
+    stage length and external driver guiding field gradient to a desired number 
+    of drive beam and main beam betatron oscillation.
+    """
+
+    def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_focusing=False):
+        stage = StageHipace()
+        stage.nom_energy = nom_energy  # [eV]
+        stage.plasma_density = plasma_density  # [m^-3]
+        stage.driver_source = setup_trapezoid_driver_source()
+        stage.external_focusing = external_focusing
+
+        return stage
+
+    num_beta_osc = 4.0  # The number of betatron oscilations that the main beam is expected to perform.
+    nom_accel_gradient = 1e9  # [V/m]
+
+    stage = setup_minimal_StageHipace()
+
+    # ========== Tests without any external fields ==========
+    assert stage.external_focusing is False
+
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, 
+                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=1.0)
+    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    
+    assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
+    assert np.isclose(0.0, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, 
+                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=2.0)
+    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    
+    assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
+    assert np.isclose(0.0, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    # ========== Tests with external fields ==========
+    stage = setup_minimal_StageHipace(external_focusing=True)
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, 
+                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=1.0)
+    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    
+    assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(140.1695315279373, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(3.4268706541720553, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    stage = setup_minimal_StageHipace(external_focusing=True)
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, 
+                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=2.0)
+    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    
+    assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(574.1247168375805, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(3.3865024719148242, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    # ========== With external fields, lower stage nominal energy ==========
+    num_beta_osc2 = 8.0
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, initial_energy=stage.nom_energy, 
+                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=1.0)
+    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    
+    assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(1038.1202586404845, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1.259217324849329, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, initial_energy=stage.nom_energy, 
+                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=2.0)
+    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    
+    assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(5119.8513420067175, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1.13403339406399, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    # ========== With external fields, lower stage nominal energy, lower driver energy ==========
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
+    stage.driver_source.energy = 4.5e9  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, initial_energy=stage.nom_energy, 
+                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=1.0)
+    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    
+    assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(88.3976616856249, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1.2945695557961696, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
+    stage.driver_source.energy = 4.5e9  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, initial_energy=stage.nom_energy, 
+                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=2.0)
+    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    
+    assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(359.3285677857948, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1.2841918239865389, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+
+
+
+    

From 3fc0d16205ca23dd35f6678c58ab082911054096 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 5 Feb 2026 17:17:09 +0100
Subject: [PATCH 32/84] Added the tag StageHipace in pyproject.toml.

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index 5c756543..48b53432 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -66,6 +66,7 @@ markers = [
   "StageBasic_linac: Integration tests for StageBasic linacs",
   "StageQuasistatic2d: Tests for the StageQuasistatic2d class",
   "StageWakeT: Tests for the StageWakeT class",
+  "StageHipace: Tets for the StageHipace class",
   "presets: Tests for collider presets",
   "impactx: Tests for ImpactX",
 ]

From 6273edeb852d8ede86d60361770bf07140746f87 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 5 Feb 2026 17:47:04 +0100
Subject: [PATCH 33/84] Minor edit of Stage.length_flattop2num_beta_osc().

---
 abel/classes/stage/stage.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 9e45fd6d..266b7f02 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1472,11 +1472,15 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         
     
     # ==================================================
-    def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None, nom_accel_gradient_flattop=None, plasma_density=None, q=SI.e, m=SI.m_e):
+    def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None, nom_accel_gradient_flattop=None, plasma_density=None, q=SI.e):
         """
         Calculate the number of betatron oscillations a particle can undergo in 
         the stage (excluding ramps).
 
+        Will take into account the contribution from an external linear magnetic 
+        field B=[gy,-gx,0] if :attr:`self._external_focusing_gradient <abel.Stage.external_focusing>`
+        is not ``None``.
+
         Parameters
         ----------
         length_flattop : [m] float, optional
@@ -1500,9 +1504,6 @@ def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None,
             Particle charge. q * nom_accel_gradient_flattop must be positive. 
             Defaults to elementary charge.
 
-        m : [kg] float, optional
-            Particle mass. Defaults to electron mass.
-
             
         Returns
         -------
@@ -1537,7 +1538,7 @@ def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None,
         if nom_accel_gradient_flattop < 1e-15: # Need to treat very small gradients separately. Often the case for ramps.
             return self.phase_advance_beta_evolution()/(2*np.pi)
         else:
-            g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
+            g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
             if self._external_focusing_gradient is not None:
                 g = g + self._external_focusing_gradient
 

From 871b3c03488c97c906f2ed55b5b93be6460dd1a6 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 11 Feb 2026 21:30:53 +0100
Subject: [PATCH 34/84] Corrected docstring in Stage.matched_beta_function().

---
 abel/classes/stage/stage.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 266b7f02..c0f1d942 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1274,7 +1274,9 @@ def get_cost_breakdown(self):
     def matched_beta_function(self, energy_incoming, match_entrance=True):
         '''
         Calculates the matched beta function of the stage. If there is an 
-        upramp, the beta function is matched to the upramp by default.
+        upramp, the beta function the beta function is magnified by default so 
+        that it shrinks to the correct size when it enters the main flattop 
+        plasma stage.
     
         
         Parameters

From c497708fc5cfd25bfc8f846d2da5f65bdbbd4c6b Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 11 Feb 2026 21:31:57 +0100
Subject: [PATCH 35/84] Added StageHipace.Stage.matched_beta_function() to
 account for the correction to the matched beta function when external
 focusing fields are applied.

---
 abel/classes/stage/impl/stage_hipace.py | 49 +++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index d8973acd..2f7e4366 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -214,6 +214,7 @@ def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, drive
         self.no_plasma = no_plasma
 
         # external focusing (APL-like) [T/m]
+        self.driver_half_oscillations = 1.0  # TODO: make use of this everywhere
         self._external_focusing_gradient = None
         self.external_focusing = external_focusing
 
@@ -676,6 +677,54 @@ def get_length(self):
             #ss = density_table[:,0]
             return ss.max()-ss.min()
         return super().get_length()
+    
+
+    # ==================================================
+    def matched_beta_function(self, energy_incoming, match_entrance=True, q=SI.e):
+        '''
+        Calculates the matched beta function of the stage. If there is an 
+        upramp, the beta function is magnified by default so that it shrinks to 
+        the correct size when it enters the main flattop plasma stage. Also 
+        takes into account external focusing field B=[gy,-gx,0] if present. 
+    
+        
+        Parameters
+        ----------
+        energy_incoming : [eV] float
+            The energy used for matching.
+
+        match_entrance : bool, optional
+            Matches the beta function to the upramp or the stage entrance if 
+            ``True``. Otherwise, will match the beta function to the downramp. 
+            Default set to ``True``.
+
+        q : [C] float, optional
+            Particle charge. Defaults to elementary charge.
+            
+        Returns
+        -------
+        beta_function : [m], float
+            The matched beta function.
+        '''
+        
+        energy_incoming = energy_incoming*SI.e  # [J]
+
+        g = SI.e*self.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
+        if self._external_focusing_gradient is not None:  # Add contribution from external field
+            g = g + self._external_focusing_gradient
+
+        k_beta = np.sqrt(np.abs(q)*g*SI.c/energy_incoming)  # [m^-1], betatron wavenumber.
+
+        if match_entrance:
+            if self.upramp is not None and self.upramp.ramp_beta_mag is not None:
+                return 1/k_beta * self.upramp.ramp_beta_mag
+            else:
+                return 1/k_beta
+        else:
+            if self.downramp.ramp_beta_mag is not None:
+                return 1/k_beta * self.downramp.ramp_beta_mag
+            else:
+                raise ValueError('Downramp ramp_beta_mag not defined.')
 
 
     # =============================================

From 9ac14f124113a7a60f134d958fbeffe29eb11bc4 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 11 Feb 2026 21:45:45 +0100
Subject: [PATCH 36/84] Edited class methods of StageHipace to use the class
 attribute driver_half_oscillations when determining the number of half
 oscillations the driver should perform under driver guiding.

---
 abel/classes/stage/impl/stage_hipace.py | 27 ++++++++++++-------------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 2f7e4366..0a389872 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -214,7 +214,7 @@ def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, drive
         self.no_plasma = no_plasma
 
         # external focusing (APL-like) [T/m]
-        self.driver_half_oscillations = 1.0  # TODO: make use of this everywhere
+        self.driver_half_oscillations = 1.0 
         self._external_focusing_gradient = None
         self.external_focusing = external_focusing
 
@@ -752,14 +752,14 @@ def external_focusing(self, enable_external_focusing : bool | None):
                 stage_copy = self
             
             if self.get_length() is not None:
-                self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=1)  # [T/m]
+                self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=self.driver_half_oscillations)  # [T/m]
             else:
                 self._external_focusing_gradient = None
 
     _external_focusing = False
 
 
-    def calc_external_focusing_gradient(self, num_half_oscillations=1, L=None):
+    def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None):
         """
         Calculate the external focusing gradient g for an azimuthal magnetic 
         field B=[gy,-gx,0] that gives ``num_half_oscillations`` half 
@@ -770,6 +770,10 @@ def calc_external_focusing_gradient(self, num_half_oscillations=1, L=None):
                 raise ValueError('Stage length is not set.')
             L = self.length_flattop
         #return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.get_length())**2  # [T/m]
+
+        if num_half_oscillations is None:
+            num_half_oscillations = self.driver_half_oscillations
+
         return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/L)**2  # [T/m]
 
 
@@ -842,15 +846,7 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
             raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')
         
         g = self._external_focusing_gradient  # [T/m]
-        if g is None:
-            g = 0.0
-            num_half_oscillations = 1
-        elif g < 1e-15:
-            num_half_oscillations = 1
-        else:
-            #num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.get_length()
-            num_half_oscillations = np.sqrt(g*SI.c/stage_copy.driver_source.energy)/np.pi*stage_copy.length_flattop
-        ds = self.length_flattop/num_half_oscillations/num_steps_per_half_osc  # [m], step size
+        ds = self.length_flattop/self.driver_half_oscillations/num_steps_per_half_osc  # [m], step size
 
         prop_length = 0
         s_trajectory = np.array([0.0])
@@ -892,7 +888,7 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
     
 
     # ==================================================
-    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, driver_half_oscillations=1.0, q=SI.e):
+    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, driver_half_oscillations=None, q=SI.e):
         """
         Calculate the stage length that gives ``num_beta_osc`` betatron 
         oscillations for a particle with given initial energy ``initial_energy`` 
@@ -927,7 +923,7 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
 
         driver_half_oscillations : float, optional
             Number of half betatron oscillations that the drive beam is 
-            expected to perform. Defaults to 1.0.
+            expected to perform. Defaults to ``self.driver_half_oscillations``.
 
         q : [C] float, optional
             Particle charge. q * nom_accel_gradient must be positive. Defaults 
@@ -964,6 +960,9 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
 
         if num_beta_osc < 0:
             raise ValueError('Number of input betatron oscillations must be positive.')
+        
+        if driver_half_oscillations is None:
+            driver_half_oscillations = self.driver_half_oscillations
 
         def rhs(L):
             g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient

From 23d0d09c56adf09d65ea1d4cdd9a2dac4a7f9f98 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 12 Feb 2026 12:04:57 +0100
Subject: [PATCH 37/84] Minor docstring correction in
 tests/test_StageHipace.py.

---
 tests/test_StageHipace.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py
index 96b3b6e9..f907a516 100644
--- a/tests/test_StageHipace.py
+++ b/tests/test_StageHipace.py
@@ -53,7 +53,7 @@ def setup_trapezoid_driver_source(enable_xy_jitter=False, enable_xpyp_jitter=Fal
 @pytest.mark.StageHipace
 def test_external_focusing():
     """
-    Tests for StageHipace.calc_length_num_beta_osc() for accurately matching the 
+    Tests for ``StageHipace.calc_length_num_beta_osc()`` for accurately matching the 
     stage length and external driver guiding field gradient to a desired number 
     of drive beam and main beam betatron oscillation.
     """

From 2f8b1e3c2f53f8cd98749dd3ae0c836dad6ccf4d Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 13 Feb 2026 16:51:58 +0100
Subject: [PATCH 38/84] Added Stage.get_ramp_length().

---
 abel/classes/stage/stage.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index c0f1d942..a5819ae3 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -510,6 +510,26 @@ def _calc_ramp_length(self, ramp : Self) -> float:
         return ramp_length
     
 
+     # ==================================================
+    def get_ramp_length(self) -> float:
+        """
+        Get the length of the ramps if the stage has ramps. Returns 0.0 
+        otherwise.
+        """
+
+        if self.has_ramp():
+            if self.upramp is not None and self.upramp.length is not None:
+                upramp_length = self.upramp.length
+            else:
+                upramp_length = 0.0
+            if self.downramp is not None and self.downramp.length is not None:
+                downramp_length = self.downramp.length
+            else:
+                downramp_length = 0.0
+            return upramp_length + downramp_length
+        else:
+            return 0.0
+
     # ==================================================
     @property
     def parent(self) -> Self | None:
@@ -1400,7 +1420,7 @@ def calc_flattop_num_beta_osc(self, num_beta_osc):
         """
         For a given total number of betatron oscillations ``num_beta_osc`` that 
         an electron with energy ``self.nom_energy`` will undergo across the 
-        whole plasma stage inclusing its uniform ramps, this function calculates 
+        whole plasma stage including its uniform ramps, this function calculates 
         the number of betatron oscillations that should be performed in the main 
         flattop plasma stage by subtracting the contributions from the ramps.
 

From d92000c1c66e7a2ddd3c24529ca4fe56c883e4cf Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 13 Feb 2026 17:07:39 +0100
Subject: [PATCH 39/84] Edited class methods in StageHipace for taking into
 account the length of ramps when calculating the external field gradient for
 driver guiding.

---
 abel/classes/stage/impl/stage_hipace.py | 68 ++++++++++++++++++++++---
 1 file changed, 61 insertions(+), 7 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 0a389872..a0399aaf 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -751,7 +751,7 @@ def external_focusing(self, enable_external_focusing : bool | None):
             else: 
                 stage_copy = self
             
-            if self.get_length() is not None:
+            if stage_copy.get_length() is not None:
                 self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=self.driver_half_oscillations)  # [T/m]
             else:
                 self._external_focusing_gradient = None
@@ -764,12 +764,46 @@ def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None):
         Calculate the external focusing gradient g for an azimuthal magnetic 
         field B=[gy,-gx,0] that gives ``num_half_oscillations`` half 
         oscillations for the drive beam over the length of the stage.
+
+        Parameters
+        ----------
+        num_half_oscillations : float, optional
+            Number of half betatron oscillations that the drive beam is 
+            intended to perform. If ``None``, will use ``self.driver_half_oscillations``.
+            Defaults to ``None``.
+
+        L : [m] float, optional
+            The length over which the driver will be guided. If ``None``, will 
+            extract the value using ``self.get_length()``. If the stage does 
+            have ramps that have not been fully set up, a deepcopy of the stage 
+            is created to set up its ramps using 
+            :func:`Stage._prepare_ramps() <abel.Stage._prepare_ramps>` so that 
+            the stage total length is defined.
+
+        Returns
+        -------
+        g : [T/m] float
+            The gradient for the azimuthal magnetic field.
         """
         if L is None:
-            if self.get_length() is None:
+
+            # Make a copy of the stage and set up its ramps if they are not set up
+            ramps_not_set_up = (
+                (self.upramp is not None and self.upramp.length is None) or
+                (self.downramp is not None and self.downramp.length is None)
+            )
+            if ramps_not_set_up:
+                stage_copy = copy.deepcopy(self)
+                stage_copy._prepare_ramps()
+                L = stage_copy.get_length()
+            else: 
+                stage_copy = self
+                L = stage_copy.get_length()
+                if L is None:
+                    L = stage_copy.length_flattop # If there are no ramps, can use either length or legnth_flattop.
+            
+            if L is None:
                 raise ValueError('Stage length is not set.')
-            L = self.length_flattop
-        #return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.get_length())**2  # [T/m]
 
         if num_half_oscillations is None:
             num_half_oscillations = self.driver_half_oscillations
@@ -923,7 +957,8 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
 
         driver_half_oscillations : float, optional
             Number of half betatron oscillations that the drive beam is 
-            expected to perform. Defaults to ``self.driver_half_oscillations``.
+            intended to perform. If ``None``, will use ``self.driver_half_oscillations``.
+            Defaults to ``None``.
 
         q : [C] float, optional
             Particle charge. q * nom_accel_gradient must be positive. Defaults 
@@ -964,11 +999,30 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
         if driver_half_oscillations is None:
             driver_half_oscillations = self.driver_half_oscillations
 
+        # Determine whether the ramps have been set up
+        ramps_not_set_up = (
+            (self.upramp is not None and self.upramp.length is None) or
+            (self.downramp is not None and self.downramp.length is None)
+        )
+        
+        # Make a copy of the stage
+        stage_copy = copy.deepcopy(self)
+
+        # The function to be used for solving the equation for phase advance numerically
         def rhs(L):
             g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
+            L_ramps = 0.0
+
+            # Set up the ramps using the stage copy
+            if ramps_not_set_up:
+                stage_copy.nom_energy_gain_flattop = nom_accel_gradient * L[0]  # L is an ndarray with one element
+                stage_copy.length_flattop = L[0]
+                stage_copy._prepare_ramps()
+                L_ramps = stage_copy.get_ramp_length()
 
             if self.external_focusing:  # Add contribution from external field used for driver guiding
-                g = g + self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=L)
+                g_ext = stage_copy.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=L+L_ramps)
+                g = g + g_ext
 
             prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient)
             energy_scaling = np.sqrt(initial_energy + q*nom_accel_gradient*L) - np.sqrt(initial_energy)
@@ -980,7 +1034,7 @@ def rhs(L):
 
         # Set the external focusing gradient for the driver guiding field if not already set
         if self.external_focusing and self._external_focusing_gradient is None:
-            self._external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=length)
+            self._external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=length+stage_copy.get_ramp_length())
 
         return length
     

From e98e361c56e8e5785bc43ef3506dfad9f543997f Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 16 Feb 2026 00:01:18 +0100
Subject: [PATCH 40/84] Changed the class attribute
 StageHipace._external_focusing_gradient default from 0.0 to None when
 StageHipace._external_focusing is set to False. Also corrected some of the
 comments.

---
 abel/classes/stage/impl/stage_hipace.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index a0399aaf..26450630 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -736,11 +736,10 @@ def external_focusing(self, enable_external_focusing : bool | None):
         self._external_focusing = bool(enable_external_focusing)
 
         if self._external_focusing is False:
-            self._external_focusing_gradient = 0.0  # TODO: set to None instead?
+            self._external_focusing_gradient = None
         elif self._external_focusing_gradient is None or self._external_focusing_gradient < 1e-15:
-            #if self.get_length() is None:
 
-            # Make a copy of the stage and set up its ramps if they are not set yp
+            # Make a copy of the stage and set up its ramps if they are not set up
             ramps_not_set_up = (
                 (self.upramp is not None and self.upramp.length is None) or
                 (self.downramp is not None and self.downramp.length is None)
@@ -800,7 +799,7 @@ def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None):
                 stage_copy = self
                 L = stage_copy.get_length()
                 if L is None:
-                    L = stage_copy.length_flattop # If there are no ramps, can use either length or legnth_flattop.
+                    L = stage_copy.length_flattop # If there are no ramps, can use either length or length_flattop.
             
             if L is None:
                 raise ValueError('Stage length is not set.')
@@ -862,7 +861,7 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
         if q * dacc_gradient > 0.0:
             raise ValueError('Drive beam charge * decceleration gradient must be negative.')
         
-        # Make a copy of the stage and set up its ramps if they are not set yp
+        # Make a copy of the stage and set up its ramps if they are not set up
         ramps_not_set_up = (
             (self.upramp is not None and self.upramp.length is None) or
             (self.downramp is not None and self.downramp.length is None)

From 74acd7c93e4bec0750789623e788d21c94ba6dc0 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 16 Feb 2026 11:22:04 +0100
Subject: [PATCH 41/84] Made corrections in
 tests/test_StageHipace.py::test_external_focusing to be compatible with
 stage._external_focusing_gradient set to None as default.

---
 tests/test_StageHipace.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py
index f907a516..b3513506 100644
--- a/tests/test_StageHipace.py
+++ b/tests/test_StageHipace.py
@@ -80,7 +80,7 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
     stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
-    assert np.isclose(0.0, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert stage._external_focusing_gradient is None
     assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
@@ -89,7 +89,7 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
     stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
-    assert np.isclose(0.0, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert stage._external_focusing_gradient is None
     assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 

From f3d3f25c8e089b4367292a2270afa66001c796ef Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 16 Feb 2026 16:47:19 +0100
Subject: [PATCH 42/84] Edited StageHipace.driver_guiding_trajectory() to use
 the stage total length instead of flattop length.

---
 abel/classes/stage/impl/stage_hipace.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 26450630..8892991a 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -872,8 +872,7 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
         else: 
             stage_copy = self
         
-        #L = stage_copy.get_length()  # [m]
-        L = stage_copy.length_flattop  # [m]
+        L = stage_copy.get_length()  # [m]
         
         if pz0 + q * dacc_gradient * L/SI.c < pz_thres:
             raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')

From 3e5eac05645b229ffaac14134fe67d694eeefe46 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Feb 2026 11:34:52 +0100
Subject: [PATCH 43/84] Moved StageHipace.driver_guiding_trajectory() further
 down the file.

---
 abel/classes/stage/impl/stage_hipace.py | 218 ++++++++++++------------
 1 file changed, 109 insertions(+), 109 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 8892991a..6e7ebb7a 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -808,115 +808,6 @@ def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None):
             num_half_oscillations = self.driver_half_oscillations
 
         return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/L)**2  # [T/m]
-
-
-    # =============================================
-    def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_half_osc=100):
-        """
-        Estimate the trajectory that the drive beam will follow when driver 
-        guiding with an external linear azimuthal magnetic field is applied to a 
-        drive beam with an initial angular offset. The calculations  are done by 
-        integrating simplified equations of motion.
-
-        Parameters
-        ----------
-        driver : ``Beam``
-            The drive beam.
-        
-        dacc_gradient : [V/m] float, optional
-            The decceleration gradient. Drive beam charge * decceleration 
-            gradient must be negative. Defaults to 0.0.
-
-        num_steps_per_half_osc : int, optional
-            Number of calcualtion steps per half-oscillation of the drive beam. 
-            The number of half-oscillations is set in 
-            :meth:`StageHipace.calc_external_focusing_gradient() <abel.StageHipace.calc_external_focusing_gradient>`. 
-            Defaults to 100.
-        
-
-        Returns
-        -------
-        s_trajectory : [m] float
-            Longitudinal coordinate of the drive beam trajectory. Reference is 
-            set at the start of the plasma stage.
-
-        x_trajectory : [m] float
-            x-coordinate of the drive beam trajectory.
-        
-        y_trajectory : [m] float
-            y-coordinate of the drive beam trajectory.
-        """
-
-        from abel.utilities.relativity import energy2momentum
-        from abel.utilities.statistics import weighted_mean
-
-        energy_thres = 10*driver.particle_mass*SI.c**2/SI.e  # [eV], 10 * particle rest energy. Gives beta=0.995.
-        pz_thres = energy2momentum(energy_thres, unit='eV', m=driver.particle_mass)
-        pz0 = energy2momentum(driver.energy(), unit='eV', m=driver.particle_mass)
-
-        if pz0 < pz_thres:
-            raise ValueError('This estimate is only valid for a relativistic beam.')
-        
-        q = driver.particle_charge()  # [C], particle charge including charge sign.
-        if q * dacc_gradient > 0.0:
-            raise ValueError('Drive beam charge * decceleration gradient must be negative.')
-        
-        # Make a copy of the stage and set up its ramps if they are not set up
-        ramps_not_set_up = (
-            (self.upramp is not None and self.upramp.length is None) or
-            (self.downramp is not None and self.downramp.length is None)
-        )
-        if ramps_not_set_up:
-            stage_copy = copy.deepcopy(self)
-            stage_copy._prepare_ramps()
-        else: 
-            stage_copy = self
-        
-        L = stage_copy.get_length()  # [m]
-        
-        if pz0 + q * dacc_gradient * L/SI.c < pz_thres:
-            raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')
-        
-        g = self._external_focusing_gradient  # [T/m]
-        ds = self.length_flattop/self.driver_half_oscillations/num_steps_per_half_osc  # [m], step size
-
-        prop_length = 0
-        s_trajectory = np.array([0.0])
-        x0 = driver.x_offset()
-        x_trajectory = np.array([x0])  # [m], records the trajectory
-        x = x0
-        y0 = driver.y_offset()
-        y_trajectory = np.array([y0])  # [m], records the trajectory
-        y = y0
-        px = weighted_mean(driver.pxs(), driver.weightings(), clean=False)
-        py = weighted_mean(driver.pys(), driver.weightings(), clean=False)
-        pz = pz0 # Can add option for deceleration using a gradient
-
-        while prop_length < L:
-
-            # Drift
-            prop_length = prop_length + 1/2*ds
-            x = x + px/pz*1/2*ds
-            y = y + py/pz*1/2*ds
-
-            # Kick
-            dpx = q*g*x*ds
-            px = px + dpx
-            dpy = q*g*y*ds
-            py = py + dpy
-            pz = pz0 + q * dacc_gradient * prop_length/SI.c # dacc_gradient>0
-
-            # Drift
-            prop_length = prop_length + 1/2*ds
-            x = x + px/pz*1/2*ds
-            y = y + py/pz*1/2*ds
-            s_trajectory = np.append(s_trajectory, prop_length)
-            x_trajectory = np.append(x_trajectory, x)
-            y_trajectory = np.append(y_trajectory, y)
-
-        s_trajectory = s_trajectory + driver.z_offset()
-
-        return s_trajectory, x_trajectory, y_trajectory
     
 
     # ==================================================
@@ -1037,6 +928,115 @@ def rhs(L):
         return length
     
 
+    # =============================================
+    def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_half_osc=100):
+        """
+        Estimate the trajectory that the drive beam will follow when driver 
+        guiding with an external linear azimuthal magnetic field is applied to a 
+        drive beam with an initial angular offset. The calculations  are done by 
+        integrating simplified equations of motion.
+
+        Parameters
+        ----------
+        driver : ``Beam``
+            The drive beam.
+        
+        dacc_gradient : [V/m] float, optional
+            The decceleration gradient. Drive beam charge * decceleration 
+            gradient must be negative. Defaults to 0.0.
+
+        num_steps_per_half_osc : int, optional
+            Number of calcualtion steps per half-oscillation of the drive beam. 
+            The number of half-oscillations is set in 
+            :meth:`StageHipace.calc_external_focusing_gradient() <abel.StageHipace.calc_external_focusing_gradient>`. 
+            Defaults to 100.
+        
+
+        Returns
+        -------
+        s_trajectory : [m] float
+            Longitudinal coordinate of the drive beam trajectory. Reference is 
+            set at the start of the plasma stage.
+
+        x_trajectory : [m] float
+            x-coordinate of the drive beam trajectory.
+        
+        y_trajectory : [m] float
+            y-coordinate of the drive beam trajectory.
+        """
+
+        from abel.utilities.relativity import energy2momentum
+        from abel.utilities.statistics import weighted_mean
+
+        energy_thres = 10*driver.particle_mass*SI.c**2/SI.e  # [eV], 10 * particle rest energy. Gives beta=0.995.
+        pz_thres = energy2momentum(energy_thres, unit='eV', m=driver.particle_mass)
+        pz0 = energy2momentum(driver.energy(), unit='eV', m=driver.particle_mass)
+
+        if pz0 < pz_thres:
+            raise ValueError('This estimate is only valid for a relativistic beam.')
+        
+        q = driver.particle_charge()  # [C], particle charge including charge sign.
+        if q * dacc_gradient > 0.0:
+            raise ValueError('Drive beam charge * decceleration gradient must be negative.')
+        
+        # Make a copy of the stage and set up its ramps if they are not set up
+        ramps_not_set_up = (
+            (self.upramp is not None and self.upramp.length is None) or
+            (self.downramp is not None and self.downramp.length is None)
+        )
+        if ramps_not_set_up:
+            stage_copy = copy.deepcopy(self)
+            stage_copy._prepare_ramps()
+        else: 
+            stage_copy = self
+        
+        L = stage_copy.get_length()  # [m]
+        
+        if pz0 + q * dacc_gradient * L/SI.c < pz_thres:
+            raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')
+        
+        g = self._external_focusing_gradient  # [T/m]
+        ds = self.length_flattop/self.driver_half_oscillations/num_steps_per_half_osc  # [m], step size
+
+        prop_length = 0
+        s_trajectory = np.array([0.0])
+        x0 = driver.x_offset()
+        x_trajectory = np.array([x0])  # [m], records the trajectory
+        x = x0
+        y0 = driver.y_offset()
+        y_trajectory = np.array([y0])  # [m], records the trajectory
+        y = y0
+        px = weighted_mean(driver.pxs(), driver.weightings(), clean=False)
+        py = weighted_mean(driver.pys(), driver.weightings(), clean=False)
+        pz = pz0 # Can add option for deceleration using a gradient
+
+        while prop_length < L:
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
+            y = y + py/pz*1/2*ds
+
+            # Kick
+            dpx = q*g*x*ds
+            px = px + dpx
+            dpy = q*g*y*ds
+            py = py + dpy
+            pz = pz0 + q * dacc_gradient * prop_length/SI.c # dacc_gradient>0
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
+            y = y + py/pz*1/2*ds
+            s_trajectory = np.append(s_trajectory, prop_length)
+            x_trajectory = np.append(x_trajectory, x)
+            y_trajectory = np.append(y_trajectory, y)
+
+        s_trajectory = s_trajectory + driver.z_offset()
+
+        return s_trajectory, x_trajectory, y_trajectory
+    
+
     # ==================================================
     # Apply waterfall function to all beam dump files
     def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsigma=None, args=None):

From 77e25926a9d5b07f94cfbe59b280dec288489210 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Feb 2026 12:38:47 +0100
Subject: [PATCH 44/84] Edited StageHipace.calc_length_num_beta_osc() to only
 use the nominal acceleration gradient set inside the stage instead of
 accepting it as an input parameter.

---
 abel/classes/stage/impl/stage_hipace.py | 71 ++++++++++++++++++++-----
 1 file changed, 58 insertions(+), 13 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 6e7ebb7a..363e20b6 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -811,12 +811,12 @@ def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None):
     
 
     # ==================================================
-    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, driver_half_oscillations=None, q=SI.e):
+    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_density=None, driver_half_oscillations=None, q=SI.e):
         """
         Calculate the stage length that gives ``num_beta_osc`` betatron 
         oscillations for a particle with given initial energy ``initial_energy`` 
-        in a uniform plasma stage (excluding ramps) with nominal acceleration 
-        gradient ``nom_accel_gradient`` and plasma density ``plasma_density``.
+        in a uniform plasma stage (excluding ramps) with defined nominal 
+        acceleration gradient and plasma density ``plasma_density``.
 
         Will take into account the contribution from an external linear magnetic 
         field B=[gy,-gx,0] if :attr:`self._external_focusing_gradient <abel.Stage.StageHipace.external_focusing>` 
@@ -836,10 +836,6 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             The initial energy of the particle at the start of the plasma stage. 
             Defaults to ``self.nom_energy``.
 
-        nom_accel_gradient : [V/m] float, optional
-            Nominal accelerating gradient of the plasma stage exclusing ramps. 
-            Defaults to ``self.nom_accel_gradient_flattop``.
-
         plasma_density : [m^-3] float, optional
             The plasma density of the plasma stage. Defaults to 
             ``self.plasma_density``.
@@ -869,10 +865,9 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             initial_energy = self.nom_energy
         initial_energy = initial_energy*SI.e  # [J]
 
-        if nom_accel_gradient is None:
-            if self.nom_accel_gradient_flattop is None:
-                raise ValueError('Stage.nom_accel_gradient_flattop not set.')
-            nom_accel_gradient = self.nom_accel_gradient_flattop
+        if self.nom_accel_gradient_flattop is None:
+            raise ValueError('Stage.nom_accel_gradient_flattop not set.')
+        nom_accel_gradient = self.nom_accel_gradient_flattop
 
         if q * nom_accel_gradient < 0:
             raise ValueError('q * nom_accel_gradient must be positive.')
@@ -904,8 +899,7 @@ def rhs(L):
 
             # Set up the ramps using the stage copy
             if ramps_not_set_up:
-                stage_copy.nom_energy_gain_flattop = nom_accel_gradient * L[0]  # L is an ndarray with one element
-                stage_copy.length_flattop = L[0]
+                stage_copy.length_flattop = L[0]  # L is an ndarray with one element
                 stage_copy._prepare_ramps()
                 L_ramps = stage_copy.get_ramp_length()
 
@@ -928,6 +922,57 @@ def rhs(L):
         return length
     
 
+    # ==================================================
+    def match_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, driver_half_oscillations=None, q=SI.e):
+        """
+        Calculate the stage length that gives ``num_beta_osc`` betatron 
+        oscillations for a particle with given initial energy ``initial_energy`` 
+        in a uniform plasma stage (excluding ramps) with nominal acceleration 
+        gradient ``nom_accel_gradient`` and plasma density ``plasma_density``.
+
+        Will take into account the contribution from an external linear magnetic 
+        field B=[gy,-gx,0] if :attr:`self._external_focusing_gradient <abel.Stage.StageHipace.external_focusing>` 
+        is set to ``True`` before calling this method.
+
+        Also set :attr:`self._external_focusing_gradient <abel.Stage.StageHipace._external_focusing_gradient>`
+        if it is not already set, and :attr:`self._external_focusing_gradient <abel.Stage.StageHipace.external_focusing>` 
+        is ``True``. 
+
+        Parameters
+        ----------
+        num_beta_osc : float
+            Total number of design betatron oscillations that the electron 
+            should perform through the plasma stage excluding ramps. 
+
+        initial_energy : [eV] float, optional
+            The initial energy of the particle at the start of the plasma stage. 
+            Defaults to ``self.nom_energy``.
+
+        nom_accel_gradient : [V/m] float, optional
+            Nominal accelerating gradient of the plasma stage exclusing ramps. 
+            Defaults to ``self.nom_accel_gradient_flattop``.
+
+        plasma_density : [m^-3] float, optional
+            The plasma density of the plasma stage. Defaults to 
+            ``self.plasma_density``.
+
+        driver_half_oscillations : float, optional
+            Number of half betatron oscillations that the drive beam is 
+            intended to perform. If ``None``, will use ``self.driver_half_oscillations``.
+            Defaults to ``None``.
+
+        q : [C] float, optional
+            Particle charge. q * nom_accel_gradient must be positive. Defaults 
+            to elementary charge.
+
+            
+        Returns
+        -------
+        length : [m] float
+            Length of the plasma stage excluding ramps matched to the given 
+            number of betatron oscillations.
+        """
+
     # =============================================
     def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_half_osc=100):
         """

From 821f4b5369ff3f594315f8276b6a15e773979aca Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Feb 2026 12:40:25 +0100
Subject: [PATCH 45/84] Edited
 tests/test_StageHipace.py::test_external_focusing to be compatible with the
 the changes in StageHipace.calc_length_num_beta_osc() to only use the nominal
 acceleration gradient set inside the stage instead of accepting it as an
 input parameter.

---
 tests/test_StageHipace.py | 36 ++++++++++--------------------------
 1 file changed, 10 insertions(+), 26 deletions(-)

diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py
index b3513506..d9c244e5 100644
--- a/tests/test_StageHipace.py
+++ b/tests/test_StageHipace.py
@@ -58,35 +58,31 @@ def test_external_focusing():
     of drive beam and main beam betatron oscillation.
     """
 
-    def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_focusing=False):
+    def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_focusing=False, nom_accel_gradient_flattop=1e9):
         stage = StageHipace()
         stage.nom_energy = nom_energy  # [eV]
         stage.plasma_density = plasma_density  # [m^-3]
         stage.driver_source = setup_trapezoid_driver_source()
         stage.external_focusing = external_focusing
+        stage.nom_accel_gradient_flattop = nom_accel_gradient_flattop  # [V/m]
 
         return stage
 
     num_beta_osc = 4.0  # The number of betatron oscilations that the main beam is expected to perform.
-    nom_accel_gradient = 1e9  # [V/m]
 
     stage = setup_minimal_StageHipace()
 
     # ========== Tests without any external fields ==========
     assert stage.external_focusing is False
 
-    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, 
-                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=1.0)
-    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, driver_half_oscillations=1.0)
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
     assert stage._external_focusing_gradient is None
     assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
-    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, 
-                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=2.0)
-    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, driver_half_oscillations=2.0)
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
     assert stage._external_focusing_gradient is None
@@ -95,9 +91,7 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
 
     # ========== Tests with external fields ==========
     stage = setup_minimal_StageHipace(external_focusing=True)
-    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, 
-                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=1.0)
-    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, driver_half_oscillations=1.0)
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(140.1695315279373, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -105,9 +99,7 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
 
 
     stage = setup_minimal_StageHipace(external_focusing=True)
-    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, 
-                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=2.0)
-    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, driver_half_oscillations=2.0)
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(574.1247168375805, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -117,9 +109,7 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
     # ========== With external fields, lower stage nominal energy ==========
     num_beta_osc2 = 8.0
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
-    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, initial_energy=stage.nom_energy, 
-                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=1.0)
-    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=1.0)
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(1038.1202586404845, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -127,9 +117,7 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
 
 
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
-    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, initial_energy=stage.nom_energy, 
-                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=2.0)
-    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=2.0)
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(5119.8513420067175, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -139,9 +127,7 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
     # ========== With external fields, lower stage nominal energy, lower driver energy ==========
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
     stage.driver_source.energy = 4.5e9  # [eV]
-    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, initial_energy=stage.nom_energy, 
-                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=1.0)
-    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=1.0)
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(88.3976616856249, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -150,9 +136,7 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
 
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
     stage.driver_source.energy = 4.5e9  # [eV]
-    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, initial_energy=stage.nom_energy, 
-                                                       nom_accel_gradient=nom_accel_gradient, driver_half_oscillations=2.0)
-    stage.nom_energy_gain = nom_accel_gradient * stage.length_flattop  # [eV]
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=2.0)
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(359.3285677857948, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)

From 754789b4510dcbb777c9485322c33e0853f75ee3 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Feb 2026 12:45:31 +0100
Subject: [PATCH 46/84] Edited Stage.calc_length_num_beta_osc() to only use the
 nominal acceleration gradient set inside the stage instead of accepting it as
 an input parameter

---
 abel/classes/stage/stage.py | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 9b758647..e604f5c2 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1362,12 +1362,12 @@ def matched_beta_function_flattop(self, energy):
 
 
     # ==================================================
-    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, q=SI.e):
+    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_density=None, q=SI.e):
         """
         Calculate the stage length that gives ``num_beta_osc`` betatron 
         oscillations for a particle with given initial energy ``initial_energy`` 
-        in a uniform plasma stage (excluding ramps) with nominal acceleration 
-        gradient ``nom_accel_gradient`` and plasma density ``plasma_density``.
+        in a uniform plasma stage (excluding ramps) with defined nominal 
+        acceleration gradient and plasma density ``plasma_density``.
 
         Parameters
         ----------
@@ -1379,10 +1379,6 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             The initial energy of the particle at the start of the plasma stage. 
             Defaults to ``self.nom_energy``.
 
-        nom_accel_gradient : [V/m] float, optional
-            Nominal accelerating gradient of the plasma stage exclusing ramps. 
-            Defaults to ``self.nom_accel_gradient_flattop``.
-
         plasma_density : [m^-3] float, optional
             The plasma density of the plasma stage. Defaults to 
             ``self.plasma_density``.
@@ -1405,10 +1401,9 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_
             initial_energy = self.nom_energy
         initial_energy = initial_energy*SI.e  # [J]
 
-        if nom_accel_gradient is None:
-            if self.nom_accel_gradient_flattop is None:
-                raise ValueError('Stage.nom_accel_gradient_flattop not set.')
-            nom_accel_gradient = self.nom_accel_gradient_flattop
+        if self.nom_accel_gradient_flattop is None:
+            raise ValueError('Stage.nom_accel_gradient_flattop not set.')
+        nom_accel_gradient = self.nom_accel_gradient_flattop
 
         if q * nom_accel_gradient < 0:
             raise ValueError('q * nom_accel_gradient must be positive.')

From 2d9d2805a933c48c12dd3261ab9cd190c3f130e3 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Feb 2026 16:17:42 +0100
Subject: [PATCH 47/84] Added
 StageHipace.match_length_guiding_2_num_beta_osc().

---
 abel/classes/stage/impl/stage_hipace.py | 116 ++++++++++++++++--------
 1 file changed, 77 insertions(+), 39 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 363e20b6..860c237a 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -730,6 +730,7 @@ def matched_beta_function(self, energy_incoming, match_entrance=True, q=SI.e):
     # =============================================
     @property
     def external_focusing(self) -> bool:
+        "Flag for enabling driver guiding using an external magnetic field."
         return self._external_focusing
     @external_focusing.setter
     def external_focusing(self, enable_external_focusing : bool | None):
@@ -758,6 +759,7 @@ def external_focusing(self, enable_external_focusing : bool | None):
     _external_focusing = False
 
 
+    # =============================================
     def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None):
         """
         Calculate the external focusing gradient g for an azimuthal magnetic 
@@ -819,12 +821,8 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_den
         acceleration gradient and plasma density ``plasma_density``.
 
         Will take into account the contribution from an external linear magnetic 
-        field B=[gy,-gx,0] if :attr:`self._external_focusing_gradient <abel.Stage.StageHipace.external_focusing>` 
-        is set to ``True`` before calling this method.
-
-        Also set :attr:`self._external_focusing_gradient <abel.Stage.StageHipace._external_focusing_gradient>`
-        if it is not already set, and :attr:`self._external_focusing_gradient <abel.Stage.StageHipace.external_focusing>` 
-        is ``True``. 
+        field B=[gy,-gx,0] if :attr:`self.external_focusing <abel.StageHipace.external_focusing>` 
+        is set to ``True``.
 
         Parameters
         ----------
@@ -842,7 +840,7 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_den
 
         driver_half_oscillations : float, optional
             Number of half betatron oscillations that the drive beam is 
-            intended to perform. If ``None``, will use ``self.driver_half_oscillations``.
+            intended to perform. If ``None``, will use :attr:`StageHipace.driver_half_oscillations <abel.StageHipace.driver_half_oscillations>`.
             Defaults to ``None``.
 
         q : [C] float, optional
@@ -882,8 +880,10 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_den
         
         if driver_half_oscillations is None:
             driver_half_oscillations = self.driver_half_oscillations
+        if driver_half_oscillations < 0:
+            raise ValueError('Number of driver oscillations must be positive.')
 
-        # Determine whether the ramps have been set up
+        # Assess whether the ramps have been set up
         ramps_not_set_up = (
             (self.upramp is not None and self.upramp.length is None) or
             (self.downramp is not None and self.downramp.length is None)
@@ -915,28 +915,27 @@ def rhs(L):
         solution = fsolve(lambda L: rhs(L) - 2*np.pi*num_beta_osc , x0=1)
         length = solution[0]
 
-        # Set the external focusing gradient for the driver guiding field if not already set
-        if self.external_focusing and self._external_focusing_gradient is None:
-            self._external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=length+stage_copy.get_ramp_length())
-
         return length
     
 
     # ==================================================
-    def match_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel_gradient=None, plasma_density=None, driver_half_oscillations=None, q=SI.e):
+    def match_length_guiding_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=None, q=SI.e, set_consistent_params=True):
         """
-        Calculate the stage length that gives ``num_beta_osc`` betatron 
-        oscillations for a particle with given initial energy ``initial_energy`` 
-        in a uniform plasma stage (excluding ramps) with nominal acceleration 
-        gradient ``nom_accel_gradient`` and plasma density ``plasma_density``.
+        Set :attr:`self.length_flattop <abel.Stage.length_flattop>` for a 
+        uniform plasma stage such that a particle with initial energy 
+        :attr:`self.nom_energy <abel.Stage.nom_energy>` will perform
+        ``num_beta_osc`` betatron oscillations through the stage (including any 
+        existing ramps). 
 
-        Will take into account the contribution from an external linear magnetic 
-        field B=[gy,-gx,0] if :attr:`self._external_focusing_gradient <abel.Stage.StageHipace.external_focusing>` 
-        is set to ``True`` before calling this method.
+        The stage length calculation is performed using 
+        :meth:`StageHipace.calc_flattop_num_beta_osc() <abel.StageHipace.calc_flattop_num_beta_osc>`.
+
+        Also set :attr:`self._external_focusing_gradient <abel.Stage._external_focusing_gradient>`
+        and :attr:`self.driver_half_oscillations <abel.StageHipace.driver_half_oscillations>`
+        to be consistent with the calculated stage length
+        if ``set_consistent_params`` and set :attr:`self.external_focusing <abel.StageHipace.external_focusing>` 
+        are ``True``. 
 
-        Also set :attr:`self._external_focusing_gradient <abel.Stage.StageHipace._external_focusing_gradient>`
-        if it is not already set, and :attr:`self._external_focusing_gradient <abel.Stage.StageHipace.external_focusing>` 
-        is ``True``. 
 
         Parameters
         ----------
@@ -944,35 +943,74 @@ def match_length_num_beta_osc(self, num_beta_osc, initial_energy=None, nom_accel
             Total number of design betatron oscillations that the electron 
             should perform through the plasma stage excluding ramps. 
 
-        initial_energy : [eV] float, optional
-            The initial energy of the particle at the start of the plasma stage. 
-            Defaults to ``self.nom_energy``.
-
-        nom_accel_gradient : [V/m] float, optional
-            Nominal accelerating gradient of the plasma stage exclusing ramps. 
-            Defaults to ``self.nom_accel_gradient_flattop``.
-
-        plasma_density : [m^-3] float, optional
-            The plasma density of the plasma stage. Defaults to 
-            ``self.plasma_density``.
-
         driver_half_oscillations : float, optional
             Number of half betatron oscillations that the drive beam is 
-            intended to perform. If ``None``, will use ``self.driver_half_oscillations``.
+            intended to perform. If ``None``, will use :attr:`StageHipace.driver_half_oscillations <abel.StageHipace.driver_half_oscillations>`
             Defaults to ``None``.
 
         q : [C] float, optional
             Particle charge. q * nom_accel_gradient must be positive. Defaults 
             to elementary charge.
 
+        set_consistent_params : bool, optional
+            Flag for setting :attr:`self._external_focusing_gradient <abel.Stage._external_focusing_gradient>`
+            and :attr:`self.driver_half_oscillations <abel.StageHipace.driver_half_oscillations>`
+            to be consistent with the calculated stage length. Defaults to 
+            ``True``.
+
             
         Returns
         -------
-        length : [m] float
-            Length of the plasma stage excluding ramps matched to the given 
-            number of betatron oscillations.
+        None
         """
 
+        # Assess whether length flattop can be set
+        if self._length_flattop_calc is not None and self._length_flattop is None:
+            from abel.classes.stage.stage import VariablesOverspecifiedError
+            raise VariablesOverspecifiedError("Stage length already known/calculateable, cannot set.")
+
+        if self.has_ramp():
+            # Calculate the number of betatron oscillations that the main beam 
+            # should perform in the flattop:
+            num_beta_osc_flattop = self.calc_flattop_num_beta_osc(num_beta_osc)
+        else:
+            num_beta_osc_flattop = num_beta_osc
+
+        if driver_half_oscillations is None:
+            driver_half_oscillations = self.driver_half_oscillations
+
+        # Calculate the length of the flattop stage
+        length_flattop = self.calc_length_num_beta_osc(num_beta_osc=num_beta_osc_flattop, 
+                                                       initial_energy=self.nom_energy, 
+                                                       plasma_density=self.plasma_density, 
+                                                       driver_half_oscillations=driver_half_oscillations, 
+                                                       q=q)
+
+        # Set the length of the flattop stage
+        self.length_flattop = length_flattop
+
+        # Set the external focusing gradient for the driver guiding field and 
+        # the number of half betatron oscillations that the drive beam is 
+        # intended to perform.
+        if self.external_focusing and set_consistent_params:
+            # Assess whether the ramps have been set up
+            ramps_not_set_up = (
+                (self.upramp is not None and self.upramp.length is None) or
+                (self.downramp is not None and self.downramp.length is None)
+            )
+
+            if ramps_not_set_up:
+                # Make a copy of the stage
+                stage_copy = copy.deepcopy(self)
+                stage_copy._prepare_ramps()
+            else:
+                stage_copy = self
+
+            length = length_flattop + stage_copy.get_ramp_length()
+            self._external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=length)
+            self.driver_half_oscillations = driver_half_oscillations
+
+
     # =============================================
     def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_half_osc=100):
         """

From d0e6b3633f30fb84b8527bc2e97b39d8e23ecad7 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Feb 2026 16:21:57 +0100
Subject: [PATCH 48/84] Changed the order of the input parameters in
 StageHipace.match_length_guiding_2_num_beta_osc().

---
 abel/classes/stage/impl/stage_hipace.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 860c237a..5b7a763e 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -919,7 +919,7 @@ def rhs(L):
     
 
     # ==================================================
-    def match_length_guiding_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=None, q=SI.e, set_consistent_params=True):
+    def match_length_guiding_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=None, set_consistent_params=True, q=SI.e):
         """
         Set :attr:`self.length_flattop <abel.Stage.length_flattop>` for a 
         uniform plasma stage such that a particle with initial energy 
@@ -948,16 +948,16 @@ def match_length_guiding_2_num_beta_osc(self, num_beta_osc, driver_half_oscillat
             intended to perform. If ``None``, will use :attr:`StageHipace.driver_half_oscillations <abel.StageHipace.driver_half_oscillations>`
             Defaults to ``None``.
 
-        q : [C] float, optional
-            Particle charge. q * nom_accel_gradient must be positive. Defaults 
-            to elementary charge.
-
         set_consistent_params : bool, optional
             Flag for setting :attr:`self._external_focusing_gradient <abel.Stage._external_focusing_gradient>`
             and :attr:`self.driver_half_oscillations <abel.StageHipace.driver_half_oscillations>`
             to be consistent with the calculated stage length. Defaults to 
             ``True``.
 
+        q : [C] float, optional
+            Particle charge. q * nom_accel_gradient must be positive. Defaults 
+            to elementary charge.
+
             
         Returns
         -------

From d13968fc0a9a0539f21c593c1b06615728936d18 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Feb 2026 17:06:54 +0100
Subject: [PATCH 49/84] Made correction in
 StageHipace.match_length_guiding_2_num_beta_osc() for cases involving ramps.

---
 abel/classes/stage/impl/stage_hipace.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 5b7a763e..3ba593e4 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -927,10 +927,13 @@ def match_length_guiding_2_num_beta_osc(self, num_beta_osc, driver_half_oscillat
         ``num_beta_osc`` betatron oscillations through the stage (including any 
         existing ramps). 
 
-        The stage length calculation is performed using 
+        - Assumes that each of the (uniform) ramps are configured to give pi/2 
+        phase advance for the main beam.
+
+        - The stage length calculation is performed using 
         :meth:`StageHipace.calc_flattop_num_beta_osc() <abel.StageHipace.calc_flattop_num_beta_osc>`.
 
-        Also set :attr:`self._external_focusing_gradient <abel.Stage._external_focusing_gradient>`
+        - Also set :attr:`self._external_focusing_gradient <abel.Stage._external_focusing_gradient>`
         and :attr:`self.driver_half_oscillations <abel.StageHipace.driver_half_oscillations>`
         to be consistent with the calculated stage length
         if ``set_consistent_params`` and set :attr:`self.external_focusing <abel.StageHipace.external_focusing>` 
@@ -972,7 +975,11 @@ def match_length_guiding_2_num_beta_osc(self, num_beta_osc, driver_half_oscillat
         if self.has_ramp():
             # Calculate the number of betatron oscillations that the main beam 
             # should perform in the flattop:
-            num_beta_osc_flattop = self.calc_flattop_num_beta_osc(num_beta_osc)
+            if self.upramp.ramp_shape != 'uniform' or self.downramp.ramp_shape != 'uniform':
+                raise ValueError('This method assumes uniform ramps.')
+            if self.upramp.length_flattop is not None or self.downramp.length_flattop is not None:
+                raise ValueError('This method assumes uniform ramps with length set to give pi/2 phase advance for the main beam.')
+            num_beta_osc_flattop = num_beta_osc - 0.5  # The ramps are by default set up to give pi/2 phase advance for the main beam.
         else:
             num_beta_osc_flattop = num_beta_osc
 

From 578744b23ecadfe673539e94e559c74b52c13639 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Feb 2026 18:26:15 +0100
Subject: [PATCH 50/84] Edited
 tests/test_StageHipace.py::test_external_focusing to be compatible with
 StageHipace.calc_length_num_beta_osc() no longer setting parameters inside
 stage. Also added
 tests/test_StageHipace.py::test_match_length_guiding_2_num_beta_osc().

---
 tests/test_StageHipace.py | 151 ++++++++++++++++++++++++++++++++++----
 1 file changed, 135 insertions(+), 16 deletions(-)

diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py
index d9c244e5..23bf74e7 100644
--- a/tests/test_StageHipace.py
+++ b/tests/test_StageHipace.py
@@ -47,7 +47,15 @@ def setup_trapezoid_driver_source(enable_xy_jitter=False, enable_xpyp_jitter=Fal
     return driver
 
 
+def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_focusing=False, nom_accel_gradient_flattop=1e9):
+    stage = StageHipace()
+    stage.nom_energy = nom_energy  # [eV]
+    stage.plasma_density = plasma_density  # [m^-3]
+    stage.driver_source = setup_trapezoid_driver_source()
+    stage.external_focusing = external_focusing
+    stage.nom_accel_gradient_flattop = nom_accel_gradient_flattop  # [V/m]
 
+    return stage
 
 
 @pytest.mark.StageHipace
@@ -55,25 +63,16 @@ def test_external_focusing():
     """
     Tests for ``StageHipace.calc_length_num_beta_osc()`` for accurately matching the 
     stage length and external driver guiding field gradient to a desired number 
-    of drive beam and main beam betatron oscillation.
+    of drive beam and main beam betatron oscillations.
     """
 
-    def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_focusing=False, nom_accel_gradient_flattop=1e9):
-        stage = StageHipace()
-        stage.nom_energy = nom_energy  # [eV]
-        stage.plasma_density = plasma_density  # [m^-3]
-        stage.driver_source = setup_trapezoid_driver_source()
-        stage.external_focusing = external_focusing
-        stage.nom_accel_gradient_flattop = nom_accel_gradient_flattop  # [V/m]
-
-        return stage
-
     num_beta_osc = 4.0  # The number of betatron oscilations that the main beam is expected to perform.
 
     stage = setup_minimal_StageHipace()
 
     # ========== Tests without any external fields ==========
     assert stage.external_focusing is False
+    assert stage._external_focusing_gradient is None
 
     stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, driver_half_oscillations=1.0)
     
@@ -92,17 +91,19 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
     # ========== Tests with external fields ==========
     stage = setup_minimal_StageHipace(external_focusing=True)
     stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, driver_half_oscillations=1.0)
+    assert stage._external_focusing_gradient is None
+    stage._external_focusing_gradient = 140.1695315279373
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
-    assert np.isclose(140.1695315279373, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(3.4268706541720553, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
     stage = setup_minimal_StageHipace(external_focusing=True)
     stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, driver_half_oscillations=2.0)
+    assert stage._external_focusing_gradient is None
+    stage._external_focusing_gradient = 574.1247168375805
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
-    assert np.isclose(574.1247168375805, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(3.3865024719148242, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
@@ -110,17 +111,19 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
     num_beta_osc2 = 8.0
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
     stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=1.0)
+    assert stage._external_focusing_gradient is None
+    stage._external_focusing_gradient = 1038.1202586404845
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
-    assert np.isclose(1038.1202586404845, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.259217324849329, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
     stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=2.0)
+    assert stage._external_focusing_gradient is None
+    stage._external_focusing_gradient = 5119.8513420067175
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
-    assert np.isclose(5119.8513420067175, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.13403339406399, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
@@ -128,22 +131,138 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
     stage.driver_source.energy = 4.5e9  # [eV]
     stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=1.0)
+    assert stage._external_focusing_gradient is None
+    stage._external_focusing_gradient = 88.3976616856249
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
-    assert np.isclose(88.3976616856249, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.2945695557961696, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
     stage.driver_source.energy = 4.5e9  # [eV]
     stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc2, driver_half_oscillations=2.0)
+    assert stage._external_focusing_gradient is None
+    stage._external_focusing_gradient = 359.3285677857948
+    
+    assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(1.2841918239865389, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+
+@pytest.mark.StageHipace
+def test_match_length_guiding_2_num_beta_osc():
+    """
+    Tests for ``StageHipace.match_length_guiding_2_num_beta_osc()`` for 
+    accurately matching the stage length and external driver guiding field 
+    gradient to a desired number of drive beam and main beam betatron 
+    oscillations.
+    """
+
+    num_beta_osc = 4.0  # The number of betatron oscilations that the main beam is expected to perform.
+    stage = setup_minimal_StageHipace()
+
+    # ========== Tests without any external fields ==========
+    assert stage.external_focusing is False
+    assert stage._external_focusing_gradient is None
+    assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0)
+
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=1.0, set_consistent_params=True)
+    assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0)
+    assert stage._external_focusing_gradient is None
+    assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
+    assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=2.0, set_consistent_params=True)
+    assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0) # Should still be the default value 1.0, since stage.external_focusing is False.
+    assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
+    assert stage._external_focusing_gradient is None
+    assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    # ========== Tests with external fields ==========
+    stage = setup_minimal_StageHipace(external_focusing=True)
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc)
+    
+    assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(140.1695315279373, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(3.4268706541720553, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    stage = setup_minimal_StageHipace(external_focusing=True)
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=2.0)
+
+    assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0)
+    assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(574.1247168375805, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(3.3865024719148242, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    # ========== With external fields, lower stage nominal energy ==========
+    num_beta_osc2 = 8.0
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc2)
+    
+    assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(1038.1202586404845, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1.259217324849329, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_half_oscillations=2.0)
+    
+    assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(5119.8513420067175, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1.13403339406399, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    # ========== With external fields, lower stage nominal energy, lower driver energy ==========
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
+    stage.driver_source.energy = 4.5e9  # [eV]
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc2)
+    
+    assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
+    assert np.isclose(88.3976616856249, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1.2945695557961696, stage.length_flattop, rtol=1e-5, atol=0.0)
+
+
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
+    stage.driver_source.energy = 4.5e9  # [eV]
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_half_oscillations=2.0)
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(359.3285677857948, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.2841918239865389, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
+    # ========== Tests with external fields, with ramps ==========
+    num_beta_osc3 = 5.5
+    stage = setup_minimal_StageHipace(nom_energy=10e9, external_focusing=True)
+    stage.driver_source.energy = 20e9  # [eV]
+    stage.ramp_beta_mag = 2.0
+    stage.upramp = PlasmaRamp()
+    stage.downramp = PlasmaRamp()
 
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_half_oscillations=1.0)
 
+    assert np.isclose(263.5820977342094, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1.3834379291754226, stage.length_flattop, rtol=1e-5, atol=0.0)
+    assert np.isclose(num_beta_osc3 - 0.5, stage.length_flattop2num_beta_osc(), rtol=1e-4, atol=0.0)
+    
 
+    stage = setup_minimal_StageHipace(nom_energy=10e9, external_focusing=True)
+    stage.driver_source.energy = 20e9  # [eV]
+    stage.ramp_beta_mag = 2.0
+    stage.upramp = PlasmaRamp()
+    stage.downramp = PlasmaRamp()
+
+    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_half_oscillations=2.0)
+
+    assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0)
+    assert np.isclose(1097.699359490811, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1.3520401015419283, stage.length_flattop, rtol=1e-5, atol=0.0)
+    assert np.isclose(num_beta_osc3 - 0.5, stage.length_flattop2num_beta_osc(), rtol=1e-4, atol=0.0)
+    
     
+
+

From 12f92698b66bb56765620b1a63375b5a1ea8440e Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 19 Feb 2026 14:25:23 +0100
Subject: [PATCH 51/84] Edited
 tests/test_StageReducedModels_beamline.py::test_ramped_linac() to be
 compatible with recent changes.

---
 tests/test_StageReducedModels_beamline.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tests/test_StageReducedModels_beamline.py b/tests/test_StageReducedModels_beamline.py
index 3dd213c8..964594a7 100644
--- a/tests/test_StageReducedModels_beamline.py
+++ b/tests/test_StageReducedModels_beamline.py
@@ -77,10 +77,11 @@ def setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_be
     stage = StageReducedModels()
     stage.time_step_mod = 0.03*2                                                    # In units of betatron wavelengths/c.
     stage.length_flattop = length_flattop                                           # [m]
-    stage.nom_energy_gain = stage.length_flattop*1e9                                # [eV]
+    if length_flattop is not None:
+        stage.nom_energy_gain = stage.length_flattop*1e9                            # [eV]
     stage.plasma_density = plasma_density                                           # [m^-3]
     stage.driver_source = driver_source
-    stage.main_source = main_source
+    #stage.main_source = main_source
     stage.ramp_beta_mag = ramp_beta_mag
     stage.enable_tr_instability = enable_tr_instability 
     stage.enable_radiation_reaction = enable_radiation_reaction
@@ -226,13 +227,16 @@ def test_ramped_linac():
     driver_source.z_offset = 1602e-6                                                # [m]
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=3.0e9)
 
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag,  enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag,  length_flattop=None, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+
+    stage.nom_energy = main_source.energy
+    stage.nom_accel_gradient_flattop = 1e9  # [V/m]
 
     # Adjust the lengths of the two stages in the linac to match the number of betatron oscillation 
-    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc=9.5, initial_energy=main_source.energy, nom_accel_gradient=1e9)
-    stage.nom_energy_gain = stage.length_flattop*1e9
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc=9.5)
+
     last_stage = copy.deepcopy(stage)
-    last_stage.length_flattop = last_stage.calc_length_num_beta_osc(num_beta_osc=9.5, initial_energy=main_source.energy+stage.nom_energy_gain, nom_accel_gradient=1e9)
+    last_stage.length_flattop = last_stage.calc_length_num_beta_osc(num_beta_osc=9.5, initial_energy=main_source.energy+stage.nom_energy_gain_flattop)
     last_stage.nom_energy_gain = last_stage.length_flattop*1e9
 
     # Set up the interstage

From 05301e852ab8fee171f2b8bf1267963ca1fca275 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 19 Feb 2026 14:26:49 +0100
Subject: [PATCH 52/84] Removed main source from setup_StageReducedModels() in
 tests/test_StageReducedModels_beamline.py, as it is no longer needed.

---
 tests/test_StageReducedModels_beamline.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/test_StageReducedModels_beamline.py b/tests/test_StageReducedModels_beamline.py
index 964594a7..0b7a5a52 100644
--- a/tests/test_StageReducedModels_beamline.py
+++ b/tests/test_StageReducedModels_beamline.py
@@ -72,7 +72,7 @@ def setup_basic_main_source(plasma_density, ramp_beta_mag=1.0, energy=361.8e9):
     return main
 
 
-def setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, length_flattop=1.56, enable_tr_instability=True, enable_radiation_reaction=True, enable_ion_motion=False, use_ramps=False, drive_beam_update_period=0, save_final_step=False):
+def setup_StageReducedModels(plasma_density, driver_source, ramp_beta_mag, length_flattop=1.56, enable_tr_instability=True, enable_radiation_reaction=True, enable_ion_motion=False, use_ramps=False, drive_beam_update_period=0, save_final_step=False):
     
     stage = StageReducedModels()
     stage.time_step_mod = 0.03*2                                                    # In units of betatron wavelengths/c.
@@ -81,7 +81,6 @@ def setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_be
         stage.nom_energy_gain = stage.length_flattop*1e9                            # [eV]
     stage.plasma_density = plasma_density                                           # [m^-3]
     stage.driver_source = driver_source
-    #stage.main_source = main_source
     stage.ramp_beta_mag = ramp_beta_mag
     stage.enable_tr_instability = enable_tr_instability 
     stage.enable_radiation_reaction = enable_radiation_reaction

From 03745f82f71920333dd4acc4a8b64ea87f8e3fdd Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 19 Feb 2026 16:27:38 +0100
Subject: [PATCH 53/84] Renamed
 StageHipace.match_length_guiding_2_num_beta_osc() to
 StageHipace.match_length_2_num_beta_osc().

---
 abel/classes/stage/stage.py | 62 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 61 insertions(+), 1 deletion(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index e604f5c2..a23e0925 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1579,8 +1579,68 @@ def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None,
             num_beta_osc = prefactor * energy_scaling / (2*np.pi)
 
             return num_beta_osc
+        
+
+    # ==================================================
+    def match_length_2_num_beta_osc(self, num_beta_osc, q=SI.e):
+        """
+        Set :attr:`self.length_flattop <abel.Stage.length_flattop>` for a 
+        uniform plasma stage such that a particle with initial energy 
+        :attr:`self.nom_energy <abel.Stage.nom_energy>` will perform
+        ``num_beta_osc`` betatron oscillations through the stage (including any 
+        existing ramps). 
+
+        - Assumes that each of the (uniform) ramps are configured to give pi/2 
+        phase advance for the main beam.
+
+        - The stage length calculation is performed using 
+        :meth:`Stage.calc_flattop_num_beta_osc() <abel.Stage.calc_flattop_num_beta_osc>`.
+
+
+        Parameters
+        ----------
+        num_beta_osc : float
+            Total number of design betatron oscillations that the electron 
+            should perform through the plasma stage excluding ramps. 
+
+        q : [C] float, optional
+            Particle charge. q * nom_accel_gradient must be positive. Defaults 
+            to elementary charge.
+
+            
+        Returns
+        -------
+        None
+        """
+
+        # Assess whether length flattop can be set
+        if self._length_flattop_calc is not None and self._length_flattop is None:
+            from abel.classes.stage.stage import VariablesOverspecifiedError
+            raise VariablesOverspecifiedError("Stage length already known/calculateable, cannot set.")
+
+        if self.has_ramp():
+            # Calculate the number of betatron oscillations that the main beam 
+            # should perform in the flattop:
+            if self.upramp.ramp_shape != 'uniform' or self.downramp.ramp_shape != 'uniform':
+                raise ValueError('This method assumes uniform ramps.')
+            if self.upramp.length_flattop is not None or self.downramp.length_flattop is not None:
+                raise ValueError('This method assumes uniform ramps with length set to give pi/2 phase advance for the main beam.')
+            
+            num_beta_osc_flattop = num_beta_osc - 0.5  # The ramps are by default set up to give pi/2 phase advance for the main beam.
+            # num_beta_osc_flattop = self.calc_flattop_num_beta_osc(num_beta_osc)
+        else:
+            num_beta_osc_flattop = num_beta_osc
+
+        # Calculate the length of the flattop stage
+        length_flattop = self.calc_length_num_beta_osc(num_beta_osc=num_beta_osc_flattop, 
+                                                       initial_energy=self.nom_energy, 
+                                                       plasma_density=self.plasma_density, 
+                                                       q=q)
+
+        # Set the length of the flattop stage
+        self.length_flattop = length_flattop
     
-    
+
     # ==================================================
     def phase_advance_beta_evolution(self, beta0=None):
         """

From d272cd457167d085f14f51b55c7ac8794dae16aa Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 19 Feb 2026 16:30:20 +0100
Subject: [PATCH 54/84] Edited
 tests/test_StageHipace.py::test_match_length_2_num_beta_osc due to the name
 change from StageHipace.match_length_guiding_2_num_beta_osc() to
 StageHipace.match_length_2_num_beta_osc().

---
 tests/test_StageHipace.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py
index 23bf74e7..165798d5 100644
--- a/tests/test_StageHipace.py
+++ b/tests/test_StageHipace.py
@@ -150,9 +150,9 @@ def test_external_focusing():
 
 
 @pytest.mark.StageHipace
-def test_match_length_guiding_2_num_beta_osc():
+def test_match_length_2_num_beta_osc():
     """
-    Tests for ``StageHipace.match_length_guiding_2_num_beta_osc()`` for 
+    Tests for ``StageHipace.match_length_2_num_beta_osc()`` for 
     accurately matching the stage length and external driver guiding field 
     gradient to a desired number of drive beam and main beam betatron 
     oscillations.
@@ -166,14 +166,14 @@ def test_match_length_guiding_2_num_beta_osc():
     assert stage._external_focusing_gradient is None
     assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0)
 
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=1.0, set_consistent_params=True)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=1.0, set_consistent_params=True)
     assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0)
     assert stage._external_focusing_gradient is None
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
     assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=2.0, set_consistent_params=True)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=2.0, set_consistent_params=True)
     assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0) # Should still be the default value 1.0, since stage.external_focusing is False.
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
     assert stage._external_focusing_gradient is None
@@ -182,7 +182,7 @@ def test_match_length_guiding_2_num_beta_osc():
 
     # ========== Tests with external fields ==========
     stage = setup_minimal_StageHipace(external_focusing=True)
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc)
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(140.1695315279373, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -190,7 +190,7 @@ def test_match_length_guiding_2_num_beta_osc():
 
 
     stage = setup_minimal_StageHipace(external_focusing=True)
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=2.0)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=2.0)
 
     assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0)
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
@@ -201,7 +201,7 @@ def test_match_length_guiding_2_num_beta_osc():
     # ========== With external fields, lower stage nominal energy ==========
     num_beta_osc2 = 8.0
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc2)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc2)
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(1038.1202586404845, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -209,7 +209,7 @@ def test_match_length_guiding_2_num_beta_osc():
 
 
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_half_oscillations=2.0)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_half_oscillations=2.0)
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(5119.8513420067175, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -219,7 +219,7 @@ def test_match_length_guiding_2_num_beta_osc():
     # ========== With external fields, lower stage nominal energy, lower driver energy ==========
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
     stage.driver_source.energy = 4.5e9  # [eV]
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc2)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc2)
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(88.3976616856249, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -228,7 +228,7 @@ def test_match_length_guiding_2_num_beta_osc():
 
     stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True)
     stage.driver_source.energy = 4.5e9  # [eV]
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_half_oscillations=2.0)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_half_oscillations=2.0)
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(359.3285677857948, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
@@ -243,7 +243,7 @@ def test_match_length_guiding_2_num_beta_osc():
     stage.upramp = PlasmaRamp()
     stage.downramp = PlasmaRamp()
 
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_half_oscillations=1.0)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_half_oscillations=1.0)
 
     assert np.isclose(263.5820977342094, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.3834379291754226, stage.length_flattop, rtol=1e-5, atol=0.0)
@@ -256,7 +256,7 @@ def test_match_length_guiding_2_num_beta_osc():
     stage.upramp = PlasmaRamp()
     stage.downramp = PlasmaRamp()
 
-    stage.match_length_guiding_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_half_oscillations=2.0)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_half_oscillations=2.0)
 
     assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0)
     assert np.isclose(1097.699359490811, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)

From dfafa9bfb93b5542c99a2dc6abf898e46a0ef626 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 19 Feb 2026 16:31:03 +0100
Subject: [PATCH 55/84] Renamed
 StageHipace.match_length_guiding_2_num_beta_osc() to
 StageHipace.match_length_2_num_beta_osc().

---
 abel/classes/stage/impl/stage_hipace.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 3ba593e4..f0d29f25 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -919,7 +919,7 @@ def rhs(L):
     
 
     # ==================================================
-    def match_length_guiding_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=None, set_consistent_params=True, q=SI.e):
+    def match_length_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=None, set_consistent_params=True, q=SI.e):
         """
         Set :attr:`self.length_flattop <abel.Stage.length_flattop>` for a 
         uniform plasma stage such that a particle with initial energy 

From 1c06ad936af27f71ac690f96425c97973881fb56 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Thu, 19 Feb 2026 16:54:35 +0100
Subject: [PATCH 56/84] Correction in tests/test_StageReducedModels_beamline.py
 due to the removal of main source from setup_StageReducedModels().

---
 tests/test_StageReducedModels_beamline.py | 26 +++++++++++------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/tests/test_StageReducedModels_beamline.py b/tests/test_StageReducedModels_beamline.py
index 0b7a5a52..9978dc2c 100644
--- a/tests/test_StageReducedModels_beamline.py
+++ b/tests/test_StageReducedModels_beamline.py
@@ -139,7 +139,7 @@ def test_baseline_linac():
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag=1.0)
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=1.0, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=1.0, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -226,7 +226,7 @@ def test_ramped_linac():
     driver_source.z_offset = 1602e-6                                                # [m]
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=3.0e9)
 
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag,  length_flattop=None, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag,  length_flattop=None, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
 
     stage.nom_energy = main_source.energy
     stage.nom_accel_gradient_flattop = 1e9  # [V/m]
@@ -382,7 +382,7 @@ def test_ramped_linac():
 
 #     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
 #     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-#     stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, length_flattop=length_flattop, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, drive_beam_update_period)
+#     stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, length_flattop=length_flattop, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, drive_beam_update_period)
 #     interstage = setup_InterstageImpactX(stage)
 
 #     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -416,7 +416,7 @@ def test_ramped_linac():
 
 #     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
 #     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-#     stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, length_flattop=length_flattop, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, drive_beam_update_period)
+#     stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, length_flattop=length_flattop, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, drive_beam_update_period)
 #     interstage = setup_InterstageImpactX(stage)
 
 #     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -454,7 +454,7 @@ def test_angular_jitter_linac():
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -491,7 +491,7 @@ def test_angular_jitter_ramped_linac():
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -529,7 +529,7 @@ def test_trInstability_linac():
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
     main_source.energy = 3.0e9                                                      # [eV], HALHF v2 start energy
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -563,7 +563,7 @@ def test_jitter_trInstability_ramped_linac():
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=3.0e9)
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -661,7 +661,7 @@ def test_jitter_trInstability_ramped_linac():
 
 #     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
 #     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag)
-#     stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+#     stage = setup_StageReducedModels(plasma_density, driver_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
 #     interstage = setup_InterstageImpactX(stage)
 
 #     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -694,7 +694,7 @@ def test_ionMotion_linac():
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=361.8e9)
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -728,7 +728,7 @@ def test_jitter_trInstability_ionMotion_linac():
 
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=361.8e9)
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps)
     interstage = setup_InterstageImpactX(stage)
 
     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)
@@ -767,7 +767,7 @@ def test_jitter_trInstability_ionMotion_ramped_linac():
     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
     assert driver_source.align_beam_axis is False
     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=81.0e9)  # Choosing an energy that gives a sensible number of time steps.
-    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, main_source=main_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, save_final_step=True)
+    stage = setup_StageReducedModels(plasma_density=plasma_density, driver_source=driver_source, ramp_beta_mag=ramp_beta_mag, enable_tr_instability=enable_tr_instability, enable_radiation_reaction=enable_radiation_reaction, enable_ion_motion=enable_ion_motion, use_ramps=use_ramps, save_final_step=True)
     interstage = setup_InterstageImpactX(stage)
 
     assert stage.driver_source.align_beam_axis is True
@@ -868,7 +868,7 @@ def test_jitter_trInstability_ionMotion_ramped_linac():
 
 #     driver_source = setup_trapezoid_driver_source(enable_xy_jitter, enable_xpyp_jitter)
 #     main_source = setup_basic_main_source(plasma_density, ramp_beta_mag, energy=3.0e9)
-#     stage = setup_StageReducedModels(plasma_density, driver_source, main_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
+#     stage = setup_StageReducedModels(plasma_density, driver_source, ramp_beta_mag, enable_tr_instability, enable_radiation_reaction, enable_ion_motion, use_ramps, drive_beam_update_period)
 #     interstage = setup_InterstageImpactX(stage)
 
 #     linac = PlasmaLinac(source=main_source, stage=stage, interstage=interstage, num_stages=num_stages, alternate_interstage_polarity=False)

From cfa0d1c548aabcc2548c9b0dee8a7e9d56e06b34 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 20 Feb 2026 15:33:58 +0100
Subject: [PATCH 57/84] Patched StageHipace.external_focusing setter to only
 calculate StageHipace._external_focusing_gradient when ramps can be set up
 for the case with ramps.

---
 abel/classes/stage/impl/stage_hipace.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index f0d29f25..79358dda 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -745,7 +745,10 @@ def external_focusing(self, enable_external_focusing : bool | None):
                 (self.upramp is not None and self.upramp.length is None) or
                 (self.downramp is not None and self.downramp.length is None)
             )
-            if ramps_not_set_up:
+
+            can_set_up_ramps = (self.nom_energy_gain_flattop is not None and self.nom_energy is not None)
+
+            if ramps_not_set_up and can_set_up_ramps:
                 stage_copy = copy.deepcopy(self)
                 stage_copy._prepare_ramps()
             else: 

From 796303ed9969b54b4208faa26c2e1707e4302b6e Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 20 Feb 2026 17:07:44 +0100
Subject: [PATCH 58/84] Added external_focusing_gradient as a property with
 setter and getter in StageHipace.

---
 abel/classes/stage/impl/stage_hipace.py | 69 ++++++++++++++++---------
 1 file changed, 45 insertions(+), 24 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 79358dda..63bc0a26 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -214,8 +214,7 @@ def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, drive
         self.no_plasma = no_plasma
 
         # external focusing (APL-like) [T/m]
-        self.driver_half_oscillations = 1.0 
-        self._external_focusing_gradient = None
+        self.driver_half_oscillations = 1.0
         self.external_focusing = external_focusing
 
         # plasma profile
@@ -339,10 +338,10 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
         filename_input = 'input_file'
         path_input = tmpfolder + filename_input
 
-        if self.external_focusing and self._external_focusing_gradient is None:
-            self._external_focusing_gradient = self.calc_external_focusing_gradient() # Set the gradient for external focusing fields if not already set.
+        if self.external_focusing and self.external_focusing_gradient is None:
+            self.external_focusing_gradient = self.calc_external_focusing_gradient() # Set the gradient for external focusing fields if not already set.
 
-        hipace_write_inputs(path_input, filename_beam, filename_driver, self.plasma_density, self.num_steps, time_step, box_range_z, box_size_xy, ion_motion=self.ion_motion, ion_species=self.ion_species, beam_ionization=self.beam_ionization, radiation_reaction=self.radiation_reaction, output_period=output_period, num_cell_xy=self.num_cell_xy, num_cell_z=num_cell_z, driver_only=self.driver_only, density_table_file=density_table_file, no_plasma=self.no_plasma, external_focusing_gradient=self._external_focusing_gradient, mesh_refinement=self.mesh_refinement, do_spin_tracking=self.do_spin_tracking)
+        hipace_write_inputs(path_input, filename_beam, filename_driver, self.plasma_density, self.num_steps, time_step, box_range_z, box_size_xy, ion_motion=self.ion_motion, ion_species=self.ion_species, beam_ionization=self.beam_ionization, radiation_reaction=self.radiation_reaction, output_period=output_period, num_cell_xy=self.num_cell_xy, num_cell_z=num_cell_z, driver_only=self.driver_only, density_table_file=density_table_file, no_plasma=self.no_plasma, external_focusing_gradient=self.external_focusing_gradient, mesh_refinement=self.mesh_refinement, do_spin_tracking=self.do_spin_tracking)
         
         
         ## RUN SIMULATION
@@ -685,7 +684,7 @@ def matched_beta_function(self, energy_incoming, match_entrance=True, q=SI.e):
         Calculates the matched beta function of the stage. If there is an 
         upramp, the beta function is magnified by default so that it shrinks to 
         the correct size when it enters the main flattop plasma stage. Also 
-        takes into account external focusing field B=[gy,-gx,0] if present. 
+        takes into account external focusing field B=[g_ext*y, -g_ext*x, 0] if present. 
     
         
         Parameters
@@ -710,8 +709,8 @@ def matched_beta_function(self, energy_incoming, match_entrance=True, q=SI.e):
         energy_incoming = energy_incoming*SI.e  # [J]
 
         g = SI.e*self.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
-        if self._external_focusing_gradient is not None:  # Add contribution from external field
-            g = g + self._external_focusing_gradient
+        if self.external_focusing_gradient is not None:  # Add contribution from external field
+            g = g + self.external_focusing_gradient
 
         k_beta = np.sqrt(np.abs(q)*g*SI.c/energy_incoming)  # [m^-1], betatron wavenumber.
 
@@ -735,19 +734,32 @@ def external_focusing(self) -> bool:
     @external_focusing.setter
     def external_focusing(self, enable_external_focusing : bool | None):
         self._external_focusing = bool(enable_external_focusing)
+    _external_focusing = False
 
-        if self._external_focusing is False:
-            self._external_focusing_gradient = None
-        elif self._external_focusing_gradient is None or self._external_focusing_gradient < 1e-15:
 
-            # Make a copy of the stage and set up its ramps if they are not set up
+    # =============================================
+    @property
+    def external_focusing_gradient(self) -> float:
+        """
+        The external focusing gradient g_ext [T/m] for an azimuthal magnetic 
+        field B = [g_ext*y, -g_ext*x, 0].
+        """
+
+        if self.external_focusing is True and self._external_focusing_gradient is None:
+            # If external focusing is enabled, but the gradient of the external 
+            # focusing field is not yet set, try to calculate it:
+
+            # Check whether the ramps have been set up if they exist
             ramps_not_set_up = (
                 (self.upramp is not None and self.upramp.length is None) or
                 (self.downramp is not None and self.downramp.length is None)
             )
 
-            can_set_up_ramps = (self.nom_energy_gain_flattop is not None and self.nom_energy is not None)
+            # Check if there is enough information to set up the ramps
+            can_set_up_ramps = (self.get_nom_energy_gain(ignore_ramps_if_undefined=True) is not None 
+                                and self.nom_energy is not None)
 
+            # Make a copy of the stage and set up its ramps if they are not set up
             if ramps_not_set_up and can_set_up_ramps:
                 stage_copy = copy.deepcopy(self)
                 stage_copy._prepare_ramps()
@@ -755,19 +767,28 @@ def external_focusing(self, enable_external_focusing : bool | None):
                 stage_copy = self
             
             if stage_copy.get_length() is not None:
-                self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=self.driver_half_oscillations)  # [T/m]
-            else:
-                self._external_focusing_gradient = None
+                self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=self.driver_half_oscillations)
 
-    _external_focusing = False
+        return self._external_focusing_gradient
+    
+    @external_focusing_gradient.setter
+    def external_focusing_gradient(self, g_ext : float | None):
+        if g_ext is not None and not isinstance(g_ext, float):
+            raise ValueError("External focusing gradient must be a float or None.")
+        if self.external_focusing is False:
+            self._external_focusing_gradient = None
+            raise ValueError("Cannot set external focusing gradient when self.external_focusing is False.")
+        self._external_focusing_gradient = g_ext
+
+    _external_focusing_gradient = None
 
 
     # =============================================
     def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None):
         """
-        Calculate the external focusing gradient g for an azimuthal magnetic 
-        field B=[gy,-gx,0] that gives ``num_half_oscillations`` half 
-        oscillations for the drive beam over the length of the stage.
+        Calculate the external focusing gradient g_ext for an azimuthal magnetic 
+        field B = [g_ext*y, -g_ext*x, 0] that gives ``num_half_oscillations`` 
+        half oscillations for the drive beam over the length of the stage.
 
         Parameters
         ----------
@@ -786,7 +807,7 @@ def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None):
 
         Returns
         -------
-        g : [T/m] float
+        g_ext : [T/m] float
             The gradient for the azimuthal magnetic field.
         """
         if L is None:
@@ -824,7 +845,7 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_den
         acceleration gradient and plasma density ``plasma_density``.
 
         Will take into account the contribution from an external linear magnetic 
-        field B=[gy,-gx,0] if :attr:`self.external_focusing <abel.StageHipace.external_focusing>` 
+        field B=[g_ext*y, -g_ext*x, 0] if :attr:`self.external_focusing <abel.StageHipace.external_focusing>` 
         is set to ``True``.
 
         Parameters
@@ -1017,7 +1038,7 @@ def match_length_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=Non
                 stage_copy = self
 
             length = length_flattop + stage_copy.get_ramp_length()
-            self._external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=length)
+            self.external_focusing_gradient = self.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=length)
             self.driver_half_oscillations = driver_half_oscillations
 
 
@@ -1088,7 +1109,7 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
         if pz0 + q * dacc_gradient * L/SI.c < pz_thres:
             raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')
         
-        g = self._external_focusing_gradient  # [T/m]
+        g = self.external_focusing_gradient  # [T/m]
         ds = self.length_flattop/self.driver_half_oscillations/num_steps_per_half_osc  # [m], step size
 
         prop_length = 0

From e1f0474d249c71838092751cd0bb7c6b846d831a Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 23 Feb 2026 13:04:30 +0100
Subject: [PATCH 59/84] Adited tests/test_StageHipace.py to be compatible with
 the new StageHipace external_focusing_gradient getter and setter.

---
 tests/test_StageHipace.py | 108 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 108 insertions(+)

diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py
index 165798d5..7c55dfd8 100644
--- a/tests/test_StageHipace.py
+++ b/tests/test_StageHipace.py
@@ -60,6 +60,98 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
 
 @pytest.mark.StageHipace
 def test_external_focusing():
+    """
+    Tests for ``StageHipace.external_focusing()`` for setting and calculating 
+    the gradient g_ext [T/m] for an external azimuthal magnetic field 
+    B = [g_ext*y, -g_ext*x, 0].
+    """
+
+    # ========== Tests withexternal focusing disabled ==========
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=False, nom_accel_gradient_flattop=None)
+    assert stage.external_focusing is False
+    assert stage._external_focusing_gradient is None
+    assert stage.external_focusing_gradient is None
+    with pytest.raises(ValueError):
+        stage.external_focusing_gradient = 3.14  # Cannot set any value when stage.external_focusing is False.
+
+
+    # ========== Tests with external focusing enabled ==========
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True, nom_accel_gradient_flattop=None)
+    assert stage.external_focusing is True
+    assert stage._external_focusing_gradient is None
+    assert stage.external_focusing_gradient is None
+    with pytest.raises(ValueError):
+        stage.external_focusing_gradient = 'test'
+    stage.external_focusing_gradient = 3.14
+    assert np.isclose(3.14, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)
+    stage.external_focusing_gradient = 314.0
+    assert np.isclose(314.0, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)
+
+
+    # ========== Test the external gradient calculation without ramps ==========
+    stage = setup_minimal_StageHipace(nom_energy=3e9, external_focusing=True, nom_accel_gradient_flattop=1e9)
+    stage.driver_source.energy = 4.5e9  # [eV]
+    assert stage.external_focusing is True
+    assert stage._external_focusing_gradient is None
+    assert stage.external_focusing_gradient is None
+    assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0)
+    stage.length = 1.2945695557961696  # [m]
+    assert np.isclose(1.2945695557961696*1e9, stage.nom_energy_gain, rtol=1e-10, atol=0.0)
+    assert np.isclose(1.2945695557961696*1e9, stage.nom_energy_gain_flattop, rtol=1e-10, atol=0.0)
+    assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)
+
+    stage.external_focusing_gradient = 3.14  # Overwrite the focusing gradient
+    assert np.isclose(3.14, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)
+    stage.external_focusing_gradient = None  # Overwrite the focusing gradient
+    assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)  # Should again automatically calculate a focusing gradient.
+
+    # Match the focusing gradient to more driver betatron oscillations
+    stage.driver_half_oscillations = 2.0
+    stage.length = 1.2841918239865389
+    assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)  # Still the previous focusing gradient.
+    stage.external_focusing_gradient = None  # Overwrite the focusing gradient
+    assert np.isclose(359.3285677857948, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)  # Should again automatically calculate a focusing gradient.
+
+
+    # ========== Test the external gradient calculation with ramps ==========
+    stage = setup_minimal_StageHipace(nom_energy=10e9, external_focusing=True, nom_accel_gradient_flattop=1e9)
+    stage.driver_source.energy = 20e9  # [eV]
+    stage.ramp_beta_mag = 2.0
+    stage.upramp = PlasmaRamp()
+    stage.downramp = PlasmaRamp()
+    assert stage.external_focusing is True
+    assert stage._external_focusing_gradient is None
+    assert stage.external_focusing_gradient is None
+    assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0)
+    stage.length_flattop = 1.3834379291754226  # [m]
+
+    assert stage._external_focusing_gradient is None
+    assert np.isclose(263.5820977342094, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(263.5820977342094, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+
+    stage.external_focusing_gradient = 3.14  # Overwrite the focusing gradient
+    assert np.isclose(3.14, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)
+    stage.external_focusing_gradient = None  # Overwrite the focusing gradient
+    assert np.isclose(263.5820977342094, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)  # Should again automatically calculate a focusing gradient.
+
+    # Match the focusing gradient to more driver betatron oscillations
+    stage.driver_half_oscillations = 2.0
+    stage.length_flattop = 1.3520401015419283  # [m]
+    assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0)
+    assert np.isclose(263.5820977342094, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)  # Still the previous focusing gradient.
+    stage.external_focusing_gradient = None  # Overwrite the focusing gradient
+    assert np.isclose(1097.699359490811, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)  # Should again automatically calculate a focusing gradient.
+    assert np.isclose(1097.699359490811, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    
+
+    # stage.length = 1.2841918239865389
+    # assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)  # Still the previous focusing gradient.
+    # stage.external_focusing_gradient = None  # Overwrite the focusing gradient
+    # assert np.isclose(359.3285677857948, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)  # Should again automatically calculate a focusing gradient.
+    
+
+@pytest.mark.StageHipace
+def test_calc_length_num_beta_osc():
     """
     Tests for ``StageHipace.calc_length_num_beta_osc()`` for accurately matching the 
     stage length and external driver guiding field gradient to a desired number 
@@ -164,17 +256,20 @@ def test_match_length_2_num_beta_osc():
     # ========== Tests without any external fields ==========
     assert stage.external_focusing is False
     assert stage._external_focusing_gradient is None
+    assert stage.external_focusing_gradient is None
     assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0)
 
     stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=1.0, set_consistent_params=True)
     assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0)
     assert stage._external_focusing_gradient is None
+    assert stage.external_focusing_gradient is None
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
     assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
     stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_half_oscillations=2.0, set_consistent_params=True)
     assert np.isclose(1.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0) # Should still be the default value 1.0, since stage.external_focusing is False.
+    assert stage.external_focusing_gradient is None
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-5, atol=0.0)
     assert stage._external_focusing_gradient is None
     assert np.isclose(3.440220555221998, stage.length_flattop, rtol=1e-5, atol=0.0)
@@ -182,10 +277,16 @@ def test_match_length_2_num_beta_osc():
 
     # ========== Tests with external fields ==========
     stage = setup_minimal_StageHipace(external_focusing=True)
+    assert stage._external_focusing_gradient is None
+    assert stage.external_focusing_gradient is None
+    stage.external_focusing_gradient = 3.14  # Set a random gradient which will be overwritten later
+    assert np.isclose(3.14, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
+
     stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc)
     
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(140.1695315279373, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(140.1695315279373, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(3.4268706541720553, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
@@ -195,6 +296,7 @@ def test_match_length_2_num_beta_osc():
     assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-15, atol=0.0)
     assert np.isclose(num_beta_osc, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(574.1247168375805, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(574.1247168375805, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(3.3865024719148242, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
@@ -205,6 +307,7 @@ def test_match_length_2_num_beta_osc():
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(1038.1202586404845, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1038.1202586404845, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.259217324849329, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
@@ -213,6 +316,7 @@ def test_match_length_2_num_beta_osc():
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(5119.8513420067175, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(5119.8513420067175, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.13403339406399, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
@@ -223,6 +327,7 @@ def test_match_length_2_num_beta_osc():
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(88.3976616856249, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.2945695557961696, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
@@ -232,6 +337,7 @@ def test_match_length_2_num_beta_osc():
     
     assert np.isclose(num_beta_osc2, stage.length_flattop2num_beta_osc(), rtol=1e-10, atol=0.0)
     assert np.isclose(359.3285677857948, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(359.3285677857948, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.2841918239865389, stage.length_flattop, rtol=1e-5, atol=0.0)
 
 
@@ -246,6 +352,7 @@ def test_match_length_2_num_beta_osc():
     stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_half_oscillations=1.0)
 
     assert np.isclose(263.5820977342094, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(263.5820977342094, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.3834379291754226, stage.length_flattop, rtol=1e-5, atol=0.0)
     assert np.isclose(num_beta_osc3 - 0.5, stage.length_flattop2num_beta_osc(), rtol=1e-4, atol=0.0)
     
@@ -260,6 +367,7 @@ def test_match_length_2_num_beta_osc():
 
     assert np.isclose(2.0, stage.driver_half_oscillations, rtol=1e-10, atol=0.0)
     assert np.isclose(1097.699359490811, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
+    assert np.isclose(1097.699359490811, stage.external_focusing_gradient, rtol=1e-5, atol=0.0)
     assert np.isclose(1.3520401015419283, stage.length_flattop, rtol=1e-5, atol=0.0)
     assert np.isclose(num_beta_osc3 - 0.5, stage.length_flattop2num_beta_osc(), rtol=1e-4, atol=0.0)
     

From 48460131422d3a51841e6b23dc81efc0681f1900 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 23 Feb 2026 13:05:05 +0100
Subject: [PATCH 60/84] Edited docstrings in tests/test_StageHipace.py.

---
 tests/test_StageHipace.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py
index 7c55dfd8..d0389c38 100644
--- a/tests/test_StageHipace.py
+++ b/tests/test_StageHipace.py
@@ -61,9 +61,9 @@ def setup_minimal_StageHipace(nom_energy=100e9, plasma_density=6e20, external_fo
 @pytest.mark.StageHipace
 def test_external_focusing():
     """
-    Tests for ``StageHipace.external_focusing()`` for setting and calculating 
-    the gradient g_ext [T/m] for an external azimuthal magnetic field 
-    B = [g_ext*y, -g_ext*x, 0].
+    Tests for ``StageHipace.external_focusing_gradient`` for setting and 
+    calculating the gradient g_ext [T/m] for an external azimuthal magnetic 
+    field B = [g_ext*y, -g_ext*x, 0].
     """
 
     # ========== Tests withexternal focusing disabled ==========
@@ -144,12 +144,6 @@ def test_external_focusing():
     assert np.isclose(1097.699359490811, stage._external_focusing_gradient, rtol=1e-5, atol=0.0)
     
 
-    # stage.length = 1.2841918239865389
-    # assert np.isclose(88.3976616856249, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)  # Still the previous focusing gradient.
-    # stage.external_focusing_gradient = None  # Overwrite the focusing gradient
-    # assert np.isclose(359.3285677857948, stage.external_focusing_gradient, rtol=1e-10, atol=0.0)  # Should again automatically calculate a focusing gradient.
-    
-
 @pytest.mark.StageHipace
 def test_calc_length_num_beta_osc():
     """

From cc1c158e82a6676fb7d10192b58309a4ab941e63 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 23 Feb 2026 13:05:33 +0100
Subject: [PATCH 61/84] Edited docstrings in
 abel/classes/stage/impl/stage_hipace.py.

---
 abel/classes/stage/impl/stage_hipace.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 63bc0a26..9a57b6c8 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -741,13 +741,13 @@ def external_focusing(self, enable_external_focusing : bool | None):
     @property
     def external_focusing_gradient(self) -> float:
         """
-        The external focusing gradient g_ext [T/m] for an azimuthal magnetic 
+        The focusing gradient g_ext [T/m] for an external azimuthal magnetic 
         field B = [g_ext*y, -g_ext*x, 0].
         """
 
         if self.external_focusing is True and self._external_focusing_gradient is None:
             # If external focusing is enabled, but the gradient of the external 
-            # focusing field is not yet set, try to calculate it:
+            # focusing field is not yet set, try to calculate and set it:
 
             # Check whether the ramps have been set up if they exist
             ramps_not_set_up = (
@@ -759,7 +759,8 @@ def external_focusing_gradient(self) -> float:
             can_set_up_ramps = (self.get_nom_energy_gain(ignore_ramps_if_undefined=True) is not None 
                                 and self.nom_energy is not None)
 
-            # Make a copy of the stage and set up its ramps if they are not set up
+            # Make a copy of the stage and set up its ramps if they are not set 
+            # up and can be set up
             if ramps_not_set_up and can_set_up_ramps:
                 stage_copy = copy.deepcopy(self)
                 stage_copy._prepare_ramps()

From 0ce451041fba8f624993504413ff5ca3f5e0e93f Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 23 Feb 2026 13:07:40 +0100
Subject: [PATCH 62/84] Edited abel/classes/stage/stage.py to return
 Stage.external_focusing_gradient as None by default for Stage subclasses not
 supporting external driver guiding fields.

---
 abel/classes/stage/stage.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index e7b1dbc7..560701e0 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -127,8 +127,6 @@ def __init__(self, nom_accel_gradient, nom_energy_gain, plasma_density, driver_s
         self.ramp_beta_mag = ramp_beta_mag
         
         self.stage_number = None
-
-        self._external_focusing_gradient = None
         
         # nominal initial energy
         self.nom_energy = None 
@@ -1559,7 +1557,7 @@ def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None,
         the stage (excluding ramps).
 
         Will take into account the contribution from an external linear magnetic 
-        field B=[gy,-gx,0] if :attr:`self._external_focusing_gradient <abel.Stage.external_focusing>`
+        field B=[gy,-gx,0] if :attr:`self.external_focusing_gradient <abel.Stage.external_focusing>`
         is not ``None``.
 
         Parameters
@@ -1620,8 +1618,8 @@ def length_flattop2num_beta_osc(self, length_flattop=None, initial_energy=None,
             return self.phase_advance_beta_evolution()/(2*np.pi)
         else:
             g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
-            if self._external_focusing_gradient is not None:
-                g = g + self._external_focusing_gradient
+            if self.external_focusing_gradient is not None:
+                g = g + self.external_focusing_gradient
 
             prefactor = 2*np.sqrt(np.abs(q)*g*SI.c) / (q*nom_accel_gradient_flattop)
             energy_scaling = np.sqrt(initial_energy*SI.e + q*nom_accel_gradient_flattop*length_flattop) - np.sqrt(initial_energy*SI.e)
@@ -1689,6 +1687,16 @@ def match_length_2_num_beta_osc(self, num_beta_osc, q=SI.e):
 
         # Set the length of the flattop stage
         self.length_flattop = length_flattop
+
+
+    # =============================================
+    @property
+    def external_focusing_gradient(self) -> float:
+        """
+        Return None by default for ``Stage`` subclasses not supporting external 
+        focusing fields.
+        """
+        return None
     
 
     # ==================================================

From 88a8d86011a6b6ffbb832c2843e27c3aeff4f4ee Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 27 Feb 2026 15:54:32 +0100
Subject: [PATCH 63/84] Correction in StageHipace.driver_guiding_trajectory().

---
 abel/classes/stage/impl/stage_hipace.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 9a57b6c8..b78eb23b 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -1111,6 +1111,8 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
             raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')
         
         g = self.external_focusing_gradient  # [T/m]
+        if g is None:
+            g = 0.0
         ds = self.length_flattop/self.driver_half_oscillations/num_steps_per_half_osc  # [m], step size
 
         prop_length = 0

From da77605224c19c1ee582ab4a7f9724810eb0fce9 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 2 Mar 2026 11:43:06 +0100
Subject: [PATCH 64/84] Edited
 abel/wrappers/hipace/hipace_wrapper.py::hipace_write_inputs to set
 external_focusing_gradient to 0.0 when it is None.

---
 abel/wrappers/hipace/hipace_wrapper.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/abel/wrappers/hipace/hipace_wrapper.py b/abel/wrappers/hipace/hipace_wrapper.py
index 7a0b9177..a1d2c00b 100644
--- a/abel/wrappers/hipace/hipace_wrapper.py
+++ b/abel/wrappers/hipace/hipace_wrapper.py
@@ -156,7 +156,9 @@ def hipace_write_inputs(filename_input, filename_beam, filename_driver, plasma_d
         print('>> HiPACE++: Changing from', num_cell_xy, 'to', new_num_cell_xy, ' (i.e., 2^n-1) for better performance.')
         num_cell_xy = new_num_cell_xy
 
-    # plasma-density profile from file
+    # gradient for external magnetic field
+    if external_focusing_gradient is None:
+        external_focusing_gradient = 0.0
     if abs(external_focusing_gradient) > 0:
         external_focusing_comment = ''
     else:

From 3e3dba062a6d29c5018620246b40babacb715d53 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 6 Mar 2026 15:12:48 +0100
Subject: [PATCH 65/84] Added StageHipace.estimate_beam_trajectory() for
 estimating main beam trajectory taking into the effect of acceleration,
 driver guiding and ion background focusing.

---
 abel/classes/stage/impl/stage_hipace.py | 125 ++++++++++++++++++++++++
 1 file changed, 125 insertions(+)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 54c178f7..599c3131 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -1162,6 +1162,131 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
         s_trajectory = s_trajectory + driver.z_offset()
 
         return s_trajectory, x_trajectory, y_trajectory
+
+
+    # =============================================
+    def estimate_beam_trajectory(self, beam, num_steps=None):
+        """
+        Estimate the trajectory that the drive beam will follow when driver 
+        guiding with an external linear azimuthal magnetic field is applied to a 
+        drive beam with an initial angular offset. The calculations  are done by 
+        integrating simplified equations of motion.
+
+        Parameters
+        ----------
+        beam : ``Beam``
+            The main beam.
+        
+        dacc_gradient : [V/m] float, optional
+            The decceleration gradient. Drive beam charge * decceleration 
+            gradient must be negative. Defaults to 0.0.
+
+        num_steps : int, optional
+            ...
+            Defaults to ``None``.
+        
+
+        Returns
+        -------
+        s_trajectory : [m] float
+            Longitudinal coordinate of the drive beam trajectory. Reference is 
+            set at the start of the plasma stage.
+
+        x_trajectory : [m] float
+            x-coordinate of the drive beam trajectory.
+        
+        y_trajectory : [m] float
+            y-coordinate of the drive beam trajectory.
+        """
+
+        from abel.utilities.relativity import energy2momentum
+        from abel.utilities.statistics import weighted_mean
+
+        
+        #pz_thres = energy2momentum(energy_thres, unit='eV', m=driver.particle_mass)
+        pz0 = energy2momentum(beam.energy(), unit='eV', m=beam.particle_mass)
+
+        #if pz0 < pz_thres:
+        #    raise ValueError('This estimate is only valid for a relativistic beam.')
+        
+        q = beam.particle_charge()  # [C], particle charge including charge sign.
+        if q * self.nom_accel_gradient_flattop > 0.0:
+            raise ValueError('Beam charge * self.nom_accel_gradient_flattop  gradient must be negative.')
+        
+        # Make a copy of the stage and set up its ramps if they are not set up
+        ramps_not_set_up = (
+            (self.upramp is not None and self.upramp.length is None) or
+            (self.downramp is not None and self.downramp.length is None)
+        )
+        if ramps_not_set_up:
+            stage_copy = copy.deepcopy(self)
+            stage_copy._prepare_ramps()
+        else: 
+            stage_copy = self
+        
+        L = stage_copy.get_length()  # [m]
+        
+        g0 = SI.e*self.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
+        g = g0
+        if self.external_focusing_gradient is not None:
+            g = g0 + self.external_focusing_gradient
+
+        if num_steps is None:
+            matched_beta = self.matched_beta_function(beam.energy())
+            num_steps = int(L/(matched_beta/20))
+        ds = self.length_flattop/num_steps # [m], step size
+
+        driver_s_trajectory, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(self.driver_source.track(), dacc_gradient=0.0, num_steps_per_half_osc=num_steps)
+
+        prop_length = 0
+        s_trajectory = np.full(num_steps, None, dtype=object)
+        s_trajectory[0] = 0.0
+        x0 = beam.x_offset()
+        #x_trajectory = np.array([x0])  # [m], records the trajectory
+        x_trajectory = np.full(num_steps, None, dtype=object)
+        x_trajectory[0] = x0
+        x = x0
+        y0 = beam.y_offset()
+        #y_trajectory = np.array([y0])  # [m], records the trajectory
+        y_trajectory = np.full(num_steps, None, dtype=object)
+        y_trajectory[0] = y0
+        y = y0
+        px = weighted_mean(beam.pxs(), beam.weightings(), clean=False)
+        py = weighted_mean(beam.pys(), beam.weightings(), clean=False)
+        pz = pz0 # Can add option for deceleration using a gradient
+
+        i = 0
+
+        while prop_length < L:
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
+            y = y + py/pz*1/2*ds
+
+            # Kick
+            dpx = q*g*x*ds - q*g0*driver_x_trajectory[i]*ds
+            px = px + dpx
+            dpy = q*g*y*ds - q*g0*driver_y_trajectory[i]*ds
+            py = py + dpy
+            pz = pz0 - q * self.nom_accel_gradient_flattop * prop_length/SI.c
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
+            y = y + py/pz*1/2*ds
+
+            i = i + 1
+            s_trajectory[i] = prop_length
+            x_trajectory[i] = x
+            y_trajectory[i] = y
+            # s_trajectory = np.append(s_trajectory, prop_length)
+            # x_trajectory = np.append(x_trajectory, x)
+            # y_trajectory = np.append(y_trajectory, y)
+
+        s_trajectory = s_trajectory + beam.location
+
+        return s_trajectory, x_trajectory, y_trajectory
     
 
     # ==================================================

From 8a5658bbb8778748f210103db4e8bc03ed01e0a9 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Fri, 6 Mar 2026 15:27:07 +0100
Subject: [PATCH 66/84] Simplified StageHipace.driver_guiding_trajectory().

---
 abel/classes/stage/impl/stage_hipace.py | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 599c3131..cd0ca6a7 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -1054,7 +1054,7 @@ def match_length_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=Non
 
 
     # =============================================
-    def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_half_osc=100):
+    def driver_guiding_trajectory(self, dacc_gradient=0.0, num_steps=100):
         """
         Estimate the trajectory that the drive beam will follow when driver 
         guiding with an external linear azimuthal magnetic field is applied to a 
@@ -1063,18 +1063,12 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
 
         Parameters
         ----------
-        driver : ``Beam``
-            The drive beam.
-        
         dacc_gradient : [V/m] float, optional
             The decceleration gradient. Drive beam charge * decceleration 
             gradient must be negative. Defaults to 0.0.
 
-        num_steps_per_half_osc : int, optional
-            Number of calcualtion steps per half-oscillation of the drive beam. 
-            The number of half-oscillations is set in 
-            :meth:`StageHipace.calc_external_focusing_gradient() <abel.StageHipace.calc_external_focusing_gradient>`. 
-            Defaults to 100.
+        num_steps : int, optional
+            Number of time steps. Defaults to 100.
         
 
         Returns
@@ -1093,6 +1087,8 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
         from abel.utilities.relativity import energy2momentum
         from abel.utilities.statistics import weighted_mean
 
+        driver = self.driver_source.track()
+
         energy_thres = 10*driver.particle_mass*SI.c**2/SI.e  # [eV], 10 * particle rest energy. Gives beta=0.995.
         pz_thres = energy2momentum(energy_thres, unit='eV', m=driver.particle_mass)
         pz0 = energy2momentum(driver.energy(), unit='eV', m=driver.particle_mass)
@@ -1123,7 +1119,7 @@ def driver_guiding_trajectory(self, driver, dacc_gradient=0.0, num_steps_per_hal
         g = self.external_focusing_gradient  # [T/m]
         if g is None:
             g = 0.0
-        ds = self.length_flattop/self.driver_half_oscillations/num_steps_per_half_osc  # [m], step size
+        ds = self.length_flattop/num_steps  # [m], step size
 
         prop_length = 0
         s_trajectory = np.array([0.0])
@@ -1236,7 +1232,7 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
             num_steps = int(L/(matched_beta/20))
         ds = self.length_flattop/num_steps # [m], step size
 
-        driver_s_trajectory, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(self.driver_source.track(), dacc_gradient=0.0, num_steps_per_half_osc=num_steps)
+        _, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(dacc_gradient=0.0, num_steps=num_steps)
 
         prop_length = 0
         s_trajectory = np.full(num_steps, None, dtype=object)

From fffe736414710e949ad7927e7ea25c3a3fbde914 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Sat, 7 Mar 2026 14:40:06 +0100
Subject: [PATCH 67/84] Corrected a bug in
 StageHipace.estimate_beam_trajectory().

---
 abel/classes/stage/impl/stage_hipace.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index cd0ca6a7..ac1e4012 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -1230,7 +1230,7 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
         if num_steps is None:
             matched_beta = self.matched_beta_function(beam.energy())
             num_steps = int(L/(matched_beta/20))
-        ds = self.length_flattop/num_steps # [m], step size
+        ds = self.length_flattop/(num_steps+1) # [m], step size
 
         _, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(dacc_gradient=0.0, num_steps=num_steps)
 
@@ -1253,7 +1253,7 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
 
         i = 0
 
-        while prop_length < L:
+        while i < num_steps-1:
 
             # Drift
             prop_length = prop_length + 1/2*ds

From 6da11541acc638eb85c84ca7ce08c89913c76693 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 9 Mar 2026 14:36:51 +0100
Subject: [PATCH 68/84] Made corrections in
 StageHipace.estimate_beam_trajectory() to get the correct longitudinal
 coordinate.

---
 abel/classes/stage/impl/stage_hipace.py | 29 ++++++++++---------------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index ac1e4012..dc2ed0bc 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -1054,7 +1054,7 @@ def match_length_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=Non
 
 
     # =============================================
-    def driver_guiding_trajectory(self, dacc_gradient=0.0, num_steps=100):
+    def driver_guiding_trajectory(self, num_steps=100, dacc_gradient=0.0):
         """
         Estimate the trajectory that the drive beam will follow when driver 
         guiding with an external linear azimuthal magnetic field is applied to a 
@@ -1063,12 +1063,12 @@ def driver_guiding_trajectory(self, dacc_gradient=0.0, num_steps=100):
 
         Parameters
         ----------
+        num_steps : int, optional
+            Number of time steps. Defaults to 100.
+
         dacc_gradient : [V/m] float, optional
             The decceleration gradient. Drive beam charge * decceleration 
             gradient must be negative. Defaults to 0.0.
-
-        num_steps : int, optional
-            Number of time steps. Defaults to 100.
         
 
         Returns
@@ -1155,7 +1155,7 @@ def driver_guiding_trajectory(self, dacc_gradient=0.0, num_steps=100):
             x_trajectory = np.append(x_trajectory, x)
             y_trajectory = np.append(y_trajectory, y)
 
-        s_trajectory = s_trajectory + driver.z_offset()
+        #s_trajectory = s_trajectory + driver.z_offset()
 
         return s_trajectory, x_trajectory, y_trajectory
 
@@ -1172,10 +1172,6 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
         ----------
         beam : ``Beam``
             The main beam.
-        
-        dacc_gradient : [V/m] float, optional
-            The decceleration gradient. Drive beam charge * decceleration 
-            gradient must be negative. Defaults to 0.0.
 
         num_steps : int, optional
             ...
@@ -1220,19 +1216,18 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
         else: 
             stage_copy = self
         
-        L = stage_copy.get_length()  # [m]
-        
         g0 = SI.e*self.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
         g = g0
         if self.external_focusing_gradient is not None:
             g = g0 + self.external_focusing_gradient
 
+        L = stage_copy.get_length()  # [m]
         if num_steps is None:
             matched_beta = self.matched_beta_function(beam.energy())
-            num_steps = int(L/(matched_beta/20))
-        ds = self.length_flattop/(num_steps+1) # [m], step size
+            num_steps = int(L /(matched_beta/20))
+        ds = L /(num_steps-1) # [m], step size
 
-        _, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(dacc_gradient=0.0, num_steps=num_steps)
+        _, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(num_steps=num_steps, dacc_gradient=0.0)
 
         prop_length = 0
         s_trajectory = np.full(num_steps, None, dtype=object)
@@ -1249,10 +1244,9 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
         y = y0
         px = weighted_mean(beam.pxs(), beam.weightings(), clean=False)
         py = weighted_mean(beam.pys(), beam.weightings(), clean=False)
-        pz = pz0 # Can add option for deceleration using a gradient
+        pz = pz0
 
         i = 0
-
         while i < num_steps-1:
 
             # Drift
@@ -1280,7 +1274,8 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
             # x_trajectory = np.append(x_trajectory, x)
             # y_trajectory = np.append(y_trajectory, y)
 
-        s_trajectory = s_trajectory + beam.location
+        #print(i, y_trajectory[i], y_trajectory[-1])
+        #s_trajectory = s_trajectory + beam.location
 
         return s_trajectory, x_trajectory, y_trajectory
     

From 240e2d709ce0b48f21d110ca6d1884bb39fb1f97 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 9 Mar 2026 19:42:00 +0100
Subject: [PATCH 69/84] Major changes in StageHipace.estimate_beam_trajectory()
 for adding support for ramps.

---
 abel/classes/stage/impl/stage_hipace.py | 174 ++++++++++++++++++------
 1 file changed, 131 insertions(+), 43 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index dc2ed0bc..492b01e3 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -1163,15 +1163,53 @@ def driver_guiding_trajectory(self, num_steps=100, dacc_gradient=0.0):
     # =============================================
     def estimate_beam_trajectory(self, beam, num_steps=None):
         """
-        Estimate the trajectory that the drive beam will follow when driver 
+        ...
+        """
+
+        from abel.utilities.statistics import weighted_mean
+        from abel.utilities.relativity import energy2momentum
+
+        # Prepare parameters
+        x0 = beam.x_offset()
+        y0 = beam.y_offset()
+        weights = beam.weightings()
+        px0 = weighted_mean(beam.pxs(), weights, clean=False)
+        py0 = weighted_mean(beam.pys(), weights, clean=False)
+        pz0 = energy2momentum(beam.energy(), unit='eV', m=beam.particle_mass)
+        q = beam.particle_charge()
+
+        L = self.get_length()  # [m], total length including any ramps.
+        if num_steps is None:
+            matched_beta = self.matched_beta_function(beam.energy())
+            num_steps = int(L /(matched_beta/20))
+
+        # Actual calculations
+        s_trajectory, x_trajectory, y_trajectory, _, _ = self._estimate_beam_trajectory(s0=0.0, 
+                                                        x0=x0, 
+                                                        y0=y0, 
+                                                        px0=px0, 
+                                                        py0=py0, 
+                                                        pz0=pz0, 
+                                                        q=q, 
+                                                        num_steps=num_steps, 
+                                                        driver_x_trajectory=None, 
+                                                        driver_y_trajectory=None)
+
+        return s_trajectory, x_trajectory, y_trajectory
+
+    # =============================================
+    def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None, driver_x_trajectory=None, driver_y_trajectory=None):
+        """
+        Estimate the trajectory that the main beam will follow when driver 
         guiding with an external linear azimuthal magnetic field is applied to a 
         drive beam with an initial angular offset. The calculations  are done by 
         integrating simplified equations of motion.
 
+        ...
+
         Parameters
         ----------
-        beam : ``Beam``
-            The main beam.
+        
 
         num_steps : int, optional
             ...
@@ -1189,19 +1227,9 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
         
         y_trajectory : [m] float
             y-coordinate of the drive beam trajectory.
-        """
-
-        from abel.utilities.relativity import energy2momentum
-        from abel.utilities.statistics import weighted_mean
 
-        
-        #pz_thres = energy2momentum(energy_thres, unit='eV', m=driver.particle_mass)
-        pz0 = energy2momentum(beam.energy(), unit='eV', m=beam.particle_mass)
-
-        #if pz0 < pz_thres:
-        #    raise ValueError('This estimate is only valid for a relativistic beam.')
-        
-        q = beam.particle_charge()  # [C], particle charge including charge sign.
+        ...
+        """
         if q * self.nom_accel_gradient_flattop > 0.0:
             raise ValueError('Beam charge * self.nom_accel_gradient_flattop  gradient must be negative.')
         
@@ -1216,38 +1244,86 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
         else: 
             stage_copy = self
         
-        g0 = SI.e*self.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
+        # Calculate the focusing field gradient
+        g0 = SI.e*stage_copy.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
         g = g0
-        if self.external_focusing_gradient is not None:
-            g = g0 + self.external_focusing_gradient
+        if stage_copy.external_focusing_gradient is not None:
+            g = g0 + stage_copy.external_focusing_gradient
 
-        L = stage_copy.get_length()  # [m]
-        if num_steps is None:
-            matched_beta = self.matched_beta_function(beam.energy())
-            num_steps = int(L /(matched_beta/20))
-        ds = L /(num_steps-1) # [m], step size
+        # Only calculate the time step size when calling estimate_beam_trajectory() from a flattop
+        if not stage_copy.is_upramp() and not stage_copy.is_downramp():
+            L = stage_copy.get_length()  # [m], total length including any ramps.
+            ds = L /(num_steps-1) # [m], step size
 
-        _, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(num_steps=num_steps, dacc_gradient=0.0)
+        # Calculate the drive beam trajectory
+        if driver_x_trajectory is None or driver_y_trajectory is None:
+            _, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(num_steps=num_steps, dacc_gradient=0.0)
 
-        prop_length = 0
+        # Initialise arrays
         s_trajectory = np.full(num_steps, None, dtype=object)
-        s_trajectory[0] = 0.0
-        x0 = beam.x_offset()
-        #x_trajectory = np.array([x0])  # [m], records the trajectory
         x_trajectory = np.full(num_steps, None, dtype=object)
-        x_trajectory[0] = x0
-        x = x0
-        y0 = beam.y_offset()
-        #y_trajectory = np.array([y0])  # [m], records the trajectory
         y_trajectory = np.full(num_steps, None, dtype=object)
-        y_trajectory[0] = y0
+        px_trajectory = np.full(num_steps, None, dtype=object)
+        py_trajectory = np.full(num_steps, None, dtype=object)
+
+        # Recursive call from the upramp
+        if stage_copy.upramp is not None:
+            upramp = stage_copy.convert_PlasmaRamp(stage_copy.upramp)
+            if self.external_focusing:
+                upramp.external_focusing_gradient = stage_copy.external_focusing_gradient
+
+            num_steps_upramp = int(upramp.length_flattop/ds)
+
+            s_trajectory_upramp, x_trajectory_upramp, y_trajectory_upramp, px_trajectory_upramp, py_trajectory_upramp = upramp._estimate_beam_trajectory(s0, x0, y0, px0, py0, pz0, q, num_steps_upramp, driver_x_trajectory, driver_y_trajectory)
+
+            # Initial parameters for the flattop
+            prop_length = s_trajectory_upramp[-1]
+            s_trajectory[:len(s_trajectory_upramp)] = s_trajectory_upramp
+            x0 = x_trajectory_upramp[-1]
+            x_trajectory[:len(x_trajectory_upramp)] = x_trajectory_upramp
+            y0 = y_trajectory_upramp[-1]
+            y_trajectory[:len(y_trajectory_upramp)] = y_trajectory_upramp
+            px0 = px_trajectory_upramp[-1]
+            px_trajectory[:len(px_trajectory_upramp)] = px_trajectory_upramp
+            py0 = py_trajectory_upramp[-1]
+            py_trajectory[:len(py_trajectory_upramp)] = py_trajectory_upramp
+            pz0 = pz0 - q * stage_copy.upramp.nom_accel_gradient_flattop * prop_length/SI.c
+
+            i = num_steps_upramp - 1
+
+            if stage_copy.downramp is not None:
+                num_steps_downramp = int(stage_copy.downramp.length_flattop/ds)
+                i_end = num_steps - num_steps_downramp - 1
+            else:
+                i_end = num_steps - 1
+
+        # No ramps
+        else:
+            prop_length = s0
+            s_trajectory[0] = prop_length
+            x_trajectory[0] = x0
+            y_trajectory[0] = y0
+            px_trajectory[0] = px0
+            py_trajectory[0] = py0
+
+            i = 0
+            i_end = num_steps - 1
+
+            if self.is_downramp():
+                
+                # Extract the part of drive beam trajectory for the downramp
+                driver_index = num_steps - num_steps_downramp - 1
+                driver_x_trajectory = driver_x_trajectory[driver_index:]
+                driver_y_trajectory = driver_y_trajectory[driver_index:]
+
+        # Set the initial conditions
+        x = x0
         y = y0
-        px = weighted_mean(beam.pxs(), beam.weightings(), clean=False)
-        py = weighted_mean(beam.pys(), beam.weightings(), clean=False)
+        px = px0
+        py = py0
         pz = pz0
 
-        i = 0
-        while i < num_steps-1:
+        while i < i_end:
 
             # Drift
             prop_length = prop_length + 1/2*ds
@@ -1270,14 +1346,26 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
             s_trajectory[i] = prop_length
             x_trajectory[i] = x
             y_trajectory[i] = y
-            # s_trajectory = np.append(s_trajectory, prop_length)
-            # x_trajectory = np.append(x_trajectory, x)
-            # y_trajectory = np.append(y_trajectory, y)
+            px_trajectory[i] = px
+            py_trajectory[i] = py
 
-        #print(i, y_trajectory[i], y_trajectory[-1])
-        #s_trajectory = s_trajectory + beam.location
+        # Recursive call from the downramp
+        if stage_copy.downramp is not None:
+            downramp = stage_copy.convert_PlasmaRamp(stage_copy.downramp)
+            if self.external_focusing:
+                downramp.external_focusing_gradient = stage_copy.external_focusing_gradient
 
-        return s_trajectory, x_trajectory, y_trajectory
+            num_steps_downramp = int(downramp.length_flattop/ds)
+
+            s_trajectory_downramp, x_trajectory_downramp, y_trajectory_downramp, px_trajectory_downramp, py_trajectory_downramp = downramp._estimate_beam_trajectory(prop_length, x, y, px, py, pz, q, num_steps_downramp, driver_x_trajectory, driver_y_trajectory)
+
+            s_trajectory[-len(s_trajectory_downramp):] = s_trajectory_downramp
+            x_trajectory[-len(x_trajectory_downramp):] = x_trajectory_downramp
+            y_trajectory[-len(y_trajectory_downramp):] = y_trajectory_downramp
+            px_trajectory[-len(px_trajectory_downramp):] = px_trajectory_downramp
+            py_trajectory[-len(py_trajectory_downramp):] = py_trajectory_downramp
+
+        return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory
     
 
     # ==================================================

From 74d9f48cfd7e0651888021dcedac68a70f05402c Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 10 Mar 2026 15:51:11 +0100
Subject: [PATCH 70/84] Correction in StageHipace._estimate_beam_trajectory for
 calculating the beam trajectory when ramps are included.

---
 abel/classes/stage/impl/stage_hipace.py | 119 ++++++++++++++++++++++--
 1 file changed, 111 insertions(+), 8 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 492b01e3..89e50e7d 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -1230,6 +1230,8 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
 
         ...
         """
+        from abel.utilities.other import find_closest_value_in_arr
+
         if q * self.nom_accel_gradient_flattop > 0.0:
             raise ValueError('Beam charge * self.nom_accel_gradient_flattop  gradient must be negative.')
         
@@ -1252,8 +1254,60 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
 
         # Only calculate the time step size when calling estimate_beam_trajectory() from a flattop
         if not stage_copy.is_upramp() and not stage_copy.is_downramp():
+
+            #print('not stage_copy.is_upramp() and not stage_copy.is_downramp():', not stage_copy.is_upramp() and not stage_copy.is_downramp()) # TODO: delete
+
+            # Set the step size
             L = stage_copy.get_length()  # [m], total length including any ramps.
-            ds = L /(num_steps-1) # [m], step size
+            stage_copy.ds = L /(num_steps-1) # [m], step size
+
+
+            if stage_copy.has_ramp():
+                # Set the number of time steps in the upramp and downramp
+                #ss_helper = np.linspace(0.0, L , num_steps)
+                ss_helper = np.arange(num_steps) * stage_copy.ds
+                num_steps_upramp, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length)
+                num_steps_upramp = num_steps_upramp + 1
+                idx_flat_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length+stage_copy.length_flattop)
+                
+                num_steps_downramp = num_steps - idx_flat_end
+                stage_copy.idx_flat_end = idx_flat_end
+                stage_copy.num_steps_upramp = num_steps_upramp
+                stage_copy.num_steps_downramp = num_steps_downramp
+
+
+
+            ds = stage_copy.ds # TODO: delete
+
+            
+            # TODO: delete
+            #print('np.max(np.diff(ss_helper))/ds', np.max(np.diff(ss_helper))/ ds)
+
+            #num_flat = idx_flat_end - stage_copy.num_steps_upramp
+            #num_down = num_steps - idx_flat_end
+
+
+            
+
+            # print('num_steps_upramp, num_steps_flattop..., :', 
+            #       stage_copy.num_steps_upramp , 
+            #       num_flat, 
+            #       stage_copy.num_steps_downramp, 
+            #       num_steps, 
+            #       stage_copy.num_steps_upramp + num_flat + stage_copy.num_steps_downramp) # TODO: delete
+
+            #print('ss_helper[stage_copy.num_steps_upramp ]', ss_helper[stage_copy.num_steps_upramp-1 ], stage_copy.upramp.length_flattop, '\n')
+
+            #print('ss_helper[stage_copy.idx_flat_end]:', ss_helper[stage_copy.idx_flat_end], stage_copy.upramp.length + stage_copy.length_flattop, stage_copy.length - stage_copy.downramp.length, '\n')
+
+            #print('ss_helper[-1], ss_helper[num_up+num_flat+num_down]', ss_helper[-1], ss_helper[stage_copy.num_steps_upramp+num_flat+stage_copy.num_steps_downramp-1], stage_copy.length, ss_helper[-1]/stage_copy.length, '\n')
+
+            
+
+        #ds = stage_copy.get_length() /(num_steps-1) # [m], step size. Determined by the total length of either ramp or flattop+ramps.
+        ds = stage_copy.ds
+
+        
 
         # Calculate the drive beam trajectory
         if driver_x_trajectory is None or driver_y_trajectory is None:
@@ -1269,10 +1323,12 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
         # Recursive call from the upramp
         if stage_copy.upramp is not None:
             upramp = stage_copy.convert_PlasmaRamp(stage_copy.upramp)
+            stage_copy.upramp = upramp
             if self.external_focusing:
                 upramp.external_focusing_gradient = stage_copy.external_focusing_gradient
 
-            num_steps_upramp = int(upramp.length_flattop/ds)
+            #num_steps_upramp = int(upramp.length/ds)
+            num_steps_upramp = stage_copy.num_steps_upramp 
 
             s_trajectory_upramp, x_trajectory_upramp, y_trajectory_upramp, px_trajectory_upramp, py_trajectory_upramp = upramp._estimate_beam_trajectory(s0, x0, y0, px0, py0, pz0, q, num_steps_upramp, driver_x_trajectory, driver_y_trajectory)
 
@@ -1289,11 +1345,16 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
             py_trajectory[:len(py_trajectory_upramp)] = py_trajectory_upramp
             pz0 = pz0 - q * stage_copy.upramp.nom_accel_gradient_flattop * prop_length/SI.c
 
+            print('s_trajectory_upramp[-1]/stage_copy.upramp.length', s_trajectory_upramp[-1]/stage_copy.upramp.length)
+
             i = num_steps_upramp - 1
 
             if stage_copy.downramp is not None:
-                num_steps_downramp = int(stage_copy.downramp.length_flattop/ds)
-                i_end = num_steps - num_steps_downramp - 1
+                #num_steps_downramp = int(stage_copy.downramp.length_flattop/ds)
+                # num_steps_downramp = stage_copy.num_steps_downramp
+                # print('INSIDE stage_copy.downramp is not None:', num_steps_downramp) #TODO: DELETE
+                # i_end = num_steps - num_steps_downramp - 1
+                i_end = stage_copy.idx_flat_end
             else:
                 i_end = num_steps - 1
 
@@ -1309,12 +1370,25 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
             i = 0
             i_end = num_steps - 1
 
+
+            # # TODO: delete
+            # s_test = np.full(num_steps, ds, dtype=float)
+            # s_test[0] = 0.0
+            # s_test = np.cumsum(s_test)
+            # print('s_test[-1]: ...', s_test[-1], self.length, s_test[-1]/self.length)
+
             if self.is_downramp():
                 
                 # Extract the part of drive beam trajectory for the downramp
-                driver_index = num_steps - num_steps_downramp - 1
-                driver_x_trajectory = driver_x_trajectory[driver_index:]
-                driver_y_trajectory = driver_y_trajectory[driver_index:]
+                #num_steps_downramp = int(self.length_flattop/ds)
+                #num_steps_downramp = len(s_trajectory)
+                num_steps_downramp = stage_copy.num_steps_downramp 
+                driver_index = num_steps_downramp
+                driver_x_trajectory = driver_x_trajectory[-driver_index:]
+                driver_y_trajectory = driver_y_trajectory[-driver_index:]
+
+                
+                print('inside self.is_downramp():', len(driver_y_trajectory), driver_index, len(s_trajectory), num_steps) # TODO: delete
 
         # Set the initial conditions
         x = x0
@@ -1323,6 +1397,9 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
         py = py0
         pz = pz0
 
+        #print('ds:', ds) #TODO:delete
+
+        # Solve the equations of motion
         while i < i_end:
 
             # Drift
@@ -1352,19 +1429,45 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
         # Recursive call from the downramp
         if stage_copy.downramp is not None:
             downramp = stage_copy.convert_PlasmaRamp(stage_copy.downramp)
+            stage_copy.downramp = downramp
             if self.external_focusing:
                 downramp.external_focusing_gradient = stage_copy.external_focusing_gradient
 
-            num_steps_downramp = int(downramp.length_flattop/ds)
+            #num_steps_downramp = int(downramp.length_flattop/ds)
+            num_steps_downramp = stage_copy.num_steps_downramp 
+
+            #print(downramp.is_downramp(), stage_copy.downramp.is_downramp(), num_steps_downramp ) # TODO: delete
 
             s_trajectory_downramp, x_trajectory_downramp, y_trajectory_downramp, px_trajectory_downramp, py_trajectory_downramp = downramp._estimate_beam_trajectory(prop_length, x, y, px, py, pz, q, num_steps_downramp, driver_x_trajectory, driver_y_trajectory)
 
+
+
             s_trajectory[-len(s_trajectory_downramp):] = s_trajectory_downramp
             x_trajectory[-len(x_trajectory_downramp):] = x_trajectory_downramp
             y_trajectory[-len(y_trajectory_downramp):] = y_trajectory_downramp
             px_trajectory[-len(px_trajectory_downramp):] = px_trajectory_downramp
             py_trajectory[-len(py_trajectory_downramp):] = py_trajectory_downramp
 
+            # TODO: delete
+
+            print('s_trajectory_downramp[-1]/stage_copy.length...', s_trajectory_downramp[-1]/stage_copy.length, 
+                  s_trajectory[-1]/stage_copy.length, '\n')
+            
+            #print('s_trajectory_downramp[0]/s_trajectory[stage_copy.idx_flat_end]', s_trajectory_downramp[0]/s_trajectory[stage_copy.idx_flat_end], '\n')
+            
+            print('(s_trajectory_downramp[-1]-s_trajectory_downramp[0])/stage_copy.downramp.length..', (s_trajectory_downramp[-1]-s_trajectory_downramp[0])/stage_copy.downramp.length,
+                  (s_trajectory_downramp[-1]-s_trajectory[-len(s_trajectory_downramp)])/stage_copy.downramp.length,
+                   '\n')
+            
+            print('s_trajectory[stage_copy.num_steps_upramp]/stage_copy.upramp.length', 
+                  s_trajectory[stage_copy.num_steps_upramp-1]/stage_copy.upramp.length, '\n')
+
+            print(s_trajectory[-len(s_trajectory_downramp)]/(stage_copy.length-stage_copy.downramp.length), 
+                  
+                  s_trajectory[-len(s_trajectory_downramp)]/(stage_copy.length_flattop+stage_copy.upramp.length)) #TODO: delete
+            
+
+
         return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory
     
 

From 331727809acafbda6144ba1533391b154d44327d Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 10 Mar 2026 18:37:00 +0100
Subject: [PATCH 71/84] Deleted most of the printing diagnostics in
 StageHipace._estimate_beam_trajectory(), added doc strings to
 StageHipace._estimate_beam_trajectory(),
 StageHipace.estimate_beam_trajectory() and edited
 StageHipace.driver_guiding_trajectory() so that its solver is based on the
 total number of time steps.

---
 abel/classes/stage/impl/stage_hipace.py | 271 +++++++++++++-----------
 1 file changed, 153 insertions(+), 118 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index 89e50e7d..bbd7f8de 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -1054,17 +1054,19 @@ def match_length_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=Non
 
 
     # =============================================
-    def driver_guiding_trajectory(self, num_steps=100, dacc_gradient=0.0):
+    def driver_guiding_trajectory(self, num_steps=None, dacc_gradient=0.0):
         """
         Estimate the trajectory that the drive beam will follow when driver 
         guiding with an external linear azimuthal magnetic field is applied to a 
-        drive beam with an initial angular offset. The calculations  are done by 
+        drive beam with an initial angular offset. The calculations are done by 
         integrating simplified equations of motion.
 
         Parameters
         ----------
         num_steps : int, optional
-            Number of time steps. Defaults to 100.
+            Number of time steps. If ``None``, will calculate the number of time 
+            steps such that the step size is a small fraction of the matched 
+            beta function of the drive beam. Defaults to ``None``.
 
         dacc_gradient : [V/m] float, optional
             The decceleration gradient. Drive beam charge * decceleration 
@@ -1073,14 +1075,14 @@ def driver_guiding_trajectory(self, num_steps=100, dacc_gradient=0.0):
 
         Returns
         -------
-        s_trajectory : [m] float
+        s_trajectory : [m] 1D float ndarray
             Longitudinal coordinate of the drive beam trajectory. Reference is 
             set at the start of the plasma stage.
 
-        x_trajectory : [m] float
+        x_trajectory : [m] 1D float ndarray
             x-coordinate of the drive beam trajectory.
         
-        y_trajectory : [m] float
+        y_trajectory : [m] 1D float ndarray
             y-coordinate of the drive beam trajectory.
         """
 
@@ -1116,24 +1118,39 @@ def driver_guiding_trajectory(self, num_steps=100, dacc_gradient=0.0):
         if pz0 + q * dacc_gradient * L/SI.c < pz_thres:
             raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')
         
+        # Get the focusing field gradient
         g = self.external_focusing_gradient  # [T/m]
         if g is None:
             g = 0.0
-        ds = self.length_flattop/num_steps  # [m], step size
 
+        # Determine the step size
+        if num_steps is None:
+            matched_beta = self.matched_beta_function(driver.energy())
+            num_steps = int(L /(matched_beta/20))
+        
+        ds = L/(num_steps-1)  # [m], step size
+
+        # Initialise arrays
+        s_trajectory = np.full(num_steps, None, dtype=object)
+        x_trajectory = np.full(num_steps, None, dtype=object)
+        y_trajectory = np.full(num_steps, None, dtype=object)
+
+        # Set initial parameters
         prop_length = 0
-        s_trajectory = np.array([0.0])
         x0 = driver.x_offset()
-        x_trajectory = np.array([x0])  # [m], records the trajectory
         x = x0
         y0 = driver.y_offset()
-        y_trajectory = np.array([y0])  # [m], records the trajectory
         y = y0
+        s_trajectory[0] = prop_length
+        x_trajectory[0] = x0
+        y_trajectory[0] = y0
         px = weighted_mean(driver.pxs(), driver.weightings(), clean=False)
         py = weighted_mean(driver.pys(), driver.weightings(), clean=False)
         pz = pz0 # Can add option for deceleration using a gradient
 
-        while prop_length < L:
+        # Solve the equation of motion
+        i = 0
+        while i < num_steps - 1:
 
             # Drift
             prop_length = prop_length + 1/2*ds
@@ -1145,17 +1162,17 @@ def driver_guiding_trajectory(self, num_steps=100, dacc_gradient=0.0):
             px = px + dpx
             dpy = q*g*y*ds
             py = py + dpy
-            pz = pz0 + q * dacc_gradient * prop_length/SI.c # dacc_gradient>0
+            pz = pz0 + q * dacc_gradient * prop_length/SI.c # dacc_gradient > 0
 
             # Drift
             prop_length = prop_length + 1/2*ds
             x = x + px/pz*1/2*ds
             y = y + py/pz*1/2*ds
-            s_trajectory = np.append(s_trajectory, prop_length)
-            x_trajectory = np.append(x_trajectory, x)
-            y_trajectory = np.append(y_trajectory, y)
 
-        #s_trajectory = s_trajectory + driver.z_offset()
+            i = i + 1
+            s_trajectory[i] = prop_length
+            x_trajectory[i] = x
+            y_trajectory[i] = y
 
         return s_trajectory, x_trajectory, y_trajectory
 
@@ -1163,7 +1180,36 @@ def driver_guiding_trajectory(self, num_steps=100, dacc_gradient=0.0):
     # =============================================
     def estimate_beam_trajectory(self, beam, num_steps=None):
         """
-        ...
+        Estimate the trajectory for the main beam following the trajectory of a 
+        drive beam generated by ``self.driver_source``.
+        
+        Effects such as driver guiding with an external linear azimuthal 
+        magnetic field, background ion focusing and uniform plasma density ramps 
+        are taken into account. The calculations are done by integrating 
+        simplified equations of motion.
+
+        Parameters
+        ----------
+        beam : ``Beam``
+            The main beam to be tracked.
+
+        num_steps : int, optional
+            Number of time steps. If ``None``, will calculate the number of time 
+            steps such that the step size is a small fraction of the matched 
+            beta function of the main beam. Defaults to ``None``.
+
+
+        Returns
+        -------
+        s_trajectory : [m] 1D float ndarray
+            Longitudinal coordinate of the main beam trajectory. Reference is 
+            set at the start of the plasma stage.
+
+        x_trajectory : [m] 1D float ndarray
+            x-coordinate of the main beam trajectory.
+        
+        y_trajectory : [m] 1D float ndarray
+            y-coordinate of the main beam trajectory.
         """
 
         from abel.utilities.statistics import weighted_mean
@@ -1197,38 +1243,74 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
 
         return s_trajectory, x_trajectory, y_trajectory
 
+
     # =============================================
-    def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None, driver_x_trajectory=None, driver_y_trajectory=None):
+    def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, driver_x_trajectory=None, driver_y_trajectory=None):
         """
-        Estimate the trajectory that the main beam will follow when driver 
-        guiding with an external linear azimuthal magnetic field is applied to a 
-        drive beam with an initial angular offset. The calculations  are done by 
-        integrating simplified equations of motion.
-
-        ...
+        Helper function for estimating the trajectory for the main beam 
+        following the trajectory of a drive beam defined by 
+        ``driver_x_trajectory`` and ``driver_y_trajectory``.
+        
+        Effects such as driver guiding with an external linear azimuthal 
+        magnetic field, background ion focusing and uniform plasma density ramps 
+        are taken into account. The calculations are done by integrating 
+        simplified equations of motion.
+         
 
         Parameters
         ----------
-        
+        s0 : [m] float
+            The initial longitudinal coordinate of the main beam.
 
-        num_steps : int, optional
-            ...
-            Defaults to ``None``.
+        x0 : [m] float
+            The intial x-coordinate of the main beam.
+
+        y0 : [m] float
+            The intial y-coordinate of the main beam.
+
+        px0 : [kg m/s] float
+            The intial x-momentum of the main beam.
+
+        py0 : [kg m/s] float
+            The intial y-momentum of the main beam.
+
+        pz0 : [kg m/s] float
+            The intial z-momentum of the main beam.
+
+        q : [C] flloat
+            The particle charge of the main beam.
+
+        num_steps : int
+            Number of time steps.
+
+        driver_x_trajectory : [m] 1D float ndarray, optional
+            The x-coordinate of trajectory of the drive beam. The length of 
+            ``driver_x_trajectory`` must be the same as ``num_steps``. Is 
+            automatically calculated if ``None``. Defaults to ``None``.
+
+        driver_y_trajectory : [m] 1D float ndarray, optional
+            The y-coordinate of trajectory of the drive beam. The length of 
+            ``driver_y_trajectory`` must be the same as ``num_steps``. Is 
+            automatically calculated if ``None``. Defaults to ``None``.
         
 
         Returns
         -------
-        s_trajectory : [m] float
-            Longitudinal coordinate of the drive beam trajectory. Reference is 
+        s_trajectory : [m] 1D float ndarray
+            Longitudinal coordinate of the main beam trajectory. Reference is 
             set at the start of the plasma stage.
 
-        x_trajectory : [m] float
-            x-coordinate of the drive beam trajectory.
+        x_trajectory : [m] 1D float ndarray
+            x-coordinate of the main beam trajectory.
         
-        y_trajectory : [m] float
-            y-coordinate of the drive beam trajectory.
+        y_trajectory : [m] 1D float ndarray
+            y-coordinate of the main beam trajectory.
+
+        px_trajectory : [kg m/s] 1D float ndarray
+            Mean x-component of the beam momentum along the trajectory.
 
-        ...
+        py_trajectory : [kg m/s] 1D float ndarray
+            Mean y-component of the beam momentum along the trajectory.
         """
         from abel.utilities.other import find_closest_value_in_arr
 
@@ -1255,63 +1337,30 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
         # Only calculate the time step size when calling estimate_beam_trajectory() from a flattop
         if not stage_copy.is_upramp() and not stage_copy.is_downramp():
 
-            #print('not stage_copy.is_upramp() and not stage_copy.is_downramp():', not stage_copy.is_upramp() and not stage_copy.is_downramp()) # TODO: delete
-
             # Set the step size
             L = stage_copy.get_length()  # [m], total length including any ramps.
-            stage_copy.ds = L /(num_steps-1) # [m], step size
-
+            stage_copy.ds = L / (num_steps-1) # [m], step size
 
             if stage_copy.has_ramp():
                 # Set the number of time steps in the upramp and downramp
-                #ss_helper = np.linspace(0.0, L , num_steps)
                 ss_helper = np.arange(num_steps) * stage_copy.ds
-                num_steps_upramp, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length)
-                num_steps_upramp = num_steps_upramp + 1
-                idx_flat_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length+stage_copy.length_flattop)
-                
-                num_steps_downramp = num_steps - idx_flat_end
+                idx_upramp_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length)
+                num_steps_upramp = idx_upramp_end + 1  # Number of time steps for the upramp
+                idx_flat_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length+stage_copy.length_flattop)  # Index marking the end of the flattop.
+                num_steps_downramp = num_steps - idx_flat_end  # Number of time steps for the downramp
                 stage_copy.idx_flat_end = idx_flat_end
                 stage_copy.num_steps_upramp = num_steps_upramp
                 stage_copy.num_steps_downramp = num_steps_downramp
 
-
-
-            ds = stage_copy.ds # TODO: delete
-
-            
-            # TODO: delete
-            #print('np.max(np.diff(ss_helper))/ds', np.max(np.diff(ss_helper))/ ds)
-
-            #num_flat = idx_flat_end - stage_copy.num_steps_upramp
-            #num_down = num_steps - idx_flat_end
-
-
-            
-
-            # print('num_steps_upramp, num_steps_flattop..., :', 
-            #       stage_copy.num_steps_upramp , 
-            #       num_flat, 
-            #       stage_copy.num_steps_downramp, 
-            #       num_steps, 
-            #       stage_copy.num_steps_upramp + num_flat + stage_copy.num_steps_downramp) # TODO: delete
-
-            #print('ss_helper[stage_copy.num_steps_upramp ]', ss_helper[stage_copy.num_steps_upramp-1 ], stage_copy.upramp.length_flattop, '\n')
-
-            #print('ss_helper[stage_copy.idx_flat_end]:', ss_helper[stage_copy.idx_flat_end], stage_copy.upramp.length + stage_copy.length_flattop, stage_copy.length - stage_copy.downramp.length, '\n')
-
-            #print('ss_helper[-1], ss_helper[num_up+num_flat+num_down]', ss_helper[-1], ss_helper[stage_copy.num_steps_upramp+num_flat+stage_copy.num_steps_downramp-1], stage_copy.length, ss_helper[-1]/stage_copy.length, '\n')
-
-            
-
-        #ds = stage_copy.get_length() /(num_steps-1) # [m], step size. Determined by the total length of either ramp or flattop+ramps.
         ds = stage_copy.ds
 
-        
-
         # Calculate the drive beam trajectory
         if driver_x_trajectory is None or driver_y_trajectory is None:
             _, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(num_steps=num_steps, dacc_gradient=0.0)
+        
+        if not stage_copy.is_upramp() and not stage_copy.is_downramp():
+            if len(driver_x_trajectory) != num_steps or len(driver_y_trajectory) != num_steps:
+                raise ValueError('The length of driver_x_trajectory and driver_y_trajectory must be the same as num_steps.')
 
         # Initialise arrays
         s_trajectory = np.full(num_steps, None, dtype=object)
@@ -1327,7 +1376,6 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
             if self.external_focusing:
                 upramp.external_focusing_gradient = stage_copy.external_focusing_gradient
 
-            #num_steps_upramp = int(upramp.length/ds)
             num_steps_upramp = stage_copy.num_steps_upramp 
 
             s_trajectory_upramp, x_trajectory_upramp, y_trajectory_upramp, px_trajectory_upramp, py_trajectory_upramp = upramp._estimate_beam_trajectory(s0, x0, y0, px0, py0, pz0, q, num_steps_upramp, driver_x_trajectory, driver_y_trajectory)
@@ -1345,15 +1393,11 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
             py_trajectory[:len(py_trajectory_upramp)] = py_trajectory_upramp
             pz0 = pz0 - q * stage_copy.upramp.nom_accel_gradient_flattop * prop_length/SI.c
 
-            print('s_trajectory_upramp[-1]/stage_copy.upramp.length', s_trajectory_upramp[-1]/stage_copy.upramp.length)
+            #print('s_trajectory_upramp[-1]/stage_copy.upramp.length', s_trajectory_upramp[-1]/stage_copy.upramp.length)  # TODO: delete
 
             i = num_steps_upramp - 1
 
             if stage_copy.downramp is not None:
-                #num_steps_downramp = int(stage_copy.downramp.length_flattop/ds)
-                # num_steps_downramp = stage_copy.num_steps_downramp
-                # print('INSIDE stage_copy.downramp is not None:', num_steps_downramp) #TODO: DELETE
-                # i_end = num_steps - num_steps_downramp - 1
                 i_end = stage_copy.idx_flat_end
             else:
                 i_end = num_steps - 1
@@ -1370,25 +1414,27 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
             i = 0
             i_end = num_steps - 1
 
+            if self.is_downramp():
 
-            # # TODO: delete
-            # s_test = np.full(num_steps, ds, dtype=float)
-            # s_test[0] = 0.0
-            # s_test = np.cumsum(s_test)
-            # print('s_test[-1]: ...', s_test[-1], self.length, s_test[-1]/self.length)
+                # #TODO: delete
+                # import matplotlib.pyplot as plt
+                # plt.figure()
+                # plt.plot(driver_s_trajectory, driver_y_trajectory*1e6)
 
-            if self.is_downramp():
-                
-                # Extract the part of drive beam trajectory for the downramp
-                #num_steps_downramp = int(self.length_flattop/ds)
-                #num_steps_downramp = len(s_trajectory)
-                num_steps_downramp = stage_copy.num_steps_downramp 
-                driver_index = num_steps_downramp
-                driver_x_trajectory = driver_x_trajectory[-driver_index:]
-                driver_y_trajectory = driver_y_trajectory[-driver_index:]
+                # ss_driver = driver_s_trajectory[stage_copy.idx_flat_end:]
+                # plt.plot(ss_driver, driver_y_trajectory[stage_copy.idx_flat_end:]*1e6)
+
+                # print('len(driver_s_trajectory)', len(driver_s_trajectory), num_steps)
+                # print('driver_s_trajectory endpoint control', driver_s_trajectory[-1]/self.parent.length)
+                # print('driver downramp starting point', ss_driver[0], driver_s_trajectory[stage_copy.idx_flat_end])
+                # print((ss_driver[-1]-ss_driver[0]), self.length)
+                # print('(ss_driver[-1]-ss_driver[0])/ stage_copy.downramp.length',(ss_driver[-1]-ss_driver[0])/ self.length)
 
                 
-                print('inside self.is_downramp():', len(driver_y_trajectory), driver_index, len(s_trajectory), num_steps) # TODO: delete
+                
+                # Extract the part of drive beam trajectory for the downramp
+                driver_x_trajectory = driver_x_trajectory[stage_copy.idx_flat_end:]
+                driver_y_trajectory = driver_y_trajectory[stage_copy.idx_flat_end:]
 
         # Set the initial conditions
         x = x0
@@ -1397,8 +1443,6 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
         py = py0
         pz = pz0
 
-        #print('ds:', ds) #TODO:delete
-
         # Solve the equations of motion
         while i < i_end:
 
@@ -1433,15 +1477,10 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
             if self.external_focusing:
                 downramp.external_focusing_gradient = stage_copy.external_focusing_gradient
 
-            #num_steps_downramp = int(downramp.length_flattop/ds)
-            num_steps_downramp = stage_copy.num_steps_downramp 
-
-            #print(downramp.is_downramp(), stage_copy.downramp.is_downramp(), num_steps_downramp ) # TODO: delete
+            num_steps_downramp = stage_copy.num_steps_downramp
 
             s_trajectory_downramp, x_trajectory_downramp, y_trajectory_downramp, px_trajectory_downramp, py_trajectory_downramp = downramp._estimate_beam_trajectory(prop_length, x, y, px, py, pz, q, num_steps_downramp, driver_x_trajectory, driver_y_trajectory)
 
-
-
             s_trajectory[-len(s_trajectory_downramp):] = s_trajectory_downramp
             x_trajectory[-len(x_trajectory_downramp):] = x_trajectory_downramp
             y_trajectory[-len(y_trajectory_downramp):] = y_trajectory_downramp
@@ -1449,25 +1488,21 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps=None
             py_trajectory[-len(py_trajectory_downramp):] = py_trajectory_downramp
 
             # TODO: delete
-
-            print('s_trajectory_downramp[-1]/stage_copy.length...', s_trajectory_downramp[-1]/stage_copy.length, 
-                  s_trajectory[-1]/stage_copy.length, '\n')
+            print('upramp length control    s_trajectory[stage_copy.num_steps_upramp]/stage_copy.upramp.length', 
+                  s_trajectory[stage_copy.num_steps_upramp-1]/stage_copy.upramp.length, '\n')
             
-            #print('s_trajectory_downramp[0]/s_trajectory[stage_copy.idx_flat_end]', s_trajectory_downramp[0]/s_trajectory[stage_copy.idx_flat_end], '\n')
+            print('flattop length control   ', (s_trajectory[stage_copy.idx_flat_end]-s_trajectory[stage_copy.num_steps_upramp-1])/ stage_copy.length_flattop, '\n')
             
-            print('(s_trajectory_downramp[-1]-s_trajectory_downramp[0])/stage_copy.downramp.length..', (s_trajectory_downramp[-1]-s_trajectory_downramp[0])/stage_copy.downramp.length,
+            print('downramp starting point   s_trajectory_downramp[0]/s_trajectory[stage_copy.idx_flat_end]', s_trajectory_downramp[0]/s_trajectory[stage_copy.idx_flat_end], '\n')
+            
+            print('downramp length control   (s_trajectory_downramp[-1]-s_trajectory_downramp[0])/stage_copy.downramp.length..', (s_trajectory_downramp[-1]-s_trajectory_downramp[0])/stage_copy.downramp.length,
                   (s_trajectory_downramp[-1]-s_trajectory[-len(s_trajectory_downramp)])/stage_copy.downramp.length,
                    '\n')
             
-            print('s_trajectory[stage_copy.num_steps_upramp]/stage_copy.upramp.length', 
-                  s_trajectory[stage_copy.num_steps_upramp-1]/stage_copy.upramp.length, '\n')
-
-            print(s_trajectory[-len(s_trajectory_downramp)]/(stage_copy.length-stage_copy.downramp.length), 
-                  
-                  s_trajectory[-len(s_trajectory_downramp)]/(stage_copy.length_flattop+stage_copy.upramp.length)) #TODO: delete
+            print('stage end point control   s_trajectory_downramp[-1]/stage_copy.length...', s_trajectory_downramp[-1]/stage_copy.length, 
+                  s_trajectory[-1]/stage_copy.length, '\n')
             
 
-
         return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory
     
 

From 17f30a67794c13b7150d631dda375af7b7f4d140 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 10 Mar 2026 18:39:45 +0100
Subject: [PATCH 72/84] Removed the rest of the printing diagnostics in
 StageHipace._estimate_beam_trajectory().

---
 abel/classes/stage/impl/stage_hipace.py | 35 -------------------------
 1 file changed, 35 deletions(-)

diff --git a/abel/classes/stage/impl/stage_hipace.py b/abel/classes/stage/impl/stage_hipace.py
index bbd7f8de..5e8d02b2 100644
--- a/abel/classes/stage/impl/stage_hipace.py
+++ b/abel/classes/stage/impl/stage_hipace.py
@@ -1393,8 +1393,6 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
             py_trajectory[:len(py_trajectory_upramp)] = py_trajectory_upramp
             pz0 = pz0 - q * stage_copy.upramp.nom_accel_gradient_flattop * prop_length/SI.c
 
-            #print('s_trajectory_upramp[-1]/stage_copy.upramp.length', s_trajectory_upramp[-1]/stage_copy.upramp.length)  # TODO: delete
-
             i = num_steps_upramp - 1
 
             if stage_copy.downramp is not None:
@@ -1415,23 +1413,6 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
             i_end = num_steps - 1
 
             if self.is_downramp():
-
-                # #TODO: delete
-                # import matplotlib.pyplot as plt
-                # plt.figure()
-                # plt.plot(driver_s_trajectory, driver_y_trajectory*1e6)
-
-                # ss_driver = driver_s_trajectory[stage_copy.idx_flat_end:]
-                # plt.plot(ss_driver, driver_y_trajectory[stage_copy.idx_flat_end:]*1e6)
-
-                # print('len(driver_s_trajectory)', len(driver_s_trajectory), num_steps)
-                # print('driver_s_trajectory endpoint control', driver_s_trajectory[-1]/self.parent.length)
-                # print('driver downramp starting point', ss_driver[0], driver_s_trajectory[stage_copy.idx_flat_end])
-                # print((ss_driver[-1]-ss_driver[0]), self.length)
-                # print('(ss_driver[-1]-ss_driver[0])/ stage_copy.downramp.length',(ss_driver[-1]-ss_driver[0])/ self.length)
-
-                
-                
                 # Extract the part of drive beam trajectory for the downramp
                 driver_x_trajectory = driver_x_trajectory[stage_copy.idx_flat_end:]
                 driver_y_trajectory = driver_y_trajectory[stage_copy.idx_flat_end:]
@@ -1487,22 +1468,6 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
             px_trajectory[-len(px_trajectory_downramp):] = px_trajectory_downramp
             py_trajectory[-len(py_trajectory_downramp):] = py_trajectory_downramp
 
-            # TODO: delete
-            print('upramp length control    s_trajectory[stage_copy.num_steps_upramp]/stage_copy.upramp.length', 
-                  s_trajectory[stage_copy.num_steps_upramp-1]/stage_copy.upramp.length, '\n')
-            
-            print('flattop length control   ', (s_trajectory[stage_copy.idx_flat_end]-s_trajectory[stage_copy.num_steps_upramp-1])/ stage_copy.length_flattop, '\n')
-            
-            print('downramp starting point   s_trajectory_downramp[0]/s_trajectory[stage_copy.idx_flat_end]', s_trajectory_downramp[0]/s_trajectory[stage_copy.idx_flat_end], '\n')
-            
-            print('downramp length control   (s_trajectory_downramp[-1]-s_trajectory_downramp[0])/stage_copy.downramp.length..', (s_trajectory_downramp[-1]-s_trajectory_downramp[0])/stage_copy.downramp.length,
-                  (s_trajectory_downramp[-1]-s_trajectory[-len(s_trajectory_downramp)])/stage_copy.downramp.length,
-                   '\n')
-            
-            print('stage end point control   s_trajectory_downramp[-1]/stage_copy.length...', s_trajectory_downramp[-1]/stage_copy.length, 
-                  s_trajectory[-1]/stage_copy.length, '\n')
-            
-
         return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory
     
 

From 6a500f6b078be3ec08cf4400c5fc69e20eba3f90 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 11 Mar 2026 00:34:58 +0100
Subject: [PATCH 73/84] Edited Stage.phase_advance_beta_evolution() to take
 into the account of external driver guiding field.

---
 abel/classes/stage/stage.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/abel/classes/stage/stage.py b/abel/classes/stage/stage.py
index 6ba98263..87e91381 100644
--- a/abel/classes/stage/stage.py
+++ b/abel/classes/stage/stage.py
@@ -1726,7 +1726,9 @@ def phase_advance_beta_evolution(self, beta0=None):
         from abel.utilities.beam_physics import evolve_beta_function
         from abel.utilities.beam_physics import phase_advance
         
-        g_ion = SI.e*self.plasma_density/(2*SI.epsilon_0)
+        g = SI.e*self.plasma_density/(2*SI.epsilon_0)
+        if self.external_focusing_gradient is not None:
+            g = g + self.external_focusing_gradient * SI.c
         p0 = np.sqrt((self.nom_energy*SI.e)**2-(SI.m_e*SI.c**2)**2)/SI.c
         if beta0 is None:
             if self.is_upramp():
@@ -1737,7 +1739,7 @@ def phase_advance_beta_evolution(self, beta0=None):
                 beta0 = beta_matched(self.plasma_density, self.nom_energy)
 
         ls = np.array([self.length_flattop])
-        ks = np.array([g_ion*SI.e/SI.c/p0])
+        ks = np.array([g*SI.e/SI.c/p0])
         _, _, beta_evolution = evolve_beta_function(ls=ls, ks=ks, beta0=beta0, fast=False, plot=False)
         return phase_advance(beta_evolution[0,:], beta_evolution[1,:])
 

From a6b41f79576120f0011bd96e80ea770a1034110a Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 11 Mar 2026 12:43:16 +0100
Subject: [PATCH 74/84] Added class methods for calculating the external
 focusing gradient and drive beam and main beam trajectories to StageHipace.

---
 src/abel/classes/stage/impl/stage_hipace.py | 544 +++++++++++++++++++-
 1 file changed, 535 insertions(+), 9 deletions(-)

diff --git a/src/abel/classes/stage/impl/stage_hipace.py b/src/abel/classes/stage/impl/stage_hipace.py
index 931fe804..297d5bd7 100644
--- a/src/abel/classes/stage/impl/stage_hipace.py
+++ b/src/abel/classes/stage/impl/stage_hipace.py
@@ -214,8 +214,8 @@ def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, drive
         self.no_plasma = no_plasma
 
         # external focusing (APL-like) [T/m]
+        self.driver_half_oscillations = 1.0
         self.external_focusing = external_focusing
-        self._external_focusing_gradient = None
 
         # plasma profile
         self.plasma_profile = SimpleNamespace()
@@ -268,13 +268,6 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
         self._prepare_ramps()
         self._make_ramp_profile(tmpfolder)
         
-        # set external focusing
-        if self.external_focusing == False:
-            self._external_focusing_gradient = 0
-        if self.external_focusing == True and self._external_focusing_gradient is None:
-            num_half_oscillations = 1
-            self._external_focusing_gradient = self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/self.get_length())**2  # [T/m]
-        
         beam0 = beam_incoming
         driver0 = driver_incoming
         
@@ -350,7 +343,11 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
         # input file
         filename_input = 'input_file'
         path_input = tmpfolder + filename_input
-        hipace_write_inputs(path_input, filename_beam, filename_driver, self.plasma_density, self.num_steps, time_step, box_range_z, box_size_xy, ion_motion=self.ion_motion, ion_species=self.ion_species, beam_ionization=self.beam_ionization, radiation_reaction=self.radiation_reaction, output_period=output_period, num_cell_xy=self.num_cell_xy, num_cell_z=num_cell_z, driver_only=self.driver_only, density_table_file=density_table_file, no_plasma=self.no_plasma, external_focusing_gradient=self._external_focusing_gradient, mesh_refinement=self.mesh_refinement, do_spin_tracking=self.do_spin_tracking)
+        
+        if self.external_focusing and self.external_focusing_gradient is None:
+            self.external_focusing_gradient = self.calc_external_focusing_gradient() # Set the gradient for external focusing fields if not already set.
+
+        hipace_write_inputs(path_input, filename_beam, filename_driver, self.plasma_density, self.num_steps, time_step, box_range_z, box_size_xy, ion_motion=self.ion_motion, ion_species=self.ion_species, beam_ionization=self.beam_ionization, radiation_reaction=self.radiation_reaction, output_period=output_period, num_cell_xy=self.num_cell_xy, num_cell_z=num_cell_z, driver_only=self.driver_only, density_table_file=density_table_file, no_plasma=self.no_plasma, external_focusing_gradient=self.external_focusing_gradient, mesh_refinement=self.mesh_refinement, do_spin_tracking=self.do_spin_tracking)
         
         
         ## RUN SIMULATION
@@ -689,8 +686,537 @@ def get_length(self):
             #ss = density_table[:,0]
             return ss.max()-ss.min()
         return super().get_length()
+    
+
+    # =============================================
+    @property
+    def external_focusing(self) -> bool:
+        "Flag for enabling driver guiding using an external magnetic field."
+        return self._external_focusing
+    @external_focusing.setter
+    def external_focusing(self, enable_external_focusing : bool | None):
+        self._external_focusing = bool(enable_external_focusing)
+    _external_focusing = False
+
+
+    # =============================================
+    @property
+    def external_focusing_gradient(self) -> float:
+        """
+        The focusing gradient g_ext [T/m] for an external azimuthal magnetic 
+        field B = [g_ext*y, -g_ext*x, 0].
+        """
+
+        if self.external_focusing is True and self._external_focusing_gradient is None:
+            # If external focusing is enabled, but the gradient of the external 
+            # focusing field is not yet set, try to calculate and set it:
+
+            # Check whether the ramps have been set up if they exist
+            ramps_not_set_up = (
+                (self.upramp is not None and self.upramp.length is None) or
+                (self.downramp is not None and self.downramp.length is None)
+            )
+
+            # Check if there is enough information to set up the ramps
+            can_set_up_ramps = (self.get_nom_energy_gain(ignore_ramps_if_undefined=True) is not None 
+                                and self.nom_energy is not None)
+
+            # Make a copy of the stage and set up its ramps if they are not set 
+            # up and can be set up
+            if ramps_not_set_up and can_set_up_ramps:
+                stage_copy = copy.deepcopy(self)
+                stage_copy._prepare_ramps()
+            else: 
+                stage_copy = self
+            
+            if stage_copy.get_length() is not None:
+                self._external_focusing_gradient = stage_copy.calc_external_focusing_gradient(num_half_oscillations=self.driver_half_oscillations)
+
+        return self._external_focusing_gradient
+    
+    @external_focusing_gradient.setter
+    def external_focusing_gradient(self, g_ext : float | None):
+        if g_ext is not None and not isinstance(g_ext, float):
+            raise ValueError("External focusing gradient must be a float or None.")
+        if self.external_focusing is False:
+            self._external_focusing_gradient = None
+            raise ValueError("Cannot set external focusing gradient when self.external_focusing is False.")
+        self._external_focusing_gradient = g_ext
+
+    _external_focusing_gradient = None
+
+
+    # =============================================
+    def calc_external_focusing_gradient(self, num_half_oscillations=None, L=None):
+        """
+        Calculate the external focusing gradient g_ext for an azimuthal magnetic 
+        field B = [g_ext*y, -g_ext*x, 0] that gives ``num_half_oscillations`` 
+        half oscillations for the drive beam over the length of the stage.
+
+        Parameters
+        ----------
+        num_half_oscillations : float, optional
+            Number of half betatron oscillations that the drive beam is 
+            intended to perform. If ``None``, will use ``self.driver_half_oscillations``.
+            Defaults to ``None``.
+
+        L : [m] float, optional
+            The length over which the driver will be guided. If ``None``, will 
+            extract the value using ``self.get_length()``. If the stage does 
+            have ramps that have not been fully set up, a deepcopy of the stage 
+            is created to set up its ramps using 
+            :func:`Stage._prepare_ramps() <abel.Stage._prepare_ramps>` so that 
+            the stage total length is defined.
+
+        Returns
+        -------
+        g_ext : [T/m] float
+            The gradient for the azimuthal magnetic field.
+        """
+        if L is None:
+
+            # Make a copy of the stage and set up its ramps if they are not set up
+            ramps_not_set_up = (
+                (self.upramp is not None and self.upramp.length is None) or
+                (self.downramp is not None and self.downramp.length is None)
+            )
+            if ramps_not_set_up:
+                stage_copy = copy.deepcopy(self)
+                stage_copy._prepare_ramps()
+                L = stage_copy.get_length()
+            else: 
+                stage_copy = self
+                L = stage_copy.get_length()
+                if L is None:
+                    L = stage_copy.length_flattop # If there are no ramps, can use either length or length_flattop.
+            
+            if L is None:
+                raise ValueError('Stage length is not set.')
+
+        if num_half_oscillations is None:
+            num_half_oscillations = self.driver_half_oscillations
+
+        return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/L)**2  # [T/m]
 
     
+    # =============================================
+    def driver_guiding_trajectory(self, num_steps=None, dacc_gradient=0.0):
+        """
+        Estimate the trajectory that the drive beam will follow when driver 
+        guiding with an external linear azimuthal magnetic field is applied to a 
+        drive beam with an initial angular offset. The calculations are done by 
+        integrating simplified equations of motion.
+
+        Parameters
+        ----------
+        num_steps : int, optional
+            Number of time steps. If ``None``, will calculate the number of time 
+            steps such that the step size is a small fraction of the matched 
+            beta function of the drive beam. Defaults to ``None``.
+
+        dacc_gradient : [V/m] float, optional
+            The decceleration gradient. Drive beam charge * decceleration 
+            gradient must be negative. Defaults to 0.0.
+        
+
+        Returns
+        -------
+        s_trajectory : [m] 1D float ndarray
+            Longitudinal coordinate of the drive beam trajectory. Reference is 
+            set at the start of the plasma stage.
+
+        x_trajectory : [m] 1D float ndarray
+            x-coordinate of the drive beam trajectory.
+        
+        y_trajectory : [m] 1D float ndarray
+            y-coordinate of the drive beam trajectory.
+        """
+
+        from abel.utilities.relativity import energy2momentum
+        from abel.utilities.statistics import weighted_mean
+
+        driver = self.driver_source.track()
+
+        energy_thres = 10*driver.particle_mass*SI.c**2/SI.e  # [eV], 10 * particle rest energy. Gives beta=0.995.
+        pz_thres = energy2momentum(energy_thres, unit='eV', m=driver.particle_mass)
+        pz0 = energy2momentum(driver.energy(), unit='eV', m=driver.particle_mass)
+
+        if pz0 < pz_thres:
+            raise ValueError('This estimate is only valid for a relativistic beam.')
+        
+        q = driver.particle_charge()  # [C], particle charge including charge sign.
+        if q * dacc_gradient > 0.0:
+            raise ValueError('Drive beam charge * decceleration gradient must be negative.')
+        
+        # Make a copy of the stage and set up its ramps if they are not set up
+        ramps_not_set_up = (
+            (self.upramp is not None and self.upramp.length is None) or
+            (self.downramp is not None and self.downramp.length is None)
+        )
+        if ramps_not_set_up:
+            stage_copy = copy.deepcopy(self)
+            stage_copy._prepare_ramps()
+        else: 
+            stage_copy = self
+        
+        L = stage_copy.get_length()  # [m]
+        
+        if pz0 + q * dacc_gradient * L/SI.c < pz_thres:
+            raise ValueError('The energy depletion will be too severe. This estimate is only valid for a relativistic beam.')
+        
+        # Get the focusing field gradient
+        g = self.external_focusing_gradient  # [T/m]
+        if g is None:
+            g = 0.0
+
+        # Determine the step size
+        if num_steps is None:
+            matched_beta = self.matched_beta_function(driver.energy())
+            num_steps = int(L /(matched_beta/20))
+        
+        ds = L/(num_steps-1)  # [m], step size
+
+        # Initialise arrays
+        s_trajectory = np.full(num_steps, None, dtype=object)
+        x_trajectory = np.full(num_steps, None, dtype=object)
+        y_trajectory = np.full(num_steps, None, dtype=object)
+
+        # Set initial parameters
+        prop_length = 0
+        x0 = driver.x_offset()
+        x = x0
+        y0 = driver.y_offset()
+        y = y0
+        s_trajectory[0] = prop_length
+        x_trajectory[0] = x0
+        y_trajectory[0] = y0
+        px = weighted_mean(driver.pxs(), driver.weightings(), clean=False)
+        py = weighted_mean(driver.pys(), driver.weightings(), clean=False)
+        pz = pz0 # Can add option for deceleration using a gradient
+
+        # Solve the equation of motion
+        i = 0
+        while i < num_steps - 1:
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
+            y = y + py/pz*1/2*ds
+
+            # Kick
+            dpx = q*g*x*ds
+            px = px + dpx
+            dpy = q*g*y*ds
+            py = py + dpy
+            pz = pz0 + q * dacc_gradient * prop_length/SI.c # dacc_gradient > 0
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
+            y = y + py/pz*1/2*ds
+
+            i = i + 1
+            s_trajectory[i] = prop_length
+            x_trajectory[i] = x
+            y_trajectory[i] = y
+
+        return s_trajectory, x_trajectory, y_trajectory
+
+
+    # =============================================
+    def estimate_beam_trajectory(self, beam, num_steps=None):
+        """
+        Estimate the trajectory for the main beam following the trajectory of a 
+        drive beam generated by ``self.driver_source``.
+        
+        Effects such as driver guiding with an external linear azimuthal 
+        magnetic field, background ion focusing and uniform plasma density ramps 
+        are taken into account. The calculations are done by integrating 
+        simplified equations of motion.
+
+        Parameters
+        ----------
+        beam : ``Beam``
+            The main beam to be tracked.
+
+        num_steps : int, optional
+            Number of time steps. If ``None``, will calculate the number of time 
+            steps such that the step size is a small fraction of the matched 
+            beta function of the main beam. Defaults to ``None``.
+
+
+        Returns
+        -------
+        s_trajectory : [m] 1D float ndarray
+            Longitudinal coordinate of the main beam trajectory. Reference is 
+            set at the start of the plasma stage.
+
+        x_trajectory : [m] 1D float ndarray
+            x-coordinate of the main beam trajectory.
+        
+        y_trajectory : [m] 1D float ndarray
+            y-coordinate of the main beam trajectory.
+        """
+
+        from abel.utilities.statistics import weighted_mean
+        from abel.utilities.relativity import energy2momentum
+
+        # Prepare parameters
+        x0 = beam.x_offset()
+        y0 = beam.y_offset()
+        weights = beam.weightings()
+        px0 = weighted_mean(beam.pxs(), weights, clean=False)
+        py0 = weighted_mean(beam.pys(), weights, clean=False)
+        pz0 = energy2momentum(beam.energy(), unit='eV', m=beam.particle_mass)
+        q = beam.particle_charge()
+
+        L = self.get_length()  # [m], total length including any ramps.
+        if num_steps is None:
+            matched_beta = self.matched_beta_function(beam.energy())
+            num_steps = int(L /(matched_beta/20))
+
+        # Actual calculations
+        s_trajectory, x_trajectory, y_trajectory, _, _ = self._estimate_beam_trajectory(s0=0.0, 
+                                                        x0=x0, 
+                                                        y0=y0, 
+                                                        px0=px0, 
+                                                        py0=py0, 
+                                                        pz0=pz0, 
+                                                        q=q, 
+                                                        num_steps=num_steps, 
+                                                        driver_x_trajectory=None, 
+                                                        driver_y_trajectory=None)
+
+        return s_trajectory, x_trajectory, y_trajectory
+
+
+    # =============================================
+    def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, driver_x_trajectory=None, driver_y_trajectory=None):
+        """
+        Helper function for estimating the trajectory for the main beam 
+        following the trajectory of a drive beam defined by 
+        ``driver_x_trajectory`` and ``driver_y_trajectory``.
+        
+        Effects such as driver guiding with an external linear azimuthal 
+        magnetic field, background ion focusing and uniform plasma density ramps 
+        are taken into account. The calculations are done by integrating 
+        simplified equations of motion.
+         
+
+        Parameters
+        ----------
+        s0 : [m] float
+            The initial longitudinal coordinate of the main beam.
+
+        x0 : [m] float
+            The intial x-coordinate of the main beam.
+
+        y0 : [m] float
+            The intial y-coordinate of the main beam.
+
+        px0 : [kg m/s] float
+            The intial x-momentum of the main beam.
+
+        py0 : [kg m/s] float
+            The intial y-momentum of the main beam.
+
+        pz0 : [kg m/s] float
+            The intial z-momentum of the main beam.
+
+        q : [C] flloat
+            The particle charge of the main beam.
+
+        num_steps : int
+            Number of time steps.
+
+        driver_x_trajectory : [m] 1D float ndarray, optional
+            The x-coordinate of trajectory of the drive beam. The length of 
+            ``driver_x_trajectory`` must be the same as ``num_steps``. Is 
+            automatically calculated if ``None``. Defaults to ``None``.
+
+        driver_y_trajectory : [m] 1D float ndarray, optional
+            The y-coordinate of trajectory of the drive beam. The length of 
+            ``driver_y_trajectory`` must be the same as ``num_steps``. Is 
+            automatically calculated if ``None``. Defaults to ``None``.
+        
+
+        Returns
+        -------
+        s_trajectory : [m] 1D float ndarray
+            Longitudinal coordinate of the main beam trajectory. Reference is 
+            set at the start of the plasma stage.
+
+        x_trajectory : [m] 1D float ndarray
+            x-coordinate of the main beam trajectory.
+        
+        y_trajectory : [m] 1D float ndarray
+            y-coordinate of the main beam trajectory.
+
+        px_trajectory : [kg m/s] 1D float ndarray
+            Mean x-component of the beam momentum along the trajectory.
+
+        py_trajectory : [kg m/s] 1D float ndarray
+            Mean y-component of the beam momentum along the trajectory.
+        """
+        from abel.utilities.other import find_closest_value_in_arr
+
+        if q * self.nom_accel_gradient_flattop > 0.0:
+            raise ValueError('Beam charge * self.nom_accel_gradient_flattop  gradient must be negative.')
+        
+        # Make a copy of the stage and set up its ramps if they are not set up
+        ramps_not_set_up = (
+            (self.upramp is not None and self.upramp.length is None) or
+            (self.downramp is not None and self.downramp.length is None)
+        )
+        if ramps_not_set_up:
+            stage_copy = copy.deepcopy(self)
+            stage_copy._prepare_ramps()
+        else: 
+            stage_copy = self
+        
+        # Calculate the focusing field gradient
+        g0 = SI.e*stage_copy.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m]
+        g = g0
+        if stage_copy.external_focusing_gradient is not None:
+            g = g0 + stage_copy.external_focusing_gradient
+
+        # Only calculate the time step size when calling estimate_beam_trajectory() from a flattop
+        if not stage_copy.is_upramp() and not stage_copy.is_downramp():
+
+            # Set the step size
+            L = stage_copy.get_length()  # [m], total length including any ramps.
+            stage_copy.ds = L / (num_steps-1) # [m], step size
+
+            if stage_copy.has_ramp():
+                # Set the number of time steps in the upramp and downramp
+                ss_helper = np.arange(num_steps) * stage_copy.ds
+                idx_upramp_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length)
+                num_steps_upramp = idx_upramp_end + 1  # Number of time steps for the upramp
+                idx_flat_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length+stage_copy.length_flattop)  # Index marking the end of the flattop.
+                num_steps_downramp = num_steps - idx_flat_end  # Number of time steps for the downramp
+                stage_copy.idx_flat_end = idx_flat_end
+                stage_copy.num_steps_upramp = num_steps_upramp
+                stage_copy.num_steps_downramp = num_steps_downramp
+
+        ds = stage_copy.ds
+
+        # Calculate the drive beam trajectory
+        if driver_x_trajectory is None or driver_y_trajectory is None:
+            _, driver_x_trajectory, driver_y_trajectory = self.driver_guiding_trajectory(num_steps=num_steps, dacc_gradient=0.0)
+        
+        if not stage_copy.is_upramp() and not stage_copy.is_downramp():
+            if len(driver_x_trajectory) != num_steps or len(driver_y_trajectory) != num_steps:
+                raise ValueError('The length of driver_x_trajectory and driver_y_trajectory must be the same as num_steps.')
+
+        # Initialise arrays
+        s_trajectory = np.full(num_steps, None, dtype=object)
+        x_trajectory = np.full(num_steps, None, dtype=object)
+        y_trajectory = np.full(num_steps, None, dtype=object)
+        px_trajectory = np.full(num_steps, None, dtype=object)
+        py_trajectory = np.full(num_steps, None, dtype=object)
+
+        # Recursive call from the upramp
+        if stage_copy.upramp is not None:
+            upramp = stage_copy.convert_PlasmaRamp(stage_copy.upramp)
+            stage_copy.upramp = upramp
+            if self.external_focusing:
+                upramp.external_focusing_gradient = stage_copy.external_focusing_gradient
+
+            num_steps_upramp = stage_copy.num_steps_upramp 
+
+            s_trajectory_upramp, x_trajectory_upramp, y_trajectory_upramp, px_trajectory_upramp, py_trajectory_upramp = upramp._estimate_beam_trajectory(s0, x0, y0, px0, py0, pz0, q, num_steps_upramp, driver_x_trajectory, driver_y_trajectory)
+
+            # Initial parameters for the flattop
+            prop_length = s_trajectory_upramp[-1]
+            s_trajectory[:len(s_trajectory_upramp)] = s_trajectory_upramp
+            x0 = x_trajectory_upramp[-1]
+            x_trajectory[:len(x_trajectory_upramp)] = x_trajectory_upramp
+            y0 = y_trajectory_upramp[-1]
+            y_trajectory[:len(y_trajectory_upramp)] = y_trajectory_upramp
+            px0 = px_trajectory_upramp[-1]
+            px_trajectory[:len(px_trajectory_upramp)] = px_trajectory_upramp
+            py0 = py_trajectory_upramp[-1]
+            py_trajectory[:len(py_trajectory_upramp)] = py_trajectory_upramp
+            pz0 = pz0 - q * stage_copy.upramp.nom_accel_gradient_flattop * prop_length/SI.c
+
+            i = num_steps_upramp - 1
+
+            if stage_copy.downramp is not None:
+                i_end = stage_copy.idx_flat_end
+            else:
+                i_end = num_steps - 1
+
+        # No ramps
+        else:
+            prop_length = s0
+            s_trajectory[0] = prop_length
+            x_trajectory[0] = x0
+            y_trajectory[0] = y0
+            px_trajectory[0] = px0
+            py_trajectory[0] = py0
+
+            i = 0
+            i_end = num_steps - 1
+
+            if self.is_downramp():
+                # Extract the part of drive beam trajectory for the downramp
+                driver_x_trajectory = driver_x_trajectory[stage_copy.idx_flat_end:]
+                driver_y_trajectory = driver_y_trajectory[stage_copy.idx_flat_end:]
+
+        # Set the initial conditions
+        x = x0
+        y = y0
+        px = px0
+        py = py0
+        pz = pz0
+
+        # Solve the equations of motion
+        while i < i_end:
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
+            y = y + py/pz*1/2*ds
+
+            # Kick
+            dpx = q*g*x*ds - q*g0*driver_x_trajectory[i]*ds
+            px = px + dpx
+            dpy = q*g*y*ds - q*g0*driver_y_trajectory[i]*ds
+            py = py + dpy
+            pz = pz0 - q * self.nom_accel_gradient_flattop * prop_length/SI.c
+
+            # Drift
+            prop_length = prop_length + 1/2*ds
+            x = x + px/pz*1/2*ds
+            y = y + py/pz*1/2*ds
+
+            i = i + 1
+            s_trajectory[i] = prop_length
+            x_trajectory[i] = x
+            y_trajectory[i] = y
+            px_trajectory[i] = px
+            py_trajectory[i] = py
+
+        # Recursive call from the downramp
+        if stage_copy.downramp is not None:
+            downramp = stage_copy.convert_PlasmaRamp(stage_copy.downramp)
+            stage_copy.downramp = downramp
+            if self.external_focusing:
+                downramp.external_focusing_gradient = stage_copy.external_focusing_gradient
+
+            num_steps_downramp = stage_copy.num_steps_downramp
+
+            s_trajectory_downramp, x_trajectory_downramp, y_trajectory_downramp, px_trajectory_downramp, py_trajectory_downramp = downramp._estimate_beam_trajectory(prop_length, x, y, px, py, pz, q, num_steps_downramp, driver_x_trajectory, driver_y_trajectory)
+
+            s_trajectory[-len(s_trajectory_downramp):] = s_trajectory_downramp
+            x_trajectory[-len(x_trajectory_downramp):] = x_trajectory_downramp
+            y_trajectory[-len(y_trajectory_downramp):] = y_trajectory_downramp
+            px_trajectory[-len(px_trajectory_downramp):] = px_trajectory_downramp
+            py_trajectory[-len(py_trajectory_downramp):] = py_trajectory_downramp
+
+        return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory
+
+
     # ==================================================
     # Apply waterfall function to all beam dump files
     def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsigma=None, args=None):

From 951bc38bbda0c1d367dd465cb2b830d302135f0d Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 11 Mar 2026 12:44:23 +0100
Subject: [PATCH 75/84] Added getter for external focusing gradient to Stage.

---
 src/abel/classes/stage/stage.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/abel/classes/stage/stage.py b/src/abel/classes/stage/stage.py
index 45b61dde..fce84f43 100644
--- a/src/abel/classes/stage/stage.py
+++ b/src/abel/classes/stage/stage.py
@@ -1387,6 +1387,16 @@ def matched_beta_function_flattop(self, energy):
         return beta_matched(self.plasma_density, energy)
     
 
+    # =============================================
+    @property
+    def external_focusing_gradient(self) -> float:
+        """
+        Return None by default for ``Stage`` subclasses not supporting external 
+        focusing fields.
+        """
+        return None
+    
+    
     # ==================================================
     def energy_usage(self):
         return self.driver_source.energy_usage()

From fe31038673503a2f2bc61a51760cc9d9edcbdd72 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 11 Mar 2026 12:45:33 +0100
Subject: [PATCH 76/84] Edited src/abel/wrappers/hipace/hipace_wrapper.py for
 setting the external focusing gradient.

---
 src/abel/wrappers/hipace/hipace_wrapper.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/abel/wrappers/hipace/hipace_wrapper.py b/src/abel/wrappers/hipace/hipace_wrapper.py
index 7a0b9177..90ce9487 100644
--- a/src/abel/wrappers/hipace/hipace_wrapper.py
+++ b/src/abel/wrappers/hipace/hipace_wrapper.py
@@ -156,7 +156,9 @@ def hipace_write_inputs(filename_input, filename_beam, filename_driver, plasma_d
         print('>> HiPACE++: Changing from', num_cell_xy, 'to', new_num_cell_xy, ' (i.e., 2^n-1) for better performance.')
         num_cell_xy = new_num_cell_xy
 
-    # plasma-density profile from file
+    # gradient for external magnetic field
+    if external_focusing_gradient is None:
+        external_focusing_gradient = '#'
     if abs(external_focusing_gradient) > 0:
         external_focusing_comment = ''
     else:

From 0d8ee16da4b3a0ad0acae771b4d6a515dc7d338d Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 11 Mar 2026 12:57:16 +0100
Subject: [PATCH 77/84] Removed extra entry of Stage.external_focusing_gradient
 after mergeing with StageHipace_beam_trajectory_estimate.

---
 src/abel/classes/stage/stage.py | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/src/abel/classes/stage/stage.py b/src/abel/classes/stage/stage.py
index cadc9c6e..55df1e99 100644
--- a/src/abel/classes/stage/stage.py
+++ b/src/abel/classes/stage/stage.py
@@ -1471,16 +1471,6 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_den
 
         return length
     
-
-    # =============================================
-    @property
-    def external_focusing_gradient(self) -> float:
-        """
-        Return None by default for ``Stage`` subclasses not supporting external 
-        focusing fields.
-        """
-        return None
-    
     
     # ==================================================
     def calc_flattop_num_beta_osc(self, num_beta_osc):

From 2639c059e601d24601652d1266e9a0c228ea2079 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 11 Mar 2026 16:09:07 +0100
Subject: [PATCH 78/84] Correction in StageHipace._estimate_beam_trajectory()
 and StageHipace.driver_guiding_trajectory() for array initialisation.

---
 src/abel/classes/stage/impl/stage_hipace.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/abel/classes/stage/impl/stage_hipace.py b/src/abel/classes/stage/impl/stage_hipace.py
index 297d5bd7..3fc0c494 100644
--- a/src/abel/classes/stage/impl/stage_hipace.py
+++ b/src/abel/classes/stage/impl/stage_hipace.py
@@ -877,9 +877,9 @@ def driver_guiding_trajectory(self, num_steps=None, dacc_gradient=0.0):
         ds = L/(num_steps-1)  # [m], step size
 
         # Initialise arrays
-        s_trajectory = np.full(num_steps, None, dtype=object)
-        x_trajectory = np.full(num_steps, None, dtype=object)
-        y_trajectory = np.full(num_steps, None, dtype=object)
+        s_trajectory = np.full(num_steps, None, dtype=float)
+        x_trajectory = np.full(num_steps, None, dtype=float)
+        y_trajectory = np.full(num_steps, None, dtype=float)
 
         # Set initial parameters
         prop_length = 0
@@ -1109,11 +1109,11 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
                 raise ValueError('The length of driver_x_trajectory and driver_y_trajectory must be the same as num_steps.')
 
         # Initialise arrays
-        s_trajectory = np.full(num_steps, None, dtype=object)
-        x_trajectory = np.full(num_steps, None, dtype=object)
-        y_trajectory = np.full(num_steps, None, dtype=object)
-        px_trajectory = np.full(num_steps, None, dtype=object)
-        py_trajectory = np.full(num_steps, None, dtype=object)
+        s_trajectory = np.full(num_steps, None, dtype=float)
+        x_trajectory = np.full(num_steps, None, dtype=float)
+        y_trajectory = np.full(num_steps, None, dtype=float)
+        px_trajectory = np.full(num_steps, None, dtype=float)
+        py_trajectory = np.full(num_steps, None, dtype=float)
 
         # Recursive call from the upramp
         if stage_copy.upramp is not None:

From 9b50bc75befede66b3e661c30dff591d0c7f494b Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Tue, 17 Mar 2026 17:38:12 +0100
Subject: [PATCH 79/84] Made a correction in Stage._calc_ramp_length() that
 takes into account any external focusing gradient.

---
 src/abel/classes/stage/stage.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/abel/classes/stage/stage.py b/src/abel/classes/stage/stage.py
index 55df1e99..5e78c0cb 100644
--- a/src/abel/classes/stage/stage.py
+++ b/src/abel/classes/stage/stage.py
@@ -515,8 +515,14 @@ def _calc_ramp_length(self, ramp : Self) -> float:
             ramp_beta_mag = self.ramp_beta_mag
         else:
             raise ValueError('No ramp_beta_mag defined.')
-        
-        ramp_length = beta_matched(self.plasma_density, ramp.nom_energy)*np.pi/(2*np.sqrt(1/ramp_beta_mag))
+
+        g = SI.e*self.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
+        if self.external_focusing_gradient is not None:  # Add contribution from external field
+            g = g + self.external_focusing_gradient
+
+        k_beta = np.sqrt(g*SI.c/ramp.nom_energy)  # [m^-1], betatron wavenumber.
+        ramp_length = 1/k_beta * np.pi/2 * np.sqrt(ramp_beta_mag)  # k_beta*ramp_length = pi/2 gives pi/2 phase advance.
+
         if ramp_length < 0.0:
             raise ValueError(f"ramp_length = {ramp_length} [m] < 0.0")
         return ramp_length
@@ -525,8 +531,8 @@ def _calc_ramp_length(self, ramp : Self) -> float:
      # ==================================================
     def get_ramp_length(self) -> float:
         """
-        Get the length of the ramps if the stage has ramps. Returns 0.0 
-        otherwise.
+        Get the length of the ramps if the stage has ramps already set up. 
+        Returns 0.0 otherwise.
         """
 
         if self.has_ramp():

From a05098d9e236e23b00e003d713b018024963dab0 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Mar 2026 23:15:59 +0100
Subject: [PATCH 80/84] Correction for Stage._calc_ramp_length(), as the
 previous update involving self.external_focusing_gradient may cause an
 infinite loop.

---
 src/abel/classes/stage/stage.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/abel/classes/stage/stage.py b/src/abel/classes/stage/stage.py
index 5e78c0cb..f1237d51 100644
--- a/src/abel/classes/stage/stage.py
+++ b/src/abel/classes/stage/stage.py
@@ -517,8 +517,8 @@ def _calc_ramp_length(self, ramp : Self) -> float:
             raise ValueError('No ramp_beta_mag defined.')
 
         g = SI.e*self.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
-        if self.external_focusing_gradient is not None:  # Add contribution from external field
-            g = g + self.external_focusing_gradient
+        #if self.external_focusing_gradient is not None:  # Add contribution from external field
+        #    g = g + self.external_focusing_gradient  # external_focusing_gradient may itself depend on the total length, so this may cause an infinite loop.
 
         k_beta = np.sqrt(g*SI.c/ramp.nom_energy)  # [m^-1], betatron wavenumber.
         ramp_length = 1/k_beta * np.pi/2 * np.sqrt(ramp_beta_mag)  # k_beta*ramp_length = pi/2 gives pi/2 phase advance.

From f8ebfd7b69249230da6b17e33a5bd56c943fb42c Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Mar 2026 23:33:59 +0100
Subject: [PATCH 81/84] Minor correction in
 StageHIpace.calc_length_num_beta_osc() and
 StageHipace.match_length_2_num_beta_osc().

---
 src/abel/classes/stage/impl/stage_hipace.py | 128 +++++++++++++++++++-
 1 file changed, 125 insertions(+), 3 deletions(-)

diff --git a/src/abel/classes/stage/impl/stage_hipace.py b/src/abel/classes/stage/impl/stage_hipace.py
index c344bd63..3a350d70 100644
--- a/src/abel/classes/stage/impl/stage_hipace.py
+++ b/src/abel/classes/stage/impl/stage_hipace.py
@@ -930,13 +930,12 @@ def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_den
         # The function to be used for solving the equation for phase advance numerically
         def rhs(L):
             g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
-            L_ramps = 0.0
 
             # Set up the ramps using the stage copy
             if ramps_not_set_up:
                 stage_copy.length_flattop = L[0]  # L is an ndarray with one element
                 stage_copy._prepare_ramps()
-                L_ramps = stage_copy.get_ramp_length()
+            L_ramps = stage_copy.get_ramp_length()
 
             if self.external_focusing:  # Add contribution from external field used for driver guiding
                 g_ext = stage_copy.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=L+L_ramps)
@@ -953,6 +952,129 @@ def rhs(L):
         return length
     
 
+    # ==================================================
+    # def calc_length_num_beta_osc(self, num_beta_osc, beam, driver=None, driver_half_oscillations=None):
+    #     """
+    #     Calculate the stage length that gives ``num_beta_osc`` betatron 
+    #     oscillations for a particle with given initial energy ``initial_energy`` 
+    #     in a uniform plasma stage (excluding ramps) with defined nominal 
+    #     acceleration gradient and plasma density ``plasma_density``.
+
+    #     Will take into account the contribution from an external linear magnetic 
+    #     field B=[g_ext*y, -g_ext*x, 0] if :attr:`self.external_focusing <abel.StageHipace.external_focusing>` 
+    #     is set to ``True``.
+
+    #     Parameters
+    #     ----------
+    #     num_beta_osc : float
+    #         Total number of design betatron oscillations that the electron 
+    #         should perform through the plasma stage excluding ramps. 
+
+    #     initial_energy : [eV] float, optional
+    #         The initial energy of the particle at the start of the plasma stage. 
+    #         Defaults to ``self.nom_energy``.
+
+    #     plasma_density : [m^-3] float, optional
+    #         The plasma density of the plasma stage. Defaults to 
+    #         ``self.plasma_density``.
+
+    #     driver_half_oscillations : float, optional
+    #         Number of half betatron oscillations that the drive beam is 
+    #         intended to perform. If ``None``, will use :attr:`StageHipace.driver_half_oscillations <abel.StageHipace.driver_half_oscillations>`.
+    #         Defaults to ``None``.
+
+    #     q : [C] float, optional
+    #         Particle charge. q * nom_accel_gradient must be positive. Defaults 
+    #         to elementary charge.
+
+            
+    #     Returns
+    #     -------
+    #     length : [m] float
+    #         Length of the plasma stage excluding ramps matched to the given 
+    #         number of betatron oscillations.
+    #     """
+
+    #     from scipy.optimize import fsolve
+
+    #     if num_beta_osc < 0:
+    #         raise ValueError('Number of input betatron oscillations must be positive.')
+        
+    #     if driver_half_oscillations is None:
+    #         driver_half_oscillations = self.driver_half_oscillations
+    #     if driver_half_oscillations < 0:
+    #         raise ValueError('Number of driver oscillations must be positive.')
+        
+    #     if driver is None:
+    #         driver = self.driver_source.track()
+
+    #     # Assess whether the ramps have been set up
+    #     ramps_not_set_up = (
+    #         (self.upramp is not None and self.upramp.length is None) or
+    #         (self.downramp is not None and self.downramp.length is None)
+    #     )
+
+    #     # Make a copy of the stage
+    #     stage_copy = copy.deepcopy(self)
+
+    #     #def pha
+
+    #     # The function for calculating the phase advance from beam trajectories
+    #     def phase_advance_from_trajectories(L):
+
+    #         L = L[0]  # L is an ndarray with one element
+    #         #g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
+    #         L_ramps = 0.0
+
+    #         # Set up the ramps using the stage copy
+    #         if ramps_not_set_up:
+    #             stage_copy.length = L
+    #             stage_copy._prepare_ramps()
+    #             #L_ramps = stage_copy.get_ramp_length()
+
+    #         if self.external_focusing:  # Add contribution from external field used for driver guiding
+    #             g_ext = stage_copy.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=L)
+    #             print(g_ext)
+    #             stage_copy.external_focusing_gradient = g_ext
+
+    #             #g = g + g_ext
+
+    #         s_trajectory, x_trajectory, y_trajectory, driver_x_trajectory, driver_y_trajectory = stage_copy.estimate_beam_trajectory(beam, num_steps=None)
+
+    #         r_trajectory = np.sqrt(x_trajectory**2 + y_trajectory**2)
+    #         driver_r_trajectory = np.sqrt(driver_x_trajectory**2 + driver_y_trajectory**2)
+    #         straight_r_trajectory = r_trajectory - driver_r_trajectory
+
+    #         xp_trajectory = np.diff(x_trajectory)/np.diff(s_trajectory)
+    #         xp_trajectory = np.insert(xp_trajectory, 0, beam.x_angle())
+    #         yp_trajectory = np.diff(y_trajectory)/np.diff(s_trajectory)
+    #         yp_trajectory = np.insert(yp_trajectory, 0, beam.y_angle())
+    #         angle_trajectory = np.sqrt(xp_trajectory**2 + yp_trajectory**2)
+
+    #         driver_xp_trajectory = np.diff(driver_x_trajectory)/np.diff(s_trajectory)
+    #         driver_xp_trajectory = np.insert(driver_xp_trajectory, 0, driver.x_angle())
+    #         driver_yp_trajectory = np.diff(driver_y_trajectory)/np.diff(s_trajectory)
+    #         driver_yp_trajectory = np.insert(driver_yp_trajectory, 0, driver.y_angle())
+    #         driver_angle_trajectory = np.sqrt(driver_xp_trajectory**2 + driver_yp_trajectory**2)
+
+    #         straight_angle_trajectory = angle_trajectory - driver_angle_trajectory
+
+    #         # Some control that the phase space trajectory loops aroung the origin ...
+
+    #         phase = np.unwrap(np.arctan2(straight_angle_trajectory/np.max(np.abs(straight_angle_trajectory)), 
+    #                                     straight_r_trajectory/np.max(np.abs(straight_r_trajectory))))
+            
+    #         mu = np.abs((phase[-1] - phase[0]))
+
+    #         return mu
+        
+    #     # Find the length that gives num_beta_osc
+    #     solution = fsolve(lambda L: phase_advance_from_trajectories(L)/ (2*np.pi) - num_beta_osc, x0=1)
+    #     length = solution[0]
+
+    #     return length
+
+
     # ==================================================
     def match_length_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=None, set_consistent_params=True, q=SI.e):
         """
@@ -1013,7 +1135,7 @@ def match_length_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=Non
             if self.upramp.ramp_shape != 'uniform' or self.downramp.ramp_shape != 'uniform':
                 raise ValueError('This method assumes uniform ramps.')
             if self.upramp.length_flattop is not None or self.downramp.length_flattop is not None:
-                raise ValueError('This method assumes uniform ramps with length set to give pi/2 phase advance for the main beam.')
+               raise ValueError('This method assumes uniform ramps with length set to give pi/2 phase advance for the main beam. The lengths for the ramps are already set. Setting a new length for the flattop will give wrong results.')
             num_beta_osc_flattop = num_beta_osc - 0.5  # The ramps are by default set up to give pi/2 phase advance for the main beam.
         else:
             num_beta_osc_flattop = num_beta_osc

From d8b297d00568eddabdded1b704c94632e7dd2638 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Wed, 18 Mar 2026 23:38:22 +0100
Subject: [PATCH 82/84] Major corrections and expansions in
 StageHipace._estimate_beam_trajectory() and
 StageHipace.estimate_beam_trajectory(). Also removed a commented out version
 of StageHipace.calc_length_num_beta_osc().

---
 src/abel/classes/stage/impl/stage_hipace.py | 184 +++++---------------
 1 file changed, 45 insertions(+), 139 deletions(-)

diff --git a/src/abel/classes/stage/impl/stage_hipace.py b/src/abel/classes/stage/impl/stage_hipace.py
index 3a350d70..0e24b5f4 100644
--- a/src/abel/classes/stage/impl/stage_hipace.py
+++ b/src/abel/classes/stage/impl/stage_hipace.py
@@ -952,129 +952,6 @@ def rhs(L):
         return length
     
 
-    # ==================================================
-    # def calc_length_num_beta_osc(self, num_beta_osc, beam, driver=None, driver_half_oscillations=None):
-    #     """
-    #     Calculate the stage length that gives ``num_beta_osc`` betatron 
-    #     oscillations for a particle with given initial energy ``initial_energy`` 
-    #     in a uniform plasma stage (excluding ramps) with defined nominal 
-    #     acceleration gradient and plasma density ``plasma_density``.
-
-    #     Will take into account the contribution from an external linear magnetic 
-    #     field B=[g_ext*y, -g_ext*x, 0] if :attr:`self.external_focusing <abel.StageHipace.external_focusing>` 
-    #     is set to ``True``.
-
-    #     Parameters
-    #     ----------
-    #     num_beta_osc : float
-    #         Total number of design betatron oscillations that the electron 
-    #         should perform through the plasma stage excluding ramps. 
-
-    #     initial_energy : [eV] float, optional
-    #         The initial energy of the particle at the start of the plasma stage. 
-    #         Defaults to ``self.nom_energy``.
-
-    #     plasma_density : [m^-3] float, optional
-    #         The plasma density of the plasma stage. Defaults to 
-    #         ``self.plasma_density``.
-
-    #     driver_half_oscillations : float, optional
-    #         Number of half betatron oscillations that the drive beam is 
-    #         intended to perform. If ``None``, will use :attr:`StageHipace.driver_half_oscillations <abel.StageHipace.driver_half_oscillations>`.
-    #         Defaults to ``None``.
-
-    #     q : [C] float, optional
-    #         Particle charge. q * nom_accel_gradient must be positive. Defaults 
-    #         to elementary charge.
-
-            
-    #     Returns
-    #     -------
-    #     length : [m] float
-    #         Length of the plasma stage excluding ramps matched to the given 
-    #         number of betatron oscillations.
-    #     """
-
-    #     from scipy.optimize import fsolve
-
-    #     if num_beta_osc < 0:
-    #         raise ValueError('Number of input betatron oscillations must be positive.')
-        
-    #     if driver_half_oscillations is None:
-    #         driver_half_oscillations = self.driver_half_oscillations
-    #     if driver_half_oscillations < 0:
-    #         raise ValueError('Number of driver oscillations must be positive.')
-        
-    #     if driver is None:
-    #         driver = self.driver_source.track()
-
-    #     # Assess whether the ramps have been set up
-    #     ramps_not_set_up = (
-    #         (self.upramp is not None and self.upramp.length is None) or
-    #         (self.downramp is not None and self.downramp.length is None)
-    #     )
-
-    #     # Make a copy of the stage
-    #     stage_copy = copy.deepcopy(self)
-
-    #     #def pha
-
-    #     # The function for calculating the phase advance from beam trajectories
-    #     def phase_advance_from_trajectories(L):
-
-    #         L = L[0]  # L is an ndarray with one element
-    #         #g = SI.e*plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
-    #         L_ramps = 0.0
-
-    #         # Set up the ramps using the stage copy
-    #         if ramps_not_set_up:
-    #             stage_copy.length = L
-    #             stage_copy._prepare_ramps()
-    #             #L_ramps = stage_copy.get_ramp_length()
-
-    #         if self.external_focusing:  # Add contribution from external field used for driver guiding
-    #             g_ext = stage_copy.calc_external_focusing_gradient(num_half_oscillations=driver_half_oscillations, L=L)
-    #             print(g_ext)
-    #             stage_copy.external_focusing_gradient = g_ext
-
-    #             #g = g + g_ext
-
-    #         s_trajectory, x_trajectory, y_trajectory, driver_x_trajectory, driver_y_trajectory = stage_copy.estimate_beam_trajectory(beam, num_steps=None)
-
-    #         r_trajectory = np.sqrt(x_trajectory**2 + y_trajectory**2)
-    #         driver_r_trajectory = np.sqrt(driver_x_trajectory**2 + driver_y_trajectory**2)
-    #         straight_r_trajectory = r_trajectory - driver_r_trajectory
-
-    #         xp_trajectory = np.diff(x_trajectory)/np.diff(s_trajectory)
-    #         xp_trajectory = np.insert(xp_trajectory, 0, beam.x_angle())
-    #         yp_trajectory = np.diff(y_trajectory)/np.diff(s_trajectory)
-    #         yp_trajectory = np.insert(yp_trajectory, 0, beam.y_angle())
-    #         angle_trajectory = np.sqrt(xp_trajectory**2 + yp_trajectory**2)
-
-    #         driver_xp_trajectory = np.diff(driver_x_trajectory)/np.diff(s_trajectory)
-    #         driver_xp_trajectory = np.insert(driver_xp_trajectory, 0, driver.x_angle())
-    #         driver_yp_trajectory = np.diff(driver_y_trajectory)/np.diff(s_trajectory)
-    #         driver_yp_trajectory = np.insert(driver_yp_trajectory, 0, driver.y_angle())
-    #         driver_angle_trajectory = np.sqrt(driver_xp_trajectory**2 + driver_yp_trajectory**2)
-
-    #         straight_angle_trajectory = angle_trajectory - driver_angle_trajectory
-
-    #         # Some control that the phase space trajectory loops aroung the origin ...
-
-    #         phase = np.unwrap(np.arctan2(straight_angle_trajectory/np.max(np.abs(straight_angle_trajectory)), 
-    #                                     straight_r_trajectory/np.max(np.abs(straight_r_trajectory))))
-            
-    #         mu = np.abs((phase[-1] - phase[0]))
-
-    #         return mu
-        
-    #     # Find the length that gives num_beta_osc
-    #     solution = fsolve(lambda L: phase_advance_from_trajectories(L)/ (2*np.pi) - num_beta_osc, x0=1)
-    #     length = solution[0]
-
-    #     return length
-
-
     # ==================================================
     def match_length_2_num_beta_osc(self, num_beta_osc, driver_half_oscillations=None, set_consistent_params=True, q=SI.e):
         """
@@ -1332,6 +1209,21 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
         
         y_trajectory : [m] 1D float ndarray
             y-coordinate of the main beam trajectory.
+
+        px_trajectory : [kg m/s] 1D float ndarray
+            Mean x-component of the beam momentum along the trajectory.
+
+        py_trajectory : [kg m/s] 1D float ndarray
+            Mean y-component of the beam momentum along the trajectory.
+
+        pz_trajectory : [kg m/s] 1D float ndarray
+            Mean longitudinal of the beam momentum along the trajectory.
+
+        driver_x_trajectory : [m] 1D float ndarray, optional
+            The x-coordinate of trajectory of the drive beam.
+
+        driver_y_trajectory : [m] 1D float ndarray, optional
+            The y-coordinate of trajectory of the drive beam.
         """
 
         from abel.utilities.statistics import weighted_mean
@@ -1352,7 +1244,7 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
             num_steps = int(L /(matched_beta/20))
 
         # Actual calculations
-        s_trajectory, x_trajectory, y_trajectory, _, _ = self._estimate_beam_trajectory(s0=0.0, 
+        s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory, pz_trajectory, driver_x_trajectory, driver_y_trajectory = self._estimate_beam_trajectory(s0=0.0, 
                                                         x0=x0, 
                                                         y0=y0, 
                                                         px0=px0, 
@@ -1363,7 +1255,7 @@ def estimate_beam_trajectory(self, beam, num_steps=None):
                                                         driver_x_trajectory=None, 
                                                         driver_y_trajectory=None)
 
-        return s_trajectory, x_trajectory, y_trajectory
+        return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory, pz_trajectory, driver_x_trajectory, driver_y_trajectory 
 
 
     # =============================================
@@ -1433,6 +1325,15 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
 
         py_trajectory : [kg m/s] 1D float ndarray
             Mean y-component of the beam momentum along the trajectory.
+
+        pz_trajectory : [kg m/s] 1D float ndarray
+            Mean longitudinal of the beam momentum along the trajectory.
+
+        driver_x_trajectory : [m] 1D float ndarray, optional
+            The x-coordinate of trajectory of the drive beam.
+
+        driver_y_trajectory : [m] 1D float ndarray, optional
+            The y-coordinate of trajectory of the drive beam.
         """
         from abel.utilities.other import find_closest_value_in_arr
 
@@ -1469,7 +1370,7 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
                 idx_upramp_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length)
                 num_steps_upramp = idx_upramp_end + 1  # Number of time steps for the upramp
                 idx_flat_end, _ = find_closest_value_in_arr(ss_helper, stage_copy.upramp.length+stage_copy.length_flattop)  # Index marking the end of the flattop.
-                num_steps_downramp = num_steps - idx_flat_end  # Number of time steps for the downramp
+                num_steps_downramp = num_steps - idx_flat_end - 1  # Number of time steps for the downramp
                 stage_copy.idx_flat_end = idx_flat_end
                 stage_copy.num_steps_upramp = num_steps_upramp
                 stage_copy.num_steps_downramp = num_steps_downramp
@@ -1490,6 +1391,7 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
         y_trajectory = np.full(num_steps, None, dtype=float)
         px_trajectory = np.full(num_steps, None, dtype=float)
         py_trajectory = np.full(num_steps, None, dtype=float)
+        pz_trajectory = np.full(num_steps, None, dtype=float)
 
         # Recursive call from the upramp
         if stage_copy.upramp is not None:
@@ -1497,10 +1399,9 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
             stage_copy.upramp = upramp
             if self.external_focusing:
                 upramp.external_focusing_gradient = stage_copy.external_focusing_gradient
-
             num_steps_upramp = stage_copy.num_steps_upramp 
 
-            s_trajectory_upramp, x_trajectory_upramp, y_trajectory_upramp, px_trajectory_upramp, py_trajectory_upramp = upramp._estimate_beam_trajectory(s0, x0, y0, px0, py0, pz0, q, num_steps_upramp, driver_x_trajectory, driver_y_trajectory)
+            s_trajectory_upramp, x_trajectory_upramp, y_trajectory_upramp, px_trajectory_upramp, py_trajectory_upramp, pz_trajectory_upramp, _, _ = upramp._estimate_beam_trajectory(s0, x0, y0, px0, py0, pz0, q, num_steps_upramp, driver_x_trajectory, driver_y_trajectory)
 
             # Initial parameters for the flattop
             prop_length = s_trajectory_upramp[-1]
@@ -1513,8 +1414,9 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
             px_trajectory[:len(px_trajectory_upramp)] = px_trajectory_upramp
             py0 = py_trajectory_upramp[-1]
             py_trajectory[:len(py_trajectory_upramp)] = py_trajectory_upramp
-            pz0 = pz0 - q * stage_copy.upramp.nom_accel_gradient_flattop * prop_length/SI.c
-
+            pz0 = pz_trajectory_upramp[-1]
+            pz_trajectory[:len(pz_trajectory_upramp)] = pz_trajectory_upramp
+            
             i = num_steps_upramp - 1
 
             if stage_copy.downramp is not None:
@@ -1530,12 +1432,13 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
             y_trajectory[0] = y0
             px_trajectory[0] = px0
             py_trajectory[0] = py0
+            pz_trajectory[0] = pz0
 
             i = 0
             i_end = num_steps - 1
 
             if self.is_downramp():
-                # Extract the part of drive beam trajectory for the downramp
+                # Extract the part of drive beam trajectory for the downramp. Note one element longer than num_steps_downramp.
                 driver_x_trajectory = driver_x_trajectory[stage_copy.idx_flat_end:]
                 driver_y_trajectory = driver_y_trajectory[stage_copy.idx_flat_end:]
 
@@ -1559,7 +1462,8 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
             px = px + dpx
             dpy = q*g*y*ds - q*g0*driver_y_trajectory[i]*ds
             py = py + dpy
-            pz = pz0 - q * self.nom_accel_gradient_flattop * prop_length/SI.c
+            dpz = - q * self.nom_accel_gradient_flattop * ds/SI.c
+            pz = pz + dpz
 
             # Drift
             prop_length = prop_length + 1/2*ds
@@ -1572,6 +1476,7 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
             y_trajectory[i] = y
             px_trajectory[i] = px
             py_trajectory[i] = py
+            pz_trajectory[i] = pz
 
         # Recursive call from the downramp
         if stage_copy.downramp is not None:
@@ -1582,15 +1487,16 @@ def _estimate_beam_trajectory(self, s0, x0, y0, px0, py0, pz0, q, num_steps, dri
 
             num_steps_downramp = stage_copy.num_steps_downramp
 
-            s_trajectory_downramp, x_trajectory_downramp, y_trajectory_downramp, px_trajectory_downramp, py_trajectory_downramp = downramp._estimate_beam_trajectory(prop_length, x, y, px, py, pz, q, num_steps_downramp, driver_x_trajectory, driver_y_trajectory)
+            s_trajectory_downramp, x_trajectory_downramp, y_trajectory_downramp, px_trajectory_downramp, py_trajectory_downramp, pz_trajectory_downramp, _, _ = downramp._estimate_beam_trajectory(prop_length, x, y, px, py, pz, q, num_steps_downramp+1, driver_x_trajectory, driver_y_trajectory)
 
-            s_trajectory[-len(s_trajectory_downramp):] = s_trajectory_downramp
-            x_trajectory[-len(x_trajectory_downramp):] = x_trajectory_downramp
-            y_trajectory[-len(y_trajectory_downramp):] = y_trajectory_downramp
-            px_trajectory[-len(px_trajectory_downramp):] = px_trajectory_downramp
-            py_trajectory[-len(py_trajectory_downramp):] = py_trajectory_downramp
+            s_trajectory[-len(s_trajectory_downramp)+1:] = s_trajectory_downramp[1:] # Dsicarding the first element in s_trajectory_downramp
+            x_trajectory[-len(x_trajectory_downramp)+1:] = x_trajectory_downramp[1:]
+            y_trajectory[-len(y_trajectory_downramp)+1:] = y_trajectory_downramp[1:]
+            px_trajectory[-len(px_trajectory_downramp)+1:] = px_trajectory_downramp[1:]
+            py_trajectory[-len(py_trajectory_downramp)+1:] = py_trajectory_downramp[1:]
+            pz_trajectory[-len(pz_trajectory_downramp)+1:] = pz_trajectory_downramp[1:]
 
-        return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory
+        return s_trajectory, x_trajectory, y_trajectory, px_trajectory, py_trajectory, pz_trajectory, driver_x_trajectory, driver_y_trajectory 
     
 
     # ==================================================

From b53e431669379aec7280e9514cd3011941bb5a15 Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 23 Mar 2026 12:25:55 +0100
Subject: [PATCH 83/84] Added phase_advance_traj_data() to
 src/abel/utilities/beam_physics.py.

---
 src/abel/utilities/beam_physics.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/src/abel/utilities/beam_physics.py b/src/abel/utilities/beam_physics.py
index 06d6261c..bca45c0c 100644
--- a/src/abel/utilities/beam_physics.py
+++ b/src/abel/utilities/beam_physics.py
@@ -1166,6 +1166,35 @@ def phase_advance(ss, betas):
     return integrate.simpson(y=inv_betas, x=ss)
 
 
+# =============================================
+def phase_advance_traj_data(pos, angles):
+    """
+    Compute the unwrapped betatron phase from transverse position and slope 
+    data. 
+    
+    Parameters
+    ----------
+    pos : [m] 1D float ndarray
+        Transverse position along the beamline.
+        
+    angles : 1D float ndarray
+        Transverse angle dpos/ds
+    
+    Returns
+    -------
+    phase_advance : float
+        Unwrapped phase angle in radians, representing the continuous betatron
+        phase advance at the end of the beamline
+    """
+    
+    # Compute the phase space angle by normalising coordinates to a unit circle 
+    # and taking the arctangent before unwrapping
+    phase = np.unwrap(np.arctan2(angles / np.max(np.abs(angles)), 
+                                 pos / np.max(np.abs(pos))))
+    
+    return np.abs(phase[-1] - phase[0]) 
+
+
 # =============================================
 def arc_lengths(s_trajectory, x_trajectory):
     """

From 5bc8b018b960d5d20ea61a5e0b2a82b475f9af6d Mon Sep 17 00:00:00 2001
From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com>
Date: Mon, 23 Mar 2026 14:26:08 +0100
Subject: [PATCH 84/84] Edited phase_advance_traj_data in
 src/abel/utilities/beam_physics.py to include correction for driver
 trajectory.

---
 src/abel/utilities/beam_physics.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/abel/utilities/beam_physics.py b/src/abel/utilities/beam_physics.py
index bca45c0c..2f19f680 100644
--- a/src/abel/utilities/beam_physics.py
+++ b/src/abel/utilities/beam_physics.py
@@ -1167,7 +1167,7 @@ def phase_advance(ss, betas):
 
 
 # =============================================
-def phase_advance_traj_data(pos, angles):
+def phase_advance_traj_data(pos, angles, orbit_pos=None, orbit_angles=None):
     """
     Compute the unwrapped betatron phase from transverse position and slope 
     data. 
@@ -1178,7 +1178,17 @@ def phase_advance_traj_data(pos, angles):
         Transverse position along the beamline.
         
     angles : 1D float ndarray
-        Transverse angle dpos/ds
+        Transverse angle dpos/ds.
+
+    orbit_pos : [m] 1D float ndarray, optional
+        Reference orbit transverse position (e.g., driver trajectory). If 
+        provided, these are subtracted before phase calculation. Defaults to 
+        `None`.
+
+    orbit_angles : 1D float ndarray, optional
+        Reference orbit transverse angles (e.g., driver trajectory angles). If 
+        provided, these are subtracted before phase calculation. Defaults to 
+        `None`.
     
     Returns
     -------
@@ -1186,6 +1196,11 @@ def phase_advance_traj_data(pos, angles):
         Unwrapped phase angle in radians, representing the continuous betatron
         phase advance at the end of the beamline
     """
+
+    if orbit_pos is not None:
+        pos = pos - orbit_pos
+    if orbit_angles is not None:
+        angles = angles - orbit_angles
     
     # Compute the phase space angle by normalising coordinates to a unit circle 
     # and taking the arctangent before unwrapping