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)