Pytorch implementation of Conditional Optimal Transport on Function Spaces by Hosseini, Hsu and Taghvaei.
Most of this is code is for the Wasserstein Monge GAN, but we've also included the code for pCN that we used for validation, and for the low dimensional empirical pushforward experiments.
Since pytorch and fenicsx wouldn't place nicely with each other, we have separate conda environments for solving pde's and training the WaMGANS, they can both be found in the conda_envs directory.
We build the dataset using Fenicsx. Using the conda environment defined in conda_envs/fenics_env.yml, you can use
python build_dataset.py
after modifying configuration settings to your liking
See the file darcy_train.py for an example of the training process. We used wandb to log everything throughout training, and have provided the conda environment defined in conda_envs/mgan_env.yml for training and using the WaMGAN.
These generative models can take a long time to converge. We ran for 100,000 epochs, but much fewer will still produce satisfactory results for this Darcy flow example.
with significant modifications for our use case, and is licensed as such.