From d83b0c62d1edde1c792aed275e21a7512c0b3555 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Sat, 6 Jun 2026 20:59:25 +0200 Subject: [PATCH 1/9] Add functionality to overwrite from a certain trackable element. --- src/abel/classes/beamline/beamline.py | 12 ++++++++-- src/abel/classes/runnable.py | 34 ++++++++++++++++++--------- 2 files changed, 33 insertions(+), 13 deletions(-) diff --git a/src/abel/classes/beamline/beamline.py b/src/abel/classes/beamline/beamline.py index 8cf64a60..b9b8e9cd 100644 --- a/src/abel/classes/beamline/beamline.py +++ b/src/abel/classes/beamline/beamline.py @@ -42,14 +42,22 @@ def assemble_trackables(self): # perform tracking - def track(self, beam=None, savedepth=0, runnable=None, verbose=False): + def track(self, beam=None, savedepth=0, runnable=None, verbose=False, start_from=0): # assemble the trackables if self.trackables is None: self.assemble_trackables() # perform element-wise tracking - for trackable in self.trackables: + for i, trackable in enumerate(self.trackables): + + # skip to the given trackable + if i < start_from: + continue + if start_from > 0 and start_from == i: + assert beam is not None + + # track through the trackable beam = trackable.track(beam, savedepth-1, runnable, verbose) return beam diff --git a/src/abel/classes/runnable.py b/src/abel/classes/runnable.py index d52527b6..a5110974 100644 --- a/src/abel/classes/runnable.py +++ b/src/abel/classes/runnable.py @@ -30,16 +30,27 @@ def perform_shot(self, shot): if verbose_exists: print('>> SHOT ' + str(shot+1) + ' already exists and will not be overwritten.', flush=True) else: - - # clear the shot folder - self.clear_run_data(shot) - - # run tracking + + # print info if self.num_shots > 1 and self.verbose: print('>> SHOT ' + str(shot+1) + '/' + str(self.num_shots), flush=True) - # if overwrite_from is None: # TODO - self.track(beam=None, savedepth=self.savedepth, runnable=self, verbose=self.verbose) + # set starting point and load beam if it exists + if (self.overwrite_from is None or self.overwrite_from == 0) or not os.path.exists(self.object_path(shot)): + self.clear_run_data(shot) + beam = None + start_from = 0 + else: + num_outputs = self.num_outputs(shot=shot) + if self.overwrite_from < 0: + index = num_outputs + self.overwrite_from -1 + else: + index = self.overwrite_from - 1 + beam = self.get_beam(index, shot=shot) + start_from = index + 1 + + # perform tracking + self.track(beam=beam, savedepth=self.savedepth, runnable=self, verbose=self.verbose, start_from=start_from) # save object to file self.save() @@ -50,7 +61,7 @@ def is_scan(self): return self.scan_fcn is not None # scan function - def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_shots_per_step=1, step_filter=None, shot_filter=None, savedepth=2, verbose=None, overwrite=False, parallel=False, max_cores=16): + def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_shots_per_step=1, step_filter=None, shot_filter=None, savedepth=2, verbose=None, overwrite=False, parallel=False, max_cores=16, overwrite_from=None): from joblib import Parallel, delayed from joblib_progress import joblib_progress @@ -62,10 +73,11 @@ def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_sh self.run_name = run_name self.overwrite = overwrite + self.overwrite_from = overwrite_from self.verbose = verbose self.savedepth = savedepth - if self.overwrite: + if self.overwrite and self.overwrite_from is None: self.clear_run_data() self.overwrite = False @@ -123,7 +135,7 @@ def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_sh # run simulation - def run(self, run_name=None, num_shots=1, savedepth=2, verbose=None, overwrite=False, parallel=False, max_cores=16): + def run(self, run_name=None, num_shots=1, savedepth=2, verbose=None, overwrite=False, parallel=False, max_cores=16, overwrite_from=None): # define run name (generate if not given) if run_name is None: @@ -132,7 +144,7 @@ def run(self, run_name=None, num_shots=1, savedepth=2, verbose=None, overwrite=F self.run_name = run_name # perform a scan with only one step - self.scan(run_name=self.run_name, num_shots_per_step=num_shots, savedepth=savedepth, verbose=verbose, overwrite=overwrite, parallel=parallel, max_cores=max_cores) + self.scan(run_name=self.run_name, num_shots_per_step=num_shots, savedepth=savedepth, verbose=verbose, overwrite=overwrite, parallel=parallel, max_cores=max_cores, overwrite_from=overwrite_from) From 4237441b9589273eb901e70c0d1439a26bfe2bc3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Sun, 7 Jun 2026 19:14:57 +0200 Subject: [PATCH 2/9] Allow loading beams from symlinks. --- src/abel/classes/beam.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/abel/classes/beam.py b/src/abel/classes/beam.py index 6fa748e4..e62af370 100644 --- a/src/abel/classes/beam.py +++ b/src/abel/classes/beam.py @@ -6,7 +6,7 @@ import numpy as np import openpmd_api as io -import copy, warnings +import copy, warnings, os import scipy.constants as SI import scipy.sparse as sp from scipy.spatial.transform import Rotation as Rot @@ -2074,7 +2074,13 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): # load beam (from OpenPMD format) @classmethod - def load(_, filename, beam_name='beam'): + def load(_, filename0, beam_name='beam'): + + # resolve if alias + if os.path.islink(filename0): + filename = os.readlink(filename0) + else: + filename = filename0 # load file series = io.Series(filename, io.Access.read_only) From cf7f63d0e428e9fd1659d8fade5b9c510799aae0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Sun, 7 Jun 2026 19:16:29 +0200 Subject: [PATCH 3/9] Small bugfix to stop warnings from infinite stage efficiency. --- src/abel/classes/stage/stage.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/abel/classes/stage/stage.py b/src/abel/classes/stage/stage.py index 45b61dde..c3278fcb 100644 --- a/src/abel/classes/stage/stage.py +++ b/src/abel/classes/stage/stage.py @@ -1449,7 +1449,10 @@ def calculate_efficiency(self, beam0, driver0, beam, driver): Etot0_driver = driver0.total_energy() Etot_driver = driver.total_energy() self.efficiency.driver_to_wake = (Etot0_driver-Etot_driver)/Etot0_driver - self.efficiency.wake_to_beam = (Etot_beam-Etot0_beam)/(Etot0_driver-Etot_driver) + if abs(Etot0_driver-Etot_driver) > 0: + self.efficiency.wake_to_beam = (Etot_beam-Etot0_beam)/(Etot0_driver-Etot_driver) + else: + self.efficiency.wake_to_beam = np.nan self.efficiency.driver_to_beam = self.efficiency.driver_to_wake*self.efficiency.wake_to_beam if self.get_rep_rate_average() is not None: self.efficiency.dumped_power = Etot_driver*self.get_rep_rate_average() From fc40c9699c1a5f179d3eb6c881de7f36a3dc3be4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Sun, 7 Jun 2026 19:17:40 +0200 Subject: [PATCH 4/9] Allow tracking of beamlines until a certain element. Also renamed the starting point for consistency. --- src/abel/classes/beamline/beamline.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/abel/classes/beamline/beamline.py b/src/abel/classes/beamline/beamline.py index b9b8e9cd..541026cc 100644 --- a/src/abel/classes/beamline/beamline.py +++ b/src/abel/classes/beamline/beamline.py @@ -42,7 +42,7 @@ def assemble_trackables(self): # perform tracking - def track(self, beam=None, savedepth=0, runnable=None, verbose=False, start_from=0): + def track(self, beam=None, savedepth=0, runnable=None, verbose=False, track_from=0, track_until=-1): # assemble the trackables if self.trackables is None: @@ -52,11 +52,16 @@ def track(self, beam=None, savedepth=0, runnable=None, verbose=False, start_from for i, trackable in enumerate(self.trackables): # skip to the given trackable - if i < start_from: + if i < track_from: continue - if start_from > 0 and start_from == i: + if track_from > 0 and track_from == i: assert beam is not None - + if track_until < 0: + if i > len(self.trackables) + track_until: + break + elif i > track_until: + break + # track through the trackable beam = trackable.track(beam, savedepth-1, runnable, verbose) From a41211afbaf952d5cdfdaf4f7a8685f3e7de2b8f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Sun, 7 Jun 2026 19:19:55 +0200 Subject: [PATCH 5/9] Added functionality for simultaneously handling a shared run up to a certain element, as well as overwriting from a certain element. --- src/abel/classes/runnable.py | 124 ++++++++++++++++++++++++++--------- 1 file changed, 93 insertions(+), 31 deletions(-) diff --git a/src/abel/classes/runnable.py b/src/abel/classes/runnable.py index a5110974..fb3290be 100644 --- a/src/abel/classes/runnable.py +++ b/src/abel/classes/runnable.py @@ -24,36 +24,93 @@ def perform_shot(self, shot): vals_all = np.repeat(self.vals,self.num_shots_per_step) self.scan_fcn(self, vals_all[shot]) - # check if object exists - if not self.overwrite and os.path.exists(self.object_path(shot)): - verbose_exists = self.verbose - if verbose_exists: - print('>> SHOT ' + str(shot+1) + ' already exists and will not be overwritten.', flush=True) - else: + # if the run object does not exists, then clean and run normally + if not os.path.exists(self.object_path(shot)): + + # delete incomplete run data + self.clear_run_data(shot) + # start from the beginning + beam = None + track_from = 0 + # print info if self.num_shots > 1 and self.verbose: print('>> SHOT ' + str(shot+1) + '/' + str(self.num_shots), flush=True) + + else: - # set starting point and load beam if it exists - if (self.overwrite_from is None or self.overwrite_from == 0) or not os.path.exists(self.object_path(shot)): - self.clear_run_data(shot) - beam = None - start_from = 0 + # if the run object exists but should not be overwritten, then stop + if not self.overwrite: + + if self.verbose: + print('>> SHOT ' + str(shot+1) + ' already exists and will not be overwritten.', flush=True) + return + else: - num_outputs = self.num_outputs(shot=shot) - if self.overwrite_from < 0: - index = num_outputs + self.overwrite_from -1 + + # overwrite the existing shot from the start + if self.overwrite_from is None or self.overwrite_from == 0: + + # delete incomplete run data (if not continuing a shared run) + if not (self.shared_run_until is not None and not self.is_shared_run and shot == 0): + self.clear_run_data(shot) + + # start from the beginning + beam = None + track_from = 0 + + # print info + if self.num_shots > 1 and self.verbose: + print('>> SHOT ' + str(shot+1) + '/' + str(self.num_shots) + ' (overwriting)', flush=True) + else: - index = self.overwrite_from - 1 - beam = self.get_beam(index, shot=shot) - start_from = index + 1 - - # perform tracking - self.track(beam=beam, savedepth=self.savedepth, runnable=self, verbose=self.verbose, start_from=start_from) - # save object to file - self.save() + # overwrite the existing shot from a given point + if self.overwrite_from < 0: + num_outputs = self.num_outputs(shot=shot) + index = num_outputs + self.overwrite_from - 1 + else: + index = self.overwrite_from - 1 + + # load the beam from the given point + beam = self.get_beam(index, shot=shot) + track_from = index + 1 + + # print info + if self.num_shots > 1 and self.verbose: + print('>> SHOT ' + str(shot+1) + '/' + str(self.num_shots) + ' (overwriting from element #'+ str(track_from) +')', flush=True) + + # setup shared runs (up to a certain element) + track_until = -1 + if self.shared_run_until is not None: + if self.is_shared_run: + # shared run until this element + track_until = self.shared_run_until + else: + # continuing after the shared run (the other shots symlink to the first shot) + if shot > 0: + self.clear_run_data(shot) + shot_path = self.shot_path(shot) + for i, file in enumerate(self.run_data(shot=0)): + if i > self.shared_run_until: + break + symlink = os.path.join(shot_path, os.path.basename(file)) + if os.path.exists(symlink): + os.remove(symlink) + os.symlink(file, symlink) + print(f' ... #0-{self.shared_run_until} symlinked to SHOT 1') + + # set the starting point for continuing (and load the correct beam) + index = self.shared_run_until + 1 + track_from = index + beam = self.get_beam(index-1, shot=0) + + # perform tracking + self.track(beam=beam, savedepth=self.savedepth, runnable=self, verbose=self.verbose, track_from=track_from, track_until=track_until) + + # save object to file + self.save() ## SCAN FUNCTIONALITY @@ -61,7 +118,7 @@ def is_scan(self): return self.scan_fcn is not None # scan function - def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_shots_per_step=1, step_filter=None, shot_filter=None, savedepth=2, verbose=None, overwrite=False, parallel=False, max_cores=16, overwrite_from=None): + def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_shots_per_step=1, step_filter=None, shot_filter=None, savedepth=2, verbose=None, overwrite=False, parallel=False, max_cores=16, overwrite_from=None, shared_run_until=None): from joblib import Parallel, delayed from joblib_progress import joblib_progress @@ -76,11 +133,7 @@ def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_sh self.overwrite_from = overwrite_from self.verbose = verbose self.savedepth = savedepth - - if self.overwrite and self.overwrite_from is None: - self.clear_run_data() - self.overwrite = False - + # default verbosity if self.verbose is None: self.verbose = not parallel @@ -113,6 +166,15 @@ def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_sh elif shot_filter is not None: shots_to_perform = shots_to_perform[np.isin(shots_to_perform, shot_filter)] + # TODO: add functionality where a common run starts it off (single run; if parallel run this part alone first), then do the scan after this. + # The attribute could be named "common_run_until=None". Can use symlinks to shot 0 in the other shots to save space. + # TODO: FIX THIS + self.shared_run_until = shared_run_until + if shared_run_until is not None: + self.is_shared_run = True + self.perform_shot(shot=0) + self.is_shared_run = False + # perform shots (in parallel or series) if parallel: @@ -135,7 +197,7 @@ def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_sh # run simulation - def run(self, run_name=None, num_shots=1, savedepth=2, verbose=None, overwrite=False, parallel=False, max_cores=16, overwrite_from=None): + def run(self, run_name=None, num_shots=1, savedepth=2, verbose=None, overwrite=False, parallel=False, max_cores=16, overwrite_from=None, shared_run_until=None): # define run name (generate if not given) if run_name is None: @@ -144,7 +206,7 @@ def run(self, run_name=None, num_shots=1, savedepth=2, verbose=None, overwrite=F self.run_name = run_name # perform a scan with only one step - self.scan(run_name=self.run_name, num_shots_per_step=num_shots, savedepth=savedepth, verbose=verbose, overwrite=overwrite, parallel=parallel, max_cores=max_cores, overwrite_from=overwrite_from) + self.scan(run_name=self.run_name, num_shots_per_step=num_shots, savedepth=savedepth, verbose=verbose, overwrite=overwrite, parallel=parallel, max_cores=max_cores, overwrite_from=overwrite_from, shared_run_until=shared_run_until) @@ -193,7 +255,7 @@ def shot_path(self, shot=None): def run_data(self, shot=None): shot_path = self.shot_path(shot) if os.path.exists(shot_path): - filenames = [shot_path + f for f in os.listdir(shot_path) if (os.path.isfile(os.path.join(shot_path, f)) and f.startswith('beam_') and not f.endswith('.obj'))] + filenames = [shot_path + f for f in os.listdir(shot_path) if ((os.path.isfile(os.path.join(shot_path, f)) or os.path.islink(os.path.join(shot_path, f))) and f.startswith('beam_') and not f.endswith('.obj'))] filenames.sort() return filenames else: From e57df036c80846301e9563eaca7152a1e05ffebc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Sun, 7 Jun 2026 19:43:38 +0200 Subject: [PATCH 6/9] Reverted updated attribute name in Beam.load for backward compatibility. --- src/abel/classes/beam.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/abel/classes/beam.py b/src/abel/classes/beam.py index e62af370..37819555 100644 --- a/src/abel/classes/beam.py +++ b/src/abel/classes/beam.py @@ -2074,13 +2074,11 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): # load beam (from OpenPMD format) @classmethod - def load(_, filename0, beam_name='beam'): + def load(_, filename, beam_name='beam'): # resolve if alias - if os.path.islink(filename0): - filename = os.readlink(filename0) - else: - filename = filename0 + if os.path.islink(filename): + filename = os.readlink(filename) # load file series = io.Series(filename, io.Access.read_only) From c9519fe755abba75c8a14302885d5d8aa1186cac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Sun, 7 Jun 2026 22:38:59 +0200 Subject: [PATCH 7/9] Updated which element to be run in shared run (+1). --- src/abel/classes/runnable.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/abel/classes/runnable.py b/src/abel/classes/runnable.py index fb3290be..fbde882c 100644 --- a/src/abel/classes/runnable.py +++ b/src/abel/classes/runnable.py @@ -86,14 +86,14 @@ def perform_shot(self, shot): if self.shared_run_until is not None: if self.is_shared_run: # shared run until this element - track_until = self.shared_run_until + track_until = self.shared_run_until-1 else: # continuing after the shared run (the other shots symlink to the first shot) if shot > 0: self.clear_run_data(shot) shot_path = self.shot_path(shot) for i, file in enumerate(self.run_data(shot=0)): - if i > self.shared_run_until: + if i > self.shared_run_until-1: break symlink = os.path.join(shot_path, os.path.basename(file)) if os.path.exists(symlink): @@ -102,7 +102,7 @@ def perform_shot(self, shot): print(f' ... #0-{self.shared_run_until} symlinked to SHOT 1') # set the starting point for continuing (and load the correct beam) - index = self.shared_run_until + 1 + index = self.shared_run_until track_from = index beam = self.get_beam(index-1, shot=0) From 852404aad237133913f8f29c74509e027a94cc38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Sun, 7 Jun 2026 23:22:01 +0200 Subject: [PATCH 8/9] Bugfix in IP run, due to removing deletion of files in the Runnable run function. --- src/abel/classes/ip/ip.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/abel/classes/ip/ip.py b/src/abel/classes/ip/ip.py index 8b2fe0ac..536ac54d 100644 --- a/src/abel/classes/ip/ip.py +++ b/src/abel/classes/ip/ip.py @@ -104,7 +104,7 @@ def clear_run_data(self, shot1=None, shot2=None): # delete and remake folder if os.path.exists(clear_path): shutil.rmtree(clear_path) - os.makedirs(clear_path) + # generate track path def shot_path(self, shot1=None, shot2=None): From 7b23057fed2593fcd7e2b012bf0807f6fcc5b4c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Sun, 7 Jun 2026 23:22:49 +0200 Subject: [PATCH 9/9] Fixed comment. --- src/abel/classes/runnable.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/abel/classes/runnable.py b/src/abel/classes/runnable.py index fbde882c..a4584d88 100644 --- a/src/abel/classes/runnable.py +++ b/src/abel/classes/runnable.py @@ -166,9 +166,7 @@ def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_sh elif shot_filter is not None: shots_to_perform = shots_to_perform[np.isin(shots_to_perform, shot_filter)] - # TODO: add functionality where a common run starts it off (single run; if parallel run this part alone first), then do the scan after this. - # The attribute could be named "common_run_until=None". Can use symlinks to shot 0 in the other shots to save space. - # TODO: FIX THIS + # perform shared run prior to the separate shots self.shared_run_until = shared_run_until if shared_run_until is not None: self.is_shared_run = True