arbok-inspector 1.3.0__py3-none-any.whl

arbok_inspector/pages/run_view.py

@@ -0,0 +1,259 @@
+"""Run view page showing the data and plots for a specific run"""
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from datetime import datetime, timedelta
+import json
+import os
+import importlib.resources as resources
+
+from nicegui import ui, app
+
+from arbok_inspector.state import inspector
+from arbok_inspector.widgets.build_xarray_grid import build_xarray_grid
+from arbok_inspector.widgets.build_xarray_html import build_xarray_html
+from arbok_inspector.widgets.build_run_view_actions import build_run_view_actions
+from arbok_inspector.helpers.unit_formater import unit_formatter
+from arbok_inspector.classes.qcodes_run import QcodesRun
+from arbok_inspector.classes.native_run import NativeRun
+
+
+from arbok_inspector.classes.dim import Dim
+
+RUN_TABLE_COLUMNS = [
+    {'field': 'name', 'filter': 'agTextColumnFilter', 'floatingFilter': True},
+    {'field': 'size'},
+    {'field': 'x', 'checkboxSelection': True},
+    {'field': 'y', 'checkboxSelection': True},
+    {'field': 'average', 'checkboxSelection': True},
+]
+
+AXIS_OPTIONS = ['average', 'select_value', 'y-axis', 'x-axis']
+
+EXPANSION_CLASSES = 'w-full p-0 gap-1 border border-gray-400 rounded-lg no-wrap items-start pt-0 mt-0'
+TITLE_CLASSES = 'text-lg font-semibold'
+
+@ui.page('/run/{run_id}')
+async def run_page(run_id: str):
+    """
+    Page showing the details and plots for a specific run.
+
+    Args:
+        run_id (str): ID of the run to display
+    """
+    ui.page_title(f"{run_id}")
+    _ = await ui.context.client.connected()
+    if 'run' in app.storage.tab:
+        print('run already exists!')
+    if inspector.database_type == 'qcodes':
+        run = QcodesRun(int(run_id))
+    elif 'arbok_native':
+        run = NativeRun(int(run_id))
+    else:
+        raise ValueError(
+            "Database type must be 'qcodes' or 'arbok_native is:"
+            f"{inspector.database_type}")
+    
+
+    app.storage.tab["placeholders"] = {'plots': None}
+    app.storage.tab["run"] = run
+    with resources.files("arbok_inspector.configurations").joinpath("1d_plot.json").open("r") as f:
+        app.storage.tab["plot_dict_1D"] = json.load(f)
+    with resources.files("arbok_inspector.configurations").joinpath("2d_plot.json").open("r") as f:
+        app.storage.tab["plot_dict_2D"] = json.load(f)
+
+    ui.label(f'Run-ID: {run_id}').classes('text-2xl font-bold')
+    with ui.row().classes('w-full gap-4'):
+        with ui.column().classes('flex-none'):
+            with ui.card().classes('w-full gap-2'):
+                ui.label("Coordinates:").classes('text-lg font-semibold pl-2')
+                ui.separator().classes('w-full my-1')
+                for i, _ in run.parallel_sweep_axes.items():
+                    add_dim_dropdown(sweep_idx = i)
+            with ui.card().classes('w-full gap-2'):
+                ui.label("Results:").classes(TITLE_CLASSES)
+                for i, result in enumerate(run.full_data_set):
+                    value = False
+                    if result in run.plot_selection:
+                        value = True
+                    ui.checkbox(
+                        text = result.replace("__", "."),
+                        value = value,
+                        on_change = lambda e, r=result: run.update_plot_selection(e.value, r),
+                    ).classes('text-sm h-4').props('color=purple')
+            with ui.card().classes('w-full gap-2'):
+                ui.label("Actions:").classes(TITLE_CLASSES)
+                build_run_view_actions()
+            with ui.expansion('Run info', icon = 'info').classes('w-full gap-2'):
+                # ui.label("Run info:").classes(TITLE_CLASSES)
+                for column_name, conf in run.database_columns.items():
+                    value = str(conf['value'])
+                    if len(value) > 20 or value is None:
+                        continue
+                    print(column_name, type(value))
+                    if 'label' in conf:
+                        label = ui.label(f"{conf['label']}: ")
+                    else:
+                        label = ui.label(f"{column_name.upper()}: ")
+                    label.classes('font-semibold m-0 p-0"')
+                    ui.label(value).classes("m-0 p-0 ml-5")
+
+        with ui.column().classes('flex-1 min-w-0'):
+            with ui.expansion('Plots', icon='stacked_line_chart', value=True)\
+                .classes(EXPANSION_CLASSES):
+                app.storage.tab["placeholders"]["plots"] = ui.row().\
+                    classes('w-full min-h-[50vh] p-1 items-stretch')
+                build_xarray_grid()
+
+                    #.style('line-height: 1rem; padding-top: 0; padding-bottom: 0;')
+            with ui.expansion('xarray summary', icon='summarize', value=False)\
+                .classes(EXPANSION_CLASSES):
+                build_xarray_html()
+            with ui.expansion('analysis', icon='science', value=False)\
+                .classes(EXPANSION_CLASSES):
+                with ui.row():
+                    ui.label("Working on it!  -Andi").classes(TITLE_CLASSES)
+            with ui.expansion('metadata', icon='numbers', value=False)\
+                .classes(f"{EXPANSION_CLASSES}  overflow-x-auto"):
+                placeholder_metadata = {}
+                
+                placeholder_metadata['code'] = ui.code(
+                    content = 'Placeholder for QUA program',
+                    language = 'python')\
+                    .classes('w-full overflow-x-auto whitespace-pre')
+                ui.button(
+                    icon = 'code',
+                    text="load qua program",
+                    on_click = lambda: load_qua_code(run, placeholder_metadata),
+                )
+                ui.button(
+                    icon = 'download',
+                    text="download serialized qua program",
+                    on_click = lambda: download_qua_code(run),
+                )
+
+def add_dim_dropdown(sweep_idx: int):
+    """
+    Add a dropdown to select the dimension option for a given sweep index.
+
+    Args:
+        sweep_idx (int): Index of the sweep to add the dropdown for
+    """
+    run = app.storage.tab["run"]
+    width = 'w-full'
+    dim = run.sweep_dict[sweep_idx]
+    local_placeholder = {"slider": None}
+    #with ui.column().classes('w-full no-wrap items-center gap-1'):
+    dims_names = run.parallel_sweep_axes[sweep_idx]
+    #ui.separator().classes('w-full my-1')
+    # with ui.card().classes('w-full gap-1 px-2 py-2'):
+    ui.radio(
+        options = dims_names,
+        value=dim.name,
+        on_change = lambda e: update_sweep_dim_name(dim, e.value)
+        ).classes(f"{width}  text-xs m-0 p-0").props('dense')
+    ui_element = ui.select(
+        options = AXIS_OPTIONS,
+        value = str(dim.option),
+        label = f'{dim.name.replace("__", ".")}',
+        on_change = lambda e: update_dim_selection(
+            dim, e.value, local_placeholder["slider"])
+    ).classes(f"{width} text-sm m-0 p-0").props('dense')
+    dim.ui_selector = ui_element
+    local_placeholder["slider"] = ui.column().classes('w-full')
+    if dim.option == 'select_value':
+        build_dim_slider(run, dim, local_placeholder["slider"])
+
+def update_dim_selection(dim: Dim, value: str, slider_placeholder):
+    """
+    Update the dimension/sweep selection and rebuild the plot grid.
+
+    Args:
+        dim (Dim): The dimension object to update
+        value (str): The new selection value
+        slider_placeholder: The UI placeholder to update
+    """
+    run = app.storage.tab["run"]
+    if slider_placeholder is not None:
+        slider_placeholder.clear()
+    print(value)
+    if value == 'average':
+        run.update_subset_dims(dim, 'average')
+        dim.option = 'average'
+    if value == 'select_value':  
+        with slider_placeholder:
+            build_dim_slider(run, dim, slider_placeholder)
+    else:
+        run.update_subset_dims(dim, value)
+        dim.option = value
+    build_xarray_grid()
+
+def build_dim_slider(run: Runm, dim: Dim, slider_placeholder):
+    """
+    Build a slider for selecting the index of a dimension.
+
+    Args:
+        dim (Dim): The dimension object
+        slider_placeholder: The UI placeholder to add the slider to
+    """
+    dim_size = run.full_data_set.sizes[dim.name]
+    with ui.row().classes("w-full items-center"):
+        with ui.column().classes('flex-grow'):
+            slider = ui.slider(
+                min=0, max=dim_size - 1, step=1, value=0,
+                on_change=lambda e: run.update_subset_dims(dim, 'select_value', e.value),
+                ).classes('flex-grow')\
+                .props('color="purple" markers label-always')
+        label = ui.html('').classes('shrink-0 text-right px-2 py-1 bg-purple text-white rounded-lg text-xs font-normal text-center')
+        update_value_from_dim_slider(label, slider, dim, plot = False)
+        slider.on(
+            'update:model-value',
+            lambda e: update_value_from_dim_slider(label, slider, dim),
+            throttle=0.2, leading_events=False)
+
+def update_value_from_dim_slider(label, slider, dim: Dim, plot = True):
+    """
+    Update the label next to the slider with the current value and unit.
+    
+    Args:
+        label: The UI label to update
+        slider: The UI slider to get the value from
+        dim (Dim): The dimension object
+    """
+    run = app.storage.tab["run"]
+    label_txt = f' {unit_formatter(run, dim, slider.value)} '
+    label.set_content(label_txt)
+    if plot:
+        build_xarray_grid()
+
+def update_sweep_dim_name(dim: Dim, new_name: str):
+    """
+    Update the name of the dimension in the sweep dict and the dim object.
+    
+    Args:
+        dim (Dim): The dimension object to update
+        new_name (str): The new name for the dimension
+    """
+    run = app.storage.tab["run"]
+    dim.name = new_name
+    dim.ui_selector.label = new_name.replace("__", ".")
+    build_xarray_grid()
+
+def load_qua_code(run: Run, placeholder: dict):
+    """Load and display the QUA code for the given run."""
+    try:
+        qua_code = run.get_qua_code(as_string = True)
+        qua_code = qua_code.split("config = {")[0]
+        placeholder['code'].set_content(qua_code)
+    except Exception as e:
+        ui.notify(f'Error loading QUA code: {str(e)}', type='negative')
+        raise e
+
+def download_qua_code(run: Run) -> None:
+    """Download the serialized QUA code for the given run."""
+    try:
+        qua_code_bytes = run.get_qua_code(as_string = False)
+        ui.download(qua_code_bytes, 'test.py')
+        #os.remove(file_name)
+    except Exception as e:
+        ui.notify(f'Error downloading QUA code: {str(e)}', type='negative')
+        raise e

