def sigma_opt(net,sigma=1.5,loss='variance'):
'''Optimize the sigma of the rbf.'''
net.wf.rbf.sigma.data[:] = sigma
# do not optimize the weights of fc
net.wf.fc.weight.requires_grad = False
# # optimize the position of the centers
# do not optimize the std of the gaussian
net.wf.rbf.centers.requires_grad = False
net.wf.rbf.sigma.requires_grad = True
# variable
obs_dict = {'local_energy':[],
'atomic_distance':[],
'get_sigma':[]}
# train
pos,obs_dict = net.train(500,
batchsize=1000,
pos = None,
obs_dict = obs_dict,
resample=100,
resample_every=1,
resample_from_last = True,
ntherm=-1,
loss = loss)
boundary = 5.
domain = {'xmin':-boundary,'xmax':boundary,
'ymin':-boundary,'ymax':boundary,
'zmin':-boundary,'zmax':boundary}
ncenter = [11,11,11]
plot_results(net,obs_dict,domain,ncenter,isoval=0.02,hist=True)
def mo_opt(net, x=0.69, sigma=1.20, loss='variance'):
'''Optimize the MO coefficients.'''
# Change the mo coeffs
q, _ = np.linalg.qr(np.array([[0.8, 0.2], [0.2, -0.8]]))
mo_coeff = torch.tensor(q).float()
net.wf.mo.weight = nn.Parameter(mo_coeff.transpose(0, 1))
# set the other params
net.wf.rbf.centers.data[0, 2] = -x
net.wf.rbf.centers.data[1, 2] = x
net.wf.rbf.sigma.data[:] = s
# optimize the weights of mo
net.wf.mo.weight.requires_grad = True
# do not optimize the position of the centers
# do not optimize the std of the gaussian
net.wf.rbf.centers.requires_grad = False
net.wf.rbf.sigma.requires_grad = False
# track obs
obs_dict = {
'local_energy': [],
'atomic_distance': [],
'get_sigma': [],
'get_mos': []
}
# train
pos, obs_dict = net.train(150,
batchsize=500,
pos=None,
obs_dict=obs_dict,
resample=100,
resample_from_last=True,
resample_every=1,
ntherm=-1,
loss=loss)
# domain for the RBF Network
boundary = 5.
domain = {
'xmin': -boundary,
'xmax': boundary,
'ymin': -boundary,
'ymax': boundary,
'zmin': -boundary,
'zmax': boundary
}
ncenter = [11, 11, 11]
# plot the result
plot_results(net, obs_dict, domain, ncenter, hist=True)
def geo_opt(net, x=1.25, sigma=1.20, loss='energy'):
'''Optimize the geometry of the mol.'''
# fix the mo coefficients
mo_coeff = torch.sqrt(torch.tensor([1. / 2.])) * torch.tensor([[1., 1.],
[1., -1.]])
net.wf.mo.weight = nn.Parameter(mo_coeff.transpose(0, 1))
net.wf.mo.weight.requires_grad = False
# set the positions and sigmas
net.wf.rbf.centers.data[0, 2] = -x
net.wf.rbf.centers.data[1, 2] = x
net.wf.rbf.sigma.data[:] = sigma
# do not optimize the weights of fc
net.wf.fc.weight.requires_grad = False
# optimize the position of the centers
# do not optimize the std of the gaussian
net.wf.rbf.centers.requires_grad = True
net.wf.rbf.sigma.requires_grad = False
# variable
obs_dict = {'local_energy': [], 'atomic_distance': [], 'get_sigma': []}
# train
pos, obs_dict = net.train(100,
batchsize=1000,
pos=None,
obs_dict=obs_dict,
resample=250,
resample_every=1,
resample_from_last=True,
ntherm=-1,
loss=loss)
boundary = 5.
domain = {
'xmin': -boundary,
'xmax': boundary,
'ymin': -boundary,
'ymax': boundary,
'zmin': -boundary,
'zmax': boundary
}
ncenter = [11, 11, 11]
plot_results(net, obs_dict, domain, ncenter, isoval=0.02, hist=True)
step_size=3.,
nelec=wf.nelec,
ndim=wf.ndim,
domain={
'min': -5,
'max': 5
})
# optimizer
opt = optim.Adam(wf.parameters(), lr=0.01)
# network
net = DeepQMC(wf=wf, sampler=sampler, optimizer=opt)
pos, obs_dict = None, None
# set up the plotter
plt.ion()
plot2D = plotter2d(wf, domain, [25, 25], sol=ho2d_sol, kinetic=False)
# train the wavefunction
pos, obs_dict = net.train(100,
batchsize=250,
pos=pos,
obs_dict=obs_dict,
resample=100,
ntherm=-1,
loss='variance',
plot=plot2D)
plot_results(net, obs_dict, domain, [25, 25], ho2d_sol, e0=1.)
e = net.wf.energy(pos)
s = net.wf.variance(pos)
print('Energy :', e)
print('Variance :', s)
sys.exit()
if 1:
plot1D = plotter1d(wf, domain, ncenter, sol=h2plus_sol)
# do not optimize the weights of fc
net.wf.fc.weight.requires_grad = False
# optimize the position of the centers
# do not optimize the std of the gaussian
net.wf.rbf.sigma.requires_grad = False
net.wf.rbf.centers.requires_grad = True
# train
pos, obs_dict = net.train(250,
batchsize=500,
pos=None,
obs_dict=obs_dict,
resample=200,
ntherm=-1,
loss='energy',
plot=plot1D)
plot_results(net, obs_dict, domain, ncenter)
