arbok-inspector 0.1.0__py3-none-any.whl

arbok_inspector/__init__.py

@@ -0,0 +1 @@
+# from .database_browser import database_browser_page

arbok_inspector/analysis/analysis_base.py

@@ -0,0 +1,29 @@
+"""Module containing AnalysisBase class"""
+
+class AnalysisBase:
+    """Base class for analysis classes"""
+    run_id = None
+    xr_data = None
+
+    def find_axis_from_keyword(self, keyword: str) -> str:
+        """
+        Find the axis corresponding to a keyword in the analysis
+        Args:
+            keyword (str): Keyword to search for
+        Returns:
+            axis (int): Axis corresponding to keyword
+        """
+        axes = []
+        for axis in self.xr_data.dims:
+            if keyword in axis:
+                axes.append(axis)
+        if len(axes) == 0:
+            raise ValueError(
+                f"Axis not found for keyword {keyword}. "
+                f"Dims are {self.xr_data.dims}"
+                )
+        elif len(axes) > 1:
+            raise ValueError(
+                f"More than one axis found for keyword {keyword}: {axes}")
+        else:
+            return axes[0]

arbok_inspector/analysis/prepare_data.py

@@ -0,0 +1,118 @@
+"""Module containing prepare_data function for analysis tools"""
+
+from matplotlib.pylab import f
+from qcodes.dataset.data_set import load_by_id, DataSet
+import xarray as xr
+import numpy as np
+import matplotlib.pyplot as plt
+
+def prepare_and_avg_data(
+    run: int | DataSet | xr.Dataset | xr.DataArray,
+    readout_name: str,
+    avg_axes: str | list = 'auto'
+    ) -> tuple[int | None, xr.DataArray, np.ndarray]:
+    """
+    Prepares the data for plotting. Takes either a run id, a qcodes dataset,
+    an xarray dataset or an xarray data-array and returns the run id, the xarray
+    data-array and the numpy data-array.
+    This is done to allow different input types for the data while keeping the
+    same output format.
+    
+    Args:
+        run (int | DataSet | xr.Dataset | xr.DataArray): Run id, qcodes dataset'
+            xarray dataset or xarray data-array
+        readout_name (str): Name of the readout observable
+    """
+    xdata_array = None
+    if avg_axes is None:
+        avg_axes = []
+    if isinstance(run, int):
+        data = load_by_id(run)
+        xdataset = data.to_xarray_dataset()
+        run_id = run
+    elif isinstance(run, DataSet):
+        data = run
+        run_id = data.run_id
+        xdataset = data.to_xarray_dataset()
+    elif isinstance(run, xr.Dataset):
+        xdataset = run
+        run_id = xdataset.attrs['run_id']
+    elif isinstance(run, xr.DataArray):
+        xdataset = None
+        xdata_array = run
+        run_id = None
+    else:
+        raise ValueError(
+            "Invalid input type for run. "
+            "Must be run-ID, DataSet or xr.Dataset or xr.DataArray. "
+            f"Is {type(run)}"
+            )
+    if xdataset is not None:
+        if readout_name not in xdataset.data_vars:
+            readout_name = find_data_variable_from_keyword(xdataset, readout_name)
+        xdata_array = xdataset[readout_name]
+    ### Average over specified axes
+    xdata_array = avg_dataarray(xdata_array, avg_axes)
+    np_data = xdata_array.to_numpy()
+    return run_id, xdata_array, np_data
+
+def find_data_variable_from_keyword(
+        xdata_array: xr.DataArray, keyword: str | tuple) -> str:
+    """
+    Find the data variable corresponding to a keyword in the data-array.
+    
+    Args:
+        xdata_array (xr.DataArray): xarray data-array to search in
+        keyword (str): Keyword to search for
+    Returns:
+        data_variable (str): Data variable corresponding to keyword
+    """
+    if isinstance(keyword, str):
+        keyword = (keyword,)
+    if not isinstance(keyword, tuple):
+        raise ValueError(
+            f"Keyword must be a string or a tuple. Is {type(keyword)}")
+    data_variables = []
+    for data_variable in xdata_array.data_vars:
+        if all([subkey in str(data_variable) for subkey in keyword]):
+            data_variables.append(data_variable)
+    if len(data_variables) == 0:
+        raise ValueError(
+            f"Data variable not found for keyword {keyword}. "
+            f"Data variables are {xdata_array.data_vars}"
+            )
+    elif len(data_variables) > 1:
+        raise ValueError(
+            f"More than one data variable found for keyword {keyword}: "
+            f"{[str(var) for var in data_variables]}")
+    else:
+        return data_variables[0]
+
+def avg_dataarray(xdata_array: xr.DataArray, avg_axes: str | list = 'auto'):
+    """
+    Averages the data-array over the specified axes. If no axes are specified
+    the data-array is averaged over all axes.
+    
+    Args:
+        xdata_array (xr.DataArray): xarray data-array to be averaged
+        avg_axes (str | list): Axes to average over
+    """
+    if avg_axes is None:
+        avg_axes = []
+    if isinstance(avg_axes, str):
+        ### If 'auto' is given, find all axes with 'iteration' in the name
+        if avg_axes == 'auto':
+            avg_axes = []
+            for dim in xdata_array.dims:
+                if 'iteration' in dim:
+                    avg_axes.append(dim)
+        else:
+            avg_axes = [avg_axes]
+    ### Average over specified axes
+    for axis in avg_axes:
+        if hasattr(xdata_array, axis):
+            xdata_array = xdata_array.mean(axis)
+        else:
+            raise KeyError(
+                f"Avg. axis {axis} not found in xarray data-array")
+    return xdata_array

