DeepSolid
An implementation of the algorithm given in "Ab initio calculation of real solids via neural network ansatz". A periodic neural network is proposed as wavefunction ansatz for solid quantum Monte Carlo and achieves unprecedented accuracy compared with other state-of-the-art methods. This repository is developed upon FermiNet and PyQMC.
Installation
DeepSolid can be installed via the supplied setup.py file.
pip3 install -e.If GPU is available, we recommend you to install jax and jaxlib with cuda 11.4+. Our experiments were carried out with jax==0.2.26 and jaxlib==0.1.75.
Usage
Ml_collection package is used for system definition. Below is a simple example of H10 in PBC:
deepsolid --config=PATH/TO/DeepSolid/config/two_hydrogen_cell.py:H,5,1,1,2.0,0,ccpvdz --config.batch_size 4096
Customize your system
Simulation system can be customized in config.py file, such as
import numpy as np
from pyscf.pbc import gto
from DeepSolid import base_config
from DeepSolid import supercell
def get_config(input_str):
symbol, S = input_str.split(',')
cfg = base_config.default()
# Set up cell.
cell = gto.Cell()
# Define the atoms in the primitive cell.
cell.atom = f"""
{symbol} 0.000000000000 0.000000000000 0.000000000000
"""
# Define the pretrain basis.
cell.basis = "ccpvdz"
# Define the lattice vectors of the primitive cell.
# In this example it's a simple cubic.
cell.a = np.array([[3.0, 0.0, 0.0],
[0.0, 3.0, 0.0],
[0.0, 0.0, 3.0]])
# Define the unit used in cell definition, only support Bohr now.
cell.unit = "B"
cell.verbose = 5
# Define the threshold to discard gaussian basis used in pretrain.
cell.exp_to_discard = 0.1
cell.build()
# Define the supercell for QMC, S specifies how to tile the primitive cell.
S = np.eye(3) * int(S)
simulation_cell = supercell.get_supercell(cell, S)
# Assign the defined supercell to cfg.
cfg.system.pyscf_cell = simulation_cell
return cfgAfter defining the config file, simply use the following command to launch the simulation:
deepsolid --config=PATH/TO/config.py:He,1 --config.batch_size 4096Read structure from poscar file
We also support reading structure from poscar file, which is commonly used. Simply use the following command
deepsolid --config=DeepSolid/config/read_poscar.py:PATH/TO/POSCAR/bcc_li.vasp,1,ccpvdzDistributed training
Present released code doesn't support multi-node training. See this link for help.
Tricks to accelerate
The bottleneck of DeepSolid is the laplacian evaluation of the neural network. We recommend the users to use partition mode instead, simply adding two more flags:
deepsolid --config=PATH/TO/config.py --config.optim.laplacian_mode=partition --config.optim.partition_number=3Partition mode will try to parallelize the calculation of laplacian and partition number must be a factor of (electron number * 3). Note that partition mode will require a lot of GPU memory.
Precision
DeepSolid supports both FP32 and FP64. However, we recommend the users turn off the TF32 mode which is automatically adopted in A100 if FP32 is chosen. TF32 can be turned off using the following command:
NVIDIA_TF32_OVERRIDE=0 deepsolid --config.use_x64=FalseGiving Credit
If you use this code in your work, please cite the associated paper.
@article{li2022ab,
title={Ab initio calculation of real solids via neural network ansatz},
author={Li, Xiang and Li, Zhe and Chen, Ji},
journal={Nature Communications},
volume={13},
number={1},
pages={7895},
year={2022},
publisher={Nature Publishing Group UK London}
}
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent