From b6acf0fb15b5b914bcaadaa8250569fde920af5d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Teemu=20J=C3=A4rvinen?= <teemu.j.jarvinen@gmail.com>
Date: Thu, 4 May 2023 16:43:27 +0300
Subject: [PATCH] Some fixes for GPU (#114)

---
 Project.toml                      |  2 +-
 src/HKQM.jl                       |  2 +-
 src/operators.jl                  | 11 +++++++++++
 src/states.jl                     |  3 ++-
 src/submodules/ToroidalCurrent.jl |  5 -----
 src/tensor_operations.jl          |  7 +++++++
 6 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/Project.toml b/Project.toml
index fca0157..bc6810c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "HKQM"
 uuid = "95871d0a-b0d4-4f82-9ec8-411977f2e218"
 authors = ["Teemu Järvinen <teemu.jarvinen@helsinki.fi> and contributors"]
-version = "0.2.1"
+version = "0.2.2-DEV"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
diff --git a/src/HKQM.jl b/src/HKQM.jl
index ce4a683..ce483cf 100644
--- a/src/HKQM.jl
+++ b/src/HKQM.jl
@@ -136,12 +136,12 @@ include("integrations.jl")
 include("accuracytests.jl")
 include("scf.jl")
 include("initial_states.jl")
+include("submodules/ToroidalCurrent.jl")
 
 
 function __init__()
     @require TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" begin
         include("tensor_operations.jl")
-        include("submodules/ToroidalCurrent.jl")
         @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("tensor_operations_cuda.jl")
     end
 end
diff --git a/src/operators.jl b/src/operators.jl
index 18956f1..230e9de 100644
--- a/src/operators.jl
+++ b/src/operators.jl
@@ -744,6 +744,10 @@ struct HamiltonOperatorFreeParticle{T} <: AbstractHamiltonOperator
         @assert dimension(m) == dimension(u"kg")
         new{typeof(m)}(ceg, LaplaceOperator(ceg), m)
     end
+    function HamiltonOperatorFreeParticle(T, ceg::AbstractElementGrid;m=1u"me_au")
+        @assert dimension(m) == dimension(u"kg")
+        new{typeof(m)}(ceg, LaplaceOperator(T, ceg), m)
+    end
 end
 
 Unitful.unit(H::HamiltonOperatorFreeParticle) =u"hartree*bohr^2"*unit(H.∇²)
@@ -778,6 +782,13 @@ struct HamiltonOperator{TF,TV} <: AbstractHamiltonOperator
         T = HamiltonOperatorFreeParticle(get_elementgrid(V); m=m)
         new{typeof(T.m),typeof(V)}(get_elementgrid(V), T, V)
     end
+    function HamiltonOperator(Tt, V::AbstractOperator{1}; m=1u"me_au")
+        @assert dimension(V) == dimension(u"J")
+        @assert dimension(m) == dimension(u"kg")
+        T = HamiltonOperatorFreeParticle(Tt, get_elementgrid(V); m=m)
+        vv = convert_array_type(Tt, V)
+        new{typeof(T.m),typeof(vv)}(get_elementgrid(V), T, vv)
+    end
 end
 
 Unitful.unit(H::HamiltonOperator) = unit(H.T)
diff --git a/src/states.jl b/src/states.jl
index 2b26400..62c64d8 100644
--- a/src/states.jl
+++ b/src/states.jl
@@ -139,7 +139,8 @@ julia> ψ⋆ψ;
 """
 function ketbra(qs1::QuantumState{<:Any,<:Any}, qs2::QuantumState{<:Any,<:Any})
     @assert size(qs1) == size(qs2)
-    conj(qs1).psi .* qs2.psi
+    _f(x,y) = conj(x) * y
+    _f.(qs1.psi, qs2.psi)
 end
 
 const ⋆ = ketbra
diff --git a/src/submodules/ToroidalCurrent.jl b/src/submodules/ToroidalCurrent.jl
index 1ecff99..80e0580 100644
--- a/src/submodules/ToroidalCurrent.jl
+++ b/src/submodules/ToroidalCurrent.jl
@@ -6,7 +6,6 @@ using DelimitedFiles
 using Distributed
 using Distances
 using Interpolations
-using TensorOperations
 
 using ..HKQM
 
@@ -192,10 +191,6 @@ function greens_function_2d(ceg; eps=1E-6u"Å")
     return reshape(out, i,j, i,j, i,j, i,j)
 end
 
-function solve_2d_poisson(G, psi)
-    @tensoropt tmp[x,y,z,X,Y,Z] := G[x,X,y,Y,i,I,j,J] * psi.psi[i,j,z,I,J,Z]
-    return QuantumState(get_elementgrid(psi), tmp, unit(psi)*u"bohr^2")
-end
 
 function solve_2d_poisson(psi)
     return solve_2d_poisson( greens_function_2d( get_elementgrid(psi) ) , psi)
diff --git a/src/tensor_operations.jl b/src/tensor_operations.jl
index 752b2d6..a7fe0bb 100644
--- a/src/tensor_operations.jl
+++ b/src/tensor_operations.jl
@@ -76,4 +76,11 @@ end
 function operate_z!(dψ::AbstractArray{<:Any,6}, dt::DerivativeTensor, ψ::Array{<:Any,6})
     @tensor dψ[i,j,k,I,J,K] = dt.values[k,l] * ψ[i,j,l,I,J,K]
     return dψ
+end
+
+##
+
+function ToroidalCurrent.solve_2d_poisson(G, psi)
+    @tensoropt tmp[x,y,z,X,Y,Z] := G[x,X,y,Y,i,I,j,J] * psi.psi[i,j,z,I,J,Z]
+    return QuantumState(get_elementgrid(psi), tmp, unit(psi)*u"bohr^2")
 end
\ No newline at end of file