# Discovering Lie Groups with Flow Matching

This repo implements the code for the paper "Discovering Lie Groups with Flow Matching".

## Requirements

### Python dependencies

* Python 3.11+ (tested on 3.11)
* Install python dependencies (See `pyproject.toml`)
```sh
pip install -e .
```

This repo uses [hydra](https://hydra.cc/) for configuration management and [WandB](https://wandb.ai/) for experiment tracking.

## Experiments
The `conf` directory contains all configs and hyperparameters.

To train, run a command of the following form:

```bash
python experiments/${EXP_FILE} dataset=${DATASET} model=${MODEL} seed=${SEED}
```
Possible values for `$EXP_FILE`, `$DATASET`, and `$MODEL` are shown in the table.
|   Experiment   |             EXP_FILE            |    DATASET    |                            MODEL                           |
|:--------------:|:-------------------------------:|:-------------:|:----------------------------------------------------------:|
|       2D       | flow_matching_2d.py             | C4_arrow      | SO2_to_C4_arrow                                            |
|                |                                 | C4_arrow      | GL2_to_C4_arrow                                            |
|                |                                 | D4_half_arrow | GL2_to_D4_half_arrow                                       |
|       3D       | flow_matching_3d.py             | Tet_irreg_tet | SO3_to_Tet_irreg_tet, SO3_to_Tet_irreg_tet_time_power_dist |
|                |                                 | Oct_irreg_tet | SO3_to_Oct_irreg_tet, SO3_to_Oct_irreg_tet_time_power_dist |
|                |                                 | Ico_irreg_tet | SO3_to_Ico_irreg_tet, SO3_to_Ico_irreg_tet_time_power_dist |
| Group Elements | flow_matching_angle_to_angle.py | C4_angle      | SO2_to_C4_angle_to_angle                                   |

Please replace the values of `$EXP_FILE`, `$DATASET`, and `$MODEL` for the desired experiment. See the `conf` directory for more details.

Example
```bash
python experiments/flow_matching_2d.py dataset=C4_arrow model=GL2_to_C4_arrow 
```

The outputs and the model checkpoints are stored in the `outputs` directory.
