/
operators.jl
165 lines (154 loc) · 6.08 KB
/
operators.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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
### Linear operators operating on quantities in real or fourier space
# This is the optimized low-level interface. Use the functions in Hamiltonian for high-level usage
"""
Linear operators that act on tuples (real, fourier)
The main entry point is `apply!(out, op, in)` which performs the operation out += op*in
where out and in are named tuples (real, fourier)
They also implement mul! and Matrix(op) for exploratory use.
"""
abstract type RealFourierOperator end
# RealFourierOperator contain fields `basis` and `kpoint`
# Unoptimized fallback, intended for exploratory use only.
# For performance, call through Hamiltonian which saves FFTs.
function LinearAlgebra.mul!(Hψ::AbstractVector, op::RealFourierOperator, ψ::AbstractVector)
ψ_real = G_to_r(op.basis, op.kpoint, ψ)
Hψ_fourier = similar(ψ)
Hψ_real = similar(ψ_real)
Hψ_fourier .= 0
Hψ_real .= 0
apply!((real=Hψ_real, fourier=Hψ_fourier), op, (real=ψ_real, fourier=ψ))
Hψ .= Hψ_fourier .+ r_to_G(op.basis, op.kpoint, Hψ_real)
Hψ
end
function LinearAlgebra.mul!(Hψ::AbstractMatrix, op::RealFourierOperator, ψ::AbstractMatrix)
@views for i = 1:size(ψ, 2)
mul!(Hψ[:, i], op, ψ[:, i])
end
Hψ
end
Base.:*(op::RealFourierOperator, ψ) = mul!(similar(ψ), op, ψ)
"""
Noop operation: don't do anything.
Useful for energy terms that don't depend on the orbitals at all (eg nuclei-nuclei interaction).
"""
struct NoopOperator{T <: Real} <: RealFourierOperator
basis::PlaneWaveBasis{T}
kpoint::Kpoint{T}
end
apply!(Hψ, op::NoopOperator, ψ) = nothing
function Matrix(op::NoopOperator)
n_Gk = length(G_vectors(op.basis, op.kpoint))
zeros(eltype(op.basis), n_Gk, n_Gk)
end
"""
Real space multiplication by a potential:
(Hψ)(r) V(r) ψ(r)
"""
struct RealSpaceMultiplication{T <: Real, AT <: AbstractArray} <: RealFourierOperator
basis::PlaneWaveBasis{T}
kpoint::Kpoint{T}
potential::AT
end
function apply!(Hψ, op::RealSpaceMultiplication, ψ)
Hψ.real .+= op.potential .* ψ.real
end
function Matrix(op::RealSpaceMultiplication)
# V(G, G') = <eG|V|eG'> = 1/sqrt(Ω) <e_{G-G'}|V>
pot_fourier = r_to_G(op.basis, op.potential)
n_G = length(G_vectors(op.basis, op.kpoint))
H = zeros(complex(eltype(op.basis)), n_G, n_G)
for (i, G) in enumerate(G_vectors(op.basis, op.kpoint))
for (j, G′) in enumerate(G_vectors(op.basis, op.kpoint))
# G_vectors(basis)[ind_ΔG] = G - G′
ind_ΔG = index_G_vectors(op.basis, G - G′)
if isnothing(ind_ΔG)
error("For full matrix construction, the FFT size must be " *
"large enough so that Hamiltonian applications are exact")
end
H[i, j] = pot_fourier[ind_ΔG] / sqrt(op.basis.model.unit_cell_volume)
end
end
H
end
"""
Fourier space multiplication, like a kinetic energy term:
(Hψ)(G) = multiplier(G) ψ(G)
"""
struct FourierMultiplication{T <: Real, AT <: AbstractArray} <: RealFourierOperator
basis::PlaneWaveBasis{T}
kpoint::Kpoint{T}
multiplier::AT
end
function apply!(Hψ, op::FourierMultiplication, ψ)
Hψ.fourier .+= op.multiplier .* ψ.fourier
end
Matrix(op::FourierMultiplication) = Array(Diagonal(op.multiplier))
"""
Nonlocal operator in Fourier space in Kleinman-Bylander format,
defined by its projectors P matrix and coupling terms D:
Hψ = PDP' ψ
"""
struct NonlocalOperator{T <: Real, PT, DT} <: RealFourierOperator
basis::PlaneWaveBasis{T}
kpoint::Kpoint{T}
# not typed, can be anything that supports PDP'ψ
P::PT
D::DT
end
function apply!(Hψ, op::NonlocalOperator, ψ)
Hψ.fourier .+= op.P * (op.D * (op.P' * ψ.fourier))
end
Matrix(op::NonlocalOperator) = op.P * op.D * op.P'
"""
Magnetic field operator A⋅(-i∇).
"""
struct MagneticFieldOperator{T <: Real, AT} <: RealFourierOperator
basis::PlaneWaveBasis{T}
kpoint::Kpoint{T}
Apot::AT # Apot[α][i,j,k] is the A field in (cartesian) direction α
end
function apply!(Hψ, op::MagneticFieldOperator, ψ)
# TODO this could probably be better optimized
for α = 1:3
iszero(op.Apot[α]) && continue
pα = [Gk[α] for Gk in Gplusk_vectors_cart(op.basis, op.kpoint)]
∂αψ_fourier = pα .* ψ.fourier
∂αψ_real = G_to_r(op.basis, op.kpoint, ∂αψ_fourier)
Hψ.real .+= op.Apot[α] .* ∂αψ_real
end
end
# TODO Implement Matrix(op::MagneticFieldOperator)
@doc raw"""
Nonlocal "divAgrad" operator ``-½ ∇ ⋅ (A ∇)`` where ``A`` is a scalar field on the
real-space grid. The ``-½`` is included, such that this operator is a generalisation of the
kinetic energy operator (which is obtained for ``A=1``).
"""
struct DivAgradOperator{T <: Real, AT} <: RealFourierOperator
basis::PlaneWaveBasis{T}
kpoint::Kpoint{T}
A::AT
end
function apply!(Hψ, op::DivAgradOperator, ψ,
ψ_scratch=zeros(complex(eltype(op.basis)),
op.basis.fft_size...))
# TODO: Performance improvements: Unscaled plans, avoid remaining allocations
# (which are only on the small k-point-specific Fourier grid
G_plus_k = [[Gk[α] for Gk in Gplusk_vectors_cart(op.basis, op.kpoint)] for α in 1:3]
for α = 1:3
∂αψ_real = G_to_r!(ψ_scratch, op.basis, op.kpoint, im .* G_plus_k[α] .* ψ.fourier)
A∇ψ = r_to_G(op.basis, op.kpoint, ∂αψ_real .* op.A)
Hψ.fourier .-= im .* G_plus_k[α] .* A∇ψ ./ 2
end
end
# TODO Implement Matrix(op::DivAgrad)
# Optimize RFOs by combining terms that can be combined
function optimize_operators_(ops)
ops = [op for op in ops if !(op isa NoopOperator)]
RSmults = [op for op in ops if op isa RealSpaceMultiplication]
isempty(RSmults) && return ops
nonRSmults = [op for op in ops if !(op isa RealSpaceMultiplication)]
combined_RSmults = RealSpaceMultiplication(RSmults[1].basis,
RSmults[1].kpoint,
sum([op.potential for op in RSmults]))
[nonRSmults..., combined_RSmults]
end