Coverage for hiphive/calculators/zbl.py: 67%

1from ase.calculators.calculator import Calculator, all_changes 

2from ase.neighborlist import NeighborList 

3import numpy as np 

4 

5import numba 

6from numba.typed import List 

7 

8 

9alpha = 1 / 137.035999046 # no unit 

10hbar = 6.582119569e-16 # eV * s 

11c = 299792458 * 1e10 # m/s * A/m 

12prefactor = alpha * hbar * c # eV * A 

13 

14 

15@numba.njit 

16def a(Z_i, Z_j): 

17 return 0.46850 / (Z_i**0.23 + Z_j**0.23) 

18 

19 

20@numba.njit 

21def phi(x): 

22 return (0.18175 * np.exp(-3.19980 * x) 

23 + 0.50986 * np.exp(-0.94229 * x) 

24 + 0.28022 * np.exp(-0.40290 * x) 

25 + 0.02817 * np.exp(-0.20162 * x)) 

26 

27 

28@numba.njit 

29def phi_derivative(x): 

30 return (- 3.19980 * 0.18175 * np.exp(-3.19980 * x) 

31 - 0.94229 * 0.50986 * np.exp(-0.94229 * x) 

32 - 0.40290 * 0.28022 * np.exp(-0.40290 * x) 

33 - 0.20162 * 0.02817 * np.exp(-0.20162 * x)) 

34 

35 

36@numba.njit 

37def zbl_energy(Z_i, Z_j, r_ij): 

38 return prefactor * Z_i * Z_j / r_ij * phi(r_ij / a(Z_i, Z_j)) 

39 

40 

41@numba.njit 

42def zbl_force(Z_i, Z_j, r_ij): 

43 return (-prefactor 

44 * (Z_i * Z_j) 

45 * (phi_derivative(r_ij / a(Z_i, Z_j)) / (r_ij * a(Z_i, Z_j)) 

46 - phi(r_ij / a(Z_i, Z_j)) / r_ij**2)) 

47 

48 

49@numba.njit 

50def np_linalg_norm_axis1(mat): 

51 norm = np.empty(len(mat), dtype=np.float64) 

52 for i in range(len(mat)): 

53 norm[i] = np.linalg.norm(mat[i]) 

54 return norm 

55 

56 

57@numba.njit 

58def inner_loop(ai, positions, cell, offsets, neighbors, forces, energy, numbers): 

59 cells = np.dot(offsets, cell) 

60 v_ij = positions[neighbors] + cells - positions[ai] 

61 r_ij = np_linalg_norm_axis1(v_ij) 

62 for aj, v, r in zip(neighbors, v_ij, r_ij): 

63 energy[0] += zbl_energy(numbers[ai], numbers[aj], r) 

64 force_j = v / r * zbl_force(numbers[ai], numbers[aj], r) 

65 forces[aj] += force_j 

66 forces[ai] -= force_j 

67 

68 

69@numba.njit 

70def outer_loop(positions, cell, offsets_list, neighbors_list, forces, energy, numbers): 

71 for ai in range(len(positions)): 

72 neighbors = neighbors_list[ai] 

73 offsets = offsets_list[ai] 

74 inner_loop(ai, positions, cell, offsets, neighbors, forces, energy, numbers) 

75 

76 

77class ZBLCalculator(Calculator): 

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

79 

80 def __init__(self, cutoff, skin=1.0, **kwargs): 

81 self._cutoff = cutoff 

82 self._skin = skin 

83 super().__init__(**kwargs) 

84 

85 def calculate(self, atoms=None, properties=['energy'], 

86 system_changes=all_changes): 

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

88 

89 if not ('forces' in properties or 'energy' in properties): 89 ↛ 90line 89 didn't jump to line 90, because the condition on line 89 was never true

90 return 

91 

92 n_atoms = len(self.atoms) 

93 

94 if 'numbers' in system_changes: 

95 self.nl = NeighborList([self._cutoff / 2] * n_atoms, 

96 self_interaction=False, 

97 skin=self._skin) 

98 

99 self.nl.update(self.atoms) 

100 

101 positions = self.atoms.positions 

102 numbers = self.atoms.numbers 

103 cell = self.atoms.cell.array 

104 

105 neighbors_list, offsets_list = List(), List() 

106 for ai in range(n_atoms): 

107 neighbors, offsets = self.nl.get_neighbors(ai) 

108 neighbors_list.append(neighbors) 

109 offsets_list.append(offsets.astype(np.float64)) 

110 

111 energy = np.array([0.0]) 

112 forces = np.zeros((n_atoms, 3)) 

113 

114 outer_loop(positions, cell, offsets_list, neighbors_list, forces, energy, numbers) 

115 

116 self.results['energy'] = energy[0] 

117 self.results['forces'] = forces