diff --git a/src/badger/archive.py b/src/badger/archive.py index 3f9d242b..ef14cafb 100644 --- a/src/badger/archive.py +++ b/src/badger/archive.py @@ -2,6 +2,7 @@ import time import warnings import logging +import yaml from badger.utils import ts_float_to_str from badger.settings import init_settings @@ -156,7 +157,16 @@ def load_run(run_fname: str) -> Routine: # TODO: create utility function to catch warnings to remove code # duplication with warnings.catch_warnings(record=True) as caught_warnings: - routine = Routine.from_file(filename) + # Note: replacing Routine.from_file with yaml.safe_load for loading routine yaml file + # for formulas update, since info in subdicts was not being loaded. + with open(filename, "r") as f: + data = yaml.safe_load(f) + routine = Routine(**data) + logger.debug( + f"yaml.safe_load: {routine.formulas}, routine.vocs: {routine.vocs}" + ) + + # routine = Routine.from_file(filename) # Check if any user warnings were caught for warning in caught_warnings: diff --git a/src/badger/formula_utils.py b/src/badger/formula_utils.py new file mode 100644 index 00000000..3d7f375e --- /dev/null +++ b/src/badger/formula_utils.py @@ -0,0 +1,310 @@ +import re +import ast +import math +import statistics +from typing import Set, Dict, Any, Tuple, List +import logging + +logger = logging.getLogger(__name__) + +_UNSAFE_MATH_FUNCS = { + "factorial", + "gamma", + "lgamma", + "comb", + "perm", +} + +_SAFE_NUMPY_FUNCS = { + "mean", + "median", + "std", + "min", + "max", + "sum", + "abs", + "sqrt", + "square", + "exp", + "log", + "log10", + "log2", + "sin", + "cos", + "tan", +} + +_SAFE_BUILTINS = { + "abs", + "sum", + "min", + "max", + "round", + "int", + "float", + "bool", +} + +_ALLOWED_FUNC_NAMES: set[str] = ( + { + name + for name in dir(math) + if callable(getattr(math, name)) + and not name.startswith("_") + and name not in _UNSAFE_MATH_FUNCS + } + | { + name + for name in dir(statistics) + if callable(getattr(statistics, name)) and not name.startswith("_") + } + | _SAFE_NUMPY_FUNCS + | _SAFE_BUILTINS + | { + "pi", + "e", + "tau", + } +) + +_ALLOWED_NODES = ( + ast.Expression, + ast.Constant, + ast.UnaryOp, + ast.UAdd, + ast.USub, + ast.BinOp, + ast.Add, + ast.Sub, + ast.Mult, + ast.Div, + ast.Pow, + ast.Mod, + ast.Call, + ast.Name, + ast.Load, + ast.BitXor, + ast.List, + ast.Tuple, +) + + +def validate_formula(expr: str, allowed_symbols: Set[str]) -> None: + """ + Validates a formula expression for safe evaluation. + Adapted from https://github.com/slaclab/trace/blob/main/trace/utilities/formula_validation.py#L28 + + This function parses the expression using Python's AST and validates that: + - Only allowed operators and functions are used + - All variable names are in the allowed symbols set + - The expression is syntactically valid + + Parameters + ---------- + expr : str + The mathematical expression to validate + allowed_symbols : Set[str] + Set of allowed variable names in the expression + + Raises: + ValueError: If the expression contains: + - Disallowed AST node types or operators + - Lists/tuples with complex nested expressions (only simple values allowed) + - Lists/tuples exceeding 50 elements + - Function calls to non-allowed functions + - References to undefined symbols (not in allowed_symbols or allowed functions) + SyntaxError: If the expression cannot be parsed as valid Python. + """ + tree = ast.parse(expr, mode="eval") + + for node in ast.walk(tree): + if not isinstance(node, _ALLOWED_NODES): + raise ValueError(f'Operator "{type(node).__name__}" not allowed') + + if isinstance(node, (ast.List, ast.Tuple)): + for element in node.elts: + if not isinstance( + element, (ast.Name, ast.Constant, ast.BinOp, ast.Call) + ): + raise ValueError( + f"Lists/tuples can only contain simple values, not {type(element).__name__}" + ) + + if len(node.elts) > 50: # arbitrary + raise ValueError(f"List/tuple too large: {len(node.elts)} elements") + + if isinstance(node, ast.Call): + fn = node.func.id if isinstance(node.func, ast.Name) else None + if fn not in _ALLOWED_FUNC_NAMES: + raise ValueError(f'Function "{fn}" not permitted') + + if isinstance(node, ast.Name): + if node.id not in allowed_symbols and node.id not in _ALLOWED_FUNC_NAMES: + raise ValueError( + f'Unknown symbol "{node.id}". Use `backticks` around variable names' + ) + + +def sanitize_for_validation(expr: str) -> tuple[str, set[str]]: + """Replace backtick-quoted variables like `PV1` with temp identifiers (v0, v1, ...). + + Returns (python_expr, allowed_syms). `allowed_syms` should be passed to validate_formula. + """ + mapping: dict[str, str] = {} + + def _repl(match: re.Match) -> str: + var = match.group(1) + if var not in mapping: + mapping[var] = f"v{len(mapping)}" + return mapping[var] + + # Match `...` (no backticks inside); preserve everything else unchanged + python_expr = re.sub(r"`([^`]+)`", _repl, expr) + return python_expr, set(mapping.values()) + + +VAR = re.compile(r"`([^`]+)`") # find backticked `var` tokens + + +def expanded_formula_mapping( + data: dict, +) -> Tuple[Dict[str, str], Dict[str, str], List[str]]: + """ + Provide a forward and reverse mapping for formula names to expanded formula strings in terms of + environment variables. Non-formulas will map to themselves. Also returns a list of base_variables + within the formulas. + + Note this generally functions well in my testing, but may have a few bugs and would + probably be better to rewrite to use dfs dict traversal rather than recusion + + Variable substitution rules + --------------------------- + Substitutes variables marked with backticks. If the backticked name is present in + "variable_mapping", recursively expands the mapped formula. + + Otherwise, it's treated as a leaf node/base observable if: + - the name is not in node["variable_mapping"] + - the mapping value is falsy (None, "") + + Parameters + ---------- + - data: dict + Badger routine "data" + Note: this probably doesn't need the full data dict, only + data["vocs"] and data["formulas"] if present + + Returns + ------- + - forward: Dict[str, str] + Mapping from each output name to either its expanded formula string (for formulas) + or to itself (for base observables/non-formulas). + - reverse: Dict[str, str] + Reverse lookup mapping from expanded formula string (or identity name) + to the first corresponding key in forward + - base_observables: List[str] + Names of leaf/base observables referenced during expansion within formulas + + Additional Notes + ---------------- + - caches expanded names to avoid expanding the same named formula multiple times. + - Detects circular references using stack (may have bug) + + """ + cache: Dict[str, str] = {} + stack = set() + + formulas = data.get("formulas", {}) or {} + + output_names = list(data["vocs"].output_names) + selected_formulas = { + sel_name: formulas[sel_name] + for sel_name in output_names + if sel_name in formulas + } + base_observables = [] # keep trach of observables within formulas + logger.debug(f"start expanding formulas: {selected_formulas}") + + def expand_node(node: Dict[str, Any]) -> str: + logger.debug(f"expanding node: {node}") + s = node["formula_str"] + mapping = node.get("variable_mapping") or {} + + def sub(m: re.Match) -> str: + var = m.group(1) + if var not in mapping: + base_observables.append(var) + return f"`{var}`" # treat as base variable + # raise KeyError(f"Missing mapping for `{var}` in formula: {s!r}") + target = mapping[var] + if not target: # base variable + base_observables.append(var) + return f"`{var}`" + return f"({expand_node(target)})" + + return VAR.sub(sub, s) + + def expand_name(name: str) -> str: + if name in cache: + # If it has already been expanded, use the cached version + return cache[name] + if name in stack: + # Don't allow circular formulas! + raise ValueError(f"Cycle detected while expanding {name!r}") + if name not in formulas: + raise KeyError(f"Unknown formula name: {name!r}") + + stack.add(name) + out = expand_node(selected_formulas[name]) + stack.remove(name) + cache[name] = out + + return out + + forward = {} + + for name in selected_formulas: + if not selected_formulas[name]["formula_str"]: + # if no formula_str treat it as not a formula and map to itself. + forward[name] = name + continue + expanded_name = expand_name(name) + forward[name] = expanded_name + + for output_name in output_names: + if output_name not in formulas: + # it is not a formula, should map to itself + forward[output_name] = output_name + + for obs_name in base_observables: + if obs_name not in forward: + forward[obs_name] = obs_name + + reverse: Dict[str, str] = {} + for name, expanded in forward.items(): + reverse.setdefault(expanded, name) + + return forward, reverse, base_observables + + +def stat_key_from_expr(expr: str) -> str: + # Extract statistic key from an expression like + # "std(`PV1`)/mean(`PV1`)" or "percentile(`PV1`, 90)" + + # use string parsing to find what the stat function is + s = re.sub(r"\s+", "", expr) + + ident = r"`[^`]+`" # content with atleast 1 character inside backticks + + if re.fullmatch(rf"std\({ident}\)/mean\({ident}\)", s): + return "std_rel" + if re.fullmatch(rf"mean\({ident}\)", s): + return "mean" + if re.fullmatch(rf"std\({ident}\)", s): + return "std" + + m = re.fullmatch(rf"percentile\({ident},(\d+)\)", s) + if m: + p = int(m.group(1)) + return "median" if p == 50 else f"p{p}" + + raise ValueError(f"Unrecognized expression: {expr!r}") diff --git a/src/badger/gui/components/editable_table_2.py b/src/badger/gui/components/editable_table_2.py new file mode 100644 index 00000000..1018bb5b --- /dev/null +++ b/src/badger/gui/components/editable_table_2.py @@ -0,0 +1,973 @@ +from dataclasses import dataclass, field +from PyQt5.QtWidgets import ( + QWidget, + QHBoxLayout, + QVBoxLayout, + QScrollArea, + QCheckBox, + QLabel, + QComboBox, + QSizePolicy, + QFrame, + QLineEdit, + QMessageBox, +) +from PyQt5.QtCore import pyqtSignal +from typing import Any +import re +from enum import Enum, auto + +import logging + +logger = logging.getLogger(__name__) + + +class FormulaNameLabel(QLineEdit): + """Custom QLineEdit that emits a signal on double-click.""" + + double_clicked = pyqtSignal() + + def mouseDoubleClickEvent(self, event): + """Override to emit signal on double-click.""" + self.double_clicked.emit() + + # propagating signal can lead to unexpected UI behavior + # super().mouseDoubleClickEvent(event) + + +class Origin(Enum): + """ + Enum class for identifying observable origin as + defined in environment or user added. + + This makes it easy to differentiate behavior if desired, + for example making user-added observables editable + """ + + ENVIRONMENT = auto() + USER = auto() + + +@dataclass +class ObservableItem: + """ + Represents a Badger Observable item row in the table. + + An ObservableItem encapsulates the data displayed for each observable + (subclassed as objective and constraint) on the GUI, including + its checked state, name, origin, optional formula for calculations, and optional + statistics operator. + + Attributes: + checked (bool): Whether the item is currently selected. + name (str): The display name of the observable item. + origin (Origin): The source of the observable item, either ENVIRONMENT or USER. + Defaults to Origin.ENVIRONMENT. + formula (dict[str, Any]): A dictionary containing formula specifications with: + - "formula_str" (str): The formula expression as a string + - "variable_mapping" (dict): Mapping of variables used in the formula + Defaults to an empty formula with no variables. + stat (str): The statistical method to apply (e.g., "mean", "sum", "max") to + apply if the observable data is expected as an array. + Defaults to "none" indicating no aggregation. + """ + + checked: bool + name: str + + # identify whether defined in environment or by user + origin: Origin = Origin.ENVIRONMENT # default to environment + + formula: dict[str, Any] = field( + default_factory=lambda: {"formula_str": "", "variable_mapping": {}} + ) + stat: str = "none" + + +@dataclass +class ObjectiveItem(ObservableItem): + """ + Dataclass to represent an objective row. + + Inherits from ObservableItem. + + Attributes: + - rule (str): MINIMIZE or MAXIMIZE, direction of optimization + + """ + + rule: str = "MINIMIZE" + + +@dataclass +class ConstraintItem(ObservableItem): + """ + Dataclass to represent a constraint row. + + Inherits from ObservableItem + + Attributes: + - relation (str): option from ["<",">"] + - threshold (float): value for comparison + - ciritical (bool): whether to mark this as a critical constraint + + """ + + relation: str = "<" + threshold: float = 0.0 + critical: bool = False + + +class ObjectiveRowWidget(QWidget): + """ + A widget for representing a single objective row with checkbox, name, and rule combobox. + + The ObjectiveRowWidget is the UI representaiton of the objective table rows. It uses + an ObjectiveItem self.item to store the data, updates the ObjetiveItem when the user + interacts with the UI, and updates the UI when the ObjectiveItem is modified. + + """ + + item_renamed = pyqtSignal(str, str) # Emits (old_name, new_name) + formula_updated = pyqtSignal( + str, ObservableItem + ) # Emits (new_formula_str, ObservableItem) + formula_double_clicked = pyqtSignal( + QWidget + ) # Emitted when formula name is double-clicked + obs_double_clicked = pyqtSignal(QWidget) + + def __init__(self, objective_item: ObjectiveItem, row_index: int, parent=None): + super().__init__(parent) + self.item = objective_item # ObjectiveItem containing + self.row_index = row_index + # disable stats by default for better backwards compatibility + self.stats_enabled: bool = False + + self._init_ui() + self._connect_signals() + self._apply_style() + + def _init_ui(self): + """Initialize the UI components.""" + layout = QHBoxLayout(self) + layout.setContentsMargins(4, 1, 1, 4) + layout.setSpacing(1) + self.setStyleSheet("border-radius: 0px;") + + # Checkbox (for selecting objective) + self.checkbox = QCheckBox() + self.checkbox.setChecked(self.item.checked) + self.checkbox.setFixedWidth(20) + layout.addWidget(self.checkbox) + + # Name input field + self.name_input = FormulaNameLabel(self.item.name) + self.name_input.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred) + self.name_input.setCursorPosition(0) + self.name_input.setReadOnly(True) + layout.addWidget(self.name_input) + + # Rule combobox + self.rule_combo = QComboBox() + self.rule_combo.addItems(["MINIMIZE", "MAXIMIZE"]) + self.rule_combo.setCurrentText(self.item.rule) + self.rule_combo.setFixedWidth(120) + layout.addWidget(self.rule_combo) + + # statistic combobox + # self.stat_combo = QComboBox() + # self.stat_combo.addItems( + # ["none","mean","std","std_rel","p80","p75","median","p25"] + # ) + # self.stat_combo.setCurrentText(self.item.stat) + # self.stat_combo.setFixedWidth(120) + # layout.addWidget(self.stat_combo) # Don't add stat combo to GUI + + def _apply_style(self): + """ + Apply alternating row colors to the table background. + + Also sets stylesheet for the QLineEdit depending on + whether the item is a formula, user-added observable, + or environment observable. + """ + # set background row color + if self.row_index % 2 == 0: + self.setStyleSheet("alternate-background-color: #262E38;") + else: + self.setStyleSheet("background-color: #262E38;") + + # set QLineEdit stylesheet + if self.item.formula["formula_str"]: + # styling for formula items + self.name_input.setStyleSheet(""" + QLineEdit { + color: LightSeaGreen; + border: 1px solid transparent; + } + QLineEdit:hover { + border: 1px solid DarkCyan; + } + """) + elif not self.stats_enabled: + # if stats disabled, don't show indicator on mouse hover + self.name_input.setStyleSheet(""" + QLineEdit { + color: lightGray; + border: 1px solid transparent; + } + """) + if self.item.origin == Origin.USER: + self.name_input.setStyleSheet(""" + QLineEdit { + color: darkGray; + border: 1px solid transparent; + } + """) + else: + # Styling for non-formula observables + self.name_input.setStyleSheet(""" + QLineEdit { + color: lightGray; + border: 1px solid transparent; + } + QLineEdit:hover { + border: 1px solid Gray; + } + """) + if self.item.origin == Origin.USER: + self.name_input.setStyleSheet(""" + QLineEdit { + color: darkGray; + border: 1px solid transparent; + } + QLineEdit:hover { + border: 1px solid Gray; + } + """) + + self.update_tooltip() + + def update_tooltip(self): + """Update the tooltip to display the current formula or statistic""" + if self.item.formula["formula_str"]: + self.name_input.setToolTip(f"Formula: {self.item.formula['formula_str']}") + elif self.item.stat and self.stats_enabled: + self.name_input.setToolTip(f"Statistic: {self.item.stat}") + + def enable_statistics(self): + """ + Call this function to enable statistics. + + This feature is disabled by default. + If enabled supports defining 'plain' observables in the environment, which are + expected to be either arrays or single datapoints, and choosing how + that data is processed in the GUI. + """ + self.stats_enabled = True + self._apply_style() + + def _connect_signals(self): + """Connect UI signals to data updates.""" + self.checkbox.stateChanged.connect(self._on_checkbox_changed) + self.rule_combo.currentTextChanged.connect(self._on_rule_changed) + # self.stat_combo.currentTextChanged.connect(self._on_stat_changed) + # currently name_input is disabled, name is edited from edit popup + self.name_input.returnPressed.connect(self._on_name_changed) + self.name_input.editingFinished.connect(self._on_name_changed) # on focus loss + + if self.item.formula["formula_str"]: + self.name_input.double_clicked.connect( + lambda: self.formula_double_clicked.emit(self) + ) + else: # if self.item.origin == Origin.USER: + self.name_input.double_clicked.connect( + lambda: self.obs_double_clicked.emit(self) + ) + + def _on_checkbox_changed(self): + """Update item when checkbox state changes.""" + self.item.checked = self.checkbox.isChecked() + + def _on_rule_changed(self): + """Update item when rule selection changes.""" + self.item.rule = self.rule_combo.currentText() + + def update_item_name(self, new_name: str): + """ + Update the name of the item. This is called to update the + name externally (e.g. from FormulaEditor rather than name_input QLineEdit). + emits item_renamed signal so that the parent can update references. + """ + old_name = self.item.name + self.item.name = new_name + self.name_input.setText(new_name) + self.item_renamed.emit(old_name, new_name) + + def _on_name_changed(self): + """ + Update item when name text changes. This is called from the + item_name QLineEdit. Updates the item name with the current text + and emits item_renamed signal so that the parent can update references. + """ + old_name = self.item.name + new_name = self.name_input.text() + + if old_name != new_name: + self.item.name = new_name + self.item_renamed.emit(old_name, new_name) + + +class ObjectiveInsertRowWidget(QWidget): + """A custom widget for inserting new objectives with a QLineEdit for the name.""" + + item_requested = pyqtSignal(str) # Emits the name when user presses Enter + + def __init__(self, parent=None): + super().__init__(parent) + self._init_ui() + self._connect_signals() + + def _init_ui(self): + """Initialize the UI components.""" + layout = QHBoxLayout(self) + layout.setContentsMargins(1, 1, 1, 1) + layout.setSpacing(0) + + # Checkbox (disabled, just for alignment) + checkbox_spacer = QLabel("") + # self.checkbox.setEnabled(False) + checkbox_spacer.setFixedWidth(20) + checkbox_spacer.setStyleSheet(""" + border-radius: 0px; + """) + layout.addWidget(checkbox_spacer) + + # Name input field + self.name_input = QLineEdit() + self.name_input.setPlaceholderText("Enter new objective name...") + self.name_input.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred) + self.name_input.setStyleSheet(""" + border-radius: 0px; + """) + layout.addWidget(self.name_input) + + # Spacer to align with rule column + spacer = QLabel("") + spacer.setFixedWidth(120) + spacer.setStyleSheet(""" + border-radius: 0px; + """) + layout.addWidget(spacer) + + # Style + self.setStyleSheet(""" + background-color: #303A45; + border-top: 1px solid #455364; + """) + + def _connect_signals(self): + """Connect signals for inserting new items.""" + self.name_input.returnPressed.connect(self._on_return_pressed) + + def _on_return_pressed(self): + """Handle Enter key press.""" + name = self.name_input.text().strip() + if name: + if "`" in name: + self.show_name_warning() + return + self.item_requested.emit(name) + self.name_input.clear() + self.name_input.setFocus() + + def show_name_warning(self) -> None: + QMessageBox.warning( + self, + "Use 'add formula' button to add equations!", + "Use 'add formula' button to add equations!", + ) + + +class HeaderWidget(QWidget): + """A custom widget representing the table header.""" + + def __init__(self, parent=None, additional_columns: list[str] = []): + super().__init__(parent) + self._init_ui(additional_columns) + + def _init_ui(self, additional_columns: list[str]) -> None: + """Initialize the header UI.""" + layout = QHBoxLayout(self) + layout.setContentsMargins(1, 1, 1, 4) + layout.setSpacing(1) + + # Checkbox column header (empty or with a master checkbox) + checkbox_header = QLabel("") + checkbox_header.setFixedWidth(20) + layout.addWidget(checkbox_header) + + # Name column header + name_header = QLabel("Name") + name_header.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred) + # name_header.setStyleSheet("font-weight: bold;") + layout.addWidget(name_header) + + # Add additional column headers + for column_name in additional_columns: + col_header = QLabel(column_name) + col_header.setFixedWidth(120) + layout.addWidget(col_header) + + # Style the header + self.setStyleSheet(""" + background-color: #455364; + border-bottom: 1px solid #a0a0a0; + border-radius: 0px; + """) + self.setFixedHeight(30) + + +class ObjectivesListView(QScrollArea): + """ + A scrollable list view for displaying objectives as row widgets, with filtering support. + + """ + + data_changed = pyqtSignal() # Signal to indicate that data has changed + formula_double_clicked = pyqtSignal(ObjectiveRowWidget) + obs_double_clicked = pyqtSignal(ObjectiveRowWidget) + + def __init__(self, parent=None): + super().__init__(parent) + self.setWidgetResizable(True) + self.setFrameShape(QFrame.StyledPanel) + self.setMinimumHeight(200) + + # Main container widget + main_container = QWidget() + main_layout = QVBoxLayout(main_container) + main_layout.setContentsMargins(0, 0, 0, 0) + main_layout.setSpacing(0) + + # Add header + self.header = HeaderWidget(additional_columns=["Rule"]) # , "Statistic"]) + main_layout.addWidget(self.header) + + # Container widget to hold all row widgets + self.container = QWidget() + self.container_layout = QVBoxLayout(self.container) + self.container_layout.setContentsMargins(0, 0, 0, 0) + self.container_layout.setSpacing(0) + + main_layout.addWidget(self.container) + + self.setWidget(main_container) + + # Apply table-like styling + self.setStyleSheet(""" + QScrollArea { + border: 1px solid #455364; + } + """) + + # Store all items + self._all_items: list[ObjectiveItem] = [] + + # Store currently displayed row widgets + self.row_widgets: list[ObjectiveRowWidget] = [] + + # Filter state + self._filter_text = "" + self._show_checked_only = False + + # Create insert row widget (will be added to layout in _rebuild_view) + self.insert_row = ObjectiveInsertRowWidget(parent=self) + self.insert_row.item_requested.connect(self.add_item) + self.insert_row.hide() + + # flag for enable/disable stats depending on environment specifications + self.stats_enabled = False + + def enable_statistics(self): + """ + flag to enable stats function (remains disabled by default + unless called) + """ + self.stats_enabled = True + + def update_items( + self, + objectives: list[dict[str, Any]], + status: dict[str, bool], + formulas: dict[str, dict[str, Any]] | None = None, + vocs_signal: bool = False, + env_observables: list[str] = [], + ) -> None: + """ + Repopulate the list with new objectives data. + This will clear the table and update all items + to match the specified objectives and formulas. + + Parameters + ---------- + objectives : list + list of objecitves specified as a dictionary with expected keys + 'name' and 'rule_list'. "name" is expected as a string, "rule_list' + is expected as a list with the first item the 'rule' ("MINIMIZE" or + "MAXIMIZE") and the second (optional) index a statistic. + status : dict + Status information for each objective, whether it is selected or not + formulas : dict: + Dictionary of formulas + vocs_signal : bool + Whether to emit a signal (not used currently) + env_observables : List[str] + List of which observables are defined in the environment, to differentiate + user-added observables + """ + # Clear all items + self._all_items.clear() + + # Create new items from objectives + for objective in objectives: + for name, rule_list in objective.items(): + rule = rule_list[0] if rule_list else "MINIMIZE" + stat = rule_list[1] if len(rule_list) > 1 else "none" + item = ObjectiveItem( + checked=status.get(name, False), + name=name, + rule=rule, + stat=stat, + origin=Origin.USER + if env_observables and name not in env_observables + else Origin.ENVIRONMENT, + ) + self._all_items.append(item) + + # Update formulas if provided + if formulas is not None: + for name in formulas: + if name in self.item_names: + self.items[name].formula = formulas[name] + if "variable_mapping" not in self.items[name].formula: + self.items[name].formula["variable_mapping"] = {} + + else: + # If formula name is not in items + logger.warning("Name not found in items") + + # Rebuild view with current filters + self._rebuild_view() + + if vocs_signal: + self.update_vocs() + + def _rebuild_view(self) -> None: + """Rebuild the visible row widgets based on current filters.""" + # Clear existing displayed widgets + for widget in self.row_widgets: + widget.item_renamed.disconnect() + widget.formula_double_clicked.disconnect() + try: + widget.obs_double_clicked.disconnect() + except TypeError: # no signals connected + pass + widget.deleteLater() + self.row_widgets.clear() + + # Remove old layout items, skip filtering for + while self.container_layout.count() > 0: + item = self.container_layout.takeAt(self.container_layout.count() - 1) + if item.spacerItem(): + pass + elif item.widget() == self.insert_row: + pass + + all_names = self.item_names + sorted_names = sorted(all_names) + items = self.items + + # Filter and display items + row_index = 0 + for name in sorted_names: + item = items[name] + if self._passes_filters(item): + row_widget = ObjectiveRowWidget(item, row_index, parent=self) + row_widget.item_renamed.connect(self._on_item_renamed) + row_widget.formula_updated.connect( + lambda formula_str, i=item: self.update_item_formula(i, formula_str) + ) + row_widget.formula_double_clicked.connect( + self.formula_double_clicked.emit + ) + # if stats not enabled, don't connect to enable editing stats + if self.stats_enabled: + row_widget.obs_double_clicked.connect(self.obs_double_clicked.emit) + row_widget.enable_statistics() + self.row_widgets.append(row_widget) + self.container_layout.addWidget(row_widget) + row_index += 1 + + # Show and add insert row + self.insert_row.show() + self.container_layout.addWidget(self.insert_row) + + # Add stretch at the end to push items to top + self.container_layout.addStretch() + + def _passes_filters(self, item: ObjectiveItem) -> bool: + """Check if an item passes all active filters.""" + # Filter by text (case-insensitive substring match) + if self._filter_text: + if self._filter_text.lower() not in item.name.lower(): + return False + + # Filter by checked status + if self._show_checked_only and not item.checked: + return False + + return True + + def set_filter(self, text: str) -> None: + """Filter items by name substring (case-insensitive). + + Parameters + ---------- + text : str + Substring to filter by. Empty string clears the filter. + """ + self._filter_text = text + self._rebuild_view() + + def show_checked_only(self, enabled: bool) -> None: + """Show only checked items. + + Parameters + ---------- + enabled : bool + If True, only checked items will be displayed. + """ + self._show_checked_only = enabled + self._rebuild_view() + + def get_items(self) -> list[ObjectiveItem]: + """Get the unfiltered list of all ObjectiveItem objects. + + Returns + ------- + list[ObjectiveItem] + All items, regardless of current filter state. + """ + return self._all_items + + def get_selected_items(self) -> list[ObjectiveItem]: + """Get the list of currently selected (checked) items. + + Returns + ------- + list[ObjectiveItem] + List of checked items. + """ + return [item for item in self._all_items if item.checked] + + def export_data(self): + selected_items = self.get_selected_items() + logger.info( + f"Exporting data for selected items from table: {[item.name for item in selected_items]}" + ) + return [ + {item.name: {"rule": item.rule, "stat": item.stat}} + for item in selected_items + ] + + @property + def item_names(self) -> list[str]: + """Get a list of all item names.""" + # print(f"item names: {[item.name for item in self._all_items]}") + return [item.name for item in self._all_items] + + @property + def items(self) -> dict[str, ObjectiveItem]: + """Get a dictionary of items keyed by name.""" + # print("items property called") + return {item.name: item for item in self._all_items} + + @property + def formulas(self) -> dict: + _formulas = {} + for item in self._all_items: + if ( + item.formula["formula_str"] or item.origin == Origin.USER + ): # not sure if this is a good approach + _formulas[item.name] = item.formula + + return _formulas + + @property + def formula_strs(self) -> list[str]: + return [ + item.formula["formula_str"] + for item in self._all_items + if item.formula["formula_str"] + ] + + def show_duplicate_warning(self, name: str) -> None: + QMessageBox.warning( + self, + "Item already exists!", + f"Item {name} already exists!", + ) + + def add_item( + self, name: str, formula_str: str = None, checked: bool = False + ) -> None: + """ + Add a new observable item to the list. This is called + either by formula dialog (add formula, with formula_str) + or by adding a new observable from the new item line. + + Parameters + ---------- + name : str + The name of the new observable. + formula_str: str + (optional) formula string + checked: bool + Whether the item should be selected + """ + if name in self.item_names: + # If an item with the same name already exists, show a warning and do not add + self.show_duplicate_warning(name) + return + + # check for duplicate names as formulas + if name in self.formula_strs: + self.show_duplicate_warning(name) + return + + # check for duplicate formulas + if formula_str in self.formula_strs: + self.show_duplicate_warning(formula_str) + return + + # Create new objective item with default rule and formula + new_item = ObjectiveItem( + checked=checked, + name=name, + rule="MINIMIZE", + origin=Origin.USER, + ) + + if formula_str: + # add item as formula + new_item.formula["formula_str"] = formula_str + new_item.formula["variable_mapping"] = self.get_variable_mapping( + formula_str + ) + else: + # Item added from UI line as new observable with no formula. + if "`" in name: + new_item.formula["formula_str"] = name + new_item.formula["variable_mapping"] = self.get_variable_mapping(name) + + logger.debug(f"end add_item: {new_item.name}, {new_item.formula}") + + # Add to items list + self._all_items.append(new_item) + + # Rebuild view to display new item + self._rebuild_view() + + # self.update_vocs() + + def default_objective(self): + return ["MINIMIZE"] + + def get_variable_mapping(self, formula_str: str) -> dict[str, str]: + matches = self.check_for_var_references(formula_str) # find variable references + # matches is a list of variable name strings + visited = set() + + variable_mapping = {} + for match in matches: + if match in visited: + variable_mapping[match] = None + continue + + item = self.items[match] + print(f"match: {item.name}, {item.formula}") + if item.formula["formula_str"]: + # If item is formula, get full formula with variable mapping for later expansion + variable_mapping[item.name] = item.formula + else: + # If item is not formula, map to var name + variable_mapping[item.name] = item.formula["formula_str"] + + visited.add(match) + + logger.info(f"Return variable mapping: {variable_mapping}") + return variable_mapping + + def check_for_var_references(self, expr: str) -> list[str]: + """ + Find references to other variables in formula string + + Parameters + ---------- + - expr: str + string formula expression, with variables backticked + + Returns + ------- + - List[str] of referenced variables + """ + + if not self.item_names: + return [] + + alts = "|".join(map(re.escape, self.item_names)) + + # Allowed token separators + left_sep = r"[()\+\-*/\s`]" + right_sep = r"[()\+\-*/\s`]" + + # regex pattern to look for matches in formula string + # formula names will be matched unless they are + # immediately before a "." or "(", so you + # can name a func "mean" and still do + # mean(`f`) without matching mean + # this will match substrings which are: + pat = re.compile( + rf""" + (?:(?<=^)|(?<={left_sep})) # start OR preceded by left_sep + (?:{alts}) # match name + (?![.(]) # not followed by . or ( + (?=$|{right_sep}) # end OR right_sep or "**" after + """, + re.VERBOSE, + ) + + # (? None: + """ + Handle renaming of an item and update references in other items. + + Parameters + ---------- + old_name : str + The previous name of the item + new_name : str + The new name of the item + """ + # if new_name in self.item_names: + # self.show_duplicate_warning(new_name) + # return + + # Update formula_str in all items that reference the old name + for item in self._all_items: + if item.formula["formula_str"]: + # Only match variable names within backticks + old_pattern = rf"`{re.escape(old_name)}`" + item.formula["formula_str"] = re.sub( + old_pattern, f"`{new_name}`", item.formula["formula_str"] + ) + + # Update variable_mapping keys + if old_name in item.formula["variable_mapping"]: + item.formula["variable_mapping"][new_name] = item.formula[ + "variable_mapping" + ].pop(old_name) + + # Update tooltips for all visible row widgets + for row_widget in self.row_widgets: + row_widget.update_tooltip() + + def update_item_formula(self, item: ObjectiveItem, new_formula_str: str) -> None: + """ + Handle updating an item's formula and recalculate its variable mapping. + + Parameters + ---------- + item : ObjectiveItem + The item whose formula is being updated + new_formula_str : str + The new formula string + """ + + # self.check_for_circular_reference(item, new_formula_str) + + # check for duplicate formulas + if new_formula_str in self.formula_strs: + self.show_duplicate_warning(new_formula_str) + return + + item.formula["formula_str"] = new_formula_str + # Recalculate variable mapping for the updated formula + mapping = self.get_variable_mapping(new_formula_str) + item.formula["variable_mapping"] = mapping + + # update variable_mapping for any items that reference this item as a variable + for other_item in self._all_items: + if other_item.formula["formula_str"]: + # replace old variable mapping with new formula for this item + if item.name in other_item.formula["variable_mapping"]: + other_item.formula["variable_mapping"][item.name] = item.formula + + # Update tooltips for all visible row widgets + for row_widget in self.row_widgets: + row_widget.update_tooltip() + + def _on_item_formula_updated( + self, item: ObjectiveItem, new_formula_str: str + ) -> None: + """ """ + # check for duplicate formulas + if new_formula_str in self.formula_strs: + self.show_duplicate_warning(new_formula_str) + return + + logger.info( + f"Updating formula for item {item.name}, new formula: {new_formula_str}" + ) + item.formula["formula_str"] = new_formula_str + # Recalculate variable mapping for the updated formula + item.formula["variable_mapping"] = self.get_variable_mapping(new_formula_str) + + # update variable_mapping for any items that reference this item as a variable + for other_item in self._all_items: + if other_item.formula["formula_str"]: + # print(f"other item: {other_item.name}, formula: {other_item.formula}") + # replace old variable mapping with new formula for this item + if item.name in other_item.formula["variable_mapping"]: + other_item.formula["variable_mapping"][item.name] = item.formula + + # Update tooltips for all visible row widgets + for row_widget in self.row_widgets: + row_widget.update_tooltip() + + """def check_for_circular_reference( + self, item: ObjectiveItem, new_formula_str: str + ) -> bool: + # Check for circular references + print("check for circular references") + print(self.get_variable_mapping(new_formula_str)) + + def find_refs(var_mapping: dict): + for name, mapping in var_mapping.items(): + print(name, mapping) + if name == item.name: + raise ValueError("Circular reference detected!") + if isinstance(mapping, dict): + find_refs(mapping) + + find_refs(self.get_variable_mapping(new_formula_str))""" + + def update_vocs(self): + logging.debug("Emitting data_changed signal from editable_table") + self.data_changed.emit() diff --git a/src/badger/gui/components/env_cbox.py b/src/badger/gui/components/env_cbox.py index 944f272d..09b2771e 100644 --- a/src/badger/gui/components/env_cbox.py +++ b/src/badger/gui/components/env_cbox.py @@ -26,6 +26,11 @@ from badger.gui.components.con_table import ConstraintTable from badger.gui.components.obs_table import ObservableTable from badger.gui.components.data_table import init_data_table +from badger.gui.windows.formula_dialog import ( + BadgerFormulaDialog, + FormulaEdit, + ObservableEdit, +) from badger.settings import init_settings from badger.gui.utils import ( MouseWheelWidgetAdjustmentGuard, @@ -34,6 +39,10 @@ from badger.utils import strtobool from xopt.vocs import VOCS, ConstraintEnum import logging +from badger.gui.components.editable_table_2 import ( + ObjectivesListView, + ObjectiveRowWidget, +) LABEL_WIDTH = 96 ENV_PARAMS_BTN = 1 # use button or collapsible box for env parameters @@ -409,14 +418,22 @@ def init_ui(self): self.edit_obj = edit_obj = QLineEdit() edit_obj.setPlaceholderText("Filter objectives...") edit_obj.setFixedWidth(192) + self.add_formula_obj = add_formula_obj = QPushButton("Add Formula") + self.check_only_obj = check_only_obj = QCheckBox("Show Checked Only") check_only_obj.setChecked(False) hbox_action_obj.addWidget(edit_obj) hbox_action_obj.addStretch() + hbox_action_obj.addWidget(add_formula_obj) hbox_action_obj.addWidget(check_only_obj) self.obj_table = ObjectiveTable() - vbox_obj_edit.addWidget(self.obj_table) + # vbox_obj_edit.addWidget(self.obj_table) + + # New objectives list view with row widgets + self.objectives_list_view = ObjectivesListView() + vbox_obj_edit.addWidget(self.objectives_list_view) + hbox_obj.addWidget(edit_obj_col) cbox_more = CollapsibleBox(self, " More") @@ -503,6 +520,7 @@ def config_logic(self): self.edit_var.textChanged.connect(self.filter_var) self.check_only_var.stateChanged.connect(self.toggle_var_show_mode) self.edit_obj.textChanged.connect(self.filter_obj) + self.add_formula_obj.clicked.connect(self.add_formula) self.check_only_obj.stateChanged.connect(self.toggle_obj_show_mode) self.edit_con.textChanged.connect(self.filter_con) self.check_only_con.stateChanged.connect(self.toggle_con_show_mode) @@ -516,6 +534,12 @@ def config_logic(self): self.sta_table.data_changed.connect(lambda: self.update_vocs("sta_table")) self.var_table.data_changed.connect(lambda: self.update_vocs("var_table")) + self.objectives_list_view.data_changed.connect( + lambda: self.update_vocs("objectives_list_view") + ) + self.objectives_list_view.formula_double_clicked.connect(self.edit_formula) + self.objectives_list_view.obs_double_clicked.connect(self.edit_obs) + def update_vocs(self, origin: str): logger.debug(f"Emitting vocs_updated signal from env_cbox: {origin}") vocs, _ = self.compose_vocs() @@ -559,9 +583,13 @@ def filter_var(self): def toggle_obj_show_mode(self, _): self.obj_table.update_show_selected_only(self.check_only_obj.isChecked()) + self.objectives_list_view.show_checked_only( + self.check_only_obj.isChecked() + ) # new table def filter_obj(self): self.obj_table.update_keyword(self.edit_obj.text()) + self.objectives_list_view.set_filter(self.edit_obj.text()) # new table def toggle_con_show_mode(self, _): self.con_table.update_show_selected_only(self.check_only_con.isChecked()) @@ -607,15 +635,72 @@ def update_stylesheets(self, environment=""): stylesheet = "" self.setStyleSheet(stylesheet) + def add_formula(self): + dlg = BadgerFormulaDialog( + parent=self, + table=self.objectives_list_view, + ) + self.tc_dialog = dlg + try: + dlg.exec() + finally: + self.tc_dialog = None + + def edit_formula(self, row_widget: ObjectiveRowWidget): + dlg = FormulaEdit( + parent=self, + table=self.objectives_list_view, + row_widget=row_widget, + ) + self.tc_dialog = dlg + try: + dlg.exec() + finally: + self.tc_dialog = None + + def edit_obs(self, row_widget: ObjectiveRowWidget): + dlg = ObservableEdit( + parent=self, + table=self.objectives_list_view, + row_widget=row_widget, + ) + self.tc_dialog = dlg + try: + dlg.exec() + finally: + self.tc_dialog = None + + def _construct_obs_func_str(self, operation: str, obj_name: str): + stats_mapping = { + "mean": lambda x: f"mean(`{x}`)", + "std": lambda x: f"std(`{x}`)", + "p80": lambda x: f"percentile(`{x}`, 80)", + "p75": lambda x: f"percentile(`{x}`, 75)", + "p50": lambda x: f"percentile(`{x}`, 50)", + "median": lambda x: f"percentile(`{x}`, 50)", + "p25": lambda x: f"percentile(`{x}`, 25)", + "std_rel": lambda x: f"std(`{x}`)/mean(`{x}`)", + } + + if operation in stats_mapping: + new_obj_name = stats_mapping[operation](obj_name) + return new_obj_name + def compose_vocs(self) -> tuple[VOCS, list[str]]: # Compose the VOCS settings variables = self.var_table.export_variables() objectives: dict[str, Any] = {} - for objective in self.obj_table.export_data(): + for objective in self.objectives_list_view.export_data(): obj_name = next(iter(objective)) - (rule,) = objective[obj_name] - objectives[obj_name] = rule + + rule = objective[obj_name]["rule"] + stat = objective[obj_name]["stat"] + + if stat not in ["none", "formula"]: + obj_name = self._construct_obs_func_str(stat, obj_name) + + objectives[obj_name] = rule # [0] constraints: dict[str, list[float | ConstraintEnum]] = {} critical_constraints: list[str] = [] diff --git a/src/badger/gui/components/routine_page.py b/src/badger/gui/components/routine_page.py index 07e94bfe..bd7f6a26 100644 --- a/src/badger/gui/components/routine_page.py +++ b/src/badger/gui/components/routine_page.py @@ -54,6 +54,8 @@ ) from badger.factory import list_generators, list_env, get_env from badger.routine import Routine +from badger.formula import extract_variable_keys +from badger.formula_utils import stat_key_from_expr from badger.settings import init_settings from datetime import datetime from badger.utils import ( @@ -483,9 +485,11 @@ def set_options_from_template(self, template_dict: dict[str, Any]): formulas = {} objectives = [] status = {} - objectives_names_full = self.configs["observations"] + list(formulas.keys()) + objectives_names_full = list(self.configs["observations"]) + list( + formulas.keys() + ) for name in objectives_names_full: - obj = {name: ["MINIMIZE"]} + obj = {name: self.env_box.objectives_list_view.default_info()} status[name] = False # selected objectives.append(obj) for name, val in vocs.objectives.items(): @@ -504,9 +508,9 @@ def set_options_from_template(self, template_dict: dict[str, Any]): self.env_box.check_only_obj.blockSignals(True) self.env_box.check_only_obj.setChecked(True) self.env_box.check_only_obj.blockSignals(False) - self.env_box.obj_table.show_selected_only = True + self.env_box.objectives_list_view.show_selected_only = True - self.env_box.obj_table.update_items(objectives, status, formulas) + self.env_box.objectives_list_view.update_items(objectives, status, formulas) # set constraints # Initialize the constraints table with env observables @@ -516,9 +520,11 @@ def set_options_from_template(self, template_dict: dict[str, Any]): formulas = {} constraints = [] status = {} - constraints_names_full = self.configs["observations"] + list(formulas.keys()) + constraints_names_full = list(self.configs["observations"]) + list( + formulas.keys() + ) for name in constraints_names_full: - cons = {name: ["<", 0.0, False]} + cons = {name: self.env_box.con_table.default_info()} status[name] = False # selected constraints.append(cons) for name, val in vocs.constraints.items(): @@ -556,7 +562,7 @@ def set_options_from_template(self, template_dict: dict[str, Any]): observables = [] status = {} observables_names_full = ( - var_names + self.configs["observations"] + list(formulas.keys()) + var_names + list(self.configs["observations"]) + list(formulas.keys()) ) for name in observables_names_full: obs = {name: []} @@ -611,7 +617,7 @@ def generate_template_dict_from_gui(self): "vrange_limit_options": self.ratio_var_ranges, "vrange_hard_limit": self.var_hard_limit, "additional_variables": self.env_box.var_table.addtl_vars, - "formulas": self.env_box.obj_table.formulas, + "formulas": self.env_box.objectives_list_view.formulas, "constraint_formulas": self.env_box.con_table.formulas, "observable_formulas": self.env_box.sta_table.formulas, "initial_point_actions": self.init_table_actions, @@ -833,34 +839,84 @@ def refresh_ui(self, routine: Routine | None = None, silent: bool = False): objectives = [] status = {} - objectives_names_full = self.configs["observations"] + list(formulas.keys()) + objectives_names_full = list( + set( + list(self.configs["observations"]) # env observables + + list(formulas.keys()) # routine formulas + ) + ) + # print(f"refresh_ui: full_objectives: {objectives_names_full}") + # print(f"refresh_ui: configs['observations']: {self.configs['observations']}") + # print(f"refresh_ui: vocs objectives: {list(routine.vocs.objectives.keys())}") for name in objectives_names_full: - obj = {name: ["MINIMIZE"]} - status[name] = False # selected - objectives.append(obj) + try: + # get defaults from env configs + rule = self.configs["observations"][name]["defaults"].get( + "rule", "MINIMIZE" + ) + stat = self.configs["observations"][name]["defaults"].get( + "stat", "none" + ) + except (KeyError, TypeError): + # TypeError: self.configs["observations"] is a list + # KeyError: env observations does not specify default values + + # Either way no default values are specified in env configs, + # set to general defaults + rule, stat = "MINIMIZE", "none" + + status[name] = False # start with not selected + objectives.append({name: [rule, stat]}) # add objective with default values + for name, val in routine.vocs.objectives.items(): + # iterate through objectives defined in routine.vocs + # this will overwrite default values with those used + # in the routine, as well as capture new observables + # which may have been added. + + stat = None # reset per objective rule = val + # unwrap if not known and has backtick + if name not in objectives_names_full and "`" in name: + stat = stat_key_from_expr(name) + name = extract_variable_keys(name)[0] + + # if still new, add it + if name not in objectives_names_full: + objectives_names_full.append(name) + # start by adding defaults + objectives.append({name: ["MINIMIZE", "none"]}) + status[name] = False + idx = objectives_names_full.index(name) + if idx == -1: raise BadgerRoutineError( f"Objective {name} not found in the routine's observables." ) else: - objectives[idx] = {name: [rule]} + if stat is None: + stat = "none" + + objectives[idx] = {name: [rule, stat]} status[name] = True # Show selected objectives only - self.env_box.check_only_obj.blockSignals(True) - self.env_box.check_only_obj.setChecked(True) - self.env_box.check_only_obj.blockSignals(False) self.env_box.edit_obj.blockSignals(True) self.env_box.edit_obj.setText("") self.env_box.edit_obj.blockSignals(False) - self.env_box.obj_table.keyword = "" - self.env_box.obj_table.show_selected_only = True - with BlockSignalsContext(self.env_box.obj_table): - self.env_box.obj_table.update_items(objectives, status, formulas) + # self.env_box.obj_table.keyword = "" + # self.env_box.obj_table.show_selected_only = True + # with BlockSignalsContext(self.env_box.obj_table): + # self.env_box.obj_table.update_items(objectives, status, formulas) + with BlockSignalsContext(self.env_box.objectives_list_view): + self.env_box.objectives_list_view.update_items( + objectives, + status, + formulas, + env_observables=self.configs["observations"], + ) # Initialize the constraints table with env observables try: @@ -869,7 +925,9 @@ def refresh_ui(self, routine: Routine | None = None, silent: bool = False): formulas = {} constraints = [] status = {} - constraints_names_full = self.configs["observations"] + list(formulas.keys()) + constraints_names_full = list(self.configs["observations"]) + list( + formulas.keys() + ) for name in constraints_names_full: cons = {name: ["<", 0.0, False]} status[name] = False # selected @@ -913,7 +971,7 @@ def refresh_ui(self, routine: Routine | None = None, silent: bool = False): observables = [] status = {} observables_names_full = ( - var_names + self.configs["observations"] + list(formulas.keys()) + var_names + list(self.configs["observations"]) + list(formulas.keys()) ) for name in observables_names_full: obs = {name: []} @@ -1130,19 +1188,52 @@ def select_env(self, i: int): self.add_var() ) + if isinstance(self.configs["observations"], dict): + self.env_box.objectives_list_view.enable_statistics() + # Stats is disabled by default for the objective table for now, this will enable it + # If obs is not a dictionary in env, should disable setting statistics + # as old environments will not need this feature, and it may add confusion + objectives = [] status = {} + formulas = {} for name in self.configs["observations"]: - cons = {name: ["MINIMIZE"]} + default_obj = self.env_box.objectives_list_view.default_objective() + # If defaults defined for name in environment use instead + try: + rule = self.configs["observations"][name]["defaults"].get( + "rule", "MINIMIZE" + ) + stat = self.configs["observations"][name]["defaults"].get( + "stat", "none" + ) + default_obj = [rule, stat] + except (KeyError, TypeError): + # Improve + logger.info("no defaults for {name}") + pass + + obj = {name: default_obj} status[name] = False # selected - objectives.append(cons) + + # If configs is a dict, see if there are default formulas + # defined for observables in environemnt + if ( + isinstance(self.configs["observations"], dict) + and "formula" in self.configs["observations"][name] + ): + formulas[name] = self.configs["observations"][name]["formula"] + + objectives.append(obj) self.env_box.check_only_obj.blockSignals(True) self.env_box.check_only_obj.setChecked(False) self.env_box.check_only_obj.blockSignals(False) - self.env_box.obj_table.show_selected_only = False + # self.env_box.obj_table.show_selected_only = False # with BlockSignalsContext(self.env_box.obj_table): - self.env_box.obj_table.update_items( - objectives, status, formulas={}, vocs_signal=False + # self.env_box.obj_table.update_items(objectives, status, vocs_signal=False) + + self.env_box.objectives_list_view.update_items( + objectives, status, formulas, vocs_signal=False ) # Initialize the constraints table with env observables @@ -1169,7 +1260,8 @@ def select_env(self, i: int): var_names = [] # do not show var names in observables until we have a fix to get_observables else: var_names = [] - for name in var_names + self.configs["observations"]: + + for name in var_names + list(self.configs["observations"]): obs = {name: []} status[name] = False # selected observables.append(obs) @@ -1178,9 +1270,7 @@ def select_env(self, i: int): self.env_box.check_only_sta.blockSignals(False) self.env_box.sta_table.show_selected_only = False # with BlockSignalsContext(self.env_box.sta_table): - self.env_box.sta_table.update_items( - observables, status, formulas={}, vocs_signal=False - ) + self.env_box.sta_table.update_items(observables, status, vocs_signal=False) self.env_box.fit_content() # self.routine = None @@ -1844,10 +1934,13 @@ def _compose_routine(self) -> Routine: vrange_hard_limit=vrange_hard_limit, initial_point_actions=initial_point_actions, additional_variables=self.env_box.var_table.addtl_vars, - formulas=self.env_box.obj_table.formulas, + # additonal_observables=self.env_box.objectives_list_view.additional_observables, + formulas=self.env_box.objectives_list_view.formulas, constraint_formulas=self.env_box.con_table.formulas, observable_formulas=self.env_box.sta_table.formulas, ) + logger.info(f"Compose routine: vocs: {routine.vocs}") + logger.info(f"Compose routine: formulas: {routine.formulas}") # Check if any user warnings were caught for warning in caught_warnings: diff --git a/src/badger/gui/windows/formula_dialog.py b/src/badger/gui/windows/formula_dialog.py new file mode 100644 index 00000000..7b7687bd --- /dev/null +++ b/src/badger/gui/windows/formula_dialog.py @@ -0,0 +1,572 @@ +from PyQt5.QtWidgets import ( + QDialog, + QWidget, + QHBoxLayout, + QPushButton, + QVBoxLayout, + QLabel, + QTextEdit, + QLineEdit, + QCompleter, + QComboBox, + QMessageBox, +) +from PyQt5.QtCore import Qt +from PyQt5.QtGui import QTextCursor +import re + +from badger.formula_utils import sanitize_for_validation, validate_formula +from badger.gui.components.editable_table_2 import ( + ObjectivesListView, + ObjectiveRowWidget, + Origin, +) + + +def _surround_with_backticks(string_list, text): + """ + Surrounds any strings from string_list found in text with backticks, + but only if not already backticked. + + Args: + string_list: List of strings to search for + text: The string to search in and modify + + Returns: + Modified text with matching strings surrounded by backticks + """ + result = text + + # Sort by length (longest first) to handle overlapping matches correctly + sorted_list = sorted(string_list, key=len, reverse=True) + + for string in sorted_list: + # Escape the string for use in regex + escaped = re.escape(string) + + # Pattern: match string NOT preceded or followed by backtick + # (? None: + """ + Initialize the user interface. + """ + + self.setWindowTitle("Add formula") + self.setMinimumWidth(370) + + root_vbox = QVBoxLayout(self) + + # Header and labels + header = QWidget() + header_hbox = QHBoxLayout(header) + header_hbox.setContentsMargins(0, 0, 0, 0) + + label = QLabel("test label info would be here") + label.setMinimumWidth(360) + + header_hbox.addWidget(label) + + content_widget = QWidget() + hbox_content = QHBoxLayout(content_widget) + hbox_content.setContentsMargins(0, 0, 0, 0) + + formula_widget = self.build_formula_input() + self.help_widget = self.build_help_widget() + self.help_widget.hide() + + hbox_content.addWidget(formula_widget) + hbox_content.addWidget(self.help_widget) + + # vbox.addWidget(header) + root_vbox.addWidget(content_widget) + # root_vbox.addWidget(button_set) + + def build_formula_input(self) -> QWidget: + formula_widget = QWidget() + formula_layout = QVBoxLayout(formula_widget) + formula_layout.setContentsMargins(0, 0, 0, 0) + formula_widget.setMinimumWidth(360) + + name_widget = QWidget() + name_layout = QVBoxLayout(name_widget) + name_layout.setContentsMargins(0, 0, 0, 0) + + name_header = QWidget() + name_header_layout = QHBoxLayout(name_header) + name_header_layout.setContentsMargins(0, 0, 0, 0) + + name_label = QLabel("Name: ") + + self.info_button = QPushButton("Show Info >") + self.info_button.setCheckable(True) + self.info_button.setFixedWidth(85) + + name_header_layout.addWidget(name_label) + name_header_layout.addStretch() + name_header_layout.addWidget(self.info_button) + + self.name_edit = QLineEdit() + self.name_edit.setPlaceholderText("Enter objective name") + name_layout.addWidget(name_header) + name_layout.addWidget(self.name_edit) + + formula_edit_widget = QWidget() + formula_edit_layout = QVBoxLayout(formula_edit_widget) + formula_edit_layout.setContentsMargins(0, 0, 0, 0) + formula_label = QLabel("Formula: ") + + self.formula_edit = CompleterTextEdit() + self.formula_edit.setPlaceholderText( + "Enter formula, for example mean(`f`) or std(`f`)**2\n" + + "Formula syntax:\n" + + " - Enter variable names in backticks: `f`\n" + + " - Use python.statistics or numpy expressions: \n" + + " - mean(`f`), std(`f`), percentile(`f`, 80)\n" + + " - operators including *, +, -, /, **\n" + ) + self.formula_edit.setStyleSheet(""" + color: darkGray; + """) + completer = QCompleter(self.table.item_names, self.formula_edit) + completer.setCaseSensitivity(Qt.CaseInsensitive) # ignore case + completer.setFilterMode(Qt.MatchStartsWith) # default behavior + + # Button set + button_set = QWidget() + hbox_set = QHBoxLayout(button_set) + hbox_set.setContentsMargins(0, 0, 0, 0) + self.btn_cancel = QPushButton("Cancel") + self.btn_add = QPushButton("Add") + self.btn_cancel.setFixedSize(96, 24) + self.btn_add.setFixedSize(96, 24) + hbox_set.addStretch() + hbox_set.addWidget(self.btn_cancel) + hbox_set.addWidget(self.btn_add) + # hbox_set.addStretch() + + self.formula_edit.setCompleter(completer) + formula_edit_layout.addWidget(formula_label) + formula_edit_layout.addWidget(self.formula_edit) + + formula_layout.addWidget(name_widget) + formula_layout.addWidget(formula_edit_widget) + formula_layout.addWidget(button_set) + + return formula_widget + + def build_help_widget(self) -> QWidget: + help_widget = QWidget() + help_layout = QVBoxLayout(help_widget) + help_layout.setContentsMargins(0, 0, 0, 0) + help_widget.setFixedWidth(230) + help_widget.setStyleSheet(""" + border: 1px solid #455364; + background-color: #37414F; + color: LightGray; + """) + + help_label = QLabel( + "Helpful Formula Info: \n" + " \n" + " - Variable names are backticked:\n" + " `f`, `equation_b`, `PV:NAME` \n" + " \n" + " - Use any expression from numpy,\n" + " python.statistics , or \n" + " python.math such as: \n" + " - mean(`f`), std(`f`) \n" + " - max([`f`, `g`, `h`]) \n" + " - percentile(`f`, 80) \n" + " - percentile(`f`, 50) \n" + " \n" + " - Supported operators: \n" + " +, -, *, /, ** \n" + ) + help_label.setAlignment(Qt.AlignLeft | Qt.AlignTop) + + help_layout.addWidget(help_label) + + return help_widget + + def config_logic(self) -> None: + self.btn_cancel.clicked.connect(self.cancel) + self.btn_add.clicked.connect(self.construct_formula_str) + self.info_button.clicked.connect(self.show_info_panel) + + def show_info_panel(self): + """ + Show/hide side panel with helpful info on writing formulas + """ + old_width = self.width() + if self.info_button.isChecked(): + self.info_button.setText("Hide Info < ") + self.setMinimumWidth(615) + self.resize(old_width + 235, self.height()) + self.help_widget.setVisible(True) + else: + self.info_button.setText("Show Info >") + self.setMinimumWidth(370) + self.resize(old_width - 235, self.height()) + self.help_widget.setVisible(False) + + def construct_formula_str(self): + name = self.name_edit.text() + formula_str = self.formula_edit.toPlainText().strip() + if self.validate_formula(formula_str): # make sure formula is valid + self.table.add_item(name, formula_str, checked=True) + self.close() + else: + print(f"invalid formula: {formula_str}") + + def validate_formula(self, expr: str) -> bool: + """ + Validate the formula expression by sanitizing it and checking if it can be parsed + using allowed symbols + """ + matches = self.table.check_for_var_references(expr) + # I don't think this is still needed, allowed separately + # referencing vars in backticks and formulas without + expr = _surround_with_backticks(matches, expr) + + python_expr, allowed = sanitize_for_validation(expr) + try: + validate_formula(python_expr, allowed_symbols=allowed) + return True + except TypeError: + return False + + def cancel(self): + self.close() + + +class FormulaEdit(BadgerFormulaDialog): + def __init__( + self, + parent: QWidget, + table: ObjectivesListView, + row_widget: ObjectiveRowWidget, + ): + """ + Initialize the dialog. + + """ + super().__init__(parent, table) + if not isinstance(row_widget, ObjectiveRowWidget): + raise ValueError("row_widget must be an instance of ObjectiveRowWidget") + self.row_widget = row_widget + self.item = item = row_widget.item + if item: + self.name_edit.setText(item.name) + self.formula_edit.setText(item.formula["formula_str"]) + + def init_ui(self) -> None: + super().init_ui() + self.setWindowTitle("Edit formula") + self.btn_add.setText("Save") + + def config_logic(self): + super().config_logic() + self.btn_add.clicked.disconnect() + self.btn_add.clicked.connect(self.construct_formula_str) + + def construct_formula_str(self): + name = self.name_edit.text() + formula_str = " ".join( + self.formula_edit.toPlainText().split() + ) # strip newlines, tabs, extra spaces + if self.validate_formula(formula_str): # make sure formula is valid + if formula_str != self.item.formula["formula_str"]: + # only update formula + self.table.update_item_formula(self.row_widget.item, formula_str) + if name != self.item.name: + # only update name + if name in self.items: + QMessageBox.warning( + self, + "Item already exists!", + f"Item {name} already exists!", + ) + return + self.row_widget.update_item_name(name) + self.close() + else: + print(f"invalid formula: {formula_str}") + + +class ObservableEdit(QDialog): + """ + Dialog for adding formula observable in Badger. + """ + + def __init__( + self, + parent: QWidget, + table: ObjectivesListView, + row_widget: ObjectiveRowWidget, + ): + """ + Initialize the dialog. + + """ + super().__init__(parent) + self.setWindowFlags(self.windowFlags() | Qt.WindowMaximizeButtonHint) + + self.table = table + + self.items = self.table.items + + self.row_widget = row_widget + self.item = item = row_widget.item + + self.init_ui() + + if not isinstance(row_widget, ObjectiveRowWidget): + raise ValueError("row_widget must be an instance of ObjectiveRowWidget") + + if item: + self.name_edit.setText(item.name) + if item.origin != Origin.USER: + self.name_edit.setEnabled(False) + + self.config_logic() + + def config_logic(self) -> None: + self.btn_cancel.clicked.connect(self.cancel) + self.btn_add.clicked.connect(self.construct_formula_str) + + def init_ui(self) -> None: + """ + Initialize the user interface. + """ + + self.setWindowTitle("Observable") + self.setMinimumWidth(360) + + root_vbox = QVBoxLayout(self) + + # Header and labels + header = QWidget() + header_hbox = QHBoxLayout(header) + header_hbox.setContentsMargins(0, 0, 0, 0) + + label = QLabel("test label info would be here") + label.setMinimumWidth(360) + header_hbox.addWidget(label) + + content_widget = QWidget() + hbox_content = QHBoxLayout(content_widget) + hbox_content.setContentsMargins(0, 0, 0, 0) + + formula_widget = self.build_formula_input() + hbox_content.addWidget(formula_widget) + root_vbox.addWidget(content_widget) + + def build_formula_input(self) -> QWidget: + formula_widget = QWidget() + formula_layout = QVBoxLayout(formula_widget) + formula_layout.setContentsMargins(0, 0, 0, 0) + formula_widget.setMinimumWidth(320) + + name_widget = QWidget() + name_layout = QVBoxLayout(name_widget) + name_layout.setContentsMargins(0, 0, 0, 0) + + name_header = QWidget() + name_header_layout = QHBoxLayout(name_header) + name_header_layout.setContentsMargins(0, 0, 0, 0) + + name_label = QLabel("Name: ") + name_label.setToolTip("This is what shows up on the GUI") + + self.name_edit = QLineEdit() + self.name_edit.setPlaceholderText("Enter objective name") + self.name_edit.setFixedWidth(320) + + name_header_layout.addWidget(name_label) + name_header_layout.addWidget(self.name_edit) + name_header_layout.addStretch() + name_layout.addWidget(name_header) + + formula_edit_widget = QWidget() + formula_edit_layout = QHBoxLayout(formula_edit_widget) + formula_edit_layout.setContentsMargins(0, 0, 0, 0) + formula_label = QLabel("Statistic: ") + formula_edit_layout.addWidget(formula_label) + + self.stat_combo = QComboBox() + self.stat_combo.addItems( + [ + "none", + "mean", + "std", + "std_rel", + "p80", + "p75", + "median", + "p25", + ] + ) + + self.stat_combo.setCurrentText(self.item.stat) + self.stat_combo.setFixedWidth(120) + formula_edit_layout.addWidget(self.stat_combo) + formula_edit_layout.addStretch() + + # Button set + button_set = QWidget() + hbox_set = QHBoxLayout(button_set) + hbox_set.setContentsMargins(0, 0, 0, 0) + self.btn_cancel = QPushButton("Cancel") + self.btn_add = QPushButton("Set") + self.btn_cancel.setFixedSize(96, 24) + self.btn_add.setFixedSize(96, 24) + hbox_set.addSpacing(114) + hbox_set.addWidget(self.btn_cancel) + hbox_set.addWidget(self.btn_add) + hbox_set.addStretch() + + formula_layout.addWidget(name_widget) + formula_layout.addWidget(formula_edit_widget) + formula_layout.addWidget(button_set) + + return formula_widget + + def _on_stat_changed(self): + """Update item when statistic selection changes.""" + self.item.formula["stat"] = self.stat_combo.currentText() + self.item.stat = self.stat_combo.currentText() + + def construct_formula_str(self): + name = self.name_edit.text() + stat = self.stat_combo.currentText() + if stat != self.item.stat: + self._on_stat_changed() + if name != self.item.name: + # only update name + if name in self.items: + print(f"Observable name {name} already exists!") + return + self.row_widget.update_item_name(name) + + self.close() + + def cancel(self): + self.close() diff --git a/src/badger/routine.py b/src/badger/routine.py index a0030ccd..24f471e9 100644 --- a/src/badger/routine.py +++ b/src/badger/routine.py @@ -20,6 +20,7 @@ from badger.utils import curr_ts from badger.environment import BaseEnvironment, instantiate_env from badger.factory import get_env +from badger.formula_utils import expanded_formula_mapping logger = logging.getLogger(__name__) @@ -116,11 +117,36 @@ def validate_model(cls, data: Any): # create evaluator env = data["environment"] + # get formulas, expand and map to output names + formulas = data.get("formulas", {}) or {} + output_names = list(data["vocs"].output_names) + + # Create a mapping of names as defined in the routine to observables + # for the environment. This will expand any formulas in terms of base + # observables, and map them back to the original name. Also includes + # a list of base_obs, so that the observables within any formulas will + # be included in the recorded data. + expanded_names, reverse_map, base_obs = expanded_formula_mapping(data) + selected_observables = [ + expanded_names[output_name] + for output_name in set(output_names + base_obs) + ] + logger.debug(f"Creating formula mapping: {formulas}") + logger.debug(f"Selected observables: {selected_observables}") + def evaluate_point(point: dict): logger.debug(f"Evaluating point: {point}") point = pd.Series(point).explode().to_dict() env.set_variables(point) - obs = env.get_observables(data["vocs"].output_names) + + # Get observables from env + obs = env.get_observables(selected_observables) + + # map observables back to output names from routine + for expanded_name, original_name in reverse_map.items(): + if expanded_name in obs: + obs[original_name] = obs.pop(expanded_name) + ts = curr_ts() obs["timestamp"] = ts.timestamp() obs["live"] = 1