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

workbench/model_scripts/custom_models/uq_models/feature_space_proximity.py

@@ -0,0 +1,194 @@
+import pandas as pd
+import numpy as np
+from sklearn.preprocessing import StandardScaler
+from sklearn.neighbors import NearestNeighbors
+from typing import List, Optional
+import logging
+
+# Workbench Imports
+from workbench.algorithms.dataframe.proximity import Proximity
+from workbench.algorithms.dataframe.projection_2d import Projection2D
+
+# Set up logging
+log = logging.getLogger("workbench")
+
+
+class FeatureSpaceProximity(Proximity):
+    """Proximity computations for numeric feature spaces using Euclidean distance."""
+
+    def __init__(
+        self,
+        df: pd.DataFrame,
+        id_column: str,
+        features: List[str],
+        target: Optional[str] = None,
+        include_all_columns: bool = False,
+    ):
+        """
+        Initialize the FeatureSpaceProximity class.
+
+        Args:
+            df: DataFrame containing data for neighbor computations.
+            id_column: Name of the column used as the identifier.
+            features: List of feature column names to be used for neighbor computations.
+            target: Name of the target column. Defaults to None.
+            include_all_columns: Include all DataFrame columns in neighbor results. Defaults to False.
+        """
+        # Validate and filter features before calling parent init
+        self._raw_features = features
+        super().__init__(
+            df, id_column=id_column, features=features, target=target, include_all_columns=include_all_columns
+        )
+
+    def _prepare_data(self) -> None:
+        """Filter out non-numeric features and drop NaN rows."""
+        # Validate features
+        self.features = self._validate_features(self.df, self._raw_features)
+
+        # Drop NaN rows for the features we're using
+        self.df = self.df.dropna(subset=self.features).copy()
+
+    def _validate_features(self, df: pd.DataFrame, features: List[str]) -> List[str]:
+        """Remove non-numeric features and log warnings."""
+        non_numeric = [f for f in features if f not in df.select_dtypes(include=["number"]).columns]
+        if non_numeric:
+            log.warning(f"Non-numeric features {non_numeric} aren't currently supported, excluding them")
+        return [f for f in features if f not in non_numeric]
+
+    def _build_model(self) -> None:
+        """Standardize features and fit Nearest Neighbors model."""
+        self.scaler = StandardScaler()
+        X = self.scaler.fit_transform(self.df[self.features])
+        self.nn = NearestNeighbors().fit(X)
+
+    def _transform_features(self, df: pd.DataFrame) -> np.ndarray:
+        """Transform features using the fitted scaler."""
+        return self.scaler.transform(df[self.features])
+
+    def _project_2d(self) -> None:
+        """Project the numeric features to 2D for visualization."""
+        if len(self.features) >= 2:
+            self.df = Projection2D().fit_transform(self.df, features=self.features)
+
+
+# Testing the FeatureSpaceProximity class
+if __name__ == "__main__":
+
+    pd.set_option("display.max_columns", None)
+    pd.set_option("display.width", 1000)
+
+    # Create a sample DataFrame
+    data = {
+        "ID": [1, 2, 3, 4, 5],
+        "Feature1": [0.1, 0.2, 0.3, 0.4, 0.5],
+        "Feature2": [0.5, 0.4, 0.3, 0.2, 0.1],
+        "Feature3": [2.5, 2.4, 2.3, 2.3, np.nan],
+    }
+    df = pd.DataFrame(data)
+
+    # Test the FeatureSpaceProximity class
+    features = ["Feature1", "Feature2", "Feature3"]
+    prox = FeatureSpaceProximity(df, id_column="ID", features=features)
+    print(prox.neighbors(1, n_neighbors=2))
+
+    # Test the neighbors method with radius
+    print(prox.neighbors(1, radius=2.0))
+
+    # Test with Features list
+    prox = FeatureSpaceProximity(df, id_column="ID", features=["Feature1"])
+    print(prox.neighbors(1))
+
+    # Create a sample DataFrame
+    data = {
+        "id": ["a", "b", "c", "d", "e"],  # Testing string IDs
+        "Feature1": [0.1, 0.2, 0.3, 0.4, 0.5],
+        "Feature2": [0.5, 0.4, 0.3, 0.2, 0.1],
+        "target": [1, 0, 1, 0, 5],
+    }
+    df = pd.DataFrame(data)
+
+    # Test with String Ids
+    prox = FeatureSpaceProximity(
+        df,
+        id_column="id",
+        features=["Feature1", "Feature2"],
+        target="target",
+        include_all_columns=True,
+    )
+    print(prox.neighbors(["a", "b"]))
+
+    # Test duplicate IDs
+    data = {
+        "id": ["a", "b", "c", "d", "d"],  # Duplicate ID (d)
+        "Feature1": [0.1, 0.2, 0.3, 0.4, 0.5],
+        "Feature2": [0.5, 0.4, 0.3, 0.2, 0.1],
+        "target": [1, 0, 1, 0, 5],
+    }
+    df = pd.DataFrame(data)
+    prox = FeatureSpaceProximity(df, id_column="id", features=["Feature1", "Feature2"], target="target")
+    print(df.equals(prox.df))
+
+    # Test on real data from Workbench
+    from workbench.api import FeatureSet, Model
+
+    fs = FeatureSet("aqsol_features")
+    model = Model("aqsol-regression")
+    features = model.features()
+    df = fs.pull_dataframe()
+    prox = FeatureSpaceProximity(df, id_column=fs.id_column, features=model.features(), target=model.target())
+    print("\n" + "=" * 80)
+    print("Testing Neighbors...")
+    print("=" * 80)
+    test_id = df[fs.id_column].tolist()[0]
+    print(f"\nNeighbors for ID {test_id}:")
+    print(prox.neighbors(test_id))
+
+    print("\n" + "=" * 80)
+    print("Testing isolated_compounds...")
+    print("=" * 80)
+
+    # Test isolated data in the top 1%
+    isolated_1pct = prox.isolated(top_percent=1.0)
+    print(f"\nTop 1% most isolated compounds (n={len(isolated_1pct)}):")
+    print(isolated_1pct)
+
+    # Test isolated data in the top 5%
+    isolated_5pct = prox.isolated(top_percent=5.0)
+    print(f"\nTop 5% most isolated compounds (n={len(isolated_5pct)}):")
+    print(isolated_5pct)
+
+    print("\n" + "=" * 80)
+    print("Testing target_gradients...")
+    print("=" * 80)
+
+    # Test with different parameters
+    gradients_1pct = prox.target_gradients(top_percent=1.0, min_delta=1.0)
+    print(f"\nTop 1% target gradients (min_delta=5.0) (n={len(gradients_1pct)}):")
+    print(gradients_1pct)
+
+    gradients_5pct = prox.target_gradients(top_percent=5.0, min_delta=5.0)
+    print(f"\nTop 5% target gradients (min_delta=5.0) (n={len(gradients_5pct)}):")
+    print(gradients_5pct)
+
+    # Test proximity_stats
+    print("\n" + "=" * 80)
+    print("Testing proximity_stats...")
+    print("=" * 80)
+    stats = prox.proximity_stats()
+    print(stats)
+
+    # Plot the distance distribution using pandas
+    print("\n" + "=" * 80)
+    print("Plotting distance distribution...")
+    print("=" * 80)
+    prox.df["nn_distance"].hist(bins=50, figsize=(10, 6), edgecolor="black")
+
+    # Visualize the 2D projection
+    print("\n" + "=" * 80)
+    print("Visualizing 2D Projection...")
+    print("=" * 80)
+    from workbench.web_interface.components.plugin_unit_test import PluginUnitTest
+    from workbench.web_interface.components.plugins.scatter_plot import ScatterPlot
+
+    unit_test = PluginUnitTest(ScatterPlot, input_data=prox.df[:1000], x="x", y="y", color=model.target())
+    unit_test.run()

