diff --git a/pyproject.toml b/pyproject.toml
index 41263101..777d938e 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",
 ]
diff --git a/src/abel/classes/stage/impl/stage_hipace.py b/src/abel/classes/stage/impl/stage_hipace.py
index 931fe804..74c42302 100644
--- a/src/abel/classes/stage/impl/stage_hipace.py
+++ b/src/abel/classes/stage/impl/stage_hipace.py
@@ -87,7 +87,11 @@ 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``.
+
+    driver_num_half_beta_osc : float, optional
+        Number of half betatron oscillations that the drive beam is intended 
+        to perform. Defaults to 1.0.
 
     mesh_refinement : bool, optional
         Enable HiPACE++ mesh refinement. See the 
@@ -113,7 +117,7 @@ class StageHipace(Stage):
     .. [1] HiPACE++ User Guide, https://hipace.readthedocs.io/
     """
     
-    def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, driver_source=None, ramp_beta_mag=None, keep_data=False, save_drivers=False, output=None, ion_motion=True, ion_species='H', beam_ionization=True, radiation_reaction=False, num_nodes=1, num_cell_xy=511, driver_only=False, plasma_density_from_file=None, no_plasma=False, external_focusing=False, mesh_refinement=True, do_spin_tracking=False, run_path=None):
+    def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, driver_source=None, ramp_beta_mag=None, keep_data=False, save_drivers=False, output=None, ion_motion=True, ion_species='H', beam_ionization=True, radiation_reaction=False, num_nodes=1, num_cell_xy=511, driver_only=False, plasma_density_from_file=None, no_plasma=False, driver_num_half_beta_osc=1.0, external_focusing=False, mesh_refinement=True, do_spin_tracking=False, run_path=None):
         """
         The constructor for the ``StageHipace`` class.
 
@@ -182,6 +186,10 @@ def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, drive
         no_plasma : bool, optional
             If ``True``, runs the stage without plasma. Defaults to ``False``.
 
+        driver_num_half_beta_osc : float, optional
+            Number of half betatron oscillations that the drive beam is intended 
+            to perform. Defaults to 1.0.
+
         external_focusing : bool, optional
             Flag for enabling drive beam guiding by applying a linear transverse 
             external magnetic field across the beams. If ``True``, the field 
@@ -213,9 +221,10 @@ def __init__(self, length=None, nom_energy_gain=None, plasma_density=None, drive
         self.save_drivers = save_drivers
         self.no_plasma = no_plasma
 
+        self.driver_num_half_beta_osc = driver_num_half_beta_osc
+
         # external focusing (APL-like) [T/m]
         self.external_focusing = external_focusing
-        self._external_focusing_gradient = None
 
         # plasma profile
         self.plasma_profile = SimpleNamespace()
@@ -268,13 +277,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 +352,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 +695,321 @@ 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_num_half_beta_osc)
+
+        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_num_half_beta_osc``.
+            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_num_half_beta_osc
+
+        return self.driver_source.energy/SI.c*(num_half_oscillations*np.pi/L)**2  # [T/m]
+    
+
+    # ==================================================
+    def calc_length_num_beta_osc(self, num_beta_osc, initial_energy=None, plasma_density=None, driver_num_half_beta_osc=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 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_num_half_beta_osc : float, optional
+            Number of half betatron oscillations that the drive beam is 
+            intended to perform. If ``None``, will use :attr:`StageHipace.driver_num_half_beta_osc <abel.StageHipace.driver_num_half_beta_osc>`.
+            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 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 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.')
+        
+        if driver_num_half_beta_osc is None:
+            driver_num_half_beta_osc = self.driver_num_half_beta_osc
+        if driver_num_half_beta_osc < 0:
+            raise ValueError('Number of driver oscillations must be positive.')
+
+        # 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)
+
+        # 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
+
+            # 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()
+
+            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_num_half_beta_osc, 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)
+            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]
+
+        return length
+    
+
+    # ==================================================
+    def match_length_2_num_beta_osc(self, num_beta_osc, driver_num_half_beta_osc=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 
+        :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.
+
+        - Also set :attr:`self._external_focusing_gradient <abel.Stage._external_focusing_gradient>`
+        and :attr:`self.driver_num_half_beta_osc <abel.StageHipace.driver_num_half_beta_osc>`
+        to be consistent with the calculated stage length
+        if ``set_consistent_params`` and set :attr:`self.external_focusing <abel.StageHipace.external_focusing>` 
+        are ``True``. 
+
+
+        Parameters
+        ----------
+        num_beta_osc : float
+            Total number of design betatron oscillations that the electron 
+            should perform through the plasma stage excluding ramps. 
+
+        driver_num_half_beta_osc : float, optional
+            Number of half betatron oscillations that the drive beam is 
+            intended to perform. If ``None``, will use :attr:`StageHipace.driver_num_half_beta_osc <abel.StageHipace.driver_num_half_beta_osc>`
+            Defaults to ``None``.
+
+        set_consistent_params : bool, optional
+            Flag for setting :attr:`self._external_focusing_gradient <abel.Stage._external_focusing_gradient>`
+            and :attr:`self.driver_num_half_beta_osc <abel.StageHipace.driver_num_half_beta_osc>`
+            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
+        -------
+        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. 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
+
+        if driver_num_half_beta_osc is None:
+            driver_num_half_beta_osc = self.driver_num_half_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, 
+                                                       driver_num_half_beta_osc=driver_num_half_beta_osc, 
+                                                       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_num_half_beta_osc, L=length)
+            self.driver_num_half_beta_osc = driver_num_half_beta_osc 
     
