workbench 0.8.202__py3-none-any.whl → 0.8.220__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

@@ -1,11 +1,3 @@
 # Requirements for ChemProp model scripts
-# Note: These are the local dev requirements. The Docker images have their own requirements.txt
-chemprop==2.2.1
-rdkit==2025.9.1
-torch>=2.0.0
-lightning>=2.0.0
-pandas>=2.0.0
-numpy>=1.24.0
-scikit-learn>=1.3.0
-awswrangler>=3.0.0
-joblib>=1.3.0
+# 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

@@ -1,31 +1,48 @@
-"""Molecular fingerprint computation utilities"""
+"""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 pandas as pd
 
-# Molecular Descriptor Imports
-from rdkit import Chem
-from rdkit.Chem import rdFingerprintGenerator
+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=2, n_bits=2048, counts=True) -> pd.DataFrame:
-    """Compute and add Morgan fingerprints to the DataFrame.
+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 (pd.DataFrame): Input DataFrame containing SMILES strings.
-        radius (int): Radius for the Morgan fingerprint.
-        n_bits (int): Number of bits for the fingerprint.
-        counts (bool): Count simulation for the fingerprint.
+        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: The input DataFrame with the Morgan fingerprints added as bit strings.
+        pd.DataFrame: Input DataFrame with 'fingerprint' column added.
+                      Values are comma-separated uint8 counts.
 
     Note:
-        See: https://greglandrum.github.io/rdkit-blog/posts/2021-07-06-simulating-counts.html
+        Count fingerprints outperform binary for ADMET prediction.
+        See: https://pubs.acs.org/doi/10.1021/acs.est.3c02198
     """
     delete_mol_column = False
 
@@ -39,7 +56,7 @@ def compute_morgan_fingerprints(df: pd.DataFrame, radius=2, n_bits=2048, counts=
         log.warning("Detected serialized molecules in 'molecule' column. Removing...")
         del df["molecule"]
 
-    # Convert SMILES to RDKit molecule objects (vectorized)
+    # Convert SMILES to RDKit molecule objects
     if "molecule" not in df.columns:
         log.info("Converting SMILES to RDKit Molecules...")
         delete_mol_column = True
@@ -47,23 +64,32 @@ def compute_morgan_fingerprints(df: pd.DataFrame, radius=2, n_bits=2048, counts=
         # Make sure our molecules are not None
         failed_smiles = df[df["molecule"].isnull()][smiles_column].tolist()
         if failed_smiles:
-            log.error(f"Failed to convert the following SMILES to molecules: {failed_smiles}")
-        df = df.dropna(subset=["molecule"])
+            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
     )
 
-    # Create a Morgan fingerprint generator
-    if counts:
-        n_bits *= 4  # Multiply by 4 to simulate counts
-    morgan_generator = rdFingerprintGenerator.GetMorganGenerator(radius=radius, fpSize=n_bits, countSimulation=counts)
+    def mol_to_count_string(mol):
+        """Convert molecule to comma-separated count fingerprint string."""
+        if mol is None:
+            return pd.NA
 
-    # Compute Morgan fingerprints (vectorized)
-    fingerprints = largest_frags.apply(
-        lambda mol: (morgan_generator.GetFingerprint(mol).ToBitString() if mol else 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
@@ -71,59 +97,62 @@ def compute_morgan_fingerprints(df: pd.DataFrame, radius=2, n_bits=2048, counts=
     # Drop the intermediate 'molecule' column if it was added
     if delete_mol_column:
         del df["molecule"]
+
     return df
 
 
 if __name__ == "__main__":
-    print("Running molecular fingerprint tests...")
-    print("Note: This requires molecular_screening module to be available")
+    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
+        "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 Fingerprints
-    print("\n1. Testing Morgan fingerprint generation...")
+    # 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(), radius=2, n_bits=512, counts=False)
+    fp_df = compute_morgan_fingerprints(test_df.copy())
 
     print("   Fingerprint generation results:")
     for _, row in fp_df.iterrows():
         fp = row.get("fingerprint", "N/A")
-        fp_len = len(fp) if fp != "N/A" else 0
-        print(f"   {row['name']:15} → {fp_len} bits")
+        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 fingerprint parameters
-    print("\n2. Testing different fingerprint parameters...")
+    # Test 2: Different parameters
+    print("\n2. Testing with different parameters (radius=3, n_bits=1024)...")
 
-    # Test with counts enabled
-    fp_counts_df = compute_morgan_fingerprints(test_df.copy(), radius=3, n_bits=256, counts=True)
+    fp_df_custom = compute_morgan_fingerprints(test_df.copy(), radius=3, n_bits=1024)
 
-    print("   With count simulation (256 bits * 4):")
-    for _, row in fp_counts_df.iterrows():
+    for _, row in fp_df_custom.iterrows():
         fp = row.get("fingerprint", "N/A")
-        fp_len = len(fp) if fp != "N/A" else 0
-        print(f"   {row['name']:15} → {fp_len} bits")
+        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", ""]})
-    try:
-        fp_invalid = compute_morgan_fingerprints(invalid_df.copy())
-        print(f"   ✓ Invalid SMILES handled: {len(fp_invalid)} valid molecules")
-    except Exception as e:
-        print(f"   ✓ Invalid SMILES properly raised error: {type(e).__name__}")
+    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()
@@ -131,4 +160,16 @@ if __name__ == "__main__":
     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!")