# Multi-Modal Medical Prediction with EXCAP

[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
[![PyTorch](https://img.shields.io/badge/PyTorch-1.12+-red.svg)](https://pytorch.org/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

A comprehensive framework for multi-modal medical prediction using the EXCAP (Efficient Adaptive Multi-modal Causal) architecture. This repository contains the implementation for the paper "Multi-Modal Medical Prediction with Efficient Adaptive Causal Attention".

## 🚀 Quick Start

### Installation

```bash
# Navigate to release directory
cd release

# Install dependencies
pip install -r requirements.txt
pip install -e .
```

### One-Click Training

We provide a one-click training script that automatically detects and uses offline sample data:

```bash
# Run with default settings (auto-detects offline data)
./run_multi_task_training.sh

# Or with custom parameters
./run_multi_task_training.sh --gpu 0 --epochs 20 --batch_size 256
```

**Features:**
- ✅ Automatically detects offline sample data (no database required)
- ✅ All parameters configurable via command line
- ✅ See `QUICK_START.md` for detailed usage

### Manual Training

**Multi-task training (Mortality + LOS):**
```bash
python train_mortality_los_complete.py \
    --gpu 0 \
    --epochs 20 \
    --lr 5e-5 \
    --batch_size 256 \
    --num_workers 1 \
    --max_samples 80000 \
    --max_val_samples 10000 \
    --experiment_name demo_multi_task
```

#### Step 2: Monitor Training

Training logs and TensorBoard files will be saved to `runs/classify/{experiment_name}/`.

View training progress:
```bash
tensorboard --logdir runs/classify
```

#### Step 3: Check Results

After training, results are saved in:
- Model checkpoints: `runs/classify/{experiment_name}/best_model.pth`
- Metrics: `runs/classify/{experiment_name}/final_report.json`
- Test results: `runs/classify/{experiment_name}/test_results.json`

## 📋 Command Line Arguments

### Common Arguments

- `--gpu`: GPU ID (default: '0')
- `--epochs`: Number of training epochs (default: 20)
- `--lr`: Learning rate (default: 5e-4)
- `--batch_size`: Batch size (default: 256)
- `--num_workers`: Number of data loading workers (default: 1)
- `--experiment_name`: Name for this experiment (default: auto-generated)
- `--log_dir`: Directory for logs (default: 'runs/classify')
- `--max_samples`: Maximum training samples (0 = all, default: 80000)
- `--max_val_samples`: Maximum validation samples (0 = all, default: 10000)
- `--max_test_samples`: Maximum test samples (0 = all, default: 10000)

## 🔧 Configuration

The framework uses YAML configuration files. Key configuration options:

```yaml
# Model configuration
model:
  shared_dim: 128
  image_dim: 128
  text_dim: 128
  ts_dim: 64
  static_dim: 10
  use_excap: true
  max_segments: 8
  segment_len: 64

```

Edit `exp/mimic_data/exp_mortality_24h48h_los.yaml` to configure your model parameters.

## 📊 Expected Output

During training, you will see:
- Training progress with loss values for each task
- Validation metrics: AUROC, AUPRC, Accuracy, F1-score
- Model checkpoints saved automatically

## 🏗️ Project Structure

```
release/
├── README.md                          # Main documentation
├── requirements.txt                   # Dependencies
├── train_mortality_los_complete.py   # Multi-task training script
├── run_multi_task_training.sh        # One-click training script
├── baseline/                          # Baseline models
├── model/                             # Model implementations
├── exp/                               # Configuration files
├── datapress/                         # Data processing
└── data_dir/                          # Sample data
```


## 📊 Baseline Models

We provide baseline implementations for comparison:

### Available Baselines

- **TimesNet**: Time series forecasting model
- **iTransformer**: Inverted Transformer for time series
- **PatchTST**: Patch-based Time Series Transformer

### Training Baselines

```bash
cd baseline
python3 train_baseline_models.py --model TimesNet --gpu 0 --epochs 20
```

**Available models**: `TimesNet`, `iTransformer`, `PatchTST`

**Example:**
```bash
# Train TimesNet
python3 train_baseline_models.py --model TimesNet --epochs 20 --batch_size 256

# Train iTransformer with custom settings
python3 train_baseline_models.py --model iTransformer --epochs 50 --lr 5e-5

# Train PatchTST for quick testing
python3 train_baseline_models.py --model PatchTST --max_samples 1000
```

The baseline script uses the same data splits and evaluation metrics as the main model for fair comparison. See `baseline/README.md` for detailed usage.

## 📈 Performance

The model reports the following metrics:
- **AUROC**: Area Under ROC Curve
- **AUPRC**: Area Under Precision-Recall Curve
- **Accuracy**: Classification accuracy
- **F1-score**: F1 score
- **Confusion Matrix**: TP, TN, FP, FN

## 🔍 Troubleshooting

**Common issues:**
- **Import errors**: Ensure you're in the release directory and dependencies are installed
- **CUDA out of memory**: Reduce `--batch_size` or `--max_samples`
- **Data loading fails**: Try `--num_workers 0` to disable multiprocessing
- **Config file not found**: The script auto-detects configs; ensure you're in the release directory

See `QUICK_START.md` for detailed troubleshooting.

## 📚 Key Features

- Multi-modal fusion (images, text, time-series, static features)
- Efficient adaptive causal attention (EXCAP)
- Dynamic time-series segmentation
- Robust missing data handling
- Automatic feature normalization

## 📄 Citation

```bibtex
@article{excap2024,
  title={Multi-Modal Medical Prediction with Efficient Adaptive Causal Attention},
  author={Anonymous},
  journal={Anonymous},
  year={2024}
}
```

## 📚 Additional Documentation

- `QUICK_START.md` - Detailed usage guide
- `COMPLETE_DOCUMENTATION.md` - Technical documentation
- `OFFLINE_DATA_README.md` - Offline data usage
- `baseline/README.md` - Baseline models guide

## 📝 License

This project is licensed under the MIT License.
