xcoll 0.5.12__py3-none-any.whl → 0.6.1__py3-none-any.whl

xcoll/scattering_routines/engine.py

@@ -1,122 +1,115 @@
 # copyright ############################### #
 # This file is part of the Xcoll Package.   #
-# Copyright (c) CERN, 2024.                 #
+# Copyright (c) CERN, 2025.                 #
 # ######################################### #
 
 import os
 import numpy as np
-import shutil
+from numbers import Number
+from functools import wraps
 
 import xobjects as xo
-import xpart as xp
 import xtrack as xt
+import xtrack.particles.pdg as pdg
+
 try:
     # TODO: once xaux is in Xsuite keep only this
-    from xaux import ClassProperty, ClassPropertyMeta, FsPath, singleton
-except ImportError:
-    from ..xaux import ClassProperty, ClassPropertyMeta, FsPath, singleton
-
+    from xaux import FsPath, ranID
+except (ImportError, ModuleNotFoundError):
+    from ..xaux import FsPath, ranID
 
-class BaseEngineMeta(xo.hybrid_class.MetaHybridClass, ClassPropertyMeta):
-    pass
 
-@singleton
-class BaseEngine(xo.HybridClass, metaclass=BaseEngineMeta):
+class BaseEngine(xo.HybridClass):
     _xofields = {
-        '_particle_ref': xp.Particles,
+        '_particle_ref': xt.Particles._XoStruct,
         '_seed':         xo.UInt64,
         '_capacity':     xo.Int64,
     }
 
     _int32 = False
-    _element_classes = None
     _only_protons = False
+    _element_classes = None
     _uses_input_file = False
+    _num_input_files = 0
     _uses_run_folder = False
 
     def __init__(self, **kwargs):
-        if not self._initialised:
-            if np.any([key[0] != '_' for key in self._xofields.keys()]):
-                raise ValueError(f"All fields in `{self.__class__.__name__}._xofields` have "
-                               + f"to start with an underscore! This is to ensure to work "
-                               + f"correctly with `ClassProperty`.")
-            if '_xobject' not in kwargs:
-                # Initialise defaults
-                self._cwd = None
-                self._line = None
-                self._verbose = False
-                self._input_file = None
-                self._element_dict = {}
-                self._warning_given = False
-                self._tracking_initialised = False
-                kwargs.setdefault('_particle_ref', xp.Particles())
-                kwargs.setdefault('_seed', 0)
-                kwargs.setdefault('_capacity', 0)
-            filtered_kwargs = {}
-            remaining_kwargs = {}
-            for key, value in kwargs.items():
-                if key in self._xofields.keys() or key == '_xobject':
-                    filtered_kwargs[key] = value
-                else:
-                    remaining_kwargs[key] = value
-            super().__init__(**filtered_kwargs)
-            kwargs = remaining_kwargs
-            self._initialised = True
-        # Apply kwargs
-        for kk, vv in kwargs.items():
-            if not hasattr(self.__class__, kk):
-                raise ValueError(f"Invalid attribute {kk} for {self.__class__.__name__}!")
-            setattr(self, kk, vv)
+        if self._element_classes is None:
+            raise NotImplementedError(f"{self.__class__.__name__} needs to define `_element_classes`!")
+        # Initialise defaults
+        self._cwd = None
+        self._line = None
+        self._verbose = False
+        self._input_file = None
+        self._element_dict = {}
+        self._warning_given = False
+        self._environment = None
+        self._element_index = 0
+        self._tracking_initialised = False
+        self._deactivated_elements = {}
+        kwargs.setdefault('_particle_ref', xt.Particles())
+        kwargs.setdefault('_seed', 0)
+        kwargs.setdefault('_capacity', 0)
+        super().__init__(**kwargs)
 
     def __del__(self, *args, **kwargs):
         self.stop(warn=False)
 
-
     def _warn(self, error=None):
         if not self._warning_given:
             print(f"Warning: Failed to import {self.__class__.__name__} environment "
-                + f" (did you compile?).\n{self.__class__.__name__.replace('Engine', '')} "
-                + f"elements will be installed but are not trackable.\n", flush=True)
+                + f"(did you compile?).\n{self.name.capitalize()} elements can be installed "
+                + f"but are not trackable.", flush=True)
             if error:
-                print(error, flush=True)
+                print(f"Error: {error}", flush=True)
             self._warning_given = True