+
     # ==================================================
     # Apply waterfall function to all beam dump files
     def __waterfall_fcn(self, fcns, edges, data_dir, species='beam', remove_halo_nsigma=None, args=None):
@@ -698,7 +1017,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``, 
@@ -727,7 +1046,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 
@@ -803,7 +1122,7 @@ def plot_waterfalls(self, data_dir, species='beam', remove_halo_nsigma=20, save_
 
 
         Returns
-        ----------
+        -------
         ``None``
         '''
 
diff --git a/src/abel/classes/stage/impl/stage_reduced_models.py b/src/abel/classes/stage/impl/stage_reduced_models.py
index ae36a750..995bfb8a 100644
--- a/src/abel/classes/stage/impl/stage_reduced_models.py
+++ b/src/abel/classes/stage/impl/stage_reduced_models.py
@@ -682,9 +682,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
@@ -782,9 +782,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
diff --git a/src/abel/classes/stage/stage.py b/src/abel/classes/stage/stage.py
index 45b61dde..a73eeb32 100644
--- a/src/abel/classes/stage/stage.py
+++ b/src/abel/classes/stage/stage.py
@@ -16,6 +16,7 @@
 import scipy.constants as SI
 from types import SimpleNamespace
 from abel.utilities.plasma_physics import beta_matched
+from abel.utilities.beam_physics import length2num_beta_osc
 from typing import Self
 
 class Stage(Trackable, CostModeled):
@@ -125,6 +126,7 @@ def __init__(self, nom_accel_gradient, nom_energy_gain, plasma_density, driver_s
         self.plasma_density = plasma_density
         self.driver_source = driver_source
         self.ramp_beta_mag = ramp_beta_mag
+        self.external_focusing = False
         
         self.stage_number = None
         
@@ -515,13 +517,39 @@ 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  # 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.
+
         if ramp_length < 0.0:
             raise ValueError(f"ramp_length = {ramp_length} [m] < 0.0")
         return ramp_length
     
 
+     # ==================================================
+    def get_ramp_length(self) -> float:
+        """
+        Get the length of the ramps if the stage has ramps already set up. 
+        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:
@@ -1333,10 +1361,12 @@ def get_cost_breakdown(self):
 
 
     # ==================================================
-    def matched_beta_function(self, energy_incoming, match_entrance=True):
+    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 matched to the upramp by default.
+        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=[g_ext*y, -g_ext*x, 0] if present. 
     
         
         Parameters
@@ -1348,21 +1378,32 @@ def matched_beta_function(self, energy_incoming, match_entrance=True):
             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 beta_matched(self.plasma_density, energy_incoming)*self.upramp.ramp_beta_mag
+                return 1/k_beta * self.upramp.ramp_beta_mag
             else:
-                return beta_matched(self.plasma_density, energy_incoming)
+                return 1/k_beta
         else:
             if self.downramp.ramp_beta_mag is not None:
-                return beta_matched(self.plasma_density, energy_incoming)*self.downramp.ramp_beta_mag
+                return 1/k_beta * self.downramp.ramp_beta_mag
             else:
                 raise ValueError('Downramp ramp_beta_mag not defined.')
 
@@ -1385,6 +1426,651 @@ def matched_beta_function_flattop(self, energy):
         '''
 
         return beta_matched(self.plasma_density, energy)
+
+
+    # ==================================================
+    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 defined nominal 
+        acceleration 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``.
+
+        plasma_density : [m^-3] float, optional
+            The plasma density of the plasma stage. Defaults to 
+            ``self.plasma_density``.
+
+        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.
+        """
+
+        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 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.')
+
+        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)
+
+        length = ((2*np.pi*num_beta_osc/gradient_prefactor + np.sqrt(initial_energy))**2 - initial_energy) / (q*nom_accel_gradient)
+
+        return length
+        
+    
+    # ==================================================
+    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_gradient>`
+        is not ``None``.
+
+        Parameters
+        ----------
+        length_flattop : [m] float, optional
+            Length of a plasma stage excluding ramps that the particle can 
+            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.nom_energy``.
+
+        nom_accel_gradient_flattop : [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``.
+
+        q : [C] float, optional
+            Particle charge. q * nom_accel_gradient_flattop must be positive. 
+            Defaults to elementary charge.
+
+            
+        Returns
+        -------
+        num_beta_osc : float
+            Total number of betatron oscillations that the particle will perform
+            across the plasma stage.
+        """
+
+        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:
+            plasma_density = self.plasma_density
+
+        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
+
+        return length2num_beta_osc(length_flattop, initial_energy, nom_accel_gradient_flattop, g, q=SI.e)
+        
+
+    # ==================================================
+    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.
+
+
+        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.
+        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
+
+
+    # =============================================
+    @property
+    def external_focusing_gradient(self) -> float:
+        """
+        Return None by default for ``Stage`` subclasses not supporting external 
+        focusing fields.
+        """
+        return None
+    @external_focusing_gradient.setter
+    def external_focusing_gradient(self, g_ext : float | None):
+        self._external_focusing_gradient = None  # Driver guiding with external field is not implemented yet, so always set None.
+    
+
+    # =============================================
+    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=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
+        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.
+
+        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
+        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, 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, 
+                                                        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, px_trajectory, py_trajectory, pz_trajectory, driver_x_trajectory, driver_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.
+
+        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
+
+        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 - 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
+
+        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=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)
+        pz_trajectory = np.full(num_steps, None, dtype=float)
+
+        # 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, 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]
+            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 = 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:
+                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
+            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. 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:]
+
+        # 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
+            dpz = - q * self.nom_accel_gradient_flattop * ds/SI.c
+            pz = pz + dpz
+
+            # 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
+            pz_trajectory[i] = pz
+
+        # 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, 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)+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, pz_trajectory, driver_x_trajectory, driver_y_trajectory 
     
 
     # ==================================================
