workbench 0.8.205__py3-none-any.whl → 0.8.212__py3-none-any.whl

workbench/model_scripts/xgb_model/model_script_utils.py

@@ -0,0 +1,335 @@
+"""Shared utility functions for model training scripts (templates).
+
+These functions are used across multiple model templates (XGBoost, PyTorch, ChemProp)
+to reduce code duplication and ensure consistent behavior.
+"""
+
+from io import StringIO
+import json
+import numpy as np
+import pandas as pd
+from sklearn.metrics import (
+    confusion_matrix,
+    mean_absolute_error,
+    median_absolute_error,
+    precision_recall_fscore_support,
+    r2_score,
+    root_mean_squared_error,
+)
+from scipy.stats import spearmanr
+
+
+def check_dataframe(df: pd.DataFrame, df_name: str) -> None:
+    """Check if the provided dataframe is empty and raise an exception if it is.
+
+    Args:
+        df: DataFrame to check
+        df_name: Name of the DataFrame (for error message)
+
+    Raises:
+        ValueError: If the DataFrame is empty
+    """
+    if df.empty:
+        msg = f"*** The training data {df_name} has 0 rows! ***STOPPING***"
+        print(msg)
+        raise ValueError(msg)
+
+
+def expand_proba_column(df: pd.DataFrame, class_labels: list[str]) -> pd.DataFrame:
+    """Expands a column containing a list of probabilities into separate columns.
+
+    Handles None values for rows where predictions couldn't be made.
+
+    Args:
+        df: DataFrame containing a "pred_proba" column
+        class_labels: List of class labels
+
+    Returns:
+        DataFrame with the "pred_proba" expanded into separate columns (e.g., "class1_proba")
+
+    Raises:
+        ValueError: If DataFrame does not contain a "pred_proba" column
+    """
+    proba_column = "pred_proba"
+    if proba_column not in df.columns:
+        raise ValueError('DataFrame does not contain a "pred_proba" column')
+
+    proba_splits = [f"{label}_proba" for label in class_labels]
+    n_classes = len(class_labels)
+
+    # Handle None values by replacing with list of NaNs
+    proba_values = []
+    for val in df[proba_column]:
+        if val is None:
+            proba_values.append([np.nan] * n_classes)
+        else:
+            proba_values.append(val)
+
+    proba_df = pd.DataFrame(proba_values, columns=proba_splits)
+
+    # Drop any existing proba columns and reset index for concat
+    df = df.drop(columns=[proba_column] + proba_splits, errors="ignore")
+    df = df.reset_index(drop=True)
+    df = pd.concat([df, proba_df], axis=1)
+    return df
+
+
+def match_features_case_insensitive(df: pd.DataFrame, model_features: list[str]) -> pd.DataFrame:
+    """Matches and renames DataFrame columns to match model feature names (case-insensitive).
+
+    Prioritizes exact matches, then case-insensitive matches.
+
+    Args:
+        df: Input DataFrame
+        model_features: List of feature names expected by the model
+
+    Returns:
+        DataFrame with columns renamed to match model features
+
+    Raises:
+        ValueError: If any model features cannot be matched
+    """
+    df_columns_lower = {col.lower(): col for col in df.columns}
+    rename_dict = {}
+    missing = []
+    for feature in model_features:
+        if feature in df.columns:
+            continue  # Exact match
+        elif feature.lower() in df_columns_lower:
+            rename_dict[df_columns_lower[feature.lower()]] = feature
+        else:
+            missing.append(feature)
+
+    if missing:
+        raise ValueError(f"Features not found: {missing}")
+
+    return df.rename(columns=rename_dict)
+
+
+def convert_categorical_types(
+    df: pd.DataFrame, features: list[str], category_mappings: dict[str, list[str]] | None = None
+) -> tuple[pd.DataFrame, dict[str, list[str]]]:
+    """Converts appropriate columns to categorical type with consistent mappings.
+
+    In training mode (category_mappings is None or empty), detects object/string columns
+    with <20 unique values and converts them to categorical.
+    In inference mode (category_mappings provided), applies the stored mappings.
+
+    Args:
+        df: The DataFrame to process
+        features: List of feature names to consider for conversion
+        category_mappings: Existing category mappings. If None or empty, training mode.
+                          If populated, inference mode.
+
+    Returns:
+        Tuple of (processed DataFrame, category mappings dictionary)
+    """
+    if category_mappings is None:
+        category_mappings = {}
+
+    # Training mode
+    if not category_mappings:
+        for col in df.select_dtypes(include=["object", "string"]):
+            if col in features and df[col].nunique() < 20:
+                print(f"Training mode: Converting {col} to category")
+                df[col] = df[col].astype("category")
+                category_mappings[col] = df[col].cat.categories.tolist()
+
+    # Inference mode
+    else:
+        for col, categories in category_mappings.items():
+            if col in df.columns:
+                print(f"Inference mode: Applying categorical mapping for {col}")
+                df[col] = pd.Categorical(df[col], categories=categories)
+
+    return df, category_mappings
+
+
+def decompress_features(
+    df: pd.DataFrame, features: list[str], compressed_features: list[str]
+) -> tuple[pd.DataFrame, list[str]]:
+    """Decompress bitstring features into individual bit columns.
+
+    Args:
+        df: The features DataFrame
+        features: Full list of feature names
+        compressed_features: List of feature names to decompress (bitstrings)
+
+    Returns:
+        Tuple of (DataFrame with decompressed features, updated feature list)
+    """
+    # 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."
+        )
+
+    # Make a copy to avoid mutating the original list
+    decompressed_features = features.copy()
+
+    for feature in compressed_features:
+        if (feature not in df.columns) or (feature not in decompressed_features):
+            print(f"Feature '{feature}' not in the features list, skipping decompression.")
+            continue
+
+        # Remove the feature from the list 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 input_fn(input_data, content_type: str) -> pd.DataFrame:
+    """Parse input data and return a DataFrame.
+
+    Args:
+        input_data: Raw input data (bytes or string)
+        content_type: MIME type of the input data
+
+    Returns:
+        Parsed DataFrame
+
+    Raises:
+        ValueError: If input is empty or content_type is not supported
+    """
+    if not input_data:
+        raise ValueError("Empty input data is not supported!")
+
+    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: pd.DataFrame, accept_type: str) -> tuple[str, str]:
+    """Convert output DataFrame to requested format.
+
+    Args:
+        output_df: DataFrame to convert
+        accept_type: Requested MIME type
+
+    Returns:
+        Tuple of (formatted output string, MIME type)
+
+    Raises:
+        RuntimeError: If accept_type is not supported
+    """
+    if "text/csv" in accept_type:
+        csv_output = output_df.fillna("N/A").to_csv(index=False)
+        return csv_output, "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.")
+
+
+def compute_regression_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, float]:
+    """Compute standard regression metrics.
+
+    Args:
+        y_true: Ground truth target values
+        y_pred: Predicted values
+
+    Returns:
+        Dictionary with keys: rmse, mae, medae, r2, spearmanr, support
+    """
+    return {
+        "rmse": root_mean_squared_error(y_true, y_pred),
+        "mae": mean_absolute_error(y_true, y_pred),
+        "medae": median_absolute_error(y_true, y_pred),
+        "r2": r2_score(y_true, y_pred),
+        "spearmanr": spearmanr(y_true, y_pred).correlation,
+        "support": len(y_true),
+    }
+
+
+def print_regression_metrics(metrics: dict[str, float]) -> None:
+    """Print regression metrics in the format expected by SageMaker metric definitions.
+
+    Args:
+        metrics: Dictionary of metric name -> value
+    """
+    print(f"rmse: {metrics['rmse']:.3f}")
+    print(f"mae: {metrics['mae']:.3f}")
+    print(f"medae: {metrics['medae']:.3f}")
+    print(f"r2: {metrics['r2']:.3f}")
+    print(f"spearmanr: {metrics['spearmanr']:.3f}")
+    print(f"support: {metrics['support']}")
+
+
+def compute_classification_metrics(
+    y_true: np.ndarray, y_pred: np.ndarray, label_names: list[str], target_col: str
+) -> pd.DataFrame:
+    """Compute per-class classification metrics.
+
+    Args:
+        y_true: Ground truth labels
+        y_pred: Predicted labels
+        label_names: List of class label names
+        target_col: Name of the target column (for DataFrame output)
+
+    Returns:
+        DataFrame with columns: target_col, precision, recall, f1, support
+    """
+    scores = precision_recall_fscore_support(y_true, y_pred, average=None, labels=label_names)
+    return pd.DataFrame(
+        {
+            target_col: label_names,
+            "precision": scores[0],
+            "recall": scores[1],
+            "f1": scores[2],
+            "support": scores[3],
+        }
+    )
+
+
+def print_classification_metrics(score_df: pd.DataFrame, target_col: str, label_names: list[str]) -> None:
+    """Print per-class classification metrics in the format expected by SageMaker.
+
+    Args:
+        score_df: DataFrame from compute_classification_metrics
+        target_col: Name of the target column
+        label_names: List of class label names
+    """
+    metrics = ["precision", "recall", "f1", "support"]
+    for t in label_names:
+        for m in metrics:
+            value = score_df.loc[score_df[target_col] == t, m].iloc[0]
+            print(f"Metrics:{t}:{m} {value}")
+
+
+def print_confusion_matrix(y_true: np.ndarray, y_pred: np.ndarray, label_names: list[str]) -> None:
+    """Print confusion matrix in the format expected by SageMaker.
+
+    Args:
+        y_true: Ground truth labels
+        y_pred: Predicted labels
+        label_names: List of class label names
+    """
+    conf_mtx = confusion_matrix(y_true, y_pred, labels=label_names)
+    for i, row_name in enumerate(label_names):
+        for j, col_name in enumerate(label_names):
+            value = conf_mtx[i, j]
+            print(f"ConfusionMatrix:{row_name}:{col_name} {value}")

