workbench 0.8.171__py3-none-any.whl → 0.8.173__py3-none-any.whl

workbench/model_scripts/custom_models/chem_info/mol_standardize.py

@@ -0,0 +1,428 @@
+"""
+mol_standardize.py - Molecular Standardization for ADMET Preprocessing
+Following ChEMBL structure standardization pipeline
+
+Purpose:
+    Standardizes chemical structures to ensure consistent molecular representations
+    for ADMET modeling. Handles tautomers, salts, charges, and structural variations
+    that can cause the same compound to be represented differently.
+
+Standardization Pipeline:
+    1. Cleanup
+       - Removes explicit hydrogens
+       - Disconnects metal atoms from organic fragments
+       - Normalizes functional groups (e.g., nitro, sulfoxide representations)
+
+    2. Fragment Parent Selection (optional, controlled by extract_salts parameter)
+       - Identifies and keeps the largest organic fragment
+       - Removes salts, solvents, and counterions
+       - Example: [Na+].CC(=O)[O-] → CC(=O)O (keeps acetate, removes sodium)
+
+    3. Charge Neutralization (optional, controlled by extract_salts parameter)
+       - Neutralizes charges where possible
+       - Only applied when extract_salts=True (following ChEMBL pipeline)
+       - Skipped when extract_salts=False to preserve ionic character
+       - Example: CC(=O)[O-] → CC(=O)O
+
+    4. Tautomer Canonicalization (optional, default=True)
+       - Generates canonical tautomer form for consistency
+       - Example: Oc1ccccn1 → O=c1cccc[nH]1 (2-hydroxypyridine → 2-pyridone)
+
+Output DataFrame Columns:
+    - orig_smiles: Original input SMILES (preserved for traceability)
+    - smiles: Standardized molecule (with or without salts based on extract_salts)
+    - salt: Removed salt/counterion as SMILES (only populated if extract_salts=True)
+
+Salt Handling:
+    Salt forms can dramatically affect properties like solubility.
+    This module offers two modes for handling salts:
+
+    When extract_salts=True (default, ChEMBL standard):
+        - Removes salts/counterions to get parent molecule
+        - Neutralizes charges on the parent
+        - Records removed salts in 'salt' column
+        Input: [Na+].CC(=O)[O-]  →  Parent: CC(=O)O, Salt: [Na+]
+
+    When extract_salts=False (preserve full salt form):
+        - Keeps all fragments including salts/counterions
+        - Preserves ionic charges (no neutralization)
+
+Mixture Detection:
+    The module detects and logs potential mixtures (vs true salt forms):
+    - Multiple large neutral organic fragments indicate a mixture
+    - Mixtures are logged but NOT recorded in the salt column
+    - True salts (small/charged fragments) are properly extracted
+
+    Downstream modeling options:
+    1. Use parent only (standard approach for most ADMET properties)
+    2. Include salt as a categorical or computed feature
+    3. Model parent + salt effects hierarchically
+    4. Use full salt form for properties like solubility/formulation
+
+References:
+    - "ChEMBL Structure Pipeline" (Bento et al., 2020)
+      https://doi.org/10.1186/s13321-020-00456-1
+    - "Standardization and Validation with the RDKit" (Greg Landrum, RSC Open Science 2021)
+      https://github.com/greglandrum/RSC_OpenScience_Standardization_202104/blob/main/Standardization%20and%20Validation%20with%20the%20RDKit.ipynb
+
+Usage:
+    from mol_standardize import standardize
+
+    # Basic usage (removes salts by default, ChEMBL standard)
+    df_std = standardize(df, smiles_column='smiles')
+
+    # Keep salts in the molecule (preserve ionic forms)
+    df_std = standardize(df, extract_salts=False)
+
+    # Without tautomer canonicalization (faster, less aggressive)
+    df_std = standardize(df, canonicalize_tautomer=False)
+"""
+
+import logging
+from typing import Optional, Tuple
+import pandas as pd
+from rdkit import Chem
+from rdkit.Chem import Mol
+from rdkit.Chem.MolStandardize import rdMolStandardize
+from rdkit import RDLogger
+
+log = logging.getLogger("workbench")
+RDLogger.DisableLog("rdApp.warning")
+
+
+class MolStandardizer:
+    """
+    Streamlined molecular standardizer for ADMET preprocessing
+    Uses ChEMBL standardization pipeline with RDKit
+    """
+
+    def __init__(self, canonicalize_tautomer: bool = True, remove_salts: bool = True):
+        """
+        Initialize standardizer with ChEMBL defaults
+
+        Args:
+            canonicalize_tautomer: Whether to canonicalize tautomers (default True)
+            remove_salts: Whether to remove salts/counterions (default True)
+        """
+        self.canonicalize_tautomer = canonicalize_tautomer
+        self.remove_salts = remove_salts
+        self.params = rdMolStandardize.CleanupParameters()
+        self.tautomer_enumerator = rdMolStandardize.TautomerEnumerator(self.params)
+
+    def standardize(self, mol: Mol) -> Tuple[Optional[Mol], Optional[str]]:
+        """
+        Main standardization pipeline for ADMET
+
+        Pipeline:
+        1. Cleanup (remove Hs, disconnect metals, normalize)
+        2. Get largest fragment (optional - only if remove_salts=True)
+        3. Neutralize charges
+        4. Canonicalize tautomer (optional)
+
+        Args:
+            mol: RDKit molecule object
+
+        Returns:
+            Tuple of (standardized molecule or None if failed, salt SMILES or None)
+        """
+        if mol is None:
+            return None, None
+
+        try:
+            # Step 1: Cleanup
+            mol = rdMolStandardize.Cleanup(mol, self.params)
+            if mol is None:
+                return None, None
+
+            salt_smiles = None
+
+            # Step 2: Fragment handling (conditional based on remove_salts)
+            if self.remove_salts:
+                # Get parent molecule and extract salt information
+                parent_mol = rdMolStandardize.FragmentParent(mol, self.params)
+                if parent_mol:
+                    salt_smiles = self._extract_salt(mol, parent_mol)
+                    mol = parent_mol
+                else:
+                    return None, None
+            # If not removing salts, keep the full molecule intact
+
+            # Step 3: Neutralize charges (skip if keeping salts to preserve ionic forms)
+            if self.remove_salts:
+                mol = rdMolStandardize.ChargeParent(mol, self.params, skipStandardize=True)
+                if mol is None:
+                    return None, salt_smiles
+
+            # Step 4: Canonicalize tautomer
+            if self.canonicalize_tautomer:
+                mol = self.tautomer_enumerator.Canonicalize(mol)
+
+            return mol, salt_smiles
+
+        except Exception as e:
+            log.warning(f"Standardization failed: {e}")
+            return None, None
+
+    def _extract_salt(self, orig_mol: Mol, parent_mol: Mol) -> Optional[str]:
+        """
+        Extract salt/counterion by comparing original and parent molecules.
+
+        Detects and handles mixtures vs true salt forms:
+        - True salts: small (<= 6 heavy atoms) or charged fragments
+        - Mixtures: multiple large neutral organic fragments
+
+        Args:
+            orig_mol: Original molecule (before FragmentParent)
+            parent_mol: Parent molecule (after FragmentParent)
+
+        Returns:
+            SMILES string of salt components or None if no salts/mixture detected
+        """
+        try:
+            # Get all fragments from original molecule
+            orig_frags = Chem.GetMolFrags(orig_mol, asMols=True)
+
+            # If only one fragment, no salt
+            if len(orig_frags) <= 1:
+                return None
+
+            # Get canonical SMILES of parent for comparison
+            parent_smiles = Chem.MolToSmiles(parent_mol, canonical=True)
+
+            # Separate fragments into salts vs potential mixture components
+            salt_frags = []
+            mixture_frags = []
+
+            for frag in orig_frags:
+                frag_smiles = Chem.MolToSmiles(frag, canonical=True)
+
+                # Skip the parent fragment
+                if frag_smiles == parent_smiles:
+                    continue
+
+                # Classify fragment as salt or mixture component
+                num_heavy = frag.GetNumHeavyAtoms()
+                has_charge = any(atom.GetFormalCharge() != 0 for atom in frag.GetAtoms())
+
+                # More nuanced classification
+                if has_charge and num_heavy <= 10:  # Small charged fragment - likely a salt
+                    salt_frags.append(frag_smiles)
+                elif not has_charge and num_heavy <= 6:  # Small neutral - could be solvent/salt
+                    salt_frags.append(frag_smiles)
+                else:
+                    # Large neutral fragment - likely part of a mixture
+                    mixture_frags.append(frag_smiles)
+
+            # Check if this looks like a mixture
+            if mixture_frags:
+                # Log mixture detection
+                total_frags = len(orig_frags)
+                log.warning(
+                    f"Mixture detected: {total_frags} total fragments, "
+                    f"{len(mixture_frags)} large neutral organics. "
+                    f"Removing: {'.'.join(mixture_frags + salt_frags)}"
+                )
+                # Return None for mixtures - don't pollute the salt column
+                return None
+
+            # Return actual salts only
+            return ".".join(salt_frags) if salt_frags else None
+
+        except Exception as e:
+            log.info(f"Salt extraction failed: {e}")
+            return None
+
+
+def standardize(
+    df: pd.DataFrame,
+    canonicalize_tautomer: bool = True,
+    extract_salts: bool = True,
+) -> pd.DataFrame:
+    """
+    Standardize molecules in a DataFrame for ADMET modeling
+
+    Args:
+        df: Input DataFrame with SMILES column
+        canonicalize_tautomer: Whether to canonicalize tautomers (default: True)
+        extract_salts: Whether to remove and extract salts (default: True)
+                      If False, keeps full molecule with salts/counterions intact,
+                      skipping charge neutralization to preserve ionic character
+
+    Returns:
+        DataFrame with:
+        - orig_smiles: Original SMILES (preserved)
+        - smiles: Standardized SMILES (working column for downstream)
+        - salt: Removed salt/counterion SMILES (only if extract_salts=True)
+                None for mixtures or when no true salts present
+    """
+
+    # Check for the smiles column (any capitalization)
+    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")
+
+    # Copy input DataFrame to avoid modifying original
+    result = df.copy()
+
+    # Preserve original SMILES if not already saved
+    if "orig_smiles" not in result.columns:
+        result["orig_smiles"] = result[smiles_column]
+
+    # Initialize standardizer with salt removal control
+    standardizer = MolStandardizer(canonicalize_tautomer=canonicalize_tautomer, remove_salts=extract_salts)
+
+    def process_smiles(smiles: str) -> pd.Series:
+        """
+        Process a single SMILES string through standardization pipeline
+
+        Args:
+            smiles: Input SMILES string
+
+        Returns:
+            Series with standardized SMILES and extracted salt (if applicable)
+        """
+        # Handle missing values
+        if pd.isna(smiles) or smiles == "":
+            log.error("Encountered missing or empty SMILES string")
+            return pd.Series({"smiles": None, "salt": None})
+
+        # Parse molecule
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            log.error(f"Invalid SMILES: {smiles}")
+            return pd.Series({"smiles": None, "salt": None})
+
+        # Full standardization with optional salt removal
+        std_mol, salt_smiles = standardizer.standardize(mol)
+
+        # After standardization, validate the result
+        if std_mol is not None:
+            # Check if molecule is reasonable
+            if std_mol.GetNumAtoms() == 0 or std_mol.GetNumAtoms() > 200:  # Arbitrary limits
+                log.error(f"Unusual molecule size: {std_mol.GetNumAtoms()} atoms")
+
+        if std_mol is None:
+            return pd.Series(
+                {
+                    "smiles": None,
+                    "salt": salt_smiles,  # May have extracted salt even if full standardization failed
+                }
+            )
+
+        # Convert back to SMILES
+        return pd.Series(
+            {"smiles": Chem.MolToSmiles(std_mol, canonical=True), "salt": salt_smiles if extract_salts else None}
+        )
+
+    # Process molecules
+    processed = result[smiles_column].apply(process_smiles)
+
+    # Update the dataframe with processed results
+    for col in ["smiles", "salt"]:
+        result[col] = processed[col]
+
+    return result
+
+
+if __name__ == "__main__":
+    import time
+    from workbench.api import DataSource
+
+    # Test with DataFrame including various salt forms
+    test_data = pd.DataFrame(
+        {
+            "smiles": [
+                # Organic salts
+                "[Na+].CC(=O)[O-]",  # Sodium acetate
+                "CC(=O)O.CCN",  # Acetic acid + ethylamine (acid-base pair)
+                # Tautomers
+                "CC(=O)CC(C)=O",  # Acetylacetone - tautomer
+                "c1ccc(O)nc1",  # 2-hydroxypyridine/2-pyridone - tautomer
+                # Multi-fragment
+                "CCO.CC",  # Ethanol + methane mixture
+                # Simple organics
+                "CC(C)(C)c1ccccc1",  # tert-butylbenzene
+                # Carbonate salts
+                "[Na+].[Na+].[O-]C([O-])=O",  # Sodium carbonate
+                "[Li+].[Li+].[O-]C([O-])=O",  # Lithium carbonate
+                "[K+].[K+].[O-]C([O-])=O",  # Potassium carbonate
+                "[Mg++].[O-]C([O-])=O",  # Magnesium carbonate
+                "[Ca++].[O-]C([O-])=O",  # Calcium carbonate
+                # Drug salts
+                "CC(C)NCC(O)c1ccc(O)c(O)c1.Cl",  # Isoproterenol HCl
+                "CN1CCC[C@H]1c2cccnc2.[Cl-]",  # Nicotine HCl
+                # Tautomer with salt
+                "c1ccc(O)nc1.Cl",  # 2-hydroxypyridine with HCl
+                # Edge cases
+                None,  # Missing value
+                "INVALID",  # Invalid SMILES
+            ],
+            "compound_id": [f"C{i:03d}" for i in range(1, 17)],
+        }
+    )
+
+    # General test
+    standardize(test_data)
+
+    # Remove the last two rows to avoid errors with None and INVALID
+    test_data = test_data.iloc[:-2].reset_index(drop=True)
+
+    # Test WITHOUT salt removal (keeps full molecule)
+    print("\nStandardization KEEPING salts (extract_salts=False):")
+    print("This preserves the full molecule including counterions")
+    result_keep = standardize(test_data, extract_salts=False, canonicalize_tautomer=True)
+    display_cols = ["compound_id", "orig_smiles", "smiles", "salt"]
+    print(result_keep[display_cols].to_string())
+
+    # Test WITH salt removal
+    print("\n" + "=" * 70)
+    print("Standardization REMOVING salts (extract_salts=True):")
+    print("This extracts parent molecule and records salt information")
+    result_remove = standardize(test_data, extract_salts=True, canonicalize_tautomer=True)
+    print(result_remove[display_cols].to_string())
+
+    # Test WITHOUT tautomerization (keeping salts)
+    print("\n" + "=" * 70)
+    print("Standardization KEEPING salts, NO tautomerization:")
+    result_no_taut = standardize(test_data, extract_salts=False, canonicalize_tautomer=False)
+    print(result_no_taut[display_cols].to_string())
+
+    # Show the difference for salt-containing molecules
+    print("\n" + "=" * 70)
+    print("Comparison showing differences:")
+    for idx, row in result_keep.iterrows():
+        keep_smiles = row["smiles"]
+        remove_smiles = result_remove.loc[idx, "smiles"]
+        no_taut_smiles = result_no_taut.loc[idx, "smiles"]
+        salt = result_remove.loc[idx, "salt"]
+
+        # Show differences when they exist
+        if keep_smiles != remove_smiles or keep_smiles != no_taut_smiles:
+            print(f"\n{row['compound_id']} ({row['orig_smiles']}):")
+            if keep_smiles != no_taut_smiles:
+                print(f"  With salt + taut:    {keep_smiles}")
+                print(f"  With salt, no taut:  {no_taut_smiles}")
+            if keep_smiles != remove_smiles:
+                print(f"  Parent only + taut:  {remove_smiles}")
+            if salt:
+                print(f"  Extracted salt:      {salt}")
+
+    # Summary statistics
+    print("\n" + "=" * 70)
+    print("Summary:")
+    print(f"Total molecules: {len(result_remove)}")
+    print(f"Molecules with salts: {result_remove['salt'].notna().sum()}")
+    unique_salts = result_remove["salt"].dropna().unique()
+    print(f"Unique salts found: {unique_salts[:5].tolist()}")
+
+    # Get a real dataset from Workbench and time the standardization
+    ds = DataSource("aqsol_data")
+    df = ds.pull_dataframe()[["id", "smiles"]]
+    start_time = time.time()
+    std_df = standardize(df, extract_salts=True, canonicalize_tautomer=True)
+    end_time = time.time()
+    print(f"\nStandardized {len(std_df)} molecules from Workbench in {end_time - start_time:.2f} seconds")
+    print(std_df.head())
+    print(f"Molecules with salts: {std_df['salt'].notna().sum()}")
+    unique_salts = std_df["salt"].dropna().unique()
+    print(f"Unique salts found: {unique_salts[:5].tolist()}")