workbench/model_scripts/custom_models/uq_models/gaussian_process.template

@@ -9,7 +9,7 @@ from sklearn.model_selection import train_test_split
 TEMPLATE_PARAMS = {
     "features": "{{feature_list}}",
     "target": "{{target_column}}",
-    "train_all_data": "{{train_all_data}}"
+    "train_all_data": "{{train_all_data}}",
 }
 
 from io import StringIO
@@ -33,7 +33,7 @@ def match_features_case_insensitive(df: pd.DataFrame, model_features: list) -> p
     """
     Matches and renames DataFrame columns to match model feature names (case-insensitive).
     Prioritizes exact matches, then case-insensitive matches.
-    
+
     Raises ValueError if any model features cannot be matched.
     """
     df_columns_lower = {col.lower(): col for col in df.columns}
@@ -46,7 +46,7 @@ def match_features_case_insensitive(df: pd.DataFrame, model_features: list) -> p
             rename_dict[df_columns_lower[feature.lower()]] = feature
         else:
             missing.append(feature)
-    
+
     if missing:
         raise ValueError(f"Features not found: {missing}")
 
@@ -76,10 +76,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     # Load training data from the specified directory
-    training_files = [
-        os.path.join(args.train, file)
-        for file in os.listdir(args.train) if file.endswith(".csv")
-    ]
+    training_files = [os.path.join(args.train, file) for file in os.listdir(args.train) if file.endswith(".csv")]
     df = pd.concat([pd.read_csv(file, engine="python") for file in training_files])
 
     # Check if the DataFrame is empty
