workbench 0.8.217__py3-none-any.whl → 0.8.224__py3-none-any.whl

workbench/utils/meta_model_simulator.py

@@ -61,6 +61,13 @@ class MetaModelSimulator:
             df["abs_residual"] = df["residual"].abs()
             self._dfs[name] = df
 
+        # Find common rows across all models
+        id_sets = {name: set(df[self.id_column]) for name, df in self._dfs.items()}
+        common_ids = set.intersection(*id_sets.values())
+        sizes = ", ".join(f"{name}: {len(ids)}" for name, ids in id_sets.items())
+        log.info(f"Row counts before alignment: {sizes} -> common: {len(common_ids)}")
+        self._dfs = {name: df[df[self.id_column].isin(common_ids)] for name, df in self._dfs.items()}
+
         # Align DataFrames by sorting on id column
         self._dfs = {name: df.sort_values(self.id_column).reset_index(drop=True) for name, df in self._dfs.items()}
         log.info(f"Loaded {len(self._dfs)} models, {len(list(self._dfs.values())[0])} samples each")
@@ -372,13 +379,13 @@ class MetaModelSimulator:
         return weight_df
 
     def ensemble_failure_analysis(self) -> dict:
-        """Compare ensemble vs best overall model (not per-row oracle).
+        """Compare best ensemble strategy vs best individual model.
 
         Returns:
             Dict with comparison statistics
         """
         print("\n" + "=" * 60)
-        print("ENSEMBLE VS BEST MODEL COMPARISON")
+        print("BEST ENSEMBLE VS BEST MODEL COMPARISON")
         print("=" * 60)
 
         model_names = list(self._dfs.keys())
@@ -393,35 +400,55 @@ class MetaModelSimulator:
             combined[f"{name}_abs_err"] = df["abs_residual"].values
 
         pred_cols = [f"{name}_pred" for name in model_names]
+        conf_cols = [f"{name}_conf" for name in model_names]
+        pred_arr = combined[pred_cols].values
+        conf_arr = combined[conf_cols].values
 
-        # Calculate ensemble prediction (inverse-MAE weighted)
         mae_scores = {name: self._dfs[name]["abs_residual"].mean() for name in model_names}
         inv_mae_weights = np.array([1.0 / mae_scores[name] for name in model_names])
         inv_mae_weights = inv_mae_weights / inv_mae_weights.sum()
-        pred_arr = combined[pred_cols].values
-        combined["ensemble_pred"] = (pred_arr * inv_mae_weights).sum(axis=1)
-        combined["ensemble_abs_err"] = (combined["ensemble_pred"] - combined["target"]).abs()
 
-        # Find best overall model (lowest MAE)
+        # Compute all ensemble strategies (true ensembles that combine multiple models)
+        ensemble_strategies = {}
+        ensemble_strategies["Simple Mean"] = combined[pred_cols].mean(axis=1)
+        conf_sum = conf_arr.sum(axis=1, keepdims=True) + 1e-8
+        ensemble_strategies["Confidence-Weighted"] = (pred_arr * (conf_arr / conf_sum)).sum(axis=1)
+        ensemble_strategies["Inverse-MAE Weighted"] = (pred_arr * inv_mae_weights).sum(axis=1)
+        scaled_conf = conf_arr * inv_mae_weights
+        scaled_conf_sum = scaled_conf.sum(axis=1, keepdims=True) + 1e-8
+        ensemble_strategies["Scaled Conf-Weighted"] = (pred_arr * (scaled_conf / scaled_conf_sum)).sum(axis=1)
+        worst_model = max(mae_scores, key=mae_scores.get)
+        remaining = [n for n in model_names if n != worst_model]
+        remaining_cols = [f"{n}_pred" for n in remaining]
+        # Only add Drop Worst if it still combines multiple models
+        if len(remaining) > 1:
+            ensemble_strategies[f"Drop Worst ({worst_model})"] = combined[remaining_cols].mean(axis=1)
+
+        # Find best individual model
         best_model = min(mae_scores, key=mae_scores.get)
         combined["best_model_abs_err"] = combined[f"{best_model}_abs_err"]
+        best_model_mae = mae_scores[best_model]
 
-        # Compare ensemble vs best model
+        # Find best true ensemble strategy
+        strategy_maes = {name: (preds - combined["target"]).abs().mean() for name, preds in ensemble_strategies.items()}
+        best_strategy = min(strategy_maes, key=strategy_maes.get)
+        combined["ensemble_pred"] = ensemble_strategies[best_strategy]
+        combined["ensemble_abs_err"] = (combined["ensemble_pred"] - combined["target"]).abs()
+        ensemble_mae = strategy_maes[best_strategy]
+
+        # Compare
         combined["ensemble_better"] = combined["ensemble_abs_err"] < combined["best_model_abs_err"]
         n_better = combined["ensemble_better"].sum()
         n_total = len(combined)
 
-        ensemble_mae = combined["ensemble_abs_err"].mean()
-        best_model_mae = mae_scores[best_model]
-
         print(f"\nBest individual model: {best_model} (MAE={best_model_mae:.4f})")
-        print(f"Ensemble MAE: {ensemble_mae:.4f}")
+        print(f"Best ensemble strategy: {best_strategy} (MAE={ensemble_mae:.4f})")
         if ensemble_mae < best_model_mae:
             improvement = (best_model_mae - ensemble_mae) / best_model_mae * 100
             print(f"Ensemble improves over best model by {improvement:.1f}%")
         else:
             degradation = (ensemble_mae - best_model_mae) / best_model_mae * 100
-            print(f"Ensemble is worse than best model by {degradation:.1f}%")
+            print(f"No ensemble benefit: best single model outperforms all ensemble strategies by {degradation:.1f}%")
 
         print("\nPer-row comparison:")
         print(f"  Ensemble wins: {n_better}/{n_total} ({100*n_better/n_total:.1f}%)")
@@ -443,6 +470,7 @@ class MetaModelSimulator:
 
         return {
             "ensemble_mae": ensemble_mae,
+            "best_strategy": best_strategy,
             "best_model": best_model,
             "best_model_mae": best_model_mae,
             "ensemble_win_rate": n_better / n_total,

workbench/utils/model_utils.py

@@ -173,7 +173,6 @@ def fingerprint_prox_model_local(
         include_all_columns=include_all_columns,
         radius=radius,
         n_bits=n_bits,
-        counts=counts,
     )
 
 

workbench/utils/plot_utils.py

@@ -1,14 +1,18 @@
 """Plot Utilities for Workbench"""
 
+import logging
 import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
+from dash import html
+
+log = logging.getLogger("workbench")
 
 
 # For approximating beeswarm effect
 def beeswarm_offsets(values, point_size=0.05, precision=2, max_offset=0.3):
     """
-    Generate optimal beeswarm offsets with a maximum limit.
+    Generate beeswarm offsets using random jitter with collision avoidance.
 
     Args:
         values: Array of positions to be adjusted
@@ -22,42 +26,55 @@ def beeswarm_offsets(values, point_size=0.05, precision=2, max_offset=0.3):
     values = np.asarray(values)
     rounded = np.round(values, precision)
     offsets = np.zeros_like(values, dtype=float)
-
-    # Sort indices by original values
-    sorted_idx = np.argsort(values)
+    rng = np.random.default_rng(42)  # Fixed seed for reproducibility
 
     for val in np.unique(rounded):
         # Get indices belonging to this group
-        group_idx = sorted_idx[np.isin(sorted_idx, np.where(rounded == val)[0])]
+        group_mask = rounded == val
+        group_idx = np.where(group_mask)[0]
 
         if len(group_idx) > 1:
             # Track occupied positions for collision detection
             occupied = []
 
             for idx in group_idx:
-                # Find best position with no collision
-                offset = 0
-                direction = 1
-                step = 0
-
-                while True:
-                    # Check if current offset position is free
-                    collision = any(abs(offset - pos) < point_size for pos in occupied)
-
-                    if not collision or abs(offset) >= max_offset:
-                        # Accept position if no collision or max offset reached
-                        if abs(offset) > max_offset:
-                            # Clamp to maximum
-                            offset = max_offset * (1 if offset > 0 else -1)
-                        break
-
-                    # Switch sides with increasing distance
-                    step += 0.25
-                    direction *= -1
-                    offset = direction * step * point_size
-
-                offsets[idx] = offset
-                occupied.append(offset)
+                # Try random positions, starting near center and expanding outward
+                best_offset = 0
+                found = False
+
+                # First point goes to center
+                if not occupied:
+                    found = True
+                else:
+                    # Try random positions with increasing spread
+                    for attempt in range(50):
+                        # Gradually increase the range of random offsets
+                        spread = min(max_offset, point_size * (1 + attempt * 0.5))
+                        offset = rng.uniform(-spread, spread)
+
+                        # Check for collision with occupied positions
+                        if not any(abs(offset - pos) < point_size for pos in occupied):
+                            best_offset = offset
+                            found = True
+                            break
+
+                    # If no free position found after attempts, find the least crowded spot
+                    if not found:
+                        # Try a grid of positions and pick one with most space
+                        candidates = np.linspace(-max_offset, max_offset, 20)
+                        rng.shuffle(candidates)
+                        for candidate in candidates:
+                            if not any(abs(candidate - pos) < point_size * 0.8 for pos in occupied):
+                                best_offset = candidate
+                                found = True
+                                break
+
+                        # Last resort: just use a random position within bounds
+                        if not found:
+                            best_offset = rng.uniform(-max_offset, max_offset)
+
+                offsets[idx] = best_offset
+                occupied.append(best_offset)
 
     return offsets
 
@@ -179,6 +196,107 @@ def prediction_intervals(df, figure, x_col):
     return figure
 
 
+def molecule_hover_tooltip(smiles: str, mol_id: str = None, width: int = 300, height: int = 200) -> list:
+    """Generate a molecule hover tooltip from a SMILES string.
+
+    This function creates a visually appealing tooltip with a dark background
+    that displays the molecule ID at the top and structure below when hovering
+    over scatter plot points.
+
+    Args:
+        smiles: SMILES string representing the molecule
+        mol_id: Optional molecule ID to display at the top of the tooltip
+        width: Width of the molecule image in pixels (default: 300)
+        height: Height of the molecule image in pixels (default: 200)
+
+    Returns:
+        list: A list containing an html.Div with the ID header and molecule SVG,
+              or an html.Div with an error message if rendering fails
+    """
+    try:
+        from workbench.utils.chem_utils.vis import svg_from_smiles
+        from workbench.utils.theme_manager import ThemeManager
+
+        # Get the background color from the current theme
+        background = ThemeManager().background()
+
+        # Generate the SVG image from SMILES
+        img = svg_from_smiles(smiles, width, height, background=background)
+
+        if img is None:
+            log.warning(f"Could not render molecule for SMILES: {smiles}")
+            return [
+                html.Div(
+                    "Invalid SMILES",
+                    className="custom-tooltip",
+                    style={
+                        "padding": "10px",
+                        "color": "rgb(255, 140, 140)",
+                        "width": f"{width}px",
+                        "height": f"{height}px",
+                        "display": "flex",
+                        "alignItems": "center",
+                        "justifyContent": "center",
+                    },
+                )
+            ]
+
+        # Build the tooltip with ID header and molecule image
+        children = []
+
+        # Add ID header if provided
+        if mol_id is not None:
+            children.append(
+                html.Div(
+                    str(mol_id),
+                    style={
+                        "textAlign": "center",
+                        "padding": "8px",
+                        "color": "rgb(200, 200, 200)",
+                        "fontSize": "14px",
+                        "fontWeight": "bold",
+                        "borderBottom": "1px solid rgba(128, 128, 128, 0.5)",
+                    },
+                )
+            )
+
+        # Add molecule image
+        children.append(
+            html.Img(
+                src=img,
+                style={"padding": "0px", "margin": "0px", "display": "block"},
+                width=str(width),
+                height=str(height),
+            )
+        )
+
+        return [
+            html.Div(
+                children,
+                className="custom-tooltip",
+                style={"padding": "0px", "margin": "0px"},
+            )
+        ]
+
+    except ImportError as e:
+        log.error(f"RDKit not available for molecule rendering: {e}")
+        return [
+            html.Div(
+                "RDKit not installed",
+                className="custom-tooltip",
+                style={
+                    "padding": "10px",
+                    "color": "rgb(255, 195, 140)",
+                    "width": f"{width}px",
+                    "height": f"{height}px",
+                    "display": "flex",
+                    "alignItems": "center",
+                    "justifyContent": "center",
+                },
+            )
+        ]
+
+
 if __name__ == "__main__":
     """Exercise the Plot Utilities"""
     import plotly.express as px

workbench/utils/shap_utils.py

@@ -9,6 +9,7 @@ from typing import Optional, List, Tuple, Dict, Union
 from workbench.utils.xgboost_model_utils import xgboost_model_from_s3
 from workbench.utils.model_utils import load_category_mappings_from_s3
 from workbench.utils.pandas_utils import convert_categorical_types
+from workbench.model_script_utils.model_script_utils import decompress_features
 
 # Set up the log
 log = logging.getLogger("workbench")
