Interstage pre-alignment

src/abel/classes/beam.py

@@ -1038,8 +1038,11 @@ def eigen_emittance_max(self):
     def eigen_emittance_min(self):
         return np.sqrt(self.norm_emittance_x()*self.norm_emittance_y()) - self.angular_momentum()
 
-    def norm_amplitude_x(self, plasma_density=None, clean=False):
-        if plasma_density is not None:
+    def norm_amplitude_x(self, beta0=None, plasma_density=None, clean=False):
+        if beta0 is not None:
+            beta_x = beta0
+            alpha_x = 0
+        elif plasma_density is not None:
             beta_x = beta_matched(plasma_density, self.energy())
             alpha_x = 0
         else:
@@ -1049,8 +1052,11 @@ def norm_amplitude_x(self, plasma_density=None, clean=False):
             alpha_x = -covx[1,0]/emgx
         return np.sqrt(self.gamma()/beta_x)*np.sqrt(self.x_offset()**2 + (self.x_offset()*alpha_x + self.x_angle()*beta_x)**2)
 
-    def norm_amplitude_y(self, plasma_density=None, clean=False):
-        if plasma_density is not None:
+    def norm_amplitude_y(self, beta0=None, plasma_density=None, clean=False):
+        if beta0 is not None:
+            beta_y = beta0
+            alpha_y = 0
+        elif plasma_density is not None:
             beta_y = beta_matched(plasma_density, self.energy())
             alpha_y = 0
         else:

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
+from abel.utilities.relativity import energy2gamma
 
 class InterstagePlasmaLens(Interstage, ABC):
     """
@@ -548,4 +549,72 @@ def print_summary(self):
         print(f'         Total bend angle:           {np.rad2deg(self.total_bend_angle()):.2f} deg')
         print('------------------------------------------------')
 
-
+
+
+    ## PRE-ALIGNMENT ROUTINE
+
+    def pre_align(self, beam0, verbose=True):
+
+        # 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()
+
+        # 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]])
+
+        # invert the response matrix to find the optimal lens offsets
+        inv_response_matrix = np.linalg.inv(response_matrix)
+        dx_lenses = np.matmul(inv_response_matrix, -X_beam)
+
+        # pre-run to find offsets with zero lens offsets
+        self.lens1_offset_x = dx_lenses[0]
+        self.lens2_offset_x = dx_lenses[1]
+
+        # 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'>> Ideal offset: lens 1 = {dx_lenses[0]*1e6:.2f} µm, lens 2 = {dx_lenses[1]*1e6:.2f} µm')
+
+        return dx_lenses
+
+
+

tests/test_impactx.py

@@ -152,6 +152,43 @@ def test_interstage_plasma_lens_impactx_spin_tracking():
     assert np.isclose(beam_pl_z.spin_polarization_y(), 0, atol=0.02)
     assert np.isclose(beam_pl_z.spin_polarization_z(), 0.42291824, atol=0.02)
 
+
+
+@pytest.mark.impactx
+def test_interstage_plasma_lens_prealign():
+    """Test the pre-alignment procedure of the plasma-lens-based interstage."""
 
+    # define electron source
+    source = make_source()
 
-
+    # define interstage (plasma-lens-based)
+    interstage_pl = InterstagePlasmaLensImpactX()
+    interstage_pl.nom_energy = source.energy
+    interstage_pl.R56 = 0.0
+    interstage_pl.beta0 = source.beta_x
+    interstage_pl.length_dipole = 1.0
+    interstage_pl.field_dipole = 1.0
+    interstage_pl.cancel_chromaticity = True
+    interstage_pl.cancel_sec_order_dispersion = True
+    interstage_pl.use_apertures = False
+    interstage_pl.enable_csr = True
+    interstage_pl.enable_isr = True
+
+    # make a beam
+    beam0 = source.track()
+
+    # track through the pre-aligned interstage
+    beam_before = interstage_pl.track(beam0, verbose=False)
+
+    assert not np.isclose(abs(beam_before.x_offset()), 0.0, atol=1e-8)
+    assert not np.isclose(abs(beam_before.x_angle()), 0.0, atol=1e-7)
+
+    # do the pre-alignment
+    interstage_pl.pre_align(beam0)
+
+    # track through the pre-aligned interstage
+    beam_after = interstage_pl.track(beam0, verbose=False)
+
+    assert np.isclose(abs(beam_after.x_offset()), 0.0, atol=1e-8)
+    assert np.isclose(abs(beam_after.x_angle()), 0.0, atol=1e-7)
+