@@ -112,10 +109,7 @@ if __name__ == "__main__":
     )
 
     # Create a Pipeline with StandardScaler
-    model = Pipeline([
-        ("scaler", StandardScaler()),
-        ("model", model)
-    ])
+    model = Pipeline([("scaler", StandardScaler()), ("model", model)])
 
     # Prepare features and targets for training
     X_train = df_train[features]

workbench/model_scripts/custom_models/uq_models/ngboost.template

@@ -3,11 +3,8 @@ from ngboost import NGBRegressor
 from sklearn.model_selection import train_test_split
 
 # Model Performance Scores
-from sklearn.metrics import (
-    mean_absolute_error,
-    r2_score,
-    root_mean_squared_error
-)
+from sklearn.metrics import mean_absolute_error, median_absolute_error, r2_score, root_mean_squared_error
+from scipy.stats import spearmanr
 
 from io import StringIO
 import json
@@ -21,7 +18,7 @@ import pandas as pd
 TEMPLATE_PARAMS = {
     "features": "{{feature_list}}",
     "target": "{{target_column}}",
-    "train_all_data": "{{train_all_data}}"
+    "train_all_data": "{{train_all_data}}",
 }
 
 
@@ -87,10 +84,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     # Load training data from the specified directory
-    training_files = [
-        os.path.join(args.train, file)
-        for file in os.listdir(args.train) if file.endswith(".csv")
-    ]
+    training_files = [os.path.join(args.train, file) for file in os.listdir(args.train) if file.endswith(".csv")]
     print(f"Training Files: {training_files}")
 
     # Combine files and read them all into a single pandas dataframe
@@ -136,11 +130,16 @@ if __name__ == "__main__":
     # Calculate various model performance metrics (regression)
     rmse = root_mean_squared_error(y_validate, preds)
     mae = mean_absolute_error(y_validate, preds)
+    medae = median_absolute_error(y_validate, preds)
     r2 = r2_score(y_validate, preds)
-    print(f"RMSE: {rmse:.3f}")
-    print(f"MAE: {mae:.3f}")
-    print(f"R2: {r2:.3f}")
-    print(f"NumRows: {len(df_val)}")
+    spearman_corr = spearmanr(y_validate, preds).correlation
+    support = len(df_val)
+    print(f"rmse: {rmse:.3f}")
+    print(f"mae: {mae:.3f}")
+    print(f"medae: {medae:.3f}")
+    print(f"r2: {r2:.3f}")
+    print(f"spearmanr: {spearman_corr:.3f}")
+    print(f"support: {support}")
 
     # Save the trained NGBoost model
     joblib.dump(ngb_model, os.path.join(args.model_dir, "ngb_model.joblib"))
