Machine Learning: From Data Preparation to Survival Prediction

Machine learning on multi-omics data requires rigorous feature engineering, domain-specific statistical methods, and careful validation before clinical deployment. This case study demonstrates Lobster AI's machine learning agents across three complexity levels: preparing hepatocellular carcinoma data for ML readiness assessment, discovering breast cancer biomarkers with bootstrap feature importance and pathway enrichment, and building a complete survival prediction pipeline with Cox proportional hazards modeling and risk stratification. The progression from simple data preparation through biomarker discovery to survival analysis mirrors the real-world ML workflow for translational cancer research.

Agents and Data Sources

This analysis uses the lobster-ml package, which provides a parent-child agent hierarchy for machine learning workflows:

External data sources queried: GEO (Gene Expression Omnibus) via research_agent. Local computation handled by scikit-learn (StandardScaler, train_test_split, LassoCV), scanpy (highly variable gene selection), scikit-survival (Cox regression, log-rank test), and INDRA Discovery API (pathway enrichment).

Simple: Hepatocellular Carcinoma Data Preparation

The first scenario demonstrates ML readiness assessment and feature preparation for a 30-sample hepatocellular carcinoma (HCC) bulk RNA-seq dataset. HCC is the most common primary liver cancer and a major cause of cancer mortality worldwide. Tumor-vs-normal classification from bulk RNA-seq is a canonical ML task that validates whether expression signatures can distinguish malignant from healthy tissue.

Turn 1: Dataset Discovery and Download

The research agent searched GEO for small bulk RNA-seq cancer datasets and identified two candidates: GSE304512 (HCC, 38 samples) and GSE266302 (pseudomyxoma peritonei, 20 samples). Both datasets were queued for download with clearly labeled tumor and normal groups.

lobster query -p bedrock --session-id ml_simple \
  "Search GEO for a small bulk RNA-seq dataset comparing cancer vs normal tissue. \
   I need something with fewer than 100 samples, clearly labeled groups (e.g. tumor \
   vs normal), and expression counts. Find a suitable dataset and download it."

Turn 2: Download Recovery and ML Feature Preparation

The second turn demonstrated Lobster's resilience when encountering unavailable datasets. After two failed download attempts (GSE304512 and GSE50760 format incompatibility), the system autonomously searched for alternatives and successfully downloaded GSE135631 — a 30-sample HCC dataset with balanced tumor and normal groups.

lobster query -p bedrock --session-id ml_simple \
  "Download GSE304512. Once loaded, use the machine learning expert to check \
   ML readiness of the dataset, prepare features using highly variable gene \
   selection with 2000 features and z-score scaling, then create stratified \
   train/test splits (80/20) using the tissue type labels."

Dataset Summary (GSE135631):

Download Recovery Log:

The machine_learning_expert assessed ML readiness and flagged the 32.4% sparsity and small sample size (30 samples — risk of overfitting). Highly variable gene selection with 2,000 features failed due to infinity values in variance calculation on TPM-normalized data (a known edge case when genes have near-zero mean expression). The agent gracefully fell back to using the full 55,789-gene feature set with z-score scaling.

ML Preparation Pipeline:

Train/Test Split:

The final stratified 80/20 split produced properly balanced training (24 samples) and test (6 samples) sets ready for classifier training. This demonstrates Lobster's ability to recover autonomously from data acquisition failures and prepare ML-ready datasets even when feature selection encounters edge cases.

Medium: Breast Cancer Biomarker Discovery

The second scenario demonstrates the full ML biomarker discovery pipeline: dataset acquisition, dual-method feature selection (bootstrap feature importance + LASSO), and biological validation via pathway enrichment. Breast cancer is the most commonly diagnosed cancer worldwide, and estrogen receptor (ER) status is the most important molecular classifier for treatment decisions.

Turn 1: Dataset Discovery

The research agent searched GEO for breast cancer bulk RNA-seq datasets with at least two clear clinical groups and identified GSE240671 (122 samples, neoadjuvant chemotherapy responders vs non-responders).

