Coverage for /builds/ase/ase/ase/calculators/counterions.py: 96.72%

1# fmt: off 

2 

3import numpy as np 

4 

5from ase import units 

6from ase.calculators.calculator import Calculator 

7 

8k_c = units.Hartree * units.Bohr 

9 

10 

11class AtomicCounterIon(Calculator): 

12 implemented_properties = ['energy', 'forces'] 

13 

14 def __init__(self, charge, epsilon, sigma, sites_per_mol=1, 

15 rc=7.0, width=1.0): 

16 """ Counter Ion Calculator. 

17 

18 A very simple, nonbonded (Coulumb and LJ) 

19 interaction calculator meant for single atom ions 

20 to charge neutralize systems (and nothing else)... 

21 """ 

22 self.rc = rc 

23 self.width = width 

24 self.sites_per_mol = sites_per_mol 

25 self.epsilon = epsilon 

26 self.sigma = sigma 

27 self.charge = charge 

28 Calculator.__init__(self) 

29 

30 def add_virtual_sites(self, positions): 

31 return positions 

32 

33 def get_virtual_charges(self, atoms): 

34 charges = np.tile(self.charge, len(atoms) // self.sites_per_mol) 

35 return charges 

36 

37 def redistribute_forces(self, forces): 

38 return forces 

39 

40 def calculate(self, atoms, properties, system_changes): 

41 Calculator.calculate(self, atoms, properties, system_changes) 

42 

43 R = atoms.get_positions() 

44 charges = self.get_virtual_charges(atoms) 

45 pbc = atoms.pbc 

46 

47 energy = 0.0 

48 forces = np.zeros_like(atoms.get_positions()) 

49 

50 for m in range(len(atoms)): 

51 D = R[m + 1:] - R[m] 

52 shift = np.zeros_like(D) 

53 for i, periodic in enumerate(pbc): 

54 if periodic: 

55 L = atoms.cell.diagonal()[i] 

56 shift[:, i] = (D[:, i] + L / 2) % L - L / 2 - D[:, i] 

57 D += shift 

58 d2 = (D**2).sum(1) 

59 d = d2**0.5 

60 

61 x1 = d > self.rc - self.width 

62 x2 = d < self.rc 

63 x12 = np.logical_and(x1, x2) 

64 y = (d[x12] - self.rc + self.width) / self.width 

65 t = np.zeros(len(d)) # cutoff function 

66 t[x2] = 1.0 

67 t[x12] -= y**2 * (3.0 - 2.0 * y) 

68 dtdd = np.zeros(len(d)) 

69 dtdd[x12] -= 6.0 / self.width * y * (1.0 - y) 

70 

71 c6 = (self.sigma**2 / d2)**3 

72 c12 = c6**2 

73 e_lj = 4 * self.epsilon * (c12 - c6) 

74 e_c = k_c * charges[m + 1:] * charges[m] / d 

75 

76 energy += np.dot(t, e_lj) 

77 energy += np.dot(t, e_c) 

78 

79 F = (24 * self.epsilon * (2 * c12 - c6) / d2 * t - 

80 e_lj * dtdd / d)[:, None] * D 

81 

82 forces[m] -= F.sum(0) 

83 forces[m + 1:] += F 

84 

85 F = (e_c / d2 * t)[:, None] * D \ 

86 - (e_c * dtdd / d)[:, None] * D 

87 

88 forces[m] -= F.sum(0) 

89 forces[m + 1:] += F 

90 

91 self.results['energy'] = energy 

92 self.results['forces'] = forces