Adding WQMC objective option

[necludov]

Hi everyone,

@dpfau asked me to create a PR implementing our recent work Wasserstein Quantum Monte Carlo. I just finished the initial implementation and tested it for Li atom.

Results for Li atom. The left plot demonstrates the relative energy error, where the ground true energy is taken from the psiformer paper. The right plot demonstrates the variance of the local energy, which indicates the convergence to the ground energy eigenstate. All the values are smoothed with the running average over 100 iterations.

The commands to run the test. ferminet --config ferminet/configs/atom_test.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100 --config.optim.objective wvmc ferminet --config ferminet/configs/atom_test.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100 --config.optim.objective vmc

Nota Bene:

• I had to change the evaluation of the potential to avoid 'nans' during backpropagation. It should be equivalent, but please review it.
• I guess I didn't completely follow your style guide, even though I was trying to stick to it where I could infer it.
• I haven't yet tested this particular implementation for bigger molecules (the ones we considered in the paper). Moreover, for some hyperparameters, I observed this implementation diverging for Li, perhaps, due to the smaller batch size compared to what we used in the paper.

Please, let me know how to proceed from here, and what to improve, test, and include.

[google-cla[bot]]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[necludov]

I'm glad you enjoyed the paper! It means a lot coming from you! Let me know what else to check and fix.

[necludov]

I implemented the comments and tested them. The command to run now is ferminet --config ferminet/configs/li_wqmc.py --config.batch_size 256