Rohde Schwarz ZNB /VNA refactor common code

qcodes/instrument_drivers/rohde_schwarz/ZNB.py

@@ -3,7 +3,6 @@
 from qcodes import VisaInstrument
 from qcodes import ChannelList, InstrumentChannel
 from qcodes.utils import validators as vals
-from cmath import phase
 import numpy as np
 from qcodes import MultiParameter, ArrayParameter
 
@@ -12,27 +11,7 @@
 
 class FrequencySweepMagPhase(MultiParameter):
     """
-    Hardware controlled parameter class for Rohde Schwarz ZNB trace.
-
-    Instrument returns an list of transmission data in the form of a list of
-    complex numbers taken from a frequency sweep.
-
-    This is a multiparameter containing both amplitude and phase
-
-    Args:
-        name: parameter name
-        instrument: instrument the parameter belongs to
-        start: starting frequency of sweep
-        stop: ending frequency of sweep
-        npts: number of points in frequency sweep
-
-    Methods:
-          set_sweep(start, stop, npts): sets the shapes and
-              setpoint arrays of the parameter to correspond with the sweep
-          get(): executes a sweep and returns magnitude and phase arrays
-
-    TODO:
-      - ability to choose for linear or db in magnitude return
+    Sweep that return magnitude and phase.
     """
 
     def __init__(self, name, instrument, start, stop, npts, channel):
@@ -55,38 +34,12 @@ def set_sweep(self, start, stop, npts):
         self.setpoints = ((f,), (f,))
         self.shapes = ((npts,), (npts,))
 
-    def get(self):
-        if not self._instrument._parent.rf_power():
-            log.warning("RF output is off when getting mag phase")
-        # it is possible that the instrument and qcodes disagree about
-        # which parameter is measured on this channel
-        instrument_parameter = self._instrument.vna_parameter()
-        if instrument_parameter != self._instrument._vna_parameter:
-            raise RuntimeError("Invalid parameter. Tried to measure "
-                               "{} got {}".format(self._instrument._vna_parameter,
-                                                  instrument_parameter))
-        self._instrument.write('SENS{}:AVER:STAT ON'.format(self._channel))
-        self._instrument.write('SENS{}:AVER:CLE'.format(self._channel))
-        self._instrument._parent.cont_meas_off()
-
-        # instrument averages over its last 'avg' number of sweeps
-        # need to ensure averaged result is returned
-        for avgcount in range(self._instrument.avg()):
-            self._instrument.write('INIT{}:IMM; *WAI'.format(self._channel))
-        data_str = self._instrument.ask(
-            'CALC{}:DATA? SDAT'.format(self._channel)).split(',')
-        data_list = [float(v) for v in data_str]
-
-        # data_list of complex numbers [re1,im1,re2,im2...]
-        data_arr = np.array(data_list).reshape(int(len(data_list) / 2), 2)
-        mag_array, phase_array = [], []
-        for comp in data_arr:
-            complex_num = complex(comp[0], comp[1])
-            mag_array.append(abs(complex_num))
-            phase_array.append(phase(complex_num))
-        self._instrument._parent.cont_meas_on()
-        return mag_array, phase_array
-
+    def get_raw(self):
+        old_format = self._instrument.format()
+        self._instrument.format('Complex')
+        data = self._instrument._get_sweep_data(force_polar = True)
+        self._instrument.format(old_format)
+        return abs(data), np.angle(data)
 
 class FrequencySweep(ArrayParameter):
     """
@@ -108,15 +61,14 @@ class FrequencySweep(ArrayParameter):
           get(): executes a sweep and returns magnitude and phase arrays
 
     """
-
     def __init__(self, name, instrument, start, stop, npts, channel):
         super().__init__(name, shape=(npts,),
                          instrument=instrument,
                          unit='dB',
                          label='{} magnitude'.format(
                              instrument._vna_parameter),
                          setpoint_units=('Hz',),
-                         setpoint_names=('{}_frequency'.format(instrument._vna_parameter),))
+                         setpoint_names=('{} frequency'.format(instrument._vna_parameter),))
         self.set_sweep(start, stop, npts)
         self._channel = channel
 
@@ -127,35 +79,13 @@ def set_sweep(self, start, stop, npts):
         self.setpoints = (f,)
         self.shape = (npts,)
 