arbok_inspector/classes/dim.py

@@ -0,0 +1,26 @@
+"""Module for the Dim class."""
+
+class Dim:
+    """
+    Class representing a dimension of the data
+    """
+    def __init__(self, name):
+        """
+       Constructor for Dim class
+        
+        Args:
+            name (str): Name of the dimension
+            
+        Attributes:
+            name (str): Name of the dimension
+            option (str): Option for the dimension (average, select_value, x-axis, y-axis)
+            select_index (int): Index of the selected value for select_value option
+            ui_selector: Reference to the UI element for the dimension
+        """
+        self.name = name
+        self.option = None
+        self.select_index = 0
+        self.ui_selector = None
+
+    def __str__(self):
+        return self.name

arbok_inspector/classes/run.py

@@ -0,0 +1,238 @@
+"""
+Run class representing a single run of the experiment.
+"""
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+import ast
+import re
+import json
+from qcodes.dataset import load_by_id
+from nicegui import ui, app
+
+from arbok_inspector.classes.dim import Dim
+from arbok_inspector.widgets.build_xarray_grid import build_xarray_grid
+# from arbok_inspector.pages.database_browser import shared_data
+
+if TYPE_CHECKING:
+    from qcodes.dataset.data_set import DataSet
+    from xarray import Dataset
+AXIS_OPTIONS = ['average', 'select_value', 'y-axis', 'x-axis']
+
+
+class Run:
+    """
+    Class representing a run with its data and methods
+    """
+    def __init__(self, run_id: int):
+        """
+        Constructor for Run class
+        
+        Args:
+            run_id (int): ID of the run
+        """
+        self.run_id: int = run_id
+        self.title: str = f'Run ID: {run_id}  (-> add experiment)'
+        self.dataset: DataSet = load_by_id(run_id)
+        self.full_data_set: Dataset = self.dataset.to_xarray_dataset()
+        self.last_subset: Dataset = self.full_data_set
+
+        self.together_sweeps: bool = False
+        self.parallel_sweep_axes: dict = {}
+        self.sweep_dict: dict[int, Dim] = {}
+        self.load_sweep_dict()
+        self.dims: list[Dim] = list(self.sweep_dict.values())
+        self.dim_axis_option: dict[str, str|list[Dim]] = self.set_dim_axis_option()
+        print(self.dims)
+
+        self.plot_selection: list[str] = self.select_results_by_keywords(
+            app.storage.general["result_keywords"]
+        )
+        print(f"Initial plot selection: {self.plot_selection}")
+        self.plots_per_column: int = 2
+
+    def load_sweep_dict(self):
+        """
+        Load the sweep dictionary from the dataset
+        TODO: check metadata for sweep information!
+        Returns:
+            sweep_dict (dict): Dictionary with sweep information
+            is_together (bool): True if all sweeps are together, False otherwise
+        """
+        if "parallel_sweep_axes" in self.dataset.metadata:
+            conf = self.dataset.metadata["parallel_sweep_axes"]
+            conf = conf.replace("'", '"')  # Ensure JSON compatibility
+            print( conf)
+            conf = json.loads(conf)
+            self.parallel_sweep_axes = {int(i): sweeps for i, sweeps in conf.items()}
+            self.together_sweeps = True
+        else:
+            dims = self.full_data_set.dims
+            self.parallel_sweep_axes = {i: [dim] for i, dim in enumerate(dims)}
+            self.together_sweeps = False
+        self.sweep_dict = {
+            i: Dim(names[0]) for i, names in self.parallel_sweep_axes.items()
+            }
+        print(self.sweep_dict)
+        return self.sweep_dict
+
+    def set_dim_axis_option(self):
+        """
+        Set the default dimension options for the run in 4 steps:
+        1. Set all iteration dims to 'average'
+        2. Set the innermost dim to 'x-axis' (the last one that is not averaged)
+        3. Set the next innermost dim to 'y-axis'
+        4. Set all remaining dims to 'select_value'
+
+        Returns:
+            options (dict): Dictionary with keys 'average', 'select_value', 'y-axis',
+        """
+        options = {x: [] for x in AXIS_OPTIONS}
+        print(f"Setting average to {app.storage.general['avg_axis']}")
+        for dim in self.dims:
+            if app.storage.general["avg_axis"] in dim.name:
+                dim.option = 'average'
+                options['average'].append(dim)
+        for dim in reversed(self.dims):
+            if dim not in options['average'] and dim != options['x-axis']:
+                dim.option = "x-axis"
+                options['x-axis'] = dim
+                print(f"Setting x-axis to {dim.name}")
+                break
+        for dim in reversed(self.dims):
+            if dim not in options['average'] and dim != options['x-axis']:
+                dim.option = 'y-axis'
+                options['y-axis'] = dim
+                print(f"Setting y-axis to {dim.name}")
+                break
+        for dim in self.dims:
+            if dim not in options['average'] and dim != options['x-axis'] and dim != options['y-axis']:
+                dim.option = 'select_value'
+                options['select_value'].append(dim)
+                dim.select_index = 0
+                print(f"Setting select_value to {dim.name}")
+        return options
+
+    def select_results_by_keywords(self, keywords: list[str|tuple]) -> list[str]:
+        """
+        Select results by keywords in their name.
+        Args:
+            keywords (list): List of keywords to search for
+        Returns:
+            selected_results (list): List of selected result names
+        """
+        print(f"using keywords: {keywords}")
+        if keywords is None or len(keywords) == 0 or keywords == '':
+            return [next(iter(self.full_data_set.data_vars))]
+        s_quoted = re.sub(r'\b([a-zA-Z_][a-zA-Z0-9_]*)\b', r'"\1"', keywords)
+        try:
+            keywords = ast.literal_eval(s_quoted)
+        except (SyntaxError, ValueError):
+            print(f"Error parsing keywords: {s_quoted}")
+            keywords = []
+            ui.notify(
+                f"Error parsing result keywords: {s_quoted}. Please use a valid Python list.",
+                color='red',
+                position='top-right'
+            )
+        if not isinstance(keywords, list):
+            keywords = [keywords]
+        selected_results = []
+        print(f"using keywords: {keywords}")
+        for result in self.full_data_set.data_vars:
+            for keyword in keywords:
+                if isinstance(keyword, str) and keyword in str(result):
+                    selected_results.append(result)
+                elif isinstance(keyword, tuple) and all(
+                        subkey in str(result) for subkey in keyword):
+                    selected_results.append(result)
+        selected_results = list(set(selected_results))  # Remove duplicates
+        if len(selected_results) == 0:
+            selected_results = [next(iter(self.full_data_set.data_vars))]
+        print(f"Selected results: {selected_results}")
+        return selected_results
+    
+    def update_subset_dims(self, dim: Dim, selection: str, index = None):
+        """
+        Update the subset dimensions based on user selection.
+
+        Args:
+            dim (Dim): The dimension object to update
+            selection (str): The new selection option
+                ('average', 'select_value', 'x-axis', 'y-axis')
+            index (int, optional): The index for 'select_value' option. Defaults to None.
+        """
+        text = f'Updating subset dims: {dim.name} to {selection}'
+        print(text)
+        ui.notify(text, position='top-right')
+
+        ### First, remove old option this dim was on
+        for option in ['average', 'select_value']:
+            if dim in self.dim_axis_option[option]:
+                print(f"Removing {dim.name} from {option}")
+                self.dim_axis_option[option].remove(dim)
+                dim.option = None
+        if dim.option in ['x-axis', 'y-axis']:
+            print(f"Removing {dim.name} from {dim.option}")
+            self.dim_axis_option[dim.option] = None
+
+        ### Now, set new option
+        if selection in ['average', 'select_value']:
+            # dim.ui_selector.value = selection
+            dim.select_index = index
+            self.dim_axis_option[selection].append(dim)
+            return
+        if selection in ['x-axis', 'y-axis']:
+            old_dim = self.dim_axis_option[selection]
+            self.dim_axis_option[selection] = dim
+            if old_dim is not None:
+                # Set previous dim (having this option) to 'select_value'
+                # Required since x and y axis ahve to be unique
+                print(f"Updating {old_dim.name} to {dim.name} on {selection}")
+                if old_dim.option in ['x-axis', 'y-axis']:
+                    self.dim_axis_option['select_value'].append(old_dim)
+                    old_dim.option = 'select_value'
+                    old_dim.ui_selector.value = 'select_value'
+                    self.update_subset_dims(old_dim, 'select_value', old_dim.select_index)
+        dim.ui_selector.update()
+
+    def generate_subset(self):
+        """
+        Generate the subset of the full dataset based on the current dimension options.
+        Returns:
+            sub_set (xarray.Dataset): The subset of the full dataset
+        """
+        # TODO: take the averaging out of this! We only want to average if necessary
+        # averaging can be computationally intensive!
+        sub_set = self.full_data_set
+        for avg_axis in self.dim_axis_option['average']:
+            sub_set = sub_set.mean(dim=avg_axis.name)
+        sel_dict = {d.name: d.select_index for d in self.dim_axis_option['select_value']}
+        sub_set = sub_set.isel(**sel_dict).squeeze()
+        self.last_subset = sub_set
+        return sub_set
+
+    def update_plot_selection(self, value: bool, readout_name: str):
+        """
+        Update the plot selection based on user interaction.
+
+        Args:
+            value (bool): True if the result is selected, False otherwise
+            readout_name (str): Name of the result to update
+        """
+        print(f"{readout_name= } {value= }")
+        pretty_readout_name = readout_name.replace("__", ".")
+        if readout_name not in self.plot_selection:
+            self.plot_selection.append(readout_name)
+            ui.notify(
+                message=f'Result {pretty_readout_name} added to plot selection',
+                position='top-right'
+                )
+        else:
+            self.plot_selection.remove(readout_name)
+            ui.notify(
+                f'Result {pretty_readout_name} removed from plot selection',
+                position='top-right'
+            )
+        print(f"{self.plot_selection= }")
+        build_xarray_grid()

