-
Notifications
You must be signed in to change notification settings - Fork 83
/
orbitals.jl
86 lines (75 loc) · 3.68 KB
/
orbitals.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
using Random # Used to have a generic API for CPU and GPU computations alike: see random_orbitals
# Returns the occupied orbitals, the occupation array and optionally the eigenvalues without
# virtual states (or states with small occupation level for metals).
# threshold is a parameter to distinguish between states we want to keep and the
# others when using temperature. It is set to 0.0 by default, to treat with insulators.
function select_occupied_orbitals(basis, ψ, occupation; threshold=0.0)
N = [something(findlast(x -> x > threshold, occk), 0) for occk in occupation]
selected_ψ = [@view ψk[:, 1:N[ik]] for (ik, ψk) in enumerate(ψ)]
selected_occ = [ occk[1:N[ik]] for (ik, occk) in enumerate(occupation)]
# if we have an insulator, sanity check that the orbitals we kept are the
# occupied ones
if iszero(threshold)
model = basis.model
n_spin = model.n_spin_components
n_bands = div(model.n_electrons, n_spin * filled_occupation(model), RoundUp)
@assert n_bands == size(selected_ψ[1], 2)
end
(; ψ=selected_ψ, occupation=selected_occ)
end
# Packing routine to store all bands ψ using same-sized matrices.
# This is useful for serialisation to disk
function blockify_ψ(basis::PlaneWaveBasis, ψ::AbstractVector)
# Determine largest rectangular grid on which all bands can live
n_G_vectors = [length(kpt.mapping) for kpt in basis.kpoints]
max_n_G = mpi_max(maximum(n_G_vectors), basis.comm_kpts)
# Zero-pad each ψk to the full stored rectangle
ψ = map(ψ) do ψk
n_G, n_bands = size(ψk)
ψk_full = similar(ψk, max_n_G, n_bands)
ψk_full[1:n_G, :] = ψk
ψk_full
end
(; ψ, n_G_vectors)
end
# Unpacking routine for the above by truncating each array to the sizes
# given in the second argument.
function unblockify_ψ(ψ_padded::AbstractVector, n_G_vectors::AbstractVector)
map(n_G_vectors, ψ_padded) do n_Gk, ψk_padded
ψk_padded[1:n_Gk, :]
end
end
# Packing routines used in direct_minimization and newton algorithms.
# They pack / unpack sets of ψ's (or compatible arrays, such as hamiltonian
# applies and gradients) to make them compatible to be used in algorithms
# from IterativeSolvers.
# Some care is needed here : some operators (for instance K in newton.jl)
# are real-linear but not complex-linear. To overcome this difficulty, instead of
# seeing them as operators from C^N to C^N, we see them as
# operators from R^2N to R^2N. In practice, this is done with the
# reinterpret function from julia.
# /!\ pack_ψ does not share memory while unpack_ψ does
reinterpret_real(x) = reinterpret(real(eltype(x)), x)
reinterpret_complex(x) = reinterpret(Complex{eltype(x)}, x)
function pack_ψ(ψ)
# TODO as an optimization, do that lazily? See LazyArrays
reduce(vcat, [vec(ψk) for ψk in ψ])
end
# Returns pointers into the unpacked ψ
# /!\ The resulting array is only valid as long as the original x is still in live memory.
function unsafe_unpack_ψ(x, sizes_ψ)
lengths = prod.(sizes_ψ)
ends = cumsum(lengths)
# We unsafe_wrap the resulting array to avoid a complicated type for ψ.
map(1:length(sizes_ψ)) do ik
unsafe_wrap(Array{complex(eltype(x))},
pointer(@views x[ends[ik]-lengths[ik]+1:ends[ik]]),
sizes_ψ[ik])
end
end
unpack_ψ(x, sizes_ψ) = deepcopy(unsafe_unpack_ψ(x, sizes_ψ))
function random_orbitals(basis::PlaneWaveBasis{T}, kpt::Kpoint, howmany::Integer) where {T}
orbitals = similar(basis.G_vectors, Complex{T}, length(G_vectors(basis, kpt)), howmany)
randn!(TaskLocalRNG(), orbitals) # use the RNG on the device if we're using a GPU
ortho_qr(orbitals)
end