Schrodinet

使用径向基函数神经网络进行一维问题的量子蒙特卡洛模拟。

安装

克隆仓库并使用pip安装代码

git clone https://github.com/NLESC-JCER/Schrodinet/
cd Schrodinet
pip install .

一维简谐振子

下面的脚本说明了如何优化一维简谐振子的波函数。

import torch
from torch import optim

from schrodinet.sampler.metropolis import Metropolis
from schrodinet.wavefunction.wf_potential import Potential
from schrodinet.solver.solver_potential import SolverPotential
from schrodinet.solver.plot_potential import plot_results_1d, plotter1d

def pot_func(pos):
    '''Potential function desired.'''
    return 0.5*pos**2


def ho1d_sol(pos):
    '''Analytical solution of the 1D harmonic oscillator.'''
    return torch.exp(-0.5*pos**2)

# Define the domain and the number of RBFs

# wavefunction
wf = Potential(pot_func, domain, ncenter, fcinit='random', nelec=1, sigma=0.5)

# sampler
sampler = Metropolis(nwalkers=1000, nstep=2000,
                     step_size=1., nelec=wf.nelec,
                     ndim=wf.ndim, init={'min': -5, 'max': 5})

# optimizer
opt = optim.Adam(wf.parameters(), lr=0.05)
scheduler = optim.lr_scheduler.StepLR(opt, step_size=100, gamma=0.75)

# Solver
solver = SolverPotential(wf=wf, sampler=sampler,
                         optimizer=opt, scheduler=scheduler)

# Train the wave function
plotter = plotter1d(wf, domain, 100, sol=ho1d_sol)
solver.run(300, loss='variance', plot=plotter, save='model.pth')

# Plot the final wave function
plot_results_1d(solver, domain, 100, ho1d_sol, e0=0.5, load='model.pth')

优化后，可以轻松生成以下轨迹

可以在不同的势能上执行相同的程序。下面的图显示了该方法在简谐振子和 Morse 势能上的性能。