1from ase.units import Bohr, AUT, _me, _amu
2from gpaw.tddft.units import attosec_to_autime
3from gpaw.forces import calculate_forces
4
5###############################################################################
6# EHRENFEST DYNAMICS WITHIN THE PAW METHOD
7# WORKS WITH PAW AS LONG THERE ISN'T TOO
8# TOO MUCH OVERLAP BETWEEN THE SPHERES
9# SUPPORTS ALSO HGH PSEUDOPOTENTIALS
10###############################################################################
11
12# m a(t+dt)   = F[psi(t),x(t)]
13# x(t+dt/2)   = x(t) + v(t) dt/2 + .5 a(t) (dt/2)^2
14# vh(t+dt/2)  = v(t) + .5 a(t) dt/2
15# m a(t+dt/2) = F[psi(t),x(t+dt/2)]
16# v(t+dt/2)   = vh(t+dt/2) + .5 a(t+dt/2) dt/2
17#
18# psi(t+dt)   = U(t,t+dt) psi(t)
19#
20# m a(t+dt/2) = F[psi(t+dt),x(t+dt/2)]
21# x(t+dt)     = x(t+dt/2) + v(t+dt/2) dt/2 + .5 a(t+dt/2) (dt/2)^2
22# vh(t+dt)    = v(t+dt/2) + .5 a(t+dt/2) dt/2
23# m a(t+dt)   = F[psi(t+dt),x(t+dt)]
24# v(t+dt)     = vh(t+dt/2) + .5 a(t+dt/2) dt/2
25
26# TODO: move force corrections from forces.py to this module, as well as
27# the cg method for calculating the inverse of S from overlap.py
28
29
30class EhrenfestVelocityVerlet:
31
32    def __init__(self, calc, mass_scale=1.0, setups='paw'):
33        """Initializes the Ehrenfest MD calculator.
34
35        Parameters
36        ----------
37
38        calc: TDDFT Object
39
40        mass_scale: 1.0
41            Scaling coefficient for atomic masses
42
43        setups: {'paw', 'hgh'}
44            Type of setups to use for the calculation
45
46        Note
47        ------
48
49        Use propagator = 'EFSICN' for when creating the TDDFT object from a
50        PAW ground state calculator and propagator = 'EFSICN_HGH' for HGH
51        pseudopotentials
52
53        """
54        self.calc = calc
55        self.setups = setups
56        self.x = self.calc.atoms.positions.copy() / Bohr
57        self.xn = self.x.copy()
58        self.v = self.x.copy()
59        amu_to_aumass = _amu / _me
60        if self.calc.atoms.get_velocities() is not None:
61            self.v = self.calc.atoms.get_velocities() / (Bohr / AUT)
62        else:
63            self.v[:] = 0.0
64            self.calc.atoms.set_velocities(self.v)
65
66        self.vt = self.v.copy()
67        self.vh = self.v.copy()
68        self.time = 0.0
69
70        self.M = calc.atoms.get_masses() * amu_to_aumass * mass_scale
71
72        self.a = self.v.copy()
73        self.ah = self.a.copy()
74        self.an = self.a.copy()
75        self.F = self.a.copy()
76
77        self.calc.get_td_energy()
78        self.F = self.get_forces()
79
80        for i in range(len(self.F)):
81            self.a[i] = self.F[i] / self.M[i]
82
83    def get_forces(self):
84        return calculate_forces(self.calc.wfs,
85                                self.calc.td_density.get_density(),
86                                self.calc.td_hamiltonian.hamiltonian,
87                                self.calc.log)
88
89    def propagate(self, dt):
90        """Performs one Ehrenfest MD propagation step
91
92        Parameters
93        ----------
94
95        dt: scalar
96            Time step (in attoseconds) used for the Ehrenfest MD step
97
98        """
99        self.x = self.calc.atoms.positions.copy() / Bohr
100        self.v = self.calc.atoms.get_velocities() / (Bohr / AUT)
101
102        dt = dt * attosec_to_autime
103
104        # m a(t+dt)   = F[psi(t),x(t)]
105        self.calc.atoms.positions = self.x * Bohr
106        self.calc.set_positions(self.calc.atoms)
107        self.calc.get_td_energy()
108        self.F = self.get_forces()
109
110        for i in range(len(self.F)):
111            self.a[i] = self.F[i] / self.M[i]
112
113        # x(t+dt/2)   = x(t) + v(t) dt/2 + .5 a(t) (dt/2)^2
114        # vh(t+dt/2)  = v(t) + .5 a(t) dt/2
115        self.xh = self.x + self.v * dt / 2 + .5 * self.a * dt / 2 * dt / 2
116        self.vhh = self.v + .5 * self.a * dt / 2
117
118        # m a(t+dt/2) = F[psi(t),x(t+dt/2)a]
119        self.calc.atoms.positions = self.xh * Bohr
120        self.calc.set_positions(self.calc.atoms)
121        self.calc.get_td_energy()
122        self.F = self.get_forces()
123
124        for i in range(len(self.F)):
125            self.ah[i] = self.F[i] / self.M[i]
126
127        # v(t+dt/2)   = vh(t+dt/2) + .5 a(t+dt/2) dt/2
128        self.vh = self.vhh + 0.5 * self.ah * dt / 2
129
130        # Propagate wf
131        # psi(t+dt)   = U(t,t+dt) psi(t)
132        self.propagate_single(dt)
133
134        # m a(t+dt/2) = F[psi(t+dt),x(t+dt/2)]
135        self.calc.atoms.positions = self.xh * Bohr
136        self.calc.set_positions(self.calc.atoms)
137        self.calc.get_td_energy()
138        self.F = self.get_forces()
139
140        for i in range(len(self.F)):
141            self.ah[i] = self.F[i] / self.M[i]
142
143        # x(t+dt)     = x(t+dt/2) + v(t+dt/2) dt/2 + .5 a(t+dt/2) (dt/2)^2
144        # vh(t+dt)    = v(t+dt/2) + .5 a(t+dt/2) dt/2
145        self.xn = self.xh + self.vh * dt / 2 + 0.5 * self.ah * dt / 2 * dt / 2
146        self.vhh = self.vh + .5 * self.ah * dt / 2
147
148        # m a(t+dt)   = F[psi(t+dt),x(t+dt)]
149        self.calc.atoms.positions = self.xn * Bohr
150        self.calc.set_positions(self.calc.atoms)
151        self.calc.get_td_energy()
152        self.calc.update_eigenvalues()
153        self.F = self.get_forces()
154
155        for i in range(len(self.F)):
156            self.an[i] = self.F[i] / self.M[i]
157
158        # v(t+dt)     = vh(t+dt/2) + .5 a(t+dt/2) dt/2
159        self.vn = self.vhh + .5 * self.an * dt / 2
160
161        # update
162        self.x[:] = self.xn
163        self.v[:] = self.vn
164        self.a[:] = self.an
165
166        # update atoms
167        self.calc.atoms.set_positions(self.x * Bohr)
168        self.calc.atoms.set_velocities(self.v * Bohr / AUT)
169
170    def propagate_single(self, dt):
171        if self.setups == 'paw':
172            self.calc.propagator.propagate(self.time, dt, self.vh)
173        else:
174            self.calc.propagator.propagate(self.time, dt)
175
176    def get_energy(self):
177        """Updates kinetic, electronic and total energies"""
178
179        self.Ekin = 0.5 * (self.M * (self.v**2).sum(axis=1)).sum()
180        self.e_coulomb = self.calc.get_td_energy()
181        self.Etot = self.Ekin + self.e_coulomb
182        return self.Etot
183
184    def get_velocities_in_au(self):
185        return self.v
186
187    def set_velocities_in_au(self, v):
188        self.v[:] = v
189        self.calc.atoms.set_velocities(v * Bohr / AUT)
190