1import numpy as np
2import pytest
3from ase.data import s22
4from ase.optimize import FIRE
5from ase.constraints import FixBondLengths
6from ase.calculators.tip3p import TIP3P, epsilon0, sigma0
7from ase.calculators.combine_mm import CombineMM
8
9
10def make_atoms():
11    atoms = s22.create_s22_system('Water_dimer')
12    # rotate down in x axis:
13    center = atoms[0].position
14    atoms.translate(-center)
15    h = atoms[3].position[1] - atoms[0].position[1]
16    l = np.linalg.norm(atoms[0].position - atoms[3].position)
17    angle = np.degrees(np.arcsin(h / l))
18    atoms.rotate(angle, '-z', center=center)
19    return atoms
20
21
22def make_4mer():
23    atoms = make_atoms()
24    atoms2 = make_atoms()
25    atoms2.translate([0., 3, 0])
26    atoms2.rotate(180, 'y')
27    atoms2.translate([3, 0, 0])
28    atoms += atoms2
29    return atoms
30
31
32@pytest.mark.slow
33def test_combine_mm2(testdir):
34    # More biased initial positions for faster test. Set
35    # to false for a slower, harder test.
36    fast_test = True
37
38    atoms = make_4mer()
39    atoms.constraints = FixBondLengths([(3 * i + j, 3 * i + (j + 1) % 3)
40                                        for i in range(int(len(atoms) // 3))
41                                        for j in [0, 1, 2]])
42    atoms.calc = TIP3P(np.Inf)
43    tag = '4mer_tip3_opt.'
44    with FIRE(atoms, logfile=tag + 'log', trajectory=tag + 'traj') as opt:
45        opt.run(fmax=0.05)
46    tip3_pos = atoms.get_positions()
47
48    sig = np.array([sigma0, 0, 0])
49    eps = np.array([epsilon0, 0, 0])
50    rc = np.Inf
51    idxes = [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11],
52             list(range(6)), list(range(9)), list(range(6, 12))]
53
54    for ii, idx in enumerate(idxes):
55        atoms = make_4mer()
56        if fast_test:
57            atoms.set_positions(tip3_pos)
58        atoms.constraints = FixBondLengths([(3 * i + j, 3 * i + (j + 1) % 3)
59                                            for i in range(len(atoms) // 3)
60                                            for j in [0, 1, 2]])
61
62        atoms.calc = CombineMM(idx, 3, 3, TIP3P(rc), TIP3P(rc),
63                               sig, eps, sig, eps, rc=rc)
64
65        tag = '4mer_combtip3_opt_{0:02d}.'.format(ii)
66        with FIRE(atoms, logfile=tag + 'log', trajectory=tag + 'traj') as opt:
67            opt.run(fmax=0.05)
68        assert((abs(atoms.positions - tip3_pos) < 1e-8).all())
69        print(
70            '{0}: {1!s:>28s}: Same Geometry as TIP3P'.format(
71                atoms.calc.name, idx))
72