+        self.stop()
+        if error:
+            raise error
+
+    def _print(self, *args, **kwargs):
+        if self.verbose:
+            kwargs.setdefault('flush', True)
+            print(*args, **kwargs)
 
 
     # ==================
     # === Properties ===
     # ==================
 
+    @property
+    def environment(self):
+        return self._environment
 
-    @ClassProperty
-    def name(cls):
-        return cls.__name__.replace('Engine', '').lower()
+    @property
+    def name(self):
+        return self.__class__.__name__.replace('Engine', '').lower()
 
-    @ClassProperty
-    def verbose(cls):
-        return cls.get_self()._verbose
+    @property
+    def verbose(self):
+        return self._verbose
 
     @verbose.setter
-    def verbose(cls, val):
-        cls.get_self()._verbose = val
+    def verbose(self, val):
+        self._verbose = val
 
-    @ClassProperty
-    def line(cls):
-        return cls.get_self()._line
+    @property
+    def line(self):
+        return self._line
 
     @line.setter
-    def line(cls, val):
+    def line(self, val):
         if not val is None and not isinstance(val, xt.Line):
             raise ValueError("`line` has to be an xt.Line object!")
-        cls.get_self()._line = val
+        self._line = val
 
     @line.deleter
-    def line(cls):
-        cls.get_self()._line = None
+    def line(self):
+        self._line = None
 
-    @ClassProperty
-    def particle_ref(cls):
-        self = cls.get_self()
-        initial = xp.Particles().to_dict()
+    @property
+    def particle_ref(self):
+        initial = xt.Particles().to_dict()
         current = self._particle_ref.to_dict()
         if xt.line._dicts_equal(initial, current):
             return None
@@ -124,173 +117,274 @@ class BaseEngine(xo.HybridClass, metaclass=BaseEngineMeta):
             return self._particle_ref
 
     @particle_ref.setter
-    def particle_ref(cls, val):
-        self = cls.get_self()
+    def particle_ref(self, val):
         if val is None:
-            self._particle_ref = xp.Particles()
+            self._particle_ref = xt.Particles()
         else:
-            if not isinstance(val, xp.Particles):
-                raise ValueError("`particle_ref` has to be an xp.Particles object!")
+            if not isinstance(val, xt.Particles):
+                raise ValueError("`particle_ref` has to be an xt.Particles object!")
             if val._capacity > 1:
                 raise ValueError("`particle_ref` has to be a single particle!")
-            if val.pdg_id[0] == 0:
-                if cls._only_protons:
-                    val.pdg_id[0] = xp.get_pdg_id_from_name('proton')
+            pdg_id = val.pdg_id[0]
+            if pdg_id == 0:
+                if self._only_protons:
+                    pdg_id = pdg.get_pdg_id_from_name('proton')
                 else:
-                    raise ValueError("`particle_ref` needs to have a valid pdg_id")
-            elif cls._only_protons and val.pdg_id[0] != xp.get_pdg_id_from_name('proton'):
-                raise ValueError("{cls.__name__} only supports protons!")
+                    raise ValueError(f"{self.__class__.__name__} allows the use of particles "
+                                   + f"different than protons. Hence, `particle_ref` "
+                                   + f"needs to have a valid pdg_id.")
+            elif self._only_protons and pdg_id != pdg.get_pdg_id_from_name('proton'):
+                raise ValueError("{self.__class__.__name__} only supports protons!")
             self._particle_ref = val
+            self._particle_ref.pdg_id[0] = pdg_id
 
     @particle_ref.deleter
-    def particle_ref(cls):
-        cls.get_self()._particle_ref = xp.Particles()
+    def particle_ref(self):
+        self.particle_ref = None
 
-    @ClassProperty
-    def capacity(cls):
-        self = cls.get_self()
+    @property
+    def capacity(self):
         if self._capacity == 0:
             return None
         else:
             return int(self._capacity)
 
     @capacity.setter
-    def capacity(cls, val):
+    def capacity(self, val):
         if val is None:
             val = 0
-        cls.get_self()._capacity = int(val)
+        if not isinstance(val, Number) or val < 0:
+            raise ValueError("`capacity` has to be a positive integer!")
+        self._capacity = int(val)
 
     @capacity.deleter
-    def capacity(cls):
-        raise ValueError("Not allowed.")
+    def capacity(self):
+        self.capacity = None
 
