workbench 0.8.168__py3-none-any.whl → 0.8.193__py3-none-any.whl

workbench/scripts/ml_pipeline_sqs.py

@@ -0,0 +1,186 @@
+import argparse
+import logging
+import json
+from pathlib import Path
+
+# Workbench Imports
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+from workbench.utils.config_manager import ConfigManager
+from workbench.utils.s3_utils import upload_content_to_s3
+
+log = logging.getLogger("workbench")
+cm = ConfigManager()
+workbench_bucket = cm.get_config("WORKBENCH_BUCKET")
+
+
+def submit_to_sqs(
+    script_path: str,
+    size: str = "small",
+    realtime: bool = False,
+    dt: bool = False,
+    promote: bool = False,
+) -> None:
+    """
+    Upload script to S3 and submit message to SQS queue for processing.
+
+    Args:
+        script_path: Local path to the ML pipeline script
+        size: Job size tier - "small" (default), "medium", or "large"
+        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)
+
+    Raises:
+        ValueError: If size is invalid or script file not found
+    """
+    print(f"\n{'=' * 60}")
+    print("🚀  SUBMITTING ML PIPELINE JOB")
+    print(f"{'=' * 60}")
+    if size not in ["small", "medium", "large"]:
+        raise ValueError(f"Invalid size '{size}'. Must be 'small', 'medium', or 'large'")
+
+    # Validate script exists
+    script_file = Path(script_path)
+    if not script_file.exists():
+        raise FileNotFoundError(f"Script not found: {script_path}")
+
+    print(f"📄  Script: {script_file.name}")
+    print(f"📏  Size tier: {size}")
+    print(f"⚡  Mode: {'Real-time' if realtime else 'Serverless'} (serverless={'False' if realtime else 'True'})")
+    print(f"🔄  DynamicTraining: {dt}")
+    print(f"🆕  Promote: {promote}")
+    print(f"🪣  Bucket: {workbench_bucket}")
+    sqs = AWSAccountClamp().boto3_session.client("sqs")
+    script_name = script_file.name
+
+    # List Workbench queues
+    print("\n📋  Listing Workbench SQS queues...")
+    try:
+        queues = sqs.list_queues(QueueNamePrefix="workbench-")
+        queue_urls = queues.get("QueueUrls", [])
+        if queue_urls:
+            print(f"✅  Found {len(queue_urls)} workbench queue(s):")
+            for url in queue_urls:
+                queue_name = url.split("/")[-1]
+                print(f"   • {queue_name}")
+        else:
+            print("⚠️  No workbench queues found")
+    except Exception as e:
+        print(f"❌  Error listing queues: {e}")
+
+    # Upload script to S3
+    s3_path = f"s3://{workbench_bucket}/batch-jobs/{script_name}"
+    print("\n📤  Uploading script to S3...")
+    print(f"   Source: {script_path}")
+    print(f"   Destination: {s3_path}")
+
+    try:
+        upload_content_to_s3(script_file.read_text(), s3_path)
+        print("✅  Script uploaded successfully")
+    except Exception as e:
+        print(f"❌  Upload failed: {e}")
+        raise
+    # Get queue URL and info
+    queue_name = "workbench-ml-pipeline-queue.fifo"
+    print("\n🎯  Getting queue information...")
+    print(f"   Queue name: {queue_name}")
+
+    try:
+        queue_url = sqs.get_queue_url(QueueName=queue_name)["QueueUrl"]
+        print(f"   Queue URL: {queue_url}")
+
+        # Get queue attributes for additional info
+        attrs = sqs.get_queue_attributes(
+            QueueUrl=queue_url, AttributeNames=["ApproximateNumberOfMessages", "ApproximateNumberOfMessagesNotVisible"]
+        )
+        messages_available = attrs["Attributes"].get("ApproximateNumberOfMessages", "0")
+        messages_in_flight = attrs["Attributes"].get("ApproximateNumberOfMessagesNotVisible", "0")
+        print(f"   Messages in queue: {messages_available}")
+        print(f"   Messages in flight: {messages_in_flight}")
+
+    except Exception as e:
+        print(f"❌  Error accessing queue: {e}")
+        raise
+
+    # Prepare message
+    message = {"script_path": s3_path, "size": size}
+
+    # Set environment variables
+    message["environment"] = {
+        "SERVERLESS": "False" if realtime else "True",
+        "DT": str(dt),
+        "PROMOTE": str(promote),
+    }
+
+    # Send the message to SQS
+    try:
+        print("\n📨  Sending message to SQS...")
+        response = sqs.send_message(
+            QueueUrl=queue_url,
+            MessageBody=json.dumps(message, indent=2),
+            MessageGroupId="ml-pipeline-jobs",  # Required for FIFO
+        )
+        message_id = response["MessageId"]
+        print("✅  Message sent successfully!")
+        print(f"   Message ID: {message_id}")
+    except Exception as e:
+        print(f"❌  Failed to send message: {e}")
+        raise
+
+    # Success summary
+    print(f"\n{'=' * 60}")
+    print("✅  JOB SUBMISSION COMPLETE")
+    print(f"{'=' * 60}")
+    print(f"📄  Script: {script_name}")
+    print(f"📏  Size: {size}")
+    print(f"⚡  Mode: {'Real-time' if realtime else 'Serverless'} (SERVERLESS={'False' if realtime else 'True'})")
+    print(f"🔄  DynamicTraining: {dt}")
+    print(f"🆕  Promote: {promote}")
+    print(f"🆔  Message ID: {message_id}")
+    print("\n🔍  MONITORING LOCATIONS:")
+    print(f"   • SQS Queue: AWS Console → SQS → {queue_name}")
+    print("   • Lambda Logs: AWS Console → Lambda → Functions")
+    print("   • Batch Jobs: AWS Console → Batch → Jobs")
+    print("   • CloudWatch: AWS Console → CloudWatch → Log groups")
+    print("\n⏳  Your job should start processing soon...")
+
+
+def main():
+    """CLI entry point for submitting ML pipelines via SQS."""
+    parser = argparse.ArgumentParser(description="Submit ML pipeline to SQS queue for Batch processing")
+    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",
+    )
+    args = parser.parse_args()
+    try:
+        submit_to_sqs(
+            args.script_file,
+            args.size,
+            realtime=args.realtime,
+            dt=args.dt,
+            promote=args.promote,
+        )
+    except Exception as e:
+        print(f"\n❌  ERROR: {e}")
+        log.error(f"Error: {e}")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

