[Draft] Add functionality to the Runnable

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)

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

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):

src/abel/classes/runnable.py

@@ -24,33 +24,101 @@ 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
 
     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:

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()