# NODE Repository

## Repository Structure

```
node/
├── main.py                    # Main Hydra entry point (Training models)
├── config/
│   ├── hydra_local.yaml       # Local training config
│   ├── hydra_slurm.yaml       # SLURM cluster config
│   ├── dataset/               # Dataset configurations
│   └── hydra/launcher/        # Launcher configurations
├── utils/
│   ├── args.py               # Config processing
│   └── model_factory.py      # Model creation
├── model.py                  # Model definitions
├── data.py                   # Data loading
├── training/
│   └── train.py             # Training and Validation loops
├── metrics.py               # Evaluation metrics
├── plot.py                  # Visualization functions
└── requirements.txt         # Dependencies
```

## Installation

### 1. Create Conda Environment
```bash
# Create a new conda environment
conda create -n node python=3.9
conda activate node

# Clone the repository
git clone <repo-url>
cd node

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

### 2. Required Dependencies
The main dependencies include:
- PyTorch
- Hydra-core
- TorchDyn (for Neural ODEs)
- NumPy, Matplotlib, Seaborn
- Pandas, Scikit-learn
- Weights & Biases (optional, for logging)

## Configuration with Hydra

This repository uses [Hydra](https://hydra.cc/) for configuration management. All experiment parameters are defined in YAML configuration files located in the `config/` directory. SLURM cluster runs are also supported using the [Hydra Submitit](https://hydra.cc/docs/plugins/submitit_launcher/) Launcher.

### Configuration Structure
```
config/
├── hydra_local.yaml           # Local training configuration
├── hydra_slurm.yaml          # SLURM cluster configuration  
├── dataset/
│   ├── wang_100T.yaml        # Simulated Wang dataset
│   └── co9_*.yaml           # Real neural datasets
└── hydra/
    └── launcher/
        ├── slurm_cpu.yaml    # CPU cluster launcher
        └── slurm_gpu.yaml    # GPU cluster launcher (Untested)
```

## Basic Usage

### Local Training
To train a model, run the `main.py` script. All training parameters (including
model hyperparameters) are defined in the hydra configs.

```bash
# Train with default configuration
python main.py --config-name hydra_local
```

Variables in the hydra configs can also be overridden in the command line:
```bash
# Train with different solver and/or dynamics model
python main.py --config-name hydra_local \
    solver=continuous \
    dynamics_model=minimlp

# Train with different dataset
python main.py --config-name hydra_local \
    dataset=co9_12122023_100ms_0offset_complex

# Submit to SLURM cluster with CPU
python main.py --config-name hydra_slurm --multirun \
    solver=continuous \
    dynamics_model=minimlp \
    epochs=3000

# Submit with parameter sweeps
python main.py --config-name hydra_slurm --multirun \
    solver=discrete,continuous \
    dynamics_model=linear,mlp,minimlp \
    readout=separate,monotonic


```bash
# Validate existing model with plotting
python main.py --config-name <config_filename> \
    validate_only=true \
    checkpoint_path=<path_to_trained_model/*.pth> \
```

## Configuration Parameters

### Core Architecture
```yaml
solver: 'discrete'              # discrete | continuous
dynamics_model: 'minimlp'      # rnn | gru | linear | mlp | minimlp  
readout: 'separate'             # linear | monotonic | separate | separate_nonlinear
```

### Model Architecture
```yaml
dynamics_hidden_dim: [128]      # Hidden dimensions for MLP dynamics
dynamics_nonlinearity: 'tanh'   # tanh | relu | leaky_relu
readout_hidden_dim: 50          # Hidden dim for nonlinear readouts
```

## Architecture Options

### Solvers
- `discrete`: Classic RNN-style update loop. Used for baselines (`rnn`, `gru`) or discrete-time ODE dynamics
- `continuous`: Neural ODE solver using `torchdyn`. Integrates dynamics over continuous time

### Dynamics Models 
- `rnn`, `gru`: Standard recurrent cells for baseline models. Must be used with `discrete` solver
- `linear`: Simple linear dynamics: \(f(h) = Wh + b\)
- `mlp`: Multi-layer perceptron with configurable hidden dimensions
- `minimlp`: Interpretable, population-structured dynamics with separate pathways per neuron type

### Readout Types
- `linear`: Standard linear readout layer
- `monotonic`: Enforces positive weights via SoftPlus activation  
- `separate`: Independent monotonic readout per neuron population
- `separate_nonlinear`: Two-layer nonlinear readout per population

## Output Structure

Training outputs are automatically organized by Hydra:

```
results/
└── {dataset}/
    └── S-{solver}/
        └── D-{dynamics_model}/
            └── R-{readout}/
                └── {experiment_name}/
                    ├── checkpoints/
                    │   ├── {name}_best.pth      # Best model
                    │   └── {name}_final.pth     # Final model
                    ├── config/
                    │   └── config.yaml          # Full experiment config
                    ├── logs/
                    │   └── {name}.json          # Training logs
                    ├── metrics/
                    │   ├── training_metrics.json
                    │   └── training_curves.png
                    ├── plots/
                    │   ├── prediction_plots/
                    │   ├── lograte_plots/
                    │   └── flowfield_plots/
                    └── heatmaps/               # If heatmap=true
                        ├── best_heatmap.png
                        └── final_heatmap.png
```

## Monitoring & Logging

### Weights & Biases Integration
```yaml
use_wandb: true
wandb:
  project: 'EInode'
  entity: 'your-username'
  mode: 'online'  # online | offline | disabled
```

### Local Logging
All experiments automatically save:
- Training/validation curves
- Model checkpoints (best and final)
- Detailed metrics in JSON format
- Generated plots (if `plot=true`)


### Getting Help
- Check Hydra documentation: https://hydra.cc/
- Review example configurations in `config/` directory
- Examine training logs in experiment output directory

