workbench 0.8.213__py3-none-any.whl → 0.8.219__py3-none-any.whl

workbench/model_scripts/xgb_model/xgb_model.template

@@ -7,39 +7,30 @@
 # - Sample weights support
 # - Categorical feature handling
 # - Compressed feature decompression
+#
+# NOTE: Imports are structured to minimize serverless endpoint startup time.
+# Heavy imports (sklearn, awswrangler) are deferred to training time.
 
-import argparse
 import json
 import os
 
-import awswrangler as wr
 import joblib
 import numpy as np
 import pandas as pd
 import xgboost as xgb
-from sklearn.model_selection import KFold, StratifiedKFold, train_test_split
-from sklearn.preprocessing import LabelEncoder
 
 from model_script_utils import (
-    check_dataframe,
-    compute_classification_metrics,
-    compute_regression_metrics,
     convert_categorical_types,
     decompress_features,
     expand_proba_column,
     input_fn,
     match_features_case_insensitive,
     output_fn,
-    print_classification_metrics,
-    print_confusion_matrix,
-    print_regression_metrics,
 )
 from uq_harness import (
     compute_confidence,
     load_uq_models,
     predict_intervals,
-    save_uq_models,
-    train_uq_models,
 )
 
 # =============================================================================
@@ -49,25 +40,27 @@ DEFAULT_HYPERPARAMETERS = {
     # Training parameters
     "n_folds": 5,  # Number of CV folds (1 = single train/val split)
     # Core tree parameters
-    "n_estimators": 200,
-    "max_depth": 6,
+    "n_estimators": 300,
+    "max_depth": 7,
     "learning_rate": 0.05,
-    # Sampling parameters
-    "subsample": 0.7,
-    "colsample_bytree": 0.6,
-    "colsample_bylevel": 0.8,
-    # Regularization
-    "min_child_weight": 5,
-    "gamma": 0.2,
-    "reg_alpha": 0.5,
-    "reg_lambda": 2.0,
+    # Sampling parameters (less aggressive - ensemble provides regularization)
+    "subsample": 0.8,
+    "colsample_bytree": 0.8,
+    # Regularization (lighter - ensemble averaging reduces overfitting)
+    "min_child_weight": 3,
+    "gamma": 0.1,
+    "reg_alpha": 0.1,
+    "reg_lambda": 1.0,
     # Random seed
-    "random_state": 42,
+    "seed": 42,
 }
 
 # Workbench-specific parameters (not passed to XGBoost)
 WORKBENCH_PARAMS = {"n_folds"}
 
