"Test that atoms.center() works when adding vacuum ()"

import numpy as np
from math import pi, sqrt, cos
from ase import data
from ase.lattice.cubic import FaceCenteredCubic


def checkang(a, b, phi):
    "Check the angle between two vectors."
    cosphi = np.dot(a,b) / sqrt(np.dot(a,a) * np.dot(b,b))
    assert np.abs(cosphi - cos(phi)) < 1e-10

symb = "Cu"
Z = data.atomic_numbers[symb]
a0 = data.reference_states[Z]['a']

# (100) oriented block
atoms = FaceCenteredCubic(size=(5,5,5), symbol="Cu", pbc=(1,1,0))
assert len(atoms) == 5*5*5*4
c = atoms.get_cell()
checkang(c[0], c[1], pi/2)
checkang(c[0], c[2], pi/2)
checkang(c[1], c[2], pi/2)
assert np.abs(5 * a0 - c[2,2]) < 1e-10

# Add vacuum in one direction
vac = 10.0
atoms.center(axis=2, vacuum=vac)
c = atoms.get_cell()
checkang(c[0], c[1], pi/2)
checkang(c[0], c[2], pi/2)
checkang(c[1], c[2], pi/2)
assert np.abs(4.5 * a0 + 2* vac - c[2,2]) < 1e-10

# Add vacuum in all directions
vac = 4.0
atoms.center(vacuum=vac)
c = atoms.get_cell()
checkang(c[0], c[1], pi/2)
checkang(c[0], c[2], pi/2)
checkang(c[1], c[2], pi/2)
assert np.abs(4.5 * a0 + 2* vac - c[0,0]) < 1e-10
assert np.abs(4.5 * a0 + 2* vac - c[1,1]) < 1e-10
assert np.abs(4.5 * a0 + 2* vac - c[2,2]) < 1e-10

# Now a general unit cell
atoms = FaceCenteredCubic(size=(5,5,5), directions=[[1,0,0], [0,1,0], [1,0,1]],
                          symbol="Cu", pbc=(1,1,0))
assert len(atoms) == 5*5*5*2
c = atoms.get_cell()
checkang(c[0], c[1], pi/2)
checkang(c[0], c[2], pi/4)
checkang(c[1], c[2], pi/2)
assert np.abs(2.5 * a0 - c[2,2]) < 1e-10

# Add vacuum in one direction
vac = 10.0
atoms.center(axis=2, vacuum=vac)
c = atoms.get_cell()
checkang(c[0], c[1], pi/2)
checkang(c[0], c[2], pi/4)
checkang(c[1], c[2], pi/2)
assert np.abs(2 * a0 + 2* vac - c[2,2]) < 1e-10

# Recenter without specifying vacuum
atoms.center()
c = atoms.get_cell()
checkang(c[0], c[1], pi/2)
checkang(c[0], c[2], pi/4)
checkang(c[1], c[2], pi/2)
assert np.abs(2 * a0 + 2* vac - c[2,2]) < 1e-10

a2 = atoms.copy()

# Add vacuum in all directions
vac = 4.0
atoms.center(vacuum=vac)
c = atoms.get_cell()
checkang(c[0], c[1], pi / 2)
checkang(c[0], c[2], pi / 4)
checkang(c[1], c[2], pi / 2)
assert np.abs(4.5 * a0 + 2 * vac - c[1, 1]) < 1e-10
assert np.abs(2 * a0 + 2 * vac - c[2, 2]) < 1e-10

# One axis at the time:
for i in range(3):
    a2.center(vacuum=vac, axis=i)

assert abs(atoms.positions - a2.positions).max() < 1e-12
assert abs(atoms.cell - a2.cell).max() < 1e-12
