Chimie Théorique » QMX

#!/usr/bin/env python

"""This module defines extensible classes for running tests on molecules.

Use this to compare different calculators, XC functionals and so on by

calculating e.g. atomization energies and bond lengths across the

g2 database of small molecules.

"""

import os

import sys

import traceback

import numpy as np

from ase import PickleTrajectory, read

from ase.calculators.emt import EMT

from ase.data.molecules import molecule, atoms as g2_atoms, g1

class BatchTest:

    """Contains logic for looping over tests and file management."""

    def __init__(self, test):

        self.test = test

        self.txt = sys.stdout # ?

    def run_single_test(self, formula):

        print >> self.txt, self.test.name, formula, '...',

        self.txt.flush()

        filename = self.test.get_filename(formula)

        if os.path.exists(filename):

            print >> self.txt, 'Skipped.'

            return

        try:

            open(filename, 'w').close() # Empty file

            system, calc = self.test.setup(formula)

            self.test.run(formula, system, filename)

            print >> self.txt, 'OK!'

            self.txt.flush()

        except self.test.exceptions:

            print >> self.txt, 'Failed!'

            traceback.print_exc(file=self.txt)

            print >> self.txt

            self.txt.flush()

    def run(self, formulas):

        """Run a batch of tests.

        This will invoke the test method on each formula, printing

        status to stdout.

        Those formulas that already have test result files will

        be skipped."""

        # Create directories if necessary

        if self.test.dir and not os.path.isdir(self.test.dir):

            os.mkdir(self.test.dir) # Won't work on 'dir1/dir2', but oh well

        for formula in formulas:

            self.run_single_test(formula)

    def collect(self, formulas, verbose=False):

        """Yield results of previous calculations."""

        for formula in formulas:

            try:

                filename = self.test.get_filename(formula)

                results = self.test.retrieve(formula, filename)

                if verbose:

                    print >> self.txt, 'Loaded:', formula, filename

                yield formula, results

            except (IOError, RuntimeError, TypeError):

                # XXX which errors should we actually catch?

                if verbose:

                    print >> self.txt, 'Error:', formula, '[%s]' % filename

                    traceback.print_exc(file=self.txt)

class MoleculeTest:

    """Generic class for runnings various tests on the g2 dataset.

    Usage: instantiate MoleculeTest with desired test settings and

    invoke its run() method on the desired formulas.

    This class will use the ASE EMT calculator by default.  You can

    create a subclass using an arbitrary calculator by overriding the

    setup_calculator method.  Most methods can be overridden to

    provide highly customized behaviour.  """

    def __init__(self, name, vacuum=6.0, exceptions=None):

        """Create a molecule test.

        The name parameter will be part of all output files generated

        by this molecule test.  If name contains a '/' character, the

        preceding part will be interpreted as a directory in which to

        put files.

        The vacuum parameter is used to set the cell size.

        A tuple of exception types can be provided which will be

        caught during a batch of calculations.  Types not specified

        will be considered fatal."""

        dir, path = os.path.split(name)

        self.dir = dir

        self.name = name

        self.vacuum = vacuum

        if exceptions is None:

            exceptions = ()

        self.exceptions = exceptions

    def setup_calculator(self, system, formula):

        """Create a new calculator.

        Default is an EMT calculator.  Most implementations will want to

        override this method."""

        raise NotImplementedError

    def setup_system(self, formula):

        """Create an Atoms object from the given formula.

        By default this will be loaded from the g2 database, setting

        the cell size by means of the molecule test's vacuum parameter."""

        system = molecule(formula)

        system.center(vacuum=self.vacuum)

        return system

    def setup(self, formula):

        """Build calculator and atoms objects.

        This will invoke the setup_calculator and setup_system methods."""

        system = self.setup_system(formula)

        calc = self.setup_calculator(system, formula)

        system.set_calculator(calc)

        return system, calc

    def get_filename(self, formula, extension='traj'):

        """Returns the filename for a test result file.

        Default format is <name>.<formula>.traj

        The test may write other files, but this filename is used as a

        flag denoting whether the calculation has been done

        already."""

        return '.'.join([self.name, formula, extension])

    def run(self, formula, system, filename):

        raise NotImplementedError

    def retrieve(self, formula, filename):

        """Retrieve results of previous calculation from file.

        Default implementation returns the total energy.

        This method should be overridden whenever the test method is

        overridden to calculate something else than the total energy."""

        raise NotImplementedError

