# Neural Equivariant Eikonal equation solver 

## Requirements

To install the requirements, we use conda. We recommend creating a new environment for the project.
```
conda create -n eikonal-enf-jax python=3.10
conda activate eikonal-enf-jax
```

Install the relevant dependencies.
```
pip install -r requirements.txt
```

## Experiments
To reproduce the experiments you need to run the following commands precomposed by:
```
export XLA_FLAGS="--xla_gpu_enable_command_buffer=" && export XLA_PYTHON_CLIENT_MEM_FRACTION=0.8 && export PYTHONPATH="."
```

### OpenFWI 2D fitting autodecoding:
Modify on the config file ```open_fwi_2d.yaml``` the data set that you would want to train and then run
```
python experiments/fitting/ad_fit.py --space euclidean --experiment open_fwi_2d
```

### OpenFWI 2D fitting meta:
Modify on the config file ```open_fwi_2d_meta.yaml``` the data set that you would want to train and then run
```
python experiments/fitting/meta_ad_fit.py --space euclidean --experiment open_fwi_2d
```

### OpenFWI 2D fitting pretrained meta:
Modify on the config file ```open_fwi_2d_mix.yaml``` the data set that you would want to train, add the wandb identifyer of the autodecoding run that you want to use as initialization, and then run
```
python experiments/fitting/meta_ad_fit.py --space euclidean --experiment open_fwi_2d
```

### OpenFWI 2D ablation of autodecoding steps:
Add in the python script the wandb ids of the runs that you want to set as the ablation, and then run the script
```
python experiments/fitting/ablation_auto_steps.py
```

### OpenFWI 2D ablation equivariance:
Modify on the config file ```open_fwi_2d.yaml``` the data set that you would want to train, and set
```
group: Sym-NoEquiv
dim_orientation: 0
```
then run 
```
python experiments/fitting/ad_fit.py --space euclidean --experiment open_fwi_2d
```

### OpenFWI 3D fitting autodecoding:
Modify on the config file ```open_fwi_3d.yaml``` the data set that you would want to train (Style-b in the paper) and then run
```
python experiments/fitting/ad_fit.py --space euclidean --experiment open_fwi_3d
```

### OpenFWI 3D fitting pretrained meta:
Modify on the config file ```open_fwi_2d_mix.yaml``` the data set that you would want to train, add the wandb identifyer of the autodecoding run that you want to use as initialization, and then run
```
python experiments/fitting/mix_ad_fit.py --space euclidean --experiment open_fwi_3d
```

### OpenFWI 3D ablation:
Add in the python script the wandb ids of the runs that you want to set as the ablation, and then run the script
```
python experiments/fitting/ablation_3d.py
```
