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.