@@ -111,61 +112,6 @@ def shap_values_data(
         return result_df, feature_df
 
 
-def decompress_features(
-    df: pd.DataFrame, features: List[str], compressed_features: List[str]
-) -> Tuple[pd.DataFrame, List[str]]:
-    """Prepare features for the XGBoost model
-
-    Args:
-        df (pd.DataFrame): The features DataFrame
-        features (List[str]): Full list of feature names
-        compressed_features (List[str]): List of feature names to decompress (bitstrings)
-
-    Returns:
-        pd.DataFrame: DataFrame with the decompressed features
-        List[str]: Updated list of feature names after decompression
-
-    Raises:
-        ValueError: If any missing values are found in the specified features
-    """
-
-    # Check for any missing values in the required features
-    missing_counts = df[features].isna().sum()
-    if missing_counts.any():
-        missing_features = missing_counts[missing_counts > 0]
-        print(
-            f"WARNING: Found missing values in features: {missing_features.to_dict()}. "
-            "WARNING: You might want to remove/replace all NaN values before processing."
-        )
-
-    # Decompress the specified compressed features
-    decompressed_features = features
-    for feature in compressed_features:
-        if (feature not in df.columns) or (feature not in features):
-            print(f"Feature '{feature}' not in the features list, skipping decompression.")
-            continue
-
-        # Remove the feature from the list of features to avoid duplication
-        decompressed_features.remove(feature)
-
-        # Handle all compressed features as bitstrings
-        bit_matrix = np.array([list(bitstring) for bitstring in df[feature]], dtype=np.uint8)
-        prefix = feature[:3]
-
-        # Create all new columns at once - avoids fragmentation
-        new_col_names = [f"{prefix}_{i}" for i in range(bit_matrix.shape[1])]
-        new_df = pd.DataFrame(bit_matrix, columns=new_col_names, index=df.index)
-
-        # Add to features list
-        decompressed_features.extend(new_col_names)
-
-        # Drop original column and concatenate new ones
-        df = df.drop(columns=[feature])
-        df = pd.concat([df, new_df], axis=1)
-
-    return df, decompressed_features
-
-
 def _calculate_shap_values(workbench_model, sample_df: pd.DataFrame = None):
     """
     Internal function to calculate SHAP values for Workbench Models.

workbench/utils/theme_manager.py

@@ -76,10 +76,28 @@ class ThemeManager:
     def set_theme(cls, theme_name: str):
         """Set the current theme."""
 
-        # For 'auto', we try to grab a theme from the Parameter Store
-        # if we can't find one, we'll set the theme to the default
+        # For 'auto', we check multiple sources in priority order:
+        # 1. Browser cookie (from localStorage, for per-user preference)
+        # 2. Parameter Store (for org-wide default)
+        # 3. Default theme
         if theme_name == "auto":
-            theme_name = cls.ps.get("/workbench/dashboard/theme", warn=False) or cls.default_theme
+            theme_name = None
+
+            # 1. Check Flask request cookie (set from localStorage)
+            try:
+                from flask import request, has_request_context
+
+                if has_request_context():
+                    theme_name = request.cookies.get("wb_theme")
+            except Exception:
+                pass
+
+            # 2. Fall back to ParameterStore
+            if not theme_name:
+                theme_name = cls.ps.get("/workbench/dashboard/theme", warn=False)
+
+            # 3. Fall back to default
+            theme_name = theme_name or cls.default_theme
 
         # Check if the theme is in our available themes
         if theme_name not in cls.available_themes:
@@ -104,9 +122,27 @@ class ThemeManager:
         cls.current_theme_name = theme_name
         cls.log.info(f"Theme set to '{theme_name}'")
 
+    # Bootstrap themes that are dark mode (from Bootswatch)
+    _dark_bootstrap_themes = {"DARKLY", "CYBORG", "SLATE", "SOLAR", "SUPERHERO", "VAPOR"}
+
     @classmethod
     def dark_mode(cls) -> bool:
-        """Check if the current theme is a dark mode theme."""
+        """Check if the current theme is a dark mode theme.
+
+        Determines dark mode by checking if the Bootstrap base theme is a known dark theme.
+        Falls back to checking if 'dark' is in the theme name.
+        """
+        theme = cls.available_themes.get(cls.current_theme_name, {})
+        base_css = theme.get("base_css", "")
+
+        # Check if the base CSS URL contains a known dark Bootstrap theme
+        if base_css:
+            base_css_upper = base_css.upper()
+            for dark_theme in cls._dark_bootstrap_themes:
+                if dark_theme in base_css_upper:
+                    return True
+
+        # Fallback: check if 'dark' is in the theme name
         return "dark" in cls.current_theme().lower()
 
     @classmethod
@@ -184,30 +220,57 @@ class ThemeManager:
 
     @classmethod
     def css_files(cls) -> list[str]:
-        """Get the list of CSS files for the current theme."""
-        theme = cls.available_themes[cls.current_theme_name]
+        """Get the list of CSS files for the current theme.
+
+        Note: Uses /base.css route for dynamic theme switching instead of CDN URLs.
+        """
         css_files = []
 
-        # Add base.css or its CDN URL
-        if theme["base_css"]:
-            css_files.append(theme["base_css"])
+        # Use Flask route for base CSS (allows dynamic theme switching)
+        css_files.append("/base.css")
 
         # Add the DBC template CSS
         css_files.append(cls.dbc_css)
 
         # Add custom.css if it exists
-        if theme["custom_css"]:
-            css_files.append("/custom.css")
+        css_files.append("/custom.css")
 
         return css_files
 
+    @classmethod
+    def _get_theme_from_cookie(cls):
+        """Get the theme dict based on the wb_theme cookie, falling back to current theme."""
+        from flask import request
+
+        theme_name = request.cookies.get("wb_theme")
+        if theme_name and theme_name in cls.available_themes:
+            return cls.available_themes[theme_name], theme_name
+        return cls.available_themes[cls.current_theme_name], cls.current_theme_name
+
     @classmethod
     def register_css_route(cls, app):
-        """Register Flask route for custom.css."""
+        """Register Flask routes for CSS and before_request hook for theme switching."""
+        from flask import redirect
+
+        @app.server.before_request
+        def check_theme_cookie():
+            """Check for theme cookie on each request and update theme if needed."""
+            _, theme_name = cls._get_theme_from_cookie()
+            if theme_name != cls.current_theme_name:
+                cls.set_theme(theme_name)
+
+        @app.server.route("/base.css")
+        def serve_base_css():
+            """Redirect to the appropriate Bootstrap theme CSS based on cookie."""
+            theme, _ = cls._get_theme_from_cookie()
+            if theme["base_css"]:
+                return redirect(theme["base_css"])
+            return "", 404
 
         @app.server.route("/custom.css")
         def serve_custom_css():
-            theme = cls.available_themes[cls.current_theme_name]
+            """Serve the custom.css file based on cookie."""
+            theme, _ = cls._get_theme_from_cookie()
             if theme["custom_css"]:
                 return send_from_directory(theme["custom_css"].parent, theme["custom_css"].name)
             return "", 404
@@ -250,23 +313,25 @@ class ThemeManager:
         # Loop over each path in the theme path
         for theme_path in cls.theme_path_list:
             for theme_dir in theme_path.iterdir():
-                if theme_dir.is_dir():
-                    theme_name = theme_dir.name
-
-                    # Grab the base.css URL
-                    base_css_url = cls._get_base_css_url(theme_dir)
-
-                    # Grab the plotly template json, custom.css, and branding json
-                    plotly_template = theme_dir / "plotly.json"
-                    custom_css = theme_dir / "custom.css"
-                    branding = theme_dir / "branding.json"
-
-                    cls.available_themes[theme_name] = {
-                        "base_css": base_css_url,
-                        "plotly_template": plotly_template,
-                        "custom_css": custom_css if custom_css.exists() else None,
-                        "branding": branding if branding.exists() else None,
-                    }
+                # Skip hidden directories (e.g., .idea, .git)
+                if not theme_dir.is_dir() or theme_dir.name.startswith("."):
+                    continue
+                theme_name = theme_dir.name
+
+                # Grab the base.css URL
+                base_css_url = cls._get_base_css_url(theme_dir)
+
+                # Grab the plotly template json, custom.css, and branding json
+                plotly_template = theme_dir / "plotly.json"
+                custom_css = theme_dir / "custom.css"
+                branding = theme_dir / "branding.json"
+
+                cls.available_themes[theme_name] = {
+                    "base_css": base_css_url,
+                    "plotly_template": plotly_template,
+                    "custom_css": custom_css if custom_css.exists() else None,
+                    "branding": branding if branding.exists() else None,
+                }
 
         if not cls.available_themes:
             cls.log.warning(f"No themes found in '{cls.theme_path_list}'...")