Coverage for gpaw/new/extensions.py: 24%

1from ase.units import Hartree, Bohr 

2from ase.calculators.calculator import PropertyNotImplementedError 

3import numpy as np 

4from gpaw.mpi import serial_comm, broadcast_exception, broadcast_float 

5import uuid 

6from pathlib import Path 

7import os 

8from gpaw.dft import Extension as ExtensionParameter 

9 

10 

11class Extension: 

12 name = 'unnamed extension' 

13 

14 def get_energy_contributions(self) -> dict[str, float]: 

15 return {} 

16 

17 def force_contribution(self): 

18 raise NotImplementedError 

19 

20 def move_atoms(self, relpos_ac) -> None: 

21 raise NotImplementedError 

22 

23 def update_non_local_hamiltonian(self, 

24 D_sii, 

25 setup, 

26 atom_index, 

27 dH_sii) -> float: 

28 return 0.0 

29 

30 def build(self, atoms, comms, log): 

31 return self 

32 

33 

34class D3(ExtensionParameter): 

35 name = 'd3' 

36 

37 def __init__(self, *, xc, **kwargs): 

38 self.xc = xc 

39 self.kwargs = kwargs 

40 

41 def todict(self) -> dict: 

42 return {'xc': self.xc, **self.kwargs} 

43 

44 def build(self, atoms, communicators, log): 

45 atoms = atoms.copy() 

46 world = communicators['w'] 

47 from ase.calculators.dftd3 import PureDFTD3 

48 

49 # Since DFTD3 is filesystem based, and GPAW has no such requirements 

50 # we need to be absolutely sure that there are no race-conditions 

51 # in files. label cannot be used, because dftd3 executable still 

52 # writes gradients to fixed files, thus a unique folder needs to be 

53 # created. 

54 

55 class D3Extension(Extension): 

56 name = 'd3' 

57 

58 def __init__(self): 

59 super().__init__() 

60 self.stress_vv = np.zeros((3, 3)) * np.nan 

61 self.F_av = np.zeros_like(atoms.positions) * np.nan 

62 self.E = np.nan 

63 self._calculate(atoms) 

64 

65 def _calculate(_self, atoms): 

66 # Circumvent a DFTD3 bug for an isolated atom ASE #1672 

67 if len(atoms) == 1 and not atoms.pbc.any(): 

68 _self.stress_vv = np.zeros((3, 3)) * np.nan 

69 _self.F_av = np.zeros_like(atoms.positions) 

70 _self.E = 0.0 

71 return 

72 

73 cwd = Path.cwd() 

74 assert atoms.calc is None 

75 # Call DFTD3 only with single core due to #1671 

76 with broadcast_exception(world): 

77 if world.rank == 0: 

78 try: 

79 _self.calculate_single_core() 

80 finally: 

81 os.chdir(cwd) 

82 _self.E = broadcast_float(_self.E, world) 

83 world.broadcast(_self.F_av, 0) 

84 world.broadcast(_self.stress_vv, 0) 

85 

86 def calculate_single_core(_self): 

87 """Single core method to calculate D3 forces and stresses""" 

88 

89 label = uuid.uuid4().hex[:8] 

90 directory = Path('dftd3-ext-' + label).absolute() 

91 directory.mkdir() 

92 

93 # Due to ase #1673, relative folders are not supported 

94 # neither are absolute folders due to 80 character limit. 

95 # The only way out, is to chdir to a temporary folder here. 

96 os.chdir(directory) 

97 log('Evaluating D3 corrections at temporary' 

98 f' folder {directory}') 

99 atoms.calc = PureDFTD3(xc=self.xc, 

100 directory='.', 

101 comm=serial_comm, 

102 **self.kwargs) 

103 

104 # XXX params.xc should be taken directly from the calculator. 

105 # XXX What if this is changed via set? 

106 _self.F_av = atoms.get_forces() / Hartree * Bohr 

107 

108 try: 

109 # Copy needed because array is not c-contigous 

110 _self.stress_vv = atoms.get_stress(voigt=False).copy() \ 

111 / Hartree * Bohr**3 

112 except PropertyNotImplementedError: 

113 _self.stress_vv = np.zeros((3, 3)) * np.nan 

114 

115 _self.E = atoms.get_potential_energy() / Hartree 

116 try: 

117 os.unlink('ase_dftd3.out') 

118 os.unlink('ase_dftd3.POSCAR') 

119 os.unlink('dftd3_cellgradient') 

120 os.unlink('dftd3_gradient') 

121 os.rmdir(directory.absolute()) 

122 except OSError as e: 

123 log('Unable to remove files and folder', e) 

124 atoms.calc = None 

125 

126 def get_energy_contributions(_self) -> dict[str, float]: 

127 """Returns the energy contributions from D3 in Hartree""" 

128 return {f'D3 (xc={self.xc})': _self.E} 

129 

130 def get_energy(self) -> float: 

131 """Returns the energy contribution from D3 in eV""" 

132 return self.E * Hartree 

133 

134 def force_contribution(self): 

135 return self.F_av 

136 

137 def stress_contribution(self): 

138 if np.isnan(self.stress_vv).all(): 

139 raise PropertyNotImplementedError 

140 return self.stress_vv 

141 

142 def move_atoms(self, relpos_ac) -> None: 

143 atoms.set_scaled_positions(relpos_ac) 

144 self._calculate(atoms) 

145 

146 return D3Extension()