diff --git a/.github/workflows/pytests.yml b/.github/workflows/pytests.yml index 7183f703..57ebdb24 100644 --- a/.github/workflows/pytests.yml +++ b/.github/workflows/pytests.yml @@ -47,7 +47,10 @@ jobs: - name: Install latest ASE from pypi run: | echo PIP_CONSTRAINT $PIP_CONSTRAINT - python3 -m pip install ase + # avoid broken extxyz writing (3.25, fixed in 3.26) + # avoid broken optimizer.converged() (3.26) https://gitlab.com/ase/ase/-/issues/1744 + python3 -m pip install 'ase<3.25' + # echo -n "ASE VERSION " python3 -c "import ase; print(ase.__file__, ase.__version__)" python3 -c "import numpy; print('numpy version', numpy.__version__)" diff --git a/pyproject.toml b/pyproject.toml index 65873722..d2300b16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ dynamic = ["version"] reactions_iter_fit = "wfl.cli.reactions_iter_fit:cli" [tool.setuptools.packages.find] - exclude = [ "test*" ] + include = [ "wfl*" ] [tool.setuptools.dynamic] version = {attr = "wfl.__version__"} diff --git a/tests/local_scripts/complete_pytest.tin b/tests/local_scripts/complete_pytest.tin index 43b390dc..2d1f7060 100755 --- a/tests/local_scripts/complete_pytest.tin +++ b/tests/local_scripts/complete_pytest.tin @@ -3,6 +3,8 @@ module purge # module load compiler/gnu python/system python_extras/quippy lapack/mkl module load compiler/gnu python python_extras/quippy lapack/mkl +# for wfl dependencies +module load python_extras/wif module load python_extras/torch/cpu if [ -z "$WFL_PYTEST_EXPYRE_INFO" ]; then @@ -61,7 +63,7 @@ fi mkdir -p $pytest_dir -pytest -v -s --basetemp $pytest_dir ${runremote} --runslow --runperf -rxXs "$@" >> complete_pytest.tin.out 2>&1 +python3 -m pytest -v -s --basetemp $pytest_dir ${runremote} --runslow --runperf -rxXs "$@" >> complete_pytest.tin.out 2>&1 l=`egrep '^=.*(passed|failed|skipped|xfailed|error).* in ' complete_pytest.tin.out` diff --git a/tests/test_md.py b/tests/test_md.py index ea1f4755..ac2385b5 100644 --- a/tests/test_md.py +++ b/tests/test_md.py @@ -1,9 +1,11 @@ from pytest import approx import pytest -import numpy as np import os from pathlib import Path +import json + +import numpy as np from ase import Atoms import ase.io @@ -17,6 +19,10 @@ from wfl.configset import ConfigSet, OutputSpec from wfl.generate.md.abort import AbortOnCollision, AbortOnLowEnergy +try: + from wif.Langevin_BAOAB import Langevin_BAOAB +except ImportError: + Langevin_BAOAB = None def select_every_10_steps_for_tests_during(at): return at.info.get("MD_step", 1) % 10 == 0 @@ -52,7 +58,6 @@ def test_NVE(cu_slab): temperature=500.0, rng=np.random.default_rng(1)) atoms_traj = list(atoms_traj) - atoms_final = atoms_traj[-1] assert len(atoms_traj) == 301 @@ -68,27 +73,131 @@ def test_NVT_const_T(cu_slab): temperature=500.0, temperature_tau=30.0, rng=np.random.default_rng(1)) atoms_traj = list(atoms_traj) - atoms_final = atoms_traj[-1] assert len(atoms_traj) == 301 assert all([at.info['MD_temperature_K'] == 500.0 for at in atoms_traj]) + assert np.all(atoms_traj[0].cell == atoms_traj[-1].cell) def test_NVT_Langevin_const_T(cu_slab): - calc = EMT() inputs = ConfigSet(cu_slab) outputs = OutputSpec() atoms_traj = md.md(inputs, outputs, calculator=calc, integrator="Langevin", steps=300, dt=1.0, - temperature=500.0, temperature_tau=100/fs, rng=np.random.default_rng(1)) + temperature=500.0, temperature_tau=100/fs, rng=np.random.default_rng(1)) atoms_traj = list(atoms_traj) - atoms_final = atoms_traj[-1] assert len(atoms_traj) == 301 assert all([at.info['MD_temperature_K'] == 500.0 for at in atoms_traj]) + assert np.all(atoms_traj[0].cell == atoms_traj[-1].cell) + + +def test_NPT_Langevin_fail(cu_slab): + calc = EMT() + + inputs = ConfigSet(cu_slab) + outputs = OutputSpec() + + with pytest.raises(ValueError): + atoms_traj = md.md(inputs, outputs, calculator=calc, integrator="Langevin", steps=300, dt=1.0, + temperature=500.0, temperature_tau=100/fs, pressure=0.0, + rng=np.random.default_rng(1)) + + +def test_NPT_Berendsen_hydro_F_fail(cu_slab): + calc = EMT() + + inputs = ConfigSet(cu_slab) + outputs = OutputSpec() + + with pytest.raises(ValueError): + atoms_traj = md.md(inputs, outputs, calculator=calc, integrator="Berendsen", steps=300, dt=1.0, + temperature=500.0, temperature_tau=100/fs, pressure=0.0, hydrostatic=False, + rng=np.random.default_rng(1)) + + +def test_NPT_Berendsen_NPH_fail(cu_slab): + calc = EMT() + + inputs = ConfigSet(cu_slab) + outputs = OutputSpec() + + with pytest.raises(ValueError): + atoms_traj = md.md(inputs, outputs, calculator=calc, integrator="Berendsen", steps=300, dt=1.0, + pressure=0.0, + rng=np.random.default_rng(1)) + + +def test_NPT_Berendsen(cu_slab): + calc = EMT() + + inputs = ConfigSet(cu_slab) + outputs = OutputSpec() + + atoms_traj = md.md(inputs, outputs, calculator=calc, integrator="Berendsen", steps=300, dt=1.0, + temperature=500.0, temperature_tau=100/fs, pressure=0.0, + rng=np.random.default_rng(1)) + + atoms_traj = list(atoms_traj) + print("I cell", atoms_traj[0].cell) + print("F cell", atoms_traj[1].cell) + + assert len(atoms_traj) == 301 + assert all([at.info['MD_temperature_K'] == 500.0 for at in atoms_traj]) + assert np.any(atoms_traj[0].cell != atoms_traj[-1].cell) + + cell_f = atoms_traj[0].cell[0, 0] / atoms_traj[-1].cell[0, 0] + assert np.allclose(atoms_traj[0].cell, atoms_traj[-1].cell * cell_f) + + +@pytest.mark.skipif(Langevin_BAOAB is None, reason="No Langevin_BAOAB available") +def test_NPT_Langevin_BAOAB(cu_slab): + calc = EMT() + + inputs = ConfigSet(cu_slab) + outputs = OutputSpec() + + atoms_traj = md.md(inputs, outputs, calculator=calc, integrator="Langevin_BAOAB", steps=300, dt=1.0, + temperature=500.0, temperature_tau=100/fs, pressure=0.0, + rng=np.random.default_rng(1)) + + atoms_traj = list(atoms_traj) + print("I cell", atoms_traj[0].cell) + print("F cell", atoms_traj[1].cell) + + assert len(atoms_traj) == 301 + assert all([at.info['MD_temperature_K'] == 500.0 for at in atoms_traj]) + assert np.any(atoms_traj[0].cell != atoms_traj[-1].cell) + + cell_f = atoms_traj[0].cell[0, 0] / atoms_traj[-1].cell[0, 0] + assert np.allclose(atoms_traj[0].cell, atoms_traj[-1].cell * cell_f) + + +@pytest.mark.skipif(Langevin_BAOAB is None, reason="No Langevin_BAOAB available") +def test_NPT_Langevin_BAOAB_hydro_F(cu_slab): + calc = EMT() + + inputs = ConfigSet(cu_slab) + outputs = OutputSpec() + + atoms_traj = md.md(inputs, outputs, calculator=calc, integrator="Langevin_BAOAB", steps=300, dt=1.0, + temperature=500.0, temperature_tau=100/fs, pressure=0.0, hydrostatic=False, + rng=np.random.default_rng(1)) + + atoms_traj = list(atoms_traj) + print("I cell", atoms_traj[0].cell) + print("F cell", atoms_traj[1].cell) + + assert len(atoms_traj) == 301 + assert all([at.info['MD_temperature_K'] == 500.0 for at in atoms_traj]) + assert np.any(atoms_traj[0].cell != atoms_traj[-1].cell) + + cell_f = atoms_traj[0].cell[0, 0] / atoms_traj[-1].cell[0, 0] + assert not np.allclose(atoms_traj[0].cell, atoms_traj[-1].cell * cell_f) + def test_NVT_Langevin_const_T_per_config(cu_slab): @@ -99,7 +208,7 @@ def test_NVT_Langevin_const_T_per_config(cu_slab): outputs = OutputSpec() for at_i, at in enumerate(inputs): - at.info["WFL_MD_TEMPERATURE"] = 500 + at_i * 100 + at.info["WFL_MD_KWARGS"] = json.dumps({'temperature': 500 + at_i * 100}) n_steps = 30 diff --git a/wfl/__version__.py b/wfl/__version__.py index e19434e2..334b8995 100644 --- a/wfl/__version__.py +++ b/wfl/__version__.py @@ -1 +1 @@ -__version__ = "0.3.3" +__version__ = "0.3.4" diff --git a/wfl/generate/md/__init__.py b/wfl/generate/md/__init__.py index 479d829f..7e77e88b 100644 --- a/wfl/generate/md/__init__.py +++ b/wfl/generate/md/__init__.py @@ -1,4 +1,5 @@ import sys +import json import numpy as np from ase.md.nptberendsen import NPTBerendsen @@ -7,26 +8,31 @@ 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 ase.units import fs + +try: + from wif.Langevin_BAOAB import Langevin_BAOAB +except ImportError: + LangevinBAOAB = None from wfl.autoparallelize import autoparallelize, autoparallelize_docstring from wfl.utils.save_calc_results import at_copy_save_calc_results from wfl.utils.misc import atoms_to_list from wfl.utils.parallel import construct_calculator_picklesafe -from wfl.utils.pressure import sample_pressure from ..utils import save_config_type - -bar = 1.0e-4 * GPa +from .utils import _get_temperature, _get_pressure -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, +def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="Berendsen", temperature=None, temperature_tau=None, + pressure=None, hydrostatic=True, pressure_tau=None, pressure_mass_factor=1, 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, - logger_kwargs=None, logger_interval=None, rng=None, _autopara_per_item_info=None): + logger_interval=0, logger_kwargs=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. + Most keyword args (all except `calculator`, `steps`, `dt`, and `logger_*`) can be + overridden per-config by a json-encoded dict in each `atoms.info["WFL_MD_KWARGS"]` Parameters ---------- @@ -36,8 +42,9 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere ASE calculator or routine to call to create calculator dt: float time step (fs) - integrator: str, default "NVTBerendsen" - Select integrator. Default is Berendsen but also langevin can be used + integrator: "Berendsen" or "Langevin" or "Langevin_BAOAB", default "Berendsen" + MD time integrator. Berendsen and Langevin fall back to VelocityVerlet (NVE) if temperature_tau is None, + while Langevin_BAOAB will do NPH in that case steps: int number of steps temperature: float or (float, float, [int]]), or list of dicts default None @@ -46,16 +53,18 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere - tuple/list of float, float, [int=10]: T_init, T_final, and optional number of stages for ramp - [ {'T_i': float, 'T_f' : float, 'traj_frac' : flot, 'n_stages': int=10}, ... ] list of stages, each one a ramp, with duration defined as fraction of total number of steps - Overridden by `atoms.info["WFL_MD_TEMPERATURE"]` temperature_tau: float, default None Time scale for thermostat (fs). Directly used for Berendsen integrator, or as 1/friction for Langevin integrator. pressure: None / float / tuple Applied pressure distribution (GPa) as parsed by wfl.utils.pressure.sample_pressure(). Enables Berendsen constant P volume rescaling. - Overridden by `atoms.info["WFL_MD_PRESSURE"]` + hydrostatic: bool, default True + Allow only hydrostatic strain for variable cell (pressure not None) pressure_tau: float, default None time scale for Berendsen constant P volume rescaling (fs) ignored if pressure is None, defaults to 3*temperature_tau + pressure_mass_factor: float, default 1 + factor to multiply by default pressure mass heuristic compressibility_au: float, default None compressibility, if available, for NPTBerendsen compressibility_fd_displ: float, default 0.01 @@ -69,8 +78,7 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere Will overwrite any other properties that start with same "__", so that by default only last op's properties will be stored. verbose: bool, default False - verbose output - MD logs are not printed unless this is True + verbose output (MD logging is handled by setting logger_interval > 0) update_config_type: ["append" | "overwrite" | False], default "append" whether/how to add 'MD' to at.info['config_type'] traj_select_during_func: func(Atoms), default func(Atoms) -> bool=True @@ -84,12 +92,12 @@ 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_interval: int, default 0 + Enable logging at this interval if > 0 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. If logfile is "-", stdout will be used, and config number will be prepended to each outout line. User defined ase.md.MDLogger derived class can be provided with "logger" as key. - logger_interval: int, default None - Enable logging at this interval _autopara_per_item_info: dict INTERNALLY used by autoparallelization framework to make runs reproducible (see wfl.autoparallelize.autoparallelize() docs) @@ -98,216 +106,236 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere ------- list(Atoms) trajectories """ - assert integrator in ["NVTBerendsen", "Langevin"] - - calculator = construct_calculator_picklesafe(calculator) - - all_trajs = [] - if verbose: - logfile = '-' + return _sample_autopara_wrappable_kwargs(atoms, calculator, steps, dt, + integrator=integrator, + temperature=temperature, + temperature_tau=temperature_tau, + pressure=pressure, + hydrostatic=hydrostatic, + pressure_tau=pressure_tau, + pressure_mass_factor=pressure_mass_factor, + compressibility_au=compressibility_au, + compressibility_fd_displ=compressibility_fd_displ, + traj_step_interval=traj_step_interval, + skip_failures=skip_failures, + results_prefix=results_prefix, + verbose=verbose, + update_config_type=update_config_type, + traj_select_during_func=traj_select_during_func, + traj_select_after_func=traj_select_after_func, + abort_check=abort_check, + logger_interval=logger_interval, + logger_kwargs=logger_kwargs, + rng=rng, + _autopara_per_item_info=_autopara_per_item_info) + + +def _sample_autopara_wrappable_single(at, at_i, calculator, steps, dt, logger_interval, logger_constructor, logger_logfile, logger_kwargs, integrator="Berendsen", temperature=None, temperature_tau=None, + pressure=None, hydrostatic=True, pressure_tau=None, pressure_mass_factor=1, 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): + # 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 + temperature_use = _get_temperature(temperature, at, steps) + pressure_use, compressibility_au_use = _get_pressure(pressure, compressibility_au, compressibility_fd_displ, at, rng) + + if temperature_use is not None: + # set initial temperature + assert rng is not None + MaxwellBoltzmannDistribution(at, temperature_K=temperature_use[0]['T_i'], force_temp=True, communicator=None, rng=rng) + Stationary(at, preserve_temperature=True) + + stage_kwargs = {'timestep': dt * fs} + + if integrator == "Langevin_BAOAB": + md_constructor = Langevin_BAOAB else: - logfile = None - - if logger_interval is not None and logger_interval > 0: - if logger_kwargs is None: - logger_kwargs = {} - logger_constructor = logger_kwargs.pop("logger", MDLogger) - logger_logfile = logger_kwargs.get("logfile", "-") - - def _get_pressure(atoms): - return atoms.info.get("WFL_MD_PRESSURE", pressure) - - def _get_temperature(atoms): - temperature_use = atoms.info.get("WFL_MD_TEMPERATURE", temperature) - - if temperature_tau is None and (temperature_use is not None and not isinstance(temperature_use, (float, int, np.floating, np.integer))): - raise RuntimeError(f'NVE (temperature_tau is None) can only accept temperature=float for initial T, got {type(temperature_use)}') - - if temperature_use is not None: - # assume that dicts are already in temperature profile format - if not isinstance(temperature_use, dict): - try: - # check if it's a list, tuple, etc - len(temperature_use) - except: - # number into a list - temperature_use = [temperature_use] - if not isinstance(temperature_use[0], dict): - # create a stage dict from a constant or ramp - t_stage_data = temperature_use - # start with constant - t_stage = {'T_i': t_stage_data[0], 'T_f': t_stage_data[0], 'traj_frac': 1.0, 'n_stages': 10, 'steps': steps} - if len(t_stage_data) >= 2: - # set different final T for ramp - t_stage['T_f'] = t_stage_data[1] - if len(t_stage_data) >= 3: - # set number of stages - t_stage['n_stages'] = t_stage_data[2] - temperature_use = [t_stage] + if temperature_tau is None: + md_constructor = VelocityVerlet + else: + if integrator == 'Langevin': + md_constructor = Langevin + elif integrator == 'Berendsen': + if pressure_use is None: + md_constructor = NVTBerendsen + else: + md_constructor = NPTBerendsen else: - for t_stage in temperature_use: - if 'n_stages' not in t_stage: - t_stage['n_stages'] = 10 + raise ValueError(f'Unkown integrator {integrator}') + + # pressure arguments, relatively simple because there are no stages + if pressure_use is not None: + if integrator == 'Langevin_BAOAB': + stage_kwargs['externalstress'] = pressure_use + stage_kwargs['P_tau'] = pressure_tau * fs if pressure_tau is not None else temperature_tau * fs * 3 + stage_kwargs['P_mass_factor'] = pressure_mass_factor + stage_kwargs['hydrostatic'] = hydrostatic + elif integrator == 'Berendsen': + if temperature_tau is None: + raise ValueError('integrator Berendsen got pressure but no temperature') + if not hydrostatic: + raise ValueError('integrator Berendsen got hydrostatic False') + stage_kwargs['pressure_au'] = pressure_use + stage_kwargs['compressibility_au'] = compressibility_au_use + stage_kwargs['taup'] = pressure_tau * fs if pressure_tau is not None else temperature_tau * fs * 3 + else: + raise ValueError(f'Only Langevin_BAOAB and Berendsen integrator support pressure, got {integrator}') - return temperature_use + # temperature args except actual temperature, which is set below with stages + if temperature_tau is not None: + if integrator == 'Langevin_BAOAB': + stage_kwargs['T_tau'] = temperature_tau * fs + assert rng is not None + stage_kwargs["rng"] = rng + elif integrator == 'Berendsen': + stage_kwargs['taut'] = temperature_tau * fs + elif integrator == 'Langevin': + stage_kwargs["friction"] = 1 / (temperature_tau * fs) + assert rng is not None + stage_kwargs["rng"] = rng + else: + assert False, f'Unknown integrator {integrator}' - for at_i, at in enumerate(atoms_to_list(atoms)): - # 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") - - temperature_use = _get_temperature(at) - pressure_use = _get_pressure(at) - - at.calc = calculator - if pressure_use is not None: - pressure_use = sample_pressure(pressure_use, at, rng=rng) - at.info['MD_pressure_GPa'] = pressure_use - # convert to ASE internal units - pressure_use *= GPa - if compressibility_au is None: - E0 = at.get_potential_energy() - c0 = at.get_cell() - at.set_cell(c0 * (1.0 + compressibility_fd_displ), scale_atoms=True) - Ep = at.get_potential_energy() - at.set_cell(c0 * (1.0 - compressibility_fd_displ), scale_atoms=True) - Em = at.get_potential_energy() - at.set_cell(c0, scale_atoms=True) - d2E_dF2 = (Ep + Em - 2.0 * E0) / (compressibility_fd_displ ** 2) - compressibility_au_use = at.get_volume() / d2E_dF2 + if temperature_tau is None: + # relatively simple, one stage + all_stage_kwargs = [stage_kwargs.copy()] + all_run_kwargs = [{'steps': steps}] + else: + # set up temperature stages + all_stage_kwargs = [] + all_run_kwargs = [] + + for t_stage_i, t_stage in enumerate(temperature_use): + stage_steps = t_stage['traj_frac'] * steps + + if t_stage['T_f'] == t_stage['T_i']: + # constant T + stage_kwargs['temperature_K'] = t_stage['T_i'] + all_stage_kwargs.append(stage_kwargs.copy()) + all_run_kwargs.append({'steps': int(np.round(stage_steps))}) else: - compressibility_au_use = compressibility_au - - if temperature_use is not None: - # set initial temperature - assert rng is not None - MaxwellBoltzmannDistribution(at, temperature_K=temperature_use[0]['T_i'], force_temp=True, communicator=None, rng=rng) - Stationary(at, preserve_temperature=True) + # ramp + for T in np.linspace(t_stage['T_i'], t_stage['T_f'], t_stage['n_stages']): + stage_kwargs['temperature_K'] = T + all_stage_kwargs.append(stage_kwargs.copy()) + substage_steps = int(np.round(stage_steps / t_stage['n_stages'])) + all_run_kwargs.extend([{'steps': substage_steps}] * t_stage['n_stages']) - stage_kwargs = {'timestep': dt * fs, 'logfile': logfile} + traj = [] + cur_step = 1 + first_step_of_later_stage = False - if temperature_tau is None: - # NVE - if pressure_use is not None: - raise RuntimeError(f'Got pressure {pressure_use} but no active thermostat temperature_tau={temperature_tau}. Can only do NPT, not NPH, dynamics') - md_constructor = VelocityVerlet - # one stage, simple - all_stage_kwargs = [stage_kwargs.copy()] - all_run_kwargs = [{'steps': steps}] + def process_step(interval): + nonlocal cur_step, first_step_of_later_stage - else: - # NVT or NPT - all_stage_kwargs = [] - all_run_kwargs = [] - - if pressure_use is not None: - md_constructor = NPTBerendsen - stage_kwargs['pressure_au'] = pressure_use - stage_kwargs['compressibility_au'] = compressibility_au_use - stage_kwargs['taut'] = temperature_tau * fs - stage_kwargs['taup'] = pressure_tau * fs if pressure_tau is not None else temperature_tau * fs * 3 - else: - if integrator == "NVTBerendsen": - md_constructor = NVTBerendsen - stage_kwargs['taut'] = temperature_tau * fs + 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) - elif integrator == "Langevin": - md_constructor = Langevin - stage_kwargs["friction"] = 1 / (temperature_tau * fs) - assert rng is not None - stage_kwargs["rng"] = rng + if traj_select_during_func(at): + traj.append(at_save) - for t_stage_i, t_stage in enumerate(temperature_use): - stage_steps = t_stage['traj_frac'] * steps + if abort_check is not None: + if abort_check.stop(at): + raise RuntimeError(f"MD was stopped by the MD checker function {abort_check.__class__.__name__}") - if t_stage['T_f'] == t_stage['T_i']: - # constant T - stage_kwargs['temperature_K'] = t_stage['T_i'] - all_stage_kwargs.append(stage_kwargs.copy()) - all_run_kwargs.append({'steps': int(np.round(stage_steps))}) - else: - # ramp - for T in np.linspace(t_stage['T_i'], t_stage['T_f'], t_stage['n_stages']): - stage_kwargs['temperature_K'] = T - all_stage_kwargs.append(stage_kwargs.copy()) - substage_steps = int(np.round(stage_steps / t_stage['n_stages'])) - all_run_kwargs.extend([{'steps': substage_steps}] * t_stage['n_stages']) - - traj = [] - cur_step = 1 first_step_of_later_stage = False + cur_step += 1 + + for stage_i, (stage_kwargs, run_kwargs) in enumerate(zip(all_stage_kwargs, all_run_kwargs)): + if verbose: + print('run stage', stage_kwargs, run_kwargs) - def process_step(interval): - nonlocal cur_step, first_step_of_later_stage + # avoid double counting of steps at end of each stage and beginning of next + cur_step -= 1 - 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 temperature_tau is not None: + at.info['MD_temperature_K'] = stage_kwargs['temperature_K'] - if traj_select_during_func(at): - traj.append(at_save) + md = md_constructor(at, **stage_kwargs) - if abort_check is not None: - if abort_check.stop(at): - raise RuntimeError(f"MD was stopped by the MD checker function {abort_check.__class__.__name__}") + md.attach(process_step, 1, traj_step_interval) + if logger_interval > 0: + if logger_logfile == "-": + logger_kwargs["logfile"] = "-" + else: + logger_kwargs["logfile"] = f"{logger_logfile}.config_{item_i}" + logger_kwargs["dyn"] = md + logger_kwargs["atoms"] = at + logger = logger_constructor(**logger_kwargs) + if logger_kwargs["logfile"] == "-": + # add prefix to each line + logger.hdr = f"config {item_i} " + logger.hdr + logger.fmt = f"config {item_i} " + logger.fmt + md.attach(logger, logger_interval) + + if stage_i > 0: + first_step_of_later_stage = True + + try: + md.run(**run_kwargs) + except Exception as exc: + if skip_failures: + sys.stderr.write(f'MD failed with exception \'{exc}\'\n') + sys.stderr.flush() + break + else: + raise - first_step_of_later_stage = False - cur_step += 1 + if len(traj) == 0 or traj[-1] != at: + if traj_select_during_func(at): + at.info['MD_time_fs'] = cur_step * dt + traj.append(at_copy_save_calc_results(at, prefix=results_prefix)) - for stage_i, (stage_kwargs, run_kwargs) in enumerate(zip(all_stage_kwargs, all_run_kwargs)): - if verbose: - print('run stage', stage_kwargs, run_kwargs) + if traj_select_after_func is not None: + traj = traj_select_after_func(traj) - # avoid double counting of steps at end of each stage and beginning of next - cur_step -= 1 + for at in traj: + save_config_type(at, update_config_type, 'MD') - if temperature_tau is not None: - at.info['MD_temperature_K'] = stage_kwargs['temperature_K'] + return traj - md = md_constructor(at, **stage_kwargs) - md.attach(process_step, 1, traj_step_interval) - if logger_interval is not None and logger_interval > 0: - if logger_logfile == "-": - logger_kwargs["logfile"] = "-" - else: - logger_kwargs["logfile"] = f"{logger_logfile}.config_{item_i}" - logger_kwargs["dyn"] = md - logger_kwargs["atoms"] = at - logger = logger_constructor(**logger_kwargs) - if logger_logfile == "-": - # add prefix to each line - logger.hdr = f"config {item_i} " + logger.hdr - logger.fmt = f"config {item_i} " + logger.fmt - md.attach(logger, logger_interval) - - if stage_i > 0: - first_step_of_later_stage = True - - try: - md.run(**run_kwargs) - except Exception as exc: - if skip_failures: - sys.stderr.write(f'MD failed with exception \'{exc}\'\n') - sys.stderr.flush() - break - else: - raise +def _sample_autopara_wrappable_kwargs(atoms, calculator, steps, dt, **kwargs): + assert kwargs.get("integrator", "Berendsen") in ["Berendsen", "Langevin", "Langevin_BAOAB"] - if len(traj) == 0 or traj[-1] != at: - if traj_select_during_func(at): - at.info['MD_time_fs'] = cur_step * dt - traj.append(at_copy_save_calc_results(at, prefix=results_prefix)) + calculator = construct_calculator_picklesafe(calculator) - if traj_select_after_func is not None: - traj = traj_select_after_func(traj) + logger_interval = kwargs.pop("logger_interval", 0) + logger_kwargs = kwargs.pop("logger_kwargs", None) + if logger_kwargs is None: + logger_kwargs = {} + else: + logger_kwargs = logger_kwargs.copy() + logger_constructor = None + logger_logfile = None + if logger_interval > 0: + logger_constructor = logger_kwargs.pop("logger", MDLogger) + logger_logfile = logger_kwargs.pop("logfile", "-") + + kwargs_at = kwargs.copy() + kwargs_orig = {} - for at in traj: - save_config_type(at, update_config_type, 'MD') + all_trajs = [] + for at_i, at in enumerate(atoms_to_list(atoms)): + # reset kwargs_at to orig values + kwargs_at.update(kwargs_orig) + # set overriding values + for k, v in json.loads(at.info.get("WFL_MD_KWARGS", "{}")).items(): + if k not in kwargs_orig: + kwargs_orig[k] = kwargs[k] + kwargs_at[k] = v + + # do operation with overridden values + traj = _sample_autopara_wrappable_single(at, at_i, calculator, steps, dt, logger_interval, logger_constructor, logger_logfile, logger_kwargs, **kwargs_at) all_trajs.append(traj) diff --git a/wfl/generate/md/utils.py b/wfl/generate/md/utils.py new file mode 100644 index 00000000..f3cba840 --- /dev/null +++ b/wfl/generate/md/utils.py @@ -0,0 +1,56 @@ +import numpy as np + +from ase.units import GPa + +from wfl.utils.pressure import sample_pressure + +def _get_temperature(temperature_use, temperature_tau, steps): + if temperature_tau is None and (temperature_use is not None and not isinstance(temperature_use, (float, int, np.floating, np.integer))): + raise RuntimeError(f'NVE (temperature_tau is None) can only accept temperature=float for initial T, got {type(temperature_use)}') + + if temperature_use is not None: + # assume that dicts are already in temperature profile format + if not isinstance(temperature_use, dict): + try: + # check if it's a list, tuple, etc + len(temperature_use) + except TypeError: + # number into a list + temperature_use = [temperature_use] + if not isinstance(temperature_use[0], dict): + # create a stage dict from a constant or ramp + t_stage_data = temperature_use + # start with constant + t_stage = {'T_i': t_stage_data[0], 'T_f': t_stage_data[0], 'traj_frac': 1.0, 'n_stages': 10, 'steps': steps} + if len(t_stage_data) >= 2: + # set different final T for ramp + t_stage['T_f'] = t_stage_data[1] + if len(t_stage_data) >= 3: + # set number of stages + t_stage['n_stages'] = t_stage_data[2] + temperature_use = [t_stage] + else: + for t_stage in temperature_use: + if 'n_stages' not in t_stage: + t_stage['n_stages'] = 10 + + return temperature_use + +def _get_pressure(pressure_use, compressibility_au_use, compressibility_fd_displ, at, rng): + if pressure_use is not None: + pressure_use = sample_pressure(pressure_use, at, rng=rng) + at.info['MD_pressure_GPa'] = pressure_use + # convert to ASE internal units + pressure_use *= GPa + if compressibility_au_use is None: + E0 = at.get_potential_energy() + c0 = at.get_cell() + at.set_cell(c0 * (1.0 + compressibility_fd_displ), scale_atoms=True) + Ep = at.get_potential_energy() + at.set_cell(c0 * (1.0 - compressibility_fd_displ), scale_atoms=True) + Em = at.get_potential_energy() + at.set_cell(c0, scale_atoms=True) + d2E_dF2 = (Ep + Em - 2.0 * E0) / (compressibility_fd_displ ** 2) + compressibility_au_use = at.get_volume() / d2E_dF2 + + return pressure_use, compressibility_au_use