# NeuralMPM

![Teaser video](assets/teaser_vid.mov) 

> [!CAUTION]
> This is an early code version that was used for research purposes. It is 
> not yet ready yet for public use. We provide instructions to run some 
> experiments but expect dirty and messy code! A final version will be 
> released over the coming months.

### Installation

1. Create a clean environment with Python 3.9.
2. Install PyTorch: `pip install torch==2.2.2 torchvision==0.17.2 
   torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cpu`. 
   Otherwise, you can install the GPU version by replacing `cpu` by `cu118` 
   or by visiting https://pytorch.org 
3. Install PyTorch Geometric (see [here](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html)): `pip install torch_geometric torch_scatter torch_cluster -f https://data.pyg.org/whl/torch-2.2.0+cpu.html`
4. Install the required packages with `pip install -r requirements.txt`.
5. Install the package with `pip install -e .`.

### Downloading data

Here are the instructions for downloading the `WaterRamps` dataset:
1. Download the compressed dataset [here](https://mega.nz/file/3ewWVaqD#o5pykhYgCXEbQyvWxNnRKEYFWnD-FKjV7yKBKIvdPPw).
2. Extract it in a new `data` folder.
3. Your folder structure should look like this:
```
data
└── WaterRamps
    ├── train
    │   ├── sim_0.h5
    │   ├── sim_1.h5
    │   └── ...
    ├── test
    │   ├── sim_0.h5
    │   ├── sim_1.h5
    │   └── ...
    └── valid
        ├── sim_0.h5
        ├── sim_1.h5
        └── ...
```

### Running experiments

The `train.py` script is deprecated to run directly. Please use the 
`experiments` package. Experiments are defined in `experiments/experiments.
py`, where you can find our various pre-defined experiments, which will be 
cleaned for the final code release along with documentation and comments 
for associating the experiments with the paper. To run, you can call 
`experiments/run.py` or `nmpm-exp` from the command line.

Example:
```bash
nmpm-exp --experiment debug_exp --local
```

This will start a very small experiment that can run on CPU to debug the
installation. Feel free to try other experiments, preferably on GPU.

Note: Running without --local requires a Slurm cluster (together with `dawgz`).

### Render a pre-trained model

To render a pre-trained model, you can use the `nmpm-render` command.

We provide one of our pre-trained WaterRamps models so you can check it. 
The model and config are located in `pretrained/waterramps_0`.

Example rendering of some frames:
```bash
nmpm-render frames --run pretrained/waterramps_0 \
                   --ids 20-22,50 \
                   --gt true \
                   --steps 0,100,200,300
```
`ids` corresponds to the test simulations to render. `gt` toggles rendering 
ground-truth frames. `steps` specifies the frames to render.

This should create 4 folders in `renders`, notably you should find these 
exact images in `renders/waterramps_0/test_21/gt_100.pdf` and 
`renders/waterramps_0/test_21/pred_100.pdf`:

![GT](assets/gt_100.png) ![Pred](assets/pred_100.png)

### Other commands

- `python neural_mpm/valid.py`: compute MSE and AE on a pre-trained model.
- `python neural_mpm/metrics.py`: compute EMD for a pre-trained model.
- `python neural_mpm/rollout.py`: export .h5 files for a pre-trained model.

### Note

The code for running the multimaterial experiments is currently located on 
another branch and will be merged in the final version.