@@ -212,8 +211,8 @@ def predict_fn(df, model) -> pd.DataFrame:
     dist_params = y_dists.params
 
     # Extract mean and std from distribution parameters
-    df["prediction"] = dist_params['loc']  # mean
-    df["prediction_std"] = dist_params['scale']  # standard deviation
+    df["prediction"] = dist_params["loc"]  # mean
+    df["prediction_std"] = dist_params["scale"]  # standard deviation
 
     # Add 95% prediction intervals using ppf (percent point function)
     df["q_025"] = y_dists.ppf(0.025)  # 2.5th percentile

workbench/model_scripts/ensemble_xgb/ensemble_xgb.template

@@ -3,7 +3,7 @@ TEMPLATE_PARAMS = {
     "model_type": "{{model_type}}",
     "target_column": "{{target_column}}",
     "feature_list": "{{feature_list}}",
-    "model_metrics_s3_path": "{{model_metrics_s3_path}}"
+    "model_metrics_s3_path": "{{model_metrics_s3_path}}",
 }
 
 # Imports for XGB Model
@@ -12,11 +12,8 @@ import awswrangler as wr
 import numpy as np
 
 # Model Performance Scores
-from sklearn.metrics import (
-    mean_absolute_error,
-    r2_score,
-    root_mean_squared_error
-)
+from sklearn.metrics import mean_absolute_error, median_absolute_error, r2_score, root_mean_squared_error
+from scipy.stats import spearmanr
 
 from io import StringIO
 import json
@@ -39,6 +36,7 @@ def check_dataframe(df: pd.DataFrame, df_name: str) -> None:
         print(msg)
         raise ValueError(msg)
 