workbench/model_scripts/custom_models/chem_info/molecular_descriptors.py

@@ -7,13 +7,13 @@
 #
 import argparse
 import os
-import joblib
 from io import StringIO
 import pandas as pd
 import json
 
 # Local imports
-from local_utils import compute_molecular_descriptors
+from mol_standardize import standardize
+from mol_descriptors import compute_descriptors
 
 
 # TRAINING SECTION
@@ -32,15 +32,12 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     # This model doesn't get trained, it just a feature creation 'model'
-
-    # Sagemaker seems to get upset if we don't save a model, so we'll create a placeholder model
-    placeholder_model = {}
-    joblib.dump(placeholder_model, os.path.join(args.model_dir, "model.joblib"))
+    # So we don't need to do anything here
 
 
 # Model loading and prediction functions
 def model_fn(model_dir):
-    return joblib.load(os.path.join(model_dir, "model.joblib"))
+    return None
 
 
 def input_fn(input_data, content_type):
@@ -78,6 +75,7 @@ def output_fn(output_df, accept_type):
 # Prediction function
 def predict_fn(df, model):
 
-    # Compute the Molecular Descriptors
-    df = compute_molecular_descriptors(df)
+    # Standardize the molecule (extract salts) and then compute descriptors
+    df = standardize(df, extract_salts=True)
+    df = compute_descriptors(df)
     return df

