diff --git a/tests/test_md.py b/tests/test_md.py index 340bdc1d..fcbb99ed 100644 --- a/tests/test_md.py +++ b/tests/test_md.py @@ -2,12 +2,15 @@ import pytest import numpy as np +import os +from pathlib import Path from ase import Atoms import ase.io from ase.build import bulk from ase.calculators.emt import EMT from ase.units import fs +from ase.md.logger import MDLogger from wfl.autoparallelize import autoparainfo from wfl.generate import md @@ -216,3 +219,30 @@ def test_md_abort_function(cu_slab): rng=np.random.default_rng(1)) assert len(list(atoms_traj)) < 501 + + +def test_md_attach_logger(cu_slab): + + calc = EMT() + autopara_info = autoparainfo.AutoparaInfo(num_python_subprocesses=2, num_inputs_per_python_subprocess=1, skip_failed=False) + + inputs = ConfigSet([cu_slab, cu_slab]) + outputs = OutputSpec() + + logger_kwargs = { + "logger" : MDLogger, + "logfile" : "test_log", + } + + atoms_traj = md.md(inputs, outputs, calculator=calc, integrator="Langevin", steps=300, dt=1.0, + temperature=500.0, temperature_tau=100/fs, logger_kwargs=logger_kwargs, logger_interval=1, + rng=np.random.default_rng(1), autopara_info=autopara_info,) + + atoms_traj = list(atoms_traj) + atoms_final = atoms_traj[-1] + + workdir = Path(os.getcwd()) + + assert len(atoms_traj) == 602 + assert all([Path(workdir / "test_log.item_0").is_file(), Path(workdir / "test_log.item_1").is_file()]) + diff --git a/wfl/generate/md/__init__.py b/wfl/generate/md/__init__.py index 2d59b6d9..91402e69 100644 --- a/wfl/generate/md/__init__.py +++ b/wfl/generate/md/__init__.py @@ -6,6 +6,7 @@ from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary from ase.md.verlet import VelocityVerlet from ase.md.langevin import Langevin +from ase.md.logger import MDLogger from ase.units import GPa, fs from wfl.autoparallelize import autoparallelize, autoparallelize_docstring @@ -21,8 +22,8 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBerendsen", temperature=None, temperature_tau=None, pressure=None, pressure_tau=None, compressibility_au=None, compressibility_fd_displ=0.01, traj_step_interval=1, skip_failures=True, results_prefix='last_op__md_', verbose=False, update_config_type="append", - traj_select_during_func=lambda at: True, traj_select_after_func=None, abort_check=None, rng=None, - _autopara_per_item_info=None): + traj_select_during_func=lambda at: True, traj_select_after_func=None, abort_check=None, + logger_kwargs=None, logger_interval=None, rng=None, _autopara_per_item_info=None): """runs an MD trajectory with aggresive, not necessarily physical, integrators for sampling configs. By default calculator properties for each frame stored in keys prefixed with "last_op__md_", which may be overwritten by next operation. @@ -81,6 +82,11 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere checks the MD snapshots and aborts the simulation on some condition. rng: numpy.random.Generator, default None random number generator to use (needed for pressure sampling, initial temperature, or Langevin dynamics) + logger_kwargs: dict, default None + kwargs to MDLogger to attach to each MD run, including "logfile" as string to which + config number will be appended. User defined ase.md.MDLogger derived class can be provided with "logger" as key. + logger_interval: int, default None + interval for logger _autopara_per_item_info: dict INTERNALLY used by autoparallelization framework to make runs reproducible (see wfl.autoparallelize.autoparallelize() docs) @@ -100,6 +106,10 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere else: logfile = None + if logger_kwargs is not None: + logger_constructor = logger_kwargs.pop("logger", MDLogger) + logger_logfile = logger_kwargs["logfile"] + if temperature_tau is None and (temperature is not None and not isinstance(temperature, (float, int))): raise RuntimeError(f'NVE (temperature_tau is None) can only accept temperature=float for initial T, got {type(temperature)}') @@ -128,6 +138,7 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere # get rng from autopara_per_item info if available ("rng" arg that was passed in was # already used by autoparallelization framework to set "rng" key in per-item dict) rng = _autopara_per_item_info[at_i].get("rng") + item_i = _autopara_per_item_info[at_i].get("item_i") at.calc = calculator if pressure is not None and compressibility_au is None: @@ -211,6 +222,7 @@ def process_step(interval): if not first_step_of_later_stage and cur_step % interval == 0: at.info['MD_time_fs'] = cur_step * dt at.info['MD_step'] = cur_step + at.info["MD_current_temperature"] = at.get_temperature() at_save = at_copy_save_calc_results(at, prefix=results_prefix) if traj_select_during_func(at): @@ -234,7 +246,14 @@ def process_step(interval): at.info['MD_temperature_K'] = stage_kwargs['temperature_K'] md = md_constructor(at, **stage_kwargs) + md.attach(process_step, 1, traj_step_interval) + if logger_kwargs is not None: + logger_kwargs["logfile"] = f"{logger_logfile}.item_{item_i}" + logger_kwargs["dyn"] = md + logger_kwargs["atoms"] = at + logger = logger_constructor(**logger_kwargs) + md.attach(logger, logger_interval) if stage_i > 0: first_step_of_later_stage = True