More options for spill plots

data/common.py

@@ -45,12 +45,16 @@ CMD_BEAM_INJECTION = 283
 #################
 
 
-def avg(*data, n=100):
-    """Averages the data by combining n subsequent points into one (works on last axis)"""
+def avg(*data, n=100, function=np.mean):
+    """Averages the data by combining n subsequent points into one (works on last axis)
+    :param data: the data to average over
+    :param n: number of subsequent datapoints of intput to average into one point in the output
+    :param function: averaging function to apply to last axis of input data, e.g. np.mean, np.std, np.max etc.
+    """
     result = []
     for d in data:
         w = int(d.shape[-1]/n)
-        result.append(d[...,:w*n].reshape(*d.shape[:-1], w, n).mean(axis=-1))
+        result.append(function(d[...,:w*n].reshape(*d.shape[:-1], w, n), axis=-1))
     return result
 
 def smooth(*data, n):

data/tdc.py

@@ -83,17 +83,27 @@ class TDCSpill:
         """Delay between arrival of consecutive particles in s"""
         return np.diff(self.hittimes)
         
-    def hitcounts(self, dt):
+    def hitcounts(self, dt, dt2=None):
         """Particle counts per time bin
         
         Re-samples the timestamps into an equally spaced time series with counts per unit time
         :param dt: width in s of time bin for counting.
+        :param dt2: width in s of makro time bin. If not None, this function returns a 2D array with time resolution dt2 along the first and dt along the second axis
         :returns: (count_array, bin_edges) in units of (1, second)
         """
-        # re-sample timestamps into equally sampled time series
+        # re-sample timestamps into equally sampled time series of shape [length/dt]
         bins = int(self.length/dt)
         range = (self.time_offset, self.time_offset + bins*dt)
-        return np.histogram(self.hittimes, bins=bins, range=range)
+        counts, edges = np.histogram(self.hittimes, bins=bins, range=range)
+        
+        if dt2 is not None:
+            # make 2D array of shape [length/dt2, dt2/dt]
+            ebins = int(dt2/dt)
+            counts = counts[:int(len(counts)/ebins)*ebins].reshape((-1, ebins))
+            edges = edges[:int(len(edges)/ebins+1)*ebins:ebins]
+    
+        return counts, edges
+    
     
     def __repr__(self):
         return f'TDCSpill({self.detector}, length={self.length:g}s, counts={self.counts})'

plotting.py

@@ -275,7 +275,7 @@ def turn_or_time_range(time, turn_range=None, time_range=None):
     return slice(*turn_range)
 
 
-def plot_tbt(ax, libera_data, what='fsxy', *, over_time=True, turn_range=None, time_range=None, averaging=500, **kwargs):
+def plot_tbt(ax, libera_data, what='fsxy', *, over_time=True, turn_range=None, time_range=None, averaging=500, show_avg_std=True, show_avg_extrema=True, **kwargs):
     """Plot turn-by-turn data    
     
     :param libera_data: instance of LiberaTBTData
@@ -284,6 +284,8 @@ def plot_tbt(ax, libera_data, what='fsxy', *, over_time=True, turn_range=None, t
     :param turn_range: (start, stop) tuple of turns to plot
     :param time_range: (start, stop) tuple of time in s to plot
     :param averaging: number of consecutive turns to average over
+    :param show_avg_std: if True, plot the band of standard deviation around the averaged data
+    :param show_avg_std: if True, plot the band of min-to-max around the averaged data
     """
     assert isinstance(libera_data, LiberaTBTData), f'Expected LiberaTBTData but got {type(libera_data)}'
     
