1from itertools import count 

2import numpy as np 

3 

4from ase import Atoms 

5from ase.units import invcm, Ha 

6from ase.data import atomic_masses 

7from ase.calculators.calculator import all_changes 

8from ase.calculators.morse import MorsePotential 

9from ase.calculators.excitation_list import Excitation, ExcitationList 

10 

11"""The H2 molecule represented by Morse-Potentials for 

12gound and first 3 excited singlet states B + C(doubly degenerate)""" 

13 

14npa = np.array 

15# data from: 

16# https://webbook.nist.gov/cgi/cbook.cgi?ID=C1333740&Mask=1000#Diatomic 

17# X B C C 

18Re = npa([0.74144, 1.2928, 1.0327, 1.0327]) # eq. bond length 

19ome = npa([4401.21, 1358.09, 2443.77, 2443.77]) # vibrational frequency 

20# electronic transition energy 

21Etrans = npa([0, 91700.0, 100089.9, 100089.9]) * invcm 

22 

23# dissociation energy 

24# GS: https://aip.scitation.org/doi/10.1063/1.3120443 

25De = np.ones(4) * 36118.069 * invcm 

26# B, C separated energy E(1s) - E(2p) 

27De[1:] += Ha / 2 - Ha / 8 

28De -= Etrans 

29 

30# Morse parameter 

31m = atomic_masses[1] * 0.5 # reduced mass 

32# XXX find scaling factor 

33rho0 = Re * ome * invcm * np.sqrt(m / 2 / De) * 4401.21 / 284.55677429605862 

34 

35 

36def H2Morse(state=0): 

37 """Return H2 as a Morse-Potential with calculator attached.""" 

38 atoms = Atoms('H2', positions=np.zeros((2, 3))) 

39 atoms[1].position[2] = Re[state] 

40 atoms.calc = H2MorseCalculator(state=state) 

41 atoms.get_potential_energy() 

42 return atoms 

43 

44 

45class H2MorseCalculator(MorsePotential): 

46 """H2 ground or excited state as Morse potential""" 

47 _count = count(0) 

48 

49 def __init__(self, restart=None, state=0, rng=np.random): 

50 self.rng = rng 

51 MorsePotential.__init__(self, 

52 restart=restart, 

53 epsilon=De[state], 

54 r0=Re[state], rho0=rho0[state]) 

55 

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

57 system_changes=all_changes): 

58 if atoms is not None: 

59 assert len(atoms) == 2 

60 MorsePotential.calculate(self, atoms, properties, system_changes) 

61 

62 # determine 'wave functions' including 

63 # Berry phase (arbitrary sign) and 

64 # random orientation of wave functions perpendicular 

65 # to the molecular axis 

66 

67 # molecular axis 

68 vr = atoms[1].position - atoms[0].position 

69 r = np.linalg.norm(vr) 

70 hr = vr / r 

71 # perpendicular axes 

72 vrand = self.rng.random(3) 

73 hx = np.cross(hr, vrand) 

74 hx /= np.linalg.norm(hx) 

75 hy = np.cross(hr, hx) 

76 hy /= np.linalg.norm(hy) 

77 wfs = [1, hr, hx, hy] 

78 # Berry phase 

79 berry = (-1)**self.rng.randint(0, 2, 4) 

80 self.wfs = [wf * b for wf, b in zip(wfs, berry)] 

81 

82 def read(self, filename): 

83 ms = self 

84 with open(filename) as fd: 

85 ms.wfs = [int(fd.readline().split()[0])] 

86 for i in range(1, 4): 

87 ms.wfs.append( 

88 np.array([float(x) 

89 for x in fd.readline().split()[:4]])) 

90 ms.filename = filename 

91 return ms 

92 

93 def write(self, filename, option=None): 

94 """write calculated state to a file""" 

95 with open(filename, 'w') as fd: 

96 fd.write('{}\n'.format(self.wfs[0])) 

97 for wf in self.wfs[1:]: 

98 fd.write('{0:g} {1:g} {2:g}\n'.format(*wf)) 

99 

100 def overlap(self, other): 

101 ov = np.zeros((4, 4)) 

