workbench 0.8.234__py3-none-any.whl → 0.8.239__py3-none-any.whl

workbench/model_scripts/xgb_model/model_script_utils.py

@@ -16,6 +16,7 @@ from sklearn.metrics import (
     r2_score,
     root_mean_squared_error,
 )
+from sklearn.model_selection import GroupKFold, GroupShuffleSplit
 from scipy.stats import spearmanr
 
 
@@ -367,3 +368,227 @@ def print_confusion_matrix(y_true: np.ndarray, y_pred: np.ndarray, label_names:
         for j, col_name in enumerate(label_names):
             value = conf_mtx[i, j]
             print(f"ConfusionMatrix:{row_name}:{col_name} {value}")
+
+
+# =============================================================================
+# Dataset Splitting Utilities for Molecular Data
+# =============================================================================
+def get_scaffold(smiles: str) -> str:
+    """Extract Bemis-Murcko scaffold from a SMILES string.
+
+    Args:
+        smiles: SMILES string of the molecule
+
+    Returns:
+        SMILES string of the scaffold, or empty string if molecule is invalid
+    """
+    from rdkit import Chem
+    from rdkit.Chem.Scaffolds import MurckoScaffold
+
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return ""
+    try:
+        scaffold = MurckoScaffold.GetScaffoldForMol(mol)
+        return Chem.MolToSmiles(scaffold)
+    except Exception:
+        return ""
+
+
+def get_scaffold_groups(smiles_list: list[str]) -> np.ndarray:
+    """Assign each molecule to a scaffold group.
+
+    Args:
+        smiles_list: List of SMILES strings
+
+    Returns:
+        Array of group indices (same scaffold = same group)
+    """
+    scaffold_to_group = {}
+    groups = []
+
+    for smi in smiles_list:
+        scaffold = get_scaffold(smi)
+        if scaffold not in scaffold_to_group:
+            scaffold_to_group[scaffold] = len(scaffold_to_group)
+        groups.append(scaffold_to_group[scaffold])
+
+    n_scaffolds = len(scaffold_to_group)
+    print(f"Found {n_scaffolds} unique scaffolds from {len(smiles_list)} molecules")
+    return np.array(groups)
+
+
+def get_butina_clusters(smiles_list: list[str], cutoff: float = 0.4) -> np.ndarray:
+    """Cluster molecules using Butina algorithm on Morgan fingerprints.
+
+    Uses RDKit's Butina clustering with Tanimoto distance on Morgan fingerprints.
+    This is Pat Walters' recommended approach for creating diverse train/test splits.
+
+    Args:
+        smiles_list: List of SMILES strings
+        cutoff: Tanimoto distance cutoff for clustering (default 0.4)
+               Lower values = more clusters = more similar molecules per cluster
+
+    Returns:
+        Array of cluster indices
+    """
+    from rdkit import Chem, DataStructs
+    from rdkit.Chem.rdFingerprintGenerator import GetMorganGenerator
+    from rdkit.ML.Cluster import Butina
+
+    # Create Morgan fingerprint generator
+    fp_gen = GetMorganGenerator(radius=2, fpSize=2048)
+
+    # Generate Morgan fingerprints
+    fps = []
+    valid_indices = []
+    for i, smi in enumerate(smiles_list):
+        mol = Chem.MolFromSmiles(smi)
+        if mol is not None:
+            fp = fp_gen.GetFingerprint(mol)
+            fps.append(fp)
+            valid_indices.append(i)
+
+    if len(fps) == 0:
+        raise ValueError("No valid molecules found for clustering")
+
+    # Compute distance matrix (upper triangle only for efficiency)
+    n = len(fps)
+    dists = []
+    for i in range(1, n):
+        sims = DataStructs.BulkTanimotoSimilarity(fps[i], fps[:i])
+        dists.extend([1 - s for s in sims])
+
+    # Butina clustering
+    clusters = Butina.ClusterData(dists, n, cutoff, isDistData=True)
+
+    # Map back to original indices
+    cluster_labels = np.zeros(len(smiles_list), dtype=int)
+    for cluster_idx, cluster in enumerate(clusters):
+        for mol_idx in cluster:
+            original_idx = valid_indices[mol_idx]
+            cluster_labels[original_idx] = cluster_idx
+
+    # Assign invalid molecules to their own clusters
+    next_cluster = len(clusters)
+    for i in range(len(smiles_list)):
+        if i not in valid_indices:
+            cluster_labels[i] = next_cluster
+            next_cluster += 1
+
+    n_clusters = len(set(cluster_labels))
+    print(f"Butina clustering: {n_clusters} clusters from {len(smiles_list)} molecules (cutoff={cutoff})")
+    return cluster_labels
+
+
+def _find_smiles_column(columns: list[str]) -> str | None:
+    """Find SMILES column (case-insensitive match for 'smiles').
+
+    Args:
+        columns: List of column names
+
+    Returns:
+        The matching column name, or None if not found
+    """
+    return next((c for c in columns if c.lower() == "smiles"), None)
+
+
+def get_split_indices(
+    df: pd.DataFrame,
+    n_splits: int = 5,
+    strategy: str = "random",
+    smiles_column: str | None = None,
+    target_column: str | None = None,
+    test_size: float = 0.2,
+    random_state: int = 42,
+    butina_cutoff: float = 0.4,
+) -> list[tuple[np.ndarray, np.ndarray]]:
+    """Get train/validation split indices using various strategies.
+
+    This is a unified interface for generating splits that can be used across
+    all model templates (XGBoost, PyTorch, ChemProp).
+
+    Args:
+        df: DataFrame containing the data
+        n_splits: Number of CV folds (1 = single train/val split)
+        strategy: Split strategy - one of:
+            - "random": Standard random split (default sklearn behavior)
+            - "scaffold": Bemis-Murcko scaffold-based grouping
+            - "butina": Morgan fingerprint clustering (recommended for ADMET)
+        smiles_column: Column containing SMILES. If None, auto-detects 'smiles' (case-insensitive)
+        target_column: Column containing target values (for stratification, optional)
+        test_size: Fraction for validation set when n_splits=1 (default 0.2)
+        random_state: Random seed for reproducibility
+        butina_cutoff: Tanimoto distance cutoff for Butina clustering (default 0.4)
+
+    Returns:
+        List of (train_indices, val_indices) tuples
+
+    Note:
+        If scaffold/butina strategy is requested but no SMILES column is found,
+        automatically falls back to random split with a warning message.
+
+    Example:
+        >>> folds = get_split_indices(df, n_splits=5, strategy="scaffold")
+        >>> for train_idx, val_idx in folds:
+        ...     X_train, X_val = df.iloc[train_idx], df.iloc[val_idx]
+    """
+    from sklearn.model_selection import KFold, StratifiedKFold, train_test_split
+
+    n_samples = len(df)
+
+    # Random split (original behavior)
+    if strategy == "random":
+        if n_splits == 1:
+            indices = np.arange(n_samples)
+            train_idx, val_idx = train_test_split(indices, test_size=test_size, random_state=random_state)
+            return [(train_idx, val_idx)]
+        else:
+            if target_column and df[target_column].dtype in ["object", "category", "bool"]:
+                kfold = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
+                return list(kfold.split(df, df[target_column]))
+            else:
+                kfold = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
+                return list(kfold.split(df))
+
+    # Scaffold or Butina split requires SMILES - auto-detect if not provided
+    if smiles_column is None:
+        smiles_column = _find_smiles_column(df.columns.tolist())
+
+    # Fall back to random split if no SMILES column available
+    if smiles_column is None or smiles_column not in df.columns:
+        print(f"No 'smiles' column found for strategy='{strategy}', falling back to random split")
+        return get_split_indices(
+            df,
+            n_splits=n_splits,
+            strategy="random",
+            target_column=target_column,
+            test_size=test_size,
+            random_state=random_state,
+        )
+
+    smiles_list = df[smiles_column].tolist()
+
+    # Get group assignments
+    if strategy == "scaffold":
+        groups = get_scaffold_groups(smiles_list)
+    elif strategy == "butina":
+        groups = get_butina_clusters(smiles_list, cutoff=butina_cutoff)
+    else:
+        raise ValueError(f"Unknown strategy: {strategy}. Use 'random', 'scaffold', or 'butina'")
+
+    # Generate splits using GroupKFold or GroupShuffleSplit
+    if n_splits == 1:
+        # Single split: use GroupShuffleSplit
+        splitter = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=random_state)
+        return list(splitter.split(df, groups=groups))
+    else:
+        # K-fold: use GroupKFold (ensures no group appears in both train and val)
+        # Note: GroupKFold doesn't shuffle, so we shuffle group order first
+        unique_groups = np.unique(groups)
+        rng = np.random.default_rng(random_state)
+        shuffled_group_map = {g: i for i, g in enumerate(rng.permutation(unique_groups))}
+        shuffled_groups = np.array([shuffled_group_map[g] for g in groups])
+
+        gkf = GroupKFold(n_splits=n_splits)
+        return list(gkf.split(df, groups=shuffled_groups))