workbench/model_scripts/xgb_model/uq_harness.py

@@ -0,0 +1,278 @@
+"""UQ Harness: Uncertainty Quantification using MAPIE Conformalized Quantile Regression.
+
+This module provides a reusable UQ harness that can wrap any point predictor model
+(XGBoost, PyTorch, ChemProp, etc.) to provide calibrated prediction intervals.
+
+Usage:
+    # Training
+    uq_models, uq_metadata = train_uq_models(X_train, y_train, X_val, y_val)
+    save_uq_models(uq_models, uq_metadata, model_dir)
+
+    # Inference
+    uq_models, uq_metadata = load_uq_models(model_dir)
+    df = predict_intervals(df, X, uq_models, uq_metadata)
+    df = compute_confidence(df, uq_metadata["median_interval_width"])
+"""
+
+import json
+import os
+import numpy as np
+import pandas as pd
+import joblib
+from lightgbm import LGBMRegressor
+from mapie.regression import ConformalizedQuantileRegressor
+
+
+# Default confidence levels for prediction intervals
+DEFAULT_CONFIDENCE_LEVELS = [0.50, 0.68, 0.80, 0.90, 0.95]
+
+
+def train_uq_models(
+    X_train: pd.DataFrame | np.ndarray,
+    y_train: pd.Series | np.ndarray,
+    X_val: pd.DataFrame | np.ndarray,
+    y_val: pd.Series | np.ndarray,
+    confidence_levels: list[float] | None = None,
+) -> tuple[dict, dict]:
+    """Train MAPIE UQ models for multiple confidence levels.
+
+    Args:
+        X_train: Training features
+        y_train: Training targets
+        X_val: Validation features for conformalization
+        y_val: Validation targets for conformalization
+        confidence_levels: List of confidence levels (default: [0.50, 0.68, 0.80, 0.90, 0.95])
+
+    Returns:
+        Tuple of (uq_models dict, uq_metadata dict)
+    """
+    if confidence_levels is None:
+        confidence_levels = DEFAULT_CONFIDENCE_LEVELS
+
+    mapie_models = {}
+
+    for confidence_level in confidence_levels:
+        alpha = 1 - confidence_level
+        lower_q = alpha / 2
+        upper_q = 1 - alpha / 2
+
+        print(f"\nTraining quantile models for {confidence_level * 100:.0f}% confidence interval...")
+        print(f"  Quantiles: {lower_q:.3f}, {upper_q:.3f}, 0.500")
+
+        # Train three LightGBM quantile models for this confidence level
+        quantile_estimators = []
+        for q in [lower_q, upper_q, 0.5]:
+            print(f"    Training model for quantile {q:.3f}...")
+            est = LGBMRegressor(
+                objective="quantile",
+                alpha=q,
+                n_estimators=1000,
+                max_depth=6,
+                learning_rate=0.01,
+                num_leaves=31,
+                min_child_samples=20,
+                subsample=0.8,
+                colsample_bytree=0.8,
+                random_state=42,
+                verbose=-1,
+                force_col_wise=True,
+            )
+            est.fit(X_train, y_train)
+            quantile_estimators.append(est)
+
+        # Create MAPIE CQR model for this confidence level
+        print(f"  Setting up MAPIE CQR for {confidence_level * 100:.0f}% confidence...")
+        mapie_model = ConformalizedQuantileRegressor(
+            quantile_estimators, confidence_level=confidence_level, prefit=True
+        )
+
+        # Conformalize the model with validation data
+        print("  Conformalizing with validation data...")
+        mapie_model.conformalize(X_val, y_val)
+
+        # Store the model
+        model_name = f"mapie_{confidence_level:.2f}"
+        mapie_models[model_name] = mapie_model
+
+        # Validate coverage for this confidence level
+        y_pred, y_pis = mapie_model.predict_interval(X_val)
+        coverage = np.mean((y_val >= y_pis[:, 0, 0]) & (y_val <= y_pis[:, 1, 0]))
+        print(f"  Coverage: Target={confidence_level * 100:.0f}%, Empirical={coverage * 100:.1f}%")
+
+    # Compute median interval width for confidence calculation (using 80% CI = q_10 to q_90)
+    print("\nComputing normalization statistics for confidence scores...")
+    model_80 = mapie_models["mapie_0.80"]
+    _, y_pis_80 = model_80.predict_interval(X_val)
+    interval_width = np.abs(y_pis_80[:, 1, 0] - y_pis_80[:, 0, 0])
+    median_interval_width = float(np.median(interval_width))
+    print(f"  Median interval width (q_10-q_90): {median_interval_width:.6f}")
+
+    # Analyze interval widths across confidence levels
+    print("\nInterval Width Analysis:")
+    for conf_level in confidence_levels:
+        model = mapie_models[f"mapie_{conf_level:.2f}"]
+        _, y_pis = model.predict_interval(X_val)
+        widths = y_pis[:, 1, 0] - y_pis[:, 0, 0]
+        print(f"  {conf_level * 100:.0f}% CI: Mean width={np.mean(widths):.3f}, Std={np.std(widths):.3f}")
+
+    uq_metadata = {
+        "confidence_levels": confidence_levels,
+        "median_interval_width": median_interval_width,
+    }
+
+    return mapie_models, uq_metadata
+
+
+def save_uq_models(uq_models: dict, uq_metadata: dict, model_dir: str) -> None:
+    """Save UQ models and metadata to disk.
+
+    Args:
+        uq_models: Dictionary of MAPIE models keyed by name (e.g., "mapie_0.80")
+        uq_metadata: Dictionary with confidence_levels and median_interval_width
+        model_dir: Directory to save models
+    """
+    # Save each MAPIE model
+    for model_name, model in uq_models.items():
+        joblib.dump(model, os.path.join(model_dir, f"{model_name}.joblib"))
+
+    # Save median interval width
+    with open(os.path.join(model_dir, "median_interval_width.json"), "w") as fp:
+        json.dump(uq_metadata["median_interval_width"], fp)
+
+    # Save UQ metadata
+    with open(os.path.join(model_dir, "uq_metadata.json"), "w") as fp:
+        json.dump(uq_metadata, fp, indent=2)
+
+    print(f"Saved {len(uq_models)} UQ models to {model_dir}")
+
+
+def load_uq_models(model_dir: str) -> tuple[dict, dict]:
+    """Load UQ models and metadata from disk.
+
+    Args:
+        model_dir: Directory containing saved models
+
+    Returns:
+        Tuple of (uq_models dict, uq_metadata dict)
+    """
+    # Load UQ metadata
+    uq_metadata_path = os.path.join(model_dir, "uq_metadata.json")
+    if os.path.exists(uq_metadata_path):
+        with open(uq_metadata_path) as fp:
+            uq_metadata = json.load(fp)
+    else:
+        # Fallback for older models that only have median_interval_width.json
+        uq_metadata = {"confidence_levels": DEFAULT_CONFIDENCE_LEVELS}
+        median_width_path = os.path.join(model_dir, "median_interval_width.json")
+        if os.path.exists(median_width_path):
+            with open(median_width_path) as fp:
+                uq_metadata["median_interval_width"] = json.load(fp)
+
+    # Load all MAPIE models
+    uq_models = {}
+    for conf_level in uq_metadata["confidence_levels"]:
+        model_name = f"mapie_{conf_level:.2f}"
+        model_path = os.path.join(model_dir, f"{model_name}.joblib")
+        if os.path.exists(model_path):
+            uq_models[model_name] = joblib.load(model_path)
+
+    return uq_models, uq_metadata
+
+
+def predict_intervals(
+    df: pd.DataFrame,
+    X: pd.DataFrame | np.ndarray,
+    uq_models: dict,
+    uq_metadata: dict,
+) -> pd.DataFrame:
+    """Add prediction intervals to a DataFrame.
+
+    Args:
+        df: DataFrame to add interval columns to
+        X: Features for prediction (must match training features)
+        uq_models: Dictionary of MAPIE models
+        uq_metadata: Dictionary with confidence_levels
+
+    Returns:
+        DataFrame with added quantile columns (q_025, q_05, ..., q_975)
+    """
+    confidence_levels = uq_metadata["confidence_levels"]
+
+    for conf_level in confidence_levels:
+        model_name = f"mapie_{conf_level:.2f}"
+        model = uq_models[model_name]
+
+        # Get conformalized predictions
+        y_pred, y_pis = model.predict_interval(X)
+
+        # Map confidence levels to quantile column names
+        if conf_level == 0.50:  # 50% CI
+            df["q_25"] = y_pis[:, 0, 0]
+            df["q_75"] = y_pis[:, 1, 0]
+            df["q_50"] = y_pred  # Median prediction
+        elif conf_level == 0.68:  # 68% CI (~1 std)
+            df["q_16"] = y_pis[:, 0, 0]
+            df["q_84"] = y_pis[:, 1, 0]
+        elif conf_level == 0.80:  # 80% CI
+            df["q_10"] = y_pis[:, 0, 0]
+            df["q_90"] = y_pis[:, 1, 0]
+        elif conf_level == 0.90:  # 90% CI
+            df["q_05"] = y_pis[:, 0, 0]
+            df["q_95"] = y_pis[:, 1, 0]
+        elif conf_level == 0.95:  # 95% CI
+            df["q_025"] = y_pis[:, 0, 0]
+            df["q_975"] = y_pis[:, 1, 0]
+
+    # Calculate pseudo-standard deviation from the 68% interval width
+    if "q_84" in df.columns and "q_16" in df.columns:
+        df["prediction_std"] = (df["q_84"] - df["q_16"]).abs() / 2.0
+
+    # Reorder quantile columns for easier reading
+    quantile_cols = ["q_025", "q_05", "q_10", "q_16", "q_25", "q_50", "q_75", "q_84", "q_90", "q_95", "q_975"]
+    existing_q_cols = [c for c in quantile_cols if c in df.columns]
+    other_cols = [c for c in df.columns if c not in quantile_cols]
+    df = df[other_cols + existing_q_cols]
+
+    return df
+
+
+def compute_confidence(
+    df: pd.DataFrame,
+    median_interval_width: float,
+    lower_q: str = "q_10",
+    upper_q: str = "q_90",
+    alpha: float = 1.0,
+    beta: float = 1.0,
+) -> pd.DataFrame:
+    """Compute confidence scores (0.0 to 1.0) based on prediction interval width.
+
+    Uses exponential decay based on:
+    1. Interval width relative to median (alpha weight)
+    2. Distance from median prediction (beta weight)
+
+    Args:
+        df: DataFrame with 'prediction', 'q_50', and quantile columns
+        median_interval_width: Pre-computed median interval width from training data
+        lower_q: Lower quantile column name (default: 'q_10')
+        upper_q: Upper quantile column name (default: 'q_90')
+        alpha: Weight for interval width term (default: 1.0)
+        beta: Weight for distance from median term (default: 1.0)
+
+    Returns:
+        DataFrame with added 'confidence' column
+    """
+    # Interval width
+    interval_width = (df[upper_q] - df[lower_q]).abs()
+
+    # Distance from median, normalized by interval width
+    distance_from_median = (df["prediction"] - df["q_50"]).abs()
+    normalized_distance = distance_from_median / (interval_width + 1e-6)
+
+    # Cap the distance penalty at 1.0
+    normalized_distance = np.minimum(normalized_distance, 1.0)
+
+    # Confidence using exponential decay
+    interval_term = interval_width / median_interval_width
+    df["confidence"] = np.exp(-(alpha * interval_term + beta * normalized_distance))
+
+    return df