Add NWA data class and plotting

data.py

@@ -307,3 +307,43 @@ class LiberaTBTData(LiberaData):
     """Container for turn-by-turn data from libera-ireg dump
     """
     pass
+
+
+
+@dataclass
+class NWAData(Trace):
+    """Container for data from Network Analyser save
+    
+    f - frequency in Hz
+    m - magnitude
+    p - phase
+    """
+    f: np.ndarray
+    m: np.ndarray
+    p: np.ndarray
+    f_unit: str = 'Hz'
+    m_unit: str = 'a.u.'
+    p_unit: str = 'a.u.'
+    
+    @classmethod
+    def from_file(cls, magfile, phasefile=None, unwrap=False, verbose=0):
+        """Read in data from CSV file for magnitude and phase
+        
+        :param magfile: path to filename with magnitude trace data
+        :param phasefile: path to filename with phase trace data. If None, phase data is assumed to be the second column in magfile.
+        :param unwrap: if true, relative phase is unwraped to absolute phase centered around zero
+        """
+        f, m, *other = np.loadtxt(magfile, skiprows=3, delimiter=',').T
+        if phasefile is None:
+            p = other[0]
+        else:
+            fp, p, *_ = np.loadtxt(phasefile, skiprows=3, delimiter=',').T
+            if np.any(f != fp):
+                raise ValueError('Files provided do not have equal frequency components.')
+        
+        if unwrap:
+            p = np.unwrap(p)
+            p -= np.mean(p)
+        
+        return NWAData(f, m, p)
+

plotting.py

@@ -114,6 +114,14 @@ def fiberplot(ax, datasets, *, labels=[None, None], vertical=True):
         ax.set(xticks=yticks)
     else:
         ax.set(yticks=yticks)
+
+def format_axis_radians(yaxis):
+    yaxis.set_major_locator(mpl.ticker.MultipleLocator(np.deg2rad(360)))
+    yaxis.set_minor_locator(mpl.ticker.MultipleLocator(np.deg2rad(180)))
+    yaxis.set_major_formatter(mpl.ticker.FuncFormatter(lambda x, p: '0' if x==0 else '$\\pi$' if x==np.pi else '$-\\pi$' if x==-np.pi else f'${x/np.pi:g}\\pi$'))
+    yaxis.set_minor_formatter(mpl.ticker.FuncFormatter(lambda x, p: '0' if x==0 else '$\\pi$' if x==np.pi else '$-\\pi$' if x==-np.pi else None))
+
+
     
 
 
@@ -146,6 +154,16 @@ def plot_tbt(ax, libera_data, over_time=True, turn_range=None, time_range=None):
     ax2.set(ylabel='Pickup sum / a.u.', ylim=(0,1e8))
     ax2.legend([lf,ls], ['Revolution frequency', 'Pickup sum signal'], loc='center right', fontsize='small')
 
+def plot_btf(axf, axp, data, *, frev=None, **kwargs):
+    
+    f = data.f*(1/frev if frev else 1)
+    axf.plot(f, data.m, **kwargs)
+    axp.plot(f, data.p, **kwargs)
+
+    axf.set(ylabel=f'Magnitude / {data.m_unit}', xlim=(np.min(f), np.max(f)))
+    axp.set(ylabel=f'Phase / {data.p_unit}', xlabel='Stimulus tune' if frev else f'Stimulus frequency / {data.f_unit}')
+    
+    return fig, (af, ap)
 
 
 def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tune_range=None, fit=False, smoothing=None, return_spectrum=False, **kwargs):