Efficient optimization of deep neural quantum states toward machine precision: data and code

Efficient optimization of deep neural quantum states toward machine precision

Paper: arXiv:2302.01941

Data

The data in the paper is included in the `data` folder as `csv` files. The wave functions in the symmetry basis of 6x6 lattice are stored as `npy` files in the `wave_function` folder.

Variational wave functions

The parameters for 64-layer neural quantum states are stored in the `params` file. The NQS can be defined as shown in `energy.py`. In the case of an L x L lattice and Ns Monte Carlo samples, the input spin configurations \(\sigma\) should have a shape of (Ns, L x L)$ with element values +1 for spin-up and -1 for spin-down. The network output, given by `state.evaluate(spins)`, is a vector with Ns elements, each representing a wave function amplitude \(\psi_\sigma\).

Code

The main code for the neural quantum states in this paper is included in the `minsr` folder. The code requires python>=3.8, 0.3.9<=jax<=0.3.25 and 0.4.1<=flax<=0.6.5.

The `energy.py` file is a short explanation of how to use the codes and also a program for evaluating variational energies of the 64-layer neural quantum states. Run `python3 energy.py --help` for all available configurations. By default, it evaluates the variational energy on 6x6 Heisenberg model. After a thermalization process, the variational energy and its uncertainty will be printed after every 1000 samples. The running time is roughly a day for the largest 16x16 lattice on an RTX3090 GPU.