#!/usr/bin/env python3
"""
TDC Data Loading and Splitting Template
This script demonstrates how to load TDC datasets and apply different
splitting strategies for model training and evaluation.
Usage:
python load_and_split_data.py
"""
from tdc.single_pred import ADME
from tdc.multi_pred import DTI
from tdc import Evaluator
import pandas as pd
def load_single_pred_example():
"""
Example: Loading and splitting a single-prediction dataset (ADME)
"""
print("=" * 60)
print("Example 1: Single-Prediction Task (ADME)")
print("=" * 60)
# Load Caco2 dataset (intestinal permeability)
print("\nLoading Caco2_Wang dataset...")
data = ADME(name='Caco2_Wang')
# Get basic dataset info
print(f"\nDataset size: {len(data.get_data())} molecules")
data.print_stats()
# Method 1: Scaffold split (default, recommended)
print("\n--- Scaffold Split ---")
split = data.get_split(method='scaffold', seed=42, frac=[0.7, 0.1, 0.2])
train = split['train']
valid = split['valid']
test = split['test']
print(f"Train: {len(train)} molecules")
print(f"Valid: {len(valid)} molecules")
print(f"Test: {len(test)} molecules")
# Display sample data
print("\nSample training data:")
print(train.head(3))
# Method 2: Random split
print("\n--- Random Split ---")
split_random = data.get_split(method='random', seed=42, frac=[0.8, 0.1, 0.1])
print(f"Train: {len(split_random['train'])} molecules")
print(f"Valid: {len(split_random['valid'])} molecules")
print(f"Test: {len(split_random['test'])} molecules")
return split
def load_multi_pred_example():
"""
Example: Loading and splitting a multi-prediction dataset (DTI)
"""
print("\n" + "=" * 60)
print("Example 2: Multi-Prediction Task (DTI)")
print("=" * 60)
# Load BindingDB Kd dataset (drug-target interactions)
print("\nLoading BindingDB_Kd dataset...")
data = DTI(name='BindingDB_Kd')
# Get basic dataset info
full_data = data.get_data()
print(f"\nDataset size: {len(full_data)} drug-target pairs")
print(f"Unique drugs: {full_data['Drug_ID'].nunique()}")
print(f"Unique targets: {full_data['Target_ID'].nunique()}")
# Method 1: Random split
print("\n--- Random Split ---")
split_random = data.get_split(method='random', seed=42)
print(f"Train: {len(split_random['train'])} pairs")
print(f"Valid: {len(split_random['valid'])} pairs")
print(f"Test: {len(split_random['test'])} pairs")
# Method 2: Cold drug split (unseen drugs in test)
print("\n--- Cold Drug Split ---")
split_cold_drug = data.get_split(method='cold_drug', seed=42)
train = split_cold_drug['train']
test = split_cold_drug['test']
# Verify no drug overlap
train_drugs = set(train['Drug_ID'])
test_drugs = set(test['Drug_ID'])
overlap = train_drugs & test_drugs
print(f"Train: {len(train)} pairs, {len(train_drugs)} unique drugs")
print(f"Test: {len(test)} pairs, {len(test_drugs)} unique drugs")
print(f"Drug overlap: {len(overlap)} (should be 0)")
# Method 3: Cold target split (unseen targets in test)
print("\n--- Cold Target Split ---")
split_cold_target = data.get_split(method='cold_target', seed=42)
train = split_cold_target['train']
test = split_cold_target['test']
train_targets = set(train['Target_ID'])
test_targets = set(test['Target_ID'])
overlap = train_targets & test_targets
print(f"Train: {len(train)} pairs, {len(train_targets)} unique targets")
print(f"Test: {len(test)} pairs, {len(test_targets)} unique targets")
print(f"Target overlap: {len(overlap)} (should be 0)")
# Display sample data
print("\nSample DTI data:")
print(full_data.head(3))
return split_cold_drug
def evaluation_example(split):
"""
Example: Evaluating model predictions with TDC evaluators
"""
print("\n" + "=" * 60)
print("Example 3: Model Evaluation")
print("=" * 60)
test = split['test']
# For demonstration, create dummy predictions
# In practice, replace with your model's predictions
import numpy as np
np.random.seed(42)
# Simulate predictions (replace with model.predict(test['Drug']))
y_true = test['Y'].values
y_pred = y_true + np.random.normal(0, 0.5, len(y_true)) # Add noise
# Evaluate with different metrics
print("\nEvaluating predictions...")
# Regression metrics
mae_evaluator = Evaluator(name='MAE')
mae = mae_evaluator(y_true, y_pred)
print(f"MAE: {mae:.4f}")
rmse_evaluator = Evaluator(name='RMSE')
rmse = rmse_evaluator(y_true, y_pred)
print(f"RMSE: {rmse:.4f}")
r2_evaluator = Evaluator(name='R2')
r2 = r2_evaluator(y_true, y_pred)
print(f"R²: {r2:.4f}")
spearman_evaluator = Evaluator(name='Spearman')
spearman = spearman_evaluator(y_true, y_pred)
print(f"Spearman: {spearman:.4f}")
def custom_split_example():
"""
Example: Creating custom splits with different fractions
"""
print("\n" + "=" * 60)
print("Example 4: Custom Split Fractions")
print("=" * 60)
data = ADME(name='HIA_Hou')
# Custom split fractions
custom_fracs = [
([0.6, 0.2, 0.2], "60/20/20 split"),
([0.8, 0.1, 0.1], "80/10/10 split"),
([0.7, 0.15, 0.15], "70/15/15 split")
]
for frac, description in custom_fracs:
split = data.get_split(method='scaffold', seed=42, frac=frac)
print(f"\n{description}:")
print(f" Train: {len(split['train'])} ({frac[0]*100:.0f}%)")
print(f" Valid: {len(split['valid'])} ({frac[1]*100:.0f}%)")
print(f" Test: {len(split['test'])} ({frac[2]*100:.0f}%)")
def main():
"""
Main function to run all examples
"""
print("\n" + "=" * 60)
print("TDC Data Loading and Splitting Examples")
print("=" * 60)
# Example 1: Single prediction with scaffold split
split = load_single_pred_example()
# Example 2: Multi prediction with cold splits
dti_split = load_multi_pred_example()
# Example 3: Model evaluation
evaluation_example(split)
# Example 4: Custom split fractions
custom_split_example()
print("\n" + "=" * 60)
print("Examples completed!")
print("=" * 60)
if __name__ == "__main__":
main()