arbok-inspector 1.3.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (33) hide show
  1. arbok_inspector/__init__.py +1 -0
  2. arbok_inspector/analysis/analysis_base.py +29 -0
  3. arbok_inspector/analysis/prepare_data.py +118 -0
  4. arbok_inspector/classes/base_run.py +275 -0
  5. arbok_inspector/classes/dim.py +26 -0
  6. arbok_inspector/classes/native_run.py +172 -0
  7. arbok_inspector/classes/qcodes_run.py +65 -0
  8. arbok_inspector/cli.py +4 -0
  9. arbok_inspector/configurations/1d_plot.json +49 -0
  10. arbok_inspector/configurations/2d_plot.json +60 -0
  11. arbok_inspector/dev.py +19 -0
  12. arbok_inspector/helpers/string_formaters.py +37 -0
  13. arbok_inspector/helpers/unit_formater.py +29 -0
  14. arbok_inspector/main.py +15 -0
  15. arbok_inspector/pages/__init__.py +2 -0
  16. arbok_inspector/pages/database_browser.py +139 -0
  17. arbok_inspector/pages/greeter.py +93 -0
  18. arbok_inspector/pages/run_view.py +259 -0
  19. arbok_inspector/state.py +101 -0
  20. arbok_inspector/test.db +0 -0
  21. arbok_inspector/test_main.py +65 -0
  22. arbok_inspector/widgets/build_run_selecter.py +163 -0
  23. arbok_inspector/widgets/build_run_view_actions.py +104 -0
  24. arbok_inspector/widgets/build_xarray_grid.py +145 -0
  25. arbok_inspector/widgets/build_xarray_html.py +57 -0
  26. arbok_inspector/widgets/json_plot_settings_dialog.py +77 -0
  27. arbok_inspector/widgets/update_day_selecter.py +64 -0
  28. arbok_inspector-1.3.0.dist-info/METADATA +90 -0
  29. arbok_inspector-1.3.0.dist-info/RECORD +33 -0
  30. arbok_inspector-1.3.0.dist-info/WHEEL +5 -0
  31. arbok_inspector-1.3.0.dist-info/entry_points.txt +2 -0
  32. arbok_inspector-1.3.0.dist-info/licenses/LICENSE +21 -0
  33. arbok_inspector-1.3.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1 @@
