1struct Kinetic1D{T<:Real}
  2    basis::PlaneWaveBasis{T}
  3    kpoint::Kpoint{T}
  4end
  5function Base.Matrix(op::Kinetic1D)
  6    basis = op.basis
  7    kpoint = op.kpoint
  8
  9    n_G = length(G_vectors(basis, kpoint))
 10    TC = complex(eltype(basis))
 11    H = zeros(TC, n_G, n_G)
 12    for (i, Gpk) in enumerate(Gplusk_vectors_cart(basis, kpoint))
 13        H[i, i] = Gpk[1]
 14    end
 15    H
 16end
 17function kinetic1d(basis, ukX, kpt, i, j)
 18    ui = ukX[:, i]
 19    uj = ukX[:, j]
 20    Gk_vecs = [Gk[1] for Gk in G_vectors_cart(basis, kpt)]
 21    braket = sum(Gk_vecs .* conj(ui) .* uj)
 22    if i == j
 23        braket += (basis.model.recip_lattice * kpt.coordinate)[1, 1, 1]
 24    end
 25    braket
 26end
 27
 28struct Potential1D{T<:Real, A<:AbstractArray{Complex{T}, 3}}
 29    basis::PlaneWaveBasis{T}
 30    kpoint::Kpoint{T}
 31    # Specific storage
 32    potential::A
 33end
 34function compute_v2_four(basis::PlaneWaveBasis{T}; positions=basis.model.positions,
 35                         q=zero(Vec3{T})) where {T}
 36    # pot_fourier is <e_G|V|e_G'> expanded in a basis of e_{G-G'}
 37    # Since V is a sum of radial functions located at atomic
 38    # positions, this involves a form factor (`local_potential_fourier`)
 39    # and a structure factor e^{-i G·r}
 40    model = basis.model
 41    Gqs_cart = [model.recip_lattice * (G + q) for G in G_vectors(basis)]
 42    # TODO Bring Gqs_cart on the CPU for compatibility with the pseudopotentials which
 43    #      are not isbits ... might be able to solve this by restructuring the loop
 44    disregistry_positions = [[i] for i in length(positions)÷2+1:length(positions)]
 45
 46    # Pre-compute the form factors at unique values of |G| to speed up
 47    # the potential Fourier transform (by a lot). Using a hash map gives O(1)
 48    # lookup.
 49    form_factors = IdDict{Tuple{Int,T},T}()  # IdDict for Dual compatibility
 50    for G in Gqs_cart
 51        p = norm(G)
 52        for (igroup, group) in enumerate(disregistry_positions)
 53            if !haskey(form_factors, (igroup, p))
 54                element = model.atoms[first(group)]
 55                form_factors[(igroup, p)] = DFTK.local_potential_fourier(element, p)
 56            end
 57        end
 58    end
 59
 60    Gqs = [G + q for G in G_vectors(basis)]  # TODO Again for GPU compatibility
 61    pot_fourier_1D = map(enumerate(Gqs)) do (iG, G)
 62        p = norm(Gqs_cart[iG])
 63        pot = sum(enumerate(disregistry_positions)) do (igroup, group)
 64            structure_factor = sum(r -> DFTK.cis2pi(-dot(G, r)), @view positions[group])
 65            form_factors[(igroup, p)] * structure_factor
 66        end
 67        pot / sqrt(model.unit_cell_volume)
 68    end
 69    # Explicitly reshape the 1D vector to the 3D grid size required by DFTK's FFT routines.
 70    return reshape(pot_fourier_1D, basis.fft_size)
 71end
 72function Potential1D(basis::PlaneWaveBasis{T}, kpoint::Kpoint{T}, ε::T, basis_fn::Function) where {T}
 73    basis = basis_fn(ε)
 74    V2_four = compute_v2_four(basis)
 75    prefac_m = map(G_vectors(basis)) do G
 76        -2π * im * G[1]
 77    end
 78    V2_four_m = zeros(complex(T), size(V2_four))
 79    for i in axes(V2_four, 1)
 80        V2_four_m[i, :, :] = V2_four[i, :, :] .* prefac_m[i, :, :]
 81    end
 82    potential = ifft(basis, V2_four_m)
 83
 84    Potential1D(basis, kpoint, potential)
 85end
 86function Base.Matrix(op::Potential1D{T}) where {T}
 87    # V(G, G') = <eG|V|eG'> = 1/sqrt(Ω) <e_{G-G'}|V>
 88    pot_fourier = fft(op.basis, op.potential)
 89    n_G = length(G_vectors(op.basis, op.kpoint))
 90    H = zeros(complex(eltype(op.basis)), n_G, n_G)
 91    for (j, G′) in enumerate(G_vectors(op.basis, op.kpoint))
 92        for (i, G) in enumerate(G_vectors(op.basis, op.kpoint))
 93            # G_vectors(basis)[ind_ΔG] = G - G'
 94            ind_ΔG = DFTK.index_G_vectors(op.basis, G - G′)
 95            if isnothing(ind_ΔG)
 96                error("For full matrix construction, the FFT size must be " *
 97                      "large enough so that Hamiltonian applications are exact")
 98            end
 99            H[i, j] = pot_fourier[ind_ΔG] / sqrt(op.basis.model.unit_cell_volume)
100        end
101    end
102    H
103end
104function compute_δv2(basis::PlaneWaveBasis{T}) where {T}
105    V2_four = compute_v2_four(basis)
106    prefac_m = map(G -> -2π * im * G[1], G_vectors(basis))
107    # Reshape the 1D prefactor to the 3D grid size for consistent broadcasting
108    prefac_m_3D = reshape(prefac_m, basis.fft_size)
109
110    # Now this is a correct 3D broadcast operation
111    V2_four_m = V2_four .* prefac_m_3D
112
113    return ifft(basis, V2_four_m)[:, 1, 1]
114end
115function compute_δ²v2(basis::PlaneWaveBasis{T}) where {T}
116    V2_four = compute_v2_four(basis)
117    prefac_m = map(G -> -2π * im * G[1], G_vectors(basis))
118    # Reshape the 1D prefactor to the 3D grid size for consistent broadcasting
119    prefac_m_3D = reshape(prefac_m, basis.fft_size)
120
121    # Now this is a correct 3D broadcast operation
122    V2_four_m = V2_four .* prefac_m_3D .* prefac_m_3D
123
124    return ifft(basis, V2_four_m)[:, 1, 1]
125end
126function potential1d(basis::PlaneWaveBasis{T}, ukX::AbstractArray, kpt::Kpoint{T}, potential::AbstractArray, i::Int, j::Int; ui::AbstractArray, uj::AbstractArray) where {T}
127    ifft!(ui, basis, kpt, ukX[:, i])
128    ifft!(uj, basis, kpt, ukX[:, j])
129    sum(potential .* conj(ui) .* uj) ./ length(ui)
130end
131function potential1d_nos(basis::PlaneWaveBasis{T}, ukX::AbstractArray, kpt::Kpoint{T}, potential::AbstractArray, i::Int, j::Int) where {T}
132    ui = ifft(basis, kpt, ukX[:, i])
133    uj = ifft(basis, kpt, ukX[:, j])
134    sum(potential .* conj(ui) .* uj) ./ length(ui)
135end