arbok_inspector/state.py

@@ -0,0 +1,101 @@
+from nicegui import ui
+from pathlib import Path
+import asyncio
+from typing import Optional
+
+from qcodes.dataset import initialise_or_create_database_at
+
+import fsspec
+import sqlite3
+from sqlalchemy import create_engine, select, func
+
+class ArbokInspector:
+    def __init__(self):
+        self.qcodes_database_path: Optional[Path] = None
+        self.initial_dialog = None
+        self.database_type = None  # 'qcodes' or 'arbok'
+
+        self.qcodes_database_path = None
+        self.conn = None
+        self.cursor = None
+        
+        self.database_engine = None
+        self.minio_filesystem = None
+        self.minio_bucket = None
+        
+    def connect_qcodes_database(self):
+        self.conn = sqlite3.connect(self.qcodes_database_path)
+        self.conn.row_factory = sqlite3.Row
+        self.cursor = self.conn.cursor()
+        initialise_or_create_database_at(self.qcodes_database_path)
+
+    def connect_to_qcodes_database(self, path_input) -> None:
+        """Connect to a QCoDeS database given a file path input widget."""
+        self.database_engine = None
+        self.minio_filesystem = None
+        self.minio_bucket = None
+        if path_input.value is None:
+            ui.notify('Please enter a file path', type='warning')
+            return
+        try:
+            file_path = Path(path_input.value)
+            if file_path.exists():
+                self.qcodes_database_path = file_path
+                ui.notify(f'Database path set: {file_path.name}', type='positive')
+                try:
+                    self.connect_qcodes_database()
+                    if self.initial_dialog:
+                        self.initial_dialog.close()
+                    self.database_type = 'qcodes'
+                    ui.navigate.to('/browser')
+                except sqlite3.Error as e:
+                    ui.notify(f'Error connecting to database: {str(e)}', type='negative')
+            else:
+                ui.notify('File does not exist', type='negative')
+        except Exception as ex:
+            ui.notify(f'Error: {str(ex)}', type='negative')
+        self.database_type = 'qcodes'
+    
+    def connect_to_arbok_database(
+        self,
+        database_url: str,
+        minio_url: str,
+        minio_user: str,
+        minio_password: str,
+        minio_bucket: str) -> None:
+        """
+        Connect to a native Arbok database given connection parameters.
+        
+        Args:
+            database_url (str): The database connection URL.
+            minio_url (str): The MinIO server URL.
+            minio_user (str): The MinIO username.
+            minio_password (str): The MinIO password.
+            minio_bucket (str): The MinIO bucket name.
+        """
+        self.qcodes_database_path = None
+        self.conn = None
+        self.cursor = None
+        try:
+            self.database_engine = create_engine(database_url)
+        except Exception as ex:
+            ui.notify(f'Error creating database engine: {str(ex)}', type='negative')
+            return
+
+        try:
+            self.minio_filesystem = fsspec.filesystem(
+                protocol = "s3",
+                client_kwargs={"endpoint_url": minio_url},
+                key=minio_user,
+                secret=minio_password
+                )
+        except Exception as ex:
+            ui.notify(f'Error connecting to MinIO: {str(ex)}', type='negative')
+            return
+        if self.initial_dialog:
+            self.initial_dialog.close()
+        self.database_type = 'native_arbok'
+        self.minio_bucket = minio_bucket
+        ui.navigate.to('/browser')
+    
+inspector = ArbokInspector()

