Coverage for gpaw/test/test_extensions.py: 13%

1import pytest 

2from gpaw.new.extensions import Extension 

3from ase.units import Hartree, Bohr 

4import numpy as np 

5 

6 

7class Spring: 

8 name = 'spring' 

9 

10 def __init__(self, *, a1, a2, l, k): 

11 self.a1, self.a2, self.l, self.k = a1, a2, l, k 

12 

13 def build(self, atoms, domain_comm, log): 

14 atoms = atoms.copy() 

15 log('Building Spring') 

16 

17 class EnergyAdder(Extension): 

18 name = 'spring' 

19 

20 @property 

21 def _name(_self): 

22 return f'Spring k={self.k}' 

23 

24 def __init__(self): 

25 self._calculate(atoms) 

26 

27 def _calculate(_self, atoms): 

28 D = atoms.get_distance(self.a1, self.a2) - self.l 

29 v = atoms.positions[self.a1] - atoms.positions[self.a2] 

30 v /= np.linalg.norm(v) 

31 F = self.k * D / Hartree * Bohr 

32 _self.E = 1 / 2 * self.k * D**2 / Hartree 

33 _self.F_av = np.zeros((len(atoms), 3)) 

34 _self.F_av[self.a1, :] = -v * F 

35 _self.F_av[self.a2, :] = v * F 

36 

37 def force_contribution(self): 

38 return self.F_av 

39 

40 def get_energy_contributions(self): 

41 return {self._name: self.E} 

42 

43 def move_atoms(self, relpos_ac): 

44 atoms.set_scaled_positions(relpos_ac) 

45 self._calculate(atoms) 

46 

47 return EnergyAdder() 

48 

49 def todict(self): 

50 return dict(a1=self.a1, a2=self.a2, l=self.l, k=self.k) 

51 

52 

53@pytest.mark.parametrize('parallel', [(1, 1), (1, 2), (2, 1)]) 

54@pytest.mark.parametrize('mode', [{'name': 'pw', 'ecut': 300}, 'lcao']) 

55def test_extensions(mode, parallel, in_tmp_dir, gpaw_new): 

56 if not gpaw_new: 

57 pytest.skip('Only GPAW new') 

58 ktot = 20 

59 

60 from gpaw.new.ase_interface import GPAW 

61 from gpaw import restart 

62 from gpaw.mpi import world 

63 domain, band = parallel 

64 if world.size < domain * band: 

65 pytest.skip('Not enough cores for this test.') 

66 if world.size > domain * band * 2: 

67 pytest.skip('Too many cores for this test.') 

68 

69 # 1. Create a calculation with a particular list of extensions. 

70 def get_atoms(): 

71 from ase.build import molecule 

72 atoms = molecule('H2') 

73 atoms.center(vacuum=3) 

74 atoms.set_pbc((True, True, True)) 

75 return atoms 

76 

77 atoms = get_atoms() 

78 

79 def get_calc(atoms): 

80 # To test multiple extensions, create two sprigs which add 

81 # up to k=ktot, which is what is tested in this test 

82 calc = GPAW(extensions=[Spring(a1=0, a2=1, l=2, k=4), 

83 Spring(a1=0, a2=1, l=2, k=ktot - 4)], 

84 symmetry='off', 

85 parallel={'band': band, 'domain': domain}, 

86 kpts=(2, 1, 1), 

87 convergence={'density': 1e-6}, 

88 mode=mode) 

89 atoms.calc = calc 

90 return calc 

91 

92 calc = get_calc(atoms) 

93 

94 E, F = atoms.get_potential_energy(), atoms.get_forces() 

95 

96 # Write the GPW file for the restart test later on (4.) 

97 print('Wrote the potential energy', E) 

98 calc.write('calc.gpw') 

99 

100 # 2. Test that moving the atoms works after an SFC convergence 

101 atoms.positions[0, 2] -= 0.1 

102 movedE, movedF = atoms.get_potential_energy(), atoms.get_forces() 

103 

104 # Reset atoms to their original positions 

105 atoms.positions[0, 2] += 0.1 

106 

107 # 3. Calculate a reference result without extensions 

108 calc = GPAW(mode=mode, 

109 kpts=(2, 1, 1), 

110 convergence={'density': 1e-6}, 

111 symmetry='off') 

112 atoms.calc = calc 

113 

114 E0, F0 = atoms.get_potential_energy(), atoms.get_forces() 

115 

116 # Manually evaluate the spring energy, and compare forces 

117 l = atoms.get_distance(0, 1) 

118 assert E == pytest.approx(E0 + 1 / 2 * ktot * (l - 2)**2) 

119 assert F[0, 2] == pytest.approx(F0[0, 2] - ktot * (l - 2)) 

120 

121 # Evaluate the reference energy and forces also for the moved atoms 

122 atoms.positions[0, 2] -= 0.1 

123 movedE0, movedF0 = atoms.get_potential_energy(), atoms.get_forces() 

124 l = atoms.get_distance(0, 1) 

125 assert movedE == pytest.approx(movedE0 + 1 / 2 * ktot * (l - 2)**2) 

126 assert movedF[0, 2] == pytest.approx(movedF0[0, 2] - ktot * (l - 2)) 

127 

128 def hook(extensions): 

129 return [Spring(**{k: v for k, v in dct.items() if k != 'name'}) 

130 if dct['name'] == 'spring' 

131 else dct 

132 for dct in extensions] 

133 

134 # 4. Test restarting from a file 

135 atoms, calc = restart( 

136 'calc.gpw', 

137 Class=GPAW, 

138 object_hooks={'extensions': hook}) 

139 

140 # Make sure the cached energies and forces are correct 

141 # without a new calculation 

142 assert E == pytest.approx(atoms.get_potential_energy()) 

143 assert F == pytest.approx(atoms.get_forces()) 

144 

145 if mode == 'lcao': 

146 # See issue #1369 

147 return 

148 

149 # Make sure the recalculated energies are forces are correct 

150 atoms.set_positions(atoms.get_positions() + 1e-10) 

151 assert E == pytest.approx(atoms.get_potential_energy(), abs=1e-5) 

152 assert F == pytest.approx(atoms.get_forces(), abs=1e-5) 

153 

154 # 5. Test full blown relaxation. 

155 from ase.optimize import BFGS 

156 atoms = get_atoms() 

157 calc = get_calc(atoms) 

158 relax = BFGS(atoms) 

159 relax.run() 

160 nsteps = relax.nsteps 

161 assert atoms.get_distance(0, 1) == pytest.approx(1.8483, abs=1e-2) 

162 Egs = atoms.get_potential_energy() 

163 L = atoms.get_distance(0, 1) 

164 

165 # 6. Test restarting from a relaxation. 

166 atoms = get_atoms() 

167 calc = get_calc(atoms) 

168 relax = BFGS(atoms, restart='relax_restart') 

169 for _, _ in zip(relax.irun(), range(3)): 

170 pass 

171 calc.write('restart_relax.gpw') 

172 atoms, calc = restart('restart_relax.gpw', 

173 Class=GPAW, 

174 object_hooks={'extensions': hook}) 

175 relax = BFGS(atoms, restart='relax_restart') 

176 relax.run() 

177 

178 assert relax.nsteps + 3 == nsteps 

179 assert atoms.get_distance(0, 1) == pytest.approx(L, abs=1e-2) 

180 assert atoms.get_potential_energy() == pytest.approx(Egs, abs=1e-4)