workbench/model_scripts/custom_models/uq_models/generated_model_script.py

@@ -1,6 +1,7 @@
 # Model: NGBoost Regressor with Distribution output
 from ngboost import NGBRegressor
-from xgboost import XGBRegressor  # Base Estimator
+from ngboost.distns import Cauchy, T
+from xgboost import XGBRegressor  # Point Estimator
 from sklearn.model_selection import train_test_split
 
 # Model Performance Scores
@@ -26,12 +27,12 @@ from proximity import Proximity
 
 # Template Placeholders
 TEMPLATE_PARAMS = {
-    "id_column": "udm_mol_bat_id",
-    "target": "udm_asy_res_intrinsic_clearance_ul_per_min_per_mg_protein",
-    "features": ['bcut2d_logplow', 'numradicalelectrons', 'smr_vsa5', 'fr_lactam', 'fr_morpholine', 'fr_aldehyde', 'slogp_vsa1', 'fr_amidine', 'bpol', 'fr_ester', 'fr_azo', 'kappa3', 'peoe_vsa5', 'fr_ketone_topliss', 'vsa_estate9', 'estate_vsa9', 'bcut2d_mrhi', 'fr_ndealkylation1', 'numrotatablebonds', 'minestateindex', 'fr_quatn', 'peoe_vsa3', 'fr_epoxide', 'fr_aniline', 'minpartialcharge', 'fr_nitroso', 'fpdensitymorgan2', 'fr_oxime', 'fr_sulfone', 'smr_vsa1', 'kappa1', 'fr_pyridine', 'numaromaticrings', 'vsa_estate6', 'molmr', 'estate_vsa1', 'fr_dihydropyridine', 'vsa_estate10', 'fr_alkyl_halide', 'chi2n', 'fr_thiocyan', 'fpdensitymorgan1', 'fr_unbrch_alkane', 'slogp_vsa9', 'chi4n', 'fr_nitro_arom', 'fr_al_oh', 'fr_furan', 'fr_c_s', 'peoe_vsa8', 'peoe_vsa14', 'numheteroatoms', 'fr_ndealkylation2', 'maxabspartialcharge', 'vsa_estate2', 'peoe_vsa7', 'apol', 'numhacceptors', 'fr_tetrazole', 'vsa_estate1', 'peoe_vsa9', 'naromatom', 'bcut2d_chghi', 'fr_sh', 'fr_halogen', 'slogp_vsa4', 'fr_benzodiazepine', 'molwt', 'fr_isocyan', 'fr_prisulfonamd', 'maxabsestateindex', 'minabsestateindex', 'peoe_vsa11', 'slogp_vsa12', 'estate_vsa5', 'numaliphaticcarbocycles', 'bcut2d_mwlow', 'slogp_vsa7', 'fr_allylic_oxid', 'fr_methoxy', 'fr_nh0', 'fr_coo2', 'fr_phenol', 'nacid', 'nbase', 'chi3v', 'fr_ar_nh', 'fr_nitrile', 'fr_imidazole', 'fr_urea', 'bcut2d_mrlow', 'chi1', 'smr_vsa6', 'fr_aryl_methyl', 'narombond', 'fr_alkyl_carbamate', 'fr_piperzine', 'exactmolwt', 'qed', 'chi0n', 'fr_sulfonamd', 'fr_thiazole', 'numvalenceelectrons', 'fr_phos_acid', 'peoe_vsa12', 'fr_nh1', 'fr_hdrzine', 'fr_c_o_nocoo', 'fr_lactone', 'estate_vsa6', 'bcut2d_logphi', 'vsa_estate7', 'peoe_vsa13', 'numsaturatedcarbocycles', 'fr_nitro', 'fr_phenol_noorthohbond', 'rotratio', 'fr_barbitur', 'fr_isothiocyan', 'balabanj', 'fr_arn', 'fr_imine', 'maxpartialcharge', 'fr_sulfide', 'slogp_vsa11', 'fr_hoccn', 'fr_n_o', 'peoe_vsa1', 'slogp_vsa6', 'heavyatommolwt', 'fractioncsp3', 'estate_vsa8', 'peoe_vsa10', 'numaliphaticrings', 'fr_thiophene', 'maxestateindex', 'smr_vsa10', 'labuteasa', 'smr_vsa2', 'fpdensitymorgan3', 'smr_vsa9', 'slogp_vsa10', 'numaromaticheterocycles', 'fr_nh2', 'fr_diazo', 'chi3n', 'fr_ar_coo', 'slogp_vsa5', 'fr_bicyclic', 'fr_amide', 'estate_vsa10', 'fr_guanido', 'chi1n', 'numsaturatedrings', 'fr_piperdine', 'fr_term_acetylene', 'estate_vsa4', 'slogp_vsa3', 'fr_coo', 'fr_ether', 'estate_vsa7', 'bcut2d_chglo', 'fr_oxazole', 'peoe_vsa6', 'hallkieralpha', 'peoe_vsa2', 'chi2v', 'nocount', 'vsa_estate5', 'fr_nhpyrrole', 'fr_al_coo', 'bertzct', 'estate_vsa11', 'minabspartialcharge', 'slogp_vsa8', 'fr_imide', 'kappa2', 'numaliphaticheterocycles', 'numsaturatedheterocycles', 'fr_hdrzone', 'smr_vsa4', 'fr_ar_n', 'nrot', 'smr_vsa8', 'slogp_vsa2', 'chi4v', 'fr_phos_ester', 'fr_para_hydroxylation', 'smr_vsa3', 'nhohcount', 'estate_vsa2', 'mollogp', 'tpsa', 'fr_azide', 'peoe_vsa4', 'numhdonors', 'fr_al_oh_notert', 'fr_c_o', 'chi0', 'fr_nitro_arom_nonortho', 'vsa_estate3', 'fr_benzene', 'fr_ketone', 'vsa_estate8', 'smr_vsa7', 'fr_ar_oh', 'fr_priamide', 'ringcount', 'estate_vsa3', 'numaromaticcarbocycles', 'bcut2d_mwhi', 'chi1v', 'heavyatomcount', 'vsa_estate4', 'chi0v'],
+    "id_column": "udm_mol_id",
+    "target": "udm_asy_res_value",
+    "features": ['bcut2d_logplow', 'numradicalelectrons', 'smr_vsa5', 'fr_lactam', 'fr_morpholine', 'fr_aldehyde', 'slogp_vsa1', 'fr_amidine', 'bpol', 'fr_ester', 'fr_azo', 'kappa3', 'peoe_vsa5', 'fr_ketone_topliss', 'vsa_estate9', 'estate_vsa9', 'bcut2d_mrhi', 'fr_ndealkylation1', 'numrotatablebonds', 'minestateindex', 'fr_quatn', 'peoe_vsa3', 'fr_epoxide', 'fr_aniline', 'minpartialcharge', 'fr_nitroso', 'fpdensitymorgan2', 'fr_oxime', 'fr_sulfone', 'smr_vsa1', 'kappa1', 'fr_pyridine', 'numaromaticrings', 'vsa_estate6', 'molmr', 'estate_vsa1', 'fr_dihydropyridine', 'vsa_estate10', 'fr_alkyl_halide', 'chi2n', 'fr_thiocyan', 'fpdensitymorgan1', 'fr_unbrch_alkane', 'slogp_vsa9', 'chi4n', 'fr_nitro_arom', 'fr_al_oh', 'fr_furan', 'fr_c_s', 'peoe_vsa8', 'peoe_vsa14', 'numheteroatoms', 'fr_ndealkylation2', 'maxabspartialcharge', 'vsa_estate2', 'peoe_vsa7', 'apol', 'numhacceptors', 'fr_tetrazole', 'vsa_estate1', 'peoe_vsa9', 'naromatom', 'bcut2d_chghi', 'fr_sh', 'fr_halogen', 'slogp_vsa4', 'fr_benzodiazepine', 'molwt', 'fr_isocyan', 'fr_prisulfonamd', 'maxabsestateindex', 'minabsestateindex', 'peoe_vsa11', 'slogp_vsa12', 'estate_vsa5', 'numaliphaticcarbocycles', 'bcut2d_mwlow', 'slogp_vsa7', 'fr_allylic_oxid', 'fr_methoxy', 'fr_nh0', 'fr_coo2', 'fr_phenol', 'nacid', 'nbase', 'chi3v', 'fr_ar_nh', 'fr_nitrile', 'fr_imidazole', 'fr_urea', 'bcut2d_mrlow', 'chi1', 'smr_vsa6', 'fr_aryl_methyl', 'narombond', 'fr_alkyl_carbamate', 'fr_piperzine', 'exactmolwt', 'qed', 'chi0n', 'fr_sulfonamd', 'fr_thiazole', 'numvalenceelectrons', 'fr_phos_acid', 'peoe_vsa12', 'fr_nh1', 'fr_hdrzine', 'fr_c_o_nocoo', 'fr_lactone', 'estate_vsa6', 'bcut2d_logphi', 'vsa_estate7', 'peoe_vsa13', 'numsaturatedcarbocycles', 'fr_nitro', 'fr_phenol_noorthohbond', 'rotratio', 'fr_barbitur', 'fr_isothiocyan', 'balabanj', 'fr_arn', 'fr_imine', 'maxpartialcharge', 'fr_sulfide', 'slogp_vsa11', 'fr_hoccn', 'fr_n_o', 'peoe_vsa1', 'slogp_vsa6', 'heavyatommolwt', 'fractioncsp3', 'estate_vsa8', 'peoe_vsa10', 'numaliphaticrings', 'fr_thiophene', 'maxestateindex', 'smr_vsa10', 'labuteasa', 'smr_vsa2', 'fpdensitymorgan3', 'smr_vsa9', 'slogp_vsa10', 'numaromaticheterocycles', 'fr_nh2', 'fr_diazo', 'chi3n', 'fr_ar_coo', 'slogp_vsa5', 'fr_bicyclic', 'fr_amide', 'estate_vsa10', 'fr_guanido', 'chi1n', 'numsaturatedrings', 'fr_piperdine', 'fr_term_acetylene', 'estate_vsa4', 'slogp_vsa3', 'fr_coo', 'fr_ether', 'estate_vsa7', 'bcut2d_chglo', 'fr_oxazole', 'peoe_vsa6', 'hallkieralpha', 'peoe_vsa2', 'chi2v', 'nocount', 'vsa_estate5', 'fr_nhpyrrole', 'fr_al_coo', 'bertzct', 'estate_vsa11', 'minabspartialcharge', 'slogp_vsa8', 'fr_imide', 'kappa2', 'numaliphaticheterocycles', 'numsaturatedheterocycles', 'fr_hdrzone', 'smr_vsa4', 'fr_ar_n', 'nrot', 'smr_vsa8', 'slogp_vsa2', 'chi4v', 'fr_phos_ester', 'fr_para_hydroxylation', 'smr_vsa3', 'nhohcount', 'estate_vsa2', 'mollogp', 'tpsa', 'fr_azide', 'peoe_vsa4', 'numhdonors', 'fr_al_oh_notert', 'fr_c_o', 'chi0', 'fr_nitro_arom_nonortho', 'vsa_estate3', 'fr_benzene', 'fr_ketone', 'vsa_estate8', 'smr_vsa7', 'fr_ar_oh', 'fr_priamide', 'ringcount', 'estate_vsa3', 'numaromaticcarbocycles', 'bcut2d_mwhi', 'chi1v', 'heavyatomcount', 'vsa_estate4', 'chi0v', 'chiral_centers', 'r_cnt', 's_cnt', 'db_stereo', 'e_cnt', 'z_cnt', 'chiral_fp', 'db_fp'],
     "compressed_features": [],
     "train_all_data": False,
-    "track_columns": ['udm_asy_res_intrinsic_clearance_ul_per_min_per_mg_protein']
+    "track_columns": "udm_asy_res_value"
 }
 
 