-    @ClassProperty
-    def seed(cls):
-        self = cls.get_self()
+    @property
+    def seed(self):
         if self._seed == 0:
             return None
         else:
             return self._seed
 
     @seed.setter
-    def seed(cls, val):
+    def seed(self, val):
         if val is None:
             val = 0
+        if not isinstance(val, Number) or val < 0:
+            raise ValueError("`seed` has to be a positive integer!")
         val = int(val)
-        if cls._int32:
+        if self._int32:
             new_val = np.uint32(val)
         else:
             new_val = np.uint64(val)
         if new_val != val:
-            print(f"Warning: type change for seed {val}. Using {new_val}.")
-        cls.get_self()._seed = new_val
+            self._print(f"Warning: type change for seed {val}. Using {new_val}.")
+        self._seed = new_val
 
     @seed.deleter
-    def seed(cls):
-        cls.get_self()._seed = 0
+    def seed(self):
+        self.seed = None
 
-    @ClassProperty
-    def input_file(cls):
-        return cls.get_self()._input_file
+    @property
+    def input_file(self):
+        if self._uses_input_file:
+            return self._input_file
+
+    @property
+    def element_dict(self):
+        return self._element_dict
 
 
     # ======================
     # === Public Methods ===
     # ======================
 
-
-    @classmethod
-    def start(cls, *, line=None, elements=None, names=None, cwd=None, seed=None,
-              particle_ref=None, input_file=None, **kwargs):
-        self = cls.get_self(**kwargs)
-        if self.is_running() is None:
-            raise NotImplementedError(f"Need to implement `is_running` for {cls.__name__}!")
-        elif self.is_running() is True:
-            if self.verbose:
-                print("Engine already running.", flush=True)
+    def start(self, *, clean=True, input_file=None, **kwargs):
+        if not self.environment:
+            raise RuntimeError(f"{self.name.capitalize()} environment not set up! "
+                             + f"Do not manually create an instance of the engine.")
+        if not self.environment.initialised:
+            raise RuntimeError(f"{self.name.capitalize()} environment not initialised! "
+                             + f"Please set all paths in the environment before "
+                             + f"starting the engine.")
+        if not self.environment.compiled:
+            raise RuntimeError(f"{self.name.capitalize()} interface not compiled! "
+                             + f"Please compile before starting the engine.")
+        if self.is_running():
+            self._print("Engine already running.")
             return
 
-        self._starting_engine = True  # We need this to allow changing the element settings which otherwise are locked
-        self._use_seed(seed)
-        self._use_line(line)
-        self._use_particle_ref(particle_ref)
-        self._sync_line_particle_ref()
-        self._get_elements(line=line, elements=elements, names=names)
-        self._set_cwd(cwd=cwd)
-        self._use_input_file(input_file=input_file, **kwargs)
-        self.clean_output_files()
-        self._starting_engine = False
-
-
-    @classmethod
-    def stop(cls, clean=False, **kwargs):
-        self = cls.get_self(**kwargs)
-        if hasattr(self, '_old_seed'):
-            self.seed = self._old_seed
-            del self._old_seed
-        if hasattr(self, '_old_line'):
-            self.line = self._old_line
-            del self._old_line
-        if hasattr(self, '_old_particle_ref'):
-            self.particle_ref = self._old_particle_ref
-            del self._old_particle_ref
-            self._sync_line_particle_ref()
-        if hasattr(self, '_old_cwd') and self._old_cwd is not None:
-            os.chdir(self._old_cwd)
-            del self._old_cwd
+        # Clean up any leftover failed runs
+        self.stop(clean=clean)
+
+        print(f"Starting {self.__class__.__name__}...   ", flush=True, end='')
+        kwargs = self._pre_start(**kwargs)
+
+        # This needs to be set in the ChildEngine, either in _start_engine() or at the start of tracking
+        self._tracking_initialised = False
+
+        # Set all engine properties that have a setter (this will remove these properties from the kwargs)
+        kwargs = self._set_engine_properties(**kwargs)
+
+        # Create input file if needed (this will remove the kwargs relevant to the input file and physics)
+        kwargs = self._use_input_file(input_file, **kwargs)
         if clean:
-            self.clean_output_files(clean_all=True)
-        self._cwd = None
-        self._input_file = None
-        self._element_dict = {}
+            self.clean_input_files(clean_all=False)
+        self._preparing_input = False
+
+        # Start the engine in the ChildEngine
+        self._start_engine(**kwargs)
+
+        # Done starting
+        if self.verbose:
+            print(f"{self.__class__.__name__} started.", flush=True)
+        else:
+            print(f"Done.", flush=True)
+
+    def stop(self, clean=False, **kwargs):
+        kwargs = self._stop_engine(**kwargs)
+        if clean:
+            self.clean(clean_all=True, **kwargs)
+        self._restore_engine_properties(clean=clean)
         self._warning_given = False
         self._tracking_initialised = False
 
 
-    @classmethod
-    def assert_particle_ref(cls, **kwargs):
-        if cls.get_self(**kwargs).particle_ref is None:
-            raise ValueError(f"{cls.__name__} reference particle not set!")
+    def assert_particle_ref(self):
+        if self.particle_ref is None:
+            raise ValueError(f"{self.__class__.__name__} reference particle not set!")
+
+
+    def generate_input_file(self, *, clean=True, filename=None, **kwargs):
+        # This method manually generates an input file without starting the engine
+        if not self._uses_input_file:
+            raise ValueError(f"{self.__class__.__name__} does not use input files!")
+        if self._element_dict:
+            raise ValueError("Elements already assigned to engine (cannot regenerate input "
+                           + "file after starting engine)!")
+
+        # Set all engine properties that have a setter (this will remove these properties from the kwargs)
+        kwargs = self._set_engine_properties(**kwargs)
+
+        # Create input file
+        input_file, _ = self._generate_input_file(**kwargs)
+        if not hasattr(input_file, '__iter__') or isinstance(input_file, str):
+            # Some engines might need multiple input files (like Fluka)
+            input_file = [input_file]
+        if filename is None:
+            if hasattr(self, '_old_cwd') and self._old_cwd is not None:
+                input_file[0] = input_file[0].rename(self._old_cwd / input_file[0].name)
+        else:
+            input_file[0] = input_file[0].rename(filename)
+        for i, file in enumerate(input_file[1:]):
+            input_file[i+1] = file.rename(input_file[0].parent / file.name)
+        input_file = input_file[0] if self._num_input_files==1 else input_file
+
+        if clean:
+            self.clean_input_files(clean_all=False)
+        self._restore_engine_properties(clean=clean)
+
+        return input_file
+
+
+    def is_running(self):
+        if hasattr(self, '_preparing_input') and self._preparing_input:
+            # We need this to allow changing the element settings which otherwise are locked
+            return False
+        # If we get here, we cannot say if the engine is running or not and we need an
+        # implementation in the child class
+        return self._is_running()
+
+
+    def clean(self, **kwargs):
+        self.clean_input_files(**kwargs)
+        self.clean_output_files(**kwargs)
+
+    def clean_input_files(self, **kwargs):
+        kwargs = self._get_input_cwd_for_cleaning(**kwargs)
+        self._clean_input_files(**kwargs)
+
+    def clean_output_files(self, **kwargs):
+        kwargs = self._get_input_cwd_for_cleaning(**kwargs)
+        self._clean_output_files(**kwargs)
 
 
     # =======================
     # === Private Methods ===
     # =======================
 
-    # For all the following fields, they can either be set in advance on the engine,
+    # For all the engine fields, they can either be set in advance on the engine,
     # or they can be set when the engine is started. In the latter case, the values
     # are temporary and the original will be restored when the engine is stopped.
 
-    def _use_line(self, line=None):
-        self._old_line = self.line
-        self.line = line
+    def _set_engine_properties(self, **kwargs):
+        self._preparing_input = True  # We need this to allow changing the element settings which otherwise are locked
+        def _set_property(prop):
+            val = kwargs.pop(prop, None)
+            if val is not None:
+                # We only need to update the property when it is not None
+                setattr(self, f'_old_{prop}', getattr(self, prop))
+                setattr(self, prop, val)
+        # We need to set the following properties first as they are needed by the others
+        _set_property('verbose')
+        _set_property('line')
+        # The following properties have a specific logic
+        self._use_seed(kwargs.pop('seed', None))
+        self._use_particle_ref(kwargs.pop('particle_ref', None))
+        self._sync_line_particle_ref()
+        self._get_elements(kwargs.pop('elements', None), kwargs.pop('names', None))
+        self._set_cwd(kwargs.pop('cwd', None))
+        # Now we can set the rest of the properties
+        _set_property('capacity')
+        return kwargs
+
+    def _restore_engine_properties(self, clean=False):
+        self._preparing_input = False
+        # Reset particle_ref in the line
+        if hasattr(self, '_old_particle_ref'):
+            self.line.particle_ref = self._old_line_particle_ref
+            del self._old_line_particle_ref
+        # The following properties have a specific logic
+        self._reactivate_elements()
+        self._reset_cwd(clean=clean)
+        # Reset all other properties
+        self_attributes = self.__dict__.copy()
+        for kk, vv in self_attributes.items():
+            if kk.startswith('_old_'):
+                prop = kk[5:]
+                setattr(self, prop, vv)
+                delattr(self, kk)
+        self._input_file = None
+        self._element_dict = {}
 
     def _use_seed(self, seed=None):
