Feature/control parameters plotted in plot dataset

docs/changes/newsfragments/6931.improved

@@ -0,0 +1,4 @@
+Added 'parameters' argument to the 'plot_dataset' (and 'plot_by_id')
+functions. This enables a choice of which dependent parameters to plot
+(in the case that the plotting of all dependent parameters is not
+desired).

src/qcodes/dataset/plotting.py

@@ -101,10 +101,11 @@
     cutoff_percentile: tuple[float, float] | float | None = None,
     complex_plot_type: Literal["real_and_imag", "mag_and_phase"] = "real_and_imag",
     complex_plot_phase: Literal["radians", "degrees"] = "radians",
+    parameters: Sequence[str] | str | None = None,
     **kwargs: Any,
 ) -> AxesTupleList:
     """
-    Construct all plots for a given dataset
+    Construct plots for a given dataset
 
     Implemented so far:
 
@@ -150,6 +151,8 @@
         complex_plot_phase: Format of phase for plotting complex-valued data,
             either ``"radians"`` or ``"degrees"``. Applicable only for the
             cases where the dataset contains complex numbers
+        parameters: Names of dependent parameters to plot if the plotting of
+            all is not desired
         **kwargs: Keyword arguments passed to the plotting function.
 
     Returns:
@@ -206,7 +209,10 @@
         alldata, conversion=complex_plot_type, degrees=degrees
     )
 
-    nplots = len(alldata)
+    if isinstance(parameters, str):
+        parameters = [parameters]
+
+    nplots = len(alldata) if parameters is None else len(parameters)
 
     if isinstance(axes, matplotlib.axes.Axes):
         axeslist = [axes]
@@ -229,15 +235,46 @@
                 f"Provided arguments: {subplots_kwargs}"
             )
         if len(axeslist) != nplots:
-            raise RuntimeError(
-                f"Trying to make {nplots} plots, but"
-                f"received {len(axeslist)} axes objects."
-            )
+            if parameters is not None:
+                raise RuntimeError(
+                    f"Trying to plot {len(parameters)} parameters "
+                    f"but received {len(axeslist)} axes objects. "
+                )
+            else:
+                raise RuntimeError(
+                    f"Trying to make {nplots} plots, but"
+                    f"received {len(axeslist)} axes objects."
+                )
 
     if colorbars is None:
         colorbars = len(axeslist) * [None]
     new_colorbars: list[Colorbar | None] = []
 
+    if parameters is None:
+        log.debug("No data specified - plotting all.")
+
+    else:
+        # validate parameters
+        dependents = []
+        for i in range(len(dataset.dependent_parameters)):
+            dependents.append(dataset.dependent_parameters[i].name)
+        for param in parameters:
+            if param not in dependents:
+                raise ValueError(
+                    f"Invalid parameter(s) given. Received {parameters} "
+                    f"but can only accept elements from {dependents}. "
+                )
+
+        log.debug(f"Plotting data for {parameters}.")
+
+        for i, data in enumerate(alldata):
+            if len(data) == 2:  # 1D PLOTTING
+                if data[1]["name"] not in parameters:
+                    alldata.pop(i)
+            elif len(data) == 3:  # 2D PLOTTING
+                if data[2]["name"] not in parameters:
+                    alldata.pop(i)
+
     for data, ax, colorbar in zip(alldata, axeslist, colorbars):
         if len(data) == 2:  # 1D PLOTTING
             log.debug(f"Doing a 1D plot with kwargs: {kwargs}")
@@ -383,6 +420,7 @@
     cutoff_percentile: tuple[float, float] | float | None = None,
     complex_plot_type: Literal["real_and_imag", "mag_and_phase"] = "real_and_imag",
     complex_plot_phase: Literal["radians", "degrees"] = "radians",
+    parameters: Sequence[str] | str | None = None,
     **kwargs: Any,
 ) -> AxesTupleList:
     """
@@ -403,6 +441,7 @@
         cutoff_percentile,
         complex_plot_type,
         complex_plot_phase,
+        parameters,
         **kwargs,
     )
 

tests/dataset/test_plotting.py

@@ -5,6 +5,8 @@
 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
+from matplotlib.lines import Line2D
+from numpy.testing import assert_allclose
 from pytest import FixtureRequest
 
 import qcodes as qc
@@ -18,6 +20,7 @@
     plot_dataset,
 )
 from qcodes.instrument_drivers.mock_instruments import DummyInstrument
+from qcodes.parameters import Parameter
 from qcodes.plotting.axis_labels import _ENGINEERING_PREFIXES, _UNITS_FOR_RESCALING
 
 if TYPE_CHECKING:
@@ -367,3 +370,60 @@ def test__complex_to_real_preparser_complex_setpoint() -> None:
     assert measured_param["label"] == "measured voltage"
     assert measured_param["unit"] == "V"
     assert all(measured_param["data"] == np.array([0, 1, 2]))
+
+
+@pytest.mark.parametrize("params", [None, "y", "y2", ["y", "y2"]])
+def test_plot_dataset_parameters(experiment, request: FixtureRequest, params) -> None:
+    """
+    Test plotting of specified parameters with plot_dataset
+    function.
+    """
+    # Generate some data
+    x = Parameter(name="x", label="Voltage", unit="V", set_cmd=None, get_cmd=None)
+    y = Parameter(name="y", label="Voltage", unit="V", set_cmd=None, get_cmd=None)
+    y2 = Parameter(name="y2", label="Current", unit="A", set_cmd=None, get_cmd=None)
+
+    meas = Measurement()
+    meas.register_parameter(x)
+    meas.register_parameter(y, setpoints=[x])
+    meas.register_parameter(y2, setpoints=[x])
+
+    xvals = np.linspace(-5, 5, 250)
+
+    with meas.run() as datasaver:
+        for xnum in xvals:
+            datasaver.add_result((x, xnum), (y, xnum**2))
+            datasaver.add_result((x, xnum), (y2, -(xnum**2)))
+
+    dataset = datasaver.dataset
+
+    axes, _ = plot_dataset(dataset=dataset, parameters=params)
+
+    # None is default for parameters
+    # so if parameters is None all dependent parameters should be plotted
+    params = ["y", "y2"] if params is None else params
+
+    if params == ["y", "y2"]:
+        assert len(axes) == 2
+
+        for i, param in enumerate(["y", "y2"]):
+            data = dataset.get_parameter_data()[param][param]
+            artist = axes[i].get_children()[0]
+            assert isinstance(artist, Line2D)
+            plotted = artist.get_ydata()
+            # actual data and plotted not exactly equal (?)
+            # so check for very small diff
+            assert_allclose(np.array(plotted), data, rtol=1e-10)
+
+    # check only 'requested' parameter has been plotted
+    elif params == "y" or "y2":
+        assert isinstance(params, str)
+        assert len(axes) == 1
+        dsdata = dataset.get_parameter_data()
+        data = dsdata[params][params]
+        artist = axes[0].get_children()[0]
+        assert isinstance(artist, Line2D)
+        plotted = artist.get_ydata()
+        # actual data and plotted not exactly equal (?)
+        # so check for very small diff
+        assert_allclose(np.array(plotted), data, rtol=1e-10)