Improved the interstage pre-alignment procedure

src/abel/classes/beamline/impl/linac/impl/plasma_linac.py

@@ -978,6 +978,7 @@ def frameFcn(i):
             dQdxdpx, xs, pxs = beam.phase_space_density(beam.xs, beam.pxs, hbins=xs0, vbins=pxs_final)
             axs[1,0].cla()
             cax = axs[1,0].pcolor(xs*1e6, pxs*1e-6*SI.c/SI.e, -dQdxdpx, cmap=CONFIG.default_cmap, shading='auto')
+            axs[1,0].set_ylim([-pxlim*1e-6*SI.c/SI.e, pxlim*1e-6*SI.c/SI.e])
             axs[1,0].set_ylabel("Momentum, $p_x$ [MeV/c]")
             axs[1,0].set_title('Horizontal phase space')
 

src/abel/classes/interstage/plasma_lens/__init__.py

@@ -9,6 +9,7 @@
 from types import SimpleNamespace
 import numpy as np
 import scipy.constants as SI
+import copy
 from abel.utilities.relativity import energy2gamma
 
 class InterstagePlasmaLens(Interstage, ABC):
@@ -394,7 +395,7 @@ def minfun_beta(params):
 
         # match the beta function
         from scipy.optimize import minimize
-        result_beta = minimize(minfun_beta, k_lens0, tol=1e-16, options={'maxiter': 200})
+        result_beta = minimize(minfun_beta, k_lens0, tol=1e-16, options={'maxiter': 1000})
         self._strength_plasma_lens = result_beta.x[0]*self.length_plasma_lens
 
 
@@ -436,7 +437,7 @@ def minfun_dispersion_R56(params):
 
         # match the beta function
         from scipy.optimize import minimize
-        result_dispersion_R56 = minimize(minfun_dispersion_R56, [B_chic1_guess, B_chic2_guess], tol=1e-8, options={'maxiter': 50})
+        result_dispersion_R56 = minimize(minfun_dispersion_R56, [B_chic1_guess, B_chic2_guess], tol=1e-16, options={'maxiter': 500})
         self._field_ratio_chicane_dipole1 = result_dispersion_R56.x[0]/self.field_dipole
         self._field_ratio_chicane_dipole2 = result_dispersion_R56.x[1]/self.field_dipole
 
@@ -477,7 +478,7 @@ def minfun_W(params):
 
         # match the beta function
         from scipy.optimize import minimize
-        result_W = minimize(minfun_W, tau_lens0, tol=1e-16, options={'maxiter': 100})
+        result_W = minimize(minfun_W, tau_lens0, tol=1e-16, options={'maxiter': 1000})
         self._nonlinearity_plasma_lens = result_W.x[0]
 
 
@@ -514,7 +515,7 @@ def minfun_second_order_dispersion(params):
 
         # match the beta function
         from scipy.optimize import minimize
-        result_dispersion = minimize(minfun_second_order_dispersion, m_sext0, method='Nelder-Mead', tol=1e-20, options={'maxiter': 50})
+        result_dispersion = minimize(minfun_second_order_dispersion, m_sext0, method='Nelder-Mead', tol=1e-16, options={'maxiter': 500})
         self._strength_sextupole = result_dispersion.x[0]*self.length_central_gap_or_sextupole
 
 
@@ -553,52 +554,70 @@ def print_summary(self):
 
     ## PRE-ALIGNMENT ROUTINE
 
-    def pre_align(self, beam0, verbose=True):
+    def pre_align(self, source=None, dx=None, verbose=False, parallel=True):
 
+        # copy and adjust the source
+        source_copy = copy.deepcopy(source)
+        source_copy.num_particles = 10000
+        source_copy.symmetrize = True
+        source_copy.x_offset = 0
+        source_copy.y_offset = 0
+        source_copy.x_angle = 0
+        source_copy.y_angle = 0
+        source_copy.energy = self.nom_energy
+        source_copy.beta_x = self.beta0
+        source_copy.beta_y = self.beta0
+
+        # make the beam from the source
+        beam0 = source_copy.track()
+
         # pre-run to find offsets with zero lens offsets
         self.lens1_offset_x = 0
         self.lens2_offset_x = 0
         self.lens1_offset_y = 0
         self.lens2_offset_y = 0
-        deltaE = self.nom_energy-beam0.energy()
-        beam0.accelerate(deltaE)
-        beam0.apply_betatron_damping(deltaE)
-        beam_before = self.track(beam0)
-        X_beam = np.array([beam_before.x_offset(), beam_before.x_angle()])
-
-        # calculate characteristic scale of offsets from normalized amplitude
-        A0 = beam_before.norm_amplitude_x(beta0=self.beta0)
-        offset_scale = np.sqrt(A0**2*self.beta0/energy2gamma(self.nom_energy))/2
-
+
         # print info
         if verbose:
             print('Interstage pre-alignment: building response matrix...')
