Add option to include a z-offset in multi-bunch plots

Sometimes an offset to the starting position is required to ensure no particles have negative starting position. In such cases, the z-location of particles should be adjusted to enable direct comparison of experiment and simulation.

Specifying a z-offset makes the plots as if the particles had started at z=0 rather than at the offset location.

Also changed titles for multi-bunch energy spectra plots, from e.g. "BPM 110" to "at z = 0.6 mm"

GUI/src/ImpactTPlot.py

@@ -912,6 +912,14 @@ def __init__(self, parent):
         box = self.subfig.get_position()
         self.subfig.set_position([box.x0*1.4, box.y0, box.width, box.height])
         self.plot()
+    def create_option_frame(self, Nbunch, per_bunch=True):
+        """Set up options (override base class, add offset option)."""
+        self.option_frame = tk.Frame(self)
+        self.option_frame.pack()
+        self.create_bunch_selector(Nbunch, per_bunch)
+        self.create_xaxis_selector()
+        self.create_offset_selector()
+        self.create_plot_button()
     def create_bunch_selector(self, Nbunch, per_bunch=True):
         """Add selector (override base class, give 'from' and 'to' bunches)."""
         if per_bunch and Nbunch > 1:
@@ -935,6 +943,13 @@ def create_bunch_selector(self, Nbunch, per_bunch=True):
                                                 width=6,
                                                 values=self.to_bunch_list)
             self.to_bunch_select.pack(fill='both', expand=1, side='left')
+    def create_offset_selector(self):
+        """Add entry box for a z-offset."""
+        self.offset_label = tk.Label(self.option_frame, text='z offset (mm):')
+        self.offset_label.pack(side='left')
+        self.offset_box = tk.Entry(self.option_frame, width=6)
+        self.offset_box.insert(0, '0')
+        self.offset_box.pack(fill='both', expand=1, side='left')
     def get_from_bunch(self):
         """Get branch selected in options frame."""
         if hasattr(self, "from_bunch_select"):
@@ -950,6 +965,12 @@ def get_to_bunch(self):
     def get_bunch_list(self):
         """Find which bunches are selected in the option frame."""
         return list(range(self.get_from_bunch(), self.get_to_bunch()+1))
+    def get_offset(self):
+        """Get z offset as set in options frame."""
+        if hasattr(self, "offset_box"):
+            return float(self.offset_box.get())/1000
+        else:
+            return None
 
 class PlotMBBeamSizeFrame(PlotMultiBunchBaseFrame):
     """Frame to plot rms beam sizes for selected bunches."""
@@ -964,7 +985,8 @@ def plot(self):
             print('Continuing without experimental data.')
             experimental_results = None
         bunch_list = self.get_bunch_list()
-        xdata, ydata = MultiBunchPlot.load_statistics_data(bunch_list)
+        xdata, ydata = MultiBunchPlot.load_statistics_data(bunch_list,
+                                                           self.get_offset())
         # TODO: Can we calculate rdata from xdata and ydata?
         self.subfig.cla()
         MultiBunchPlot.plot_beam_size(self.subfig.axes, xdata, bunch_list,
@@ -979,7 +1001,7 @@ def __init__(self, parent):
     def plot(self):
         """Plot bunch counts up to the selected max bunch."""
         self.subfig.cla()
-        self.data = MultiBunchPlot.load_bunch_count_data()
+        self.data = MultiBunchPlot.load_bunch_count_data(self.get_offset())
         MultiBunchPlot.plot_bunch_count(self.subfig, self.data,
                                         self.get_bunch_list(),
                                         xaxis=self.get_selected_xaxis())
@@ -992,7 +1014,8 @@ def __init__(self, parent):
     def plot(self):
         """Plot average of x and y emittance for selected bunches."""
         bunch_list = self.get_bunch_list()
-        xdata, ydata = MultiBunchPlot.load_statistics_data(bunch_list)
+        xdata, ydata = MultiBunchPlot.load_statistics_data(bunch_list,
+                                                           self.get_offset())
         self.subfig.cla()
         MultiBunchPlot.plot_emittance(self.subfig.axes,
                                       xdata, ydata, bunch_list,
@@ -1006,7 +1029,8 @@ def __init__(self, parent):
     def plot(self):
         """Plot growth of average x and y emittance for selected bunches."""
         bunch_list = self.get_bunch_list()
-        xdata, ydata = MultiBunchPlot.load_statistics_data(bunch_list)
+        xdata, ydata = MultiBunchPlot.load_statistics_data(bunch_list,
+                                                           self.get_offset())
         self.subfig.cla()
         MultiBunchPlot.plot_emittance_growth(self.subfig.axes,
                                              xdata, ydata, bunch_list,
@@ -1019,13 +1043,15 @@ def __init__(self, parent):
         PlotBaseFrame.__init__(self, parent, per_bunch=True)
         self.last_filenumber = 0
         self.last_bunch_count = 0
+        self.last_offset = 0
         self.plot()
     def create_option_frame(self, Nbunch, per_bunch=True):
         """Add options (override base class)."""
         self.option_frame = tk.Frame(self)
         self.option_frame.pack()
         self.create_slice_selector()
         self.create_bunch_selector(Nbunch)
+        self.create_offset_selector()
         self.create_bins_box()
         self.create_plot_button()
     def create_slice_selector(self):
@@ -1040,6 +1066,13 @@ def create_slice_selector(self):
                                          width=6,
                                          values=list(self.slice_list))
         self.slice_select.pack(fill='both', expand=1, side='left')
+    def update_slice_selector(self):
+        """Reset the list of slice positions when the z offset changes."""
+        filenumber = self.get_filenumber()
+        self.slice_list = self.get_slice_list()
+        new_slice = [k for k,v in self.slice_list.items() if v == filenumber][0]
+        self.slice_select['values'] = list(self.slice_list)
+        self.slice_select.set(new_slice)
     def create_bins_box(self):
         """Add text box to set number of bins for plot sampling."""
         self.bins_label = tk.Label(self.option_frame, text='Bins: ')
@@ -1050,7 +1083,7 @@ def create_bins_box(self):
     def get_slice_list(self):
         """Get list of available phase space slices, including BPMs."""
         lattice = MultiBunchPlot.get_lattice()
-        bpm_list = MultiBunchPlot.get_bpms(lattice)
+        bpm_list = MultiBunchPlot.get_bpms(lattice, self.get_offset())
         slice_list = {'Initial': IMPACT_T_initial_slice}
         slice_list.update(bpm_list)
         slice_list['Final'] = IMPACT_T_final_slice
@@ -1077,13 +1110,18 @@ def build_title(self, base_title, filenumber):
             return (f'{base_title} at z = {matches[0]} '.capitalize()
                     + f'for {bunch_text}')
     def refresh_data(self):
-        """Reload data (list of all bunch data) when filenumber is changed."""
+        """Reload all bunch data when filenumber or offset is changed."""
         filenumber = self.get_filenumber()
-        if filenumber != self.last_filenumber:
+        z_offset = self.get_offset()
+        if z_offset != self.last_offset:
+            self.update_slice_selector()
+        if filenumber != self.last_filenumber or z_offset != self.last_offset:
             full_bunch_list = list(range(1, self.Nbunch+1))
             self.data = MultiBunchPlot.load_phase_space_data(filenumber,
-                                                             full_bunch_list)
+                                                             full_bunch_list,
+                                                             self.get_offset())
             self.last_filenumber = filenumber
+            self.last_offset = z_offset
     def get_selected_data(self):
         """Return datasets for only the selected bunches."""
         self.refresh_data()

GUI/src/MultiBunchPlot.py

@@ -24,11 +24,26 @@ def get_lattice():
     input = read_input_file(get_input_filename(1))
     return [line.split() for line in input[9:]]
 
-def get_bpms(lattice):
+def get_bpms(lattice, z_offset=None):
     """Get the location and file number of all BPMs as a list of tuples."""
-    return [(str(float(elem[4])*1000) + ' mm', int(elem[2]))
+    return [(get_position_as_text(float(elem[4]), z_offset), int(elem[2]))
             for elem in lattice if elem[3]=='-2']
 
+def get_position_as_text(position, z_offset=None):
+    """Return a text description of a numerical position."""
+    if z_offset:
+        position = position - z_offset
+    if position > 1:
+        unit = 'm'
+    else:
+        position = position*1000
+        unit = 'mm'
+    if position >= 1000:
+        position = round(position)
+    else:
+        position = float(f'{position:.4g}')
+    return f'{position} {unit}'
+
 def get_bunch_counts(bunch_list):
     """Get the particle counts for each bunch as an array."""
     input_filenames = [get_input_filename(bunch) for bunch in bunch_list]
@@ -89,28 +104,43 @@ def load_experimental_results(filename='experimental_data.txt'):
     """Load z, rms beam size and error values (m) from file as an array."""
     return load_data_from_file(filename, header=1)
 
-def load_statistics_data(bunch_list):
+def load_statistics_data(bunch_list, z_offset=None):
     """Load x and y data per bunch from statistics files as arrays."""
     xdata = []
     ydata = []
     for bunch in bunch_list:
-        xdata.append(load_data_from_file(f'fort.{bunch}024'))
-        ydata.append(load_data_from_file(f'fort.{bunch}025'))
+        x = load_data_from_file(f'fort.{bunch}024')
+        y = load_data_from_file(f'fort.{bunch}024')
+        if z_offset:
+            x = shift_z(x, z_offset, 1)
+            y = shift_z(y, z_offset, 1)
+        xdata.append(x)
+        ydata.append(y)
     return numpy.array(xdata), numpy.array(ydata)
 
-def load_bunch_count_data():
+def load_bunch_count_data(z_offset=None):
     """Load bunch counts vs time from `fort.11` as an array."""
     data = load_data_from_file('fort.11')
-    return numpy.delete(data, 0, 1) # remove first column (time-step #)
+    data = numpy.delete(data, 0, 1) # remove first column (time-step #)
+    if z_offset:
+        data = shift_z(data, z_offset, 1)
+    return data
 
-def load_phase_space_data(filenumber, bunch_list):
+def load_phase_space_data(filenumber, bunch_list, z_offset=None):
     """Load phase space data per bunch as a list of arrays."""
     data = []
     for bunch in bunch_list:
         this_data = load_data_from_file(f'fort.{filenumber+bunch-1}')
+        if z_offset:
+            this_data = shift_z(this_data, z_offset, 4)
         data.append(calculate_energies(this_data, bunch))
     return data
 
+def shift_z(data, z_offset, z_col=1):
+    """Shift all z values by the given offset."""
+    data[:,z_col] = data[:,z_col] - z_offset
+    return data
+
 def calculate_energies(data, bunch):
     """Calculate the energies for per-bunch data and save as an extra column."""
     gamma = numpy.sqrt(1 + numpy.square(data.T[5]))
@@ -374,7 +404,7 @@ def add_plot_margins(axes, margin):
     axes.set_xlim(xmin - xmargin, xmax + xmargin)
     axes.set_ylim(ymin - ymargin, ymax + ymargin)
 
-def plot_all(bunch_list):
+def plot_all(bunch_list, z_offset=None):
     """Run and save all plots consecutively."""
     full_bunch_list = list(range(1, int(get_bunch_count()) + 1))
     bunch_list, invalid_bunches = check_bunch_list(bunch_list)
@@ -389,7 +419,7 @@ def plot_all(bunch_list):
         experimental_results = None
     print('Loading statistical data...')
     try:
-        xdata, ydata = load_statistics_data(bunch_list)
+        xdata, ydata = load_statistics_data(bunch_list, z_offset)
     except FileNotFoundError as err:
         print(f'Statistical data file not found: {err}')
         print('Skipping statistical plots.')
@@ -423,7 +453,7 @@ def plot_all(bunch_list):
         matplotlib.pyplot.close(figure)
     print('Loading bunch count data...')
     try:
-        data = load_bunch_count_data()
+        data = load_bunch_count_data(z_offset)
     except FileNotFoundError as err:
         print(f'Bunch count data file not found: {err}')
         print('Skipping bunch count plot.')
@@ -435,7 +465,7 @@ def plot_all(bunch_list):
         matplotlib.pyplot.close(figure)
     print('Loading initial phase space data...')
     try:
-        full_data = load_phase_space_data(40, full_bunch_list)
+        full_data = load_phase_space_data(40, full_bunch_list, z_offset)
     except FileNotFoundError as err:
         print(f'Phase space data file not found: {err}')
         print('Skipping initial phase space step.')
@@ -469,7 +499,7 @@ def plot_all(bunch_list):
         matplotlib.pyplot.close(figure)
     print('Loading final phase space data...')
     try:
-        full_data = load_phase_space_data(50, full_bunch_list)
+        full_data = load_phase_space_data(50, full_bunch_list, z_offset)
     except FileNotFoundError as err:
         print(f'Phase space data file not found: {err}')
         print('Skipping final phase space step.')
@@ -504,15 +534,16 @@ def plot_all(bunch_list):
     print('Getting list of BPMs...')
     try:
         lattice = get_lattice()
-        bpm_list = get_bpms(lattice)
+        bpm_list = get_bpms(lattice, z_offset)
     except FileNotFoundError as err:
         print(f'Input file not found: {err}')
         print('Skipping BPM plot steps.')
     else:
         for location, filenumber in bpm_list:
             print(f'Loading BPM {filenumber} phase space data...')
             try:
-                full_data = load_phase_space_data(filenumber, full_bunch_list)
+                full_data = load_phase_space_data(filenumber, full_bunch_list,
+                                                  z_offset)
             except FileNotFoundError as err:
                 print(f'BPM data file not found: {err}')
                 print('Skipping this BPM plot step.')
@@ -539,13 +570,13 @@ def plot_all(bunch_list):
                 print(f'Plotting BPM {filenumber} energy spectra...')
                 figure, axes = matplotlib.pyplot.subplots(dpi=300)
                 plot_bunch_energies(axes, data, bunch_list, bins=300,
-                                    title=(f'BPM {filenumber} energy spectra '
+                                    title=(f'Energy spectra at z = {location} '
                                            f'for {bunch_text(bunch_list)}'))
                 figure.savefig(f'energies-{filenumber}')
                 matplotlib.pyplot.close(figure)
                 figure, axes = matplotlib.pyplot.subplots(dpi=300)
                 plot_total_energy(axes, combined_data, bunch_list, bins=300,
-                                  title=(f'BPM {filenumber} energy spectrum '
+                                  title=(f'Energy spectrum at z = {location} '
                                          f'for {bunch_text(bunch_list)}'))
                 figure.savefig(f'energy-{filenumber}')
                 matplotlib.pyplot.close(figure)
@@ -563,4 +594,12 @@ def plot_all(bunch_list):
                 bunch_list = list(range(1, int(get_bunch_count()) + 1))
     else:
         bunch_list = list(range(1, int(get_bunch_count()) + 1))
-    plot_all(bunch_list)
+    if len(sys.argv) > 2:
+        z_offset = sys.argv[2]
+        try:
+            z_offset = float(z_offset)
+        except:
+            z_offset = None
+    else:
+        z_offset = None
+    plot_all(bunch_list, z_offset)