diff --git a/src/abel/utilities/beam_physics.py b/src/abel/utilities/beam_physics.py
index 43ac3a82..fbb03265 100644
--- a/src/abel/utilities/beam_physics.py
+++ b/src/abel/utilities/beam_physics.py
@@ -5,6 +5,7 @@
 # License: GPL-3.0-or-later
 
 import numpy as np
+import scipy.constants as SI
 
 
 # =============================================
@@ -1152,3 +1153,113 @@ def evolve_chromatic_amplitude(ls, inv_rhos, ks, ms, taus, beta0, alpha0=0, Dx0=
 
     return W, evolution
 
+
+# =============================================
+def phase_advance_integrate_betas(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)
+
+
+# =============================================
+def phase_advance_traj_data(pos, angles, orbit_pos=None, orbit_angles=None):
+    """
+    Compute the unwrapped betatron phase advance from transverse position and 
+    angle data by mapping the phase-space trajectory onto a normalised unit 
+    circle and evaluating the continuous phase angle evolution.
+    
+    Parameters
+    ----------
+    pos : [m] 1D float ndarray
+        Transverse position along the beamline.
+        
+    angles : 1D float ndarray
+        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
+    -------
+    phase_advance : float
+        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
+    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 length2num_beta_osc(length, initial_energy, accel_gradient, foc_gradient, 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 : [m] float
+            Length of a plasma stage excluding ramps that the particle can 
+            perform betatron oscillations in. 
+
+        initial_energy : [eV] float
+            The initial energy of the particle at the start of the plasma stage. 
+
+        accel_gradient : [V/m] float
+            Nominal accelerating gradient of the plasma stage exclusing ramps. 
+
+        foc_gradient : [T/m] float
+            The gradient of the total transverse focusing field.
+
+        q : [C] float, optional
+            Particle charge. q * accel_gradient must be positive. 
+            Defaults to elementary charge.
+
+            
+        Returns
+        -------
+        num_beta_osc : float
+            Total number of betatron oscillations that the particle will perform
+            across the plasma stage.
+        """
+
+        if q * accel_gradient < 0:
+            raise ValueError('q * accel_gradient must be positive.')
+
+        if accel_gradient < 1e-15: # Need to treat very small gradients separately. Often the case for plasma density ramps.
+            return np.sqrt(np.abs(q)*SI.c*foc_gradient/(initial_energy*SI.e)) * length/(2*np.pi)
+        else:
+
+            prefactor = 2*np.sqrt(np.abs(q)*foc_gradient*SI.c) / (q*accel_gradient)
+            energy_scaling = np.sqrt(initial_energy*SI.e + q*accel_gradient*length) - np.sqrt(initial_energy*SI.e)
+
+            num_beta_osc = prefactor * energy_scaling / (2*np.pi)
+
+            return num_beta_osc
+
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:
diff --git a/tests/test_StageHipace.py b/tests/test_StageHipace.py
new file mode 100644
index 00000000..5b07bfe9
--- /dev/null
+++ b/tests/test_StageHipace.py
@@ -0,0 +1,370 @@
+# -*- 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
+
+
+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
+def test_external_focusing():
+    """
+    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 ==========
+    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_num_half_beta_osc, 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_num_half_beta_osc = 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_num_half_beta_osc, 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_num_half_beta_osc = 2.0
+    stage.length_flattop = 1.3520401015419283  # [m]
+    assert np.isclose(2.0, stage.driver_num_half_beta_osc, 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)
+    
+
+@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 
+    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
+
+    stage.length_flattop = stage.calc_length_num_beta_osc(num_beta_osc, initial_energy=stage.nom_energy, driver_num_half_beta_osc=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, driver_num_half_beta_osc=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
+    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, driver_num_half_beta_osc=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(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_num_half_beta_osc=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(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, driver_num_half_beta_osc=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(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_num_half_beta_osc=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(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, driver_num_half_beta_osc=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(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_num_half_beta_osc=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_2_num_beta_osc():
+    """
+    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.
+    """
+
+    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 stage.external_focusing_gradient is None
+    assert np.isclose(1.0, stage.driver_num_half_beta_osc, rtol=1e-15, atol=0.0)
+
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_num_half_beta_osc=1.0, set_consistent_params=True)
+    assert np.isclose(1.0, stage.driver_num_half_beta_osc, 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_num_half_beta_osc=2.0, set_consistent_params=True)
+    assert np.isclose(1.0, stage.driver_num_half_beta_osc, 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)
+
+
+    # ========== 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)
+
+
+    stage = setup_minimal_StageHipace(external_focusing=True)
+    stage.match_length_2_num_beta_osc(num_beta_osc=num_beta_osc, driver_num_half_beta_osc=2.0)
+
+    assert np.isclose(2.0, stage.driver_num_half_beta_osc, 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)
+
+
+    # ========== 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_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(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_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_num_half_beta_osc=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(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_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(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_2_num_beta_osc(num_beta_osc=num_beta_osc2, driver_num_half_beta_osc=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(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_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_num_half_beta_osc=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)
+    
+
+    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_2_num_beta_osc(num_beta_osc=num_beta_osc3, driver_num_half_beta_osc=2.0)
+
+    assert np.isclose(2.0, stage.driver_num_half_beta_osc, 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)
+    
+    
+
+
diff --git a/tests/test_StageReducedModels_beamline.py b/tests/test_StageReducedModels_beamline.py
index e9fccef4..305e5d6f 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,15 +72,15 @@ 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, 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]
+    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.ramp_beta_mag = ramp_beta_mag
     stage.enable_tr_instability = enable_tr_instability 
     stage.enable_radiation_reaction = enable_radiation_reaction
@@ -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, 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)
@@ -203,6 +202,10 @@ def test_baseline_linac():
     shutil.rmtree(linac.run_path())
 
 
+def no_accel_phase_advance(length, initial_energy, foc_gradient):
+    return np.sqrt(SI.e*SI.c*foc_gradient/(initial_energy*SI.e)) * length
+
+
 @pytest.mark.StageReducedModels_linac
 def test_ramped_linac():
     """
@@ -211,6 +214,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]
@@ -221,14 +225,28 @@ 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, 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)
+
+    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_flattop)
+    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)
@@ -239,15 +257,16 @@ def test_ramped_linac():
     assert linac.stage.driver_source.align_beam_axis is True
     assert stages[0].driver_source.align_beam_axis is True
     assert stages[1].driver_source.align_beam_axis is True
