# LucidAtlas 🔍

**Learning Uncertainty-Aware, Covariate-Disentangled, Individualized Atlas Representation**

LucidAtlas is a novel framework for building interpretable, uncertainty-aware atlas representations that capture both population trends and individual variations in spatial data. By combining neural additive models with principled marginalization techniques, LucidAtlas addresses key challenges in medical image analysis and computational anatomy.

## � Key Innovations

- **Dependency-Aware Covariate Disentanglement**: Unlike traditional neural additive models, LucidAtlas properly accounts for statistical dependencies between covariates, providing more accurate interpretations for scientific discovery
- **Spatially-Aware Uncertainty Quantification**: Captures heteroscedastic uncertainty that varies across both spatial locations and covariates
- **Individualized Prediction**: Enables personalized forecasting for longitudinal analysis using cross-sectional data
- **Out-of-Distribution Detection**: Identifies anatomically abnormal cases with high sensitivity and specificity
- **Marginalization Framework**: Supports arbitrary covariate subset conditioning through a principled mathematical approach

## 📈 Capabilities

LucidAtlas simultaneously addresses questions at multiple levels:

- **Covariate Level**: How do factors like age and weight influence anatomical structures?
- **Population Level**: How does anatomy vary across a population? Where is uncertainty highest?
- **Subject Level**: How can we make individualized predictions? Is this subject's anatomy atypical?
- **Spatial Level**: Do covariate effects vary by anatomical location?

## 📋 Requirements

```bash
# Install dependencies from requirements.txt
pip install -r requirements.txt
```

## 🔍 Project Structure

```
LucidAtlas/
├── arg_utils.py                  # Argument utilities for command-line interfaces
├── calculate_global_importance.py # Global feature importance calculation
├── vis_utils.py                  # Visualization utilities
├── commands/                     # Shell scripts and Python command files for experiments
├── configs/                      # Configuration files for different datasets
├── hetreg/                       # Core heteroscedastic regression models
├── model/                        # Neural network model definitions
├── pipeline/                     # Training and evaluation pipeline
└── utilities/                    # General utility functions
```

## 🧠 Core Methodology

LucidAtlas extends neural additive models (NAMs) with several key innovations:

1. **Spatial Dependency Modeling**: Captures how anatomy varies across spatial locations
2. **DeepSets-based Covariate Dependency Modeling**: Models statistical relationships between covariates 
3. **Principled Marginalization**: Enables interpretation of how arbitrary subsets of covariates influence outcomes
4. **Prior Knowledge Integration**: Incorporates domain knowledge (e.g., monotonicity constraints)

## 🏁 Getting Started

### Training a Model

```bash
# Basic training command
python -m pipeline.train --experiment configs/path/to/config.json

# Combined train, test, and visualization with a complete configuration
python pipeline/train_test_vis.py -e configs/airways/log_cv5_es/airway_lucidatlas_full.json
```

Training parameters in the configuration files include:
- Model architecture (NAM structure, monotonicity constraints)
- Dataset selection and preprocessing
- Marginalization settings
- Training hyperparameters (learning rate, batch size)
- Uncertainty modeling options

### Evaluating and Interpreting Models

```bash
# Evaluate population trends and uncertainties
python -m pipeline.test_comp_single --experiment configs/path/to/config.json

# Analyze local feature importance
python -m pipeline.test_local_feat --experiment configs/path/to/config.json

# Calculate global feature importance
python -m pipeline.test_global_feat --experiment configs/path/to/config.json
```

## � Analysis Capabilities

### Marginalized Covariate Effects

Explore how individual covariates or combinations of covariates influence the response variable:

```python
# Analyze effect of a single covariate (e.g., age)
python -m pipeline.analyze_covariate_effect --covariate age --experiment configs/path/to/config.json

# Analyze joint effects of multiple covariates with dependency modeling
python -m pipeline.analyze_joint_effects --covariates age,weight --experiment configs/path/to/config.json
```

### Out-of-Distribution Detection

Detect anatomically abnormal cases using our uncertainty-aware approach:

```bash
# Run OOD detection on a test dataset
python -m pipeline.test_ood_detection --experiment configs/path/to/config.json --test_data path/to/test_data
```

### Individualized Prediction

Make personalized predictions for longitudinal analysis:

```bash
# Predict future anatomy given initial observation
python -m pipeline.predict_individual --subject_id SUBJ001 --time_point T1 --experiment configs/path/to/config.json
```

## 📊 Visualization

The framework includes comprehensive visualization utilities in `vis_utils.py`:

```python
# Visualize marginalized covariate effects
from vis_utils import plot_marginalized_effect
plot_marginalized_effect(model, covariate="age", spatial_location=0.5)

# Visualize heteroscedastic uncertainty
from vis_utils import plot_uncertainty_distribution
plot_uncertainty_distribution(model, covariates={"age": 10, "weight": 30}, spatial_range=(0, 1))
```

Output visualizations are saved in the experiment directory specified in the configuration file.

## 🔧 Advanced Usage

### Cross-validation Experiments

```bash
# Run 5-fold cross-validation on airway dataset
bash commands/run_airway_cv5.sh

# Run with covariate imputation and monotonicity constraints
bash commands/run_airway_cv5_ind.sh

# Run on brain volume dataset
bash commands/run_brain_cv5.sh
```

### Batch Processing

Process multiple experiments efficiently using our command scripts:

```bash
# Run comprehensive airway analysis pipeline
bash commands/run_airway.sh

# Run multiple configurations with varying hyperparameters
bash commands/run_test_configs.sh
```


## 📄 License

This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
