Achromatic spectrometer

abel/__init__.py

@@ -42,6 +42,7 @@
 from .classes.source.impl.source_basic import *
 from .classes.source.impl.source_trapezoid import SourceTrapezoid
 from .classes.source.impl.source_combiner import *
+from .classes.source.impl.source_energy_comb import *
 from .classes.source.impl.source_from_file import *
 from .classes.source.impl.source_flattop import *
 from .classes.source.impl.source_capsule import *
@@ -72,6 +73,8 @@
 from .classes.bds.impl.bds_fbt import *
 from .classes.spectrometer.impl.spectrometer_facet_ocelot import *
 from .classes.spectrometer.impl.spectrometer_basic_clear import *
+from .classes.spectrometer.impl.achromatic.spectrometer_achromatic import SpectrometerAchromatic
+from .classes.spectrometer.impl.achromatic.impl.spectrometer_achromatic_basic import SpectrometerAchromaticBasic
 from .classes.ip.impl.ip_basic import *
 from .classes.ip.impl.ip_guineapig import *
 

abel/classes/beam.py

@@ -1364,7 +1364,7 @@ def plot_trace_space_y(self):
         cb = fig.colorbar(p)
         cb.ax.set_ylabel('Charge density (pC/um/mrad)')
 
-    def plot_transverse_profile(self):
+    def plot_transverse_profile(self, xlims=(None, None), ylims=(None, None), title=''):
         dQdxdy, xs, ys = self.phase_space_density(self.xs, self.ys)
 
         fig, ax = plt.subplots()
@@ -1379,6 +1379,9 @@ def plot_transverse_profile(self):
         cb = fig.colorbar(p)
         cb.ax.set_ylabel('Charge density (pC/um^2)')
 
+        ax.set_xlim(xlims)
+        ax.set_ylim(ylims)
+        ax.set_title(title)
 
     # TODO: unfinished!
     # def plot_bunch_pattern(self):

abel/classes/beamline/impl/experiment/experiment.py

@@ -9,33 +9,34 @@
 import warnings
 
 class Experiment(Beamline):
-    
+
     def __init__(self, linac=None, component=None, spectrometer=None, num_bunches_in_train=1, bunch_separation=0, rep_rate_trains=10):
         super().__init__(num_bunches_in_train, bunch_separation, rep_rate_trains)
-        
+
         self.linac = linac
         self.component = component
         self.spectrometer = spectrometer
-        
-    
-    
+
+
+
     # assemble the trackables
     def assemble_trackables(self):
 
         self.trackables = []
-        
+
         # add the linac/source
         assert(isinstance(self.linac, Linac) or isinstance(self.linac, Source))
         self.trackables.append(self.linac)
 
         # add the experimental component
-        self.trackables.append(self.component)
+        if self.component is not None:
+            self.trackables.append(self.component)
 
         # add the spectrometer (if it exists
         if self.spectrometer is not None:
             assert(isinstance(self.spectrometer, Spectrometer))
             self.trackables.append(self.spectrometer)
-        
+
         # set the bunch train pattern etc.
         super().assemble_trackables()
 
@@ -45,10 +46,10 @@ def energy_usage(self):
 
     def get_cost_breakdown():
         return ('Experiment', 0)
-    
+
     # density plots
     def plot_spectrometer_screen(self, xlims=None, ylims=None, E_calib = False, diverg = None, plot_m12 = False, savefig = None):
-        
+
         # load phase space
         beam = self.final_beam
 