-    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)
@@ -255,7 +274,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)
@@ -284,7 +303,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()
@@ -297,6 +316,54 @@ 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
+    g_upramp1 = SI.e*upramp1.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
+    if upramp1.external_focusing_gradient is not None:
+        g_upramp1 = g_upramp1 + upramp1.external_focusing_gradient
+    assert np.isclose(no_accel_phase_advance(upramp1.length, upramp1.nom_energy, g_upramp1), 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.
+
+    g_downramp1 = SI.e*downramp1.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
+    if downramp1.external_focusing_gradient is not None:
+        g_downramp1 = g_downramp1 + downramp1.external_focusing_gradient
+    assert np.isclose(no_accel_phase_advance(downramp1.length, downramp1.nom_energy, g_downramp1), 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)
+
+    g_upramp2 = SI.e*upramp2.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
+    if upramp2.external_focusing_gradient is not None:
+        g_upramp2 = g_upramp2 + upramp2.external_focusing_gradient
+    assert np.isclose(no_accel_phase_advance(upramp2.length, upramp2.nom_energy, g_upramp2), 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.
+
+    g_downramp2 = SI.e*downramp2.plasma_density/(2*SI.epsilon_0*SI.c)  # [T/m], ion background focusing gradient
+    if downramp2.external_focusing_gradient is not None:
+        g_downramp2 = g_downramp2 + downramp2.external_focusing_gradient
+    assert np.isclose(no_accel_phase_advance(downramp2.length, downramp2.nom_energy, g_downramp2), 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()
@@ -333,7 +400,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, 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)
@@ -367,7 +434,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, 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)
@@ -405,7 +472,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, 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)
@@ -439,11 +506,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, 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)
@@ -477,12 +543,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, 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)
@@ -516,7 +581,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, 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)
@@ -614,7 +679,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)
@@ -644,11 +709,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, 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)
@@ -679,11 +743,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, 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)
@@ -718,12 +781,11 @@ 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)
     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, 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, 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
@@ -824,7 +886,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)
