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
This code needs Jax and Flax. Python 3.9 is recommended.
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.