-
Notifications
You must be signed in to change notification settings - Fork 89
/
local.jl
130 lines (109 loc) · 4.19 KB
/
local.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
## Local potentials. Can be provided from external potentials, or from `model.atoms`.
# a local potential term. Must have the field `potential`, storing the
# potential in real space on the grid. If the potential is different in the α and β
# components then it should be a 4d-array with the last axis running over the
# two spin components.
abstract type TermLocalPotential <: Term end
@timing "ene_ops: local" function ene_ops(term::TermLocalPotential, ψ, occ; kwargs...)
basis = term.basis
T = eltype(basis)
potview(data, spin) = ndims(data) == 4 ? (@view data[:, :, :, spin]) : data
ops = [RealSpaceMultiplication(basis, kpoint, potview(term.potential, kpoint.spin))
for kpoint in basis.kpoints]
if :ρ in keys(kwargs)
E = sum(total_density(kwargs[:ρ]) .* term.potential) * term.basis.dvol
else
E = T(Inf)
end
(E=E, ops=ops)
end
## External potentials
struct TermExternal <: TermLocalPotential
basis::PlaneWaveBasis
potential::AbstractArray
end
"""
External potential from an analytic function `V` (in cartesian coordinates).
No low-pass filtering is performed.
"""
struct ExternalFromReal{T <: Function}
V::T
end
function (external::ExternalFromReal)(basis::PlaneWaveBasis{T}) where {T}
potential = external.V.(r_vectors_cart(basis))
TermExternal(basis, potential)
end
"""
External potential from the (unnormalized) Fourier coefficients `V(G)`
G is passed in cartesian coordinates
"""
struct ExternalFromFourier{T <: Function}
V::T
end
function (external::ExternalFromFourier)(basis::PlaneWaveBasis)
unit_cell_volume = basis.model.unit_cell_volume
pot_fourier = [complex(external.V(G) / sqrt(unit_cell_volume))
for G in G_vectors_cart(basis)]
pot_real = G_to_r(basis, pot_fourier)
TermExternal(basis, real(pot_real))
end
## Atomic local potential
struct TermAtomicLocal <: TermLocalPotential
basis::PlaneWaveBasis
potential::AbstractArray
end
"""
Atomic local potential defined by `model.atoms`.
"""
struct AtomicLocal end
function (E::AtomicLocal)(basis::PlaneWaveBasis{T}) where {T}
model = basis.model
# pot_fourier is <e_G|V|e_G'> expanded in a basis of e_{G-G'}
# Since V is a sum of radial functions located at atomic
# positions, this involves a form factor (`local_potential_fourier`)
# and a structure factor e^{-i Gr}
pot_fourier = zeros(Complex{T}, basis.fft_size)
for (iG, G) in enumerate(G_vectors(basis))
pot = zero(T)
for (elem, positions) in model.atoms
form_factor::T = local_potential_fourier(elem, norm(model.recip_lattice * G))
for r in positions
pot += cis(-2T(π) * dot(G, r)) * form_factor
end
end
pot_fourier[iG] = pot / sqrt(model.unit_cell_volume)
end
pot_real = G_to_r(basis, pot_fourier)
TermAtomicLocal(basis, real(pot_real))
end
@timing "forces: local" function compute_forces(term::TermAtomicLocal, ψ, occ; ρ, kwargs...)
T = eltype(term.basis)
atoms = term.basis.model.atoms
recip_lattice = term.basis.model.recip_lattice
unit_cell_volume = term.basis.model.unit_cell_volume
ρ_fourier = r_to_G(term.basis, total_density(ρ))
# energy = sum of form_factor(G) * struct_factor(G) * rho(G)
# where struct_factor(G) = cis(-2π G⋅r)
forces = [zeros(Vec3{T}, length(positions)) for (el, positions) in atoms]
for (iel, (el, positions)) in enumerate(atoms)
form_factors = [Complex{T}(local_potential_fourier(el, norm(recip_lattice * G)))
for G in G_vectors(term.basis)]
for (ir, r) in enumerate(positions)
forces[iel][ir] = _force_local_internal(term.basis, ρ_fourier, form_factors, r)
end
end
forces
end
# function barrier to work around various type instabilities
function _force_local_internal(basis, ρ_fourier, form_factors, r)
T = real(eltype(ρ_fourier))
f = zero(Vec3{T})
for (iG, G) in enumerate(G_vectors(basis))
f -= real(conj(ρ_fourier[iG])
.* form_factors[iG]
.* cis(-2T(π) * dot(G, r))
.* (-2T(π)) .* G .* im
./ sqrt(basis.model.unit_cell_volume))
end
f
end