1
+ # from .database_browser import database_browser_page
@@ -0,0 +1,29 @@
1
+ """Module containing AnalysisBase class"""
2
+
3
+ class AnalysisBase:
4
+ """Base class for analysis classes"""
5
+ run_id = None
6
+ xr_data = None
7
+
8
+ def find_axis_from_keyword(self, keyword: str) -> str:
9
+ """
10
+ Find the axis corresponding to a keyword in the analysis
11
+ Args:
12
+ keyword (str): Keyword to search for
13
+ Returns:
14
+ axis (int): Axis corresponding to keyword
15
+ """
16
+ axes = []
17
+ for axis in self.xr_data.dims:
18
+ if keyword in axis:
19
+ axes.append(axis)
20
+ if len(axes) == 0:
21
+ raise ValueError(
22
+ f"Axis not found for keyword {keyword}. "
23
+ f"Dims are {self.xr_data.dims}"
24
+ )
25
+ elif len(axes) > 1:
26
+ raise ValueError(
27
+ f"More than one axis found for keyword {keyword}: {axes}")
28
+ else:
29
+ return axes[0]
@@ -0,0 +1,118 @@
1
+ """Module containing prepare_data function for analysis tools"""
2
+
3
+ from matplotlib.pylab import f
4
+ from qcodes.dataset.data_set import load_by_id, DataSet
5
+ import xarray as xr
6
+ import numpy as np
7
+ import matplotlib.pyplot as plt
8
+
9
+ def prepare_and_avg_data(
10
+ run: int | DataSet | xr.Dataset | xr.DataArray,
11
+ readout_name: str,
12
+ avg_axes: str | list = 'auto'
13
+ ) -> tuple[int | None, xr.DataArray, np.ndarray]:
14
+ """
15
+ Prepares the data for plotting. Takes either a run id, a qcodes dataset,
16
+ an xarray dataset or an xarray data-array and returns the run id, the xarray
17
+ data-array and the numpy data-array.
18
+ This is done to allow different input types for the data while keeping the
19
+ same output format.
20
+
21
+ Args:
22
+ run (int | DataSet | xr.Dataset | xr.DataArray): Run id, qcodes dataset'
23
+ xarray dataset or xarray data-array
24
+ readout_name (str): Name of the readout observable
25
+ """
26
+ xdata_array = None
27
+ if avg_axes is None:
28
+ avg_axes = []
29
+ if isinstance(run, int):
30
+ data = load_by_id(run)
31
+ xdataset = data.to_xarray_dataset()
32
+ run_id = run
33
+ elif isinstance(run, DataSet):
34
+ data = run
35
+ run_id = data.run_id
36
+ xdataset = data.to_xarray_dataset()
37
+ elif isinstance(run, xr.Dataset):
38
+ xdataset = run
39
+ run_id = xdataset.attrs['run_id']
40
+ elif isinstance(run, xr.DataArray):
41
+ xdataset = None
42
+ xdata_array = run
43
+ run_id = None
44
+ else:
45
+ raise ValueError(
46
+ "Invalid input type for run. "
47
+ "Must be run-ID, DataSet or xr.Dataset or xr.DataArray. "
48
+ f"Is {type(run)}"
49
+ )
50
+ if xdataset is not None:
51
+ if readout_name not in xdataset.data_vars:
52
+ readout_name = find_data_variable_from_keyword(xdataset, readout_name)
53
+ xdata_array = xdataset[readout_name]
54
+ ### Average over specified axes
55
+ xdata_array = avg_dataarray(xdata_array, avg_axes)
56
+ np_data = xdata_array.to_numpy()
57
+ return run_id, xdata_array, np_data
58
+
59
+ def find_data_variable_from_keyword(
60
+ xdata_array: xr.DataArray, keyword: str | tuple) -> str:
61
+ """
62
+ Find the data variable corresponding to a keyword in the data-array.
63
+
64
+ Args:
65
+ xdata_array (xr.DataArray): xarray data-array to search in
66
+ keyword (str): Keyword to search for
67
+ Returns:
68
+ data_variable (str): Data variable corresponding to keyword
69
+ """
70
+ if isinstance(keyword, str):
71
+ keyword = (keyword,)
72
+ if not isinstance(keyword, tuple):
73
+ raise ValueError(
74
+ f"Keyword must be a string or a tuple. Is {type(keyword)}")
75
+ data_variables = []
76
+ for data_variable in xdata_array.data_vars:
77
+ if all([subkey in str(data_variable) for subkey in keyword]):
78
+ data_variables.append(data_variable)
79
+ if len(data_variables) == 0:
80
+ raise ValueError(
81
+ f"Data variable not found for keyword {keyword}. "
82
+ f"Data variables are {xdata_array.data_vars}"
83
+ )
84
+ elif len(data_variables) > 1:
85
+ raise ValueError(
86
+ f"More than one data variable found for keyword {keyword}: "
87
+ f"{[str(var) for var in data_variables]}")
88
+ else:
89
+ return data_variables[0]
90
+
91
+ def avg_dataarray(xdata_array: xr.DataArray, avg_axes: str | list = 'auto'):
92
+ """
93
+ Averages the data-array over the specified axes. If no axes are specified
94
+ the data-array is averaged over all axes.
95
+
96
+ Args:
97
+ xdata_array (xr.DataArray): xarray data-array to be averaged
98
+ avg_axes (str | list): Axes to average over
99
+ """
100
+ if avg_axes is None:
101
+ avg_axes = []
102
+ if isinstance(avg_axes, str):
103
+ ### If 'auto' is given, find all axes with 'iteration' in the name
104
+ if avg_axes == 'auto':
105
+ avg_axes = []
106
+ for dim in xdata_array.dims:
107
+ if 'iteration' in dim:
108
+ avg_axes.append(dim)
109
+ else:
110
+ avg_axes = [avg_axes]
111
+ ### Average over specified axes
112
+ for axis in avg_axes:
113
+ if hasattr(xdata_array, axis):
114
+ xdata_array = xdata_array.mean(axis)
115
+ else:
116
+ raise KeyError(
117
+ f"Avg. axis {axis} not found in xarray data-array")
118
+ return xdata_array
@@ -0,0 +1,275 @@
1
+ """
2
+ Run class representing a single run of the experiment.
3
+ """
4
+ from __future__ import annotations
5
+ from typing import TYPE_CHECKING
6
+
7
+ from abc import ABC, abstractmethod
8
+ import ast
9
+ import re
10
+ import io
11
+ import json
12
+ from qcodes.dataset import load_by_id
13
+ from nicegui import ui, app
14
+ import xarray as xr
15
+
16
+
17
+
18
+ from arbok_inspector.classes.dim import Dim
19
+ from arbok_inspector.widgets.build_xarray_grid import build_xarray_grid
20
+ from arbok_inspector.state import ArbokInspector, inspector
21
+
22
+ if TYPE_CHECKING:
23
+ from qcodes.dataset.data_set import DataSet
24
+ from xarray import Dataset
25
+
26
+ AXIS_OPTIONS = ['average', 'select_value', 'y-axis', 'x-axis']
27
+
28
+ class BaseRun(ABC):
29
+ """
30
+ Class representing a run with its data and methods
31
+ """
32
+ def __init__(self, run_id: int):
33
+ """
34
+ Constructor for Run class
35
+
36
+ Args:
37
+ run_id (int): ID of the run
38
+ """
39
+ self.run_id: int = run_id
40
+ self.title: str = f'Run ID: {run_id} (-> add experiment)'
41
+ self.inspector: ArbokInspector = inspector
42
+ self._database_columns = self._get_database_columns()
43
+ self.full_data_set: Dataset = self._load_dataset()
44
+ self.last_subset: Dataset = self.full_data_set
45
+
46
+ self.parallel_sweep_axes: dict = {}
47
+ self.sweep_dict: dict[int, Dim] = {}
48
+ self.load_sweep_dict()
49
+ self.dims: list[Dim] = list(self.sweep_dict.values())
50
+ self.dim_axis_option: dict[str, str|list[Dim]] = self.set_dim_axis_option()
51
+ print(self.dims)
52
+
53
+ self.plot_selection: list[str] = self.select_results_by_keywords(
54
+ app.storage.general["result_keywords"]
55
+ )
56
+ print(f"Initial plot selection: {self.plot_selection}")
57
+ self.plots_per_column: int = 2
58
+
59
+ @property
60
+ def database_columns(self) -> dict[str, dict[str, str]]:
61
+ """Column names of database, with their values and shown labels"""
62
+ return self._database_columns
63
+
64
+ @abstractmethod
65
+ def _get_database_columns(self) -> dict[str, dict[str, str]]:
66
+ pass
67
+
68
+ @abstractmethod
69
+ def _load_dataset(self) -> Dataset:
70
+ """
71
+ Load the dataset for the given run ID from the appropriate database type.
72
+
73
+ Args:
74
+ run_id (int): ID of the run
75
+ database_type (str): Type of the database ('qcodes' or 'arbok')
76
+ Returns:
77
+ DataSet: Loaded dataset
78
+ """
79
+ pass
80
+
81
+ @abstractmethod
82
+ def get_qua_code(self, as_string: bool = False) -> str:
83
+ """
84
+ Retrieve the QUA code associated with this run.
85
+
86
+ Returns:
87
+ qua_code (str): The QUA code as a string
88
+ """
89
+ pass
90
+
91
+ def load_sweep_dict(self):
92
+ """
93
+ Load the sweep dictionary from the dataset
94
+ TODO: check metadata for sweep information!
95
+ Returns:
96
+ sweep_dict (dict): Dictionary with sweep information
97
+ is_together (bool): True if all sweeps are together, False otherwise
98
+ """
99
+ self.parallel_sweep_axes = {}
100
+ dims = self.full_data_set.dims
101
+ for i, dim in enumerate(dims):
102
+ dependent_coords = [
103
+ name for name, coord in self.full_data_set.coords.items() if dim in coord.dims]
104
+ self.parallel_sweep_axes[i] = dependent_coords
105
+ self.sweep_dict = {
106
+ i: Dim(names[0]) for i, names in self.parallel_sweep_axes.items()
107
+ }
108
+ return self.sweep_dict
109
+
110
+ def set_dim_axis_option(self):
111
+ """
112
+ Set the default dimension options for the run in 4 steps:
113
+ 1. Set all iteration dims to 'average'
114
+ 2. Set the innermost dim to 'x-axis' (the last one that is not averaged)
115
+ 3. Set the next innermost dim to 'y-axis'
116
+ 4. Set all remaining dims to 'select_value'
117
+
118
+ Returns:
119
+ options (dict): Dictionary with keys 'average', 'select_value', 'y-axis',
120
+ """
121
+ options = {x: [] for x in AXIS_OPTIONS}
122
+ print(f"Setting average to {app.storage.general['avg_axis']}")
123
+ for dim in self.dims:
124
+ if app.storage.general["avg_axis"] is None:
125
+ break
126
+ if app.storage.general["avg_axis"] in dim.name:
127
+ dim.option = 'average'
128
+ options['average'].append(dim)
129
+ for dim in reversed(self.dims):
130
+ if dim not in options['average'] and dim != options['x-axis']:
131
+ dim.option = "x-axis"
132
+ options['x-axis'] = dim
133
+ print(f"Setting x-axis to {dim.name}")
134
+ break
135
+ for dim in reversed(self.dims):
136
+ if dim not in options['average'] and dim != options['x-axis']:
137
+ dim.option = 'y-axis'
138
+ options['y-axis'] = dim
139
+ print(f"Setting y-axis to {dim.name}")
140
+ break
141
+ for dim in self.dims:
142
+ if dim not in options['average'] and dim != options['x-axis'] and dim != options['y-axis']:
143
+ dim.option = 'select_value'
144
+ options['select_value'].append(dim)
145
+ dim.select_index = 0
146
+ print(f"Setting select_value to {dim.name}")
147
+ return options
148
+
149
+ def select_results_by_keywords(self, keywords: list[str|tuple]) -> list[str]:
150
+ """
151
+ Select results by keywords in their name.
152
+ Args:
153
+ keywords (list): List of keywords to search for
154
+ Returns:
155
+ selected_results (list): List of selected result names
156
+
157
+ TODO: simplify this! way too complicated
158
+ """
159
+ print(f"using keywords: {keywords}")
160
+ if keywords is None or len(keywords) == 0 or keywords == '':
161
+ return [next(iter(self.full_data_set.data_vars))]
162
+ s_quoted = re.sub(r'\b([a-zA-Z_][a-zA-Z0-9_]*)\b', r'"\1"', keywords)
163
+ try:
164
+ keywords = ast.literal_eval(s_quoted)
165
+ except (SyntaxError, ValueError):
166
+ print(f"Error parsing keywords: {s_quoted}")
167
+ keywords = []
168
+ ui.notify(
169
+ f"Error parsing result keywords: {s_quoted}. Please use a valid Python list.",
170
+ color='red',
171
+ position='top-right'
172
+ )
173
+ if not isinstance(keywords, list):
174
+ keywords = [keywords]
175
+ selected_results = []
176
+ print(f"using keywords: {keywords}")
177
+ for result in self.full_data_set.data_vars:
178
+ for keyword in keywords:
179
+ if isinstance(keyword, str) and keyword in str(result):
180
+ selected_results.append(result)
181
+ elif isinstance(keyword, tuple) and all(
182
+ subkey in str(result) for subkey in keyword):
183
+ selected_results.append(result)
184
+ selected_results = list(set(selected_results)) # Remove duplicates
185
+ if len(selected_results) == 0:
186
+ selected_results = [next(iter(self.full_data_set.data_vars))]
187
+ print(f"Selected results: {selected_results}")
188
+ return selected_results
189
+
190
+ def update_subset_dims(self, dim: Dim, selection: str, index = None):
191
+ """
192
+ Update the subset dimensions based on user selection.
193
+
194
+ Args:
195
+ dim (Dim): The dimension object to update
196
+ selection (str): The new selection option
197
+ ('average', 'select_value', 'x-axis', 'y-axis')
198
+ index (int, optional): The index for 'select_value' option. Defaults to None.
199
+ """
200
+ text = f'Updating subset dims: {dim.name} to {selection}'
201
+ print(text)
202
+ ui.notify(text, position='top-right')
203
+
204
+ ### First, remove old option this dim was on
205
+ for option in ['average', 'select_value']:
206
+ if dim in self.dim_axis_option[option]:
207
+ print(f"Removing {dim.name} from {option}")
208
+ self.dim_axis_option[option].remove(dim)
209
+ dim.option = None
210
+ if dim.option in ['x-axis', 'y-axis']:
211
+ print(f"Removing {dim.name} from {dim.option}")
212
+ self.dim_axis_option[dim.option] = None
213
+
214
+ ### Now, set new option
215
+ if selection in ['average', 'select_value']:
216
+ # dim.ui_selector.value = selection
217
+ dim.select_index = index
218
+ self.dim_axis_option[selection].append(dim)
219
+ return
220
+ if selection in ['x-axis', 'y-axis']:
221
+ old_dim = self.dim_axis_option[selection]
222
+ self.dim_axis_option[selection] = dim
223
+ if old_dim:
224
+ # Set previous dim (having this option) to 'select_value'
225
+ # Required since x and y axis have to be unique
226
+ print(old_dim)
227
+ print(f"Updating {old_dim.name} to {dim.name} on {selection}")
228
+ if old_dim.option in ['x-axis', 'y-axis']:
229
+ self.dim_axis_option['select_value'].append(old_dim)
230
+ old_dim.option = 'select_value'
231
+ old_dim.ui_selector.value = 'select_value'
232
+ self.update_subset_dims(old_dim, 'select_value', old_dim.select_index)
233
+ dim.ui_selector.update()
234
+
235
+ def generate_subset(self):
236
+ """
237
+ Generate the subset of the full dataset based on the current dimension options.
238
+ Returns:
239
+ sub_set (xarray.Dataset): The subset of the full dataset
240
+ """
241
+ # TODO: take the averaging out of this! We only want to average if necessary
242
+ # averaging can be computationally intensive!
243
+ sub_set = self.full_data_set
244
+ for avg_axis in self.dim_axis_option['average']:
245
+ sub_set = sub_set.mean(dim=avg_axis.name)
246
+ sel_dict = {d.name: d.select_index for d in self.dim_axis_option['select_value']}
247
+ sub_set = sub_set.isel(**sel_dict).squeeze()
248
+ self.last_subset = sub_set
249
+ return sub_set
250
+
251
+ def update_plot_selection(self, value: bool, readout_name: str):
252
+ """
253
+ Update the plot selection based on user interaction.
254
+
255
+ Args:
256
+ value (bool): True if the result is selected, False otherwise
257
+ readout_name (str): Name of the result to update
258
+ """
259
+ print(f"{readout_name= } {value= }")
260
+ pretty_readout_name = readout_name.replace("__", ".")
261
+ if readout_name not in self.plot_selection:
262
+ self.plot_selection.append(readout_name)
263
+ ui.notify(
264
+ message=f'Result {pretty_readout_name} added to plot selection',
265
+ position='top-right'
266
+ )
267
+ else:
268
+ self.plot_selection.remove(readout_name)
269
+ ui.notify(
270
+ f'Result {pretty_readout_name} removed from plot selection',
271
+ position='top-right'
272
+ )
273
+ print(f"{self.plot_selection= }")
274
+ build_xarray_grid()
275
+
@@ -0,0 +1,26 @@
1
+ """Module for the Dim class."""
2
+
3
+ class Dim:
4
+ """
5
+ Class representing a dimension of the data
6
+ """
7
+ def __init__(self, name):
8
+ """
9
+ Constructor for Dim class
10
+
11
+ Args:
12
+ name (str): Name of the dimension
13
+
14
+ Attributes:
15
+ name (str): Name of the dimension
16
+ option (str): Option for the dimension (average, select_value, x-axis, y-axis)
17
+ select_index (int): Index of the selected value for select_value option
18
+ ui_selector: Reference to the UI element for the dimension
19
+ """
20
+ self.name = name
21
+ self.option = None
22
+ self.select_index = 0
23
+ self.ui_selector = None
24
+
25
+ def __str__(self):
26
+ return self.name
@@ -0,0 +1,172 @@
1
+ """Module containing NativeRun class"""
2
+ from __future__ import annotations
3
+ from typing import TYPE_CHECKING
4
+
5
+ from sqlalchemy import select
6
+ from sqlalchemy.orm import DeclarativeBase, mapped_column, Session
7
+ from sqlalchemy.inspection import inspect
8
+ import xarray as xr
9
+
10
+ from arbok_inspector.state import inspector
11
+ from arbok_inspector.classes.base_run import BaseRun
12
+
13
+ if TYPE_CHECKING:
14
+ from qcodes.dataset.data_set import DataSet
15
+ from xarray import Dataset
16
+
17
+ COLUMN_LABELS = {}
18
+
19
+ class NativeRun(BaseRun):
20
+ """"""
21
+ def __init__(
22
+ self,
23
+ run_id: int
24
+ ):
25
+ """
26
+ Constructor for NativeRun class
27
+
28
+ Args:
29
+ run_id (int): Run ID of the measurement run
30
+ """
31
+ super().__init__(run_id)
32
+ print("HERE", self.database_columns)
33
+
34
+ def _get_database_columns(self) -> dict[str, dict[str, str]]:
35
+ """
36
+ Returns column names of database, with this rows values and shown labels
37
+ """
38
+ columns_and_values = {}
39
+ with Session(inspector.database_engine) as session:
40
+ self.sql_run = session.get(SqlRun, self.run_id)
41
+ columns_and_values['experiment'] = {
42
+ 'value': self.sql_run.experiment.name}
43
+ for attr in inspect(self.sql_run).mapper.column_attrs:
44
+ value = getattr(self.sql_run, attr.key)
45
+ columns_and_values[attr.key] = {'value': value}
46
+ if attr.key in COLUMN_LABELS:
47
+ label = COLUMN_LABELS[attr.key]
48
+ columns_and_values[attr.key]['label'] = label
49
+
50
+ session.expunge(self.sql_run)
51
+ print(columns_and_values)
52
+ return columns_and_values
53
+
54
+ def _load_dataset(self) -> Dataset:
55
+ """
56
+ Loads the dataset from the minio bucket. Results with empty dimensions
57
+ are dropped
58
+ """
59
+ minio_path = inspector.minio_bucket + "/"
60
+ minio_path += f"{self.sql_run.run_id}_{self.sql_run.uuid}"
61
+ minio_path += "/data.zarr"
62
+ print(f"Loading dataset from MinIO path: {minio_path}")
63
+ store = inspector.minio_filesystem.get_mapper(minio_path)
64
+ dataset = xr.open_zarr(store, consolidated=True)
65
+ return dataset
66
+
67
+ def get_qua_code(self, as_string: bool = False) -> str:
68
+ """
69
+ Retrieve the QUA code associated with this run.
70
+
71
+ Returns:
72
+ qua_code (str): The QUA code as a string
73
+ """
74
+ if inspector.database_type == 'native_arbok':
75
+ with Session(inspector.database_engine) as session:
76
+ sql_run = session.get(SqlRun, self.run_id)
77
+
78
+ minio_path = inspector.minio_bucket + "/"
79
+ minio_path += f"{self.sql_run.run_id}_{self.sql_run.uuid}"
80
+ minio_path += "/metadata/qua_program.py"
81
+ print(f"Loading QUA code from MinIO path: {minio_path}")
82
+ if as_string:
83
+ with inspector.minio_filesystem.open(minio_path, mode="r") as f:
84
+ return f.read()
85
+ else:
86
+ with inspector.minio_filesystem.open(minio_path, mode="rb") as f:
87
+ return io.BytesIO(f.read())
88
+
89
+ ### TODO: IMPORT THOSE CLASSES FROM ELSEWHERE!
90
+ from sqlalchemy.orm import (
91
+ DeclarativeBase, Mapped, mapped_column, relationship
92
+ )
93
+
94
+ import uuid
95
+ from sqlalchemy import (
96
+ BigInteger,
97
+ Boolean,
98
+ ForeignKey,
99
+ Integer,
100
+ JSON,
101
+ String,
102
+ Text,
103
+ UniqueConstraint,
104
+ )
105
+ from sqlalchemy.dialects.postgresql import ARRAY, JSONB, UUID
106
+
107
+ class Base(DeclarativeBase):
108
+ pass
109
+
110
+ class SqlRun(Base):
111
+ __tablename__ = "runs"
112
+
113
+ ### Run metadata
114
+ run_id: Mapped[int] = mapped_column(
115
+ Integer, primary_key=True, autoincrement=True)
116
+ uuid: Mapped[str] = mapped_column(
117
+ UUID(as_uuid=True), default=uuid.uuid4, unique=True, nullable=False)
118
+ exp_id: Mapped[int] = mapped_column(ForeignKey("experiments.exp_id"))
119
+ experiment = relationship("SqlExperiment", back_populates="runs")
120
+ device_id = mapped_column(ForeignKey("devices.device_id"))
121
+ device = relationship("SqlDevice", back_populates="runs")
122
+ name: Mapped[str] = mapped_column(String)
123
+
124
+ setup: Mapped[str] = mapped_column(String)
125
+
126
+ ### Data specific to this run
127
+ result_count: Mapped[int] = mapped_column(Integer, default=0)
128
+ batch_count: Mapped[int] = mapped_column(Integer, default=0)
129
+ coords: Mapped[list[str]] = mapped_column(ARRAY(String))
130
+ sweeps: Mapped[dict[int, list[str]]] = mapped_column(JSONB)
131
+
132
+ ### Timestamps and status
133
+ start_time: Mapped[int] = mapped_column(BigInteger)
134
+ completed_time: Mapped[int] = mapped_column(
135
+ BigInteger, default=None, nullable=True)
136
+ is_completed: Mapped[bool] = mapped_column(Boolean, default=False)
137
+
138
+ measurement_exception: Mapped[str] = mapped_column(
139
+ Text, default=None, nullable=True)
140
+ parent_datasets: Mapped[list[int]] = mapped_column(
141
+ ARRAY(Integer), default=None, nullable=True)
142
+
143
+ class SqlExperiment(Base):
144
+ __tablename__ = "experiments"
145
+
146
+ exp_id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
147
+ name: Mapped[str] = mapped_column(String, unique=True, nullable=False)
148
+ creation_time: Mapped[int] = mapped_column(BigInteger)
149
+ #run_counter: Mapped[int] = mapped_column(Integer)
150
+ #format_string: Mapped[str] = mapped_column(String)
151
+
152
+ runs: Mapped[list[SqlRun]] = relationship(
153
+ "SqlRun", back_populates="experiment", cascade="all, delete-orphan"
154
+ )
155
+
156
+ __table_args__ = (
157
+ UniqueConstraint("name", name="uq_experiment_name"),
158
+ )
159
+
160
+ def __repr__(self) -> str:
161
+ return (f"<Experiment(exp_id={self.exp_id}, name='{self.name}', "
162
+ f"creation_time={self.creation_time}> ")
163
+
164
+ class SqlDevice(Base):
165
+ __tablename__ = "devices"
166
+
167
+ device_id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
168
+ name: Mapped[str] = mapped_column(String, unique=True, nullable=False)
169
+ runs: Mapped[list[SqlRun]] = relationship("SqlRun", back_populates="device")
170
+
171
+ def __repr__(self) -> str:
172
+ return f"<Device(device_id={self.device_id}, name='{self.name}')>"
@@ -0,0 +1,65 @@
1
+ """Module containing QcodesRun class"""
2
+ from __future__ import annotations
3
+ from typing import TYPE_CHECKING
4
+
5
+ import os
6
+ from pathlib import Path
7
+
8
+ from qcodes.dataset import load_by_id
9
+ from qcodes.dataset.sqlite.database import get_DB_location
10
+ from arbok_inspector.classes.base_run import BaseRun
11
+
12
+ if TYPE_CHECKING:
13
+ from xarray import Dataset
14
+
15
+ COLUMN_LABELS = {}
16
+
17
+ class QcodesRun(BaseRun):
18
+ """"""
19
+ def __init__(
20
+ self,
21
+ run_id: int
22
+ ):
23
+ """
24
+ Constructor for QcodesRun class
25
+
26
+ Args:
27
+ run_id (int): Run ID of the measurement run
28
+ """
29
+ super().__init__(run_id)
30
+
31
+ def _load_dataset(self) -> Dataset:
32
+ dataset = load_by_id(self.run_id)
33
+ dataset = dataset.to_xarray_dataset()
34
+ return dataset
35
+
36
+ def _get_database_columns(self) -> dict[str, dict[str, str]]:
37
+ self.inspector.cursor.execute(
38
+ "SELECT * FROM runs WHERE run_id = ?", (self.run_id,))
39
+ row = self.inspector.cursor.fetchone()
40
+ if row is not None:
41
+ row_dict = dict(row)
42
+ else:
43
+ raise ValueError(f'database entry not found for run-ID: {self.run_id}')
44
+ columns_and_values = {}
45
+ for key, value in row_dict.items():
46
+ columns_and_values[key] = {'value': value}
47
+ if key in COLUMN_LABELS:
48
+ label = COLUMN_LABELS[key]
49
+ columns_and_values[key]['label'] = label
50
+ return columns_and_values
51
+
52
+ def get_qua_code(self, as_string: bool = False) -> str | bytes:
53
+ db_path = os.path.abspath(get_DB_location())
54
+ db_dir = os.path.dirname(db_path)
55
+ programs_dir = Path(db_dir) / "qua_programs/"
56
+ raise NotImplementedError
57
+ ### TODO: IMPLEMENT MORE EASILY IN ARBOK THOUGH!
58
+ # if not os.path.isdir(programs_dir):
59
+ # os.makedirs(programs_dir)
60
+ # try:
61
+ # with open(save_path, 'r', encoding="utf-8") as file:
62
+ # file.write(
63
+ # generate_qua_script(qua_program, opx_config))
64
+ # except FileNotFoundError as e:
65
+ # ui.notify(f"Qua program couldnt be found next to database: {e}")