def relax_atoms(atoms,
tol=1e-3,
method='lbfgs_precon',
max_steps=1000,
traj_file=None,
**kwargs):
import model
atoms.set_calculator(model.calculator)
if hasattr(model, 'Optimizer'):
method = 'model_optimizer'
opt = model.Optimizer(atoms)
opt.run(tol, max_steps)
elif method.startswith('lbfgs') or method == 'fire' or method == 'cg_n':
if method == 'lbfgs_ASE':
from ase.optimize import LBFGS
opt = LBFGS(atoms, **kwargs)
elif method == 'cg_n':
from quippy import Minim
opt = Minim(atoms,
relax_positions=True,
relax_cell=False,
method='cg_n')
else:
from ase.optimize.precon.precon import Exp
from ase.optimize.precon.lbfgs import PreconLBFGS
precon = None
if method.endswith('precon'):
precon = Exp(3.0, recalc_mu=True)
if method.startswith('lbfgs'):
opt = PreconLBFGS(atoms, precon=precon, **kwargs)
else:
opt = FIRE(atoms, **kwargs)
if traj_file is not None and method != 'cg_n':
traj = open(traj_file, 'w')
def write_trajectory():
write(traj, atoms, format='extxyz')
opt.attach(write_trajectory)
opt.run(tol, max_steps)
try:
traj.close()
except:
pass
else:
raise ValueError('unknown method %s!' % method)
return atoms
#samples, targets = init_single_class_batch()
samples, targets = data.init_multi_class_batch(args.nsample, freqs,
args.ntraining)
print("samples", len(samples))
print("targets", len(targets))
import model
import torch
net = model.Net(args.nsample, args.nclass_out)
print("net", net)
params = list(net.parameters())
print("params", len(params))
print("input_shape", params[0].size())
optimizer = model.Optimizer(net.parameters(), lr=args.learning_rate)
criterion = model.Criterion()
losses = []
for epoch in range(args.nepoch):
print("-", epoch, "--", end=" ")
loss_accum = 0
for target, sample in zip(targets, samples):
sample = torch.tensor(sample, dtype=torch.float).unsqueeze(1)
target = torch.tensor(target, dtype=torch.long)
#print("training", target.shape, target.dtype, sample.shape, sample.dtype)
#print(target)
optimizer.zero_grad()
prediction = net(sample)
#print("prediction", prediction.size())
def relax_atoms_cell(atoms,
tol=1e-3,
stol=None,
method='lbfgs_precon',
max_steps=100,
mask=None,
traj_file=None,
hydrostatic_strain=False,
constant_volume=False,
precon_apply_positions=True,
precon_apply_cell=True,
symmetrize=False,
**kwargs):
import model
#print "relax_atoms_cell using method",method
if symmetrize:
atoms.set_calculator(SymmetrizedCalculator(model.calculator, atoms))
else:
atoms.set_calculator(model.calculator)
## print "relax_atoms_cell initial e ", atoms.get_potential_energy()
## print "relax_atoms_cell initial f ", atoms.get_forces()
## print "relax_atoms_cell initial s ", atoms.get_stress()
if hasattr(model, 'Optimizer'):
method = 'model_optimizer'
if method != 'cg_n':
atoms = UnitCellFilter(atoms,
mask=mask,
hydrostatic_strain=hydrostatic_strain,
constant_volume=constant_volume)
if method.startswith('lbfgs') or method == 'fire' or method == 'cg_n':
if method == 'cg_n':
from quippy import Minim, fzeros
atoms.info['Minim_Hydrostatic_Strain'] = hydrostatic_strain
atoms.info['Minim_Constant_Volume'] = constant_volume
if mask is not None:
atoms.info['Minim_Lattice_Fix'] = fzeros((3, 3))
if not mask[0]:
atoms.info['Minim_Lattice_Fix'][1, 1] = 1.0
if not mask[1]:
atoms.info['Minim_Lattice_Fix'][2, 2] = 1.0
if not mask[2]:
atoms.info['Minim_Lattice_Fix'][3, 3] = 1.0
if not mask[3]:
atoms.info['Minim_Lattice_Fix'][1, 2] = 1.0
atoms.info['Minim_Lattice_Fix'][2, 1] = 1.0
if not mask[4]:
atoms.info['Minim_Lattice_Fix'][2, 3] = 1.0
atoms.info['Minim_Lattice_Fix'][3, 2] = 1.0
if not mask[5]:
atoms.info['Minim_Lattice_Fix'][1, 3] = 1.0
atoms.info['Minim_Lattice_Fix'][3, 1] = 1.0
opt = Minim(atoms,
relax_positions=True,
relax_cell=True,
method='cg_n')
else:
from ase.optimize.precon.precon import Exp
from ase.optimize.precon.lbfgs import PreconLBFGS
precon = None
if method.endswith('precon'):
precon = Exp(3.0,
apply_positions=precon_apply_positions,
apply_cell=precon_apply_cell,
recalc_mu=True)
if method.startswith('lbfgs'):
opt = PreconLBFGS(atoms, precon=precon, **kwargs)
else:
opt = FIRE(atoms, **kwargs)
if traj_file is not None:
traj = open(traj_file, 'w')
def write_trajectory():
try:
write(traj, atoms.atoms, format='extxyz')
except:
write(traj, atoms, format='extxyz')
opt.attach(write_trajectory)
if method != 'cg_n' and isinstance(opt, PreconLBFGS):
opt.run(tol, max_steps, smax=stol)
else:
opt.run(tol, max_steps)
if traj_file is not None:
traj.close()
elif method == 'model_optimizer':
opt = model.Optimizer(atoms.atoms)
opt.run()
else:
raise ValueError('unknown method %s!' % method)
if isinstance(atoms, UnitCellFilter):
return atoms.atoms
else:
return atoms
def train(net, batch_size, epochs, lr):
print('-' * 30)
print(" HYPERPARAMETERS ")
print('-' * 30)
print("Batch size = ", batch_size)
print("Epochs = ", epochs)
print("Learning rate = ", lr)
print('-' * 30)
# function to call in data
train_loader, test_loader = data.getTrainingSets(filename1, filename2,
batch_size)
# Create loss and optimizer functions
loss = model.Loss()
optimizer = model.Optimizer(net, lr)
trainingStartTime = time.time()
# Start training
totalStep = len(train_loader)
losses = []
val_losses = []
for epoch in range(epochs):
runningLoss = 0.0
totalTrainLoss = 0.0
startTime = time.time()
for i, (input, target) in enumerate(train_loader):
input = Variable(input)
target = Variable(target)
outputs = net(input)
loss_size = loss(outputs, target)
losses.append(loss_size.item())
optimizer.zero_grad()
loss_size.backward()
optimizer.step()
runningLoss += loss_size.item()
totalTrainLoss += loss_size.item()
for i, (pred, real) in enumerate(test_loader):
pred = Variable(pred)
real = Variable(real)
val_outputs = net(pred)
val_loss_size = loss(val_outputs, real)
val_losses.append(val_loss_size.item())
print('Epoch: {}/{}, Loss: {:.4f}, Val loss: {:0.4f}, Time: {:0.2f}s'.
format(epoch + 1, epochs, loss_size.item(), val_loss_size.item(),
time.time() - startTime))
return losses, val_losses