lobster query -p bedrock --session-id ml_medium \
  "Search GEO for a bulk RNA-seq dataset studying breast cancer with at least \
   two clear groups (e.g. treatment responders vs non-responders, or ER-positive \
   vs ER-negative, or luminal vs basal). I need 30-80 samples with gene expression \
   counts. Find and download the best candidate."

Turn 2: Recovery and ML Pipeline

The initial download of GSE240671 failed due to unparseable file formats (likely raw FASTQ only). Lobster recovered autonomously by searching for an alternative dataset and found the SCAN-B cohort (GSE81538) — a 405-sample Swedish breast cancer gene expression dataset published in Nature Medicine.

lobster query -p bedrock --session-id ml_medium \
  "Option 1 - find an alternative breast cancer RNA-seq dataset. Specifically \
   look for one with processed count matrices available (not raw FASTQ). Once \
   found and downloaded, proceed with the ML pipeline: check readiness, prepare \
   2000 features with variance selection, then run bootstrap feature selection \
   with XGBoost (100 features, 10 bootstrap rounds) using the group \
   labels for classification."

Dataset Summary (GSE81538):

The ML pipeline proceeded end-to-end: readiness assessment confirmed excellent data quality (0% missing values, balanced ER status labels), variance selection reduced 18,802 genes to the top 2,000, and bootstrap feature importance with 10 bootstrap rounds identified 100 robust biomarkers.

ML Feature Preparation:

Top Bootstrap-Selected Biomarkers:

Turn 3: LASSO Feature Selection and Pathway Enrichment

The third turn compared LASSO feature selection against bootstrap feature importance and linked the selected biomarkers to biological pathways via INDRA enrichment.

lobster query -p bedrock --session-id ml_medium \
  "Now run LASSO feature selection on the same prepared features \
   (geo_GSE81538_ml_features) with alpha=0.1 using the ER status label, \
   then perform pathway enrichment on the bootstrap-selected features to \
   identify the biological pathways underlying ER status. Use INDRA for enrichment."

LASSO Feature Selection Results:

The 73.8% feature retention rate indicates insufficient regularization — a properly sparse LASSO would retain significantly fewer features. The alpha=0.1 parameter was the cross-validation selected value, but in practice this retention rate is closer to Ridge regression behavior. This is presented as a comparison point against the bootstrap approach rather than as an optimized sparse model.

Top LASSO-Selected Genes:

GREB1 (Growth Regulation by Estrogen in Breast cancer 1) is the canonical estrogen-responsive gene in breast cancer research. Its appearance in the top 5 LASSO-selected features adds further validation to the computational results.

Pathway Enrichment Results (INDRA Discovery API):

Pathway Enrichment Summary:

Pathway enrichment significance was assessed via INDRA. The multiple testing correction method and significance threshold should be verified when interpreting these results for publication.

Method Comparison: Bootstrap Feature Importance vs LASSO:

Hard: Breast Cancer Survival Pipeline

The third scenario demonstrates the complete survival analysis pipeline: loading survival-annotated gene expression data, dual-method feature selection, Cox proportional hazards modeling, risk stratification, and Kaplan-Meier survival analysis. This exercises the full ML agent hierarchy with parent-child delegation across multiple turns.

Turn 1: Data Loading and ML Feature Preparation

Lobster loaded a 405-patient breast cancer cohort with survival annotations and assessed ML readiness in a single turn.

lobster query -p bedrock --session-id ml_hard_final \
  "Load the file brca_survival_cohort.h5ad from the workspace data directory. \
   This is a breast cancer cohort with 405 samples, 5000 genes, and survival \
   data (survival_months, vital_status columns). Check ML readiness, then \
   prepare ML features with variance selection targeting 2000 genes and z-score scaling."

Cohort Overview:

ML Feature Preparation:

The machine_learning_expert identified complete survival data (292 events, 113 censored, 72% event rate), confirmed 0% sparsity, and prepared ML features by selecting the top 2,000 variance genes from 5,000 candidates with z-score normalization.

Turn 2: Dual Feature Selection

The feature_selection_expert (child agent) applied two independent methods to identify prognostic genes from 2,000 candidates.

lobster query -p bedrock --session-id ml_hard_final \
  "Run bootstrap feature selection on brca_survival_cohort_ml_features \
   using vital_status as the target variable. Use XGBoost with 10 bootstrap \
   rounds, selecting the top 50 prognostic genes. Then also run LASSO feature \
   selection with alpha=0.1 on the same modality using vital_status."

Bootstrap Feature Selection Results:

LASSO Selection Results:

Top Bootstrap-Selected Genes:

Top LASSO-Selected Genes:

Bootstrap feature selection identified SLC7A11 (ferroptosis regulator with established roles in cancer cell survival) as the top-ranked gene. LASSO selection identified BAALC (known poor-prognosis marker in leukemia and solid tumors) and enriched for oncofetal antigens (AFP, PLAC1) and stem cell markers (OLFM4) — biologically coherent with aggressive breast cancer subtypes.

Turn 3: Cox Model, Risk Stratification, and Kaplan-Meier Analysis

The survival_analysis_expert (child agent) trained a Cox proportional hazards model using the top 50 bootstrap-selected genes and performed risk stratification with Kaplan-Meier survival analysis.

lobster query -p bedrock --session-id ml_hard_final \
  "Survival dependencies are now installed. Proceed with the full survival \
   analysis on brca_survival_cohort_ml_features_feature_selected: \
   (1) train Cox proportional hazards model with survival_months and vital_status, \
   (2) optimize risk threshold, (3) Kaplan-Meier stratified by risk_category."

Cox Model Performance:

A C-index (the probability that the model correctly ranks two randomly selected patients by survival time; 0.5 = random, 1.0 = perfect) of 0.676 indicates good prognostic discrimination for genomic survival models. The elastic net regularization compensates for the EPV of 5.84, which is below the traditional EPV=10 rule of thumb but acceptable for genomic survival modeling.

This C-index is an apparent (training set) metric. For publication-grade survival model validation, bootstrap optimism correction (Harrell 2015) or nested cross-validation should be applied. Published genomic signatures (OncotypeDX, MammaPrint) achieve validated C-indices of 0.68-0.75 in independent cohorts.

Risk Stratification:

Log-Rank Test:

Risk stratification revealed a dramatic survival difference: high-risk patients had median survival of 23.4 months versus 74.2 months for low-risk patients (log-rank p=7.3e-20). High-risk patients have 3.2x shorter median survival than low-risk patients — a clinically meaningful difference.

Top 10 Prognostic Biomarkers (Cox model):

The Kaplan-Meier survival curves were automatically generated and exported as publication-ready PNG and PDF files, showing clear visual separation between high-risk and low-risk groups with a median survival difference of 50.8 months.

Limitations

• Bootstrap importance, not formal stability selection. The feature selection method uses 10 bootstrap rounds with importance aggregation, not the Meinbaussen and Buhlmann (2010) stability selection framework with sub-sampling at n/2 and per-family error rate control. Ten rounds provides limited statistical confidence about feature stability.
• Cox C-index is an apparent metric. The reported C-index of 0.676 was computed on the training data without nested cross-validation or bootstrap optimism correction. The true validated C-index is likely lower.
• Survival annotation circularity. Survival endpoints were derived from PAM50 subtype and NHG grade, which are themselves defined by gene expression. The model may be partially recovering the molecular subtype classification rather than discovering independent prognostic biology.
• LASSO under-regularized. With alpha=0.1, LASSO retained 73.8% of features — insufficient sparsity for a meaningful feature selection. A more aggressive regularization (higher alpha via cross-validation) would produce a sparser, more interpretable model.
• HVG selection failed on TPM data. Scanpy's highly_variable_genes produces infinity values from near-zero-mean TPM genes. The fallback to all features is functional but increases computational cost and noise.
• scikit-survival dependency. The survival analysis pipeline requires scikit-survival, which is not included in the default installation. Users must install it separately for Cox PH modeling.

What This Demonstrates