-        self._old_seed = self.seed
-        if seed is not None:
-            self.seed = seed
-        else:
+        if seed is None:
             if self.seed is None:
                 if self._int32:
                     self.seed = np.random.randint(0, int(2**32))
                 else:
                     self.seed = np.random.randint(0, int(2**64))
-        if self.verbose:
-            print(f"Using seed {self.seed}.")
+        else:
+            self._old_seed = self.seed
+            self.seed = seed
+        self._print(f"Using seed {self.seed}.")
 
     def _use_particle_ref(self, particle_ref=None):
         # Prefer: provided particle_ref > existing particle_ref > particle_ref from line
-        self._old_particle_ref = self.particle_ref
         if particle_ref is not None:
+            self._old_particle_ref = self.particle_ref
             self.particle_ref = particle_ref
         elif self.particle_ref is None:
             if self.line is None or not hasattr(self.line, 'particle_ref') \
             or self.line.particle_ref is None:
                 raise ValueError("Need to provide either a line with a reference "
                                + "particle, or `particle_ref`.")
+            self._old_particle_ref = self.particle_ref
             self.particle_ref = self.line.particle_ref
-        if self.verbose:
-            print(f"Using {xp.get_name_from_pdg_id(self.particle_ref.pdg_id[0])}.")
+        self._print(f"Using {pdg.get_name_from_pdg_id(self.particle_ref.pdg_id[0])} "
+                  + f"with momentum {self.particle_ref.p0c[0]/1.e9:.1f} GeV.")
 
     def _sync_line_particle_ref(self):
         if self.line is None:
@@ -298,68 +392,124 @@ class BaseEngine(xo.HybridClass, metaclass=BaseEngineMeta):
         if self.line.particle_ref is not None \
         and not xt.line._dicts_equal(self.line.particle_ref.to_dict(),
                                      self.particle_ref.to_dict()):
-            overwrite_particle_ref_in_line = True
-        if overwrite_particle_ref_in_line:
-            print("Warning: Found different reference particle in line! Temporarily overwritten.")
+            self._print("Found different reference particle in line. Temporarily overwritten.")
+            self._old_line_particle_ref = self.line.particle_ref
             self.line.particle_ref = self.particle_ref
 
-    def _get_elements(self, line=None, elements=None, names=None):
-        if self._element_classes is None:
-            raise NotImplementedError(f"{self.__class__.__name__} needs to define `_element_classes`!")
-        if line is None:
+    def _get_new_element_name(self):
+        name = f"{self.name}_el_{self._element_index}"
+        self._element_index += 1
+        return name
+
+    def _deactivate_element(self, el):
+        self._deactivated_elements[el.name] = [el, el.active or True]
+        if hasattr(el, 'active'):
+            el.active = False
+        self._remove_element(el)
+
+    def _assert_element(self, element):
+        if not isinstance(element, self._element_classes):
+            raise ValueError(f"Element {element.name} is not a "
+                            + ", or a ".join([c.__name__ for c in self._element_classes])
+                            + ".")
+
+    def _get_elements(self, elements=None, names=None):
+        if elements is not None and (not hasattr(elements, '__iter__') or isinstance(elements, str)):
+            elements = [elements]
+        if names is not None and (not hasattr(names, '__iter__') or isinstance(names, str)):
+            names = [names]
+        if self.line is None:
             if elements is None:
                 raise ValueError("Need to provide either `line` or `elements`.")
-            if not hasattr(elements, '__iter__') or isinstance(elements, str):
-                elements = [elements]
             if names is None:
-                names = [f"{self.__class__.name}_el_{i}" for i, _ in enumerate(elements)]
+                names = []
+                for ee in elements:
+                    if hasattr(ee, 'name'):
+                        names.append(ee.name)
+                    else:
+                        name = self._get_new_element_name()
+                        names.append(name)
+                        ee.name = name
+            elif len(names) == len(elements):
+                for ee, name in zip(elements, names):
+                    if hasattr(ee, 'name'):
+                        if ee.name != name:
+                            self._print(f"Warning: Element name {ee.name} changed to {name}.")
+                            ee.name = name
             else:
-                if not hasattr(names, '__iter__') or isinstance(names, str):
-                    names = [names]
-                if len(names) != len(elements):
-                    raise ValueError("Length of `elements` and `names` doesn't match.")
+                raise ValueError("Length of `elements` and `names` doesn't match.")
         else:
             if elements is not None:
                 raise ValueError("Cannot provide both `line` and `elements`.")
             if names is None:
-                elements, names = line.get_elements_of_type(self._element_classes)
+                elements, names = self.line.get_elements_of_type(self._element_classes)
             else:
-                if not hasattr(names, '__iter__') or isinstance(names, str):
-                    names = [names]
-                elements = [line[name] for name in names]
-        elements = [el for el in elements if el.active and el.jaw is not None]
-        names    = [name for el, name in zip(elements,names) if el.active and el.jaw is not None]
-        for el in elements:
-            if not isinstance(el, self._element_classes):
-                raise ValueError(f"Element {el} is not a "
-                                + ", or a ".join([c.__name__ for c in self._element_classes])
-                                + ".")
-        if len(elements) == 0:
+                elements = [self.line[name] for name in names]
+        this_names = []
+        this_elements = []
+        for ee, name in zip(elements, names):
+            self._assert_element(ee)
+            if ee.jaw is None:
+                self._print(f"Warning: Jaw not set for {name}. Ignoring.")
+                self._deactivate_element(ee)
+            elif not ee.active:
+                self._print(f"Warning: Element {name} is not active. Ignoring.")
+                self._deactivate_element(ee)
+            else:
+                this_names.append(name)
+                this_elements.append(ee)
+        if len(this_elements) == 0:
             raise ValueError(f"No active {self.name} elements found!")
-        self._element_dict = dict(zip(names, elements))
+        self._element_dict = dict(zip(this_names, this_elements))
+
+    def _reactivate_elements(self):
+        names = list(self._deactivated_elements.keys())
+        for name in names:
+            ee, was_active = self._deactivated_elements.pop(name)
+            self._restore_element(ee)
+            if hasattr(ee, 'active'):
+                ee.active = was_active
 
-    def _set_cwd(self, cwd):
+    def _set_cwd(self, cwd=None):
         if self._uses_run_folder:
             if cwd is not None:
                 cwd = FsPath(cwd).expanduser().resolve()
+                if cwd.exists() and cwd != FsPath.cwd():
+                    # Check if the folder already exists
+                    # If the specified folder is the current one, we do not need to rename
+                    i = 0
+                    while (cwd.parent / f'{cwd.name}_{i:0>4}').exists():
+                        i += 1
+                        if i > 9999:
+                            raise ValueError(f"Too many folders with the same "
+                                           + f"name {cwd}!")
+                    cwd = cwd.parent / f'{cwd.name}_{i:0>4}'
             else:
-                # TODO: use xaux.ranID here
-                import base64
-                ran = base64.urlsafe_b64encode(os.urandom(8)).decode('utf-8')
-                ran_str = ''.join(c if c.isalnum() else 'X' for c in ran)
+                ran_str = ranID(only_alphanumeric=True)
                 cwd = FsPath.cwd() / f'{self.name}_run_{ran_str}'
             self._cwd = cwd
             cwd.mkdir(parents=True, exist_ok=True)
             self._old_cwd = FsPath.cwd()
             os.chdir(cwd)
 
+    def _reset_cwd(self, clean=False):
+        if hasattr(self, '_old_cwd'):
+            if self._old_cwd is not None:
+                os.chdir(self._old_cwd)
+                if clean and self._cwd is not None and self._cwd != FsPath.cwd():
+                    try:
+                        self._cwd.rmdir()
+                    except:
+                        self._print(f"Warning: Failed to remove temporary folder "
+                                  + f"{self._cwd}.")
+            del self._old_cwd
+        self._cwd = None
+
+
     def _use_input_file(self, input_file=None, **kwargs):
         if self._uses_input_file:
             if input_file is None:
-                if not hasattr(self, 'generate_input_file'):
-                    raise NotImplementedError("Need to implement `generate_input_file`"
-                                              "for {cls.__name__}!")
-                input_file = self.generate_input_file(**kwargs)
+                input_file, kwargs = self._generate_input_file(**kwargs)
             if not hasattr(input_file, '__iter__') or isinstance(input_file, str):
                 # Some engines might need multiple input files (like Fluka)
                 input_file = [input_file]
@@ -368,36 +518,83 @@ class BaseEngine(xo.HybridClass, metaclass=BaseEngineMeta):
             for file in input_file:
                 if not file.exists():
                     raise ValueError(f"Input file {file.as_posix()} not found!")
-                if file.parent != FsPath.cwd():
-                    shutil.copy(file, FsPath.cwd())
+                if file.parent != FsPath.cwd() and self._uses_run_folder:
+                    file.copy_to(FsPath.cwd(), method='mount')
                     new_files.append(FsPath.cwd() / file.name)
                 else:
                     new_files.append(file)
-            self._input_file = new_files[0] if len(new_files)==1 else new_files
+            self._input_file = new_files[0] if self._num_input_files==1 else new_files
             self._match_input_file()
+        return kwargs
+
+
+    def _get_input_cwd_for_cleaning(self, **kwargs):
+        # Get the input file and the cwd
+        input_file = kwargs.get('input_file', None)
+        cwd = kwargs.get('cwd', None)
+        if input_file is None:
+            if self.input_file is not None:
+                input_file = self.input_file
+            if cwd is None and self._cwd is not None:
+                cwd = self._cwd
+        else:
+            if not hasattr(input_file, '__iter__') or isinstance(input_file, str):
+                input_file = [input_file]
+            input_file = [FsPath(ff) for ff in input_file]
+            if cwd is None:
+                cwd = input_file[0].parent
+            if len(input_file) == 1:
+                input_file = input_file[0]
+        kwargs['input_file'] = input_file
+        if self._uses_run_folder:
+            kwargs['cwd'] = cwd
+        else:
+            kwargs['cwd'] = FsPath.cwd()
+        return kwargs
 
 
-    # ===============================
-    # === Methods to be inherited ===
-    # ===============================
+    # =================================================
+    # === Methods to be overwritten by child engine ===
+    # =================================================
 
-    @classmethod
-    def is_running(cls, **kwargs):
-        self = cls.get_self(**kwargs)
-        if hasattr(self, '_starting_engine') and self._starting_engine:
-            # We need this to allow changing the element settings which otherwise are locked
-            return False
-        # If we get here, we cannot say if the engine is running or not and we need an
-        # implementation in the child class
-        return None
+    def _start_engine(self, **kwargs):
+        raise NotImplementedError(f"Need to implement `_start_engine` for "
+                                   + f"{self.__class__.__name__}!")
 
-    @classmethod
-    def clean_output_files(cls, **kwargs):
-        pass
+    def _stop_engine(self, **kwargs):
+        raise NotImplementedError(f"Need to implement `_stop_engine` for "
+                                   + f"{self.__class__.__name__}!")
+
+    def _is_running(self, **kwargs):
+        raise NotImplementedError(f"Need to implement `_is_running` for "
+                                   + f"{self.__class__.__name__}!")
 
     def _match_input_file(self, **kwargs):
-        raise NotImplementedError("Need to implement `_match_input_file` for {cls.__name__}!")
+        if self._uses_input_file:
+            raise NotImplementedError(f"Need to implement `_match_input_file` for "
+                                   + f"{self.__class__.__name__}!")
+
+    def _generate_input_file(self, **kwargs):
+        if self._uses_input_file:
+            raise NotImplementedError("Need to implement `_generate_input_file` for "
+                                   + f"{self.__class__.__name__}!")
+
+
+    # ============================================================
+    # === Methods to be optionally overwritten by child engine ===
+    # ============================================================
+
+    def _clean_input_files(self, input_file=None, cwd=None, clean_all=False, **kwargs):
+        pass
+
+    def _clean_output_files(self, input_file=None, cwd=None, clean_all=False, **kwargs):
+        pass
+
+    def _remove_element(self, el):
+        pass
+
+    def _restore_element(self, el):
+        pass
 
-    @classmethod
-    def generate_input_file(cls, **kwargs):
-        raise NotImplementedError("Need to implement `generate_input_file` for {cls.__name__}!")
+    def _pre_start(self, **kwargs):
+        return kwargs