@@ -60,15 +61,15 @@ def plot_spectrometer_screen(self, xlims=None, ylims=None, E_calib = False, dive
         if ylims is not None:
             ybins = np.linspace(min(ylims), max(ylims), num_bins)
         dQdxdy, xedges, yedges = beam.density_transverse(hbins=xbins, vbins=ybins)
-        
+
         # prepare figure
         fig, ax = plt.subplots(2, 1, gridspec_kw={'height_ratios': [2, 1]})
         fig.set_figwidth(8)
         fig.set_figheight(8)
-        
+
         # current profile
         c0 = ax[0].pcolor(xedges*1e3, yedges * 1e3, abs(dQdxdy) * 1e3, cmap=CONFIG.default_cmap, shading='auto')
-        
+
         # make plot
         ax[0].set_xlabel('x (mm)')
         ax[0].set_ylabel('y (mm)')
@@ -77,25 +78,25 @@ def plot_spectrometer_screen(self, xlims=None, ylims=None, E_calib = False, dive
         cbar0 = fig.colorbar(c0, ax=ax[0], pad = 0.015*fig.get_figwidth())
         cbar0.ax.set_ylabel('Charge density (nC ' r'$\mathrm{mm^{-2})}$')
         ax[0].set_ylim(min(yedges * 1e3), max(yedges * 1e3))
-        
+
         # calculate energy axis (E = E0*Dy/y)
         Dy_img = self.spectrometer.get_dispersion(energy=self.spectrometer.img_energy)
         E_times_Dy = self.spectrometer.img_energy*Dy_img
-        
+
         # add imaging energies
         y_img = -E_times_Dy/self.spectrometer.img_energy
         y_img_y = -E_times_Dy/self.spectrometer.img_energy_y
         ax[0].axhline(y_img*1e3, color = 'black', linestyle = '--', linewidth = 1, label = 'Imaging energy x (GeV)')
         ax[0].axhline(y_img_y*1e3, color = 'black', linestyle = ':', linewidth = 1, label = 'Imaging energy y (GeV)')
-        
+
         # add energy axis
         warnings.simplefilter('ignore', category=RuntimeWarning)
         ax2 = ax[0].secondary_yaxis('right', functions=(lambda y: -E_times_Dy/(y/1e3)/1e9, lambda y: -E_times_Dy/(y/1e3)/1e9))
         ax2.set_ylabel('Energy (GeV)')
-        
+
         if diverg is not None:
             energies = -E_times_Dy/yedges
-            m12s = self.spectrometer.get_m12(energies)        
+            m12s = self.spectrometer.get_m12(energies)
             x_pos = m12s*diverg
             x_neg = -m12s*diverg
             ax[0].plot(x_pos*1e3, yedges*1e3, color = 'orange', label = str(diverg*1e3) + ' mrad butterfly', alpha = 0.7)
@@ -104,12 +105,11 @@ def plot_spectrometer_screen(self, xlims=None, ylims=None, E_calib = False, dive
             ax[1].set_ylabel(r'$\mathrm{m_{12}}$')
             ax[1].set_xlabel('E (GeV)')
             ax[1].grid(False, which='major')
-            
+
         ax[0].set_xlim(min(xedges*1e3), max(xedges*1e3))
         ax[0].legend(loc = 'center right', fontsize = 6)
-        
+
         if plot_m12 == False:
             fig.delaxes(ax[1])
         if savefig is not None:
             plt.savefig(str(savefig), dpi = 700)
-

abel/classes/plasma_lens/impl/plasma_lens_nonlinear_thin.py

@@ -2,28 +2,52 @@
 from matplotlib import patches
 import numpy as np
 import scipy.constants as SI
+import matplotlib.pyplot as plt
 
 class PlasmaLensNonlinearThin(PlasmaLens):
-    
-    def __init__(self, length=None, radius=None, current=None, rel_nonlinearity=0):
+
+    def __init__(self, length=None, radius=None, current=None, rel_nonlinearity=0, nonlinearity_in_x=True):
 
         super().__init__(length, radius, current)
-        
+
         # set nonlinearity (defined as R/Dx)
         self.rel_nonlinearity = rel_nonlinearity
-
-
+        self.nonlinearity_in_x = nonlinearity_in_x
+
+    def remove_charge_outside(self, beam, plot=False):
+
+        mask = np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius
+
+        if plot:
+
+            fig, ax = plt.subplots()
+            ax.add_patch(patches.Circle((self.offset_x, self.offset_y), self.radius, fill=False, linestyle='--'))
+
+            ax.scatter(beam.xs(), beam.ys(), s=1, alpha=0.1, zorder=0)
+            ax.set_aspect('equal', 'box')
+            plt.show()
+
+            rel_loss = mask.sum() / mask.size
+            print(f"Removing {rel_loss*100:.2f}% of the charge outside the lens")
+
+        # remove charge outside the lens
+        del beam[mask]
+
+        return beam
+
     def track(self, beam, savedepth=0, runnable=None, verbose=False):
 
         # remove charge outside the lens (start)
-        del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius]
-
+        # del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius]
+        self.remove_charge_outside(beam, plot=False)
+
         # drift half the distance
         beam.transport(self.length/2)
-        
+
         # remove charge outside the lens (middle)
-        del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius]
-
+        # del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius]
+        self.remove_charge_outside(beam)
+
         # get particles
         xs = beam.xs()-self.offset_x
         xps = beam.xps()
@@ -34,16 +58,21 @@ def track(self, beam, savedepth=0, runnable=None, verbose=False):
         g0 = self.get_focusing_gradient()
 
         # calculate the nonlinearity
-        inv_Dx = self.rel_nonlinearity/self.radius
-
+        if self.nonlinearity_in_x:
+            inv_Dx = self.rel_nonlinearity/self.radius
+            inv_Dy = 0.0
+        else:
+            inv_Dx = 0.0
+            inv_Dy = self.rel_nonlinearity/self.radius
+
         # thin lens kick
-        Bx = g0*(ys + xs*ys*inv_Dx)
-        By = -g0*(xs + (xs**2 + ys**2)/2*inv_Dx)
+        Bx = g0*(ys + xs*ys*inv_Dx  + (xs**2 + ys**2)/2*inv_Dy)
+        By = -g0*(xs + (xs**2 + ys**2)/2*inv_Dx + xs*ys*inv_Dy)
 
         # calculate the angular kicks
         delta_xp = self.length*(By*beam.charge_sign()*SI.c/beam.Es())
         delta_yp = -self.length*(Bx*beam.charge_sign()*SI.c/beam.Es())
-        
+
         # set new beam positions and angles (shift back plasma-lens offsets)
         beam.set_xps(xps + delta_xp)
         beam.set_yps(yps + delta_yp)
@@ -52,11 +81,11 @@ def track(self, beam, savedepth=0, runnable=None, verbose=False):
         beam.transport(self.length/2)
 
         # remove charge outside the lens (end)
-        del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius]
-
-        return super().track(beam, savedepth, runnable, verbose)   
+        # del beam[np.sqrt((beam.xs()-self.offset_x)**2 + (beam.ys()-self.offset_y)**2) > self.radius]
+        self.remove_charge_outside(beam)
+
+        return super().track(beam, savedepth, runnable, verbose)
 
 
     def get_focusing_gradient(self):
         return SI.mu_0 * self.current / (2*np.pi * self.radius**2)
-

abel/classes/plasma_lens/plasma_lens.py

@@ -4,32 +4,35 @@
 import numpy as np
 
 class PlasmaLens(Trackable):
-    
+
     @abstractmethod
     def __init__(self, length, radius, current, offset_x=0, offset_y=0):
 
         super().__init__()
-        
+
         # common variables
         self.length = length
         self.radius = radius
         self.current = current
 
         self.offset_x = offset_x
         self.offset_y = offset_y
-        
-    
+
+
     @abstractmethod
     def track(self, beam, savedepth=0, runnable=None, verbose=False):
         return super().track(beam, savedepth, runnable, verbose)
 
     @abstractmethod
     def get_focusing_gradient(self):
         pass
-        
+
     def get_length(self):
         return self.length
 
     def survey_object(self):
         return patches.Rectangle((0, -1), self.get_length(), 2)
-
+
+    @property
+    def l(self):
+        return self.length