Unit autoscale matplotlib and pyqtgraph graphs

qcodes/plots/base.py

@@ -24,6 +24,9 @@ def __init__(self, interval=1, data_keys='xyz'):
         self.traces = []
         self.data_updaters = set()
         self.interval = interval
+        self.standardunits = ['V', 's', 'J', 'W', 'm', 'eV', 'A', 'K', 'g',
+                              'Hz', 'rad', 'T', 'H', 'F', 'Pa', 'C', 'Ω', 'Ohm',
+                              'S']
 
     def clear(self):
         """

qcodes/plots/pyqtgraph.py

@@ -1,6 +1,7 @@
 """
 Live plotting using pyqtgraph
 """
+from typing import Optional, Dict, Union
 import numpy as np
 import pyqtgraph as pg
 import pyqtgraph.multiprocess as pgmp
@@ -80,6 +81,7 @@ def __init__(self, *args, figsize=(1000, 600), interval=0.25,
                 raise err
         self.win.setBackground(theme[1])
         self.win.resize(*figsize)
+        self._orig_fig_size = figsize
         self.subplots = [self.add_subplot()]
 
         if args or kwargs:
@@ -144,6 +146,7 @@ def add_to_plot(self, subplot=1, **kwargs):
 
         if prev_default_title == self.win.windowTitle():
             self.win.setWindowTitle(self.get_default_title())
+        self.fixUnitScaling()
 
     def _draw_plot(self, subplot_object, y, x=None, color=None, width=None,
                    antialias=None, **kwargs):
@@ -475,3 +478,94 @@ def save(self, filename=None):
     def setGeometry(self, x, y, w, h):
         """ Set geometry of the plotting window """
         self.win.setGeometry(x, y, w, h)
+
+    def autorange(self):
+        """
+        Auto range all limits in case they were changed during interactive
+        plot. Reset colormap if changed and resize window to original size.
+        """
+        for subplot in self.subplots:
+            vBox = subplot.getViewBox()
+            vBox.enableAutoRange(vBox.XYAxes)
+        cmap = None
+        # resize histogram
+        for trace in self.traces:
+            if 'plot_object' in trace.keys():
+                if (isinstance(trace['plot_object'], dict) and
+                            'hist' in trace['plot_object'].keys()):
+                    cmap = trace['plot_object']['cmap']
+                    maxval = trace['config']['z'].max()
+                    minval = trace['config']['z'].min()
+                    trace['plot_object']['hist'].setLevels(minval, maxval)
+                    trace['plot_object']['hist'].vb.autoRange()
+        if cmap:
+            self.set_cmap(cmap)
+        # set window back to original size
+        self.win.resize(*self._orig_fig_size)
+
+    def fixUnitScaling(self, startranges: Optional[Dict[str, Dict[str, Union[float,int]]]]=None):
+        """
+        Disable SI rescaling if units are not standard units and limit
+        ranges to data if known.
+
+        Args:
+
+            startranges: The plot can automatically infer the full ranges
+                         array parameters. However it has no knowledge of the
+                         ranges or regular parameters. You can explicitly pass
+                         in the values here as a dict of the form
+                         {'paramtername': {max: value, min:value}}
+        """
+        axismapping = {'x': 'bottom',
+                       'y': 'left',
+                       'z': 'right'}
+        standardunits = self.standardunits
+        for i, plot in enumerate(self.subplots):
+            # make a dict mapping axis labels to axis positions
+            for axis in ('x', 'y', 'z'):
+                if self.traces[i]['config'].get(axis):
+                    unit = self.traces[i]['config'][axis].unit
+                    if unit not in standardunits:
+                        if axis in ('x', 'y'):
+                            ax = plot.getAxis(axismapping[axis])
+                        else:
+                            # 2D measurement
+                            # Then we should fetch the colorbar
+                            ax = self.traces[i]['plot_object']['hist'].axis
+                        ax.enableAutoSIPrefix(False)
+                        # because updateAutoSIPrefix called from
+                        # enableAutoSIPrefix doesnt actually take the
+                        # value of the argument into account we have
+                        # to manually replicate the update here
+                        ax.autoSIPrefixScale = 1.0
+                        ax.setLabel(unitPrefix='')
+                        ax.picture = None
+                        ax.update()
+
+                    # set limits either from dataset or
+                    setarr = self.traces[i]['config'][axis].ndarray
+                    arrmin = None
+                    arrmax = None
+                    if not np.all(np.isnan(setarr)):
+                        arrmax = setarr.max()
+                        arrmin = setarr.min()
+                    elif startranges is not None:
+                        try:
+                            paramname = self.traces[i]['config'][axis].full_name
+                            arrmax = startranges[paramname]['max']
+                            arrmin = startranges[paramname]['min']
+                        except (IndexError, KeyError):
+                            continue
+
+                    if axis == 'x':
+                        rangesetter = getattr(plot.getViewBox(), 'setXRange')
+                    elif axis == 'y':
+                        rangesetter = getattr(plot.getViewBox(), 'setYRange')
+                    else:
+                        rangesetter = None
+
+                    if (rangesetter is not None
+                        and arrmin is not None
+                        and arrmax is not None):
+                        rangesetter(arrmin, arrmax)
+

qcodes/plots/qcmatplotlib.py

@@ -7,6 +7,7 @@
 from functools import partial
 
 import matplotlib.pyplot as plt
+from matplotlib import ticker
 import numpy as np
 from matplotlib.transforms import Bbox
 from numpy.ma import masked_invalid, getmask
@@ -399,4 +400,58 @@ def tight_layout(self):
         Perform a tight layout on the figure. A bit of additional spacing at
         the top is also added for the title.
         """
