From f86477d222edf685a5ba5104758e25489ca9327d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Thu, 22 Oct 2026 23:18:00 +0200 Subject: [PATCH 1/2] Add individual z shifts for drivers, based on stage number. --- .../classes/beamline/impl/driver_complex.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/abel/classes/beamline/impl/driver_complex.py b/src/abel/classes/beamline/impl/driver_complex.py index 1c130f4a..895e7e8b 100644 --- a/src/abel/classes/beamline/impl/driver_complex.py +++ b/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 From 63f2e348d8c5549de6ee0ca84ca7cf97c5fc0a57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Thu, 22 Oct 2026 23:20:19 +0200 Subject: [PATCH 2/2] Add driver tracking function to Stage to allow identifying stage number. Implement in all stage classes. --- src/abel/classes/stage/impl/stage_basic.py | 2 +- src/abel/classes/stage/impl/stage_hipace.py | 2 +- .../classes/stage/impl/stage_quasistatic_2d.py | 2 +- .../classes/stage/impl/stage_reduced_models.py | 4 ++-- src/abel/classes/stage/impl/stage_wake_t.py | 2 +- src/abel/classes/stage/stage.py | 16 ++++++++++++++++ 6 files changed, 22 insertions(+), 6 deletions(-) diff --git a/src/abel/classes/stage/impl/stage_basic.py b/src/abel/classes/stage/impl/stage_basic.py index dc30c780..94a02eed 100644 --- a/src/abel/classes/stage/impl/stage_basic.py +++ b/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) diff --git a/src/abel/classes/stage/impl/stage_hipace.py b/src/abel/classes/stage/impl/stage_hipace.py index 58601455..928a8c36 100644 --- a/src/abel/classes/stage/impl/stage_hipace.py +++ b/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 diff --git a/src/abel/classes/stage/impl/stage_quasistatic_2d.py b/src/abel/classes/stage/impl/stage_quasistatic_2d.py index 529d3cf8..54411a4e 100644 --- a/src/abel/classes/stage/impl/stage_quasistatic_2d.py +++ b/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) diff --git a/src/abel/classes/stage/impl/stage_reduced_models.py b/src/abel/classes/stage/impl/stage_reduced_models.py index ae36a750..5b0ad55b 100644 --- a/src/abel/classes/stage/impl/stage_reduced_models.py +++ b/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 diff --git a/src/abel/classes/stage/impl/stage_wake_t.py b/src/abel/classes/stage/impl/stage_wake_t.py index 191625c1..bdadcc59 100644 --- a/src/abel/classes/stage/impl/stage_wake_t.py +++ b/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 diff --git a/src/abel/classes/stage/stage.py b/src/abel/classes/stage/stage.py index b66b4e48..a49e019a 100644 --- a/src/abel/classes/stage/stage.py +++ b/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