Fix ImpactX Plasma Lens

src/abel/__init__.py

@@ -62,6 +62,7 @@
 from .classes.interstage.quads.basic import InterstageQuadsBasic
 from .classes.interstage.quads.impactx import InterstageQuadsImpactX
 from .classes.plasma_lens.impl.plasma_lens_basic import PlasmaLensBasic
+from .classes.plasma_lens.impl.plasma_lens_impactx import PlasmaLensImpactX
 from .classes.plasma_lens.impl.plasma_lens_nonlinear_thin import PlasmaLensNonlinearThin
 from .classes.plasma_lens.impl.plasma_lens_nonlinear_thick import PlasmaLensNonlinearThick
 from .classes.rf_accelerator.impl.rf_accelerator_basic import RFAcceleratorBasic

src/abel/classes/plasma_lens/impl/plasma_lens_impactx.py

@@ -5,85 +5,109 @@
 # License: GPL-3.0-or-later
 
 from abel.classes.plasma_lens.plasma_lens import PlasmaLens
+from abel.utilities.relativity import energy2momentum
 import numpy as np
 import scipy.constants as SI
 
+
 class PlasmaLensImpactX(PlasmaLens):
-
-    def __init__(self, length=None, radius=None, current=None, rel_nonlinearity=0):
+    """
+    Active plasma lens tracked with ImpactX.
+
+    The lens is modelled as a drift-kick sequence of thin, transversely tapered
+    plasma-lens elements (``impactx.elements.TaperedPL``), following the same
+    construction as :class:`abel.InterstagePlasmaLensImpactX`.
+
+    Parameters
+    ----------
+    length : [m] float
+        Length of the plasma lens.
+
+    radius : [m] float
+        Radius of the plasma-lens capillary.
+
+    current : [A] float
+        Current through the plasma lens.
+
+    rel_nonlinearity : float, optional
+        Relative nonlinearity of the transverse field profile, defined as
+        ``radius / Dx`` where ``Dx`` is the targeted horizontal dispersion.
+        Defaults to 0 (a purely linear lens).
+
+    num_slices : int, optional
+        Number of thin-lens kicks used to represent the thick lens. Defaults
+        to 30.
+
+    use_apertures : bool, optional
+        If ``True``, an elliptical aperture of half-width :attr:`radius` is
+        applied at each end of the lens, so that charge outside the capillary
+        is removed. Defaults to ``False``, which reproduces the behaviour of
+        the other ``PlasmaLens`` implementations' unclipped field expansion.
+
+    Notes
+    -----
+    ``TaperedPL`` applies a polynomial expansion of the plasma-lens field. It
+    is only physical inside the capillary; particles outside :attr:`radius`
+    receive an unphysical kick unless ``use_apertures=True``.
+    """
+
+    def __init__(self, length=None, radius=None, current=None, rel_nonlinearity=0, num_slices=30, use_apertures=False):
 
         super().__init__(length, radius, current)
-        
+
         # set nonlinearity (defined as R/Dx)
         self.rel_nonlinearity = rel_nonlinearity
-
-
+
+        # simulation options
+        self.num_slices = num_slices
+        self.use_apertures = use_apertures
+
+
+    # ==================================================
     def track(self, beam0, savedepth=0, runnable=None, verbose=False):