-    def get(self):
-        if not self._instrument._parent.rf_power():
-            log.warning("RF output is off when getting mag")
-        # it is possible that the instrument and qcodes disagree about
-        # which parameter is measured on this channel
-        instrument_parameter = self._instrument.vna_parameter()
-        if instrument_parameter != self._instrument._vna_parameter:
-            raise RuntimeError("Invalid parameter. Tried to measure "
-                               "{} got {}".format(self._instrument._vna_parameter,
-                                                  instrument_parameter))
-        self._instrument.write('SENS{}:AVER:STAT ON'.format(self._channel))
-        self._instrument.write('SENS{}:AVER:CLE'.format(self._channel))
-        self._instrument._parent.cont_meas_off()
-
-        # instrument averages over its last 'avg' number of sweeps
-        # need to ensure averaged result is returned
-        for avgcount in range(self._instrument.avg()):
-            self._instrument.write('INIT{}:IMM; *WAI'.format(self._channel))
-        data_str = self._instrument.ask(
-            'CALC{}:DATA? FDAT'.format(self._channel))
-        data = np.array(data_str.rstrip().split(',')).astype('float64')
-        if self._instrument.format() in ['Polar', 'Complex',
-                                         'Smith', 'Inverse Smith']:
-            log.warning("QCoDeS Dataset does not currently support Complex "
-                        "values. Will discard the imaginary part.")
-            data = data[0::2] + 1j * data[1::2]
-        self._instrument._parent.cont_meas_on()
-        return data
-
+    def get_raw(self):
+       data = self._instrument._get_sweep_data()
+       if self._instrument.format() in ['Polar', 'Complex',
+                            'Smith', 'Inverse Smith']:
+           log.warning("QCoDeS Dataset does not currently support Complex "
+                       "values. Will discard the imaginary part.")
+       return data
 
 class ZNBChannel(InstrumentChannel):
 
@@ -184,7 +114,7 @@ def __init__(self, parent, name, channel):
                            get_cmd='SOUR{}:POW?'.format(n),
                            set_cmd='SOUR{}:POW {{:.4f}}'.format(n),
                            get_parser=lambda x: int(round(float(x))),
