diff --git a/src/abel/classes/beam.py b/src/abel/classes/beam.py index 6fa748e4..37819555 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 @@ -2075,6 +2075,10 @@ def save(self, runnable=None, filename=None, beam_name="beam", series=None): # load beam (from OpenPMD format) @classmethod def load(_, filename, beam_name='beam'): + + # resolve if alias + if os.path.islink(filename): + filename = os.readlink(filename) # load file series = io.Series(filename, io.Access.read_only) diff --git a/src/abel/classes/beamline/beamline.py b/src/abel/classes/beamline/beamline.py index 8cf64a60..541026cc 100644 --- a/src/abel/classes/beamline/beamline.py +++ b/src/abel/classes/beamline/beamline.py @@ -42,14 +42,27 @@ 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, track_from=0, track_until=-1): # 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 < track_from: + continue + 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) return beam 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): diff --git a/src/abel/classes/runnable.py b/src/abel/classes/runnable.py index d52527b6..a4584d88 100644 --- a/src/abel/classes/runnable.py +++ b/src/abel/classes/runnable.py @@ -24,25 +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: - - # clear the shot folder + # 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 - # run tracking + # print info if self.num_shots > 1 and self.verbose: print('>> SHOT ' + str(shot+1) + '/' + str(self.num_shots), flush=True) + + else: + + # 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: + + # 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: - # if overwrite_from is None: # TODO - self.track(beam=None, savedepth=self.savedepth, runnable=self, verbose=self.verbose) + # 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 - # save object to file - self.save() + # 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-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-1: + 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 + 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 @@ -50,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): + 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 @@ -62,13 +130,10 @@ 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: - self.clear_run_data() - self.overwrite = False - + # default verbosity if self.verbose is None: self.verbose = not parallel @@ -101,6 +166,13 @@ 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)] + # 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 + self.perform_shot(shot=0) + self.is_shared_run = False + # perform shots (in parallel or series) if parallel: @@ -123,7 +195,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, shared_run_until=None): # define run name (generate if not given) if run_name is None: @@ -132,7 +204,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, shared_run_until=shared_run_until) @@ -181,7 +253,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: 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()