+        "Track the plasma lens using ImpactX."
+
+        # get the lattice
+        lattice = self.get_impactx_lattice(beam0)
+
+        # run ImpactX
+        from abel.wrappers.impactx.impactx_wrapper import run_impactx
+        beam, self.evolution = run_impactx(lattice, beam0, nom_energy=beam0.energy(), verbose=verbose, runnable=runnable)
+
+        return super().track(beam, savedepth, runnable, verbose)
+
+
+    # ==================================================
+    def get_impactx_lattice(self, beam0):
+        "Set up the ImpactX plasma-lens lattice."
+
+        from impactx import elements
+
+        # integrated focusing strength [1/m], signed by the beam charge
+        # (k = L * g / (magnetic rigidity), see the ImpactX TaperedPL docs)
+        strength = beam0.charge_sign() * self.get_focusing_gradient() * self.length * SI.e / energy2momentum(beam0.energy())
+
+        # horizontal taper parameter [1/m], i.e. the inverse target dispersion
+        taper = self.rel_nonlinearity / self.radius
+
+        # drift-kick sequence: num_slices kicks separated by num_slices+1 drifts
+        ds = self.length / (self.num_slices + 1)
+        drift = elements.ExactDrift(ds=ds, nslice=1)
+
+        lattice = [drift]
+        for _ in range(self.num_slices):
+            lattice.append(elements.TaperedPL(k=strength/self.num_slices, taper=taper, unit=0, dx=self.offset_x, dy=self.offset_y))
+            lattice.append(drift)
+
+        # clip charge outside the capillary
+        if self.use_apertures:
+            aperture = elements.Aperture(aperture_x=self.radius, aperture_y=self.radius, shape='elliptical')
+            lattice = [aperture] + lattice + [aperture]
 
-        import impactx
-        from abel.wrappers.impactx.impactx_wrapper import beam2particle_container, particle_container2beam
-
-        # initialize AMReX
-        verbose_debug = False
-
-        # make simulation object
-        sim = impactx.ImpactX()
-
-        # serial run on one CPU core
-        sim.omp_threads = 1
-
-        # set ImpactX verbosity
-        sim.verbose = int(verbose_debug)
-        sim.tiny_profiler = verbose_debug
-
-        # convert to ImpactX particle container
-        _, sim = beam2particle_container(beam0, sim)
-
-        # add beam diagnostics
-        monitor = impactx.elements.BeamMonitor("monitor", backend="h5")
-
-        # TODO: include the interstage optic
-        # TODO: print the evolution (i.e., put the values into the evolution namespace, and add plotting functions in the base class)
-
-        # design the accelerator lattice
-        ns = 25  # number of slices per ds in the element
-
-        # specify thick tapered plasma lens element
-        num_cuts = 10
-        k0 = -self.get_focusing_gradient() * SI.c / beam0.energy()
-        dtaper = self.rel_nonlinearity / self.radius  # 1/(horizontal dispersion in m)
-        ds = self.length / num_cuts
-        dk = k0 / num_cuts
-        pl = impactx.elements.TaperedPL(k=dk, taper=dtaper, units=0)
-
-        # drifts appearing the drift-kick sequence
-        drift = impactx.elements.Drift(ds=ds/2, nslice=ns)
-
-        # define the lens segments
-        thick_lens = []
-        for _ in range(0, num_cuts):
-            thick_lens.extend([drift, pl, drift])
-
-        # assign the lattice
-        sim.lattice.extend(thick_lens)
-
-        # run simulation
-        sim.evolve()
-
-        # convert back to ABEL beam
-        beam = particle_container2beam(sim.particle_container())
-
-        # clean shutdown
-        sim.finalize()
-
-        # copy meta data from input beam (will be iterated by super)
-        beam.trackable_number = beam0.trackable_number
-        beam.stage_number = beam0.stage_number
-        beam.location = beam0.location
-
-        return super().track(beam, savedepth, runnable, verbose)  
+        return lattice
 
 
+    # ==================================================
     def get_focusing_gradient(self):
+        "Plasma-lens field gradient [T/m]."
         return SI.mu_0 * self.current / (2*np.pi * self.radius**2)
-

tests/test_plasma_lens.py

@@ -139,3 +139,102 @@ def test_PlasmaLensNonlinearThick():
 
 
 
