"""Checkpointing and restart functionality for scripts using ASE Atoms objects.
Initialize checkpoint object:
CP = Checkpoint('checkpoints.db')
Checkpointed code block in try ... except notation:
try:
a, C, C_err = CP.load()
except NoCheckpoint:
C, C_err = fit_elastic_constants(a)
CP.save(a, C, C_err)
Checkpoint code block, shorthand notation:
C, C_err = CP(fit_elastic_constants)(a)
Example for checkpointing within an iterative loop, e.g. for searching crack
tip position:
try:
a, converged, tip_x, tip_y = CP.load()
except NoCheckpoint:
converged = False
tip_x = tip_x0
tip_y = tip_y0
while not converged:
... do something to find better crack tip position ...
converged = ...
CP.flush(a, converged, tip_x, tip_y)
The simplest way to use checkpointing is through the CheckpointCalculator. It
wraps any calculator object and does a checkpoint whenever a calculation
is performed:
calc = ...
cp_calc = CheckpointCalculator(calc)
atoms.calc = cp_calc
e = atoms.get_potential_energy() # 1st time, does calc, writes to checkfile
# subsequent runs, reads from checkpoint
"""
from typing import Any, Dict
import numpy as np
import ase
from ase.calculators.calculator import Calculator
from ase.db import connect
class NoCheckpoint(Exception):
pass
class DevNull:
def write(str, *args):
pass
[docs]
class Checkpoint:
_value_prefix = '_values_'
def __init__(self, db='checkpoints.db', logfile=None):
self.db = db
if logfile is None:
logfile = DevNull()
self.logfile = logfile
self.checkpoint_id = [0]
self.in_checkpointed_region = False
def __call__(self, func, *args, **kwargs):
checkpoint_func_name = str(func)
def decorated_func(*args, **kwargs):
# Get the first ase.Atoms object.
atoms = None
for a in args:
if atoms is None and isinstance(a, ase.Atoms):
atoms = a
try:
retvals = self.load(atoms=atoms)
except NoCheckpoint:
retvals = func(*args, **kwargs)
if isinstance(retvals, tuple):
self.save(*retvals, atoms=atoms,
checkpoint_func_name=checkpoint_func_name)
else:
self.save(retvals, atoms=atoms,
checkpoint_func_name=checkpoint_func_name)
return retvals
return decorated_func
def _increase_checkpoint_id(self):
if self.in_checkpointed_region:
self.checkpoint_id += [1]
else:
self.checkpoint_id[-1] += 1
self.logfile.write('Entered checkpoint region '
'{}.\n'.format(self.checkpoint_id))
self.in_checkpointed_region = True
def _decrease_checkpoint_id(self):
self.logfile.write('Leaving checkpoint region '
'{}.\n'.format(self.checkpoint_id))
if not self.in_checkpointed_region:
self.checkpoint_id = self.checkpoint_id[:-1]
assert len(self.checkpoint_id) >= 1
self.in_checkpointed_region = False
assert self.checkpoint_id[-1] >= 1
def _mangled_checkpoint_id(self):
"""
Returns a mangled checkpoint id string:
check_c_1:c_2:c_3:...
E.g. if checkpoint is nested and id is [3,2,6] it returns:
'check3:2:6'
"""
return 'check' + ':'.join(str(id) for id in self.checkpoint_id)
[docs]
def load(self, atoms=None):
"""
Retrieve checkpoint data from file. If atoms object is specified, then
the calculator connected to that object is copied to all returning
atoms object.
Returns tuple of values as passed to flush or save during checkpoint
write.
"""
self._increase_checkpoint_id()
retvals = []
with connect(self.db) as db:
try:
dbentry = db.get(checkpoint_id=self._mangled_checkpoint_id())
except KeyError:
raise NoCheckpoint
data = dbentry.data
atomsi = data['checkpoint_atoms_args_index']
i = 0
while (i == atomsi or
f'{self._value_prefix}{i}' in data):
if i == atomsi:
newatoms = dbentry.toatoms()
if atoms is not None:
# Assign calculator
newatoms.calc = atoms.calc
retvals += [newatoms]
else:
retvals += [data[f'{self._value_prefix}{i}']]
i += 1
self.logfile.write('Successfully restored checkpoint '
'{}.\n'.format(self.checkpoint_id))
self._decrease_checkpoint_id()
if len(retvals) == 1:
return retvals[0]
else:
return tuple(retvals)
def _flush(self, *args, **kwargs):
data = {f'{self._value_prefix}{i}': v
for i, v in enumerate(args)}
try:
atomsi = [isinstance(v, ase.Atoms) for v in args].index(True)
atoms = args[atomsi]
del data[f'{self._value_prefix}{atomsi}']
except ValueError:
atomsi = -1
try:
atoms = kwargs['atoms']
except KeyError:
raise RuntimeError('No atoms object provided in arguments.')
try:
del kwargs['atoms']
except KeyError:
pass
data['checkpoint_atoms_args_index'] = atomsi
data.update(kwargs)
with connect(self.db) as db:
try:
dbentry = db.get(checkpoint_id=self._mangled_checkpoint_id())
del db[dbentry.id]
except KeyError:
pass
db.write(atoms, checkpoint_id=self._mangled_checkpoint_id(),
data=data)
self.logfile.write('Successfully stored checkpoint '
'{}.\n'.format(self.checkpoint_id))
[docs]
def flush(self, *args, **kwargs):
"""
Store data to a checkpoint without increasing the checkpoint id. This
is useful to continuously update the checkpoint state in an iterative
loop.
"""
# If we are flushing from a successfully restored checkpoint, then
# in_checkpointed_region will be set to False. We need to reset to True
# because a call to flush indicates that this checkpoint is still
# active.
self.in_checkpointed_region = False
self._flush(*args, **kwargs)
[docs]
def save(self, *args, **kwargs):
"""
Store data to a checkpoint and increase the checkpoint id. This closes
the checkpoint.
"""
self._decrease_checkpoint_id()
self._flush(*args, **kwargs)
def atoms_almost_equal(a, b, tol=1e-9):
return (np.abs(a.positions - b.positions).max() < tol and
(a.numbers == b.numbers).all() and
np.abs(a.cell - b.cell).max() < tol and
(a.pbc == b.pbc).all())
[docs]
class CheckpointCalculator(Calculator):
"""
This wraps any calculator object to checkpoint whenever a calculation
is performed.
This is particularly useful for expensive calculators, e.g. DFT and
allows usage of complex workflows.
Example usage:
calc = ...
cp_calc = CheckpointCalculator(calc)
atoms.calc = cp_calc
e = atoms.get_potential_energy()
# 1st time, does calc, writes to checkfile
# subsequent runs, reads from checkpoint file
"""
implemented_properties = ase.calculators.calculator.all_properties
default_parameters: Dict[str, Any] = {}
name = 'CheckpointCalculator'
property_to_method_name = {
'energy': 'get_potential_energy',
'energies': 'get_potential_energies',
'forces': 'get_forces',
'stress': 'get_stress',
'stresses': 'get_stresses'}
def __init__(self, calculator, db='checkpoints.db', logfile=None):
Calculator.__init__(self)
self.calculator = calculator
if logfile is None:
logfile = DevNull()
self.checkpoint = Checkpoint(db, logfile)
self.logfile = logfile
[docs]
def calculate(self, atoms, properties, system_changes):
Calculator.calculate(self, atoms, properties, system_changes)
try:
results = self.checkpoint.load(atoms)
prev_atoms, results = results[0], results[1:]
try:
assert atoms_almost_equal(atoms, prev_atoms)
except AssertionError:
raise AssertionError('mismatch between current atoms and '
'those read from checkpoint file')
self.logfile.write('retrieved results for {} from checkpoint\n'
.format(properties))
# save results in calculator for next time
if isinstance(self.calculator, Calculator):
if not hasattr(self.calculator, 'results'):
self.calculator.results = {}
self.calculator.results.update(dict(zip(properties, results)))
except NoCheckpoint:
if isinstance(self.calculator, Calculator):
self.logfile.write('doing calculation of {} with new-style '
'calculator interface\n'.format(properties))
self.calculator.calculate(atoms, properties, system_changes)
results = [self.calculator.results[prop]
for prop in properties]
else:
self.logfile.write('doing calculation of {} with old-style '
'calculator interface\n'.format(properties))
results = []
for prop in properties:
method_name = self.property_to_method_name[prop]
method = getattr(self.calculator, method_name)
results.append(method(atoms))
_calculator = atoms.calc
try:
atoms.calc = self.calculator
self.checkpoint.save(atoms, *results)
finally:
atoms.calc = _calculator
self.results = dict(zip(properties, results))