Multi-Agent Coordination

No single agent could produce this analysis. The research_agent handled dataset discovery and acquisition. The data_expert loaded expression matrices. The machine_learning_expert (parent) assessed ML readiness, prepared features, and delegated to child agents. The feature_selection_expert (child) ran bootstrap feature selection and LASSO. The survival_analysis_expert (child) trained Cox models, optimized risk thresholds, and generated Kaplan-Meier curves. The supervisor routed each sub-question to the appropriate specialist and synthesized results across all turns.

Autonomous Recovery

When datasets failed to download (GSE304512, GSE240671, GSE50760), Lobster autonomously searched for alternatives without user intervention. This resilience is critical for real-world bioinformatics workflows where data acquisition rarely succeeds on the first attempt.

Biological Validation

The computational results converged on known biology across all three difficulty levels:

• Simple: Successfully prepared tumor-vs-normal features despite HVG selection failure
• Medium: PGR (FDA-approved marker), AGR2/AGR3 (validated ER+ markers), and GREB1 (canonical ER gene) all appeared in top-ranked features. Pathway enrichment confirmed steroid biosynthesis, hormone receptor signaling, luminal differentiation, and lipid metabolism — the established hallmarks of ER+ breast cancer.
• Hard: OGN achieved individual statistical significance (p=0.0074) in the Cox model. Risk stratification produced a 50.8-month median survival difference with p=7.3e-20 (though survival annotations were derived from molecular features, introducing circularity).

Comparison

Estimates based on these case study sessions. Human researcher timing assumes manual workflows without Lobster.

Reproducibility

To reproduce these analyses, install the machine learning package and run the queries sequentially:

pip install 'lobster-ai[full]==1.0.12'

Simple case (HCC data preparation):

lobster query -p bedrock --session-id ml_simple \
  "Search GEO for a small bulk RNA-seq dataset comparing cancer vs normal tissue. \
   I need something with fewer than 100 samples, clearly labeled groups, and \
   expression counts. Find a suitable dataset and download it."

lobster query -p bedrock --session-id ml_simple \
  "Download GSE304512. Once loaded, use the machine learning expert to check \
   ML readiness, prepare features using highly variable gene selection with \
   2000 features and z-score scaling, then create stratified train/test splits \
   (80/20) using the tissue type labels."

Medium case (breast cancer biomarkers):

lobster query -p bedrock --session-id ml_medium \
  "Search GEO for a bulk RNA-seq dataset studying breast cancer with at least \
   two clear groups. I need 30-80 samples with gene expression counts. Find \
   and download the best candidate."

lobster query -p bedrock --session-id ml_medium \
  "Find an alternative breast cancer RNA-seq dataset with processed count matrices. \
   Once downloaded, proceed with the ML pipeline: check readiness, prepare 2000 \
   features with variance selection, then run bootstrap feature selection \
   (100 features, 10 bootstrap rounds) using the group labels."

lobster query -p bedrock --session-id ml_medium \
  "Run LASSO feature selection on the prepared features with alpha=0.1, then \
   perform pathway enrichment on the stability-selected features using INDRA."

Hard case (survival prediction):

Requires preparing a survival-annotated cohort first (see case study disclosure above), then:

lobster query -p bedrock --session-id ml_hard_final \
  "Load brca_survival_cohort.h5ad. Check ML readiness, then prepare ML features \
   with variance selection targeting 2000 genes and z-score scaling."

lobster query -p bedrock --session-id ml_hard_final \
  "Run bootstrap feature selection on the ML features using vital_status. \
   Use 10 bootstrap rounds, selecting the top 50 prognostic genes. Then also \
   run LASSO feature selection with alpha=0.1 on the same modality."

lobster query -p bedrock --session-id ml_hard_final \
  "Train a Cox proportional hazards model on the stability-selected features \
   using survival_months as time and vital_status as event. Optimize a risk \
   threshold and generate Kaplan-Meier survival curves stratified by risk category."

Session continuity via --session-id ensures each turn builds on prior context. Results are stored in the .lobster_workspace/ directory.