+# Regression-only parameters (filtered out for classifiers)
+REGRESSION_ONLY_PARAMS = {"objective"}
+
 # Template parameters (filled in by Workbench)
 TEMPLATE_PARAMS = {
     "model_type": "{{model_type}}",
@@ -80,10 +73,140 @@ TEMPLATE_PARAMS = {
 }
 
 
+# =============================================================================
+# Model Loading (for SageMaker inference)
+# =============================================================================
+def model_fn(model_dir: str) -> dict:
+    """Load XGBoost ensemble from the specified directory."""
+    # Load ensemble metadata
+    metadata_path = os.path.join(model_dir, "ensemble_metadata.json")
+    if os.path.exists(metadata_path):
+        with open(metadata_path) as f:
+            metadata = json.load(f)
+        n_ensemble = metadata["n_ensemble"]
+    else:
+        n_ensemble = 1  # Legacy single model
+
+    # Load ensemble models
+    ensemble_models = []
+    for i in range(n_ensemble):
+        model_path = os.path.join(model_dir, f"xgb_model_{i}.joblib")
+        if not os.path.exists(model_path):
+            model_path = os.path.join(model_dir, "xgb_model.joblib")  # Legacy fallback
+        ensemble_models.append(joblib.load(model_path))
+
+    print(f"Loaded {len(ensemble_models)} model(s)")
+
+    # Load label encoder (classifier only)
+    label_encoder = None
+    encoder_path = os.path.join(model_dir, "label_encoder.joblib")
+    if os.path.exists(encoder_path):
+        label_encoder = joblib.load(encoder_path)
+
+    # Load category mappings
+    category_mappings = {}
+    category_path = os.path.join(model_dir, "category_mappings.json")
+    if os.path.exists(category_path):
+        with open(category_path) as f:
+            category_mappings = json.load(f)
+
+    # Load UQ models (regression only)
+    uq_models, uq_metadata = None, None
+    uq_path = os.path.join(model_dir, "uq_metadata.json")
+    if os.path.exists(uq_path):
+        uq_models, uq_metadata = load_uq_models(model_dir)
+
+    return {
+        "ensemble_models": ensemble_models,
+        "n_ensemble": n_ensemble,
+        "label_encoder": label_encoder,
+        "category_mappings": category_mappings,
+        "uq_models": uq_models,
+        "uq_metadata": uq_metadata,
+    }
+
+
+# =============================================================================
+# Inference (for SageMaker inference)
+# =============================================================================
+def predict_fn(df: pd.DataFrame, model_dict: dict) -> pd.DataFrame:
+    """Make predictions with XGBoost ensemble."""
+    model_dir = os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+    with open(os.path.join(model_dir, "feature_columns.json")) as f:
+        features = json.load(f)
+    print(f"Model Features: {features}")
+
+    # Extract model components
+    ensemble_models = model_dict["ensemble_models"]
+    label_encoder = model_dict.get("label_encoder")
+    category_mappings = model_dict.get("category_mappings", {})
+    uq_models = model_dict.get("uq_models")
+    uq_metadata = model_dict.get("uq_metadata")
+    compressed_features = TEMPLATE_PARAMS["compressed_features"]
+
+    # Prepare features
+    matched_df = match_features_case_insensitive(df, features)
+    matched_df, _ = convert_categorical_types(matched_df, features, category_mappings)
+
+    if compressed_features:
+        print("Decompressing features for prediction...")
+        matched_df, features = decompress_features(matched_df, features, compressed_features)
+
+    X = matched_df[features]
+
+    # Collect ensemble predictions
+    all_preds = [m.predict(X) for m in ensemble_models]
+    ensemble_preds = np.stack(all_preds, axis=0)
+
+    if label_encoder is not None:
+        # Classification: average probabilities, then argmax
+        all_probs = [m.predict_proba(X) for m in ensemble_models]
+        avg_probs = np.mean(np.stack(all_probs, axis=0), axis=0)
+        class_preds = np.argmax(avg_probs, axis=1)
+
+        df["prediction"] = label_encoder.inverse_transform(class_preds)
+        df["pred_proba"] = [p.tolist() for p in avg_probs]
+        df = expand_proba_column(df, label_encoder.classes_)
+    else:
+        # Regression: average predictions
+        df["prediction"] = np.mean(ensemble_preds, axis=0)
+        df["prediction_std"] = np.std(ensemble_preds, axis=0)
+
+        # Add UQ intervals if available
+        if uq_models and uq_metadata:
+            df = predict_intervals(df, X, uq_models, uq_metadata)
+            df = compute_confidence(df, uq_metadata["median_interval_width"], "q_10", "q_90")
+
+    print(f"Inference complete: {len(df)} predictions, {len(ensemble_models)} ensemble members")
+    return df
+
+
 # =============================================================================
 # Training
 # =============================================================================
 if __name__ == "__main__":
+    # -------------------------------------------------------------------------
+    # Training-only imports (deferred to reduce serverless startup time)
+    # -------------------------------------------------------------------------
+    import argparse
+
+    import awswrangler as wr
+    from sklearn.model_selection import KFold, StratifiedKFold, train_test_split
+    from sklearn.preprocessing import LabelEncoder
+
+    from model_script_utils import (
+        check_dataframe,
+        compute_classification_metrics,
+        compute_regression_metrics,
+        print_classification_metrics,
+        print_confusion_matrix,
+        print_regression_metrics,
+    )
+    from uq_harness import (
+        save_uq_models,
+        train_uq_models,
+    )
+
     # -------------------------------------------------------------------------
     # Setup: Parse arguments and load data
     # -------------------------------------------------------------------------
@@ -123,7 +246,7 @@ if __name__ == "__main__":
         all_df, features = decompress_features(all_df, features, compressed_features)
 
     # -------------------------------------------------------------------------
-    # Classification setup: Encode target labels
+    # Classification setup
     # -------------------------------------------------------------------------
     label_encoder = None
     if model_type == "classifier":
@@ -136,6 +259,18 @@ if __name__ == "__main__":
     # -------------------------------------------------------------------------
     n_folds = hyperparameters["n_folds"]
     xgb_params = {k: v for k, v in hyperparameters.items() if k not in WORKBENCH_PARAMS}
+
+    # Map 'seed' to 'random_state' for XGBoost
+    if "seed" in xgb_params:
+        xgb_params["random_state"] = xgb_params.pop("seed")
+
+    # Handle objective: filter regression-only params for classifiers, set default for regressors
+    if model_type == "classifier":
+        xgb_params = {k: v for k, v in xgb_params.items() if k not in REGRESSION_ONLY_PARAMS}
+    else:
+        # Default to MAE (reg:absoluteerror) for regression if not specified
+        xgb_params.setdefault("objective", "reg:absoluteerror")
+
     print(f"XGBoost params: {xgb_params}")
 
     if n_folds == 1:
@@ -285,12 +420,10 @@ if __name__ == "__main__":
     # -------------------------------------------------------------------------
     # Save model artifacts
     # -------------------------------------------------------------------------
-    # Ensemble models
-    for idx, ens_model in enumerate(ensemble_models):
-        joblib.dump(ens_model, os.path.join(args.model_dir, f"xgb_model_{idx}.joblib"))
-    print(f"Saved {len(ensemble_models)} XGBoost model(s)")
+    for idx, m in enumerate(ensemble_models):
+        joblib.dump(m, os.path.join(args.model_dir, f"xgb_model_{idx}.joblib"))
+    print(f"Saved {len(ensemble_models)} model(s)")
 
-    # Metadata files
     with open(os.path.join(args.model_dir, "ensemble_metadata.json"), "w") as f:
         json.dump({"n_ensemble": len(ensemble_models), "n_folds": n_folds}, f)
 
@@ -310,125 +443,3 @@ if __name__ == "__main__":
         save_uq_models(uq_models, uq_metadata, args.model_dir)
 
     print(f"\nModel training complete! Artifacts saved to {args.model_dir}")
-
-
-# =============================================================================
-# Model Loading (for SageMaker inference)
-# =============================================================================
-def model_fn(model_dir: str) -> dict:
-    """Load XGBoost ensemble and associated artifacts.
-
-    Args:
-        model_dir: Directory containing model artifacts
-
-    Returns:
-        Dictionary with ensemble_models, label_encoder, category_mappings, uq_models, etc.
-    """
-    # Load ensemble metadata
-    metadata_path = os.path.join(model_dir, "ensemble_metadata.json")
-    if os.path.exists(metadata_path):
-        with open(metadata_path) as f:
-            metadata = json.load(f)
-        n_ensemble = metadata["n_ensemble"]
-    else:
-        n_ensemble = 1  # Legacy single model
-
-    # Load ensemble models
-    ensemble_models = []
-    for i in range(n_ensemble):
-        model_path = os.path.join(model_dir, f"xgb_model_{i}.joblib")
-        if not os.path.exists(model_path):
-            model_path = os.path.join(model_dir, "xgb_model.joblib")  # Legacy fallback
-        ensemble_models.append(joblib.load(model_path))
-
-    # Load label encoder (classifier only)
-    label_encoder = None
-    encoder_path = os.path.join(model_dir, "label_encoder.joblib")
-    if os.path.exists(encoder_path):
-        label_encoder = joblib.load(encoder_path)
-
-    # Load category mappings
-    category_mappings = {}
-    category_path = os.path.join(model_dir, "category_mappings.json")
-    if os.path.exists(category_path):
-        with open(category_path) as f:
-            category_mappings = json.load(f)
-
-    # Load UQ models (regression only)
-    uq_models, uq_metadata = None, None
-    uq_path = os.path.join(model_dir, "uq_metadata.json")
-    if os.path.exists(uq_path):
-        uq_models, uq_metadata = load_uq_models(model_dir)
-
-    return {
-        "ensemble_models": ensemble_models,
-        "n_ensemble": n_ensemble,
-        "label_encoder": label_encoder,
-        "category_mappings": category_mappings,
-        "uq_models": uq_models,
-        "uq_metadata": uq_metadata,
-    }
-
-
-# =============================================================================
-# Inference (for SageMaker inference)
-# =============================================================================
-def predict_fn(df: pd.DataFrame, models: dict) -> pd.DataFrame:
-    """Make predictions with XGBoost ensemble.
-
-    Args:
-        df: Input DataFrame with features
-        models: Dictionary from model_fn containing ensemble and metadata
-
-    Returns:
-        DataFrame with predictions added
-    """
-    # Load feature columns
-    model_dir = os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
-    with open(os.path.join(model_dir, "feature_columns.json")) as f:
-        features = json.load(f)
-    print(f"Model Features: {features}")
-
-    # Extract model components
-    ensemble_models = models["ensemble_models"]
-    label_encoder = models.get("label_encoder")
-    category_mappings = models.get("category_mappings", {})
-    uq_models = models.get("uq_models")
-    uq_metadata = models.get("uq_metadata")
-    compressed_features = TEMPLATE_PARAMS["compressed_features"]
-
-    # Prepare features
-    matched_df = match_features_case_insensitive(df, features)
-    matched_df, _ = convert_categorical_types(matched_df, features, category_mappings)
-
-    if compressed_features:
-        print("Decompressing features for prediction...")
-        matched_df, features = decompress_features(matched_df, features, compressed_features)
-
-    X = matched_df[features]
-
-    # Collect ensemble predictions
-    all_preds = [m.predict(X) for m in ensemble_models]
-    ensemble_preds = np.stack(all_preds, axis=0)
-
-    if label_encoder is not None:
-        # Classification: average probabilities, then argmax
-        all_probs = [m.predict_proba(X) for m in ensemble_models]
-        avg_probs = np.mean(np.stack(all_probs, axis=0), axis=0)
-        class_preds = np.argmax(avg_probs, axis=1)
-
-        df["prediction"] = label_encoder.inverse_transform(class_preds)
-        df["pred_proba"] = [p.tolist() for p in avg_probs]
-        df = expand_proba_column(df, label_encoder.classes_)
-    else:
-        # Regression: average predictions
-        df["prediction"] = np.mean(ensemble_preds, axis=0)
-        df["prediction_std"] = np.std(ensemble_preds, axis=0)
-
-        # Add UQ intervals if available
-        if uq_models and uq_metadata:
-            df = predict_intervals(df, X, uq_models, uq_metadata)
-            df = compute_confidence(df, uq_metadata["median_interval_width"], "q_10", "q_90")
-
-    print(f"Inference complete: {len(df)} predictions, {len(ensemble_models)} ensemble members")
-    return df

workbench/repl/workbench_shell.py

@@ -302,11 +302,6 @@ class WorkbenchShell:
             self.commands["PandasToView"] = importlib.import_module("workbench.core.views.pandas_to_view").PandasToView
             self.commands["Pipeline"] = importlib.import_module("workbench.api.pipeline").Pipeline
 
-            # Algorithms
-            self.commands["FSP"] = importlib.import_module(
-                "workbench.algorithms.dataframe.feature_space_proximity"
-            ).FeatureSpaceProximity
-
             # These are 'nice to have' imports
             self.commands["pd"] = importlib.import_module("pandas")
             self.commands["wr"] = importlib.import_module("awswrangler")

workbench/scripts/endpoint_test.py

@@ -5,7 +5,7 @@ Usage:
     python model_script_harness.py <local_script.py> <model_name>
 
 Example:
-    python model_script_harness.py pytorch.py aqsol-pytorch-reg
+    python model_script_harness.py pytorch.py aqsol-reg-pytorch
 
 This allows you to test LOCAL changes to a model script against deployed model artifacts.
 Evaluation data is automatically pulled from the FeatureSet (training = FALSE rows).
@@ -72,7 +72,7 @@ def main():
         print("Usage: python model_script_harness.py <local_script.py> <model_name>")
         print("\nArguments:")
         print("  local_script.py  - Path to your LOCAL model script to test")
-        print("  model_name       - Workbench model name (e.g., aqsol-pytorch-reg)")
+        print("  model_name       - Workbench model name (e.g., aqsol-reg-pytorch)")
         print("\nOptional: testing/env.json with additional environment variables")
         sys.exit(1)
 

workbench/scripts/meta_model_sim.py

@@ -0,0 +1,35 @@
+"""MetaModelSimulator: Simulate and analyze ensemble model performance.
+
+This class helps evaluate whether a meta model (ensemble) would outperform
+individual child models by analyzing endpoint inference predictions.
+"""
+
+import argparse
+from workbench.utils.meta_model_simulator import MetaModelSimulator
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Simulate and analyze ensemble model performance using MetaModelSimulator."
+    )
+    parser.add_argument(
+        "models",
+        nargs="+",
+        help="List of model endpoint names to include in the ensemble simulation.",
+    )
+    parser.add_argument(
+        "--id-column",
+        default="molecule_name",
+        help="Name of the ID column (default: molecule_name)",
+    )
+    args = parser.parse_args()
+    models = args.models
+    id_column = args.id_column
+
+    # Create MetaModelSimulator instance and generate report
+    sim = MetaModelSimulator(models, id_column=id_column)
+    sim.report()
+
+
+if __name__ == "__main__":
+    main()

workbench/utils/chem_utils/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!")