microsoft / microsoft/Qcodes

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

Open
#8,232 0 comments 1 reaction 2 assignees Claimed by @jenshnielsen View on GitHub
Dominant language
Python
Stars
459
Forks
359
Avg merge
3d 6h
Merged PRs (30d)
73

Description

Take the [example notebook for `ParameterWithSetpoints`](https://microsoft.github.io/Qcodes/examples/Parameters/Simple-Example-of-ParameterWithSetpoints.html).

MWE setup code

```py
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:

```py
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:

```py
>>> 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

```py
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']
),
}
)
```

>> xr_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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.