workbench/model_scripts/xgb_model/uq_harness.py

@@ -3,6 +3,25 @@
 This module provides a reusable UQ harness that can wrap any point predictor model
 (XGBoost, PyTorch, ChemProp, etc.) to provide calibrated prediction intervals.
 
+Features:
+    - Conformalized Quantile Regression (CQR) for distribution-free coverage guarantees
+    - Multiple confidence levels (50%, 68%, 80%, 90%, 95%)
+    - Confidence scoring based on interval width
+
+Why CQR without additional Z-scaling:
+    MAPIE's conformalization step already guarantees that prediction intervals achieve
+    their target coverage on the calibration set. For example, an 80% CI will contain
+    ~80% of true values. This is the core promise of conformal prediction.
+
+    Z-scaling (post-hoc interval adjustment) would only help if there's a distribution
+    shift between calibration and test data. However:
+    1. We'd compute Z-scale on the same calibration set MAPIE uses, making it redundant
+    2. Our cross-fold validation metrics confirm coverage is already well-calibrated
+    3. Adding Z-scaling would "second-guess" MAPIE's principled conformalization
+
+    Empirically, our models achieve excellent coverage (e.g., 80% CI → 80.1% coverage),
+    validating that MAPIE's approach is sufficient without additional calibration.
+
 Usage:
     # Training
     uq_models, uq_metadata = train_uq_models(X_train, y_train, X_val, y_val)
