1import numpy as np
2from ase.calculators.calculator import Calculator
3from ase import units
4
5k_c = units.Hartree * units.Bohr
6
7
8class AtomicCounterIon(Calculator):
9    implemented_properties = ['energy', 'forces']
10
11    def __init__(self, charge, epsilon, sigma, sites_per_mol=1,
12                 rc=7.0, width=1.0):
13        """ Counter Ion Calculator.
14
15        A very simple, nonbonded (Coulumb and LJ)
16        interaction calculator meant for single atom ions
17        to charge neutralize systems (and nothing else)...
18        """
19        self.rc = rc
20        self.width = width
21        self.sites_per_mol = sites_per_mol
22        self.epsilon = epsilon
23        self.sigma = sigma
24        self.charge = charge
25        Calculator.__init__(self)
26
27    def add_virtual_sites(self, positions):
28        return positions
29
30    def get_virtual_charges(self, atoms):
31        charges = np.tile(self.charge, len(atoms) // self.sites_per_mol)
32        return charges
33
34    def redistribute_forces(self, forces):
35        return forces
36
37    def calculate(self, atoms, properties, system_changes):
38        Calculator.calculate(self, atoms, properties, system_changes)
39
40        R = atoms.get_positions()
41        charges = self.get_virtual_charges(atoms)
42        pbc = atoms.pbc
43
44        energy = 0.0
45        forces = np.zeros_like(atoms.get_positions())
46
47        for m in range(len(atoms)):
48            D = R[m + 1:] - R[m]
49            shift = np.zeros_like(D)
50            for i, periodic in enumerate(pbc):
51                if periodic:
52                    L = atoms.cell.diagonal()[i]
53                    shift[:, i] = (D[:, i] + L / 2) % L - L / 2 - D[:, i]
54            D += shift
55            d2 = (D**2).sum(1)
56            d = d2**0.5
57
58            x1 = d > self.rc - self.width
59            x2 = d < self.rc
60            x12 = np.logical_and(x1, x2)
61            y = (d[x12] - self.rc + self.width) / self.width
62            t = np.zeros(len(d))  # cutoff function
63            t[x2] = 1.0
64            t[x12] -= y**2 * (3.0 - 2.0 * y)
65            dtdd = np.zeros(len(d))
66            dtdd[x12] -= 6.0 / self.width * y * (1.0 - y)
67
68            c6 = (self.sigma**2 / d2)**3
69            c12 = c6**2
70            e_lj = 4 * self.epsilon * (c12 - c6)
71            e_c = k_c * charges[m + 1:] * charges[m] / d
72
73            energy += np.dot(t, e_lj)
74            energy += np.dot(t, e_c)
75
76            F = (24 * self.epsilon * (2 * c12 - c6) / d2 * t -
77                 e_lj * dtdd / d)[:, None] * D
78
79            forces[m] -= F.sum(0)
80            forces[m + 1:] += F
81
82            F = (e_c / d2 * t)[:, None] * D \
83                - (e_c * dtdd / d)[:, None] * D
84
85            forces[m] -= F.sum(0)
86            forces[m + 1:] += F
87
88        self.results['energy'] = energy
89        self.results['forces'] = forces
90