arbok_inspector/test_main.py

@@ -0,0 +1,65 @@
+import numpy as np
+import xarray as xr
+import plotly.express as px
+from nicegui import ui
+
+# --- Create a sample 4D DataArray ---
+data = np.random.rand(5, 10, 20, 30)  # shape: time, depth, y, x
+coords = {
+    'time': np.arange(5),
+    'depth': np.linspace(0, 100, 10),
+    'y': np.linspace(-5, 5, 20),
+    'x': np.linspace(-5, 5, 30),
+}
+array = xr.DataArray(data, dims=('time', 'depth', 'y', 'x'), coords=coords)
+
+# --- Define axes to plot ---
+x_dim = 'x'
+y_dim = 'y'
+slider_dims = [dim for dim in array.dims if dim not in (x_dim, y_dim)]
+
+# --- UI State ---
+slider_values = {dim: 0 for dim in slider_dims}
+plot_container = ui.row().classes('w-full justify-center')
+
+# --- Heatmap update function ---
+def update_plot():
+    # Index the array using current slider values
+    sel = {dim: slider_values[dim] for dim in slider_dims}
+    slice_2d = array.isel(**sel)
+
+    # Convert to plotly figure
+    fig = px.imshow(
+        slice_2d.values,
+        labels={'x': x_dim, 'y': y_dim},
+        x=array.coords[x_dim].values,
+        y=array.coords[y_dim].values,
+        color_continuous_scale='Viridis',
+    )
+    fig.update_layout(title=f'{x_dim} vs {y_dim} | ' + ', '.join([f'{dim}={slider_values[dim]}' for dim in slider_dims]))
+
+    plot_container.clear()
+    with plot_container:
+        ui.plotly(fig).classes('max-w-3xl max-h-96')
+
+# --- Create sliders for all non-plotted dimensions ---
+for dim in slider_dims:
+    max_index = len(array.coords[dim]) - 1
+    def make_slider(d=dim):
+        def on_change(val):
+            slider_values[d] = int(val)
+            update_plot()
+        ui.slider(
+            min=0,
+            max=max_index,
+            value=0,
+            step=1,
+            on_change=on_change,
+            #abel=f'{d} ({array.coords[d].values[0]})'
+        ).props('label-always').classes('w-full')
+    make_slider()
+
+# --- Initial plot ---
+update_plot()
+
+ui.run()