@@ -292,7 +294,8 @@ def plot_tbt(ax, libera_data, what='fsxy', *, over_time=True, turn_range=None, t
     turn_range = turn_or_time_range(libera_data.t, turn_range, time_range)
     t = libera_data.t[turn_range]
     x = t if over_time else np.arange(0, len(t))
-    axes = {}
+    axes = {} # label: ax
+    limits = {} # ax: Bbox
     ls, labels = [], []
     for i, w in enumerate(what):
         if axlabels[w] in axes:
@@ -302,6 +305,7 @@ def plot_tbt(ax, libera_data, what='fsxy', *, over_time=True, turn_range=None, t
             if i > 0: a.spines.right.set_position(("axes", 0.9+0.1*len(axes)))
             a.set(ylabel=axlabels[w], xlabel='Time / s' if over_time else 'Turn')
             axes.update({axlabels[w]: a})
+            limits[a] = mpl.transforms.Bbox([[0,0],[0,0]])
         
         if w == 'f':
             v = 1e-3/np.diff(t) # to kHz
@@ -311,12 +315,30 @@ def plot_tbt(ax, libera_data, what='fsxy', *, over_time=True, turn_range=None, t
         args = dict(**kwargs)
         if 'c' not in args and 'color' not in args: 
             args['c'] = dict(f=cmap_petroff_10(1), s=cmap_petroff_10(3), x=cmap_petroff_10(0), y=cmap_petroff_10(2))[w]
-        l, = a.plot(*avg(x[:len(v)], v, n=averaging), **args)
+        xx, vv = avg(x[:len(v)], v, n=averaging)
+        l, = a.plot(xx, vv, **args)
         ls.append(l)
         labels.append(dict(f='Revolution frequency', s='Pickup sum signal', x='X position', y='Y position')[w])
+        limits[a].update_from_data_xy(np.vstack(l.get_data()).T, ignore=limits[a].width==limits[a].height==0)
+        if averaging > 1 and show_avg_std:
+            _, ve = avg(x[:len(v)], v, n=averaging, function=np.std)
+            a.fill_between(xx, vv-ve, vv+ve, color=l.get_color(), alpha=0.4, zorder=-1, lw=0)
+        if averaging > 1 and show_avg_extrema:
+            _, vmi = avg(x[:len(v)], v, n=averaging, function=np.min)
+            _, vma = avg(x[:len(v)], v, n=averaging, function=np.max)
+            a.fill_between(xx, vmi, vma, color=l.get_color(), alpha=0.2, zorder=-2, lw=0)
+        
+    # autosacale
+    for a, lim in limits.items():
+        a.dataLim = lim
+        a.autoscale_view()
     
     if len(what) > 1: a.legend(ls, labels, fontsize='small')
 
+
+
+
+
 def plot_btf(axf, axp, data, *, frev=None, filled=False, smoothing=None, **kwargs):
     """Plot beam transfer function
     
@@ -741,23 +763,36 @@ def plot_beam_size(ax, ipm_data, xy, time_range=None, smoothing=None, **kwargs):
 # TDC data
 ###########
 
-def plot_spill_intensity(ax, spill_data, bin_time=10e-6, *, cumulative=False, rate=False, **plot_kwargs):
+def plot_spill_intensity(ax, spill_data, bin_time=10e-6, *, rate=False, cumulative=False, relative=False, **plot_kwargs):
     """Plot spill intensity (binned particle counts) over extraction time
     
     :param ax: Axis to plot onto
     :param spill_data: Instance of TDCSpill
     :param bin_time: width of bins for particle counting in sec
+    :param rate: if true, plot particle count rate, i.e. counts per s instead of per bin
+    :param cumulative: if true, plot cumulative particle number
+    :param relative: if true, plot fraction of total particle number
     """
     nbins = int(np.ceil(spill_data.length/bin_time))
-    weights = np.ones_like(spill_data.hittimes)/bin_time if rate else None
+    
+    assert not (rate and cumulative), 'rate and cumulative plotting are mutually exclusive'
+    weights = np.ones_like(spill_data.hittimes)
+    if rate: weights /= bin_time
+    if relative: weights /= spill_data.counts
+    label = 'Extracted particle' + \
+                (' fraction' if relative else 's') + \
+                ('' if cumulative else \
+                     (' / s' if rate else f'\nper ${(bin_time*SI.s).to_compact():~gL}$ interval'))
+                
     ax.hist(spill_data.hittimes, range=(spill_data.time_offset, spill_data.time_offset+bin_time*nbins), bins=nbins,
             histtype='step', cumulative=cumulative, weights=weights, **plot_kwargs)
-    ax.set(xlabel='Extraction time / s', #xlim=(spill_data.time_offset, spill_data.time_offset+bin_time*nbins),
-           ylabel='Extracted particles' + ('' if cumulative else ' / s' if rate else f'\nper ${(bin_time*SI.s).to_compact():~gL}$ interval'))
+    ax.set(ylabel=label, xlabel='Extraction time / s', #xlim=(spill_data.time_offset, spill_data.time_offset+bin_time*nbins),
+           )
+    if relative: ax.yaxis.set_major_formatter(mpl.ticker.PercentFormatter(1, decimals=0))
 
 
 def plot_spill_duty_factor(ax, spill_data, counting_dt=10e-6, evaluation_dt=10e-3, *, show_poisson_limit=False, **kwargs):
-    """Plot spill intensity (binned particle counts) over extraction time
+    """Plot spill duty factor over extraction time
     
     :param ax: Axis to plot onto
     :param spill_data: Instance of TDCSpill
@@ -767,18 +802,16 @@ def plot_spill_duty_factor(ax, spill_data, counting_dt=10e-6, evaluation_dt=10e-
     :param show_poisson_limit: if true, show poison limit of duty factor
     
     """
-    counts, edges = spill_data.hitcounts(counting_dt)
-    ebins = int(round(evaluation_dt/counting_dt))
-    counts = counts[:int(len(counts)/ebins)*ebins].reshape((-1, ebins))
-    edges = edges[:int(len(edges)/ebins+1)*ebins:ebins]
+    
+    counts, edges = spill_data.hitcounts(counting_dt, evaluation_dt)
     
     # spill duty factor
     F = np.mean(counts, axis=1)**2 / np.mean(counts**2, axis=1)
-    ax.stairs(F, edges, **kwargs)
+    s = ax.stairs(F, edges, **kwargs)
     
     if show_poisson_limit:
         Fp = np.mean(counts, axis=1) / ( np.mean(counts, axis=1) + 1 )
-        ax.stairs(Fp, edges, color='gray', label='Poisson limit', zorder=-1, lw=1)
+        ax.stairs(Fp, edges, label='Poisson limit', color=s.get_edgecolor() or 'gray', alpha=0.5,zorder=-1, lw=1, ls=':')
         
     ax.set(xlabel='Extraction time / s', ylim=(0,1.1),
            ylabel=f'Spill duty factor F\n(${(evaluation_dt*SI.s).to_compact():~gL}$ / ${(counting_dt*SI.s).to_compact():~gL}$ intervals)'