From b5f422428bcaac5acf36705e84cddd76de1216cf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89tienne=20Polack?= <polack@hollved.net>
Date: Fri, 21 Apr 2023 11:02:29 +0200
Subject: [PATCH] hamiltonian locally

---
 src/scf/newton.jl        |  8 +++-----
 src/terms/Hamiltonian.jl | 10 ++++------
 src/terms/operators.jl   |  7 ++++---
 3 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/src/scf/newton.jl b/src/scf/newton.jl
index 28d02b31b5..e3835bc95a 100644
--- a/src/scf/newton.jl
+++ b/src/scf/newton.jl
@@ -59,8 +59,7 @@ end
 # Projections on the space tangent to ψ
 function proj_tangent_kpt!(δψk, ψk)
     # δψk = δψk - ψk * (ψk'δψk)
-    δψk_flat = reshape(δψk, :, size(δψk)[end])
-    mul!(δψk_flat, ψk[1, :, :], ψk[1, :, :]'δψk_flat, -1, 1)
+    mul!(δψk, ψk, ψk'δψk, -1, 1)
 end
 proj_tangent_kpt(δψk, ψk) = proj_tangent_kpt!(deepcopy(δψk), ψk)
 
@@ -84,7 +83,6 @@ function newton(basis::PlaneWaveBasis{T}, ψ0;
                 tol=1e-6, tol_cg=tol / 100, maxiter=20,
                 callback=ScfDefaultCallback(),
                 is_converged=ScfConvergenceDensity(tol)) where {T}
-    @assert basis.model.n_components == 1
 
     # setting parameters
     model = basis.model
@@ -95,11 +93,11 @@ function newton(basis::PlaneWaveBasis{T}, ψ0;
     filled_occ = filled_occupation(model)
     n_spin = model.n_spin_components
     n_bands = div(model.n_electrons, n_spin * filled_occ, RoundUp)
-    @assert n_bands == size(ψ0[1], 3)
+    @assert n_bands == size(ψ0[1], 2)
 
     # number of kpoints and occupation
     Nk = length(basis.kpoints)
-    occupation = [filled_occ * ones(T, n_bands) for _ in 1:Nk]
+    occupation = [filled_occ * ones(T, n_bands) for ik = 1:Nk]
 
     # iterators
     n_iter = 0
diff --git a/src/terms/Hamiltonian.jl b/src/terms/Hamiltonian.jl
index ddb6979ee1..4248e20464 100644
--- a/src/terms/Hamiltonian.jl
+++ b/src/terms/Hamiltonian.jl
@@ -157,12 +157,10 @@ end
 
             if have_divAgrad
                 @timeit to "divAgrad" begin
-                    for σ in 1:size(Hψ, 1)
-                        apply!((fourier=Hψ[σ, :, iband], real=nothing),
-                               H.divAgrad_op,
-                               (fourier=ψ[σ, :, iband], real=nothing),
-                               ψ_real[σ, :, :, :])  # ψ_real used as scratch
-                    end
+                    apply!((fourier=Hψ[:, :, iband], real=nothing),
+                           H.divAgrad_op,
+                           (fourier=ψ[:, :, iband], real=nothing),
+                           ψ_real)  # ψ_real used as scratch
                 end
             end
 
diff --git a/src/terms/operators.jl b/src/terms/operators.jl
index 5dce5a757a..01be496976 100644
--- a/src/terms/operators.jl
+++ b/src/terms/operators.jl
@@ -141,14 +141,15 @@ struct DivAgradOperator{T <: Real, AT} <: RealFourierOperator
     A::AT
 end
 function apply!(Hψ, op::DivAgradOperator, ψ,
-                ψ_scratch=zeros(complex(eltype(op.basis)), op.basis.fft_size...))
+                ψ_scratch=zeros(complex(eltype(op.basis)), op.basis.model.n_components, op.basis.fft_size...))
+    @assert op.basis.model.n_components == 1
     # 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 = ifft!(ψ_scratch, op.basis, op.kpoint, im .* G_plus_k[α] .* ψ.fourier)
+        ∂αψ_real = ifft!(ψ_scratch[1, :, :, :], op.basis, op.kpoint, im .* G_plus_k[α] .* ψ.fourier[1, :])
         A∇ψ      = fft(op.basis, op.kpoint, ∂αψ_real .* op.A)
-        Hψ.fourier .-= im .* G_plus_k[α] .* A∇ψ ./ 2
+        Hψ.fourier[1, :, :] .-= im .* G_plus_k[α] .* A∇ψ ./ 2
 
     end
 end