-        self.fig.tight_layout(rect=[0, 0, 1, 0.95])
+        self.fig.tight_layout(rect=[0, 0, 1, 0.95])
+
+    def rescale_axis(self):
+        """
+        Rescale axis and units for axis that are in standard units
+        i.e. V, s J ... to m μ, m
+        This scales units defined in BasePlot.standardunits only
+        to avoid prefixes on combined or non standard units
+        """
+        def scale_formatter(i, pos, scale):
+            return "{0:g}".format(i * scale)
+
+        for i, subplot in enumerate(self.subplots):
+            for axis in 'x', 'y', 'z':
+                if self.traces[i]['config'].get(axis):
+                    unit = self.traces[i]['config'][axis].unit
+                    label = self.traces[i]['config'][axis].label
+                    maxval = abs(self.traces[i]['config'][axis].ndarray).max()
+                    units_to_scale = self.standardunits
+
+                    # allow values up to a <1000. i.e. nV is used up to 1000 nV
+                    prefixes = ['n', 'μ', 'm', '', 'k', 'M', 'G']
+                    thresholds = [10**(-6 + 3*n) for n in range(len(prefixes))]
+                    scales = [10**(9 - 3*n) for n in range(len(prefixes))]
+
+                    if unit in units_to_scale:
+                        scale = 1
+                        new_unit = unit
+                        for prefix, threshold, trialscale in zip(prefixes,
+                                                                 thresholds,
+                                                                 scales):
+                            if maxval < threshold:
+                                scale = trialscale
+                                new_unit = prefix + unit
+                                break
+                        # special case the largest
+                        if maxval > thresholds[-1]:
+                            scale = scales[-1]
+                            new_unit = prefixes[-1] + unit
+
+                        tx = ticker.FuncFormatter(
+                            partial(scale_formatter, scale=scale))
+                        new_label = "{} ({})".format(label, new_unit)
+                        if axis in ('x', 'y'):
+                            getattr(subplot,
+                                    "{}axis".format(axis)).set_major_formatter(
+                                tx)
+                            getattr(subplot, "set_{}label".format(axis))(
+                                new_label)
+                        else:
+                            subplot.qcodes_colorbar.formatter = tx
+                            subplot.qcodes_colorbar.ax.yaxis.set_major_formatter(
+                                tx)
+                            subplot.qcodes_colorbar.set_label(new_label)
+                            subplot.qcodes_colorbar.update_ticks()