File: test_md_logger.py

package info (click to toggle)
python-ase 3.21.1-2
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 13,936 kB
  • sloc: python: 122,428; xml: 946; makefile: 111; javascript: 47
file content (84 lines) | stat: -rw-r--r-- 2,609 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
"""Test to ensure that md logger and trajectory contain same data"""
import numpy as np
import pytest

from ase.optimize import FIRE, BFGS
from ase.data import s22
from ase.calculators.tip3p import TIP3P
from ase.constraints import FixBondLengths
from ase.md.verlet import VelocityVerlet
from ase.md.langevin import Langevin
from ase.io import Trajectory
import ase.units as u

md_cls_and_kwargs = [
    (VelocityVerlet, {}),
    (Langevin, {"temperature_K": 300, "friction": 0.02}),
]

# prepare atoms object for testing
dimer = s22.create_s22_system("Water_dimer")

calc = TIP3P(rc=9.0)
dimer.constraints = FixBondLengths(
    [(3 * i + j, 3 * i + (j + 1) % 3) for i in range(2) for j in [0, 1, 2]]
)


def fmax(forces):
    return np.sqrt((forces ** 2).sum(axis=1).max())


@pytest.mark.parametrize('cls', [FIRE, BFGS])
def test_optimizer(cls, atoms=dimer, calc=calc,
                   logfile="opt.log", trajectory="opt.traj"):
    """run optimization and verify that log and trajectory coincide"""

    opt_atoms = atoms.copy()
    opt_atoms.constraints = atoms.constraints

    opt_atoms.calc = calc

    with cls(opt_atoms, logfile=logfile, trajectory=trajectory) as opt:
        opt.run(0.2)
        opt.run(0.1)

    with Trajectory(trajectory) as traj, open(logfile) as f:
        next(f)
        for _, (a, line) in enumerate(zip(traj, f)):
            fmax1 = float(line.split()[-1])
            fmax2 = fmax(a.get_forces())

            assert np.allclose(fmax1, fmax2, atol=0.01), (fmax1, fmax2)


@pytest.mark.parametrize('cls_and_kwargs', md_cls_and_kwargs)
def test_md(cls_and_kwargs, atoms=dimer, calc=calc, logfile="md.log",
            timestep=1 * u.fs, trajectory="md.traj"):
    """ run MD for 10 steps and verify that trajectory and log coincide """

    cls, kwargs = cls_and_kwargs
    if hasattr(atoms, "constraints"):
        del atoms.constraints

    atoms.calc = calc

    with cls(atoms, logfile=logfile, timestep=timestep,
             trajectory=trajectory, **kwargs) as md:
        md.run(steps=5)
        md.run(steps=5)

    # assert log file has correct length
    with open(logfile) as fd:
        length = len(fd.readlines())

    assert length == 12, length

    with Trajectory(trajectory) as traj, open(logfile) as fd:
        next(fd)
        for _, (a, line) in enumerate(zip(traj, fd)):
            Epot1, T1 = float(line.split()[-3]), float(line.split()[-1])
            Epot2, T2 = a.get_potential_energy(), a.get_temperature()

            assert np.allclose(T1, T2, atol=0.1), (T1, T2)
            assert np.allclose(Epot1, Epot2, atol=0.01), (Epot1, Epot2)