-
-        dxs_lens = np.array([-offset_scale, offset_scale])
-        x_1 = np.zeros_like(dxs_lens)
-        xp_1 = np.zeros_like(dxs_lens)
-        x_2 = np.zeros_like(dxs_lens)
-        xp_2 = np.zeros_like(dxs_lens)
-        for i, dx_lens in enumerate(dxs_lens):
-
-            # offset first lens
-            self.lens1_offset_x = dx_lens
-            self.lens2_offset_x = 0
-            beam1 = self.track(beam0)
-            x_1[i] = beam1.x_offset()
-            xp_1[i] = beam1.x_angle()
-
-            # offset second lens
-            self.lens1_offset_x = 0
-            self.lens2_offset_x = dx_lens
-            beam2 = self.track(beam0)
-            x_2[i] = beam2.x_offset()
-            xp_2[i] = beam2.x_angle()
 
+        # pre-calculate the optical setup
+        self.matrix_lattice()
+
+        # calculate characteristic scale of offsets from normalized amplitude
+        if dx is None:
+            beam_before = self.track(beam0)
+            A0 = beam_before.norm_amplitude_x(beta0=self.beta0)
+            dx = np.sqrt(A0**2*self.beta0/energy2gamma(self.nom_energy))/2
+
+        # define offsets to be simulated
+        dxs_lens = [[-dx, 0], [dx, 0], [0, -dx], [0, dx], [0, 0]]
+
+        # define offset simulation function
+        def offset_simulation(dx_lens):
+            self_copy = copy.deepcopy(self)
+            self_copy.lens1_offset_x = dx_lens[0]
+            self_copy.lens2_offset_x = dx_lens[1]
+            beam = self_copy.track(beam0)
+            return np.array([beam.x_offset(), beam.x_angle()])
+
+        # perform simulation in parallel or not
+        if parallel:
+            from joblib import Parallel, delayed
+            Xs = np.array(Parallel(n_jobs=5)(delayed(offset_simulation)(dx_lens) for dx_lens in dxs_lens))
+
+        else:
+            Xs = np.zeros_like(dxs_lens)
+            for i, dx_lens in enumerate(dxs_lens):
+                Xs[i,:] = offset_simulation(dx_lens)
+
+        # extract the zero-offset beam
+        X_beam = np.array([Xs[4,0], Xs[4,1]])
+
         # build response matrix
-        diff_dxs_lens = np.diff(dxs_lens)[0]
-        response_matrix = np.array([[np.diff(x_1)[0]/diff_dxs_lens, np.diff(x_2)[0]/diff_dxs_lens],
-                                    [np.diff(xp_1)[0]/diff_dxs_lens, np.diff(xp_2)[0]/diff_dxs_lens]])
+        diff_dx = 2*dx
+        response_matrix = np.array([[Xs[1,0]-Xs[0,0], Xs[3,0]-Xs[2,0]],
+                                    [Xs[1,1]-Xs[0,1], Xs[3,1]-Xs[2,1]]])/(2*dx)
 
         # invert the response matrix to find the optimal lens offsets
         inv_response_matrix = np.linalg.inv(response_matrix)
@@ -611,7 +630,7 @@ def pre_align(self, beam0, verbose=True):
         # print info
         if verbose:
             beam_after = self.track(beam0)
-            print(f'>> Before: x = {beam_before.x_offset()*1e6:.2f} µm, {beam_before.x_angle()*1e3:.2f} mrad  ->  After: x = {beam_after.x_offset()*1e6:.2f} µm, {beam_after.x_angle()*1e3:.2f} mrad')
+            print(f'>> Before: x = {X_beam[0]*1e6:.2f} µm, {X_beam[1]*1e6:.2f} µrad  ->  After: x = {beam_after.x_offset()*1e6:.2f} µm, {beam_after.x_angle()*1e6:.2f} µrad')
             print(f'>> Ideal offset: lens 1 = {dx_lenses[0]*1e6:.2f} µm, lens 2 = {dx_lenses[1]*1e6:.2f} µm')
 
         return dx_lenses

src/abel/classes/trackable.py