+
 def match_features_case_insensitive(df: pd.DataFrame, model_features: list) -> pd.DataFrame:
     """
     Matches and renames the DataFrame's column names to match the model's feature names (case-insensitive).
@@ -95,11 +93,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     # Read the training data into DataFrames
-    training_files = [
-        os.path.join(args.train, file)
-        for file in os.listdir(args.train)
-        if file.endswith(".csv")
-    ]
+    training_files = [os.path.join(args.train, file) for file in os.listdir(args.train) if file.endswith(".csv")]
     print(f"Training Files: {training_files}")
 
     # Combine files and read them all into a single pandas dataframe
@@ -150,7 +144,6 @@ if __name__ == "__main__":
     result_df["residual"] = result_df[target] - result_df["prediction"]
     result_df["residual_abs"] = result_df["residual"].abs()
 
-
     # Save the results dataframe to S3
     wr.s3.to_csv(
         result_df,
@@ -161,11 +154,16 @@ if __name__ == "__main__":
     # Report Performance Metrics
     rmse = root_mean_squared_error(result_df[target], result_df["prediction"])
     mae = mean_absolute_error(result_df[target], result_df["prediction"])
+    medae = median_absolute_error(result_df[target], result_df["prediction"])
     r2 = r2_score(result_df[target], result_df["prediction"])
-    print(f"RMSE: {rmse:.3f}")
-    print(f"MAE: {mae:.3f}")
-    print(f"R2: {r2:.3f}")
-    print(f"NumRows: {len(result_df)}")
+    spearman_corr = spearmanr(result_df[target], result_df["prediction"]).correlation
+    support = len(result_df)
+    print(f"rmse: {rmse:.3f}")
+    print(f"mae: {mae:.3f}")
+    print(f"medae: {medae:.3f}")
+    print(f"r2: {r2:.3f}")
+    print(f"spearmanr: {spearman_corr:.3f}")
+    print(f"support: {support}")
 
     # Now save the models
     for name, model in models.items():
@@ -210,7 +208,7 @@ def input_fn(input_data, content_type):
     """Parse input data and return a DataFrame."""
     if not input_data:
         raise ValueError("Empty input data is not supported!")
-    
+
     # Decode bytes to string if necessary
     if isinstance(input_data, bytes):
         input_data = input_data.decode("utf-8")

workbench/model_scripts/meta_model/generated_model_script.py

@@ -0,0 +1,209 @@
+# Meta Model Template for Workbench
+#
+# NOTE: This is called a "meta model" but it's really a "meta endpoint" - it aggregates
+# predictions from multiple child endpoints. We call it a "model" because Workbench
+# creates Model artifacts that get deployed as Endpoints, so this follows that pattern.
+#
+# Assumptions/Shortcuts:
+# - All child endpoints are regression models
+# - All child endpoints output 'prediction' and 'confidence' columns
+# - Aggregation uses model weights (provided at meta model creation time)
+#
+# This template:
+# - Has no real training phase (just saves metadata including model weights)
+# - At inference time, calls child endpoints and aggregates their predictions
+
+import argparse
+import json
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from io import StringIO
+
+import pandas as pd
+
+from workbench_bridges.endpoints.fast_inference import fast_inference
+
+# Template parameters (filled in by Workbench)
+TEMPLATE_PARAMS = {
+    "child_endpoints": ['logd-reg-pytorch', 'logd-reg-chemprop'],
+    "target_column": "logd",
+    "model_weights": {'logd-reg-pytorch': 0.4228205813233993, 'logd-reg-chemprop': 0.5771794186766008},
+    "model_metrics_s3_path": "s3://sandbox-sageworks-artifacts/models/logd-meta/training",
+    "aws_region": "us-west-2",
+}
+
+
+def invoke_endpoints_parallel(endpoint_names: list[str], df: pd.DataFrame) -> dict[str, pd.DataFrame]:
+    """Call multiple child endpoints in parallel and collect their results.
+
+    Args:
+        endpoint_names: List of endpoint names to call
+        df: Input DataFrame to send to each endpoint
+
+    Returns:
+        Dict mapping endpoint_name -> result DataFrame (or None if failed)
+    """
+    results = {}
+
+    def call_endpoint(name: str) -> tuple[str, pd.DataFrame | None]:
+        try:
+            return name, fast_inference(name, df)
+        except Exception as e:
+            print(f"Error calling endpoint {name}: {e}")
+            return name, None
+
+    with ThreadPoolExecutor(max_workers=len(endpoint_names)) as executor:
+        futures = {executor.submit(call_endpoint, name): name for name in endpoint_names}
+        for future in as_completed(futures):
+            name, result = future.result()
+            results[name] = result
+
+    return results
+
+
+def aggregate_predictions(results: dict[str, pd.DataFrame], model_weights: dict[str, float]) -> pd.DataFrame:
+    """Aggregate predictions from multiple endpoints using model weights.
+
+    Args:
+        results: Dict mapping endpoint_name -> predictions DataFrame
+                 Each DataFrame must have 'prediction' and 'confidence' columns
+        model_weights: Dict mapping endpoint_name -> weight
+
+    Returns:
+        DataFrame with aggregated prediction, prediction_std, and confidence
+    """
+    # Filter out failed endpoints
+    valid_results = {k: v for k, v in results.items() if v is not None}
+    if not valid_results:
+        raise ValueError("All child endpoints failed")
+
+    # Use first result as base (for id columns, etc.)
+    first_df = list(valid_results.values())[0]
+    output_df = first_df.drop(columns=["prediction", "confidence", "prediction_std"], errors="ignore").copy()
+
+    # Build DataFrames of predictions and confidences from all endpoints
+    pred_df = pd.DataFrame({name: df["prediction"] for name, df in valid_results.items()})
+    conf_df = pd.DataFrame({name: df["confidence"] for name, df in valid_results.items()})
+
+    # Apply model weights (renormalize for valid endpoints only)
+    valid_weights = {k: model_weights.get(k, 1.0) for k in valid_results}
+    weight_sum = sum(valid_weights.values())
+    normalized_weights = {k: v / weight_sum for k, v in valid_weights.items()}
+
+    # Weighted average
+    output_df["prediction"] = sum(pred_df[name] * w for name, w in normalized_weights.items())
+
+    # Ensemble std across child endpoints
+    output_df["prediction_std"] = pred_df.std(axis=1)
+
+    # Aggregated confidence: weighted mean of child confidences
+    output_df["confidence"] = sum(conf_df[name] * w for name, w in normalized_weights.items())
+
+    return output_df
+
+
+# =============================================================================
+# Model Loading (for SageMaker inference)
+# =============================================================================
+def model_fn(model_dir: str) -> dict:
+    """Load meta model configuration."""
+    with open(os.path.join(model_dir, "meta_config.json")) as f:
+        config = json.load(f)
+
+    # Set AWS_REGION for fast_inference (baked in at training time)
+    if config.get("aws_region"):
+        os.environ["AWS_REGION"] = config["aws_region"]
+
+    print(f"Meta model loaded: {len(config['child_endpoints'])} child endpoints")
+    print(f"Model weights: {config.get('model_weights')}")
+    print(f"AWS region: {config.get('aws_region')}")
+    return config
+
+
+def input_fn(input_data, content_type):
+    """Parse input data and return a DataFrame."""
+    if not input_data:
+        raise ValueError("Empty input data is not supported!")
+
+    # Decode bytes to string if necessary
+    if isinstance(input_data, bytes):
+        input_data = input_data.decode("utf-8")
+
+    if "text/csv" in content_type:
+        return pd.read_csv(StringIO(input_data))
+    elif "application/json" in content_type:
+        return pd.DataFrame(json.loads(input_data))
+    else:
+        raise ValueError(f"{content_type} not supported!")
+
+
+def output_fn(output_df, accept_type):
+    """Supports both CSV and JSON output formats."""
+    if "text/csv" in accept_type:
+        return output_df.to_csv(index=False), "text/csv"
+    elif "application/json" in accept_type:
+        return output_df.to_json(orient="records"), "application/json"
+    else:
+        raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
+
+
+# =============================================================================
+# Inference (for SageMaker inference)
+# =============================================================================
+def predict_fn(df: pd.DataFrame, config: dict) -> pd.DataFrame:
+    """Run inference by calling child endpoints and aggregating results."""
+    child_endpoints = config["child_endpoints"]
+    model_weights = config.get("model_weights", {})
+
+    print(f"Calling {len(child_endpoints)} child endpoints: {child_endpoints}")
+
+    # Call all child endpoints
+    results = invoke_endpoints_parallel(child_endpoints, df)
+
+    # Report status
+    for name, result in results.items():
+        status = f"{len(result)} rows" if result is not None else "FAILED"
+        print(f"  {name}: {status}")
+
+    # Aggregate predictions using model weights
+    output_df = aggregate_predictions(results, model_weights)
+
+    print(f"Aggregated {len(output_df)} predictions from {len(results)} endpoints")
+    return output_df
+
+
+# =============================================================================
+# Training (just saves configuration - no actual training)
+# =============================================================================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+    parser.add_argument("--output-data-dir", type=str, default=os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data"))
+    parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
+    args = parser.parse_args()
+
+    child_endpoints = TEMPLATE_PARAMS["child_endpoints"]
+    target_column = TEMPLATE_PARAMS["target_column"]
+    model_weights = TEMPLATE_PARAMS["model_weights"]
+    aws_region = TEMPLATE_PARAMS["aws_region"]
+
+    print("=" * 60)
+    print("Meta Model Configuration")
+    print("=" * 60)
+    print(f"Child endpoints: {child_endpoints}")
+    print(f"Target column: {target_column}")
+    print(f"Model weights: {model_weights}")
+    print(f"AWS region: {aws_region}")
+
+    # Save configuration for inference
+    config = {
+        "child_endpoints": child_endpoints,
+        "target_column": target_column,
+        "model_weights": model_weights,
+        "aws_region": aws_region,
+    }
+
+    with open(os.path.join(args.model_dir, "meta_config.json"), "w") as f:
+        json.dump(config, f, indent=2)
+
+    print(f"\nMeta model configuration saved to {args.model_dir}")