Implement individual z-shifts per stage in DriverComplex

src/abel/classes/beamline/impl/driver_complex.py

@@ -14,18 +14,33 @@
 
 class DriverComplex(Beamline):
 
-    def __init__(self, source=None, rf_accelerator=None, combiner_ring=None, turnaround=None, transfer_line=None, nom_energy=None, num_drivers=None, num_bunches_in_train=None, bunch_separation=None, rep_rate_trains=None):
+    def __init__(self, source=None, rf_accelerator=None, combiner_ring=None, turnaround=None, transfer_line=None, nom_energy=None, num_drivers=None, num_bunches_in_train=None, bunch_separation=None, rep_rate_trains=None, z_shifts=None):
 
         super().__init__(num_bunches_in_train, bunch_separation, rep_rate_trains)
 
         self.num_drivers = num_drivers
         self.nom_energy = None
+        self.z_shifts = z_shifts
 
         self.source = source
         self.rf_accelerator = rf_accelerator
         self.combiner_ring = combiner_ring
         self.transfer_line = transfer_line
         self.turnaround = turnaround
+
+    def track(self, beam=None, savedepth=0, runnable=None, verbose=False, stage_number=None):
+
+        # shift the driver timing (from the source)
+        if stage_number is not None:
+            if self.z_shifts is not None:
+                if stage_number < len(self.z_shifts):
+                    self.source.z_offset = self.source.z_offset + self.z_shifts[stage_number]
+
+        # assemble trackables if not already done
+        if not self.trackables:
+            self.assemble_trackables()
+
+        return super().track(beam, savedepth, runnable, verbose)
 
 
     # assemble the trackables

src/abel/classes/stage/impl/stage_basic.py

@@ -99,7 +99,7 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
         self.stage_number = beam_incoming.stage_number
 
         # get the driver
-        driver_incoming = self.driver_source.track()
+        driver_incoming = self.track_driver_source()
 
         original_driver = copy.deepcopy(driver_incoming)
         original_beam = copy.deepcopy(beam_incoming)

src/abel/classes/stage/impl/stage_hipace.py

@@ -261,7 +261,7 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
             os.mkdir(tmpfolder)
 
         # generate driver
-        driver_incoming = self.driver_source.track()
+        driver_incoming = self.track_driver_source()
 
         # Set ramp lengths, nominal energies, nominal energy gains
         # and flattop nominal energy if not already done

src/abel/classes/stage/impl/stage_quasistatic_2d.py

@@ -67,7 +67,7 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
         self.stage_number = beam_incoming.stage_number
 
         # make driver
-        driver_incoming = self.driver_source.track()
+        driver_incoming = self.track_driver_source()
         original_driver = copy.deepcopy(driver_incoming)
 
 

src/abel/classes/stage/impl/stage_reduced_models.py

@@ -381,7 +381,7 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
 
 
         # ========== Get the drive beam ==========
-        driver_incoming = self.driver_source.track()  # Generate a drive beam with jitter.
+        driver_incoming = self.track_driver_source()  # Generate a drive beam with jitter.
 
         original_driver = copy.deepcopy(driver_incoming)
         original_beam = copy.deepcopy(beam_incoming)
@@ -1305,7 +1305,7 @@ def mask_bubble_radius_spikes(self, bubble_radius, zs_bubble_radius, make_plot=F
 
         # Mask based on peaks
         if self.estm_R_blowout is None:
-            drive_beam = self.driver_source.track()
+            drive_beam = self.track_driver_source()
             R_blowout = blowout_radius(self.plasma_density, drive_beam.peak_current())
         else:
             R_blowout = self.estm_R_blowout

src/abel/classes/stage/impl/stage_wake_t.py

@@ -94,7 +94,7 @@ def track(self, beam0, savedepth=0, runnable=None, verbose=False):
 
         # ========== Make driver (and convert to WakeT bunch) ==========
         if not self.use_single_beam:
-            driver0 = self.driver_source.track()
+            driver0 = self.track_driver_source()
             driver_original = copy.deepcopy(driver0)
 
             # rotate the beam coordinate system to align with the driver

src/abel/classes/stage/stage.py

@@ -201,7 +201,23 @@ def driver_source(self, source : Source | DriverComplex | None):
         self._driver_source = source
     _driver_source = None
 
+
+    def track_driver_source(self):
+        """
+        Tracks the driver source, supplying the stage number 
+        if it exists and the source is a DriverComplex.
+
+        Returns
+        -------
+        driver : ``Beam``
+        """
+        if isinstance(self.driver_source, DriverComplex):
+            driver = self.driver_source.track(stage_number=self.stage_number)
+        else:
+            driver = self.driver_source.track()
+        return driver
 
+
     def get_actual_driver_source(self):
         """
         Return the driver source of the stage or the driver source of the