Source code for ase.calculators.loggingcalc

"""
Provides LoggingCalculator class to wrap a Calculator and record
number of enery and force calls
"""

import json
import logging
import time
from typing import Dict, Any

import numpy as np

from ase.calculators.calculator import Calculator, all_properties

logger = logging.getLogger(__name__)


[docs]class LoggingCalculator(Calculator): """Calculator wrapper to record and plot history of energy and function evaluations """ implemented_properties = all_properties default_parameters: Dict[str, Any] = {} name = 'LoggingCalculator' 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, jsonfile=None, dumpjson=False): Calculator.__init__(self) self.calculator = calculator self.fmax = {} self.walltime = {} self.energy_evals = {} self.energy_count = {} self.set_label('(none)') if jsonfile is not None: self.read_json(jsonfile) self.dumpjson = dumpjson def calculate(self, atoms, properties, system_changes): Calculator.calculate(self, atoms, properties, system_changes) if isinstance(self.calculator, Calculator): results = [self.calculator.get_property(prop, atoms) for prop in properties] else: results = [] for prop in properties: method_name = self.property_to_method_name[prop] method = getattr(self.calculator, method_name) results.append(method(atoms)) if 'energy' in properties or 'energies' in properties: self.energy_evals.setdefault(self.label, 0) self.energy_evals[self.label] += 1 try: energy = results[properties.index('energy')] except IndexError: energy = sum(results[properties.index('energies')]) logger.info('energy call count=%d energy=%.3f', self.energy_evals[self.label], energy) self.results = dict(zip(properties, results)) if 'forces' in self.results: fmax = self.fmax.setdefault(self.label, []) walltime = self.walltime.setdefault(self.label, []) forces = self.results['forces'].copy() energy_count = self.energy_count.setdefault(self.label, []) energy_evals = self.energy_evals.setdefault(self.label, 0) energy_count.append(energy_evals) for constraint in atoms.constraints: constraint.adjust_forces(atoms, forces) fmax.append(abs(forces).max()) walltime.append(time.time()) logger.info('force call fmax=%.3f', fmax[-1]) if self.dumpjson: self.write_json('dump.json') def write_json(self, filename): with open(filename, 'w') as fd: json.dump({'fmax': self.fmax, 'walltime': self.walltime, 'energy_evals': self.energy_evals, 'energy_count': self.energy_count}, fd) def read_json(self, filename, append=False, label=None): with open(filename, 'r') as fd: dct = json.load(fd) labels = dct['fmax'].keys() if label is not None and len(labels) == 1: for key in ('fmax', 'walltime', 'energy_evals', 'energy_count'): dct[key][label] = dct[key][labels[0]] del dct[key][labels[0]] if not append: self.fmax = {} self.walltime = {} self.energy_evals = {} self.energy_count = {} self.fmax.update(dct['fmax']) self.walltime.update(dct['walltime']) self.energy_evals.update(dct['energy_evals']) self.energy_count.update(dct['energy_count']) def tabulate(self): fmt1 = '%-10s %10s %10s %8s' title = fmt1 % ('Label', '# Force', '# Energy', 'Walltime/s') print(title) print('-' * len(title)) fmt2 = '%-10s %10d %10d %8.2f' for label in sorted(self.fmax.keys()): print(fmt2 % (label, len(self.fmax[label]), len(self.energy_count[label]), self.walltime[label][-1] - self.walltime[label][0])) def plot(self, fmaxlim=(1e-2, 1e2), forces=True, energy=True, walltime=True, markers=None, labels=None, **kwargs): import matplotlib.pyplot as plt if markers is None: markers = [c + s for c in ['r', 'g', 'b', 'c', 'm', 'y', 'k'] for s in ['.-', '.--']] nsub = sum([forces, energy, walltime]) nplot = 0 if labels is not None: fmax_values = [v for (k, v) in sorted(zip(self.fmax.keys(), self.fmax.values()))] self.fmax = dict(zip(labels, fmax_values)) energy_count_values = [v for (k, v) in sorted(zip(self.energy_count.keys(), self.energy_count.values()))] self.energy_count = dict(zip(labels, energy_count_values)) walltime_values = [v for (k, v) in sorted(zip(self.walltime.keys(), self.walltime.values()))] self.walltime = dict(zip(labels, walltime_values)) if forces: nplot += 1 plt.subplot(nsub, 1, nplot) for label, color in zip(sorted(self.fmax.keys()), markers): fmax = np.array(self.fmax[label]) idx = np.arange(len(fmax)) plt.semilogy(idx, fmax, color, label=label, **kwargs) plt.xlabel('Number of force evaluations') plt.ylabel('Maximum force / eV/A') plt.ylim(*fmaxlim) plt.legend() if energy: nplot += 1 plt.subplot(nsub, 1, nplot) for label, color in zip(sorted(self.energy_count.keys()), markers): energy_count = np.array(self.energy_count[label]) fmax = np.array(self.fmax[label]) plt.semilogy(energy_count, fmax, color, label=label, **kwargs) plt.xlabel('Number of energy evaluations') plt.ylabel('Maximum force / eV/A') plt.ylim(*fmaxlim) plt.legend() if walltime: nplot += 1 plt.subplot(nsub, 1, nplot) for label, color in zip(sorted(self.walltime.keys()), markers): walltime = np.array(self.walltime[label]) fmax = np.array(self.fmax[label]) walltime -= walltime[0] plt.semilogy(walltime, fmax, color, label=label, **kwargs) plt.xlabel('Walltime / s') plt.ylabel('Maximum force / eV/A') plt.ylim(*fmaxlim) plt.legend() plt.subplots_adjust(hspace=0.33)