From f481467e952a279b37722f29512240d00ae96241 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Mon, 8 Jun 2026 15:47:31 +0200 Subject: [PATCH 1/6] Add functionality for giving beta directly in Beam's normalized amplitude. --- src/abel/classes/beam.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/abel/classes/beam.py b/src/abel/classes/beam.py index 6fa748e4..b25e5900 100644 --- a/src/abel/classes/beam.py +++ b/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: From 75a16f6a364700952e6967cc95a33df3223ca902 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Mon, 8 Jun 2026 15:49:40 +0200 Subject: [PATCH 2/6] Add pre-alignment procedure for the plasma-lens interstage. --- .../interstage/plasma_lens/__init__.py | 71 ++++++++++++++++++- 1 file changed, 70 insertions(+), 1 deletion(-) diff --git a/src/abel/classes/interstage/plasma_lens/__init__.py b/src/abel/classes/interstage/plasma_lens/__init__.py index 1c26a52e..1517dd53 100644 --- a/src/abel/classes/interstage/plasma_lens/__init__.py +++ b/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('------------------------------------------------') - \ No newline at end of file + + + ## PRE-ALIGNMENT ROUTINE + + def pre_align(self, beam0): + + # 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 + + 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) + print('Building response matrix by tracking beams in offset lenses...') + 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 + response_matrix = np.eye(2) + response_matrix[0,0] = np.diff(x_1)/np.diff(dxs_lens) + response_matrix[0,1] = np.diff(x_2)/np.diff(dxs_lens) + response_matrix[1,0] = np.diff(xp_1)/np.diff(dxs_lens) + response_matrix[1,1] = np.diff(xp_2)/np.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] + beam_after = self.track(beam0) + + print(beam_before.x_offset()*1e6, 'µm ->', beam_after.x_offset()*1e6, 'µm') + print(beam_before.x_angle()*1e3, 'mrad ->', beam_after.x_angle()*1e3, 'mrad') + print('Ideal offset, lens 1 =', dx_lenses[0], 'm') + print(' lens 2 =', dx_lenses[1], 'm') + + + + + + \ No newline at end of file From 447cc8a3720ce6332360d208efa1ef4b28a3d583 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Tue, 9 Jun 2026 14:06:19 +0200 Subject: [PATCH 3/6] Improved the pre-alignment print messages. --- .../interstage/plasma_lens/__init__.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/src/abel/classes/interstage/plasma_lens/__init__.py b/src/abel/classes/interstage/plasma_lens/__init__.py index 1517dd53..2f715d5d 100644 --- a/src/abel/classes/interstage/plasma_lens/__init__.py +++ b/src/abel/classes/interstage/plasma_lens/__init__.py @@ -553,7 +553,7 @@ def print_summary(self): ## PRE-ALIGNMENT ROUTINE - def pre_align(self, beam0): + def pre_align(self, beam0, verbose=True): # pre-run to find offsets with zero lens offsets self.lens1_offset_x = 0 @@ -569,13 +569,16 @@ def pre_align(self, beam0): # 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) - print('Building response matrix by tracking beams in offset lenses...') for i, dx_lens in enumerate(dxs_lens): # offset first lens @@ -606,15 +609,14 @@ def pre_align(self, beam0): # pre-run to find offsets with zero lens offsets self.lens1_offset_x = dx_lenses[0] self.lens2_offset_x = dx_lenses[1] - beam_after = self.track(beam0) - - print(beam_before.x_offset()*1e6, 'µm ->', beam_after.x_offset()*1e6, 'µm') - print(beam_before.x_angle()*1e3, 'mrad ->', beam_after.x_angle()*1e3, 'mrad') - print('Ideal offset, lens 1 =', dx_lenses[0], 'm') - print(' lens 2 =', dx_lenses[1], 'm') - - + + # 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 \ No newline at end of file From 78903e7eada7a5b515a907e0ef88372e157f7d4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Tue, 9 Jun 2026 14:12:15 +0200 Subject: [PATCH 4/6] Added a test of the pre-alignment procedure. --- tests/test_impactx.py | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/test_impactx.py b/tests/test_impactx.py index ce6c1d98..d1114c02 100644 --- a/tests/test_impactx.py +++ b/tests/test_impactx.py @@ -152,6 +152,41 @@ 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 + + # 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-7) + assert not np.isclose(abs(beam_before.x_angle()), 0.0, atol=1e-6) + # 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-7) + assert np.isclose(abs(beam_after.x_angle()), 0.0, atol=1e-6) \ No newline at end of file From aa91df8bc39c860ec3f9c5cacf5ab8ae7a3f23f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Tue, 9 Jun 2026 14:31:57 +0200 Subject: [PATCH 5/6] Updated the test. --- tests/test_impactx.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/test_impactx.py b/tests/test_impactx.py index d1114c02..45755186 100644 --- a/tests/test_impactx.py +++ b/tests/test_impactx.py @@ -171,22 +171,24 @@ def test_interstage_plasma_lens_prealign(): 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-7) - assert not np.isclose(abs(beam_before.x_angle()), 0.0, atol=1e-6) + + 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) - assert np.isclose(abs(beam_after.x_offset()), 0.0, atol=1e-7) - assert np.isclose(abs(beam_after.x_angle()), 0.0, atol=1e-6) - \ No newline at end of file From b1628ac431a776773089428b1f25d05af8efefda Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carl=20A=2E=20Lindstr=C3=B8m?= Date: Tue, 9 Jun 2026 14:55:11 +0200 Subject: [PATCH 6/6] Fixed bug in response matrix (the np.diff made a vector instead of a scalar). --- src/abel/classes/interstage/plasma_lens/__init__.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/abel/classes/interstage/plasma_lens/__init__.py b/src/abel/classes/interstage/plasma_lens/__init__.py index 2f715d5d..bcc02769 100644 --- a/src/abel/classes/interstage/plasma_lens/__init__.py +++ b/src/abel/classes/interstage/plasma_lens/__init__.py @@ -596,11 +596,9 @@ def pre_align(self, beam0, verbose=True): xp_2[i] = beam2.x_angle() # build response matrix - response_matrix = np.eye(2) - response_matrix[0,0] = np.diff(x_1)/np.diff(dxs_lens) - response_matrix[0,1] = np.diff(x_2)/np.diff(dxs_lens) - response_matrix[1,0] = np.diff(xp_1)/np.diff(dxs_lens) - response_matrix[1,1] = np.diff(xp_2)/np.diff(dxs_lens) + 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)