1"""TIP3P potential."""
2
3import numpy as np
4
5import ase.units as units
6from ase.calculators.calculator import Calculator, all_changes
7
8qH = 0.417
9sigma0 = 3.15061
10epsilon0 = 0.1521 * units.kcal / units.mol
11rOH = 0.9572
12angleHOH = 104.52
13thetaHOH = 104.52 / 180 * np.pi  # we keep this for backwards compatibility
14
15
16class TIP3P(Calculator):
17    implemented_properties = ['energy', 'forces']
18    nolabel = True
19    pcpot = None
20
21    def __init__(self, rc=5.0, width=1.0):
22        """TIP3P potential.
23
24        rc: float
25            Cutoff radius for Coulomb part.
26        width: float
27            Width for cutoff function for Coulomb part.
28        """
29        self.rc = rc
30        self.width = width
31        Calculator.__init__(self)
32        self.sites_per_mol = 3
33
34    def calculate(self, atoms=None,
35                  properties=['energy'],
36                  system_changes=all_changes):
37        Calculator.calculate(self, atoms, properties, system_changes)
38
39        R = self.atoms.positions.reshape((-1, 3, 3))
40        Z = self.atoms.numbers
41        pbc = self.atoms.pbc
42        cell = self.atoms.cell.diagonal()
43        nh2o = len(R)
44
45        assert (self.atoms.cell == np.diag(cell)).all(), 'not orthorhombic'
46        assert ((cell >= 2 * self.rc) | ~pbc).all(), 'cutoff too large'  # ???
47        if Z[0] == 8:
48            o = 0
49        else:
50            o = 2
51        assert (Z[o::3] == 8).all()
52        assert (Z[(o + 1) % 3::3] == 1).all()
53        assert (Z[(o + 2) % 3::3] == 1).all()
54
55        charges = np.array([qH, qH, qH])
56        charges[o] *= -2
57
58        energy = 0.0
59        forces = np.zeros((3 * nh2o, 3))
60
61        for m in range(nh2o - 1):
62            DOO = R[m + 1:, o] - R[m, o]
63            shift = np.zeros_like(DOO)
64            for i, periodic in enumerate(pbc):
65                if periodic:
66                    L = cell[i]
67                    shift[:, i] = (DOO[:, i] + L / 2) % L - L / 2 - DOO[:, i]
68            DOO += shift
69            d2 = (DOO**2).sum(1)
70            d = d2**0.5
71            x1 = d > self.rc - self.width
72            x2 = d < self.rc
73            x12 = np.logical_and(x1, x2)
74            y = (d[x12] - self.rc + self.width) / self.width
75            t = np.zeros(len(d))  # cutoff function
76            t[x2] = 1.0
77            t[x12] -= y**2 * (3.0 - 2.0 * y)
78            dtdd = np.zeros(len(d))
79            dtdd[x12] -= 6.0 / self.width * y * (1.0 - y)
80            c6 = (sigma0**2 / d2)**3
81            c12 = c6**2
82            e = 4 * epsilon0 * (c12 - c6)
83            energy += np.dot(t, e)
84            F = (24 * epsilon0 * (2 * c12 - c6) / d2 * t -
85                 e * dtdd / d)[:, np.newaxis] * DOO
86            forces[m * 3 + o] -= F.sum(0)
87            forces[m * 3 + 3 + o::3] += F
88
89            for j in range(3):
90                D = R[m + 1:] - R[m, j] + shift[:, np.newaxis]
91                r2 = (D**2).sum(axis=2)
92                r = r2**0.5
93                e = charges[j] * charges / r * units.Hartree * units.Bohr
94                energy += np.dot(t, e).sum()
95                F = (e / r2 * t[:, np.newaxis])[:, :, np.newaxis] * D
96                FOO = -(e.sum(1) * dtdd / d)[:, np.newaxis] * DOO
97                forces[(m + 1) * 3 + o::3] += FOO
98                forces[m * 3 + o] -= FOO.sum(0)
99                forces[(m + 1) * 3:] += F.reshape((-1, 3))
100                forces[m * 3 + j] -= F.sum(axis=0).sum(axis=0)
101
102        if self.pcpot:
103            e, f = self.pcpot.calculate(np.tile(charges, nh2o),
104                                        self.atoms.positions)
105            energy += e
106            forces += f
107
108        self.results['energy'] = energy
109        self.results['forces'] = forces
110
111    def embed(self, charges):
112        """Embed atoms in point-charges."""
113        self.pcpot = PointChargePotential(charges)
114        return self.pcpot
115
116    def check_state(self, atoms, tol=1e-15):
117        system_changes = Calculator.check_state(self, atoms, tol)
118        if self.pcpot and self.pcpot.mmpositions is not None:
119            system_changes.append('positions')
120        return system_changes
121
122    def add_virtual_sites(self, positions):
123        return positions  # no virtual sites
124
125    def redistribute_forces(self, forces):
126        return forces
127
128    def get_virtual_charges(self, atoms):
129        charges = np.empty(len(atoms))
130        charges[:] = qH
131        if atoms.numbers[0] == 8:
132            charges[::3] = -2 * qH
133        else:
134            charges[2::3] = -2 * qH
135        return charges
136
137
138class PointChargePotential:
139    def __init__(self, mmcharges):
140        """Point-charge potential for TIP3P.
141
142        Only used for testing QMMM.
143        """
144        self.mmcharges = mmcharges
145        self.mmpositions = None
146        self.mmforces = None
147
148    def set_positions(self, mmpositions, com_pv=None):
149        self.mmpositions = mmpositions
150
151    def calculate(self, qmcharges, qmpositions):
152        energy = 0.0
153        self.mmforces = np.zeros_like(self.mmpositions)
154        qmforces = np.zeros_like(qmpositions)
155        for C, R, F in zip(self.mmcharges, self.mmpositions, self.mmforces):
156            d = qmpositions - R
157            r2 = (d**2).sum(1)
158            e = units.Hartree * units.Bohr * C * r2**-0.5 * qmcharges
159            energy += e.sum()
160            f = (e / r2)[:, np.newaxis] * d
161            qmforces += f
162            F -= f.sum(0)
163        self.mmpositions = None
164        return energy, qmforces
165
166    def get_forces(self, calc):
167        return self.mmforces
168