@@ -240,38 +259,37 @@ def compute_confidence(
     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)
+    Confidence is derived from the 80% prediction interval (q_10 to q_90) width:
+    - Narrower intervals → higher confidence (model is more certain)
+    - Wider intervals → lower confidence (model is less certain)
+
+    Why 80% CI (q_10/q_90)?
+        - 68% CI is too narrow and sensitive to noise
+        - 95% CI is too wide and less discriminating between samples
+        - 80% provides a good balance for ranking prediction reliability
+
+    Formula: confidence = exp(-width / median_width)
+        - When width equals median, confidence ≈ 0.37
+        - When width is half median, confidence ≈ 0.61
+        - When width is double median, confidence ≈ 0.14
+
+    This exponential decay is a common choice for converting uncertainty to
+    confidence scores, providing a smooth mapping that appropriately penalizes
+    high-uncertainty predictions.
 
     Args:
-        df: DataFrame with 'prediction', 'q_50', and quantile columns
+        df: DataFrame with quantile columns from predict_intervals()
         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
+        DataFrame with added 'confidence' column (values between 0 and 1)
     """
-    # 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))
+    df["confidence"] = np.exp(-interval_width / median_interval_width)
 
     return df

workbench/model_scripts/xgb_model/xgb_model.template

@@ -51,12 +51,18 @@ DEFAULT_HYPERPARAMETERS = {
     "gamma": 0.1,
     "reg_alpha": 0.1,
     "reg_lambda": 1.0,
+    # Split strategy: "random", "scaffold", or "butina"
+    # - random: Standard random split (default)
+    # - scaffold: Bemis-Murcko scaffold-based grouping (requires 'smiles' column in data)
+    # - butina: Morgan fingerprint clustering (requires 'smiles' column, recommended for ADMET)
+    "split_strategy": "random",
+    "butina_cutoff": 0.4,  # Tanimoto distance cutoff for Butina clustering
     # Random seed
     "seed": 42,
 }
 
 # Workbench-specific parameters (not passed to XGBoost)
-WORKBENCH_PARAMS = {"n_folds"}
+WORKBENCH_PARAMS = {"n_folds", "split_strategy", "butina_cutoff"}
 
 # Regression-only parameters (filtered out for classifiers)
 REGRESSION_ONLY_PARAMS = {"objective"}
@@ -198,6 +204,7 @@ if __name__ == "__main__":
         check_dataframe,
         compute_classification_metrics,
         compute_regression_metrics,
+        get_split_indices,
         print_classification_metrics,
         print_confusion_matrix,
         print_regression_metrics,
@@ -273,25 +280,29 @@ if __name__ == "__main__":
 
     print(f"XGBoost params: {xgb_params}")
 
-    if n_folds == 1:
-        # Single train/val split
-        if "training" in all_df.columns:
-            print("Using 'training' column for train/val split")
-            train_idx = np.where(all_df["training"])[0]
-            val_idx = np.where(~all_df["training"])[0]
-        else:
-            print("WARNING: No 'training' column found, using random 80/20 split")
-            indices = np.arange(len(all_df))
-            train_idx, val_idx = train_test_split(indices, test_size=0.2, random_state=42)
+    # Get split strategy parameters
+    split_strategy = hyperparameters.get("split_strategy", "random")
+    butina_cutoff = hyperparameters.get("butina_cutoff", 0.4)
+
+    # Check for pre-defined training column (overrides split strategy)
+    if n_folds == 1 and "training" in all_df.columns:
+        print("Using 'training' column for train/val split")
+        train_idx = np.where(all_df["training"])[0]
+        val_idx = np.where(~all_df["training"])[0]
         folds = [(train_idx, val_idx)]
     else:
-        # K-fold cross-validation
-        if model_type == "classifier":
-            kfold = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=42)
-            folds = list(kfold.split(all_df, all_df[target]))
-        else:
-            kfold = KFold(n_splits=n_folds, shuffle=True, random_state=42)
-            folds = list(kfold.split(all_df))
+        # Use unified split interface (auto-detects 'smiles' column for scaffold/butina)
+        target_col = target if model_type == "classifier" else None
+        folds = get_split_indices(
+            all_df,
+            n_splits=n_folds,
+            strategy=split_strategy,
+            target_column=target_col,
+            test_size=0.2,
+            random_state=42,
+            butina_cutoff=butina_cutoff,
+        )
+        print(f"Split strategy: {split_strategy}")
 
     print(f"Training {'single model' if n_folds == 1 else f'{n_folds}-fold ensemble'}...")
 

workbench/scripts/ml_pipeline_batch.py

@@ -44,7 +44,14 @@ def _log_cloudwatch_link(job: dict, message_prefix: str = "View logs") -> None:
         log.info("Check AWS Batch console for logs")
 
 
-def run_batch_job(script_path: str, size: str = "small") -> int:
+def run_batch_job(
+    script_path: str,
+    size: str = "small",
+    realtime: bool = False,
+    dt: bool = False,
+    promote: bool = False,
+    test_promote: bool = False,
+) -> int:
     """
     Submit and monitor an AWS Batch job for ML pipeline execution.
 
@@ -56,6 +63,10 @@ def run_batch_job(script_path: str, size: str = "small") -> int:
           - small: 2 vCPU, 4GB RAM for lightweight processing
           - medium: 4 vCPU, 8GB RAM for standard ML workloads
           - large: 8 vCPU, 16GB RAM for heavy training/inference
+        realtime: If True, sets serverless=False for real-time processing (default: False)
+        dt: If True, sets DT=True in environment (default: False)
+        promote: If True, sets PROMOTE=True in environment (default: False)
+        test_promote: If True, sets TEST_PROMOTE=True in environment (default: False)
 
     Returns:
         Exit code (0 for success/disconnected, non-zero for failure)
@@ -81,6 +92,10 @@ def run_batch_job(script_path: str, size: str = "small") -> int:
             "environment": [
                 {"name": "ML_PIPELINE_S3_PATH", "value": s3_path},
                 {"name": "WORKBENCH_BUCKET", "value": workbench_bucket},
+                {"name": "SERVERLESS", "value": "False" if realtime else "True"},
+                {"name": "DT", "value": str(dt)},
+                {"name": "PROMOTE", "value": str(promote)},
+                {"name": "TEST_PROMOTE", "value": str(test_promote)},
             ]
         },
     )
@@ -124,9 +139,39 @@ def main():
     """CLI entry point for running ML pipelines on AWS Batch."""
     parser = argparse.ArgumentParser(description="Run ML pipeline script on AWS Batch")
     parser.add_argument("script_file", help="Local path to ML pipeline script")
+    parser.add_argument(
+        "--size", default="small", choices=["small", "medium", "large"], help="Job size tier (default: small)"
+    )
+    parser.add_argument(
+        "--realtime",
+        action="store_true",
+        help="Create realtime endpoints (default is serverless)",
+    )
+    parser.add_argument(
+        "--dt",
+        action="store_true",
+        help="Set DT=True (models and endpoints will have '-dt' suffix)",
+    )
+    parser.add_argument(
+        "--promote",
+        action="store_true",
+        help="Set Promote=True (models and endpoints will use promoted naming)",
+    )
+    parser.add_argument(
+        "--test-promote",
+        action="store_true",
+        help="Set TEST_PROMOTE=True (creates test endpoint with '-test' suffix)",
+    )
     args = parser.parse_args()
     try:
-        exit_code = run_batch_job(args.script_file)
+        exit_code = run_batch_job(
+            args.script_file,
+            size=args.size,
+            realtime=args.realtime,
+            dt=args.dt,
+            promote=args.promote,
+            test_promote=args.test_promote,
+        )
         exit(exit_code)
     except Exception as e:
         log.error(f"Error: {e}")