class EnergyTest:

    def run(self, formula, system, filename):

        """Calculate energy of specified system and save to file."""

        system.get_potential_energy()

         # Won't create .bak file:

        traj = PickleTrajectory(open(filename, 'w'), 'w')

        traj.write(system)

        traj.close()

    def retrieve(self, formula, filename):

        system = read(filename)

        energy = system.get_potential_energy()

        return energy

    def calculate_atomization_energies(self, molecular_energies,

                                       atomic_energies):

        atomic_energy_dict = dict(atomic_energies)

        for formula, molecular_energy in molecular_energies:

            try:

                system = molecule(formula)

                atomic = [atomic_energy_dict[s]

                          for s in system.get_chemical_symbols()]            

                atomization_energy = molecular_energy - sum(atomic)

                yield formula, atomization_energy

            except KeyError:

                pass

class BondLengthTest:

    def run(self, formula, system, filename):

        """Calculate bond length of a dimer.

        This will calculate total energies for varying atomic

        separations close to the g2 bond length, allowing

        determination of bond length by fitting.

        """

        if len(system) != 2:

            raise ValueError('Not a dimer')

        traj = PickleTrajectory(open(filename, 'w'), 'w')

        pos = system.positions

        d = np.linalg.norm(pos[1] - pos[0])

        for x in range(-2, 3):

            system.set_distance(0, 1, d * (1.0 + x * 0.02))

            traj.write(system)

        traj.close()

    def retrieve(self, formula, filename):

        traj = PickleTrajectory(filename, 'r')

        distances = np.array([np.linalg.norm(a.positions[1] - a.positions[0])

                              for a in traj])

        energies = np.array([a.get_potential_energy() for a in traj])

        polynomial = np.polyfit(distances, energies, 2) # or maybe 3rd order?

        # With 3rd order it is not always obvious which root is right

        pderiv = np.polyder(polynomial, 1)

        d0 = np.roots(pderiv)

        e0 = np.polyval(energies, d0)

        return distances, energies, d0, e0, polynomial

class EMTTest(MoleculeTest):

    def setup_calculator(self, system, calculator):

        return EMT()

class EMTEnergyTest(EnergyTest, EMTTest):

    pass

class EMTBondLengthTest(BondLengthTest, EMTTest):

    pass

def main():

    supported_elements = 'Ni, C, Pt, Ag, H, Al, O, N, Au, Pd, Cu'.split(', ')

    formulas = [formula for formula in g1

                if np.all([symbol in supported_elements

                           for symbol

                           in molecule(formula).get_chemical_symbols()])]

    atoms = [symbol for symbol in g2_atoms if symbol in supported_elements]

    dimers = [formula for formula in formulas if len(molecule(formula)) == 2]

    name1 = 'testfiles/energy'

    name2 = 'testfiles/bond'

    test1 = BatchTest(EMTEnergyTest(name1, vacuum=3.0))

    test2 = BatchTest(EMTBondLengthTest(name2, vacuum=3.0))

    print 'Energy test'

    print '-----------'

    test1.run(formulas + atoms)

    print

    print 'Bond length test'

    print '----------------'

    test2.run(dimers)

    print

    print 'Atomization energies'

    print '--------------------'

    atomic_energies = dict(test1.collect(atoms))

    molecular_energies = dict(test1.collect(formulas))

    atomization_energies = {}

    for formula, energy in molecular_energies.iteritems():

        system = molecule(formula)

        atomic = [atomic_energies[s] for s in system.get_chemical_symbols()]

        atomization_energy = energy - sum(atomic)

        atomization_energies[formula] = atomization_energy

        print formula.rjust(10), '%.02f' % atomization_energy

    print

    print 'Bond lengths'

    print '------------'

    for formula, (d_i, e_i, d0, e0, poly) in test2.collect(dimers):

        system = molecule(formula)

        bref = np.linalg.norm(system.positions[1] - system.positions[0])

        print formula.rjust(10), '%6.3f' % d0, '  g2ref =', '%2.3f' % bref

if __name__ == '__main__':

    main()