-                           vals=vals.Numbers(-150, 25))
+                           vals=vals.Numbers(-60, 25))
         self.add_parameter(name='bandwidth',
                            label='Bandwidth',
                            unit='Hz',
@@ -297,10 +227,8 @@ def _strip(self, var):
 
     def _set_start(self, val):
         channel = self._instrument_channel
-        self.write('SENS{}:FREQ:START {:.4f}'.format(channel, val))
+        self.write('SENS{}:FREQ:START {:.7f}'.format(channel, val))
         stop = self.stop()
-        npts = self.npts()
-
         if val >= stop:
             raise ValueError(
                 "Stop frequency must be larger than start frequency.")
@@ -309,54 +237,90 @@ def _set_start(self, val):
         if val != start:
             log.warning(
                 "Could not set start to {} setting it to {}".format(val, start))
-        # update setpoints for FrequencySweep param
-        self.trace.set_sweep(start, stop, npts)
-        self.trace_mag_phase.set_sweep(start, stop, npts)
+        self._update_traces()
 
     def _set_stop(self, val):
         channel = self._instrument_channel
         start = self.start()
-        npts = self.npts()
         if val <= start:
             raise ValueError(
                 "Stop frequency must be larger than start frequency.")
-        self.write('SENS{}:FREQ:STOP {:.4f}'.format(channel, val))
+        self.write('SENS{}:FREQ:STOP {:.7f}'.format(channel, val))
         # we get stop as the vna may not be able to set it to the exact value provided
         stop = self.stop()
         if val != stop:
             log.warning(
                 "Could not set stop to {} setting it to {}".format(val, stop))
-        # update setpoints for FrequencySweep param
-        self.trace.set_sweep(start, stop, npts)
-        self.trace_mag_phase.set_sweep(start, stop, npts)
+        self._update_traces()
 
     def _set_npts(self, val):
         channel = self._instrument_channel
-        self.write('SENS{}:SWE:POIN {:.4f}'.format(channel, val))
-        start = self.start()
-        stop = self.stop()
-        # update setpoints for FrequencySweep param
-        self.trace.set_sweep(start, stop, val)
-        self.trace_mag_phase.set_sweep(start, stop, val)
+        self.write('SENS{}:SWE:POIN {:.7f}'.format(channel, val))
+        self._update_traces()
 
     def _set_span(self, val):
         channel = self._instrument_channel
-        self.write('SENS{}:FREQ:SPAN {:.4f}'.format(channel, val))
-        start = self.start()
-        stop = self.stop()
-        npts = self.npts()
-        self.trace.set_sweep(start, stop, npts)
-        self.trace_mag_phase.set_sweep(start, stop, npts)
+        self.write('SENS{}:FREQ:SPAN {:.7f}'.format(channel, val))
+        self._update_traces()
 
     def _set_center(self, val):
         channel = self._instrument_channel
-        self.write('SENS{}:FREQ:CENT {:.4f}'.format(channel, val))
+        self.write('SENS{}:FREQ:CENT {:.7f}'.format(channel, val))
+        self._update_traces()
+
+    def _update_traces(self):
+        """ updates start, stop and npts of all trace parameters"""
         start = self.start()
         stop = self.stop()
         npts = self.npts()
-        self.trace.set_sweep(start, stop, npts)
-        self.trace_mag_phase.set_sweep(start, stop, npts)
+        for _, parameter in self.parameters.items():
+            if isinstance(parameter, (ArrayParameter, MultiParameter)):
+                try:
+                    parameter.set_sweep(start, stop, npts)
+                except AttributeError:
+                    pass
 
+    def _get_sweep_data(self, force_polar=False):
+
+        if not self._parent.rf_power():
+            log.warning("RF output is off when getting sweep data")
+        # it is possible that the instrument and qcodes disagree about
+        # which parameter is measured on this channel
+        instrument_parameter = self.vna_parameter()
+        if instrument_parameter != self._vna_parameter:
+            raise RuntimeError("Invalid parameter. Tried to measure "
+                               "{} got {}".format(self._vna_parameter,
+                                                  instrument_parameter))
+        self.write('SENS{}:AVER:STAT ON'.format(self._instrument_channel))
+        self.write('SENS{}:AVER:CLE'.format(self._instrument_channel))
+
+        # preserve original state of the znb
+        initial_state = self.status()
+        self.status(1)
+        self._parent.cont_meas_off()
+        try:
+            # if force polar is set, the SDAT data format will be used. Here
+            # the data will be transfered as a complex number independet of
+            # the set format in the instrument.
+            if force_polar:
+                data_format_command = 'SDAT'
+            else:
+                data_format_command = 'FDAT'
+            # instrument averages over its last 'avg' number of sweeps
+            # need to ensure averaged result is returned
+            for avgcount in range(self.avg()):
+                self.write('INIT{}:IMM; *WAI'.format(self._instrument_channel))
+            data_str = self.ask(
+                'CALC{}:DATA? {}'.format(self._instrument_channel,
+                                         data_format_command))
+            data = np.array(data_str.rstrip().split(',')).astype('float64')
+            if self.format() in ['Polar', 'Complex',
+                                 'Smith', 'Inverse Smith']:
+                data = data[0::2] + 1j * data[1::2]
+        finally:
+            self._parent.cont_meas_on()
+            self.status(initial_state)
+        return data
 
 class ZNB(VisaInstrument):
     """
@@ -376,6 +340,8 @@ class ZNB(VisaInstrument):
     TODO:
     - check initialisation settings and test functions
     """
+    CHANNEL_CLASS = ZNBChannel
+
 
     def __init__(self, name: str, address: str, init_s_params: bool=True, **kwargs):
 
@@ -418,10 +384,13 @@ def __init__(self, name: str, address: str, init_s_params: bool=True, **kwargs):
             num_ports, num_ports))
         self.add_function('display_single_window',
                           call_cmd='DISP:LAY GRID;:DISP:LAY:GRID 1,1')
+        self.add_function('display_dual_window',
+                          call_cmd='DISP:LAY GRID;:DISP:LAY:GRID 2,1')
         self.add_function('rf_off', call_cmd='OUTP1 OFF')
         self.add_function('rf_on', call_cmd='OUTP1 ON')
+        self.reset()
         self.clear_channels()
-        channels = ChannelList(self, "VNAChannels", ZNBChannel,
+        channels = ChannelList(self, "VNAChannels", self.CHANNEL_CLASS,
                                snapshotable=True)
         self.add_submodule("channels", channels)
         if init_s_params:
@@ -446,12 +415,14 @@ def display_grid(self, rows: int, cols: int):
         """
         self.write('DISP:LAY GRID;:DISP:LAY:GRID {},{}'.format(rows, cols))
 
-    def add_channel(self, vna_parameter: str):
+    def add_channel(self, vna_parameter: str, **kwargs):
         n_channels = len(self.channels)
-        channel = ZNBChannel(self, vna_parameter, n_channels + 1)
+        channel = self.CHANNEL_CLASS(self, vna_parameter, n_channels + 1, **kwargs)
         self.channels.append(channel)
         if n_channels == 0:
             self.display_single_window()
+        if n_channels == 1:
+            self.display_dual_window()
 
     def _set_default_values(self):
         for channel in self.channels: