workbench 0.8.174__py3-none-any.whl → 0.8.227__py3-none-any.whl

workbench/utils/plot_utils.py

@@ -1,14 +1,20 @@
 """Plot Utilities for Workbench"""
 
+import logging
 import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
+from dash import html
+
+from workbench.utils.color_utils import is_dark
+
+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 +28,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
 
@@ -132,7 +151,7 @@ def prediction_intervals(df, figure, x_col):
             x=sorted_df[x_col],
             y=sorted_df["q_025"],
             mode="lines",
-            line=dict(width=1, color="rgba(99, 110, 250, 0.25)"),
+            line=dict(width=1, color="rgba(99, 110, 250, 0.5)"),
             name="2.5 Percentile",
             hoverinfo="skip",
             showlegend=False,
@@ -143,12 +162,12 @@ def prediction_intervals(df, figure, x_col):
             x=sorted_df[x_col],
             y=sorted_df["q_975"],
             mode="lines",
-            line=dict(width=1, color="rgba(99, 110, 250, 0.25)"),
+            line=dict(width=1, color="rgba(99, 110, 250, 0.5)"),
             name="97.5 Percentile",
             hoverinfo="skip",
             showlegend=False,
             fill="tonexty",
-            fillcolor="rgba(99, 110, 250, 0.2)",
+            fillcolor="rgba(99, 110, 250, 0.35)",
         )
     )
     # Add inner band (q_25 to q_75) - less transparent
@@ -157,7 +176,7 @@ def prediction_intervals(df, figure, x_col):
             x=sorted_df[x_col],
             y=sorted_df["q_25"],
             mode="lines",
-            line=dict(width=1, color="rgba(99, 250, 110, 0.25)"),
+            line=dict(width=1, color="rgba(99, 250, 110, 0.5)"),
             name="25 Percentile",
             hoverinfo="skip",
             showlegend=False,
@@ -168,17 +187,123 @@ def prediction_intervals(df, figure, x_col):
             x=sorted_df[x_col],
             y=sorted_df["q_75"],
             mode="lines",
-            line=dict(width=1, color="rgba(99, 250, 100, 0.25)"),
+            line=dict(width=1, color="rgba(99, 250, 110, 0.5)"),
             name="75 Percentile",
             hoverinfo="skip",
             showlegend=False,
             fill="tonexty",
-            fillcolor="rgba(99, 250, 110, 0.2)",
+            fillcolor="rgba(99, 250, 110, 0.35)",
         )
     )
     return figure
 
 
+def molecule_hover_tooltip(
+    smiles: str, mol_id: str = None, width: int = 300, height: int = 200, background: str = None
+) -> 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)
+        background: Optional background color (if None, uses dark gray)
+
+    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
+
+        # Use provided background or default to dark gray
+        if background is None:
+            background = "rgba(64, 64, 64, 1)"
+
+        # Generate the SVG image from SMILES (base64 encoded data URI)
+        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:
+            # Set text color based on background brightness
+            text_color = "rgb(200, 200, 200)" if is_dark(background) else "rgb(60, 60, 60)"
+            children.append(
+                html.Div(
+                    str(mol_id),
+                    style={
+                        "textAlign": "center",
+                        "padding": "8px",
+                        "color": text_color,
+                        "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/pytorch_utils.py

@@ -0,0 +1,87 @@
+"""PyTorch Tabular utilities for Workbench models."""
+
+import logging
+import os
+import tarfile
+import tempfile
+from typing import Any, Tuple
+
+import awswrangler as wr
+import pandas as pd
+
+from workbench.utils.aws_utils import pull_s3_data
+from workbench.utils.metrics_utils import compute_metrics_from_predictions
+
+log = logging.getLogger("workbench")
+
+
+def download_and_extract_model(s3_uri: str, model_dir: str) -> None:
+    """Download and extract a PyTorch model artifact from S3.
+
+    Args:
+        s3_uri: S3 URI of the model.tar.gz artifact
+        model_dir: Local directory to extract the model to
+    """
+    with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
+        tmp_path = tmp.name
+
+    try:
+        wr.s3.download(path=s3_uri, local_file=tmp_path)
+        with tarfile.open(tmp_path, "r:gz") as tar:
+            tar.extractall(model_dir)
+        log.info(f"Extracted model to {model_dir}")
+    finally:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+
+
+def pull_cv_results(workbench_model: Any) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """Pull cross-validation results from AWS training artifacts.
+
+    This retrieves the validation predictions saved during model training and
+    computes metrics directly from them. For PyTorch models trained with
+    n_folds > 1, these are out-of-fold predictions from k-fold cross-validation.
+
+    Args:
+        workbench_model: Workbench model object
+
+    Returns:
+        Tuple of:
+            - DataFrame with computed metrics
+            - DataFrame with validation predictions
+    """
+    # Get the validation predictions from S3
+    s3_path = f"{workbench_model.model_training_path}/validation_predictions.csv"
+    predictions_df = pull_s3_data(s3_path)
+
+    if predictions_df is None:
+        raise ValueError(f"No validation predictions found at {s3_path}")
+
+    log.info(f"Pulled {len(predictions_df)} validation predictions from {s3_path}")
+
+    # Compute metrics from predictions
+    target = workbench_model.target()
+    class_labels = workbench_model.class_labels()
+
+    if target in predictions_df.columns and "prediction" in predictions_df.columns:
+        metrics_df = compute_metrics_from_predictions(predictions_df, target, class_labels)
+    else:
+        metrics_df = pd.DataFrame()
+
+    return metrics_df, predictions_df
+
+
+if __name__ == "__main__":
+    from workbench.api import Model
+
+    # Test pulling CV results
+    model_name = "aqsol-reg-pytorch"
+    print(f"Loading Workbench model: {model_name}")
+    model = Model(model_name)
+    print(f"Model Framework: {model.model_framework}")
+
+    # Pull CV results from training artifacts
+    metrics_df, predictions_df = pull_cv_results(model)
+    print(f"\nMetrics:\n{metrics_df}")
+    print(f"\nPredictions shape: {predictions_df.shape}")
+    print(f"Predictions columns: {predictions_df.columns.tolist()}")

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.
@@ -212,6 +158,14 @@ def _calculate_shap_values(workbench_model, sample_df: pd.DataFrame = None):
         log.error("No XGBoost model found in the artifact.")
         return None, None, None, None
 
+    # Get the booster (SHAP requires the booster, not the sklearn wrapper)
+    if hasattr(xgb_model, "get_booster"):
+        # Full sklearn model - extract the booster
+        booster = xgb_model.get_booster()
+    else:
+        # Already a booster
+        booster = xgb_model
+
     # Load category mappings if available
     category_mappings = load_category_mappings_from_s3(model_artifact_uri)
 
@@ -229,8 +183,8 @@ def _calculate_shap_values(workbench_model, sample_df: pd.DataFrame = None):
     # Create a DMatrix with categorical support
     dmatrix = xgb.DMatrix(X, enable_categorical=True)
 
-    # Use XGBoost's built-in SHAP calculation
-    shap_values = xgb_model.predict(dmatrix, pred_contribs=True, strict_shape=True)
+    # Use XGBoost's built-in SHAP calculation (booster method, not sklearn)
+    shap_values = booster.predict(dmatrix, pred_contribs=True, strict_shape=True)
     features_with_bias = features + ["bias"]
 
     # Now we need to subset the columns based on top 10 SHAP values

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}'...")