Repository: microsoft/Qcodes

Issue:
Plot_dataset for partially completed 2D
In one of the latest updates the plotting functions doesnt handle partially completed 2D dataset. It now falls back to a scatter plot instead of a partially filled map.
Ideally we would like to keep the 0.22.0 version behavior for partially completed plots.

Version 0.25.0
![image](https://user-images.githubusercontent.com/23334244/117890066-cb52f100-b2b4-11eb-9282-6aadc3af92ce.png)

Version 0.22.0
![image](https://user-images.githubusercontent.com/23334244/117890031-bc6c3e80-b2b4-11eb-8f34-0de861f3cf63.png)


Discussion:
This is caused by the dataset now being plotted from the cache which uses nan as a filler for the empty elements rather than simply omitting them. Adding something like `dataarray = dataarray[~np.isnan(dataarray)]` to `flatten_1D_data_for_plot` seems to fix this but it is unclear if this will have problematic issues for datasets that really contain nans 
A better solution would therefor be to make use of the actual shape data if available 

Resolving patch:
diff --git a/qcodes/dataset/data_export.py b/qcodes/dataset/data_export.py
--- a/qcodes/dataset/data_export.py
+++ b/qcodes/dataset/data_export.py
@@ -1,7 +1,8 @@
 import logging
-from typing import Any, Dict, List, Sequence, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
+from typing_extensions import TypedDict
 
 from qcodes.dataset.data_set import DataSet, load_by_id
 from qcodes.dataset.descriptions.param_spec import ParamSpecBase
@@ -10,11 +11,23 @@
 log = logging.getLogger(__name__)
 
 
+class DSPlotData(TypedDict):
+    """
+    The dictionary used to represent data for use within `plot_dataset`
+    """
+    name: str
+    unit: str
+    label: str
+    data: np.ndarray
+    shape: Optional[Tuple[int, ...]]
+
+
+@deprecate(alternative="ndarray.flatten()")
 def flatten_1D_data_for_plot(rawdata: Union[Sequence[Sequence[Any]],
                                             np.ndarray]) -> np.ndarray:
     """
     Cast the return value of the database query to
-    a numpy array
+    a 1D numpy array
 
     Args:
         rawdata: The return of the get_values function
@@ -23,16 +36,12 @@ def flatten_1D_data_for_plot(rawdata: Union[Sequence[Sequence[Any]],
         A one-dimensional numpy array
 
     """
-    dataarray = np.array(rawdata)
-    shape = np.shape(dataarray)
-    dataarray = dataarray.reshape(np.product(shape))
-
+    dataarray = np.array(rawdata).flatten()
     return dataarray
 
 
 @deprecate(alternative="dataset.get_parameter_data")
-def get_data_by_id(run_id: int) -> \
-        List[List[Dict[str, Union[str, np.ndarray]]]]:
+def get_data_by_id(run_id: int) -> List[List[DSPlotData]]:
     """
     Load data from database and reshapes into 1D arrays with minimal
     name, unit and label metadata.
@@ -75,7 +84,9 @@ def get_data_by_id(run_id: int) -> \
     return output
 
 
-def _get_data_from_ds(ds: DataSet) -> List[List[Dict[str, Union[str, np.ndarray]]]]:
+def _get_data_from_ds(
+    ds: DataSet,
+) -> List[List[DSPlotData]]:
     dependent_parameters: Tuple[ParamSpecBase, ...] = ds.dependent_parameters
 
     all_data = ds.cache.data()
@@ -91,14 +102,18 @@ def _get_data_from_ds(ds: DataSet) -> List[List[Dict[str, Union[str, np.ndarray]
         dependencies = ds.description.interdeps.dependencies[dependent]
 
         for param_spec_base in dependencies + (dependent,):
-            my_data_dict: Dict[str, Union[str, np.ndarray]] = {
+            my_data_dict: DSPlotData = {
                 "name": param_spec_base.name,
                 "unit": param_spec_base.unit,
                 "label": param_spec_base.label,
-                "data": data_dict[param_spec_base.name].flatten(),
+                "data": data_dict[param_spec_base.name],
+                "shape": None,
             }
             data_dicts_list.append(my_data_dict)
 
+        if ds.description.shapes is not None:
+            data_dicts_list[-1]["shape"] = ds.description.shapes.get(dependent.name)
+
         output.append(data_dicts_list)
 
     return output
diff --git a/qcodes/dataset/plotting.py b/qcodes/dataset/plotting.py
--- a/qcodes/dataset/plotting.py
+++ b/qcodes/dataset/plotting.py
@@ -7,7 +7,7 @@
 import logging
 from contextlib import contextmanager
 from functools import partial
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, List, Optional, Sequence, Tuple, Union, cast
 
 import matplotlib
 import matplotlib.pyplot as plt
@@ -19,9 +19,9 @@
 from qcodes.utils.plotting import auto_color_scale_from_config, find_scale_and_prefix
 
 from .data_export import (
+    DSPlotData,
     _get_data_from_ds,
     _strings_as_ints,
-    flatten_1D_data_for_plot,
     get_1D_plottype,
     get_2D_plottype,
     reshape_2D_data,
@@ -33,10 +33,9 @@
 AxesTuple = Tuple[matplotlib.axes.Axes, matplotlib.colorbar.Colorbar]
 AxesTupleList = Tuple[List[matplotlib.axes.Axes],
                       List[Optional[matplotlib.colorbar.Colorbar]]]
-Number = Union[float, int]
 # NamedData is the structure _get_data_from_ds returns and that plot_by_id
 # uses internally
-NamedData = List[List[Dict[str, Union[str, np.ndarray]]]]
+NamedData = List[List[DSPlotData]]
 
 # list of kwargs for plotting function, so that kwargs can be passed to
 # :func:`plot_dataset` and will be distributed to the respective plotting func.
@@ -90,19 +89,19 @@ def heatmaphandler(**kwargs: Any) -> Any:
     yield plot_handler_mapping[plottype](**kwargs.copy())
 
 
-def plot_dataset(dataset: DataSet,
-                 axes: Optional[Union[matplotlib.axes.Axes,
-                                      Sequence[matplotlib.axes.Axes]]] = None,
-                 colorbars: Optional[Union[matplotlib.colorbar.Colorbar,
-                                           Sequence[
-                                        matplotlib.colorbar.Colorbar]]] = None,
-                 rescale_axes: bool = True,
-                 auto_color_scale: Optional[bool] = None,
-                 cutoff_percentile: Optional[Union[Tuple[Number, Number],
-                                                   Number]] = None,
-                 complex_plot_type: str = 'real_and_imag',
-                 complex_plot_phase: str = 'radians',
-                 **kwargs: Any) -> AxesTupleList:
+def plot_dataset(
+    dataset: DataSet,
+    axes: Optional[Union[matplotlib.axes.Axes, Sequence[matplotlib.axes.Axes]]] = None,
+    colorbars: Optional[
+        Union[matplotlib.colorbar.Colorbar, Sequence[matplotlib.colorbar.Colorbar]]
+    ] = None,
+    rescale_axes: bool = True,
+    auto_color_scale: Optional[bool] = None,
+    cutoff_percentile: Optional[Union[Tuple[float, float], float]] = None,
+    complex_plot_type: str = "real_and_imag",
+    complex_plot_phase: str = "radians",
+    **kwargs: Any,
+) -> AxesTupleList:
     """
     Construct all plots for a given dataset
 
@@ -218,8 +217,8 @@ def plot_dataset(dataset: DataSet,
         if len(data) == 2:  # 1D PLOTTING
             log.debug(f'Doing a 1D plot with kwargs: {kwargs}')
 
-            xpoints = cast(np.ndarray, data[0]['data'])
-            ypoints = cast(np.ndarray, data[1]['data'])
+            xpoints = data[0]["data"]
+            ypoints = data[1]["data"]
 
             plottype = get_1D_plottype(xpoints, ypoints)
             log.debug(f'Determined plottype: {plottype}')
@@ -256,17 +255,17 @@ def plot_dataset(dataset: DataSet,
         elif len(data) == 3:  # 2D PLOTTING
             log.debug(f'Doing a 2D plot with kwargs: {kwargs}')
 
-            # From the setpoints, figure out which 2D plotter to use
-            # TODO: The "decision tree" for what gets plotted how and how
-            # we check for that is still unfinished/not optimised
-
-            xpoints = flatten_1D_data_for_plot(data[0]['data'])
-            ypoints = flatten_1D_data_for_plot(data[1]['data'])
-            zpoints = flatten_1D_data_for_plot(data[2]['data'])
-
-            plottype = get_2D_plottype(xpoints, ypoints, zpoints)
-
-            log.debug(f'Determined plottype: {plottype}')
+            if data[2]["shape"] is None:
+                xpoints = data[0]["data"].flatten()
+                ypoints = data[1]["data"].flatten()
+                zpoints = data[2]["data"].flatten()
+                plottype = get_2D_plottype(xpoints, ypoints, zpoints)
+                log.debug(f"Determined plottype: {plottype}")
+            else:
+                xpoints = data[0]["data"]
+                ypoints = data[1]["data"]
+                zpoints = data[2]["data"]
+                plottype = "2D_grid"
 
             how_to_plot = {'2D_grid': plot_on_a_plain_grid,
                            '2D_equidistant': plot_on_a_plain_grid,
@@ -305,19 +304,19 @@ def plot_dataset(dataset: DataSet,
     return axeslist, new_colorbars
 
 
-def plot_by_id(run_id: int,
-               axes: Optional[Union[matplotlib.axes.Axes,
-                              Sequence[matplotlib.axes.Axes]]] = None,
-               colorbars: Optional[Union[matplotlib.colorbar.Colorbar,
-                                   Sequence[
-                                       matplotlib.colorbar.Colorbar]]] = None,
-               rescale_axes: bool = True,
-               auto_color_scale: Optional[bool] = None,
-               cutoff_percentile: Optional[Union[Tuple[Number, Number],
-                                                 Number]] = None,
-               complex_plot_type: str = 'real_and_imag',
-               complex_plot_phase: str = 'radians',
-               **kwargs: Any) -> AxesTupleList:
+def plot_by_id(
+    run_id: int,
+    axes: Optional[Union[matplotlib.axes.Axes, Sequence[matplotlib.axes.Axes]]] = None,
+    colorbars: Optional[
+        Union[matplotlib.colorbar.Colorbar, Sequence[matplotlib.colorbar.Colorbar]]
+    ] = None,
+    rescale_axes: bool = True,
+    auto_color_scale: Optional[bool] = None,
+    cutoff_percentile: Optional[Union[Tuple[float, float], float]] = None,
+    complex_plot_type: str = "real_and_imag",
+    complex_plot_phase: str = "radians",
+    **kwargs: Any,
+) -> AxesTupleList:
     """
     Construct all plots for a given `run_id`. Here `run_id` is an
     alias for `captured_run_id` for historical reasons. See the docs
@@ -338,9 +337,9 @@ def plot_by_id(run_id: int,
                         **kwargs)
 
 
-def _complex_to_real_preparser(alldata: NamedData,
-                               conversion: str,
-                               degrees: bool = False) -> NamedData:
+def _complex_to_real_preparser(
+    alldata: Sequence[Sequence[DSPlotData]], conversion: str, degrees: bool = False
+) -> NamedData:
     """
     Convert complex-valued parameters to two real-valued parameters, either
     real and imaginary part or phase and magnitude part
@@ -358,7 +357,7 @@ def _complex_to_real_preparser(alldata: NamedData,
                          'but can only accept "real_and_imag" or '
                          '"mag_and_phase".')
 
-    newdata = []
+    newdata: NamedData = []
 
     # we build a new NamedData object from the given `alldata` input.
     # Note that the length of `newdata` will be larger than that of `alldata`
@@ -371,11 +370,11 @@ def _complex_to_real_preparser(alldata: NamedData,
         new_group = []
         new_groups: NamedData = [[], []]
         for index, parameter in enumerate(group):
-            data = cast(np.ndarray, parameter['data'])
-            if data.dtype.kind == 'c':
-                p1, p2 = _convert_complex_to_real(parameter,
-                                                  conversion=conversion,
-                                                  degrees=degrees)
+            data = parameter["data"]
+            if data.dtype.kind == "c":
+                p1, p2 = _convert_complex_to_real(
+                    parameter, conversion=conversion, degrees=degrees
+                )
                 if index < len(group) - 1:
                     # if the above condition is met, we are dealing with
                     # complex setpoints
@@ -405,11 +404,8 @@ def _complex_to_real_preparser(alldata: NamedData,
 
 
 def _convert_complex_to_real(
-        parameter: Dict[str, Union[str, np.ndarray]],
-        conversion: str,
-        degrees: bool
-        ) -> Tuple[Dict[str, Union[str, np.ndarray]],
-                   Dict[str, Union[str, np.ndarray]]]:
+    parameter: DSPlotData, conversion: str, degrees: bool
+) -> Tuple[DSPlotData, DSPlotData]:
     """
     Do the actual conversion and turn one parameter into two.
     Should only be called from within _complex_to_real_preparser.
@@ -433,20 +429,26 @@ def _convert_complex_to_real(
     new_units = converters['units'][conversion](parameter['unit'])
     new_names = converters['names'][conversion](parameter['name'])
 
-    new_parameters = tuple(
-        {'name': name, 'label': label,
-         'unit': unit, 'data': data}
-        for name, label, unit, data in zip(
-            new_names, new_labels, new_units, new_data))
+    parameter1: DSPlotData = {
+        "name": new_names[0],
+        "label": new_labels[0],
+        "unit": new_units[0],
+        "data": new_data[0],
+        "shape": parameter["shape"],
+    }
 
-    # The reason we ignore the type in the return is that I cannot figure
-    # out how to get mypy to correctly infer the type of iterated values
-    # (the name, label, unit, and data above)
+    parameter2: DSPlotData = {
+        "name": new_names[1],
+        "label": new_labels[1],
+        "unit": new_units[1],
+        "data": new_data[1],
+        "shape": parameter["shape"],
+    }
 
-    return new_parameters  # type: ignore[return-value]
+    return parameter1, parameter2
 
 
-def _get_label_of_data(data_dict: Dict[str, Any]) -> str:
+def _get_label_of_data(data_dict: DSPlotData) -> str:
     return data_dict['label'] if data_dict['label'] != '' \
         else data_dict['name']
 
@@ -458,17 +460,17 @@ def _make_axis_label(label: str, unit: str) -> str:
     return label
 
 
-def _make_label_for_data_axis(data: List[Dict[str, Any]], axis_index: int
-                              ) -> str:
+def _make_label_for_data_axis(data: Sequence[DSPlotData], axis_index: int) -> str:
     label = _get_label_of_data(data[axis_index])
     unit = data[axis_index]['unit']
     return _make_axis_label(label, unit)
 
 
-def _set_data_axes_labels(ax: matplotlib.axes.Axes,
-                          data: List[Dict[str, Any]],
-                          cax: Optional[matplotlib.colorbar.Colorbar] = None
-                          ) -> None:
+def _set_data_axes_labels(
+    ax: matplotlib.axes.Axes,
+    data: Sequence[DSPlotData],
+    cax: Optional[matplotlib.colorbar.Colorbar] = None,
+) -> None:
     ax.set_xlabel(_make_label_for_data_axis(data, 0))
     ax.set_ylabel(_make_label_for_data_axis(data, 1))
 
@@ -578,24 +580,24 @@ def plot_on_a_plain_grid(x: np.ndarray,
         z_strings = np.unique(z)
         z = _strings_as_ints(z)
 
-    xrow, yrow, z_to_plot = reshape_2D_data(x, y, z)
-
-    # we use a general edge calculator,
-    # in the case of non-equidistantly spaced data
-    # TODO: is this appropriate for a log ax?
-    dxs = np.diff(xrow)/2
-    dys = np.diff(yrow)/2
-    x_edges = np.concatenate((np.array([xrow[0] - dxs[0]]),
-                              xrow[:-1] + dxs,
-                              np.array([xrow[-1] + dxs[-1]])))
-    y_edges = np.concatenate((np.array([yrow[0] - dys[0]]),
-                              yrow[:-1] + dys,
-                              np.array([yrow[-1] + dys[-1]])))
+    if x.ndim == 2 and y.ndim == 2 and z.ndim == 2:
+        if not np.logical_or(np.any(np.isnan(x)), np.any(np.isnan(y))):
+            # data is on a grid that may or may not be
+            # rectilinear. Rely on matplotlib to plot
+            # this directly
+            x_to_plot, y_to_plot, z_to_plot = x, y, z
+            num_points = x_to_plot.size
+        else:
+            x_to_plot, y_to_plot, z_to_plot = _clip_nan_from_shaped_data(x, y, z)
+            num_points = x_to_plot.size * y_to_plot.size
+    else:
+        x_to_plot, y_to_plot, z_to_plot = reshape_2D_data(x, y, z)
+        num_points = x_to_plot.size * y_to_plot.size
+
     if 'rasterized' in kwargs.keys():
         rasterized = kwargs.pop('rasterized')
     else:
-        rasterized = len(x_edges) * len(y_edges) \
-                      > qc.config.plotting.rasterize_threshold
+        rasterized = num_points > qc.config.plotting.rasterize_threshold
 
     cmap = kwargs.pop('cmap') if 'cmap' in kwargs else None
 
@@ -603,11 +605,15 @@ def plot_on_a_plain_grid(x: np.ndarray,
         name = cmap.name if hasattr(cmap, 'name') else 'viridis'
         cmap = matplotlib.cm.get_cmap(name, len(z_strings))
 
-    colormesh = ax.pcolormesh(x_edges, y_edges,
-                              np.ma.masked_invalid(z_to_plot),
-                              rasterized=rasterized,
-                              cmap=cmap,
-                              **kwargs)
+    colormesh = ax.pcolormesh(
+        x_to_plot,
+        y_to_plot,
+        np.ma.masked_invalid(z_to_plot),
+        rasterized=rasterized,
+        cmap=cmap,
+        shading="nearest",
+        **kwargs,
+    )
 
     if x_is_stringy:
         ax.set_xticks(np.arange(len(np.unique(x_strings))))
@@ -631,6 +637,50 @@ def plot_on_a_plain_grid(x: np.ndarray,
     return ax, colorbar
 
 
+def _clip_nan_from_shaped_data(
+    x: np.ndarray, y: np.ndarray, z: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    def _on_rectilinear_grid_except_nan(x_data: np.ndarray, y_data: np.ndarray) -> bool:
+        """
+        check that data is on a rectilinear grid. e.g. all points are the same as the  first
+        row and column with the exception of nans. Those represent points not yet measured.
+        """
+        x_row = x_data[:, 0:1]
+        y_row = y_data[0:1, :]
+        return (
+            np.nanmax(np.abs(x_data - x_row)) == 0
+            and np.nanmax(np.abs(y_data - y_row)) == 0
+        )
+
+    if _on_rectilinear_grid_except_nan(x, y):
+        # clip any row or column where there are nans in the first row
+        # or column. Since we fill from here we assume that means that no data
+        # has been measured for this row
+        x_to_plot, y_to_plot = x[:, 0], y[0, :]
+
+        filter_x = ~np.isnan(x_to_plot)
+        filter_y = ~np.isnan(y_to_plot)
+
+        x_to_plot = x_to_plot[filter_x]
+        y_to_plot = y_to_plot[filter_y]
+        z_to_plot = z[filter_x, :]
+        z_to_plot = z_to_plot[:, filter_y].transpose()
+    else:
+        # fallback to flattening the data and use the same path as
+        # non shaped data after filtering the nans.
+        # this is not ideal as we loose the shape data but
+        # not clear how to do this better. Either return a ragged
+        # array or clip all inner dims that have nans completely
+        x = x.flatten()
+        y = y.flatten()
+        z = z.flatten()
+        filter_nans = np.logical_and(~np.isnan(x), ~np.isnan(y))
+        x_to_plot, y_to_plot, z_to_plot = reshape_2D_data(
+            x[filter_nans], y[filter_nans], z[filter_nans]
+        )
+    return x_to_plot, y_to_plot, z_to_plot
+
+
 def _scale_formatter(tick_value: float, pos: int, factor: float) -> str:
     """
     Function for matplotlib.ticker.FuncFormatter that scales the tick values
@@ -639,8 +689,9 @@ def _scale_formatter(tick_value: float, pos: int, factor: float) -> str:
     return f"{tick_value*factor:g}"
 
 
-def _make_rescaled_ticks_and_units(data_dict: Dict[str, Any]) \
-        -> Tuple[matplotlib.ticker.FuncFormatter, str]:
+def _make_rescaled_ticks_and_units(
+    data_dict: DSPlotData,
+) -> Tuple[matplotlib.ticker.FuncFormatter, str]:
     """
     Create a ticks formatter and a new label for the data that is to be used
     on the axes where the data is plotted.
@@ -684,10 +735,11 @@ def _make_rescaled_ticks_and_units(data_dict: Dict[str, Any]) \
     return ticks_formatter, new_label
 
 
-def _rescale_ticks_and_units(ax: matplotlib.axes.Axes,
-                             data: List[Dict[str, Any]],
-                             cax: matplotlib.colorbar.Colorbar = None
-                             ) -> None:
+def _rescale_ticks_and_units(
+    ax: matplotlib.axes.Axes,
+    data: Sequence[DSPlotData],
+    cax: matplotlib.colorbar.Colorbar = None,
+) -> None:
     """
     Rescale ticks and units for the provided axes as described in
     :func:`~_make_rescaled_ticks_and_units`


Test patch:
diff --git a/qcodes/tests/dataset/test__get_data_from_ds.py b/qcodes/tests/dataset/test__get_data_from_ds.py
--- a/qcodes/tests/dataset/test__get_data_from_ds.py
+++ b/qcodes/tests/dataset/test__get_data_from_ds.py
@@ -1,13 +1,13 @@
-import pytest
-import numpy as np
-from numpy.testing import assert_allclose
 import hypothesis.strategies as hst
+import numpy as np
+import pytest
 from hypothesis import HealthCheck, given, settings
+from numpy.testing import assert_allclose
 
 import qcodes as qc
-from qcodes.dataset.data_export import get_data_by_id, _get_data_from_ds
-from qcodes.dataset.descriptions.param_spec import ParamSpecBase
+from qcodes.dataset.data_export import _get_data_from_ds, get_data_by_id
 from qcodes.dataset.descriptions.dependencies import InterDependencies_
+from qcodes.dataset.descriptions.param_spec import ParamSpecBase
 from qcodes.dataset.measurements import Measurement
 from qcodes.utils.deprecate import QCoDeSDeprecationWarning
 
@@ -179,7 +179,10 @@ def test_datasaver_array_parameters_channel(channel_array_instrument,
     datadicts = datadicts[0]
     assert len(datadicts) == len(meas.parameters)
     for datadict in datadicts:
-        assert datadict['data'].shape == (N * M,)
+        if storage_type == "array":
+            assert datadict["data"].shape == (N, M)
+        else:
+            assert datadict["data"].shape == (N * M,)
 
 
 @settings(max_examples=5, deadline=None,
@@ -219,15 +222,15 @@ def test_datasaver_array_parameters_array(channel_array_instrument, DAC, N,
     assert len(datadicts) == len(meas.parameters)
     for datadict in datadicts:
         if datadict['name'] == 'dummy_dac_ch1':
-            expected_data = np.repeat(dac_datapoints, M)
+            expected_data = np.repeat(dac_datapoints, M).reshape(N, M)
         if datadict['name'] == dependency_name:
-            expected_data = np.tile(np.linspace(5, 9, 5), N)
+            expected_data = np.tile(np.linspace(5, 9, 5), (N, 1))
         if datadict['name'] == 'dummy_channel_inst_ChanA_dummy_array_parameter':
-            expected_data = np.empty(N * M)
+            expected_data = np.empty((N, M))
             expected_data[:] = 2.
         assert_allclose(datadict['data'], expected_data)
 
-        assert datadict['data'].shape == (N * M,)
+        assert datadict["data"].shape == (N, M)
 
 
 @pytest.mark.parametrize("bg_writing", [True, False])
@@ -267,11 +270,11 @@ def test_datasaver_multidim_array(experiment, bg_writing):
     for datadict_list in datadicts:
         assert len(datadict_list) == 3
         for datadict in datadict_list:
-            dataindex = data_mapping[datadict['name']]
-            expected_data = data[dataindex, :, :].ravel()
-            assert_allclose(datadict['data'], expected_data)
+            dataindex = data_mapping[datadict["name"]]
+            expected_data = data[dataindex : dataindex + 1, :, :]
+            assert_allclose(datadict["data"], expected_data)
 
-            assert datadict['data'].shape == (size1 * size2,)
+            assert datadict["data"].shape == (1, size1, size2)
 
 
 @pytest.mark.parametrize("bg_writing", [True, False])
diff --git a/qcodes/tests/dataset/test_plotting.py b/qcodes/tests/dataset/test_plotting.py
--- a/qcodes/tests/dataset/test_plotting.py
+++ b/qcodes/tests/dataset/test_plotting.py
@@ -1,38 +1,57 @@
+from typing import List
+
 import numpy as np
-from hypothesis import given, example, assume, settings, HealthCheck
-from hypothesis.strategies import text, sampled_from, floats, lists, data, \
-    one_of, just
+import pytest
+from hypothesis import HealthCheck, assume, example, given, settings
+from hypothesis.strategies import data, floats, just, lists, one_of, sampled_from, text
+from matplotlib.collections import QuadMesh
 
 import qcodes as qc
-from qcodes.dataset.plotting import _make_rescaled_ticks_and_units
-from qcodes.utils.plotting import _ENGINEERING_PREFIXES, _UNITS_FOR_RESCALING
-
-from qcodes.dataset.plotting import (plot_by_id, _appropriate_kwargs,
-    _complex_to_real_preparser)
+from qcodes.dataset.data_export import DSPlotData
+from qcodes.dataset.descriptions.detect_shapes import detect_shape_of_measurement
 from qcodes.dataset.measurements import Measurement
+from qcodes.dataset.plotting import (
+    _appropriate_kwargs,
+    _complex_to_real_preparser,
+    _make_rescaled_ticks_and_units,
+    plot_by_id,
+    plot_dataset,
+)
 from qcodes.tests.instrument_mocks import DummyInstrument
+from qcodes.utils.plotting import _ENGINEERING_PREFIXES, _UNITS_FOR_RESCALING
 
 
-@given(param_name=text(min_size=1, max_size=10),
-       param_label=text(min_size=0, max_size=15),
-       scale=sampled_from(sorted(list(_ENGINEERING_PREFIXES.keys()))),
-       unit=sampled_from(sorted(list(
-           _UNITS_FOR_RESCALING.union(
-               ['', 'unit', 'kg', '%', 'permille', 'nW'])))),
-       data_strategy=data()
-       )
-@example(param_name='huge_param',
-         param_label='Larger than the highest scale',
-         scale=max(list(_ENGINEERING_PREFIXES.keys())),
-         unit='V',
-         data_strategy=np.random.random((5,))
-                       * 10 ** (3 + max(list(_ENGINEERING_PREFIXES.keys()))))
-@example(param_name='small_param',
-         param_label='Lower than the lowest scale',
-         scale=min(list(_ENGINEERING_PREFIXES.keys())),
-         unit='V',
-         data_strategy=np.random.random((5,))
-                       * 10 ** (-3 + min(list(_ENGINEERING_PREFIXES.keys()))))
+class TerminateLoopException(Exception):
+    pass
+
+
+@given(
+    param_name=text(min_size=1, max_size=10),
+    param_label=text(min_size=0, max_size=15),
+    scale=sampled_from(sorted(list(_ENGINEERING_PREFIXES.keys()))),
+    unit=sampled_from(
+        sorted(
+            list(_UNITS_FOR_RESCALING.union(["", "unit", "kg", "%", "permille", "nW"]))
+        )
+    ),
+    data_strategy=data(),
+)
+@example(
+    param_name="huge_param",
+    param_label="Larger than the highest scale",
+    scale=max(list(_ENGINEERING_PREFIXES.keys())),
+    unit="V",
+    data_strategy=np.random.random((5,))
+    * 10 ** (3 + max(list(_ENGINEERING_PREFIXES.keys()))),
+)
+@example(
+    param_name="small_param",
+    param_label="Lower than the lowest scale",
+    scale=min(list(_ENGINEERING_PREFIXES.keys())),
+    unit="V",
+    data_strategy=np.random.random((5,))
+    * 10 ** (-3 + min(list(_ENGINEERING_PREFIXES.keys()))),
+)
 @settings(suppress_health_check=[HealthCheck.too_slow])
 def test_rescaled_ticks_and_units(scale, unit,
                                   param_name, param_label, data_strategy):
@@ -105,23 +124,97 @@ def test_plot_by_id_line_and_heatmap(experiment, request):
     meas = Measurement()
     meas.register_parameter(inst.s1)
     meas.register_parameter(inst.s2)
-    meas.register_parameter(inst.m2, setpoints=(inst.s1, inst.s2) )
+    meas.register_parameter(inst.m2, setpoints=(inst.s1, inst.s2))
     meas.register_parameter(inst.m1, setpoints=(inst.s1,))
 
     with meas.run() as datasaver:
         for outer in range(10):
-            datasaver.add_result((inst.s1, outer),
-                                (inst.m1, inst.m1()))
+            datasaver.add_result((inst.s1, outer), (inst.m1, inst.m1()))
             for inner in range(10):
-                datasaver.add_result((inst.s1, outer),
-                                    (inst.s2, inner),
-                                    (inst.m2, inst.m2()))
+                datasaver.add_result(
+                    (inst.s1, outer), (inst.s2, inner), (inst.m2, inst.m2())
+                )
 
     dataid = datasaver.run_id
     plot_by_id(dataid)
     plot_by_id(dataid, cmap='bone')
 
 
+@pytest.mark.parametrize("nan_setpoints", [True, False])
+@pytest.mark.parametrize("shifted", [True, False])
+def test_plot_dataset_2d_shaped(experiment, request, nan_setpoints, shifted):
+    """
+    Test plotting of preshaped data on a grid that may or may not be shifted
+    with and without nans in the set points.
+    """
+    inst = DummyInstrument("dummy", gates=["s1", "m1", "s2"])
+    request.addfinalizer(inst.close)
+
+    inst.m1.get = np.random.randn
+
+    meas = Measurement()
+    meas.register_parameter(inst.s1)
+    meas.register_parameter(inst.s2)
+    meas.register_parameter(inst.m1, setpoints=(inst.s1, inst.s2))
+
+    outer_shape = 10
+    inner_shape = 20
+
+    meas.set_shapes(detect_shape_of_measurement((inst.m1,), (outer_shape, inner_shape)))
+
+    shift = 0
+
+    with meas.run() as datasaver:
+        try:
+            for outer in np.linspace(0, 9, outer_shape):
+                for inner in np.linspace(0 + shift, 10 + shift, inner_shape):
+                    datasaver.add_result(
+                        (inst.s1, outer), (inst.s2, inner), (inst.m1, inst.m1())
+                    )
+                    if inner > 7 and outer > 6 and nan_setpoints:
+                        raise TerminateLoopException
+                if shifted:
+                    shift += 1
+        except TerminateLoopException:
+            pass
+
+    axes, cbs = plot_dataset(datasaver.dataset)
+    xlims = axes[0].get_xlim()
+    ylims = axes[0].get_ylim()
+
+    # check that this generates a QuadMesh which is the expected output of pcolormesh
+    assert any(isinstance(mplobj, QuadMesh) for mplobj in axes[0].get_children())
+
+    if nan_setpoints and shifted:
+        assert xlims[0] == -0.5
+        assert xlims[1] == 7.5
+        assert ylims[0] < 0
+        assert ylims[0] > -1.0
+        assert ylims[1] > 16
+        assert ylims[1] < 17
+    elif not nan_setpoints and shifted:
+        assert xlims[0] == -0.5
+        assert xlims[1] == 9.5
+        assert ylims[0] < 0
+        assert ylims[0] > -1.0
+        assert ylims[1] > 19
+        assert ylims[1] < 20
+    elif nan_setpoints and not shifted:
+        assert xlims[0] == -0.5
+        assert xlims[1] == 7.5
+        assert ylims[0] < 0
+        assert ylims[0] > -1.0
+        assert ylims[1] > 10
+        assert ylims[1] < 11
+    else:
+        assert xlims[0] == -0.5
+        assert xlims[1] == 9.5
+        assert ylims[0] < 0
+        assert ylims[0] > -1.0
+        assert ylims[1] > 10
+        assert ylims[1] < 11
+
+
 def test_appropriate_kwargs():
 
     kwargs = {'cmap': 'bone'}
@@ -152,16 +245,26 @@ def test_appropriate_kwargs():
         assert ap_kwargs['cmap'] == qc.config.plotting.default_color_map
 
 
-def test__complex_to_real_preparser_complex_toplevel_param():
-
-    data_in = [[{'data': np.array([0, 1, 2]),
-                 'name': 'voltage',
-                 'label': 'swept voltage',
-                 'unit': 'V'},
-                {'data': np.array([0+0j, 1+2j, -1+1j]),
-                 'name': 'signal',
-                 'label': 'complex signal',
-                 'unit': 'Ohm'}]]
+def test__complex_to_real_preparser_complex_toplevel_param() -> None:
+
+    data_in: List[List[DSPlotData]] = [
+        [
+            {
+                "data": np.array([0, 1, 2]),
+                "name": "voltage",
+                "label": "swept voltage",
+                "unit": "V",
+                "shape": None,
+            },
+            {
+                "data": np.array([0 + 0j, 1 + 2j, -1 + 1j]),
+                "name": "signal",
+                "label": "complex signal",
+                "unit": "Ohm",
+                "shape": None,
+            },
+        ]
+    ]
 
     data_out = _complex_to_real_preparser(data_in, conversion='real_and_imag')
 
@@ -169,10 +272,10 @@ def test__complex_to_real_preparser_complex_toplevel_param():
     assert data_out[0][0] == data_in[0][0]
 
     real_param = data_out[0][1]
-    assert real_param['name'] == 'signal_real'
-    assert real_param['label'] =='complex signal [real]'
-    assert all(real_param['data'] == np.array([0, 1, -1]))
-    assert real_param['unit'] == 'Ohm'
+    assert real_param["name"] == "signal_real"
+    assert real_param["label"] == "complex signal [real]"
+    assert all(real_param["data"] == np.array([0, 1, -1]))
+    assert real_param["unit"] == "Ohm"
 
     imag_param = data_out[1][1]
     assert imag_param['name'] == 'signal_imag'
@@ -187,10 +290,10 @@ def test__complex_to_real_preparser_complex_toplevel_param():
     assert data_out[0][0] == data_in[0][0]
 
     phase_param = data_out[1][1]
-    assert phase_param['name'] == 'signal_phase'
-    assert phase_param['label'] =='complex signal [phase]'
-    assert all(phase_param['data'] == np.angle(np.array([0+0j, 1+2j, -1+1j])))
-    assert phase_param['unit'] == 'rad'
+    assert phase_param["name"] == "signal_phase"
+    assert phase_param["label"] == "complex signal [phase]"
+    assert all(phase_param["data"] == np.angle(np.array([0 + 0j, 1 + 2j, -1 + 1j])))
+    assert phase_param["unit"] == "rad"
 
     mag_param = data_out[0][1]
     assert mag_param['name'] == 'signal_mag'
@@ -202,24 +305,34 @@ def test__complex_to_real_preparser_complex_toplevel_param():
                                           degrees=True)
 
     phase_param = data_out[1][1]
-    assert phase_param['name'] == 'signal_phase'
-    assert phase_param['label'] =='complex signal [phase]'
-    assert all(phase_param['data'] == np.angle(np.array([0+0j, 1+2j, -1+1j]),
-                                               deg=True))
-    assert phase_param['unit'] == 'deg'
-
-
-def test__complex_to_real_preparser_complex_setpoint():
-
-    data_in = [[{'data': np.array([0+0j, 1+2j, -1+1j]),
-                 'name': 'signal',
-                 'label': 'complex signal',
-                 'unit': 'Ohm'},
-                 {'data': np.array([0, 1, 2]),
-                 'name': 'voltage',
-                 'label': 'measured voltage',
-                 'unit': 'V'}
-                ]]
+    assert phase_param["name"] == "signal_phase"
+    assert phase_param["label"] == "complex signal [phase]"
+    assert all(
+        phase_param["data"] == np.angle(np.array([0 + 0j, 1 + 2j, -1 + 1j]), deg=True)
+    )
+    assert phase_param["unit"] == "deg"
+
+
+def test__complex_to_real_preparser_complex_setpoint() -> None:
+
+    data_in: List[List[DSPlotData]] = [
+        [
+            {
+                "data": np.array([0 + 0j, 1 + 2j, -1 + 1j]),
+                "name": "signal",
+                "label": "complex signal",
+                "unit": "Ohm",
+                "shape": None,
+            },
+            {
+                "data": np.array([0, 1, 2]),
+                "name": "voltage",
+                "label": "measured voltage",
+                "unit": "V",
+                "shape": None,
+            },
+        ]
+    ]
 
     data_out = _complex_to_real_preparser(data_in, conversion='real_and_imag')
 
@@ -227,10 +340,10 @@ def test__complex_to_real_preparser_complex_setpoint():
     assert data_out[0][-1] == data_in[0][-1]
 
     real_param = data_out[0][0]
-    assert real_param['name'] == 'signal_real'
-    assert real_param['label'] =='complex signal [real]'
-    assert all(real_param['data'] == np.array([0, 1, -1]))
-    assert real_param['unit'] == 'Ohm'
+    assert real_param["name"] == "signal_real"
+    assert real_param["label"] == "complex signal [real]"
+    assert all(real_param["data"] == np.array([0, 1, -1]))
+    assert real_param["unit"] == "Ohm"
 
     imag_param = data_out[0][1]
     assert imag_param['name'] == 'signal_imag'