102 ov[0, 0] = self.wfs[0] * other.wfs[0] 

103 wfs = np.array(self.wfs[1:]) 

104 owfs = np.array(other.wfs[1:]) 

105 ov[1:, 1:] = np.dot(wfs, owfs.T) 

106 return ov 

107 

108 

109class H2MorseExcitedStatesCalculator(): 

110 """First singlet excited states of H2 from Morse potentials""" 

111 

112 def __init__(self, nstates=3): 

113 """ 

114 Parameters 

115 ---------- 

116 nstates: int 

117 Numer of states to calculate 0 < nstates < 4, default 3 

118 """ 

119 assert nstates > 0 and nstates < 4 

120 self.nstates = nstates 

121 

122 def calculate(self, atoms): 

123 """Calculate excitation spectrum 

124 

125 Parameters 

126 ---------- 

127 atoms: Ase atoms object 

128 """ 

129 # central me value and rise, unit Bohr 

130 # from DOI: 10.1021/acs.jctc.9b00584 

131 mc = [0, 0.8, 0.7, 0.7] 

132 mr = [0, 1.0, 0.5, 0.5] 

133 

134 cgs = atoms.calc 

135 r = atoms.get_distance(0, 1) 

136 E0 = cgs.get_potential_energy(atoms) 

137 

138 exl = H2MorseExcitedStates() 

139 for i in range(1, self.nstates + 1): 

140 hvec = cgs.wfs[0] * cgs.wfs[i] 

141 energy = Ha * (0.5 - 1. / 8) - E0 

142 calc = H2MorseCalculator(state=i) 

143 calc.calculate(atoms) 

144 energy += calc.get_potential_energy() 

145 

146 mur = hvec * (mc[i] + (r - Re[0]) * mr[i]) 

147 muv = mur 

148 

149 exl.append(H2Excitation(energy, i, mur, muv)) 

150 return exl 

151 

152 

153class H2MorseExcitedStates(ExcitationList): 

154 """First singlet excited states of H2""" 

155 

156 def __init__(self, nstates=3): 

157 """ 

158 Parameters 

159 ---------- 

160 nstates: int, 1 <= nstates <= 3 

161 Number of excited states to consider, default 3 

162 """ 

163 self.nstates = nstates 

164 super().__init__() 

165 

166 def overlap(self, ov_nn, other): 

167 return (ov_nn[1:len(self) + 1, 1:len(self) + 1] * 

168 ov_nn[0, 0]) 

169 

170 @classmethod 

171 def read(cls, filename, nstates=3): 

172 """Read myself from a file""" 

173 exl = cls(nstates) 

174 with open(filename, 'r') as fd: 

175 exl.filename = filename 

176 n = int(fd.readline().split()[0]) 

177 for i in range(min(n, exl.nstates)): 

178 exl.append(H2Excitation.fromstring(fd.readline())) 

179 return exl 

180 

181 def write(self, fname): 

182 with open(fname, 'w') as fd: 

183 fd.write('{0}\n'.format(len(self))) 

184 for ex in self: 

185 fd.write(ex.outstring()) 

186 

187 

188class H2Excitation(Excitation): 

189 def __eq__(self, other): 

190 """Considered to be equal when their indices are equal.""" 

191 return self.index == other.index 

192 

193 def __hash__(self): 

194 """Hash similar to __eq__""" 

195 if not hasattr(self, 'hash'): 

196 self.hash = hash(self.index) 

197 return self.hash 

198 

199 

200class H2MorseExcitedStatesAndCalculator( 

201 H2MorseExcitedStatesCalculator, H2MorseExcitedStates): 

202 """Traditional joined object for backward compatibility only""" 

203 

204 def __init__(self, calculator, nstates=3): 

205 if isinstance(calculator, str): 

206 exlist = H2MorseExcitedStates.read(calculator, nstates) 

207 else: 

208 atoms = calculator.atoms 

209 atoms.calc = calculator 

210 excalc = H2MorseExcitedStatesCalculator(nstates) 

211 exlist = excalc.calculate(atoms) 

212 H2MorseExcitedStates.__init__(self, nstates=nstates) 

213 for ex in exlist: 

214 self.append(ex)