workbench/utils/chem_utils/fingerprints.py

@@ -0,0 +1,134 @@
+"""Molecular fingerprint computation utilities"""
+
+import logging
+import pandas as pd
+
+# Molecular Descriptor Imports
+from rdkit import Chem
+from rdkit.Chem import rdFingerprintGenerator
+from rdkit.Chem.MolStandardize import rdMolStandardize
+
+# 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.
+
+    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.
+
+    Returns:
+        pd.DataFrame: The input DataFrame with the Morgan fingerprints added as bit strings.
+
+    Note:
+        See: https://greglandrum.github.io/rdkit-blog/posts/2021-07-06-simulating-counts.html
+    """
+    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 (vectorized)
+    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.error(f"Failed to convert the following SMILES to molecules: {failed_smiles}")
+        df = df.dropna(subset=["molecule"])
+
+    # 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)
+
+    # Compute Morgan fingerprints (vectorized)
+    fingerprints = largest_frags.apply(
+        lambda mol: (morgan_generator.GetFingerprint(mol).ToBitString() if mol else pd.NA)
+    )
+
+    # 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 molecular fingerprint tests...")
+    print("Note: This requires molecular_screening module to be available")
+
+    # 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
+        "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_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)
+
+    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")
+
+    # Test 2: Different fingerprint parameters
+    print("\n2. Testing different fingerprint parameters...")
+
+    # Test with counts enabled
+    fp_counts_df = compute_morgan_fingerprints(test_df.copy(), radius=3, n_bits=256, counts=True)
+
+    print("   With count simulation (256 bits * 4):")
+    for _, row in fp_counts_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")
+
+    # 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__}")
+
+    # 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")
+
+    print("\n✅ All fingerprint tests completed!")

workbench/utils/chem_utils/misc.py

@@ -0,0 +1,194 @@
+"""Miscellaneous processing functions for molecular data."""
+
+import logging
+import numpy as np
+import pandas as pd
+from typing import List, Optional
+
+# Set up the logger
+log = logging.getLogger("workbench")
+
+
+def geometric_mean(series: pd.Series) -> float:
+    """Computes the geometric mean manually to avoid using scipy."""
+    return np.exp(np.log(series).mean())
+
+
+def rollup_experimental_data(
+    df: pd.DataFrame, id: str, time: str, target: str, use_gmean: bool = False
+) -> pd.DataFrame:
+    """
+    Rolls up a dataset by selecting the largest time per unique ID and averaging the target value
+    if multiple records exist at that time. Supports both arithmetic and geometric mean.
+
+    Parameters:
+        df (pd.DataFrame): Input dataframe.
+        id (str): Column representing the unique molecule ID.
+        time (str): Column representing the time.
+        target (str): Column representing the target value.
+        use_gmean (bool): Whether to use the geometric mean instead of the arithmetic mean.
+
+    Returns:
+        pd.DataFrame: Rolled-up dataframe with all original columns retained.
+    """
+    # Find the max time per unique ID
+    max_time_df = df.groupby(id)[time].transform("max")
+    filtered_df = df[df[time] == max_time_df]
+
+    # Define aggregation function
+    agg_func = geometric_mean if use_gmean else np.mean
+
+    # Perform aggregation on all columns
+    agg_dict = {col: "first" for col in df.columns if col not in [target, id, time]}
+    agg_dict[target] = lambda x: agg_func(x) if len(x) > 1 else x.iloc[0]  # Apply mean or gmean
+
+    rolled_up_df = filtered_df.groupby([id, time]).agg(agg_dict).reset_index()
+    return rolled_up_df
+
+
+def micromolar_to_log(series_µM: pd.Series) -> pd.Series:
+    """
+    Convert a pandas Series of concentrations in µM (micromolar) to their logarithmic values (log10).
+
+    Parameters:
+    series_uM (pd.Series): Series of concentrations in micromolar.
+
+    Returns:
+    pd.Series: Series of logarithmic values (log10).
+    """
+    # Replace 0 or negative values with a small number to avoid log errors
+    adjusted_series = series_µM.clip(lower=1e-9)  # Alignment with another project
+
+    series_mol_per_l = adjusted_series * 1e-6  # Convert µM/L to mol/L
+    log_series = np.log10(series_mol_per_l)
+    return log_series
+
+
+def log_to_micromolar(log_series: pd.Series) -> pd.Series:
+    """
+    Convert a pandas Series of logarithmic values (log10) back to concentrations in µM (micromolar).
+
+    Parameters:
+    log_series (pd.Series): Series of logarithmic values (log10).
+
+    Returns:
+    pd.Series: Series of concentrations in micromolar.
+    """
+    series_mol_per_l = 10**log_series  # Convert log10 back to mol/L
+    series_µM = series_mol_per_l * 1e6  # Convert mol/L to µM
+    return series_µM
+
+
+def feature_resolution_issues(df: pd.DataFrame, features: List[str], show_cols: Optional[List[str]] = None) -> None:
+    """
+    Identify and print groups in a DataFrame where the given features have more than one unique SMILES,
+    sorted by group size (largest number of unique SMILES first).
+
+    Args:
+        df (pd.DataFrame): Input DataFrame containing SMILES strings.
+        features (List[str]): List of features to check.
+        show_cols (Optional[List[str]]): Columns to display; defaults to all columns.
+    """
+    # 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")
+
+    show_cols = show_cols if show_cols is not None else df.columns.tolist()
+
+    # Drop duplicates to keep only unique SMILES for each feature combination
+    unique_df = df.drop_duplicates(subset=[smiles_column] + features)
+
+    # Find groups with more than one unique SMILES
+    group_counts = unique_df.groupby(features).size()
+    collision_groups = group_counts[group_counts > 1].sort_values(ascending=False)
+
+    # Print each group in order of size (largest first)
+    for group, count in collision_groups.items():
+        # Get the rows for this group
+        if isinstance(group, tuple):
+            group_mask = (unique_df[features] == group).all(axis=1)
+        else:
+            group_mask = unique_df[features[0]] == group
+
+        group_df = unique_df[group_mask]
+
+        print(f"Feature Group (unique SMILES: {count}):")
+        print(group_df[show_cols])
+        print("\n")
+
+
+if __name__ == "__main__":
+    print("Running molecular processing and transformation tests...")
+    print("Note: This requires the molecular_filters module to be available")
+
+    # Test 1: Concentration conversions
+    print("\n1. Testing concentration conversions...")
+
+    # Test micromolar to log
+    test_conc = pd.Series([1.0, 10.0, 100.0, 1000.0, 0.001])
+    log_values = micromolar_to_log(test_conc)
+    back_to_uM = log_to_micromolar(log_values)
+
+    print("   µM → log10 → µM:")
+    for orig, log_val, back in zip(test_conc, log_values, back_to_uM):
+        print(f"   {orig:8.3f} µM → {log_val:6.2f} → {back:8.3f} µM")
+
+    # Test 2: Geometric mean
+    print("\n2. Testing geometric mean...")
+    test_series = pd.Series([2, 4, 8, 16])
+    geo_mean = geometric_mean(test_series)
+    arith_mean = np.mean(test_series)
+    print(f"   Series: {list(test_series)}")
+    print(f"   Arithmetic mean: {arith_mean:.2f}")
+    print(f"   Geometric mean: {geo_mean:.2f}")
+
+    # Test 3: Experimental data rollup
+    print("\n3. Testing experimental data rollup...")
+
+    # Create test data with multiple timepoints and replicates
+    test_data = pd.DataFrame(
+        {
+            "compound_id": ["A", "A", "A", "B", "B", "C", "C", "C"],
+            "time": [1, 2, 2, 1, 2, 1, 1, 2],
+            "activity": [10, 20, 22, 5, 8, 100, 110, 200],
+            "assay": ["kinase", "kinase", "kinase", "kinase", "kinase", "cell", "cell", "cell"],
+        }
+    )
+
+    # Rollup with arithmetic mean
+    rolled_arith = rollup_experimental_data(test_data, "compound_id", "time", "activity", use_gmean=False)
+    print("   Arithmetic mean rollup:")
+    print(rolled_arith[["compound_id", "time", "activity"]])
+
+    # Rollup with geometric mean
+    rolled_geo = rollup_experimental_data(test_data, "compound_id", "time", "activity", use_gmean=True)
+    print("\n   Geometric mean rollup:")
+    print(rolled_geo[["compound_id", "time", "activity"]])
+
+    # Test 4: Feature resolution issues
+    print("\n4. Testing feature resolution identification...")
+
+    # Create data with some duplicate features but different SMILES
+    resolution_df = pd.DataFrame(
+        {
+            "smiles": ["CCO", "C(C)O", "CC(C)O", "CCC(C)O", "CCCO"],
+            "assay_id": ["A1", "A1", "A2", "A2", "A3"],
+            "value": [1.0, 1.5, 2.0, 2.2, 3.0],
+        }
+    )
+
+    print("   Checking for feature collisions in 'assay_id':")
+    feature_resolution_issues(resolution_df, ["assay_id"], show_cols=["smiles", "assay_id", "value"])
+
+    # Test 7: Edge cases
+    print("\n7. Testing edge cases...")
+
+    # Zero and negative concentrations
+    edge_conc = pd.Series([0, -1, 1e-10])
+    edge_log = micromolar_to_log(edge_conc)
+    print("   Edge concentration handling:")
+    for c, l in zip(edge_conc, edge_log):
+        print(f"      {c:6.2e} µM → {l:6.2f}")
+
+    print("\n✅ All molecular processing tests completed!")