[Draft] Adding Experiment setups

abel/__init__.py

@@ -21,7 +21,8 @@
 from .classes.interstage.interstage import Interstage
 from .classes.rf_accelerator.rf_accelerator import RFAccelerator, RFAcceleratorInitializationException
 from .classes.bds.bds import BeamDeliverySystem
-from .classes.spectrometer.spectrometer import Spectrometer
+from .classes.spectrometer import Spectrometer
+from .classes.spectrometer.quad_imaging import SpectrometerQuadImaging
 from .classes.damping_ring.damping_ring import DampingRing
 from .classes.combiner_ring.combiner_ring import CombinerRing
 from .classes.turnaround.turnaround import Turnaround
@@ -30,16 +31,20 @@
 from .classes.beamline.impl.linac.linac import Linac
 from .classes.beamline.impl.linac.impl.plasma_linac import PlasmaLinac
 from .classes.beamline.impl.linac.impl.conventional_linac import ConventionalLinac
-from .classes.beamline.impl.experiment.experiment import Experiment
-from .classes.beamline.impl.experiment.impl.experiment_pwfa import ExperimentPWFA
-from .classes.beamline.impl.experiment.impl.experiment_apl import ExperimentAPL
+
+from .classes.beamline.impl.experiment import Experiment
+from .classes.beamline.impl.experiment.pwfa import ExperimentPWFA
+from .classes.beamline.impl.experiment.pwfa.preset.flashforward import FLASHForward
+from .classes.beamline.impl.experiment.pwfa.preset.facet2 import FACET2
+from .classes.beamline.impl.experiment.apl import ExperimentAPL
 
 from .classes.collider.collider import Collider
 
 from .classes.source.impl.source_basic import SourceBasic
 from .classes.source.impl.source_trapezoid import SourceTrapezoid
 from .classes.source.impl.source_combiner import SourceCombiner
 from .classes.source.impl.source_from_file import SourceFromFile
+from .classes.source.impl.source_from_profile import SourceFromProfile
 from .classes.source.impl.source_flattop import SourceFlatTop
 from .classes.source.impl.source_capsule import SourceCapsule
 from .classes.stage.impl.stage_basic import StageBasic
@@ -67,8 +72,9 @@
 from .classes.transfer_line.impl.transfer_line_basic import TransferLineBasic
 from .classes.bds.impl.bds_basic import BeamDeliverySystemBasic
 from .classes.bds.impl.bds_fbt import BeamDeliverySystemFlatBeamTransformer
-from .classes.spectrometer.impl.spectrometer_facet_ocelot import SpectrometerFacetOcelot
-from .classes.spectrometer.impl.spectrometer_basic_clear import SpectrometerBasicCLEAR
+from .classes.spectrometer.quad_imaging.preset.facet2 import SpectrometerFACET2
+#from .classes.spectrometer.impl.spectrometer_facet_ocelot import SpectrometerFacetOcelot
+#from .classes.spectrometer.impl.spectrometer_basic_clear import SpectrometerBasicCLEAR
 from .classes.ip.impl.ip_basic import InteractionPointBasic
 from .classes.ip.impl.ip_guineapig import InteractionPointGuineaPig
 

abel/apis/hipace/hipace_api.py

@@ -29,6 +29,8 @@ def hipace_write_inputs(filename_input, filename_beam, filename_driver, plasma_d
     # driver-only mode
     if driver_only:
         beam_components = 'driver'
+    elif filename_driver is None:
+        beam_components = 'beam'
     else:
         beam_components = 'driver beam'
 
@@ -159,12 +161,18 @@ def hipace_run(filename_job_script, num_steps, runfolder=None, quiet=False):
 
     # when finished, load the beam and driver
     filename = os.path.join(runfolder, "diags/hdf5/openpmd_{:06}.h5".format(int(num_steps)))
-
+
+    # attempt to load the beam
     try:
         beam = Beam.load(filename, beam_name='beam')
     except:
         beam = None
-    driver = Beam.load(filename, beam_name='driver')
+
+    # attempt to load the driver
+    try:
+        driver = Beam.load(filename, beam_name='driver')
+    except:
+        driver = None
 
     return beam, driver
 

abel/classes/beam.py

@@ -19,7 +19,7 @@
 
 class Beam():
 
-    def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1, bunch_separation=0.0):
+    def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1, bunch_separation=0.0, location=0.0, particle_mass=SI.m_e):
 
         # check the inputs
         if num_particles < 1 or not isinstance(num_particles, int):
