def setUp(self):
# optimal parameters
self.opt_r = 0.69 # the two h are at +0.69 and -0.69
self.opt_sigma = 1.24
# molecule
self.mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# wave function
self.wf = Orbital(self.mol,
kinetic='auto',
configs='single(2,2)',
use_jastrow=True)
# sampler
self.sampler = Metropolis(nwalkers=1000,
nstep=2000,
step_size=0.5,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'),
move={
'type': 'all-elec',
'proba': 'normal'
})
self.hmc_sampler = Hamiltonian(nwalkers=100,
nstep=200,
step_size=0.1,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'))
# optimizer
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
# solver
self.solver = SolverOrbital(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# ground state energy
self.ground_state_energy = -1.16
# ground state pos
self.ground_state_pos = 0.69
def __init__(self, wf=None, sampler=None, optimizer=None):
SolverOrbital.__init__(self, wf, sampler, optimizer)
# task
self.configure(task='geo_opt')
#esampling
self.resampling(ntherm=-1,
resample=100,
resample_from_last=True,
resample_every=1)
# observalbe
self.observable(['local_energy'])
# distributed model
self.conf_dist(master_address='127.0.0.1',
master_port='29500',
backend='gloo')
self.save_model = 'model.pth'
class TestH2(unittest.TestCase):
def setUp(self):
# optimal parameters
self.opt_r = 0.69 # the two h are at +0.69 and -0.69
self.opt_sigma = 1.24
# molecule
self.mol = Molecule(atom='H 0 0 -0.69; H 0 0 0.69',
unit='bohr',
calculator='pyscf',
basis='sto-3g')
# wave function
self.wf = Orbital(self.mol,
kinetic='auto',
configs='single(2,2)',
use_jastrow=True)
# sampler
self.sampler = Metropolis(nwalkers=1000,
nstep=2000,
step_size=0.5,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'),
move={
'type': 'all-elec',
'proba': 'normal'
})
self.hmc_sampler = Hamiltonian(nwalkers=100,
nstep=200,
step_size=0.1,
ndim=self.wf.ndim,
nelec=self.wf.nelec,
init=self.mol.domain('normal'))
# optimizer
self.opt = optim.Adam(self.wf.parameters(), lr=0.01)
# solver
self.solver = SolverOrbital(wf=self.wf,
sampler=self.sampler,
optimizer=self.opt)
# ground state energy
self.ground_state_energy = -1.16
# ground state pos
self.ground_state_pos = 0.69
def test_single_point(self):
self.solver.wf.ao.atom_coords[0, 2] = -self.ground_state_pos
self.solver.wf.ao.atom_coords[1, 2] = self.ground_state_pos
self.solver.sampler = self.sampler
# sample and compute observables
_, e, v = self.solver.single_point()
print('Energy :', e)
print('Variance :', v)
# assert(e>self.ground_state_energy and e<-1.)
assert (e > 2 * self.ground_state_energy and e < 0.)
assert (v > 0 and v < 5.)
def test_single_point_hmc(self):
self.solver.wf.ao.atom_coords[0, 2] = -self.ground_state_pos
self.solver.wf.ao.atom_coords[1, 2] = self.ground_state_pos
self.solver.sampler = self.hmc_sampler
# sample and compute observables
_, e, v = self.solver.single_point()
print('Energy :', e)
print('Variance :', v)
# assert(e>self.ground_state_energy and e<-1.)
assert (e > 2 * self.ground_state_energy and e < 0.)
assert (v > 0 and v < 5.)
def test_geo_opt(self):
self.solver.wf.ao.atom_coords[0, 2].data = torch.tensor(-0.37)
self.solver.wf.ao.atom_coords[1, 2].data = torch.tensor(0.37)
self.solver.configure(task='geo_opt')
self.solver.observable(['local_energy', 'atomic_distances'])
self.solver.run(50, loss='energy')
# load the best model
best_model = torch.load('model.pth')
self.solver.wf.load_state_dict(best_model['model_state_dict'])
self.solver.wf.eval()
# sample and compute variables
_, e, v = self.solver.single_point()
e = e.data.numpy()
v = v.data.numpy()
# it might be too much to assert with the ground state energy
assert (e > 2 * self.ground_state_energy and e < 0.)
assert (v > 0 and v < 2.)
'params': wf.mo.parameters(),
'lr': 1E-3
}, {
'params': wf.fc.parameters(),
'lr': 1E-3
}]
opt = Adam(lr_dict, lr=1E-3)
# opt = SGD(lr_dict, lr=1E-1, momentum=0.9)
# opt = StochasticReconfiguration(wf.parameters(), wf)
# scheduler
scheduler = optim.lr_scheduler.StepLR(opt, step_size=50, gamma=0.85)
# solver
solver = SolverOrbital(wf=wf, sampler=sampler, optimizer=opt, scheduler=None)
if 1:
pos, e, v = solver.single_point(ntherm=-1, ndecor=100)
# eloc = solver.wf.local_energy(pos)
# plt.hist(eloc.detach().numpy(), bins=50)
# plt.show()
# pos = solver.sample(ntherm=0, ndecor=10)
# obs = solver.sampling_traj(pos)
# plot_energy(obs, e0=-8.)
if 0:
solver.configure(task='wf_opt', freeze=['ao', 'mo'])
solver.observable(['local_energy'])
# define the wave function
wf = Orbital(mol)
#sampler
sampler = Metropolis(nwalkers=100,
nstep=1000,
step_size=0.5,
ndim=wf.ndim,
nelec=wf.nelec,
move='one')
# optimizer
opt = Adam(wf.parameters(), lr=0.01)
# solver
solver = SolverOrbital(wf=wf, sampler=sampler, optimizer=opt)
pos = Variable(torch.rand(100, mol.nelec * 3))
pos.requires_grad = True
solver.single_point(pos=pos)
# plot the molecule
#plot_molecule(solver)
# optimize the geometry
# solver.configure(task='geo_opt')
# solver.observable(['local_energy','atomic_distances'])
# solver.run(5,loss='energy')
# plot the data
# plot_observable(solver.obs_dict,e0=-1.16)
basis_type='gto',
basis='sto-3g')
# define the wave function
wf = Orbital(mol, kinetic_jacobi=True)
# sampler
sampler = Metropolis(nwalkers=1000,
nstep=5000,
step_size=0.5,
ndim=wf.ndim,
nelec=wf.nelec,
move='one')
# solver
solver = SolverOrbital(wf=wf, sampler=sampler)
pos, e, v = solver.single_point()
sampler.nstep = 500
angles = np.linspace(0, 90, 10)
R = rot_mat(angles)
for iA in range(len(angles)):
# define the wave function
wf = Orbital(mol, kinetic_jacobi=True)
solver.wf = wf
pos, e, v = solver.single_point(pos=pos)
# bend the mol
#mol = Molecule(atom='H 0 0 -0.37; H 0 0 0.37', basis_type='gto', basis='sto-3g')
# define the wave function
wf = OrbitalH2(mol)
#sampler
sampler = Metropolis(nwalkers=1000,
nstep=1000,
step_size=0.5,
ndim=wf.ndim,
nelec=wf.nelec,
move='one')
# optimizer
opt = Adam(wf.parameters(), lr=0.01)
# solver
solver = SolverOrbital(wf=wf, sampler=sampler, optimizer=opt)
#solver.single_point()
# plot the molecule
#plot_molecule(solver)
# optimize the geometry
solver.configure(task='geo_opt')
solver.observable(['local_energy', 'atomic_distances'])
solver.run(5, loss='energy')
# plot the data
plot_observable(solver.obs_dict, e0=-1.16)
ndim=wf.ndim,
init=mol.domain('normal'),
move={
'type': 'one-elec',
'proba': 'normal'
})
# optimizer
opt = Adam(wf.parameters(), lr=0.005)
# scheduler
scheduler = optim.lr_scheduler.StepLR(opt, step_size=20, gamma=0.75)
# solver
solver = SolverOrbital(wf=wf,
sampler=sampler,
optimizer=opt,
scheduler=scheduler)
# # single point
#pos, e, v = solver.single_point(ntherm=500, ndecor=100)
# # sampling traj
# pos = solver.sample(ntherm=0, ndecor=10)
# obs = solver.sampling_traj(pos)
# plot_energy(obs, e0=-111.)
# optimize the wave function
# solver.configure(task='wf_opt', freeze=['bas_exp'])
# solver.observable(['local_energy'])
# solver.run(10, loss='energy')
# plot_energy(solver.obs_dict, e0=-113.)
sampler = Metropolis(nwalkers=100,
nstep=200,
step_size=0.5,
ndim=wf.ndim,
nelec=wf.nelec,
init=mol.domain('normal'),
move={
'type': 'all-elec',
'proba': 'normal'
})
# optimizer
opt = Adam(wf.parameters(), lr=0.01)
# solver
solver = SolverOrbital(wf=wf, sampler=sampler, optimizer=None)
pos, _, _ = solver.single_point()
# pos = solver.sample(ntherm=0, ndecor=10)
# obs = solver.sampling_traj(pos)
# plot_energy(obs, e0=-1.16)
# optimize the wave function
solver.configure(task='wf_opt', freeze=['mo', 'bas_exp'])
solver.observable(['local_energy'])
solver.run(5, loss='energy')
# # optimize the geometry
# solver.configure(task='geo_opt')
# solver.observable(['local_energy','atomic_distances'])
# solver.run(5,loss='energy')