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

workbench/model_scripts/chemprop/model_script_utils.py

@@ -0,0 +1,339 @@
+"""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 compressed features (bitstrings or count vectors) into individual columns.
+
+    Supports two formats (auto-detected):
+        - Bitstrings: "10110010..." → individual uint8 columns (0 or 1)
+        - Count vectors: "0,3,0,1,5,..." → individual uint8 columns (0-255)
+
+    Args:
+        df: The features DataFrame
+        features: Full list of feature names
+        compressed_features: List of feature names to decompress
+
+    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)
+
+        # Auto-detect format and parse: comma-separated counts or bitstring
+        sample = str(df[feature].dropna().iloc[0]) if not df[feature].dropna().empty else ""
+        parse_fn = (lambda s: list(map(int, s.split(",")))) if "," in sample else list
+        feature_matrix = np.array([parse_fn(s) for s in df[feature]], dtype=np.uint8)
+
+        # Create new columns with prefix from feature name
+        prefix = feature[:3]
+        new_col_names = [f"{prefix}_{i}" for i in range(feature_matrix.shape[1])]
+        new_df = pd.DataFrame(feature_matrix, columns=new_col_names, index=df.index)
+
+        # Update features list and dataframe
+        decompressed_features.extend(new_col_names)
+        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/chemprop/requirements.txt

@@ -0,0 +1,3 @@
+# Requirements for ChemProp model scripts
+# Note: The training and inference images already have torch and chemprop installed.
+#       So we only need to install packages that are not already included in the images.

workbench/model_scripts/custom_models/chem_info/fingerprints.py

@@ -0,0 +1,175 @@
+"""Molecular fingerprint computation utilities for ADMET modeling.
+
+This module provides Morgan count fingerprints, the standard for ADMET prediction.
+Count fingerprints outperform binary fingerprints for molecular property prediction.
+
+References:
+    - Count vs Binary: https://pubs.acs.org/doi/10.1021/acs.est.3c02198
+    - ECFP/Morgan: https://pubs.acs.org/doi/10.1021/ci100050t
+"""
+
+import logging
+
+import numpy as np
+import pandas as pd
+from rdkit import Chem, RDLogger
+from rdkit.Chem import AllChem
+from rdkit.Chem.MolStandardize import rdMolStandardize
+
+# Suppress RDKit warnings (e.g., "not removing hydrogen atom without neighbors")
+# Keep errors enabled so we see actual problems
+RDLogger.DisableLog("rdApp.warning")
+
+# Set up the logger
+log = logging.getLogger("workbench")
+
+
+def compute_morgan_fingerprints(df: pd.DataFrame, radius: int = 2, n_bits: int = 2048) -> pd.DataFrame:
+    """Compute Morgan count fingerprints for ADMET modeling.
+
+    Generates true count fingerprints where each bit position contains the
+    number of times that substructure appears in the molecule (clamped to 0-255).
+    This is the recommended approach for ADMET prediction per 2025 research.
+
+    Args:
+        df: Input DataFrame containing SMILES strings.
+        radius: Radius for the Morgan fingerprint (default 2 = ECFP4 equivalent).
+        n_bits: Number of bits for the fingerprint (default 2048).
+
+    Returns:
+        pd.DataFrame: Input DataFrame with 'fingerprint' column added.
+                      Values are comma-separated uint8 counts.
+
+    Note:
+        Count fingerprints outperform binary for ADMET prediction.
+        See: https://pubs.acs.org/doi/10.1021/acs.est.3c02198
+    """
+    delete_mol_column = False
+
+    # Check for the SMILES column (case-insensitive)
+    smiles_column = next((col for col in df.columns if col.lower() == "smiles"), None)
+    if smiles_column is None:
+        raise ValueError("Input DataFrame must have a 'smiles' column")
+
+    # Sanity check the molecule column (sometimes it gets serialized, which doesn't work)
+    if "molecule" in df.columns and df["molecule"].dtype == "string":
+        log.warning("Detected serialized molecules in 'molecule' column. Removing...")
+        del df["molecule"]
+
+    # Convert SMILES to RDKit molecule objects
+    if "molecule" not in df.columns:
+        log.info("Converting SMILES to RDKit Molecules...")
+        delete_mol_column = True
+        df["molecule"] = df[smiles_column].apply(Chem.MolFromSmiles)
+        # Make sure our molecules are not None
+        failed_smiles = df[df["molecule"].isnull()][smiles_column].tolist()
+        if failed_smiles:
+            log.warning(f"Failed to convert {len(failed_smiles)} SMILES to molecules ({failed_smiles})")
+        df = df.dropna(subset=["molecule"]).copy()
+
+    # If we have fragments in our compounds, get the largest fragment before computing fingerprints
+    largest_frags = df["molecule"].apply(
+        lambda mol: rdMolStandardize.LargestFragmentChooser().choose(mol) if mol else None
+    )
+
+    def mol_to_count_string(mol):
+        """Convert molecule to comma-separated count fingerprint string."""
+        if mol is None:
+            return pd.NA
+
+        # Get hashed Morgan fingerprint with counts
+        fp = AllChem.GetHashedMorganFingerprint(mol, radius, nBits=n_bits)
+
+        # Initialize array and populate with counts (clamped to uint8 range)
+        counts = np.zeros(n_bits, dtype=np.uint8)
+        for idx, count in fp.GetNonzeroElements().items():
+            counts[idx] = min(count, 255)
+
+        # Return as comma-separated string
+        return ",".join(map(str, counts))
+
+    # Compute Morgan count fingerprints
+    fingerprints = largest_frags.apply(mol_to_count_string)
+
+    # Add the fingerprints to the DataFrame
+    df["fingerprint"] = fingerprints
+
+    # Drop the intermediate 'molecule' column if it was added
+    if delete_mol_column:
+        del df["molecule"]
+
+    return df
+
+
+if __name__ == "__main__":
+    print("Running Morgan count fingerprint tests...")
+
+    # Test molecules
+    test_molecules = {
+        "aspirin": "CC(=O)OC1=CC=CC=C1C(=O)O",
+        "caffeine": "CN1C=NC2=C1C(=O)N(C(=O)N2C)C",
+        "glucose": "C([C@@H]1[C@H]([C@@H]([C@H](C(O1)O)O)O)O)O",  # With stereochemistry
+        "sodium_acetate": "CC(=O)[O-].[Na+]",  # Salt (largest fragment used)
+        "benzene": "c1ccccc1",
+        "butene_e": "C/C=C/C",  # E-butene
+        "butene_z": "C/C=C\\C",  # Z-butene
+    }
+
+    # Test 1: Morgan Count Fingerprints (default parameters)
+    print("\n1. Testing Morgan fingerprint generation (radius=2, n_bits=2048)...")
+
+    test_df = pd.DataFrame({"SMILES": list(test_molecules.values()), "name": list(test_molecules.keys())})
+    fp_df = compute_morgan_fingerprints(test_df.copy())
+
+    print("   Fingerprint generation results:")
+    for _, row in fp_df.iterrows():
+        fp = row.get("fingerprint", "N/A")
+        if pd.notna(fp):
+            counts = [int(x) for x in fp.split(",")]
+            non_zero = sum(1 for c in counts if c > 0)
+            max_count = max(counts)
+            print(f"   {row['name']:15} → {len(counts)} features, {non_zero} non-zero, max={max_count}")
+        else:
+            print(f"   {row['name']:15} → N/A")
+
+    # Test 2: Different parameters
+    print("\n2. Testing with different parameters (radius=3, n_bits=1024)...")
+
+    fp_df_custom = compute_morgan_fingerprints(test_df.copy(), radius=3, n_bits=1024)
+
+    for _, row in fp_df_custom.iterrows():
+        fp = row.get("fingerprint", "N/A")
+        if pd.notna(fp):
+            counts = [int(x) for x in fp.split(",")]
+            non_zero = sum(1 for c in counts if c > 0)
+            print(f"   {row['name']:15} → {len(counts)} features, {non_zero} non-zero")
+        else:
+            print(f"   {row['name']:15} → N/A")
+
+    # Test 3: Edge cases
+    print("\n3. Testing edge cases...")
+
+    # Invalid SMILES
+    invalid_df = pd.DataFrame({"SMILES": ["INVALID", ""]})
+    fp_invalid = compute_morgan_fingerprints(invalid_df.copy())
+    print(f"   ✓ Invalid SMILES handled: {len(fp_invalid)} rows returned")
+
+    # Test with pre-existing molecule column
+    mol_df = test_df.copy()
+    mol_df["molecule"] = mol_df["SMILES"].apply(Chem.MolFromSmiles)
+    fp_with_mol = compute_morgan_fingerprints(mol_df)
+    print(f"   ✓ Pre-existing molecule column handled: {len(fp_with_mol)} fingerprints generated")
+
+    # Test 4: Verify count values are reasonable
+    print("\n4. Verifying count distribution...")
+    all_counts = []
+    for _, row in fp_df.iterrows():
+        fp = row.get("fingerprint", "N/A")
+        if pd.notna(fp):
+            counts = [int(x) for x in fp.split(",")]
+            all_counts.extend([c for c in counts if c > 0])
+
+    if all_counts:
+        print(f"   Non-zero counts: min={min(all_counts)}, max={max(all_counts)}, mean={np.mean(all_counts):.2f}")
+
+    print("\n✅ All fingerprint tests completed!")

workbench/model_scripts/custom_models/chem_info/mol_descriptors.py

@@ -99,7 +99,6 @@ from rdkit.ML.Descriptors import MoleculeDescriptors
 from mordred import Calculator as MordredCalculator
 from mordred import AcidBase, Aromatic, Constitutional, Chi, CarbonTypes
 
-
 logger = logging.getLogger("workbench")
 logger.setLevel(logging.DEBUG)
 

workbench/model_scripts/custom_models/chem_info/molecular_descriptors.py

@@ -15,7 +15,6 @@ import json
 from mol_standardize import standardize
 from mol_descriptors import compute_descriptors
 
-
 # TRAINING SECTION
 #
 # This section (__main__) is where SageMaker will execute the training job

workbench/model_scripts/custom_models/chem_info/morgan_fingerprints.py

@@ -15,8 +15,7 @@ import pandas as pd
 import json
 
 # Local imports
-from local_utils import compute_morgan_fingerprints
-
+from fingerprints import compute_morgan_fingerprints
 
 # TRAINING SECTION
 #