File: test_vasp_structure.py

package info (click to toggle)
python-ase 3.26.0-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 15,484 kB
  • sloc: python: 148,112; xml: 2,728; makefile: 110; javascript: 47
file content (178 lines) | stat: -rw-r--r-- 6,050 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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
import io
import os
import unittest

import numpy as np
import pytest

import ase
import ase.build
import ase.io
from ase.build import graphene_nanoribbon
from ase.calculators.calculator import compare_atoms
from ase.constraints import (
    FixAtoms,
    FixedLine,
    FixedPlane,
    FixScaled,
    constrained_indices,
)
from ase.io.vasp import read_vasp_xdatcar, write_vasp_xdatcar


class TestXdatcarRoundtrip(unittest.TestCase):
    def setUp(self):
        self.outfile = 'NaCl.XDATCAR'

        self.NaCl = ase.build.bulk('NaCl', 'rocksalt', a=5.64)

    def tearDown(self):
        if os.path.isfile(self.outfile):
            os.remove(self.outfile)

    def assert_atoms_almost_equal(self, atoms, other, tol=1e-15):
        """Compare two Atoms objects, raising AssertionError if different"""
        system_changes = compare_atoms(atoms, other, tol=tol)

        if len(system_changes) > 0:
            raise AssertionError(
                'Atoms objects differ by {}'.format(', '.join(system_changes))
            )

    def assert_trajectory_almost_equal(self, traj1, traj2):
        self.assertEqual(len(traj1), len(traj2))
        for image, other in zip(traj1, traj2):
            self.assert_atoms_almost_equal(image, other)

    def test_roundtrip(self):
        # Create a series of translated cells
        trajectory = [self.NaCl.copy() for _ in range(5)]
        for i, atoms in enumerate(trajectory):
            atoms.set_scaled_positions(
                atoms.get_scaled_positions() + i * np.array([0.05, 0, 0.02])
            )
            atoms.wrap()

        ase.io.write(self.outfile, trajectory, format='vasp-xdatcar')
        roundtrip_trajectory = ase.io.read(self.outfile, index=':')
        self.assert_trajectory_almost_equal(trajectory, roundtrip_trajectory)

    def test_roundtrip_single_atoms(self):
        atoms = ase.build.bulk('Ge')
        ase.io.write(self.outfile, atoms, format='vasp-xdatcar')
        roundtrip_atoms = ase.io.read(self.outfile)
        self.assert_atoms_almost_equal(atoms, roundtrip_atoms)

    def test_typeerror(self):
        with self.assertRaises(TypeError):
            atoms = ase.build.bulk('Ge')
            write_vasp_xdatcar(self.outfile, atoms)
        with self.assertRaises(TypeError):
            not_atoms = 1
            ase.io.write(self.outfile, not_atoms, format='vasp-xdatcar')
        with self.assertRaises(TypeError):
            not_traj = [True, False, False]
            ase.io.write(self.outfile, not_traj, format='vasp-xdatcar')


def test_index():
    """Test if the `index` option works correctly"""
    atoms0 = ase.build.bulk('X', 'sc', a=1.0)
    atoms1 = atoms0.copy()
    atoms1.positions += 0.1
    images = (atoms0, atoms1)
    with io.StringIO() as buf:
        write_vasp_xdatcar(buf, images)

        buf.seek(0)
        atoms = read_vasp_xdatcar(buf, index=0)  # atoms0
        np.testing.assert_allclose(atoms.positions, [[0.0, 0.0, 0.0]])

        buf.seek(0)
        atoms = read_vasp_xdatcar(buf, index=-1)  # atoms1
        np.testing.assert_allclose(atoms.positions, [[0.1, 0.1, 0.1]])

        buf.seek(0)
        images = read_vasp_xdatcar(buf, index=None)  # atoms1
        np.testing.assert_allclose(atoms.positions, [[0.1, 0.1, 0.1]])

        buf.seek(0)
        images = read_vasp_xdatcar(buf, index=':')  # (atoms0, atoms1)
        assert isinstance(images, list)


# Start of tests for constraints
# VASP supports FixAtoms and FixScaled as well as FixedLine and
# FixedPlane if the direction is along a lattice vector. Test that
# these constraints are preserved when writing and reading POSCAR
# files.
indices_to_constrain = [0, 2]


@pytest.fixture()
def graphene_atoms():
    atoms = graphene_nanoribbon(2, 2, type='armchair', saturated=False)
    atoms.cell = [[10, 0, 0], [0, 10, 0], [0, 0, 10]]
    return atoms


def poscar_roundtrip(atoms):
    """Write a POSCAR file, read it back and return the new atoms object"""
    atoms.write('POSCAR', direct=True)
    return ase.io.read('POSCAR')


@pytest.mark.parametrize('whitespace', ['\n', '   ', '   \n\n  \n'])
def test_with_whitespace(graphene_atoms, whitespace):
    graphene_atoms.write('POSCAR', direct=True)
    with open('POSCAR', 'a') as fd:
        fd.write(whitespace)
    assert str(ase.io.read('POSCAR').symbols) == str(graphene_atoms.symbols)


def test_FixAtoms(graphene_atoms):
    atoms = graphene_atoms
    atoms.set_constraint(FixAtoms(indices=indices_to_constrain))
    new_atoms = poscar_roundtrip(atoms)

    # Assert that constraints are preserved
    assert isinstance(new_atoms.constraints[0], FixAtoms)
    assert np.all(new_atoms.constraints[0].index == indices_to_constrain)


def test_FixScaled(graphene_atoms):
    atoms = graphene_atoms
    atoms.set_constraint(FixScaled(indices_to_constrain, mask=[0, 1, 1]))
    new_atoms = poscar_roundtrip(atoms)

    # Assert that constraints are preserved
    assert np.all(constrained_indices(new_atoms) == indices_to_constrain)
    assert np.all(new_atoms.constraints[0].mask == [0, 1, 1])


@pytest.mark.parametrize('ConstraintClass', [FixedLine, FixedPlane])
def test_FixedLine_and_Plane(ConstraintClass, graphene_atoms):
    atoms = graphene_atoms
    atoms.set_constraint(
        ConstraintClass(indices=indices_to_constrain, direction=[1, 0, 0])
    )
    new_atoms = poscar_roundtrip(atoms)

    # FixedLine and FixedPlane are converted to FixScaled. During
    # a relaxation the results will be the same since FixScaled
    # are equivalent to the others if the direction in FixedLine
    # or FixedPlane is along a lattice vector.

    assert np.all(constrained_indices(new_atoms) == indices_to_constrain)


def test_write_read_velocities(graphene_atoms):
    vel = np.zeros_like(graphene_atoms.positions)
    vel = np.linspace(-1, 1, 3 * len(graphene_atoms)).reshape(-1, 3)
    graphene_atoms.set_velocities(vel)

    graphene_atoms.write('CONTCAR', direct=False)
    new_atoms = ase.io.read('CONTCAR')
    new_vel = new_atoms.get_velocities()

    assert np.allclose(vel, new_vel)