@@ -41,7 +41,10 @@ def __init__(self, phasespace=None, num_particles=1000, num_bunches_in_train=1,
 
         self.trackable_number = -1 # will increase to 0 after first tracking element
         self.stage_number = 0
-        self.location = 0        
+        self.location = location   
+
+        # particle mass
+        self.particle_mass = particle_mass
 
 
     # reset phase space
@@ -1685,8 +1688,7 @@ def plot_transverse_profile(self):
         ax.set_title('Transverse profile')
         cb = fig.colorbar(p)
         cb.ax.set_ylabel('Charge density (pC/um^2)')
-
-
+
     # TODO: unfinished!
     # def plot_bunch_pattern(self):
 

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

@@ -0,0 +1,84 @@
+from abc import abstractmethod
+from abel.CONFIG import CONFIG
+from abel.classes.beamline.beamline import Beamline
+from abel.classes.beamline.impl.linac.linac import Linac
+from abel.classes.source.source import Source
+from abel.classes.spectrometer import Spectrometer
+import numpy as np
+
+class Experiment(Beamline):
+
+    def __init__(self, linac=None, test_device=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.test_device = test_device
+        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 test device
+        self.trackables.append(self.test_device)
+
+        # 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()
+
+
+    def energy_usage(self):
+        return 0
+
+    def get_nom_beam_power(self):
+        return 0
+
+    def get_cost_breakdown():
+        return ('Experiment', 0)
+
+
+    # plot the spectrometer screen
+    def plot_spectrometer_screen(self, xlims=None, ylims=None):
+
+        from matplotlib import pyplot as plt
+
+        # load phase space
+        beam = self.final_beam
+
+        # make screen projection
+        xbins, ybins = None, None
+        num_bins = round(np.sqrt(len(beam))/2)
+        if xlims is not None:
+            xbins = np.linspace(min(xlims), max(xlims), num_bins)
+        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(1, 1)
+        fig.set_figwidth(6)
+        fig.set_figheight(9)
+
+        # plot spectrometer image
+        c0 = ax.pcolor(xedges*1e3, yedges * 1e3, abs(dQdxdy) * 1e3, cmap=CONFIG.default_cmap, shading='auto')
+
+        # make plot
+        ax.set_xlabel('x (mm)')
+        ax.set_ylabel('y (mm)')
+        ax.set_title('Spectrometer screen (shot ' + str(self.shot+1) + ')')
+        ax.grid(False, which='major')
+        cbar0 = fig.colorbar(c0, ax=ax, pad = 0.015*fig.get_figwidth())
+        cbar0.ax.set_ylabel('Charge density (nC ' r'$\mathrm{mm^{-2})}$')
+        ax.set_ylim(min(yedges * 1e3), max(yedges * 1e3))
+

abel/classes/beamline/impl/experiment/impl/experiment_apl.py → abel/classes/beamline/impl/experiment/apl/__init__.py

@@ -1,21 +1,21 @@
-from abel.classes.beamline.impl.experiment.experiment import Experiment
+from abel.classes.beamline.impl.experiment import Experiment
 from abel.classes.plasma_lens.plasma_lens import PlasmaLens
 
 class ExperimentAPL(Experiment):
 
     def __init__(self, linac=None, plasma_lens=None, spectrometer=None):
         self.plasma_lens = plasma_lens
 
-        super().__init__(linac=linac, component=plasma_lens, spectrometer=spectrometer)
+        super().__init__(linac=linac, test_device=plasma_lens, spectrometer=spectrometer)
 
 
     # assemble the trackables
     def assemble_trackables(self):
 
         # check element classes, then assemble
-        if self.component is None:
-            self.component = self.plasma_lens
-        assert(isinstance(self.component, PlasmaLens))
+        if self.test_device is None:
+            self.test_device = self.plasma_lens
+        assert(isinstance(self.test_device, PlasmaLens))
 
         # run beamline constructor
         super().assemble_trackables()