arbok_inspector/cli.py

@@ -0,0 +1,4 @@
+import runpy
+
+def main():
+    runpy.run_module('arbok_inspector.main', run_name='__main__')

arbok_inspector/dev.py

@@ -0,0 +1,19 @@
+from nicegui import ui
+from pathlib import Path
+import asyncio
+from typing import Optional
+
+from arbok_inspector.state import inspector
+from arbok_inspector.pages import greeter, database_browser
+
+def main():
+    ui.run(
+        title='Arbok Inspector',
+        favicon='🐍',
+        dark=True,
+        show=True,
+        port=8090
+    )
+
+if __name__ in {"__main__", "__mp_main__"}:
+    main()

arbok_inspector/helpers/string_formaters.py

@@ -0,0 +1,33 @@
+
+import xarray as xr
+from arbok_inspector.helpers.unit_formater import unit_formatter
+
+def title_formater(run):
+    """
+    Format title string for plots based on selected dimensions.
+    
+    Args:
+        run: The Run object containing the data.
+    Returns:
+        A formatted title string.
+    """
+    title = ""
+    for dim in run.dim_axis_option["select_value"]:
+        title += f"{dim.name } = {unit_formatter(run, dim, dim.select_index)}<br>"
+    return title
+
+def axis_label_formater(ds: xr.DataArray, dim_name: str) -> str:
+    """
+    Format axis label by replacing '__' with '.' and bolding the last part.
+    
+    Args:
+        dim_name: The dimension name string.
+    Returns:
+        A formatted axis label string.
+    """
+    dim_list = dim_name.split('__')
+    print(f"{dim_list=}")
+    if len(dim_list) > 1:
+        return f"{'.'.join(dim_list[:-1])}.<b>{dim_list[-1]}</b> ({ds.coords[dim_name].unit})"
+    else:
+        return f"<b>{dim_list[0]}</b> ({ds.coords[dim_name].unit})"

