Matplotlib updates

qcodes/plots/qcmatplotlib.py

@@ -13,6 +13,8 @@
 
 
 class MatPlot(BasePlot):
+    plot_1D_kwargs = {}
+    plot_2D_kwargs = {}
     """
     Plot x/y lines or x/y/z heatmap data. The first trace may be included
     in the constructor, other traces can be added with MatPlot.add()
@@ -57,8 +59,13 @@ def _init_plot(self, subplots=None, figsize=None, num=None):
         else:
             self.fig, self.subplots = plt.subplots(*subplots, num=num,
                                                    figsize=figsize)
-        if not hasattr(self.subplots, '__len__'):
-            self.subplots = (self.subplots,)
+
+        # Test if subplots is actually a single axis
+        if not isinstance(self.subplots, np.ndarray):
+            self.subplots = np.array([self.subplots])
+
+        # Flatten subplots in case it is a 2D array
+        self.subplots = np.ndarray.flatten(self.subplots)
 
         self.title = self.fig.suptitle('')
 
@@ -87,7 +94,7 @@ def add_to_plot(self, **kwargs):
                 `x`, `y`, and `fmt` (if present) are passed as positional args
         """
         # TODO some way to specify overlaid axes?
-        ax = self._get_axes(kwargs)
+        ax = self._get_axes(**kwargs)
         if 'z' in kwargs:
             plot_object = self._draw_pcolormesh(ax, **kwargs)
         else:
@@ -105,8 +112,8 @@ def add_to_plot(self, **kwargs):
             # in case the user has updated title, don't change it anymore
             self.title.set_text(self.get_default_title())
 
-    def _get_axes(self, config):
-        return self.subplots[config.get('subplot', 1) - 1]
+    def _get_axes(self, subplot=1, **kwargs):
+        return self.subplots[subplot - 1]
 
     def _update_labels(self, ax, config):
         if 'x' in config and not ax.get_xlabel():
@@ -132,7 +139,7 @@ def update_plot(self):
                 if plot_object:
                     plot_object.remove()
 
-                ax = self._get_axes(config)
+                ax = self._get_axes(**config)
                 plot_object = self._draw_pcolormesh(ax, **config)
                 trace['plot_object'] = plot_object
 
@@ -168,23 +175,98 @@ def _draw_plot(self, ax, y, x=None, fmt=None, subplot=1, **kwargs):
         # didn't want to strip it out of kwargs earlier because it should stay
         # part of trace['config'].
         args = [arg for arg in [x, y, fmt] if arg is not None]
-        line, = ax.plot(*args, **kwargs)
+
+        full_kwargs = {**self.plot_1D_kwargs, **kwargs}
+        line, = ax.plot(*args, **full_kwargs)
         return line
 
-    def _draw_pcolormesh(self, ax, z, x=None, y=None, subplot=1, **kwargs):
+    def _draw_pcolormesh(self, ax, z, x=None, y=None, subplot=1,
+                         nticks=None, use_offset=False, **kwargs):
+        """
+        Draws a 2D color plot
+
+        Args:
+            ax (Axis): Matplotlib axis object to plot in
+            z: 2D array of data values
+            x (Array, Optional): Array of values along x-axis. Dimensions should
+                be either same as z, or equal to length along x-axis.
+            y (Array, Optional): Array of values along y-axis. Dimensions should
+                be either same as z, or equal to length along y-axis.
+            subplot (int, Optional): Deprecated, see alexj notes below
+            nticks (int, Optional): preferred number of ticks along axes
+            use_offset (bool, Optional): Whether or not axes can have an offset
+            **kwargs: Optional list of kwargs to be passed on to pcolormesh.
+                These will overwrite any of the default kwargs in plot_kwargs.
+        """
+
         # NOTE(alexj)stripping out subplot because which subplot we're in is already
         # described by ax, and it's not a kwarg to matplotlib's ax.plot. But I
         # didn't want to strip it out of kwargs earlier because it should stay
         # part of trace['config'].
-        args = [masked_invalid(arg) for arg in [x, y, z]
-                if arg is not None]
+        args_masked = [masked_invalid(arg) for arg in [x, y, z]
+                      if arg is not None]
+
+        if np.any([np.all(getmask(arg)) for arg in args_masked]):
+            # if the z array is masked, don't draw at all
+            # there's nothing to draw, and anyway it throws a warning
+            # pcolormesh does not accept masked x and y axes, so we do not need
+            # to check for them.
+            return False
+
+        if x is not None and y is not None:
+            # If x and y are provided, modify the arrays such that they
+            # correspond to grid corners instead of grid centers.
+            # This is to ensure that pcolormesh centers correctly and
+            # does not ignore edge points.
+            args = []
+            for k, arr in enumerate(args_masked[:-1]):
+                # If a two-dimensional array is provided, only consider the
+                # first row/column, depending on the axis
+                if arr.ndim > 1:
+                    arr = arr[0] if k == 0 else arr[:,0]
+
+                if np.ma.is_masked(arr[1]):
+                    # Only the first element is not nan, in this case pad with
+                    # a value, and separate their values by 1
+                    arr_pad = np.pad(arr, (1, 0), mode='symmetric')
+                    arr_pad[:2] += [-0.5, 0.5]
+                else:
+                    # Add padding on both sides equal to endpoints
+                    arr_pad = np.pad(arr, (1, 1), mode='symmetric')
+                    # Add differences to edgepoints (may be nan)
+                    arr_pad[0] += arr_pad[1] - arr_pad[2]
+                    arr_pad[-1] += arr_pad[-2] - arr_pad[-3]
+
+                    diff = np.ma.diff(arr_pad) / 2
+                    # Insert value at beginning and end of diff to ensure same
+                    # length
+                    diff = np.insert(diff, 0, diff[0])
+
+                    arr_pad += diff
+                    # Ignore final value
+                    arr_pad = arr_pad[:-1]
+                args.append(masked_invalid(arr_pad))
+            args.append(args_masked[-1])
+        else:
+            # Only the masked value of z is used as a mask
+            args = args_masked[-1:]
+
+        # Include default plotting kwargs, which can be overwritten by given
+        # kwargs
+        full_kwargs = {**self.plot_2D_kwargs, **kwargs}
+        pc = ax.pcolormesh(*args, **full_kwargs)
+
+        # Set x and y limits if arrays are provided
+        if x is not None and y is not None:
+            ax.set_xlim(np.nanmin(args[0]), np.nanmax(args[0]))
+            ax.set_ylim(np.nanmin(args[1]), np.nanmax(args[1]))
 
-        for arg in args:
-            if np.all(getmask(arg)):
-                # if any entire array is masked, don't draw at all
-                # there's nothing to draw, and anyway it throws a warning
-                return False
-        pc = ax.pcolormesh(*args, **kwargs)
+        # Specify preferred number of ticks with labels
+        if nticks:
+            ax.locator_params(nbins=nticks)
+
+        # Specify if axes can have offset or not
+        ax.ticklabel_format(useOffset=use_offset)
 
         if getattr(ax, 'qcodes_colorbar', None):
             # update_normal doesn't seem to work...
@@ -203,4 +285,9 @@ def _draw_pcolormesh(self, ax, z, x=None, y=None, subplot=1, **kwargs):
             # put this where it belongs.
             ax.qcodes_colorbar.set_label(self.get_label(z))
 
+        # Scale colors if z has elements
+        cmin = np.nanmin(args_masked[-1])
+        cmax = np.nanmax(args_masked[-1])
+        ax.qcodes_colorbar.set_clim(cmin, cmax)
+
         return pc