arbok_inspector/widgets/build_run_selecter.py

@@ -0,0 +1,163 @@
+from datetime import datetime, timedelta
+
+from nicegui import ui, app
+from sqlalchemy import text
+
+from arbok_inspector.state import inspector
+
+small_col_width = 50
+med_col_width = 60
+
+QCODES_RUN_GRID_COLUMN_DEFS = [
+    {'headerName': 'Run ID', 'field': 'run_id', "width": small_col_width},
+    {'headerName': 'Name', 'field': 'name'},
+    {'headerName': 'Experiment', 'field': 'experiment_name'},
+    {'headerName': '# Results', 'field': 'result_counter', "width": small_col_width},
+    {'headerName': 'Started', 'field': 'run_timestamp', "width": small_col_width},
+    {'headerName': 'Finish', 'field': 'completed_timestamp', "width": small_col_width},
+]
+NATIVE_RUN_GRID_COLUMN_DEFS = [
+    {'headerName': 'Run ID', 'field': 'run_id', "width": small_col_width},
+    {'headerName': 'Name', 'field': 'name'},
+    {'headerName': 'Experiment', 'field': 'experiment'},
+    {'headerName': '# results', 'field': 'result_count', "width": med_col_width},
+    {'headerName': '# batches', 'field': 'batch_count', "width": med_col_width},
+    {'headerName': 'started', 'field': 'start_time', "width": med_col_width},
+    {'headerName': 'last result', 'field': 'completed_time', "width": med_col_width},
+]
+AGGRID_STYLE = 'height: 95%; min-height: 0;'
+
+def build_run_selecter(target_day):
+    container = app.storage.tab['run_selecter']
+    container.clear()
+
+    offset_hours = app.storage.general["timezone"]
+    run_grid_rows = []
+    print(f"Showing runs from {target_day}")
+    if inspector.database_type == 'qcodes':
+        rows = get_qcodes_runs_for_day(inspector.cursor, target_day, offset_hours)
+        column_defs = QCODES_RUN_GRID_COLUMN_DEFS
+    else:
+        rows = get_native_arbok_runs_for_day(inspector.database_engine, target_day, offset_hours)
+        column_defs = NATIVE_RUN_GRID_COLUMN_DEFS
+    run_grid_rows = []
+    columns = [x['field'] for x in column_defs]
+    for run in rows:
+        run_dict = {}
+        for key in columns:
+            if key in run:
+                value = run[key]
+                if 'time' in key:
+                    if value is not None:
+                        local_dt = datetime.utcfromtimestamp(value)
+                        local_dt += timedelta(hours=offset_hours)
+                        value = local_dt.strftime('%H:%M:%S')
+                    else:
+                        value = 'N/A'
+                run_dict[key] = value
+        run_grid_rows.insert(0, run_dict)
+    with container:
+        ui.aggrid(
+                {
+                    'defaultColDef': {'flex': 1},
+                    'columnDefs': column_defs,
+                    'rowData': run_grid_rows,
+                    'rowSelection': 'multiple',
+                },
+            ).classes('ag-theme-balham-dark').style(
+                AGGRID_STYLE
+            ).on(
+                'cellClicked',
+                lambda event: open_run_page(event.args['data']['run_id'])
+            )
+    ui.notify(
+        'Run selector updated: \n'
+        f'found {len(run_grid_rows)} run(s)',
+        type='positive',
+        multi_line=True,
+        classes='multi-line-notification',
+        position = 'top-right'
+    )
+
+def get_qcodes_runs_for_day(
+    cursor, target_day: str, offset_hours: float
+) -> list[dict]:
+    """
+    Fetch runs from a QCoDeS (SQLite) database, joined with experiments,
+    excluding the 'qua_program' and 'snapshot' columns entirely.
+    """
+    hours = int(offset_hours)
+    minutes = int((offset_hours - hours) * 60)
+    offset_str = f"{'+' if offset_hours >= 0 else '-'}{abs(hours):02d}:{abs(minutes):02d}"
+
+    # get all columns except the ones we want to exclude
+    exclude_columns = {'qua_program', 'snapshot'}
+    cursor.execute("PRAGMA table_info(runs)")
+    all_columns = [col['name'] for col in cursor.fetchall() if col['name'] not in exclude_columns]
+
+    # construct SELECT statement
+    columns_str = ", ".join(f"r.{col}" for col in all_columns)
+
+    query = f"""
+        SELECT {columns_str}, e.name AS experiment_name
+        FROM runs r
+        JOIN experiments e ON r.exp_id = e.exp_id
+        WHERE DATE(datetime(r.run_timestamp, 'unixepoch', '{offset_str}')) = ?
+        ORDER BY r.run_timestamp;
+    """
+
+    cursor.execute(query, (target_day,))
+    rows = cursor.fetchall()
+    row_dicts = [dict(row) for row in rows]
+    return row_dicts
+     
+NATIVE_COLUMNS = {
+    'run_id': 'run ID',
+    'name': 'name',
+    'result_count': '# results',
+    'batch_count': '# batches',
+    'start_time': 'started',
+    'completed_time': 'last result',
+    'is_completed': 'completed'
+}
+
+def get_native_arbok_runs_for_day(
+    engine,
+    target_day: str,
+    offset_hours: float) -> list[dict]:
+    """
+    Fetch runs from a native Arbok database
+    
+    Args:
+        engine: SQLAlchemy engine connected to the database
+        target_day (str): The target day in 'YYYY-MM-DD'
+        offset_hours (float): The timezone offset in hours
+    Returns:
+        list[dict]: List of runs as dictionaries
+    """
+    query = text("""
+        SELECT r.*, e.name AS experiment_name
+        FROM runs r
+        JOIN experiments e ON r.exp_id = e.exp_id
+        WHERE (to_timestamp(r.start_time) + (:offset_hours || ' hours')::interval)::date = :target_day
+        ORDER BY r.start_time;
+    """)
+
+    with engine.connect() as conn:
+        result = conn.execute(
+            query, {"offset_hours": offset_hours, "target_day": target_day}
+        )
+        runs_filtered = [
+            {**{col: row[col] for col in NATIVE_COLUMNS.keys()},
+            "experiment": row["experiment_name"]}
+            for row in result.mappings()
+        ]
+
+    return runs_filtered
+
+def open_run_page(run_id: int):
+    app.storage.general["avg_axis"] = app.storage.tab["avg_axis_input"].value
+    app.storage.general["result_keywords"] = app.storage.tab["result_keyword_input"].value
+    print(f"Result Keywords:")
+    print(app.storage.general['result_keywords'])
+    ui.navigate.to(f'/run/{run_id}', new_tab=True)