class TestRbfNetworkHarmonicOscillator1D(unittest.TestCase):
def setUp(self):
# wavefunction
self.wf = RBF_HO1D(ndim=1, nelec=1, ncenter=5)
#sampler
self.mh_sampler = METROPOLIS(nwalkers=250,
nstep=1000,
step_size=3.,
nelec=self.wf.nelec,
ndim=self.wf.ndim,
domain={
'min': -5,
'max': 5
})
#sampler
self.hmc_sampler = HAMILTONIAN(nwalkers=250,
nstep=1000,
step_size=0.01,
nelec=self.wf.nelec,
ndim=self.wf.ndim,
domain={
'min': -5,
'max': 5
},
L=5)
# optimizer
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
# network
self.net = DeepQMC(wf=self.wf,
sampler=self.mh_sampler,
optimizer=self.opt)
def test_single_point_metropolis_hasting_sampling(self):
# initiaize the fc layer
self.net.wf.fc.weight.data.fill_(0.)
self.net.wf.fc.weight.data[0, 2] = 1.
# sample and compute observables
pos = self.net.sample(ntherm=-1, with_tqdm=False)
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
assert np.allclose([e, v], [0.5, 0], atol=1E-6)
def test_single_point_hamiltonian_mc_sampling(self):
#switch to HMC sampling
self.net.sampler = self.hmc_sampler
# initiaize the fc layer
self.net.wf.fc.weight.data.fill_(0.)
self.net.wf.fc.weight.data[0, 2] = 1.
# sample and compute observables
pos = self.net.sample(ntherm=-1, with_tqdm=False)
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
assert np.allclose([e, v], [0.5, 0], atol=1E-6)
def test_optimization(self):
#switch to MH sampling
self.net.sampler = self.mh_sampler
# optimize the weight of the FC layer
# do not optimize the pos of the centers
self.net.wf.fc.weight.requires_grad = True
self.net.wf.rbf.centers.requires_grad = False
# randomize the weights
nn.init.uniform_(self.wf.fc.weight, 0, 1)
# train
pos, obs_dict = self.net.train(250,
batchsize=250,
pos=None,
obs_dict=None,
resample=100,
resample_from_last=True,
resample_every=1,
ntherm=-1,
loss='variance',
plot=None,
save_model='best_model.pth')
# load the best model
best_model = torch.load('best_model.pth')
self.net.wf.load_state_dict(best_model['model_state_dict'])
self.net.wf.eval()
# sample and compute variables
pos = self.net.sample()
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
assert np.allclose([e, v], [0.5, 0], atol=1E-3)
class TestHydrogen(unittest.TestCase):
def setUp(self):
# wavefunction
self.wf = RBF_H(centers=torch.tensor([[0., 0., 0.]]),
sigma=torch.tensor([1.]))
#sampler
self.mh_sampler = METROPOLIS(nwalkers=1000,
nstep=1000,
step_size=3.,
nelec=self.wf.nelec,
ndim=self.wf.ndim,
domain={
'min': -5,
'max': 5
})
#sampler
self.hmc_sampler = HAMILTONIAN(nwalkers=1000,
nstep=1000,
step_size=0.01,
nelec=self.wf.nelec,
ndim=self.wf.ndim,
domain={
'min': -5,
'max': 5
},
L=5)
# optimizer
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
# network
self.net = DeepQMC(wf=self.wf,
sampler=self.mh_sampler,
optimizer=self.opt)
# ground state energy
self.ground_state_energy = -0.5
def test_single_point_metropolis_hasting_sampling(self):
# sample and compute observables
pos = self.net.sample(ntherm=-1, with_tqdm=False)
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
assert np.allclose([e, v], [self.ground_state_energy, 0], atol=1E-6)
def test_single_point_hamiltonian_mc_sampling(self):
#switch to HMC sampling
self.net.sampler = self.hmc_sampler
# sample and compute observables
pos = self.net.sample(ntherm=-1, with_tqdm=False)
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
assert np.allclose([e, v], [self.ground_state_energy, 0], atol=1E-6)
def test_optimization(self):
#switch to MH sampling
self.net.sampler = self.mh_sampler
# optimize the weight of the FC layer
# do not optimize the pos of the centers
self.net.wf.fc.weight.requires_grad = False
self.net.wf.rbf.sigma.requires_grad = True
# modify the sigma
self.net.wf.rbf.sigma.data[0] = 1.5
# train
pos, obs_dict = self.net.train(250,
batchsize=250,
pos=None,
obs_dict=None,
resample=100,
resample_from_last=True,
resample_every=1,
ntherm=-1,
loss='variance',
plot=None,
save_model='best_model.pth')
# load the best model
best_model = torch.load('best_model.pth')
self.net.wf.load_state_dict(best_model['model_state_dict'])
self.net.wf.eval()
# sample and compute variables
pos = self.net.sample()
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
assert np.allclose([e, v], [self.ground_state_energy, 0], atol=1E-3)
domain = {
'xmin': -boundary,
'xmax': boundary,
'ymin': -boundary,
'ymax': boundary,
'zmin': -boundary,
'zmax': boundary
}
ncenter = [11, 11, 11]
# network
net = DeepQMC(wf=wf, sampler=sampler, optimizer=opt)
obs_dict = {'local_energy': [], 'atomic_distance': [], 'get_sigma': []}
if 1:
pos = Variable(net.sample())
pos.requires_grad = True
e = net.wf.energy(pos)
s = net.wf.variance(pos)
print('Energy :', e)
print('Variance :', s)
if 0:
net.wf.layer_mo.weight.data = torch.eye(net.wf.nao).double()
for param in net.wf.layer_ci.parameters():
param.requires_grad = False
pos = net.sample(ntherm=0)
pos = pos.reshape(100, 100, 6)
var_ = net.observalbe(net.wf.variance, pos)
domain,
ncenter,
sol=h2plus_sol,
hist=False,
pot=False,
grad=True)
if 0:
X = np.linspace(0.1, 2, 25)
energy, var = [], []
for x in X:
net.wf.rbf.centers.data[0] = -x
net.wf.rbf.centers.data[1] = x
pos = Variable(net.sample())
pos.requires_grad = True
e = net.wf.energy(pos)
s = net.wf.variance(pos)
energy.append(e)
var.append(s)
plt.plot(X, energy)
plt.show()
plt.plot(X, var)
plt.show()
if 0:
pos = Variable(net.sample())
pos.requires_grad = True
class TestH2plus(unittest.TestCase):
def setUp(self):
# optimal parameters
self.opt_r = 0.97 # the two h are at +0.97 and -0.97
self.opt_sigma = 1.20
# wavefunction
centers = torch.tensor([[0.,0.,-self.opt_r],[0.,0.,self.opt_r]])
sigma = torch.tensor([self.opt_sigma,self.opt_sigma])
self.wf = RBF_H2p(centers = centers,
sigma = sigma )
#sampler
self.mh_sampler = METROPOLIS(nwalkers=1000, nstep=1000,
step_size = 0.5, nelec = self.wf.nelec,
ndim = self.wf.ndim, domain = {'min':-5,'max':5})
#sampler
self.hmc_sampler = HAMILTONIAN(nwalkers=1000, nstep=200,
step_size = 0.1, nelec = self.wf.nelec,
ndim = self.wf.ndim, domain = {'min':-5,'max':5}, L=5)
# optimizer
self.opt = optim.Adam(self.wf.parameters(),lr=0.001)
# network
self.net = DeepQMC(wf=self.wf,sampler=self.mh_sampler,optimizer=self.opt)
# ground state energy
self.ground_state_energy = -0.597
def test_single_point_metropolis_hasting_sampling(self):
# sample and compute observables
pos = self.net.sample(ntherm=-1,with_tqdm=False)
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
print('Energy :', e)
print('Variance :', v)
assert np.allclose([e,v],[self.ground_state_energy,0],atol=1E-1)
def test_single_point_hamiltonian_mc_sampling(self):
#switch to HMC sampling
self.net.sampler = self.hmc_sampler
# sample and compute observables
pos = self.net.sample(ntherm=-1,with_tqdm=False)
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
print('Energy :', e)
print('Variance :', v)
assert np.allclose([e,v],[self.ground_state_energy,0],atol=1E-1)
def test_energy_curve(self):
#switch to MH sampling
self.net.sampler = self.mh_sampler
# fix the sigma of the AO
s = 1.20
self.net.wf.rbf.sigma.data[:] = s
X = np.linspace(0.1,2,25)
emin = 1E3
for x in X:
# move the atoms
self.net.wf.rbf.centers.data[0,2] = -x
self.net.wf.rbf.centers.data[1,2] = x
pos = Variable(self.net.sample())
pos.requires_grad = True
e = self.net.wf.energy(pos)
if e<emin:
emin = e
assert(emin<-0.55)
def test_sigma_optimization(self):
#switch to MH sampling
self.net.sampler = self.mh_sampler
# move the atoms
x = 0.97
self.net.wf.rbf.centers.data[0,2] = -x
self.net.wf.rbf.centers.data[1,2] = x
# fix the sigma of the AO
s = 1.
self.net.wf.rbf.sigma.data[:] = s
# do not optimize the weight of the FC layer
# optimize the pos of the centers
# do not optimize the sigma of the AO
self.net.wf.fc.weight.requires_grad = False
self.net.wf.rbf.centers.requires_grad = False
self.net.wf.rbf.sigma.requires_grad = True
# define the observable we want
obs_dict = {'local_energy':[],
'atomic_distance':[],
'get_sigma':[]}
# train
pos,obs_dict = self.net.train(100,
batchsize=1000,
pos = None,
obs_dict = obs_dict,
resample = 100,
resample_from_last = True,
resample_every = 1,
ntherm = -1,
loss = 'variance',
plot = None,
save_model = 'best_model.pth')
# load the best model
best_model = torch.load('best_model.pth')
self.net.wf.load_state_dict(best_model['model_state_dict'])
self.net.wf.eval()
# sample and compute variables
pos = self.net.sample()
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
# it might be too much to assert with the ground state energy
assert (e < -0.5)
assert (v < 0.1)
def test_geo_optimization(self):
#switch to MH sampling
self.net.sampler = self.mh_sampler
# move the atoms
x = 0.5
self.net.wf.rbf.centers.data[0,2] = -x
self.net.wf.rbf.centers.data[1,2] = x
# fix the sigma of the AO
s = 1.20
self.net.wf.rbf.sigma.data[:] = s
# do not optimize the weight of the FC layer
# optimize the pos of the centers
# do not optimize the sigma of the AO
self.net.wf.fc.weight.requires_grad = False
self.net.wf.rbf.centers.requires_grad = True
self.net.wf.rbf.sigma.requires_grad = False
# define the observable we want
obs_dict = {'local_energy':[],
'atomic_distance':[],
'get_sigma':[]}
# train
pos,obs_dict = self.net.train(200,
batchsize=1000,
pos = None,
obs_dict = obs_dict,
resample = 100,
resample_from_last = True,
resample_every = 1,
ntherm = -1,
loss = 'energy',
plot = None,
save_model = 'best_model.pth')
# load the best model
best_model = torch.load('best_model.pth')
self.net.wf.load_state_dict(best_model['model_state_dict'])
self.net.wf.eval()
# sample and compute variables
pos = self.net.sample()
e = self.wf.energy(pos).detach().numpy()
v = self.wf.variance(pos).detach().numpy()
# it might be too much to assert with the ground state energy
assert (e < -0.5)
assert (v < 0.1)