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

61 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-12-10 11:04 +0000

1import numpy as np 

2 

3from ase import units 

4from ase.calculators.calculator import Calculator 

5 

6k_c = units.Hartree * units.Bohr 

7 

8 

9class AtomicCounterIon(Calculator): 

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

11 

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

13 rc=7.0, width=1.0): 

14 """ Counter Ion Calculator. 

15 

16 A very simple, nonbonded (Coulumb and LJ) 

17 interaction calculator meant for single atom ions 

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

19 """ 

20 self.rc = rc 

21 self.width = width 

22 self.sites_per_mol = sites_per_mol 

23 self.epsilon = epsilon 

24 self.sigma = sigma 

25 self.charge = charge 

26 Calculator.__init__(self) 

27 

28 def add_virtual_sites(self, positions): 

29 return positions 

30 

31 def get_virtual_charges(self, atoms): 

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

33 return charges 

34 

35 def redistribute_forces(self, forces): 

36 return forces 

37 

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

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

40 

41 R = atoms.get_positions() 

42 charges = self.get_virtual_charges(atoms) 

43 pbc = atoms.pbc 

44 

45 energy = 0.0 

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

47 

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

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

50 shift = np.zeros_like(D) 

51 for i, periodic in enumerate(pbc): 

52 if periodic: 

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

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

55 D += shift 

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

57 d = d2**0.5 

58 

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

60 x2 = d < self.rc 

61 x12 = np.logical_and(x1, x2) 

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

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

64 t[x2] = 1.0 

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

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

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

68 

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

70 c12 = c6**2 

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

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

73 

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

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

76 

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

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

79 

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

81 forces[m + 1:] += F 

82 

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

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

85 

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

87 forces[m + 1:] += F 

88 

89 self.results['energy'] = energy 

90 self.results['forces'] = forces