Source code for amp

import os
import sys
import shutil
import numpy as np
import tempfile
import platform
from getpass import getuser
from socket import gethostname
import subprocess
import warnings

import ase
from ase.calculators.calculator import Calculator, Parameters
    from ase import __version__ as aseversion
except ImportError:
    # We're on ASE 3.9 or older
    from ase.version import version as aseversion

from .utilities import (make_filename, hash_images, Logger, string2dict,
                        logo, now, assign_cores, TrainingConvergenceError)

    from amp import fmodules
except ImportError:
    warnings.warn('Did not find fortran modules.')
    fmodules_version = 9
    wrong_version = fmodules.check_version(version=fmodules_version)
    if wrong_version:
        raise RuntimeError('fortran modules are not updated. Recompile'
                           'with f2py as described in the README. '
                           'Correct version is %i.' % fmodules_version)

_ampversion = '0.5'

[docs]class Amp(Calculator, object): """Atomistic Machine-Learning Potential (Amp) ASE calculator Parameters ---------- descriptor : object Class representing local atomic environment. model : object Class representing the regression model. Can be only NeuralNetwork for now. Input arguments for NeuralNetwork are hiddenlayers, activation, weights, and scalings; for more information see docstring for the class NeuralNetwork. label : str Default prefix/location used for all files. dblabel : str Optional separate prefix/location for database files, including fingerprints, fingerprint derivatives, and neighborlists. This file location can be shared between calculator instances to avoid re-calculating redundant information. If not supplied, just uses the value from label. cores : int Can specify cores to use for parallel training; if None, will determine from environment envcommand : string For parallel processing across nodes, a command can be supplied here to load the appropriate environment before starting workers. logging : boolean Option to turn off logging; e.g., to speed up force calls. atoms : object ASE atoms objects with positions, symbols, energy, and forces in ASE format. """ implemented_properties = ['energy', 'forces'] def __init__(self, descriptor, model, label='amp', dblabel=None, cores=None, envcommand=None, logging=True, atoms=None): self.logging = logging Calculator.__init__(self, label=label, atoms=atoms) # Note self._log is set and self._printheader is called by above # call when it runs self.set_label. self._parallel = {'envcommand': envcommand} # Note the following are properties: these are setter functions. self.descriptor = descriptor self.model = model self.cores = cores # Note this calls 'assign_cores'. self.dblabel = label if dblabel is None else dblabel @property def cores(self): """ Get or set the cores for the parallel environment. Parameters ---------- cores : int or dictionary Parallel configuration. If cores is an integer, parallelizes over this many processes on machine localhost. cores can also be a dictionary of the type {'node324': 16, 'node325': 16}. If not specified, tries to determine from environment, using amp.utilities.assign_cores. """ return self._parallel['cores'] @cores.setter def cores(self, cores): self._parallel['cores'] = assign_cores(cores, log=self._log) @property def descriptor(self): """ Get or set the atomic descriptor. Parameters ---------- descriptor : object Class instance representing the local atomic environment. """ return self._descriptor @descriptor.setter def descriptor(self, descriptor): descriptor.parent = self # gives the descriptor object a reference to # the main Amp instance. Then descriptor can pull parameters directly # from Amp without needing them to be passed in each method call. self._descriptor = descriptor self.reset() # Clears any old calculations. @property def model(self): """ Get or set the machine-learning model. Parameters ---------- model : object Class instance representing the regression model. """ return self._model @model.setter def model(self, model): model.parent = self # gives the model object a reference to the main # Amp instance. Then model can pull parameters directly from Amp # without needing them to be passed in each method call. self._model = model self.reset() # Clears any old calculations. @classmethod
[docs] def load(Cls, file, Descriptor=None, Model=None, **kwargs): """Attempts to load calculators and return a new instance of Amp. Only a filename or file-like object is required, in typical cases. If using a home-rolled descriptor or model, also supply uninstantiated classes to those models, as in Model=MyModel. (Not as Model=MyModel()!) Any additional keyword arguments (such as label or dblabel) can be fed through to Amp. Parameters ---------- file : str Name of the file to load data from. Descriptor : object Class representing local atomic environment. Model : object Class representing the regression model. """ if hasattr(file, 'read'): text = else: with open(file) as f: text = # Unpack parameter dictionaries. p = string2dict(text) for key in ['descriptor', 'model']: p[key] = string2dict(p[key]) # If modules are not specified, find them. if Descriptor is None: Descriptor = importhelper(p['descriptor'].pop('importname')) if Model is None: Model = importhelper(p['model'].pop('importname')) # Key 'importname' and the value removed so that it is not splatted # into the keyword arguments used to instantiate in the next line. # Instantiate the descriptor and model. descriptor = Descriptor(**p['descriptor']) # ** sends all the key-value pairs at once. model = Model(**p['model']) # Instantiate Amp. calc = Cls(descriptor=descriptor, model=model, **kwargs) calc._log('Loaded file: %s' % file) return calc
[docs] def set(self, **kwargs): """Function to set parameters. For now, this doesn't do anything as all parameters are within the model and descriptor. """ changed_parameters = Calculator.set(self, **kwargs) if len(changed_parameters) > 0: self.reset()
[docs] def set_label(self, label): """Sets label, ensuring that any needed directories are made. Parameters ---------- label : str Default prefix/location used for all files. """ Calculator.set_label(self, label) # Create directories for output structure if needed. # Note ASE doesn't do this for us. if self.label: if ( != os.curdir and not os.path.isdir( os.makedirs( if self.logging is True: self._log = Logger(make_filename(self.label, '-log.txt')) else: self._log = Logger(None) self._printheader(self._log)
[docs] def calculate(self, atoms, properties, system_changes): """Calculation of the energy of system and forces of all atoms. """ # The inherited method below just sets the atoms object, # if specified, to self.atoms. Calculator.calculate(self, atoms, properties, system_changes) log = self._log log('Calculation requested.') images = hash_images([self.atoms]) key = images.keys()[0] if properties == ['energy']: log('Calculating potential energy...', tic='pot-energy') self.descriptor.calculate_fingerprints(images=images, log=log, calculate_derivatives=False) energy = self.model.calculate_energy( self.descriptor.fingerprints[key]) self.results['energy'] = energy log('...potential energy calculated.', toc='pot-energy') if properties == ['forces']: log('Calculating forces...', tic='forces') self.descriptor.calculate_fingerprints(images=images, log=log, calculate_derivatives=True) forces = \ self.model.calculate_forces( self.descriptor.fingerprints[key], self.descriptor.fingerprintprimes[key]) self.results['forces'] = forces log('...forces calculated.', toc='forces')
[docs] def train(self, images, overwrite=False, ): """Fits the model to the training images. Parameters ---------- images : list or str List of ASE atoms objects with positions, symbols, energies, and forces in ASE format. This is the training set of data. This can also be the path to an ASE trajectory (.traj) or database (.db) file. Energies can be obtained from any reference, e.g. DFT calculations. overwrite : bool If an output file with the same name exists, overwrite it. """ log = self._log log('\nAmp training started. ' + now() + '\n') log('Descriptor: %s\n (%s)' % (self.descriptor.__class__.__name__, self.descriptor)) log('Model: %s\n (%s)' % (self.model.__class__.__name__, self.model)) images = hash_images(images, log=log) log('\nDescriptor\n==========') train_forces = self.model.forcetraining # Derivatives of fingerprints need to be calculated if train_forces is # True. calculate_derivatives = train_forces self.descriptor.calculate_fingerprints( images=images, parallel=self._parallel, log=log, calculate_derivatives=calculate_derivatives) log('\nModel fitting\n=============') result =, descriptor=self.descriptor, log=log, parallel=self._parallel) if result is True: log('Amp successfully trained. Saving current parameters.') filename = self.label + '.amp' else: log('Amp not trained successfully. Saving current parameters.') filename = make_filename(self.label, '-untrained-parameters.amp') filename =, overwrite) log('Parameters saved in file "%s".' % filename) log("This file can be opened with `calc = Amp.load('%s')`" % filename) if result is False: raise TrainingConvergenceError('Amp did not converge upon ' 'training. See log file for' ' more information.')
[docs] def save(self, filename, overwrite=False): """Saves the calculator in a way that it can be re-opened with load. Parameters ---------- filename : str File object or path to the file to write to. overwrite : bool If an output file with the same name exists, overwrite it. """ if os.path.exists(filename): if overwrite is False: oldfilename = filename filename = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.amp').name self._log('File "%s" exists. Instead saving to "%s".' % (oldfilename, filename)) else: oldfilename = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.amp').name self._log('Overwriting file: "%s". Moving original to "%s".' % (filename, oldfilename)) shutil.move(filename, oldfilename) descriptor = self.descriptor.tostring() model = self.model.tostring() p = Parameters({'descriptor': descriptor, 'model': model}) p.write(filename) return filename
def _printheader(self, log): """Prints header to log file; inspired by that in GPAW. """ log(logo) log('Amp: Atomistic Machine-learning Package') log('Developed by Andrew Peterson, Alireza Khorshidi, and others,') log('Brown University.') log('PI Website:') log('Official repository:') log('Official documentation:') log('Citation:') log(' Alireza Khorshidi & Andrew A. Peterson,') log(' Computer Physics Communications 207: 310-324 (2016).') log('') log('=' * 70) log('User: %s' % getuser()) log('Hostname: %s' % gethostname()) log('Date: %s' % now(with_utc=True)) uname = platform.uname() log('Architecture: %s' % uname[4]) log('PID: %s' % os.getpid()) log('Amp version: %s' % _ampversion) ampdirectory = os.path.dirname(os.path.abspath(__file__)) log('Amp directory: %s' % ampdirectory) commithash, commitdate = get_git_commit(ampdirectory) log(' Last commit: %s' % commithash) log(' Last commit date: %s' % commitdate) log('Python: v{0}.{1}.{2}: %s'.format(*sys.version_info[:3]) % sys.executable) log('ASE v%s: %s' % (aseversion, os.path.dirname(ase.__file__))) log('NumPy v%s: %s' % (np.version.version, os.path.dirname(np.__file__))) # SciPy is not a strict dependency. try: import scipy log('SciPy v%s: %s' % (scipy.version.version, os.path.dirname(scipy.__file__))) except ImportError: log('SciPy: not available') # ZMQ an pxssh are only necessary for parallel calculations. try: import zmq log('ZMQ/PyZMQ v%s/v%s: %s' % (zmq.zmq_version(), zmq.pyzmq_version(), os.path.dirname(zmq.__file__))) except ImportError: log('ZMQ: not available') try: import pxssh log('pxssh: %s' % os.path.dirname(pxssh.__file__)) except ImportError: log('pxssh: Not available from pxssh.') try: from pexpect import pxssh except ImportError: log('pxssh: Not available from pexpect.') else: import pexpect log('pxssh (via pexpect v%s): %s' % (pexpect.__version__, pxssh.__file__)) log('=' * 70)
[docs]def importhelper(importname): """Manually compiled list of available modules. This is to prevent the execution of arbitrary (potentially malicious) code. However, since there is an `eval` statement in string2dict maybe this is silly. """ if importname == '.descriptor.gaussian.Gaussian': from .descriptor.gaussian import Gaussian as Module elif importname == '.descriptor.zernike.Zernike': from .descriptor.zernike import Zernike as Module elif importname == '.descriptor.bispectrum.Bispectrum': from .descriptor.bispectrum import Bispectrum as Module elif importname == '.model.neuralnetwork.NeuralNetwork': from .model.neuralnetwork import NeuralNetwork as Module elif importname == '.model.neuralnetwork.tflow': from .model.tflow import NeuralNetwork as Module elif importname == '.model.LossFunction': from .model import LossFunction as Module else: raise NotImplementedError( 'Attempt to import the module %s. Was this intended? ' 'If so, trying manually importing this module and ' 'feeding it to Amp.load. To avoid this error, this ' 'module can be added to amp.importhelper.' % importname) return Module
[docs]def get_git_commit(ampdirectory): """Attempts to get the last git commit from the amp directory. """ pwd = os.getcwd() os.chdir(ampdirectory) try: with open(os.devnull, 'w') as devnull: output = subprocess.check_output(['git', 'log', '-1', '--pretty=%H\t%ci'], stderr=devnull) except: output = 'unknown hash\tunknown date' output = output.strip() commithash, commitdate = output.split('\t') os.chdir(pwd) return commithash, commitdate