@@ -209,7 +209,7 @@ def track(self, beam : Beam, savedepth : int = 0, runnable : Runnable | None = N
                     stage_word = ''
 
                 # print string
-                print(f"{col_lgray}{tracking_word} #{str(beam.trackable_number).ljust(2)}{col_reset} {bold}{(type(self).__name__+stage_word).ljust(23)}{reset_bold} {col_lgray}(s ={col_reset} {bold}{beam.location:6.1f}{reset_bold} m{col_lgray}){col_reset} :   {col_lgray}E ={col_reset}{bold}{beam.energy()/1e9:6.1f}{reset_bold} GeV{col_lgray}, Q ={col_reset}{bold}{beam.charge()*1e9:6.2f}{reset_bold} nC{col_lgray}, σz ={col_reset} {bold}{beam.bunch_length()/1e-6:5.1f}{reset_bold} µm{col_lgray}, σE ={col_reset}{bold}{beam.rel_energy_spread()/1e-2:5.1f}%{reset_bold}{col_lgray}, ε ={col_reset}{bold}{beam.norm_emittance_x()/1e-6:6.1f}{reset_bold}/{bold}{beam.norm_emittance_y()/1e-6:.1f}{reset_bold} mm-mrad{reset}")
+                print(f"{col_lgray}{tracking_word} #{str(beam.trackable_number).ljust(2)}{col_reset} {bold}{(type(self).__name__+stage_word).ljust(27)}{reset_bold} {col_lgray}(s ={col_reset} {bold}{beam.location:6.1f}{reset_bold} m{col_lgray}){col_reset} :   {col_lgray}E ={col_reset}{bold}{beam.energy()/1e9:6.1f}{reset_bold} GeV{col_lgray}, Q ={col_reset}{bold}{beam.charge()*1e9:6.2f}{reset_bold} nC{col_lgray}, σz ={col_reset} {bold}{beam.bunch_length()/1e-6:5.1f}{reset_bold} µm{col_lgray}, σE ={col_reset}{bold}{beam.rel_energy_spread()/1e-2:5.1f}%{reset_bold}{col_lgray}, ε ={col_reset}{bold}{beam.norm_emittance_x()/1e-6:6.1f}{reset_bold}/{bold}{beam.norm_emittance_y()/1e-6:.2f}{reset_bold} mm-mrad{reset}")
 
             # save to file
             if runnable is not None:

src/abel/wrappers/impactx/impactx_wrapper.py

@@ -90,10 +90,10 @@ def run_impactx(lattice, beam0, nom_energy=None, runnable=None, keep_data=False,
 
     # create a new directory
     original_folder = os.getcwd()
-    if runnable is not None:
-        runfolder = CONFIG.temp_path + str(uuid.uuid4())
-    else:
-        runfolder = 'impactx_sims'
+    #if runnable is not None:
+    runfolder = os.path.join(CONFIG.temp_path, str(uuid.uuid4()))
+    #else:
+    #    runfolder = 'impactx_sims'
     if not os.path.exists(runfolder):
         os.makedirs(runfolder)
     os.chdir(runfolder)
@@ -135,7 +135,7 @@ def run_impactx(lattice, beam0, nom_energy=None, runnable=None, keep_data=False,
     sim.finalize()
 
     # extract evolution
-    evol = extract_evolution();
+    evol = extract_evolution()
 
     # copy meta data from input beam (will be iterated by super)
     beam.trackable_number = beam0.trackable_number
@@ -210,10 +210,11 @@ def run_envelope_impactx(lattice, distr, nom_energy=None, peak_current=None, spa
 
     # create a new directory
     original_folder = os.getcwd()
-    if runnable is not None:
-        runfolder = CONFIG.temp_path + str(uuid.uuid4())
-    else:
-        runfolder = 'impactx_sims'
+    #if runnable is not None:
+    runfolder = os.path.join(CONFIG.temp_path, str(uuid.uuid4()))
+    print(runfolder)
+    #else:
+    #    runfolder = 'impactx_sims'
     if not os.path.exists(runfolder):
         os.makedirs(runfolder)
     os.chdir(runfolder)
@@ -322,10 +323,7 @@ def extract_beams(path='', runnable=None, beam0=None):
 
     from abel.classes.beam import Beam
     import openpmd_api as io
-
-    if runnable is not None:
-        path = os.path.join(runnable.shot_path(), 'impactx_sims') # TODO: even if path is provided, this will overwrite it.
-
+
     # load OpenPMD series
     series = io.Series(os.path.join(path,'diags/openPMD/monitor.h5'), io.Access.read_only)
     steps = list(series.iterations)
@@ -408,11 +406,11 @@ def extract_evolution(path=''):
 
     # read CSV file
     try:
-        ref = pd.read_csv(path+"diags/ref_particle.0.0", delimiter=r"\s+")
-        diags = pd.read_csv(path+"diags/reduced_beam_characteristics.0.0", delimiter=r"\s+")
+        ref = pd.read_csv(os.path.join(path, 'diags/ref_particle.0.0'), delimiter=r"\s+")
+        diags = pd.read_csv(os.path.join(path, 'diags/reduced_beam_characteristics.0.0'), delimiter=r"\s+")
     except:
-        ref = pd.read_csv(path+"diags/ref_particle.0", delimiter=r"\s+")
-        diags = pd.read_csv(path+"diags/reduced_beam_characteristics.0", delimiter=r"\s+")
+        ref = pd.read_csv(os.path.join(path, 'diags/ref_particle.0'), delimiter=r"\s+")
+        diags = pd.read_csv(os.path.join(path, 'diags/reduced_beam_characteristics.0'), delimiter=r"\s+")
 
     # extract numbers
     evol = SimpleNamespace()