Merge pull request #362 from libAtoms/MD_Langevin_BAOAB

Add support for Langevin_BAOAB in wfl.generate.md, and support for overriding arbitrary kwargs per-config

Author: bernstei

.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__)"

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__"}

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`
 

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
 

wfl/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.3"
+__version__ = "0.3.4"