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
|
# fmt: off
"""Test write and read."""
import io
import re
import numpy as np
import pytest
from ase import Atoms
from ase.build import bulk
from ase.data import atomic_numbers
from ase.io.lammpsdata import read_lammps_data, write_lammps_data
def compare(atoms: Atoms, atoms_ref: Atoms):
"""Compare two `Atoms` objects"""
assert all(atoms.numbers == atoms_ref.numbers)
assert atoms.get_masses() == pytest.approx(atoms_ref.get_masses())
# Note: Raw positions cannot be compared.
# `write_lammps_data` changes cell orientation.
assert atoms.get_scaled_positions() == pytest.approx(
atoms_ref.get_scaled_positions())
@pytest.mark.parametrize('masses', [False, True])
class _Base:
def run(self, atoms_ref: Atoms, masses: bool):
"""Run tests"""
self.check_explicit_numbers(atoms_ref, masses)
if masses:
self.check_masses2numbers(atoms_ref)
def check_explicit_numbers(self, atoms_ref: Atoms, masses: bool):
"""Check if write-read is consistent when giving Z_of_type."""
buf = io.StringIO()
write_lammps_data(buf, atoms_ref, masses=masses)
buf.seek(0)
# By default, write_lammps_data assigns atom types to the elements in
# alphabetical order. To be consistent, here spiecies are also sorted.
species = sorted(set(atoms_ref.get_chemical_symbols()))
Z_of_type = {i + 1: atomic_numbers[s] for i, s in enumerate(species)}
atoms = read_lammps_data(buf, Z_of_type=Z_of_type, atom_style='atomic')
compare(atoms, atoms_ref)
def check_masses2numbers(self, atoms_ref: Atoms):
"""Check if write-read is consistent when guessing atomic numbers."""
buf = io.StringIO()
write_lammps_data(buf, atoms_ref, masses=True)
buf.seek(0)
atoms = read_lammps_data(buf, atom_style='atomic')
compare(atoms, atoms_ref)
@pytest.mark.parametrize('cubic', [False, True])
class TestCubic(_Base):
"""Test cubic structures."""
def test_bcc(self, cubic: bool, masses: bool):
"""Test bcc."""
atoms_ref = bulk('Li', 'bcc', cubic=cubic)
self.run(atoms_ref, masses)
def test_fcc(self, cubic: bool, masses: bool):
"""Test fcc."""
atoms_ref = bulk('Cu', 'fcc', cubic=cubic)
self.run(atoms_ref, masses)
def test_rocksalt(self, cubic: bool, masses: bool):
"""Test rocksalt."""
atoms_ref = bulk('NaCl', 'rocksalt', a=1.0, cubic=cubic)
self.run(atoms_ref, masses)
def test_fluorite(self, cubic: bool, masses: bool):
"""Test fluorite."""
atoms_ref = bulk('CaF2', 'fluorite', a=1.0, cubic=cubic)
self.run(atoms_ref, masses)
@pytest.mark.parametrize('orthorhombic', [False, True])
class TestOrthorhombic(_Base):
"""Test orthorhombic structures."""
def test_hcp(self, orthorhombic: bool, masses: bool):
"""Test hcp."""
atoms_ref = bulk('Mg', 'hcp', orthorhombic=orthorhombic)
self.run(atoms_ref, masses)
@pytest.mark.parametrize('atom_style', ['atomic', 'charge', 'full'])
def test_atom_style(atom_style: str):
"""Test `atom_style`"""
atoms_ref = bulk('Cu', 'fcc')
# note that we should do `masses=True` to later guess atomic numbers
buf = io.StringIO()
write_lammps_data(buf, atoms_ref, masses=True, atom_style=atom_style)
# test if the `atom_style` can be guessed from the comment
buf.seek(0)
atoms = read_lammps_data(buf, atom_style=None)
compare(atoms, atoms_ref)
# test when `atom_style` is explicitly specified
buf.seek(0)
atoms = read_lammps_data(buf, atom_style=atom_style)
compare(atoms, atoms_ref)
if atom_style not in ['atomic', 'full']:
return
# test if `atom_style` can be guessed from the length of fields
# remove comment in the 'Atoms' line
buf = io.StringIO(re.sub('.*Atoms.*#.*', 'Atoms', buf.getvalue()))
atoms = read_lammps_data(buf, atom_style=None)
compare(atoms, atoms_ref)
@pytest.mark.parametrize('atom_style', ['atomic', 'charge', 'full'])
@pytest.mark.parametrize('write_image_flags', [False, True])
def test_image_flags(write_image_flags: bool, atom_style: str):
"""Test if `wrap` and `write_image_flags` work correctly."""
atoms_ref = bulk('Ge')
# shift atomic positions
scaled_positions = atoms_ref.get_scaled_positions(wrap=False)
shift = (0.125, 1.125, -0.125)
atoms_ref.set_scaled_positions(scaled_positions + shift)
# note that we should do `masses=True` to later guess atomic numbers
buf = io.StringIO()
write_lammps_data(
buf,
atoms_ref,
write_image_flags=write_image_flags,
masses=True,
atom_style=atom_style,
)
buf.seek(0)
atoms = read_lammps_data(buf)
# The atomic positions should be wrapped.
np.testing.assert_allclose(
atoms.get_scaled_positions(wrap=False),
atoms_ref.get_scaled_positions(wrap=False),
)
def test_bonds(lammpsdata_file_path):
"""Test if writing bonds works correctly."""
atoms = read_lammps_data(lammpsdata_file_path, atom_style='full')
lammpsdata_buf = io.StringIO()
write_lammps_data(
lammpsdata_buf, atoms, atom_style='full',
masses=True, velocities=True, bonds=True)
lammpsdata_buf.seek(0)
atoms2 = read_lammps_data(lammpsdata_buf, atom_style='full')
np.testing.assert_array_equal(
atoms.arrays["bonds"], atoms2.arrays["bonds"]
)
|