"""Relax atomic structures.
By defaults read from "unrelaxed.json" from disk and relaxes
structures and saves the final relaxed structure in "structure.json".
The relax recipe has a couple of note-worthy features:
- It automatically handles structures of any dimensionality
- It tries to enforce symmetries
- It continously checks after each step that no symmetries are broken,
and raises an error if this happens.
The recipe also supports relaxing structure with vdW forces using DFTD3.
To install DFTD3 do
.. code-block:: console
$ mkdir ~/DFTD3 && cd ~/DFTD3
$ wget chemie.uni-bonn.de/pctc/mulliken-center/software/dft-d3/dftd3.tgz
$ tar -zxf dftd3.tgz
$ make
$ echo 'export ASE_DFTD3_COMMAND=$HOME/DFTD3/dftd3' >> ~/.bashrc
$ source ~/.bashrc
Examples
--------
Relax without using DFTD3
.. code-block:: console
$ ase build -x diamond Si unrelaxed.json
$ asr run "relax --nod3"
Relax using the LDA exchange-correlation functional
.. code-block:: console
$ ase build -x diamond Si unrelaxed.json
$ asr run "relax --calculator {'xc':'LDA',...}"
"""
import typing
from pathlib import Path
import numpy as np
from ase.io import write, Trajectory
from ase import Atoms
from ase.optimize.bfgs import BFGS
from ase.calculators.calculator import PropertyNotImplementedError
from asr.core import command, option, AtomsFile, DictStr, prepare_result, ASRResult
from math import sqrt
import time
[docs]
class BrokenSymmetryError(Exception):
pass
Uvalues = {}
# From [acs comb sci 2.011, 13, 383-390, Setyawan et al.]
UTM = {'Ti': 4.4, 'V': 2.7, 'Cr': 3.5, 'Mn': 4.0, 'Fe': 4.6,
'Co': 5.0, 'Ni': 5.1, 'Cu': 4.0, 'Zn': 7.5, 'Ga': 3.9,
'Nb': 2.1, 'Mo': 2.4, 'Tc': 2.7, 'Ru': 3.0, 'Rh': 3.3,
'Pd': 3.6, 'Cd': 2.1, 'In': 1.9,
'Ta': 2.0, 'W': 2.2, 'Re': 2.4, 'Os': 2.6, 'Ir': 2.8, 'Pt': 3.0}
for key, value in UTM.items():
Uvalues[key] = ':d,{},0'.format(value)
[docs]
def is_relax_done(atoms, fmax=0.01, smax=0.002,
smask=np.array([1, 1, 1, 1, 1, 1])):
f = atoms.get_forces()
s = atoms.get_stress() * smask
done = (f**2).sum(1).max() <= fmax**2 and abs(s).max() <= smax
return done
[docs]
class SpgAtoms(Atoms):
[docs]
@classmethod
def from_atoms(cls, atoms):
# Due to technicalities we cannot mess with the __init__ constructor
# -> therefore we make our own
a = cls(atoms)
a.set_symmetries([np.eye(3)], [[0, 0, 0]])
return a
[docs]
def set_symmetries(self, symmetries, translations):
self.t_sc = translations
self.op_svv = [np.linalg.inv(self.cell).dot(op_cc.T).dot(self.cell) for
op_cc in symmetries]
self.nsym = len(symmetries)
tolerance = 1e-4
spos_ac = self.get_scaled_positions()
a_sa = []
for op_cc, t_c in zip(symmetries, self.t_sc):
symspos_ac = np.dot(spos_ac, op_cc.T) + t_c
a_a = []
for s_c in symspos_ac:
diff_ac = spos_ac - s_c
diff_ac -= np.round(diff_ac)
mask_c = np.all(np.abs(diff_ac) < tolerance, axis=1)
assert np.sum(mask_c) == 1, f'Bad symmetry, {mask_c}'
ind = np.argwhere(mask_c)[0][0]
assert ind not in a_a, f'Bad symmetry {ind}, {diff_ac}'
a_a.append(ind)
a_sa.append(a_a)
self.a_sa = np.array(a_sa)
[docs]
def get_stress(self, voigt=True, *args, **kwargs):
sigma0_vv = Atoms.get_stress(self, voigt=False, *args, **kwargs)
sigma_vv = np.zeros((3, 3))
for op_vv in self.op_svv:
sigma_vv += np.dot(np.dot(op_vv, sigma0_vv), op_vv.T)
sigma_vv /= self.nsym
if voigt:
return sigma_vv.flat[[0, 4, 8, 5, 2, 1]]
return sigma_vv
[docs]
def get_forces(self, *args, **kwargs):
f0_av = Atoms.get_forces(self, *args, **kwargs)
f_av = np.zeros_like(f0_av)
for map_a, op_vv in zip(self.a_sa, self.op_svv):
for a1, a2 in enumerate(map_a):
f_av[a2] += np.dot(f0_av[a1], op_vv)
f_av /= self.nsym
return f_av
[docs]
class myBFGS(BFGS):
[docs]
def log(self, forces=None, stress=None):
if forces is None:
forces = self.atoms.atoms.get_forces()
if stress is None:
stress = self.atoms.atoms.get_stress()
fmax = sqrt((forces**2).sum(axis=1).max())
smax = abs(stress).max()
e = self.atoms.get_potential_energy(
force_consistent=self.force_consistent)
T = time.localtime()
if self.logfile is not None:
name = self.__class__.__name__
if self.nsteps == 0:
self.logfile.write(' ' * len(name)
+ ' {:<4} {:<8} {:<10} '.format('Step',
'Time',
'Energy')
+ '{:<10} {:<10}\n'.format('fmax',
'smax'))
if self.force_consistent:
self.logfile.write(
'*Force-consistent energies used in optimization.\n')
fc = '*' if self.force_consistent else ''
self.logfile.write(f'{name}: {self.nsteps:<4} '
f'{T[3]:02d}:{T[4]:02d}:{T[5]:02d} '
f'{e:<10.6f}{fc} {fmax:<10.4f} {smax:<10.4f}\n')
self.logfile.flush()
[docs]
def relax(atoms, tmp_atoms_file, emin=-np.inf, smask=None, dftd3=True,
fixcell=False, allow_symmetry_breaking=False, dft=None,
fmax=0.01, enforce_symmetry=False):
if dftd3:
from ase.calculators.dftd3 import DFTD3
nd = int(np.sum(atoms.get_pbc()))
if smask is None:
if fixcell:
smask = [0, 0, 0, 0, 0, 0]
elif nd == 3:
smask = [1, 1, 1, 1, 1, 1]
elif nd == 2:
smask = [1, 1, 0, 0, 0, 1]
else:
pbc = atoms.get_pbc()
assert pbc[2], "1D periodic axis should be the last one."
smask = [0, 0, 1, 0, 0, 0]
from asr.utils.symmetry import atoms2symmetry
dataset = atoms2symmetry(atoms,
tolerance=1e-3,
angle_tolerance=0.1).dataset
spgname = dataset['international']
number = dataset['number']
nsym = len(dataset['rotations'])
atoms = SpgAtoms.from_atoms(atoms)
if enforce_symmetry:
atoms.set_symmetries(symmetries=dataset['rotations'],
translations=dataset['translations'])
if dftd3:
calc = DFTD3(dft=dft)
else:
calc = dft
atoms.calc = calc
# We are fixing atom=0 to reduce computational effort
from ase.constraints import ExpCellFilter
filter = ExpCellFilter(atoms, mask=smask)
name = Path(tmp_atoms_file).with_suffix('').name
try:
trajfile = Trajectory(tmp_atoms_file, 'a', atoms)
opt = myBFGS(filter,
logfile=name,
trajectory=trajfile)
# fmax=0 here because we have implemented our own convergence criteria
runner = opt.irun(fmax=0)
for _ in runner:
# Check that the symmetry has not been broken
newdataset = atoms2symmetry(atoms,
tolerance=1e-3,
angle_tolerance=0.1).dataset
spgname2 = newdataset['international']
number2 = newdataset['number']
nsym2 = len(newdataset['rotations'])
msg = (f'The initial spacegroup was {spgname} {number} '
f'but it changed to {spgname2} {number2} during '
'the relaxation.')
if (not allow_symmetry_breaking
and number != number2 and nsym > nsym2):
# Log the last step
opt.log()
opt.call_observers()
errmsg = 'The symmetry was broken during the relaxation! ' + msg
raise BrokenSymmetryError(errmsg)
elif number != number2:
print('Not an error: The spacegroup has changed during relaxation. '
+ msg)
spgname = spgname2
number = number2
nsym = nsym2
if enforce_symmetry:
atoms.set_symmetries(symmetries=newdataset['rotations'],
translations=newdataset['translations'])
if is_relax_done(atoms, fmax=fmax, smax=0.002, smask=smask):
opt.log()
opt.call_observers()
break
finally:
trajfile.close()
if opt.logfile is not None:
opt.logfile.close()
return atoms
[docs]
def set_initial_magnetic_moments(atoms):
atoms.set_initial_magnetic_moments(np.ones(len(atoms), float))
[docs]
@prepare_result
class Result(ASRResult):
"""Result class for :py:func:`asr.relax.main`."""
version: int = 0
atoms: Atoms
images: typing.List[Atoms]
etot: float
edft: float
spos: np.ndarray
symbols: typing.List[str]
a: float
b: float
c: float
alpha: float
beta: float
gamma: float
key_descriptions = \
{'atoms': 'Relaxed atomic structure.',
'images': 'Path taken when relaxing structure.',
'etot': 'Total energy [eV]',
'edft': 'DFT total energy [eV]',
'spos': 'Array: Scaled positions',
'symbols': 'Array: Chemical symbols',
'a': 'Cell parameter a [Ang]',
'b': 'Cell parameter b [Ang]',
'c': 'Cell parameter c [Ang]',
'alpha': 'Cell parameter alpha [deg]',
'beta': 'Cell parameter beta [deg]',
'gamma': 'Cell parameter gamma [deg]'}
@command('asr.relax',
creates=['structure.json'],
returns=Result)
@option('-a', '--atoms', help='Atoms to be relaxed.',
type=AtomsFile(), default='unrelaxed.json')
@option('--tmp-atoms', help='File containing recent progress.',
type=AtomsFile(must_exist=False), default='relax.traj')
@option('--tmp-atoms-file', help='File to store snapshots of relaxation.',
default='relax.traj', type=str)
@option('-c', '--calculator', help='Calculator and its parameters.',
type=DictStr())
@option('--d3/--nod3', help='Relax with vdW D3.', is_flag=True)
@option('--fixcell/--dont-fixcell',
help="Don't relax stresses.",
is_flag=True)
@option('--allow-symmetry-breaking/--dont-allow-symmetry-breaking',
help='Allow symmetries to be broken during relaxation.',
is_flag=True)
@option('--fmax', help='Maximum force allowed.', type=float)
@option('--enforce-symmetry/--dont-enforce-symmetry',
help='Symmetrize forces and stresses.', is_flag=True)
def main(atoms: Atoms,
calculator: dict = {'name': 'gpaw',
'mode': {'name': 'pw', 'ecut': 800},
'xc': 'PBE',
'kpts': {'density': 6.0, 'gamma': True},
'basis': 'dzp',
'symmetry': {'symmorphic': False},
'convergence': {'forces': 1e-4},
'txt': 'relax.txt',
'occupations': {'name': 'fermi-dirac',
'width': 0.05},
'charge': 0},
tmp_atoms: typing.Optional[Atoms] = None,
tmp_atoms_file: str = 'relax.traj',
d3: bool = False,
fixcell: bool = False,
allow_symmetry_breaking: bool = False,
fmax: float = 0.01,
enforce_symmetry: bool = True) -> Result:
"""Relax atomic positions and unit cell.
The relaxed structure is saved to `structure.json` which can be processed
by other recipes.
Parameters
----------
atoms
Atomic structure to relax.
calculator
Calculator dictionary description.
tmp_atoms
Atoms from a restarted calculation.
tmp_atoms_file
Filename to save relaxed trajectory in.
d3
Relax using DFTD3.
fixcell
Fix cell when relaxing, thus only relaxing atomic positions.
allow_symmetry_breaking
Allow structure to break symmetry.
fmax
Maximum force tolerance.
enforce_symmetry
Enforce symmetries. When enabled, the atomic structure, forces and
stresses will be symmetrized at each step of the relaxation.
"""
from ase.calculators.calculator import get_calculator_class
if tmp_atoms is not None:
atoms = tmp_atoms
# Make our own copy
atoms = atoms.copy()
if not atoms.has('initial_magmoms'):
set_initial_magnetic_moments(atoms)
calculatorname = calculator.pop('name')
Calculator = get_calculator_class(calculatorname)
# Some calculator specific parameters
nd = int(np.sum(atoms.get_pbc()))
if calculatorname == 'gpaw':
if 'kpts' in calculator:
from ase.calculators.calculator import kpts2kpts
if 'density' in calculator['kpts']:
kpts = kpts2kpts(calculator['kpts'], atoms=atoms)
calculator['kpts'] = kpts
if nd == 2:
assert not atoms.get_pbc()[2], \
('The third unit cell axis should be aperiodic for '
'a 2D material!')
calculator['poissonsolver'] = {'dipolelayer': 'xy'}
calc = Calculator(**calculator)
# Relax the structure
atoms = relax(atoms, tmp_atoms_file=tmp_atoms_file, dftd3=d3,
fixcell=fixcell,
allow_symmetry_breaking=allow_symmetry_breaking,
dft=calc, fmax=fmax, enforce_symmetry=enforce_symmetry)
# If the maximum magnetic moment on all atoms is big then
try:
magmoms = atoms.get_magnetic_moments()
except PropertyNotImplementedError:
# We assume this means that the magnetic moments are zero
# for this calculator.
magmoms = np.zeros(len(atoms))
if not abs(magmoms).max() > 0.1:
atoms.set_initial_magnetic_moments([0] * len(atoms))
calc = Calculator(**calculator)
# Relax the structure
atoms = relax(atoms, tmp_atoms_file=tmp_atoms_file, dftd3=d3,
fixcell=fixcell,
allow_symmetry_breaking=allow_symmetry_breaking,
dft=calc, fmax=fmax, enforce_symmetry=enforce_symmetry)
edft = calc.get_potential_energy(atoms)
etot = atoms.get_potential_energy()
cellpar = atoms.cell.cellpar()
# XXX
# metadata = calc.get_metadata()
# Save atomic structure
write('structure.json', atoms)
trajectory = Trajectory(tmp_atoms_file, 'r')
images = []
for image in trajectory:
images.append(image)
return Result.fromdata(
atoms=atoms.copy(),
etot=etot,
edft=edft,
a=cellpar[0],
b=cellpar[1],
c=cellpar[2],
alpha=cellpar[3],
beta=cellpar[4],
gamma=cellpar[5],
spos=atoms.get_scaled_positions(),
symbols=atoms.get_chemical_symbols(),
images=images
)
if __name__ == '__main__':
main.cli()