def example():
#from theforce.calculator.engine import Engine
from torch.nn import Module, Parameter
class LJ(Module):
def __init__(self, ):
super().__init__()
self.eps = Parameter(torch.tensor(1.0), requires_grad=False)
self.sigma = Parameter(torch.tensor(1.0), requires_grad=False)
def forward(self, a, aa, r):
x = (self.sigma**2 / (r**2).sum(dim=-1))
return 4 * self.eps * (x**6 - x**3).sum() / 2
def extra_repr(self):
print(self.eps)
print(self.sigma)
from theforce.descriptor.atoms import TorchAtoms as Atoms
from ase.optimize import BFGS
atoms = Atoms(symbols='ArAr',
positions=[[0, 0, 0], [2., 0, 0]],
cell=[9., 9., 9.],
pbc=True,
cutoff=3.0)
atoms.update()
calc = Engine(rc=3.0, terms=LJ())
atoms.set_calculator(calc)
BFGS(atoms).run(fmax=1e-5)
dmin = ((atoms.positions[0] - atoms.positions[1])**2).sum()
print('isclose: {} (diff = {})'.format(np.isclose(dmin, 2**(1. / 3)),
dmin - 2**(1. / 3)))
def calculate(self,
_atoms=None,
properties=['energy'],
system_changes=all_changes):
if self.potential.is_distributed:
raise NotImplementedError(
'(Auto)forces in distributed mode is not implemented')
if type(_atoms) == ase.atoms.Atoms:
atoms = TorchAtoms(ase_atoms=_atoms)
uargs = {
'cutoff': self.potential._cutoff,
'descriptors': self.potential.gp.kern.kernels
}
else:
atoms = _atoms
uargs = {}
if _atoms is not None and self.process_group is not None:
atoms.attach_process_group(self.process_group)
Calculator.calculate(self, atoms, properties, system_changes)
self.atoms.update(posgrad=True,
cellgrad=True,
forced=True,
dont_save_grads=True,
**uargs)
# energy
energy = self.potential([self.atoms],
'energy',
enable_grad=True,
variance=self.variance,
all_reduce=self.atoms.is_distributed)
if self.variance:
energy, variance = energy
# forces
rgrad = grad(energy,
self.atoms.xyz,
retain_graph=True,
allow_unused=True)[0]
forces = torch.zeros_like(self.atoms.xyz) if rgrad is None else -rgrad
if self.atoms.is_distributed:
torch.distributed.all_reduce(forces)
# stress
stress1 = -(forces[:, None] * self.atoms.xyz[..., None]).sum(dim=0)
cellgrad, = grad(energy, self.atoms.lll, allow_unused=True)
if cellgrad is None:
cellgrad = torch.zeros_like(self.atoms.lll)
if self.atoms.is_distributed:
torch.distributed.all_reduce(cellgrad)
stress2 = (cellgrad[:, None] * self.atoms.lll[..., None]).sum(dim=0)
try:
volume = self.atoms.get_volume()
except ValueError:
volume = -2 # here stress2=0, thus trace(stress) = virial (?)
stress = (stress1 + stress2).detach().numpy() / volume
# results
self.results['energy'] = energy.detach().numpy()[0]
self.results['forces'] = forces.detach().numpy()
self.results['free_energy'] = self.results['energy']
self.results['stress'] = stress.flat[[0, 4, 8, 5, 2, 1]]
if self.variance:
self.results['energy_variance'] = variance
def calculate(self, _atoms=None, properties=['energy'], system_changes=all_changes):
# ase to torch
if type(_atoms) == ase.atoms.Atoms:
atoms = TorchAtoms(ase_atoms=_atoms, cutoff=self.potential._cutoff)
else:
atoms = _atoms
Calculator.calculate(self, atoms, properties, system_changes)
# update without descriptors
self.atoms.update(posgrad=True, cellgrad=True, forced=True, build_locals=False,
descriptors=self.potential.gp.kern.kernels)
# calculate
energy, forces, cellgrad = self.ext_calculate(
self.potential, self.atoms)
# stress
stress1 = -(forces[:, None]*atoms.positions[..., None]).sum(axis=0)
stress2 = (cellgrad[:, None]*atoms.cell[..., None]).sum(axis=0)
try:
volume = atoms.get_volume()
except ValueError:
volume = -2. # here stress2=0, thus trace(stress) = virial (?)
stress = (stress1 + stress2)/volume
# results
self.results['energy'] = energy
self.results['forces'] = forces
self.results['free_energy'] = self.results['energy']
self.results['stress'] = stress.flat[[0, 4, 8, 5, 2, 1]]
def test():
from theforce.descriptor.atoms import TorchAtoms
from theforce.descriptor.radial import RepulsiveCore
torch.set_default_tensor_type(torch.DoubleTensor)
V = PairPot(55, 55, RepulsiveCore()) + PairPot(55, 55, RepulsiveCore())
a = TorchAtoms(positions=[(0, 0, 0), (2, 0, 0), (0, 2, 0)],
numbers=[55, 55, 55],
cell=[10, 10, 10],
pbc=False)
a.update(cutoff=5., posgrad=True)
e, f = V(a, forces=True)
e.backward()
print(a.xyz.grad.allclose(-f))
print(V.state)
print(
sum([
PairPot(55, 55, RepulsiveCore()),
PairPot(55, 55, RepulsiveCore())
]))
def get_unique_lces(self, thresh=0.95):
tmp = (self.atoms.as_ase() if self.to_ase else self.atoms).copy()
tmp.calc = None
atoms = TorchAtoms(ase_atoms=tmp)
atoms.update(posgrad=False,
cellgrad=False,
dont_save_grads=True,
cutoff=self.model.cutoff,
descriptors=self.model.descriptors)
k = self.model.gp.kern(atoms, atoms)
unique = []
for i in range(k.shape[0]):
is_unique = True
for j in unique:
if k[i, j] >= thresh:
is_unique = False
break
if is_unique:
unique.append(i)
return unique
def sample_rand_lces(self, indices=None, repeat=1, extend_cov=False):
added = 0
for _ in range(repeat):
tmp = (self.atoms.as_ase() if self.to_ase else self.atoms).copy()
shape = tmp.positions.shape
tmp.positions += np.random.uniform(-0.05, 0.05, size=shape)
tmp.calc = None
atoms = TorchAtoms(ase_atoms=tmp)
atoms.update(posgrad=False,
cellgrad=False,
dont_save_grads=True,
cutoff=self.model.cutoff,
descriptors=self.model.descriptors)
if indices is None:
indices = np.random.permutation(len(atoms.loc))
for k in indices:
res = abs(self.update_lce(atoms.loc[k]))
added += res
if res > 0 and extend_cov:
cov = self.model.gp.kern(self.atoms, self.model.X[-1])
self.cov = torch.cat([self.cov, cov], dim=1)
self.log(f'added {added} randomly displaced LCEs')
def rdf(self,
rmax,
nbins,
select='all',
srange=None,
sample_size=100,
file=None):
I = self.select(select)
s = Sampler(*srange) if srange else Sampler(self.start, self.stop)
data = AtomsData(
[TorchAtoms(self[s.sample()][I]) for _ in range(sample_size)])
r, gdict = rdf(data, rmax, nbins)
if file is not None:
#header = 'r ' + ' '.join(f'{key}' for key in gdict.keys())
header = ' '.join(f'{k[0]}-{k[1]}' for k in gdict.keys())
out = np.stack([
r,
] + [gdict[key] for key in gdict.keys()]).T
np.savetxt(file, out, header=header)
return r, gdict
def calculate(self,
_atoms=None,
properties=['energy'],
system_changes=all_changes):
if self.size[1] == 0 and not self.active:
raise RuntimeError('you forgot to assign a DFT calculator!')
if type(_atoms) == ase.atoms.Atoms:
atoms = TorchAtoms(ase_atoms=_atoms, ranks=self.distrib)
uargs = {
'cutoff': self.model.cutoff,
'descriptors': self.model.gp.kern.kernels
}
self.to_ase = True
else:
atoms = _atoms
uargs = {}
self.to_ase = False
if _atoms is not None and self.process_group is not None:
atoms.attach_process_group(self.process_group)
Calculator.calculate(self, atoms, properties, system_changes)
dat1 = self.size[0]
self.atoms.update(posgrad=True,
cellgrad=True,
forced=True,
dont_save_grads=True,
**uargs)
# build a model
if self.step == 0:
if self.active and self.model.ndata == 0:
self.initiate_model()
self._update_args = dict(data=False)
# kernel
self.cov = self.model.gp.kern(self.atoms, self.model.X)
# energy/forces
self.update_results(self.active or (self.meta is not None))
# active learning
self.deltas = None
self.covlog = ''
if self.active and not self.veto():
pre = self.results.copy()
m, n = self.update(**self._update_args)
if n > 0 or m > 0:
self.update_results(self.meta is not None)
if self.step > 0:
self.deltas = {}
for quant in ['energy', 'forces', 'stress']:
self.deltas[quant] = self.results[quant] - pre[quant]
if self.size[0] == dat1:
self.distrib.unload(atoms)
else:
covloss_max = float(self.get_covloss().max())
self.covlog = f'{covloss_max}'
if covloss_max > self.ediff:
tmp = self.atoms.as_ase()
tmp.calc = None
if self.rank == 0:
ase.io.Trajectory('active_uncertain.traj', 'a').write(tmp)
energy = self.results['energy']
# test
if self.active and self.test and self.step - self._last_test > self.test:
self._test()
# meta terms
meta = ''
if self.meta is not None:
energies, kwargs = self.meta(self)
if energies is not None:
meta_energy = self.reduce(energies, **kwargs)
meta = f'meta: {meta_energy}'
# step
self.log('{} {} {} {}'.format(energy, self.atoms.get_temperature(),
self.covlog, meta))
self.step += 1
# needed for self.calculate_numerical_stress
self.results['free_energy'] = self.results['energy']
def __init__(self, dyn, gp, ediff=0.1, fdiff=float('inf'), restrict=None, calculator=None, model=None,
algorithm='ultrafast', volatile=None, logfile='leapfrog.log', skip=10, skip_volatile=5,
undo_volatile=True, free_fall=100, correct_verlet=True, tune=(None, None), group=None):
self.dyn = dyn
self.gp = PosteriorPotential(gp).gp
self._ediff = ediff
self._fdiff = fdiff
self.restrict = restrict
self.skip = skip
self.skip_volatile = skip_volatile
self.undo_volatile = undo_volatile
self.free_fall = free_fall
self.correct_verlet = correct_verlet
self._tune = tune
if type(algorithm) == str:
self.algorithm = getattr(self, 'algorithm_'+algorithm)
else:
self.algorithm = types.MethodType(algorithm, self)
# atoms
if type(dyn.atoms) == ase.Atoms:
self.to_ase = True
dyn.atoms = TorchAtoms(dyn.atoms)
else:
self.to_ase = False
if group is not None:
self.atoms.attach_process_group(group)
self.atoms.update(cutoff=self.gp.cutoff,
descriptors=self.gp.descriptors)
# calc
if calculator:
self.calculator = calculator
else:
self.calculator = dyn.atoms.calc
# volatile
self._volatile = volatile if volatile else 2 if model is None else -1
# initiate
self.step = 0
self._fp = []
self._fp_e = []
self._ext = []
self.logfile = logfile
self.log('leapfrog says Hello!', mode='w')
self.log('volatile: {}'.format(self._volatile))
self.log('species: {}'.format(self.gp.species))
self.log('restrict: {}'.format(self.restrict))
# model
if model:
if type(model) == str:
potential = PosteriorPotentialFromFolder(model, group=group)
else:
potential = model
self.log('a model is provided with {} data and {} ref(s)'.format(
len(potential.data), len(potential.X)))
else:
assert ediff is not None
snap = self.snapshot()
potential = initial_model(self.gp, snap, ediff)
self.log('update: {} data: {} inducing: {} FP: {}'.format(
True, len(potential.data), len(potential.inducing), len(self._fp)))
self.log('a model is initiated with {} data and {} ref(s)'.format(
len(potential.data), len(potential.X)))
self.atoms.set_calculator(AutoForceCalculator(potential))
self.energy = [self.atoms.get_potential_energy()]
self.temperature = [self.atoms.get_temperature()]
# for parallelism
self.fp_is_allowed = True
def mlmd(ini_atoms,
cutoff,
au,
dt,
tolerance=0.1,
pair=True,
soap=True,
ndata=10,
max_steps=100,
itrain=10 * [5],
retrain=5 * [5],
retrain_every=100,
pes=potential_energy_surface):
"""
ML-assisted-MD: a calculator must be attached to ini_atoms.
Rules of thumb:
Initial training (itrain) is crucial for correct approximation
of variances.
Hyper-parameters are sensitive to nlocals=len(inducing) thus
if you don't want to retrain gp every time the data is updated,
at least keep nlocals fixed.
"""
dftcalc = ini_atoms.get_calculator()
# run a short MD to gather some (dft) data
atoms = TorchAtoms(ase_atoms=ini_atoms.copy())
atoms.set_velocities(ini_atoms.get_velocities())
atoms.set_calculator(dftcalc)
dyn = VelocityVerlet(atoms, dt=dt, trajectory='md.traj', logfile='md.log')
md_step = ndata
dyn.run(md_step)
ndft = md_step
# train a potential
data = AtomsData(traj='md.traj', cutoff=cutoff)
inducing = data.to_locals()
V = pes(data=data,
inducing=inducing,
cutoff=cutoff,
atomic_unit=au,
pairkernel=pair,
soapkernel=soap,
train=itrain,
test=True,
caching=True)
atoms.update(cutoff=cutoff, descriptors=V.gp.kern.kernels)
mlcalc = PosteriorVarianceCalculator(V)
atoms.set_calculator(mlcalc)
# long MD
while md_step < max_steps:
md_step += 1
forces = atoms.get_forces()
var = atoms.calc.results['forces_var']
tol = np.sqrt(var.max(axis=1))
if (tol > tolerance).any():
_forces = forces
_var = var
# new dft calculation
ndft += 1
print(
'|............... new dft calculation (total={})'.format(ndft))
tmp = atoms.copy()
tmp.set_calculator(dftcalc)
true_forces = tmp.get_forces()
# add new information to data
new_data = AtomsData(X=[TorchAtoms(ase_atoms=tmp)])
new_data.update(cutoff=cutoff,
descriptors=atoms.calc.potential.gp.kern.kernels)
new_locals = new_data.to_locals()
new_locals.stage(descriptors=atoms.calc.potential.gp.kern.kernels)
data += new_data
inducing += new_locals # TODO: importance sampling
# remove old(est) information
del data.X[0]
del inducing.X[:len(new_locals)] # TODO: importance sampling
# retrain
if ndft % retrain_every == 0:
print('|............... : retraining for {} steps'.format(
retrain))
for steps in iterable(retrain):
atoms.calc.potential.train(data,
inducing=inducing,
steps=steps,
cov_loss=False)
atoms.calc.potential.gp.to_file(
'gp.chp', flag='ndft={}'.format(ndft))
# update model
print('|............... new regression')
atoms.calc.potential.set_data(data, inducing, use_caching=True)
# new forces
atoms.calc.results.clear()
forces = atoms.get_forces()
var = atoms.calc.results['forces_var']
# report
_err_pred = np.sqrt(_var).max()
_err = np.abs(_forces - true_forces).max()
err_pred = np.sqrt(var).max()
err = np.abs(forces - true_forces).max()
print('|............... : old max-error: predicted={}, true={}'.
format(_err_pred, _err))
print('|............... : new max-error: predicted={}, true={}'.
format(err_pred, err))
arrays = np.concatenate([true_forces, _forces, forces, _var, var],
axis=1)
with open('forces_var.txt', 'ab') as report:
np.savetxt(report, arrays)
print(md_step, '_')
dyn.run(1)
print('finished {} steps, used dftcalc only {} times'.format(
md_step, ndft))
def example():
from ase.calculators.lj import LennardJones
from theforce.util.flake import Hex
from ase.io.trajectory import Trajectory
from ase.optimize import BFGSLineSearch
from ase.io import Trajectory
# initial state + dft calculator
ini_atoms = TorchAtoms(positions=Hex().array(), cutoff=3.0)
dftcalc = LennardJones()
ini_atoms.set_calculator(dftcalc)
BFGSLineSearch(ini_atoms).run(fmax=0.01)
vel = np.random.uniform(-1., 1., size=ini_atoms.positions.shape) * 1.
vel -= vel.mean(axis=0)
ini_atoms.set_velocities(vel)
# use a pretrained model by writing it to the checkpoint
# (alternatively set, for example, itrain=10*[5] in mlmd)
pre = """GaussianProcessPotential([PairKernel(RBF(signal=3.2251566545458794, lengthscale=tensor([0.1040])),
0, 0, factor=PolyCut(3.0, n=2))], White(signal=0.1064043798026091, requires_grad=True))""".replace(
'\n', '')
with open('gp.chp', 'w') as chp:
chp.write(pre)
# run(md)
mlmd(ini_atoms,
3.0,
0.5,
0.03,
tolerance=0.18,
max_steps=100,
soap=False,
itrain=10 * [3],
retrain_every=5,
retrain=5)
# recalculate all with the actual calculator and compare
traj = Trajectory('md.traj')
energies = []
forces = []
dum = 0
for atoms in traj:
dum += 1
e = atoms.get_potential_energy()
f = atoms.get_forces()
dftcalc.calculate(atoms)
ee = dftcalc.results['energy']
ff = dftcalc.results['forces']
energies += [(e, ee)]
forces += [(f.reshape(-1), ff.reshape(-1))]
import pylab as plt
get_ipython().run_line_magic('matplotlib', 'inline')
fig, axes = plt.subplots(1, 2, figsize=(8, 3))
axes[0].scatter(*zip(*energies))
axes[0].set_xlabel('ml')
axes[0].set_ylabel('dft')
a, b = (np.concatenate(v) for v in zip(*forces))
axes[1].scatter(a, b)
axes[1].set_xlabel('ml')
axes[1].set_ylabel('dft')
fig.tight_layout()
fig.text(0.2, 0.8, 'energy')
fig.text(0.7, 0.8, 'forces')