From 3bb362d4aacd45ed229948b4364752eee8062028 Mon Sep 17 00:00:00 2001
From: Jens Hedegaard Nielsen <j.h.nielsen@nbi.ku.dk>
Date: Tue, 16 May 2017 10:54:08 +0200
Subject: [PATCH] Fix: Matplot interpolate x and y coords for pcolormesh

Fixes #257
---
 qcodes/plots/qcmatplotlib.py | 38 ++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/qcodes/plots/qcmatplotlib.py b/qcodes/plots/qcmatplotlib.py
index 1d5305b2a7c..c174f188679 100644
--- a/qcodes/plots/qcmatplotlib.py
+++ b/qcodes/plots/qcmatplotlib.py
@@ -109,6 +109,26 @@ 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())
 
+    @staticmethod
+    def _calc_corners(array):
+        """
+        Extrapolate coordinates of between data points from a list of data points.
+        Needed because pcolor and friends take the list of edges as x and y rather 
+        than the center positions. This should handle general non linear axis.
+        """
+        myarray = array.copy()
+        # extrapolate points before and after array
+        # so we can find steps in the end/start
+        myarray_plusone = 2 * myarray[-1] - myarray[-2]
+        myarray_minusone = 2 * myarray[0] - myarray[1]
+        myarray = np.append(myarray, myarray_plusone)
+        myarray = np.insert(myarray, 0, myarray_minusone)
+        # find the difference between steps
+        stepsizes = np.diff(myarray)
+        # calculate edges
+        myarray_corners = myarray[0:-1] + stepsizes / 2
+        return myarray_corners
+
     def _get_axes(self, config):
         return self.subplots[config.get('subplot', 1) - 1]
 
@@ -218,6 +238,24 @@ def _draw_pcolormesh(self, ax, z, x=None, y=None, subplot=1,
                          yunit=None,
                          zunit=None,
                          **kwargs):
+
+        # Polormesh and friends expect coordinates of corners not the center of datapoints
+        # here we expands the x and y coordinates to match what pcolormesh expects.
+        # first we verify that the shapes are as expected
+
+        # x is the inner most index so it is 2 dimensional.
+        # however y is outer most so it is 1 dimensional
+        if x:
+            assert x.shape == z.shape
+            new_xshape = np.array(x.shape) + 1
+            new_x = np.zeros(new_xshape)
+            for i in range(len(x[:, 0])):
+                new_x[i] = self._calc_corners(x[i])
+            new_x[i+1] = new_x[i]
+            x = new_x
+        if y:
+            assert y.shape[0] == z.shape[0] and len(y.shape)==1
+            y = self._calc_corners(y)
         # 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