arbok_inspector/helpers/unit_formater.py

@@ -0,0 +1,29 @@
+"""Helper functions for formatting units with SI prefixes."""
+
+def unit_formatter(run, dim, index: int) -> str:
+    """
+    If value is larger than 1e3, format with SI prefix.
+    Same if smaller than 1e-3.
+
+    Args:
+        run (Run): Run object containing the data
+        dim (Dim): Dimension object
+        index (int): Index of the value to format
+    """
+    unit_tuples = [
+        ('G', 1e9), ('M', 1e6), ('k', 1e3), ('m', 1e-3), ('µ', 1e-6), ('n', 1e-9)]
+    try:
+        value = run.full_data_set[dim.name].values[index]
+        unit = run.full_data_set[dim.name].unit
+        if abs(value) >= 1e3 or (abs(value) < 1e-3 and value != 0):
+            for prefix, factor in unit_tuples:
+                if abs(value) >= factor or (abs(value) < 1e-3 and value != 0 and factor < 1):
+                    scaled_value = value / factor
+                    return f'{scaled_value:.3f} {prefix}<b>{unit}</b>'
+        if unit is None or unit == '':
+            return f'{value:.3f}'
+        else:
+            return f'{value:.3f} ({unit})'
+    except Exception as e:
+        print(f"Error in unit_formatter: {e}")
+        return 'N/A'

arbok_inspector/main.py

@@ -0,0 +1,15 @@
+from nicegui import ui
+from pathlib import Path
+import asyncio
+from typing import Optional
+
+from arbok_inspector.state import inspector
+from arbok_inspector.pages import greeter, database_browser
+
+ui.run(
+    title='Arbok Inspector',
+    favicon='🐍',
+    dark=True,
+    show=True,
+    port=8090
+)

arbok_inspector/pages/__init__.py

@@ -0,0 +1,2 @@
+from .greeter import greeter_page
+from .database_browser import database_browser_page