# Physics-Preserving Plasma Compression

## Running
Runs are configured with hydra configs for both autoencoders and neural fields.
Paths to the data need to be replaced wit the actual ones, after preprocessing.

For autoencoders, for example the VQ-VAE with 77k times compression is started as:
```
python -u main.py experiment=experiments=v5/v5_vqvae_cr40960_cmse.yaml
```

While for neural fields, a sweep over trajectories is started as
```
python -u neural_fields/nf_run.py config=configs/nf/eval.yaml
```


## Project structure
```
📁 data_generation                 # Info on data generation with GKW

📁 dataset                         # Dataset utilities and preprocessing
├── 📄 cyclone_diff.py             # Gyrokinetics dataset (AE)
├── 📄 cyclone.py                  # Gyrokinetics dataset
└── 📄 preprocess.py               # Preprocessing utilities

📁 eval                            # Evaluation and plotting scripts
├── 📄 complex_metrics.py          # Compute complex evaluation metrics
├── 📄 evaluate.py                 # Evaluation runner script
├── 📄 plot_utils.py               # Plotting helper functions
└── 📄 rollout.py                  # Rollout evaluation script

📁 models                          # Model architectures
├── 📁 nd_vit                      # Nested Vision Transformer models
│   ├── 📄 patching.py             # Patching layers
│   ├── 📄 positional.py           # Positional encoding
│   ├── 📄 swin_layers.py          # Swin layers
│   └── 📄 vit_layers.py           # ViT layers
├── 📄 swin_unet.py                # Swin UNet model
├── 📄 swin_ae.py                  # Swin autoencoder
├── 📄 utils.py                    # Model utility functions
└── 📄 vector_quantize_pytorch.py  # Vector quantization

📁 train                           # Training utilities
├── 📄 integrals.py                # Gyrokinetics integrals
├── 📄 losses.py                   # Physical Loss functions
├── 📄 peft_utils.py               # PEFT utility functions
├── 📄 train.py                    # Autoencoder training function
└── 📄 utils.py                    # Training helper functions

📁 neural_fields                   # Neural fields
├── 📁 models                      # NF model implementations
│   ├── 📄 mlp.py                  # MLP model
│   ├── 📄 siren.py                # SIREN model
│   ├── 📄 utils.py                # NF model utilities
│   └── 📄 wire.py                 # Wire model
├── 📄 data.py                     # Data utilities
├── 📄 gk_losses.py                # GK loss implementations
├── 📄 nf_train.py                 # Neural field training functions
├── 📄 nf_utils.py                 # NF helper utilities and plotting
├── 📄 trad.py                     # Traditional compression methods
└── 📄 nf_run.py                   # Neural field training entry point

📁 notebooks                       # Jupyter notebooks for plotting

📄 utils.py                        # General utilities
📄 run.py                          # Run experiments script
📄 main.py                         # Autoencoder training entry point
```

## Config structure
```
📁 configs
├── 📁 autoencoder                  # Autoencoder model configs
├── 📁 dataset                      # Dataset-related configuration files
├── 📁 experiments                  # Experiment version configurations
│   ├── 📁 _compression             # Compression experiment configs
│   ├── 📁 v1                       # Experiment version 1 configs
│   ├── 📁 v2                       # Experiment version 2 configs
│   ├── 📁 v3                       # Experiment version 3 configs
│   └── 📁 v5                       # Experiment version 5 configs
├── 📁 nf                           # Neural field configs
├── 📁 training                     # Training setup configs
└── 📁 validation                   # Validation setup configs
```