From 6742951ccf79542e822902f87edb3f6135cd5167 Mon Sep 17 00:00:00 2001
From: onaumenk <oleksandr.naumenko@cern.ch>
Date: Fri, 9 Jun 2023 13:17:27 +0200
Subject: [PATCH 1/6] Added automatic default value for alpha of longitudinal
 phase space plot, such that they look good for any number of macroparticles.

---
 blond/plots/plot.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/blond/plots/plot.py b/blond/plots/plot.py
index a21b4537..38f30e0b 100644
--- a/blond/plots/plot.py
+++ b/blond/plots/plot.py
@@ -14,7 +14,7 @@
 '''
 
 import os
-
+import numpy as np
 import h5py as hp
 import matplotlib.pyplot as plt
 
@@ -158,7 +158,7 @@ class Plot:
             self.msize = format_options['markersize']
 
         if 'alpha' not in format_options:
-            self.alpha = 0.05
+                self.alpha = 10**(-np.log10(self.beam.n_macroparticles)/6)
         else:
             self.alpha = format_options['alpha']
 
-- 
GitLab


From 829d8930cbc85f5cb50b739a892287ea9a802334 Mon Sep 17 00:00:00 2001
From: onaumenk <oleksandr.naumenko@cern.ch>
Date: Wed, 14 Jun 2023 12:07:11 +0200
Subject: [PATCH 2/6] Reworked plot_impedance_vs_frequency to use current BlonD
 nomenclature and objects. Added documentation and improved naming of some
 arguments, especially concerning plot options. Improved readability of the
 plots by adding legends and labels where necessary. Added option to output
 plots to the console directly by passing the argument show_plot = True to
 functions Made the plots resize themselves without errors. This means that
 each plot might look different if significant changes are present, but each
 plot individually should look good with this change.

---
 blond/plots/plot_impedance.py | 220 +++++++++++++++++++++++++---------
 1 file changed, 165 insertions(+), 55 deletions(-)

diff --git a/blond/plots/plot_impedance.py b/blond/plots/plot_impedance.py
index 43d4439d..1963cf8d 100644
--- a/blond/plots/plot_impedance.py
+++ b/blond/plots/plot_impedance.py
@@ -21,104 +21,214 @@ import matplotlib.pyplot as plt
 import numpy as np
 
 from ..impedances.impedance_sources import InputTable
+import matplotlib.lines as mlines
 
+# Automatic layouting of plots
+plt.rc('figure', autolayout=True)
+plt.rc('savefig', bbox='tight')
+plt.rc('savefig', pad_inches=0.1)
+plt.rc('figure', figsize = [8,6])
 
-def plot_impedance_vs_frequency(counter, general_params, ind_volt_from_imp,
-                                option1="sum", option2="no_spectrum",
-                                option3="freq_fft", style='-',
+def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
+                                plot_total_impedance = True, plot_spectrum=False,
+                                plot_interpolated_impedances = False, style='-',
                                 cut_left_right=None, cut_up_down=None,
-                                dirname='fig'):
+                                dirname='fig', show_plot = False):
     """
-    Plot of impedance vs frequency.
+    Plots Impedance in frequency domain. Given an InducedVoltageFreq object, i.e. one
+    that was created from a list of impedances:
+     - can plot either the total, summed impedance or each individual one from the list
+     - can additionally plot the beam spectrum
+     - and can plot the impedances on an interpolated frequency axis, if the table
+        given to induced_voltage_freq has been interpolated on a new axis
+
+    Additionally, either outputs the plots in the console, or save them on disk.
+
+    :param induced_voltage_freq:  InducedVoltageFreq object
+        InducedVoltageFreq object which impedances are to be plotted
+    :param figure_index:  int
+        Index of the figure, e.g. the turn number if creating a plot for every turn
+        Default: 0
+    :param plot_total_impedance: bool
+        If True, plots the total, summed impedance which is calculated internally in the
+        induced_voltage_freq.
+        If False, plots the raw impedances from the table given to induced_voltage_freq
+        Default: True
+    :param plot_spectrum:  bool
+        If True, plots the beam spectrum in addition to the impedances.
+        If False, no beam spectrum is plotted
+        Default: False
+    :param plot_interpolated_impedances: bool
+        If True, the impedance table in induced_voltage_freq must have been interpolated
+        to a new frequency axis. Will then plot the impedances on that new frequency axis.
+        If False, plots the raw impedances on the original frequency axis
+        given by the table in induced_voltage_freq
+        Default: False
+    :param style: str
+        Matplotlib plot style string
+        Default: "-"
+    :param cut_left_right: 2-tuple of floats
+        x-axis limits of the plots (frequency in Hz)
+        Default: None
+    :param cut_up_down: 2-tuple of floats
+        y-axis limits of the plots (resistance in Ohm)
+        Default: None
+    :param dirname: str
+        Path to directory in which the plots are to be saved.
+        Default: "fig"
+    :param show_plot: bool
+        If True, will output the plots directly in the console
+        If False, will save the plots to directory given in dirname
     """
 
-    if option1 == "sum":
+    if plot_total_impedance:
 
         ax1 = plt.subplots()[1]
-        ax1.plot(ind_volt_from_imp.frequency_array,
-                 ind_volt_from_imp.total_impedance.real, style)
-        ax1.plot(ind_volt_from_imp.frequency_array,
-                 ind_volt_from_imp.total_impedance.imag, style)
-        if option2 == "spectrum":
+        # Renormalize Impedances to the cont. Fourier Transform values, i.e. Ohm.
+        # In the InducedVoltage Objects they are normalized to the DFT values by deviding
+        # by bin size. This is reverted here when plotting to get sensible units
+
+        ax1.plot(induced_voltage_freq.freq,
+                 induced_voltage_freq.total_impedance.real * induced_voltage_freq.profile.bin_size, style, label = "Real Impedance")
+        ax1.plot(induced_voltage_freq.freq,
+                 induced_voltage_freq.total_impedance.imag * induced_voltage_freq.profile.bin_size, style, label = "Imaginary Impedance")
+
+        ax1.set_xlim(cut_left_right)
+        ax1.set_ylim(cut_up_down)
+        ax1.set_xlabel("Frequency [Hz]")
+        ax1.set_ylabel("Impedance [Ohm]")
+        if plot_spectrum:
             ax2 = ax1.twinx()
-            ax2.plot(ind_volt_from_imp.profile.beam_spectrum_freq,
-                     np.abs(ind_volt_from_imp.profile.beam_spectrum))
-        fign = dirname + '/sum_imp_vs_freq_fft' "%d" % counter + '.png'
-        plt.show()
-        plt.savefig(fign)
-        plt.clf()
+            ax2.plot(induced_voltage_freq.profile.beam_spectrum_freq,
+                     np.abs(induced_voltage_freq.profile.beam_spectrum), label = "Beam Spectrum")
+
+            ax2.set_xlim(cut_left_right)
+            ax2.set_ylim(cut_up_down)
+            ax2.set_ylabel("Beam Spectrum [a.u.]")
+        if plot_spectrum:
+            ax2.legend()
+        else:
+            ax1.legend()
+        fign = dirname + '/sum_imp_vs_freq_fft' "%d" % figure_index + '.png'
+
+        if show_plot:
+            plt.show()
+        else:
+            plt.savefig(fign)
+            plt.clf()
 
-    elif option1 == "single":
+    else:
 
         fig0 = plt.figure(0)
         ax0 = fig0.add_subplot(111)
 
         fig1 = plt.figure(1)
         ax1 = fig1.add_subplot(111)
-
-        for i in range(len(ind_volt_from_imp.impedance_source_list)):
-            if isinstance(ind_volt_from_imp.impedance_source_list[i],
-                          InputTable) and option3 == "freq_table":
-                ax0.plot(ind_volt_from_imp.impedance_source_list[i].frequency_array_loaded,
-                         ind_volt_from_imp.impedance_source_list[i].Re_Z_array_loaded, style)
+        # The impedance sources themselves are properly normalized already.
+        for i in range(len(induced_voltage_freq.impedance_source_list)):
+            if isinstance(induced_voltage_freq.impedance_source_list[i],
+                          InputTable) and not plot_interpolated_impedances:
+                ax0.plot(induced_voltage_freq.impedance_source_list[i].frequency_array_loaded,
+                         induced_voltage_freq.impedance_source_list[i].Re_Z_array_loaded, style)
                 ax0.set_xlim(cut_left_right)
                 ax0.set_ylim(cut_up_down)
-                ax1.plot(ind_volt_from_imp.impedance_source_list[i].frequency_array_loaded,
-                         ind_volt_from_imp.impedance_source_list[i].Im_Z_array_loaded, style)
+                ax1.plot(induced_voltage_freq.impedance_source_list[i].frequency_array_loaded,
+                         induced_voltage_freq.impedance_source_list[i].Im_Z_array_loaded, style)
                 ax1.set_xlim(cut_left_right)
                 ax1.set_ylim(cut_up_down)
-            else:
-                ax0.plot(ind_volt_from_imp.frequency_array,
-                         ind_volt_from_imp.impedance_source_list[i].impedance.real, style)
+            elif plot_interpolated_impedances:
+                ax0.plot(induced_voltage_freq.impedance_source_list[i].frequency_array,
+                         induced_voltage_freq.impedance_source_list[i].impedance.real, style)
                 ax0.set_xlim(cut_left_right)
                 ax0.set_ylim(cut_up_down)
-                ax1.plot(ind_volt_from_imp.frequency_array,
-                         ind_volt_from_imp.impedance_source_list[i].impedance.imag, style)
+                ax1.plot(induced_voltage_freq.impedance_source_list[i].frequency_array,
+                         induced_voltage_freq.impedance_source_list[i].impedance.imag, style)
                 ax1.set_xlim(cut_left_right)
                 ax1.set_ylim(cut_up_down)
 
-        fign1 = dirname + '/real_imp_vs_' + option3 + '_' "%d" % counter + '.png'
-        if option2 == "spectrum":
+        if plot_interpolated_impedances:
+            fig_suffix = "interpolated_freq"
+        else:
+            fig_suffix = "table_freq"
+
+        ax0.set_xlabel("Frequency [Hz]")
+        ax0.set_ylabel("Real impedance [Ohm]")
+
+        ax1.set_xlabel("Frequency [Hz]")
+        ax1.set_ylabel("Imaginary Impedance [Ohm]")
+
+
+        fign1 = dirname + '/real_imp_vs_' + fig_suffix + '_' "%d" % figure_index + '.png'
+        if plot_spectrum:
             ax2 = ax0.twinx()
-            ax2.plot(ind_volt_from_imp.profile.beam_spectrum_freq,
-                     np.abs(ind_volt_from_imp.profile.beam_spectrum))
+            spectrum, = ax2.plot(induced_voltage_freq.profile.beam_spectrum_freq,
+                     np.abs(induced_voltage_freq.profile.beam_spectrum), label = "Beam Spectrum",color = "r")
             ax2.set_xlim(cut_left_right)
             ax2.set_ylim(cut_up_down)
-        ax0.set_xlabel("Frequency [Hz]")
-        ax0.set_ylabel("Real impedance [Ohm]")
+            ax2.set_ylabel("Beam Spectrum [a.u.]")
+            ax2.legend(handles=[spectrum])
+
+
+
         plt.figure(0)
+        if show_plot:
+            plt.show()
+        else:
+            plt.savefig(fign1)
+            plt.clf()
+        fign2 = dirname + '/imag_imp_vs_' + fig_suffix + '_' "%d" % figure_index + '.png'
 
-        plt.savefig(fign1)
-        plt.clf()
-        fign2 = dirname + '/imag_imp_vs_' + option3 + '_' "%d" % counter + '.png'
 
-        if option2 == "spectrum":
+        if plot_spectrum:
             ax3 = ax1.twinx()
-            ax3.plot(ind_volt_from_imp.profile.beam_spectrum_freq,
-                     np.abs(ind_volt_from_imp.profile.beam_spectrum))
+            spectrum, = ax3.plot(induced_voltage_freq.profile.beam_spectrum_freq,
+                     np.abs(induced_voltage_freq.profile.beam_spectrum), label = "Beam Spectrum", color = "r")
             ax3.set_xlim(cut_left_right)
             ax3.set_ylim(cut_up_down)
-        plt.figure(1)
+            ax3.set_ylabel("Beam Spectrum [a.u.]")
 
-        plt.savefig(fign2)
-        plt.clf()
+            ax3.legend(handles=[spectrum])
+        plt.figure(1)
+        if show_plot:
+            plt.show()
+        else:
+            plt.savefig(fign2)
+            plt.clf()
 
 
-def plot_induced_voltage_vs_bin_centers(counter, general_params,
-                                        total_voltage, style='-',
-                                        dirname='fig'):
+def plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='-', figure_index = 0,
+                                        dirname='fig', show_plot = False):
     """
-    Plot of induced voltage vs bin centers.
+    Plots the total induced voltage calculated in the TotalInducedVoltage object given.
+
+    :param total_induced_voltage: TotalInducedVoltage object
+        TotalInducedVoltage object containing the total induced voltage to be plotted
+    :param style: str
+        Matplotlib style string
+        Default: "-"
+    :param figure_index: int
+        Index of the figure, e.g. the turn number if creating a plot for every turn
+        Default: 0
+    :param dirname: str
+        Directory to save the plots to
+        Default: "fig"
+    :param show_plot: bool
+        If True, outputs the plots directly to console
+        If False, saved the plots to disk in given directory
+        Default: False
     """
 
     fig0 = plt.figure(0)
     fig0.set_size_inches(8, 6)
-    ax0 = plt.axes([0.15, 0.1, 0.8, 0.8])
-    plt.plot(total_voltage.profile.bin_centers, total_voltage.induced_voltage, style)
+    ax0 = plt.axes()
+    plt.plot(total_induced_voltage.profile.bin_centers, total_induced_voltage.induced_voltage, style)
     ax0.set_xlabel("Time [s]")
     ax0.set_ylabel("Induced voltage [V]")
 
     # Save plot
-    fign = dirname + '/induced_voltage_' "%d" % counter + '.png'
-    plt.savefig(fign)
-    plt.clf()
+    fign = dirname + '/induced_voltage_' "%d" % figure_index + '.png'
+    if show_plot:
+        plt.show()
+    else:
+        plt.savefig(fign)
+        plt.clf()
-- 
GitLab


From dfe5d81cb21578217ce09891c3234842063841dc Mon Sep 17 00:00:00 2001
From: onaumenk <oleksandr.naumenko@cern.ch>
Date: Wed, 14 Jun 2023 12:08:39 +0200
Subject: [PATCH 3/6] Propagated changes in plot_impedance.py to the examples.
 Increased the time resolution such that plotting the beam spectrum, if
 desired, yields sensible results. Added frequency limits to have good looking
 plots Removed beam spectrum plot as it is not properly visible if wishing to
 resolve the full impedance.

---
 .../gpu_main_files/EX_02_Main_long_ps_booster.py     | 11 +++++++----
 __EXAMPLES/main_files/EX_02_Main_long_ps_booster.py  |  8 ++++----
 .../mpi_main_files/EX_02_Main_long_ps_booster.py     | 12 +++++++-----
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py b/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py
index 4da7171c..5648b6f4 100644
--- a/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py
+++ b/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py
@@ -99,7 +99,7 @@ bigaussian(general_params, RF_sct_par, my_beam, sigma_dt, seed=1)
 
 # DEFINE SLICES----------------------------------------------------------------
 slice_beam = Profile(my_beam, CutOptions(cut_left=-5.72984173562e-7,
-                                         cut_right=5.72984173562e-7, n_slices=100))
+                                         cut_right=5.72984173562e-7, n_slices=10000))
 
 # MONITOR----------------------------------------------------------------------
 
@@ -216,10 +216,13 @@ for i in range(1, n_turns + 1):
             ind_volt_freq.to_cpu()
             slice_beam.to_cpu()
 
-        plot_impedance_vs_frequency(i, general_params, ind_volt_freq,
-                                    option1="single", style='-', option3="freq_table", option2="spectrum", dirname=this_directory + '../gpu_output_files/EX_02_fig')
+        plot_impedance_vs_frequency(ind_volt_freq, figure_index=i, cut_up_down=(0, 1000), cut_left_right=(0, 3e9),
+                                    show_plot=False,
+                                    plot_total_impedance=False, style='-', plot_interpolated_impedances=False,
+                                    plot_spectrum=False, dirname=this_directory + '../gpu_output_files/EX_02_fig')
 
-        plot_induced_voltage_vs_bin_centers(i, general_params, total_induced_voltage, style='.', dirname=this_directory + '../gpu_output_files/EX_02_fig')
+        plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='.',
+                                            dirname=this_directory + '../gpu_output_files/EX_02_fig', show_plot=False)
 
         if USE_GPU:
             bm.use_gpu()
diff --git a/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py b/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py
index 4bde1c76..44afbc9f 100644
--- a/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py
+++ b/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py
@@ -90,7 +90,7 @@ bigaussian(general_params, RF_sct_par, my_beam, sigma_dt, seed=1)
 
 # DEFINE SLICES----------------------------------------------------------------
 slice_beam = Profile(my_beam, CutOptions(cut_left=-5.72984173562e-7,
-                                         cut_right=5.72984173562e-7, n_slices=100))
+                                         cut_right=5.72984173562e-7, n_slices=10000))
 
 # MONITOR----------------------------------------------------------------------
 
@@ -183,10 +183,10 @@ for i in range(1, n_turns + 1):
     # Plots
     if (i % n_turns_between_two_plots) == 0:
 
-        plot_impedance_vs_frequency(i, general_params, ind_volt_freq,
-                                    option1="single", style='-', option3="freq_table", option2="spectrum", dirname=this_directory + '../output_files/EX_02_fig')
+        plot_impedance_vs_frequency(ind_volt_freq, figure_index=i, cut_up_down=(0,1000), cut_left_right=(0,3e9), show_plot=False,
+                                    plot_total_impedance=False, style='-', plot_interpolated_impedances=False, plot_spectrum=False, dirname=this_directory + '../output_files/EX_02_fig')
 
-        plot_induced_voltage_vs_bin_centers(i, general_params, total_induced_voltage, style='.', dirname=this_directory + '../output_files/EX_02_fig')
+        plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='.', dirname=this_directory + '../output_files/EX_02_fig',  show_plot=False)
 
 # For testing purposes
 test_string += '{:+10.10e}\t{:+10.10e}\t{:+10.10e}\t{:+10.10e}\n'.format(
diff --git a/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py b/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py
index 4067e66b..f4a4709b 100644
--- a/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py
+++ b/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py
@@ -96,7 +96,7 @@ bigaussian(general_params, RF_sct_par, my_beam, sigma_dt, seed=1)
 
 # DEFINE SLICES----------------------------------------------------------------
 slice_beam = Profile(my_beam, CutOptions(cut_left=-5.72984173562e-7,
-                                         cut_right=5.72984173562e-7, n_slices=100))
+                                         cut_right=5.72984173562e-7, n_slices=10000))
 
 
 # LOAD IMPEDANCE TABLES--------------------------------------------------------
@@ -191,11 +191,13 @@ for i in range(1, n_turns + 1):
 
     # Plots
     if (i % n_turns_between_two_plots) == 0:
+        plot_impedance_vs_frequency(ind_volt_freq, figure_index=i, cut_up_down=(0, 1000), cut_left_right=(0, 3e9),
+                                    show_plot=False,
+                                    plot_total_impedance=False, style='-', plot_interpolated_impedances=False,
+                                    plot_spectrum=False, dirname=this_directory + '../mpi_output_files/EX_02_fig')
 
-        plot_impedance_vs_frequency(i, general_params, ind_volt_freq,
-                                    option1="single", style='-', option3="freq_table", option2="spectrum", dirname=this_directory + '../mpi_output_files/EX_02_fig')
-
-        plot_induced_voltage_vs_bin_centers(i, general_params, total_induced_voltage, style='.', dirname=this_directory + '../mpi_output_files/EX_02_fig')
+        plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='.',
+                                            dirname=this_directory + '../mpi_output_files/EX_02_fig', show_plot=False)
 
 my_beam.gather()
 WORKER.finalize()
-- 
GitLab


From 205175d9e8617296597ad69206056ebc34efc9c6 Mon Sep 17 00:00:00 2001
From: onaumenk <oleksandr.naumenko@cern.ch>
Date: Wed, 14 Jun 2023 12:15:05 +0200
Subject: [PATCH 4/6] Propagated changes from plot_impedance.py to Example 5

---
 __EXAMPLES/gpu_main_files/EX_05_Wake_impedance.py | 10 +++++-----
 __EXAMPLES/main_files/EX_05_Wake_impedance.py     | 10 +++++-----
 __EXAMPLES/mpi_main_files/EX_05_Wake_impedance.py | 10 +++++-----
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/__EXAMPLES/gpu_main_files/EX_05_Wake_impedance.py b/__EXAMPLES/gpu_main_files/EX_05_Wake_impedance.py
index 126974e0..8cf113cd 100644
--- a/__EXAMPLES/gpu_main_files/EX_05_Wake_impedance.py
+++ b/__EXAMPLES/gpu_main_files/EX_05_Wake_impedance.py
@@ -278,12 +278,12 @@ for i in np.arange(1, n_turns + 1):
             tot_vol_freq.to_cpu()
             tot_vol_res.to_cpu()
 
-        plot_induced_voltage_vs_bin_centers(i, general_params, tot_vol,
+        plot_induced_voltage_vs_bin_centers(tot_vol, figure_index=i,
                                             style='.', dirname=this_directory + '../gpu_output_files/EX_05_fig/1')
-        plot_induced_voltage_vs_bin_centers(i, general_params_freq,
-                                            tot_vol_freq, style='.', dirname=this_directory + '../gpu_output_files/EX_05_fig/2')
-        plot_induced_voltage_vs_bin_centers(i, general_params_res,
-                                            tot_vol_res, style='.', dirname=this_directory + '../gpu_output_files/EX_05_fig/3')
+        plot_induced_voltage_vs_bin_centers(tot_vol_freq, figure_index=i,
+                                            style='.', dirname=this_directory + '../gpu_output_files/EX_05_fig/2')
+        plot_induced_voltage_vs_bin_centers(tot_vol_res, figure_index=i,
+                                            style='.', dirname=this_directory + '../gpu_output_files/EX_05_fig/3')
 
         if USE_GPU:
             bm.use_gpu()
diff --git a/__EXAMPLES/main_files/EX_05_Wake_impedance.py b/__EXAMPLES/main_files/EX_05_Wake_impedance.py
index 50408f0a..202be93c 100644
--- a/__EXAMPLES/main_files/EX_05_Wake_impedance.py
+++ b/__EXAMPLES/main_files/EX_05_Wake_impedance.py
@@ -226,12 +226,12 @@ for i in np.arange(1, n_turns + 1):
 
     # Plots
     if (i % dt_plt) == 0:
-        plot_induced_voltage_vs_bin_centers(i, general_params, tot_vol,
+        plot_induced_voltage_vs_bin_centers(tot_vol, figure_index=i,
                                             style='.', dirname=this_directory + '../output_files/EX_05_fig/1')
-        plot_induced_voltage_vs_bin_centers(i, general_params_freq,
-                                            tot_vol_freq, style='.', dirname=this_directory + '../output_files/EX_05_fig/2')
-        plot_induced_voltage_vs_bin_centers(i, general_params_res,
-                                            tot_vol_res, style='.', dirname=this_directory + '../output_files/EX_05_fig/3')
+        plot_induced_voltage_vs_bin_centers(tot_vol_freq, figure_index=i,
+                                            style='.', dirname=this_directory + '../output_files/EX_05_fig/2')
+        plot_induced_voltage_vs_bin_centers(tot_vol_res, figure_index=i,
+                                            style='.', dirname=this_directory + '../output_files/EX_05_fig/3')
 
 # Plotting induced voltages---------------------------------------------------
 plt.clf()
diff --git a/__EXAMPLES/mpi_main_files/EX_05_Wake_impedance.py b/__EXAMPLES/mpi_main_files/EX_05_Wake_impedance.py
index 591fb263..19cb333c 100644
--- a/__EXAMPLES/mpi_main_files/EX_05_Wake_impedance.py
+++ b/__EXAMPLES/mpi_main_files/EX_05_Wake_impedance.py
@@ -236,12 +236,12 @@ for i in np.arange(1, n_turns + 1):
 
     # Plots
     if (i % dt_plt) == 0 and (WORKER.is_master):
-        plot_induced_voltage_vs_bin_centers(i, general_params, tot_vol,
+        plot_induced_voltage_vs_bin_centers(tot_vol, figure_index=i,
                                             style='.', dirname=this_directory + '../mpi_output_files/EX_05_fig/1')
-        plot_induced_voltage_vs_bin_centers(i, general_params_freq,
-                                            tot_vol_freq, style='.', dirname=this_directory + '../mpi_output_files/EX_05_fig/2')
-        plot_induced_voltage_vs_bin_centers(i, general_params_res,
-                                            tot_vol_res, style='.', dirname=this_directory + '../mpi_output_files/EX_05_fig/3')
+        plot_induced_voltage_vs_bin_centers(tot_vol_freq, figure_index=i,
+                                            style='.', dirname=this_directory + '../mpi_output_files/EX_05_fig/2')
+        plot_induced_voltage_vs_bin_centers(tot_vol_res, figure_index=i,
+                                            style='.', dirname=this_directory + '../mpi_output_files/EX_05_fig/3')
 
 my_beam.gather()
 my_beam_freq.gather()
-- 
GitLab


From 86ce93287f51ef590d9e3e75a2d108319041add0 Mon Sep 17 00:00:00 2001
From: onaumenk <oleksandr.naumenko@cern.ch>
Date: Wed, 14 Jun 2023 16:25:00 +0200
Subject: [PATCH 5/6] corrected an inconsistency: The usage of kwarg show_plot
 was inconsistent between the plot.py module (where the kwarg was show_plots)
 and plot_impedance.py (where the kwarg was show_plot). Standardised to
 show_plots.

---
 .../gpu_main_files/EX_02_Main_long_ps_booster.py |  4 ++--
 .../main_files/EX_02_Main_long_ps_booster.py     |  4 ++--
 .../mpi_main_files/EX_02_Main_long_ps_booster.py |  4 ++--
 blond/plots/plot_impedance.py                    | 16 ++++++++--------
 4 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py b/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py
index 5648b6f4..16745214 100644
--- a/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py
+++ b/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py
@@ -217,12 +217,12 @@ for i in range(1, n_turns + 1):
             slice_beam.to_cpu()
 
         plot_impedance_vs_frequency(ind_volt_freq, figure_index=i, cut_up_down=(0, 1000), cut_left_right=(0, 3e9),
-                                    show_plot=False,
+                                    show_plots=False,
                                     plot_total_impedance=False, style='-', plot_interpolated_impedances=False,
                                     plot_spectrum=False, dirname=this_directory + '../gpu_output_files/EX_02_fig')
 
         plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='.',
-                                            dirname=this_directory + '../gpu_output_files/EX_02_fig', show_plot=False)
+                                            dirname=this_directory + '../gpu_output_files/EX_02_fig', show_plots=False)
 
         if USE_GPU:
             bm.use_gpu()
diff --git a/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py b/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py
index 44afbc9f..859921e9 100644
--- a/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py
+++ b/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py
@@ -183,10 +183,10 @@ for i in range(1, n_turns + 1):
     # Plots
     if (i % n_turns_between_two_plots) == 0:
 
-        plot_impedance_vs_frequency(ind_volt_freq, figure_index=i, cut_up_down=(0,1000), cut_left_right=(0,3e9), show_plot=False,
+        plot_impedance_vs_frequency(ind_volt_freq, figure_index=i, cut_up_down=(0,1000), cut_left_right=(0,3e9), show_plots=False,
                                     plot_total_impedance=False, style='-', plot_interpolated_impedances=False, plot_spectrum=False, dirname=this_directory + '../output_files/EX_02_fig')
 
-        plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='.', dirname=this_directory + '../output_files/EX_02_fig',  show_plot=False)
+        plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='.', dirname=this_directory + '../output_files/EX_02_fig', show_plots=False)
 
 # For testing purposes
 test_string += '{:+10.10e}\t{:+10.10e}\t{:+10.10e}\t{:+10.10e}\n'.format(
diff --git a/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py b/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py
index f4a4709b..2cec9da3 100644
--- a/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py
+++ b/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py
@@ -192,12 +192,12 @@ for i in range(1, n_turns + 1):
     # Plots
     if (i % n_turns_between_two_plots) == 0:
         plot_impedance_vs_frequency(ind_volt_freq, figure_index=i, cut_up_down=(0, 1000), cut_left_right=(0, 3e9),
-                                    show_plot=False,
+                                    show_plots=False,
                                     plot_total_impedance=False, style='-', plot_interpolated_impedances=False,
                                     plot_spectrum=False, dirname=this_directory + '../mpi_output_files/EX_02_fig')
 
         plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='.',
-                                            dirname=this_directory + '../mpi_output_files/EX_02_fig', show_plot=False)
+                                            dirname=this_directory + '../mpi_output_files/EX_02_fig', show_plots=False)
 
 my_beam.gather()
 WORKER.finalize()
diff --git a/blond/plots/plot_impedance.py b/blond/plots/plot_impedance.py
index 1963cf8d..5a1eac3f 100644
--- a/blond/plots/plot_impedance.py
+++ b/blond/plots/plot_impedance.py
@@ -33,7 +33,7 @@ def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
                                 plot_total_impedance = True, plot_spectrum=False,
                                 plot_interpolated_impedances = False, style='-',
                                 cut_left_right=None, cut_up_down=None,
-                                dirname='fig', show_plot = False):
+                                dirname='fig', show_plots = False):
     """
     Plots Impedance in frequency domain. Given an InducedVoltageFreq object, i.e. one
     that was created from a list of impedances:
@@ -76,7 +76,7 @@ def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
     :param dirname: str
         Path to directory in which the plots are to be saved.
         Default: "fig"
-    :param show_plot: bool
+    :param show_plots: bool
         If True, will output the plots directly in the console
         If False, will save the plots to directory given in dirname
     """
@@ -111,7 +111,7 @@ def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
             ax1.legend()
         fign = dirname + '/sum_imp_vs_freq_fft' "%d" % figure_index + '.png'
 
-        if show_plot:
+        if show_plots:
             plt.show()
         else:
             plt.savefig(fign)
@@ -171,7 +171,7 @@ def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
 
 
         plt.figure(0)
-        if show_plot:
+        if show_plots:
             plt.show()
         else:
             plt.savefig(fign1)
@@ -189,7 +189,7 @@ def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
 
             ax3.legend(handles=[spectrum])
         plt.figure(1)
-        if show_plot:
+        if show_plots:
             plt.show()
         else:
             plt.savefig(fign2)
@@ -197,7 +197,7 @@ def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
 
 
 def plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='-', figure_index = 0,
-                                        dirname='fig', show_plot = False):
+                                        dirname='fig', show_plots = False):
     """
     Plots the total induced voltage calculated in the TotalInducedVoltage object given.
 
@@ -212,7 +212,7 @@ def plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='-', figure
     :param dirname: str
         Directory to save the plots to
         Default: "fig"
-    :param show_plot: bool
+    :param show_plots: bool
         If True, outputs the plots directly to console
         If False, saved the plots to disk in given directory
         Default: False
@@ -227,7 +227,7 @@ def plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='-', figure
 
     # Save plot
     fign = dirname + '/induced_voltage_' "%d" % figure_index + '.png'
-    if show_plot:
+    if show_plots:
         plt.show()
     else:
         plt.savefig(fign)
-- 
GitLab


From ffe0197aead2d71e2789f0c45a7a2c7a595bcd9f Mon Sep 17 00:00:00 2001
From: onaumenk <oleksandr.naumenko@cern.ch>
Date: Wed, 14 Jun 2023 17:10:55 +0200
Subject: [PATCH 6/6] Some more naming cleanup in Ex.2 and added a
 plot_wake_vs_time function.

---
 .../EX_02_Main_long_ps_booster.py             |  18 +--
 .../main_files/EX_02_Main_long_ps_booster.py  |  18 +--
 .../EX_02_Main_long_ps_booster.py             |  18 +--
 blond/plots/plot_impedance.py                 | 114 +++++++++++++++++-
 4 files changed, 136 insertions(+), 32 deletions(-)

diff --git a/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py b/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py
index 16745214..727bac64 100644
--- a/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py
+++ b/__EXAMPLES/gpu_main_files/EX_02_Main_long_ps_booster.py
@@ -84,18 +84,18 @@ phi_offset = np.pi
 
 # DEFINE RING------------------------------------------------------------------
 
-general_params = Ring(C, momentum_compaction, sync_momentum,
-                      Proton(), n_turns)
+ring = Ring(C, momentum_compaction, sync_momentum,
+            Proton(), n_turns)
 
-RF_sct_par = RFStation(general_params, [harmonic_numbers],
+RF_sct_par = RFStation(ring, [harmonic_numbers],
                        [voltage_program], [phi_offset], n_rf_systems)
 
-my_beam = Beam(general_params, n_macroparticles, n_particles)
+my_beam = Beam(ring, n_macroparticles, n_particles)
 
 ring_RF_section = RingAndRFTracker(RF_sct_par, my_beam)
 
 # DEFINE BEAM------------------------------------------------------------------
-bigaussian(general_params, RF_sct_par, my_beam, sigma_dt, seed=1)
+bigaussian(ring, RF_sct_par, my_beam, sigma_dt, seed=1)
 
 # DEFINE SLICES----------------------------------------------------------------
 slice_beam = Profile(my_beam, CutOptions(cut_left=-5.72984173562e-7,
@@ -103,7 +103,7 @@ slice_beam = Profile(my_beam, CutOptions(cut_left=-5.72984173562e-7,
 
 # MONITOR----------------------------------------------------------------------
 
-bunchmonitor = BunchMonitor(general_params, RF_sct_par, my_beam,
+bunchmonitor = BunchMonitor(ring, RF_sct_par, my_beam,
                             this_directory + '../gpu_output_files/EX_02_output_data',
                             buffer_time=1)
 
@@ -149,10 +149,10 @@ else:
 
 # steps
 steps = InductiveImpedance(my_beam, slice_beam, 34.6669349520904 / 10e9 *
-                           general_params.f_rev, RF_sct_par, deriv_mode='diff')
+                           ring.f_rev, RF_sct_par, deriv_mode='diff')
 # direct space charge
 dir_space_charge = InductiveImpedance(my_beam, slice_beam, -376.730313462
-                                      / (general_params.beta[0] * general_params.gamma[0]**2),
+                                      / (ring.beta[0] * ring.gamma[0] ** 2),
                                       RF_sct_par)
 
 
@@ -171,7 +171,7 @@ total_induced_voltage = TotalInducedVoltage(my_beam, slice_beam,
 
 format_options = {'dirname': this_directory +
                   '../gpu_output_files/EX_02_fig', 'linestyle': '.'}
-plots = Plot(general_params, RF_sct_par, my_beam, 1, n_turns, 0,
+plots = Plot(ring, RF_sct_par, my_beam, 1, n_turns, 0,
              5.72984173562e-7, - my_beam.sigma_dE * 4.2, my_beam.sigma_dE * 4.2, xunit='s',
              separatrix_plot=True, Profile=slice_beam, h5file=this_directory + '../gpu_output_files/EX_02_output_data',
              histograms_plot=True, format_options=format_options)
diff --git a/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py b/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py
index 859921e9..95e49599 100644
--- a/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py
+++ b/__EXAMPLES/main_files/EX_02_Main_long_ps_booster.py
@@ -75,18 +75,18 @@ phi_offset = np.pi
 
 # DEFINE RING------------------------------------------------------------------
 
-general_params = Ring(C, momentum_compaction, sync_momentum,
-                      Proton(), n_turns)
+ring = Ring(C, momentum_compaction, sync_momentum,
+            Proton(), n_turns)
 
-RF_sct_par = RFStation(general_params, [harmonic_numbers],
+RF_sct_par = RFStation(ring, [harmonic_numbers],
                        [voltage_program], [phi_offset], n_rf_systems)
 
-my_beam = Beam(general_params, n_macroparticles, n_particles)
+my_beam = Beam(ring, n_macroparticles, n_particles)
 
 ring_RF_section = RingAndRFTracker(RF_sct_par, my_beam)
 
 # DEFINE BEAM------------------------------------------------------------------
-bigaussian(general_params, RF_sct_par, my_beam, sigma_dt, seed=1)
+bigaussian(ring, RF_sct_par, my_beam, sigma_dt, seed=1)
 
 # DEFINE SLICES----------------------------------------------------------------
 slice_beam = Profile(my_beam, CutOptions(cut_left=-5.72984173562e-7,
@@ -94,7 +94,7 @@ slice_beam = Profile(my_beam, CutOptions(cut_left=-5.72984173562e-7,
 
 # MONITOR----------------------------------------------------------------------
 
-bunchmonitor = BunchMonitor(general_params, RF_sct_par, my_beam,
+bunchmonitor = BunchMonitor(ring, RF_sct_par, my_beam,
                             this_directory + '../output_files/EX_02_output_data', buffer_time=1)
 
 # LOAD IMPEDANCE TABLES--------------------------------------------------------
@@ -133,10 +133,10 @@ else:
 
 # steps
 steps = InductiveImpedance(my_beam, slice_beam, 34.6669349520904 / 10e9 *
-                           general_params.f_rev, RF_sct_par, deriv_mode='diff')
+                           ring.f_rev, RF_sct_par, deriv_mode='diff')
 # direct space charge
 dir_space_charge = InductiveImpedance(my_beam, slice_beam, -376.730313462
-                                      / (general_params.beta[0] * general_params.gamma[0]**2),
+                                      / (ring.beta[0] * ring.gamma[0] ** 2),
                                       RF_sct_par)
 
 
@@ -154,7 +154,7 @@ total_induced_voltage = TotalInducedVoltage(my_beam, slice_beam,
 # PLOTS
 
 format_options = {'dirname': this_directory + '../output_files/EX_02_fig', 'linestyle': '.'}
-plots = Plot(general_params, RF_sct_par, my_beam, 1, n_turns, 0,
+plots = Plot(ring, RF_sct_par, my_beam, 1, n_turns, 0,
              5.72984173562e-7, - my_beam.sigma_dE * 4.2, my_beam.sigma_dE * 4.2, xunit='s',
              separatrix_plot=True, Profile=slice_beam, h5file=this_directory + '../output_files/EX_02_output_data',
              histograms_plot=True, format_options=format_options)
diff --git a/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py b/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py
index 2cec9da3..a3b147df 100644
--- a/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py
+++ b/__EXAMPLES/mpi_main_files/EX_02_Main_long_ps_booster.py
@@ -81,18 +81,18 @@ phi_offset = np.pi
 
 # DEFINE RING------------------------------------------------------------------
 
-general_params = Ring(C, momentum_compaction, sync_momentum,
-                      Proton(), n_turns)
+ring = Ring(C, momentum_compaction, sync_momentum,
+            Proton(), n_turns)
 
-RF_sct_par = RFStation(general_params, [harmonic_numbers],
+RF_sct_par = RFStation(ring, [harmonic_numbers],
                        [voltage_program], [phi_offset], n_rf_systems)
 
-my_beam = Beam(general_params, n_macroparticles, n_particles)
+my_beam = Beam(ring, n_macroparticles, n_particles)
 
 ring_RF_section = RingAndRFTracker(RF_sct_par, my_beam)
 
 # DEFINE BEAM------------------------------------------------------------------
-bigaussian(general_params, RF_sct_par, my_beam, sigma_dt, seed=1)
+bigaussian(ring, RF_sct_par, my_beam, sigma_dt, seed=1)
 
 # DEFINE SLICES----------------------------------------------------------------
 slice_beam = Profile(my_beam, CutOptions(cut_left=-5.72984173562e-7,
@@ -135,10 +135,10 @@ else:
 
 # steps
 steps = InductiveImpedance(my_beam, slice_beam, 34.6669349520904 / 10e9 *
-                           general_params.f_rev, RF_sct_par, deriv_mode='diff')
+                           ring.f_rev, RF_sct_par, deriv_mode='diff')
 # direct space charge
 dir_space_charge = InductiveImpedance(my_beam, slice_beam, -376.730313462
-                                      / (general_params.beta[0] * general_params.gamma[0]**2),
+                                      / (ring.beta[0] * ring.gamma[0] ** 2),
                                       RF_sct_par)
 
 
@@ -159,13 +159,13 @@ map_ = [total_induced_voltage] + [ring_RF_section] + [slice_beam]
 
 if WORKER.is_master:
     # MONITOR----------------------------------------------------------------------
-    bunchmonitor = BunchMonitor(general_params, RF_sct_par, my_beam,
+    bunchmonitor = BunchMonitor(ring, RF_sct_par, my_beam,
                                 this_directory + '../mpi_output_files/EX_02_output_data', buffer_time=1)
 
     # PLOTS
 
     format_options = {'dirname': this_directory + '../mpi_output_files/EX_02_fig', 'linestyle': '.'}
-    plots = Plot(general_params, RF_sct_par, my_beam, 1, n_turns, 0,
+    plots = Plot(ring, RF_sct_par, my_beam, 1, n_turns, 0,
                  5.72984173562e-7, - my_beam.sigma_dE * 4.2, my_beam.sigma_dE * 4.2, xunit='s',
                  separatrix_plot=True, Profile=slice_beam, h5file=this_directory + '../mpi_output_files/EX_02_output_data',
                  histograms_plot=True, format_options=format_options)
diff --git a/blond/plots/plot_impedance.py b/blond/plots/plot_impedance.py
index 5a1eac3f..825f9134 100644
--- a/blond/plots/plot_impedance.py
+++ b/blond/plots/plot_impedance.py
@@ -40,7 +40,7 @@ def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
      - can plot either the total, summed impedance or each individual one from the list
      - can additionally plot the beam spectrum
      - and can plot the impedances on an interpolated frequency axis, if the table
-        given to induced_voltage_freq has been interpolated on a new axis
+        given to induced_voltage_time has been interpolated on a new axis
 
     Additionally, either outputs the plots in the console, or save them on disk.
 
@@ -51,18 +51,18 @@ def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
         Default: 0
     :param plot_total_impedance: bool
         If True, plots the total, summed impedance which is calculated internally in the
-        induced_voltage_freq.
-        If False, plots the raw impedances from the table given to induced_voltage_freq
+        induced_voltage_time.
+        If False, plots the raw impedances from the table given to induced_voltage_time
         Default: True
     :param plot_spectrum:  bool
         If True, plots the beam spectrum in addition to the impedances.
         If False, no beam spectrum is plotted
         Default: False
     :param plot_interpolated_impedances: bool
-        If True, the impedance table in induced_voltage_freq must have been interpolated
+        If True, the impedance table in induced_voltage_time must have been interpolated
         to a new frequency axis. Will then plot the impedances on that new frequency axis.
         If False, plots the raw impedances on the original frequency axis
-        given by the table in induced_voltage_freq
+        given by the table in induced_voltage_time
         Default: False
     :param style: str
         Matplotlib plot style string
@@ -195,6 +195,110 @@ def plot_impedance_vs_frequency(induced_voltage_freq, figure_index = 0,
             plt.savefig(fign2)
             plt.clf()
 
+def plot_wake_vs_time(induced_voltage_time, figure_index = 0,
+                      plot_total_wake = True,
+                      plot_interpolated_wake = False, style='-',
+                      cut_left_right=None, cut_up_down=None,
+                      dirname='fig', show_plots = False):
+    """
+    Plots wakes in time domain. Given an InducedVoltageTime object, i.e. one
+    that was created from a list of wakes:
+     - can plot either the total, summed wake or each individual one from the list
+     - and can plot the wakes  on an interpolated wake axis, if the table
+        given to induced_voltage_time has been interpolated on a new axis
+
+    Additionally, either outputs the plots in the console, or save them on disk.
+
+    :param induced_voltage_time:  InducedVoltageTime object
+        InducedVoltageTime object which wakes are to be plotted
+    :param figure_index:  int
+        Index of the figure, e.g. the turn number if creating a plot for every turn
+        Default: 0
+    :param plot_total_wake: bool
+        If True, plots the total, summed wake which is calculated internally in the
+        induced_voltage_time.
+        If False, plots the raw wakes from the table of wakes given to induced_voltage_time
+        Default: True
+    :param plot_interpolated_wake: bool
+        If True, the wake table in induced_voltage_time must have been interpolated
+        to a new frequency axis. Will then plot the wakes on that new frequency axis.
+        If False, plots the raw wakes on the original frequency axis
+        given by the table in induced_voltage_time
+        Default: False
+    :param style: str
+        Matplotlib plot style string
+        Default: "-"
+    :param cut_left_right: 2-tuple of floats
+        x-axis limits of the plots (frequency in Hz)
+        Default: None
+    :param cut_up_down: 2-tuple of floats
+        y-axis limits of the plots (resistance in Ohm)
+        Default: None
+    :param dirname: str
+        Path to directory in which the plots are to be saved.
+        Default: "fig"
+    :param show_plots: bool
+        If True, will output the plots directly in the console
+        If False, will save the plots to directory given in dirname
+    """
+
+    if plot_total_wake:
+
+        ax1 = plt.subplots()[1]
+
+        ax1.plot(induced_voltage_time.time,
+                 induced_voltage_time.total_wake, style, label ="Wake")
+
+        ax1.set_xlim(cut_left_right)
+        ax1.set_ylim(cut_up_down)
+        ax1.set_xlabel("Time [s]")
+        ax1.set_ylabel("Wake [Ohm/s]")
+        ax1.legend()
+        fign = dirname + '/sum_wake_vs_table_times' "%d" % figure_index + '.png'
+
+        if show_plots:
+            plt.show()
+        else:
+            plt.savefig(fign)
+            plt.clf()
+
+    else:
+
+        fig0 = plt.figure(0)
+        ax0 = fig0.add_subplot(111)
+
+        for i in range(len(induced_voltage_time.wake_source_list)):
+            if isinstance(induced_voltage_time.wake_source_list[i],
+                          InputTable) and not plot_interpolated_wake:
+                ax0.plot(induced_voltage_time.wake_source_list[i].time_array,
+                         induced_voltage_time.wake_source_list[i].wake_array, style)
+                ax0.set_xlim(cut_left_right)
+                ax0.set_ylim(cut_up_down)
+            elif plot_interpolated_wake:
+                ax0.plot(induced_voltage_time.wake_source_list[i].new_time_array,
+                         induced_voltage_time.wake_source_list[i].wake, style)
+                ax0.set_xlim(cut_left_right)
+                ax0.set_ylim(cut_up_down)
+
+        if plot_interpolated_wake:
+            fig_suffix = "interpolated_times"
+        else:
+            fig_suffix = "table_times"
+
+        ax0.set_xlabel("Time [s]")
+        ax0.set_ylabel("Wake [Ohm/s]")
+
+
+        fign1 = dirname + '/wake_vs_' + fig_suffix + '_' "%d" % figure_index + '.png'
+
+
+
+        plt.figure(0)
+        if show_plots:
+            plt.show()
+        else:
+            plt.savefig(fign1)
+            plt.clf()
 
 def plot_induced_voltage_vs_bin_centers(total_induced_voltage, style='-', figure_index = 0,
                                         dirname='fig', show_plots = False):
-- 
GitLab