Improve MDLogger Design and Flexiblity

Improve MDLogger Design and Flexibility

Current State

The MDLogger currently handles basic MD quantities (time, energies, temperature) in a fixed format. Logging additional quantities requires modifying the MD classes themselves, as demonstrated by the recent need to log conserved energy in the Bussi thermostat implementation.

Previous Attempt (PR !3498 (closed))

Initially proposed adding an extra_fields attribute to Dynamics classes, allowing them to pass additional quantities to the logger. While this solved the immediate need for Bussi thermostat, discussions with @yuzie007 highlighted that this approach still tightly couples logging concerns with dynamics implementation.

Proposed Solution

After discussion, a better approach emerged:

  1. Have the MDLogger rely entirely on a new attribute fields, which contain names, callable and format of tracked quantities. A new method add_fields can be used to fully customize loggers.

    logger = MDLogger("md.log", mode="w")
logger.add_fields("Econs[eV]", lambda: atoms.get_total_energy() - dyn.transferred_energy)

  2. As a result, MDLogger does not take dyn or atoms directly as parameter anymore. A convenience method is added add_md_fields to add the standard MD quantities, this is directly called in the __init__ of MolecularDynamics.

  3. Move the header writing part outside of the __init__, ideally somewhere related to the start of the Dynamics, avoiding rewrite on restart. Consequently MDLogger should implement a method write_header.

  4. At that point, MDLogger could be nicely deprecated and link to a new class Logger, ideally moved somewhere else, so that any kind of simulations could use it (optimizers, neb, etc...)

This leads to key improvements:

  • Separate logging from dynamics implementation.
  • Allow field addition to track any quantities.
  • Let users customize logging without subclassing anything.

Please find an experimental (working) implementation of the above in this branch

Assuming correct link and depreciation notes breaking changes can be greatly limited. The only remaining would be for expert users who create MDLogger and attach it manually.

Example usage in the new framework:

# When using the `logfile` parameter, nothing changes

dyn = Bussi(
    atoms,
    ...,
    logfile="dyn.log"
)

# If attaching manually, the Logger has to be initialized to obtain the desired behavior

from ase.logger import Logger

dyn = Bussi(atoms, ...)

logger = Logger("dyn.log")
logger.add_md_fields(dyn)

dyn.attach(logger)

logger.add_fields("Econs[eV]", lambda: dyn.atoms.get_total_energy() - dyn.transferred_energy, "10.3f")
logger.write_header()

dyn.run(20)

Related Links

  • !3498 (closed) Initial attempt to revamp MDLogger with relevant discussions.
  • #1344 Optimizer logging initial consideration
  • #1439 MD restart functionality discussion (relevant for header writing)

/cc @yuzie007 @askhl

Edited by Tom Demeyere