+@pytest.mark.plasma_lens
+@pytest.mark.impactx
+def test_PlasmaLensImpactX():
+    """
+    Check that the ImpactX plasma lens agrees with the thick nonlinear lens.
+
+    Both model the same drift-kick sequence, so they should agree to well
+    within the difference between their slicing schemes.
+    """
+
+    np.random.seed(42)
+
+    # set up beam
+    source = SourceBasic()
+    source.bunch_length = 100e-6 # [m]
+    source.num_particles = 20000
+    source.charge = -SI.e * 1.0e10 # [C]
+    source.energy = 1e9 # [eV]
+    source.rel_energy_spread = 1e-5
+    source.emit_nx, source.emit_ny = 1e-6, 1e-6 # [m rad]
+    source.beta_x = 0.01 # [m]
+    source.beta_y = source.beta_x
+
+    # drift distance
+    L_drift = 1.0 # [m]
+
+    # lens length and radius
+    L_pl = 0.01
+    R_pl = 500e-6
+
+    # calculate strength required to refocus in distance L
+    f = (L_drift+L_pl/2)/2
+    k = 1/(L_pl*f)
+    g = k*source.energy/SI.c
+    I = g*(2*np.pi*R_pl**2)/SI.mu_0
+
+    beam0 = source.track()
+
+    def track(plasma_lens):
+        plasma_lens.length = L_pl
+        plasma_lens.radius = R_pl
+        plasma_lens.current = I
+        plasma_lens.rel_nonlinearity = 0.5
+        plasma_lens.offset_x = -R_pl/4
+        plasma_lens.offset_y = R_pl/8
+        beam = copy.deepcopy(beam0)
+        beam.transport(L_pl)
+        beam = plasma_lens.track(beam)
+        beam.transport(L_pl)
+        return beam
+
+    beam_thick = track(PlasmaLensNonlinearThick(num_slice=30))
+    beam_impactx = track(PlasmaLensImpactX(num_slices=30))
+
+    # charge is conserved (no apertures by default)
+    assert np.isclose(beam_impactx.charge(), beam0.charge(), rtol=1e-15)
+    assert np.isclose(beam_impactx.energy(), beam0.energy(), rtol=1e-6)
+
+    # the two implementations agree
+    assert np.isclose(beam_impactx.beam_size_x(), beam_thick.beam_size_x(), rtol=1e-3)
+    assert np.isclose(beam_impactx.beam_size_y(), beam_thick.beam_size_y(), rtol=1e-3)
+    assert np.isclose(beam_impactx.norm_emittance_x(), beam_thick.norm_emittance_x(), rtol=1e-3)
+    assert np.isclose(beam_impactx.norm_emittance_y(), beam_thick.norm_emittance_y(), rtol=1e-3)
+
+    # the transverse offsets deflect the beam by the expected amount
+    assert np.isclose(beam_impactx.x_angle(), beam_thick.x_angle(), rtol=5e-3)
+    assert np.isclose(beam_impactx.y_angle(), beam_thick.y_angle(), rtol=5e-3)
+
+
+@pytest.mark.plasma_lens
+@pytest.mark.impactx
+def test_PlasmaLensImpactX_apertures():
+    """
+    Check that the ImpactX plasma lens clips charge outside the capillary
+    when apertures are enabled.
+    """
+
+    np.random.seed(42)
+
+    source = SourceBasic()
+    source.bunch_length = 100e-6 # [m]
+    source.num_particles = 10000
+    source.charge = -SI.e * 1.0e10 # [C]
+    source.energy = 1e9 # [eV]
+    source.rel_energy_spread = 1e-5
+    source.emit_nx, source.emit_ny = 1e-6, 1e-6 # [m rad]
+    source.beta_x = 120.0 # [m] large beta => beam wider than the capillary
+    source.beta_y = source.beta_x
+
+    R_pl = 500e-6 # [m] roughly two beam sigmas
+
+    plasma_lens = PlasmaLensImpactX(length=0.01, radius=R_pl, current=1e3, use_apertures=True)
+
+    beam0 = source.track()
+    beam = plasma_lens.track(copy.deepcopy(beam0))
+
+    # charge outside the capillary has been removed
+    assert beam.abs_charge() < beam0.abs_charge()
+    assert np.all(np.sqrt(beam.xs()**2 + beam.ys()**2) <= R_pl*(1 + 1e-9))