Neural coherent states

Implementation of neural coherent states, a type of neural-network quantum states introduced in the following preprint:

Artificial neural network states for non-additive systems
Wojciech Rzadkowski, Mikhail Lemeshko, Johan H. Mentink
arXiv:2105.15193

Dependencies

This code needs Jax and Flax. Python 3.9 is recommended.

Using the code

Running python main.py will perform learning procedure for a small system with two bosonic modes. Energies at each optimization step will be written to output.txt. For simplicity, adjusting both the physics and algorithm parameters is done directly in the main file by editing physics_pars and arg_pars variables.

The code runs on GPU without change. Consult this material for running on TPU.

The tests can be run with python tests.py. No errors indicate tests passing, while AssertionErrors appearing correspond to their failure.