@@ -106,8 +107,10 @@ def convert_categorical_types(df: pd.DataFrame, features: list, category_mapping
     return df, category_mappings
 
 
-def decompress_features(df: pd.DataFrame, features: List[str], compressed_features: List[str]) -> Tuple[pd.DataFrame, List[str]]:
-    """Prepare features for the XGBoost model
+def decompress_features(
+    df: pd.DataFrame, features: List[str], compressed_features: List[str]
+) -> Tuple[pd.DataFrame, List[str]]:
+    """Prepare features for the model by decompressing bitstring features
 
     Args:
         df (pd.DataFrame): The features DataFrame
@@ -132,7 +135,7 @@ def decompress_features(df: pd.DataFrame, features: List[str], compressed_featur
         )
 
     # Decompress the specified compressed features
-    decompressed_features = features
+    decompressed_features = features.copy()
     for feature in compressed_features:
         if (feature not in df.columns) or (feature not in features):
             print(f"Feature '{feature}' not in the features list, skipping decompression.")
@@ -227,7 +230,14 @@ if __name__ == "__main__":
 
     # We're using XGBoost for point predictions and NGBoost for uncertainty quantification
     xgb_model = XGBRegressor()
-    ngb_model = NGBRegressor()
+    ngb_model = NGBRegressor()  # Dist=Cauchy) Seems to give HUGE prediction intervals
+    ngb_model = NGBRegressor(
+        Dist=T,
+        learning_rate=0.005,
+        minibatch_frac=0.1,  # Very small batches
+        col_sample=0.8  # This parameter DOES exist
+    ) # Testing this out
+    print("NGBoost using T distribution for uncertainty quantification")
 
     # Prepare features and targets for training
     X_train = df_train[features]