Export to xarray dataset fails when two data variables do not share same setpoints

[thangleiter]

Take the example notebook for ParameterWithSetpoints.

MWE setup code

import tempfile
from pathlib import Path

import numpy as np
from qcodes.dataset import Measurement
from qcodes.instrument import Instrument
from qcodes.validators import Arrays, Numbers

from qcodes.dataset import initialise_or_create_database_at, load_or_create_experiment
from qcodes.parameters import Parameter, ParameterWithSetpoints


class GeneratedSetPoints(Parameter):
    """
    A parameter that generates a setpoint array from start, stop and num points
    parameters.
    """

    def __init__(self, startparam, stopparam, numpointsparam, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._startparam = startparam
        self._stopparam = stopparam
        self._numpointsparam = numpointsparam

    def get_raw(self):
        return np.linspace(
            self._startparam(), self._stopparam(), self._numpointsparam()
        )


class DummyArray(ParameterWithSetpoints):
    def get_raw(self):
        npoints = self.root_instrument.n_points.get_latest()
        return np.random.default_rng().random(npoints)


class DummySpectrumAnalyzer(Instrument):
    def __init__(self, name, **kwargs):
        super().__init__(name, **kwargs)

        self.add_parameter(
            "f_start",
            initial_value=0,
            unit="Hz",
            label="f start",
            vals=Numbers(0, 1e3),
            get_cmd=None,
            set_cmd=None,
        )

        self.add_parameter(
            "f_stop",
            unit="Hz",
            label="f stop",
            vals=Numbers(1, 1e3),
            get_cmd=None,
            set_cmd=None,
        )

        self.add_parameter(
            "n_points",
            unit="",
            initial_value=10,
            vals=Numbers(1, 1e3),
            get_cmd=None,
            set_cmd=None,
        )

        self.add_parameter(
            "freq_axis",
            unit="Hz",
            label="Freq Axis",
            parameter_class=GeneratedSetPoints,
            startparam=self.f_start,
            stopparam=self.f_stop,
            numpointsparam=self.n_points,
            vals=Arrays(shape=(self.n_points.get_latest,)),
        )

        self.add_parameter(
            "spectrum",
            unit="dBm",
            setpoints=(self.freq_axis,),
            label="Spectrum",
            parameter_class=DummyArray,
            vals=Arrays(shape=(self.n_points.get_latest,)),
        )


tutorial_db_path = Path(tempfile.mkdtemp(), "tutorial_paramter_with_setpoints.db")
initialise_or_create_database_at(tutorial_db_path)
load_or_create_experiment(
    experiment_name="tutorial_ParameterWithSetpoints", sample_name="no sample"
)

Add a parameter that depends on, say, f_stop, and sweep f_stop during the measurement:

a = DummySpectrumAnalyzer("foobar")
a.f_start(0)
a.f_stop(500)
a.n_points(501)
a.add_parameter("bar", get_cmd=lambda: a.f_stop()**2)

meas = Measurement()
meas.register_parameter(a.f_stop)
meas.register_parameter(a.spectrum, setpoints=(a.f_stop,))
meas.register_parameter(a.bar, setpoints=(a.f_stop,))

with meas.run() as datasaver:
    for f_stop in [10, 20]:
        a.f_stop(f_stop)
        datasaver.add_result(
            (a.f_stop, f_stop),
            (a.spectrum, a.spectrum.get()),
            (a.bar, a.bar.get())
        )
    dataset = datasaver.dataset

The export to xarray fails:

>>> dataset.to_xarray_dataset()
AlignmentError: cannot align objects on coordinate 'foobar_f_stop' because of conflicting indexes
first index: PandasIndex(MultiIndex([(10.0,   0.0),
...

The dataset should be easy to map to an xarray dataset though:

xr dataset init

import xarray as xr
parameter_data = dataset.get_parameter_data()
xr_dataset = xr.Dataset(
    data_vars={
        'foobar_spectrum': (
            ['foobar_f_stop', 'foobar_freq_axis'],
            parameter_data['foobar_spectrum']['foobar_spectrum']
        ),
        'foobar_bar': (
            ['foobar_f_stop'],
            parameter_data['foobar_bar']['foobar_bar']
        )
    },
    coords={
        'foobar_f_stop': (
            'foobar_f_stop',
            parameter_data['foobar_bar']['foobar_f_stop']
         ),
        'foobar_freq_axis': (
            ('foobar_f_stop', 'foobar_freq_axis'),
            parameter_data['foobar_spectrum']['foobar_freq_axis']
        ),
    }
)

Result:

>>> xr_dataset
<xarray.Dataset> Size: 16kB
Dimensions:           (foobar_f_stop: 2, foobar_freq_axis: 501)
Coordinates:
  * foobar_f_stop     (foobar_f_stop) float64 16B 10.0 20.0
    foobar_freq_axis  (foobar_f_stop, foobar_freq_axis) float64 8kB 0.0 ... 20.0
Data variables:
    foobar_spectrum   (foobar_f_stop, foobar_freq_axis) float64 8kB 0.5419 .....
    foobar_bar        (foobar_f_stop) float64 16B 251.9 256.7

I believe this used to work, but does not anymore on QCoDeS @ main. I suspect it's because get_parameter_data() broadcasts the coordinates for spectrum instead of leaving them 1d.