Add density profile ansatz

src/Gutzwiller.jl

@@ -27,6 +27,8 @@ export MultinomialAnsatz
 include("ansatz/multinomial.jl")
 export JastrowAnsatz, RelativeJastrowAnsatz
 include("ansatz/jastrow.jl")
+export DensityProfileAnsatz
+include("ansatz/densityprofile.jl")
 export ExtendedAnsatz
 include("ansatz/combination.jl")
 

src/ansatz/densityprofile.jl

@@ -0,0 +1,28 @@
+"""
+    DensityProfileAnsatz(hamiltonian) <: AbstractAnsatz
+
+```math
+D(|f⟩; p) = exp(-∑_{i=1}^M p_i ⟨f|n_i|f⟩)
+```
+"""
+struct DensityProfileAnsatz{A,N,H} <: AbstractAnsatz{A,Float64,N}
+    hamiltonian::H
+end
+
+function DensityProfileAnsatz(hamiltonian)
+    address = starting_address(hamiltonian)
+    @assert address isa SingleComponentFockAddress
+    N = num_modes(address)
+    return DensityProfileAnsatz{typeof(address),N,typeof(hamiltonian)}(hamiltonian)
+end
+
+Rimu.build_basis(dp::DensityProfileAnsatz) = build_basis(dp.hamiltonian)
+
+function val_and_grad(dp::DensityProfileAnsatz, addr, params)
+    o = onr(addr)
+    exponent = dot(o, params)
+    val = exp(-exponent)
+    grad = -val * o
+    return val, grad
+end
+(dp::DensityProfileAnsatz)(addr, params) = exp(-dot(onr(addr), params))

src/ansatz/jastrow.jl

@@ -5,35 +5,34 @@ using Rimu.Hamiltonians: circshift_dot
     JastrowAnsatz(hamiltonian) <: AbstractAnsatz
 
 ```math
-J_i = exp(-∑_{k=1}^M ∑_{l=k}^M p_{k,l} n_k n_l)
+J(|f⟩; p) = exp(-∑_{k=1}^M ∑_{l=k}^M p_{k,l} ⟨f|n_k n_l|f⟩)
 ```
 
 With translationally invariant Hamiltonians, use [`RelativeJastrowAnsatz`](@ref) instead.
 """
-struct JastrowAnsatz{A,T,N,H} <: AbstractAnsatz{A,T,N}
+struct JastrowAnsatz{A,N,H} <: AbstractAnsatz{A,Float64,N}
     hamiltonian::H
 end
 
 function JastrowAnsatz(hamiltonian)
     address = starting_address(hamiltonian)
     @assert address isa SingleComponentFockAddress
     N = num_modes(address) * (num_modes(address) + 1) ÷ 2
-    return JastrowAnsatz{typeof(address),Float64,N,typeof(hamiltonian)}(hamiltonian)
+    return JastrowAnsatz{typeof(address),N,typeof(hamiltonian)}(hamiltonian)
 end
 
 Rimu.build_basis(j::JastrowAnsatz) = build_basis(j.hamiltonian)
 
-function val_and_grad(j::JastrowAnsatz, addr, params)
+function val_and_grad(j::JastrowAnsatz{A,N}, addr::A, params) where {A,N}
     o = onr(addr)
     M = num_modes(addr)
 
     val = 0.0
-    grad = zeros(typeof(params))
+    grad = zeros(SVector{N,Float64})
 
     p = 0
-    for i in 1:M, j in i:M
+    @inbounds for i in 1:M, j in i:M
         p += 1
-
         onproduct = o[i] * o[j]
         local_val = params[p] * onproduct
         val += local_val
@@ -66,19 +65,19 @@ For a translationally invariant Hamiltonian, this is equivalent to [`JastrowAnsa
 but has fewer parameters.
 
 ```math
-R_i = exp(-∑_{d=0}^{M÷2} p_d ∑_{k=1}^{M} n_k n_{k + d})
+J(|f⟩; p) = exp(-∑_{d=0}^{M÷2} p_d ∑_{k=1}^{M} ⟨f|n_k n_{k + d}|f⟩)
 ```
 
 """
-struct RelativeJastrowAnsatz{A,T,N,H} <: AbstractAnsatz{A,T,N}
+struct RelativeJastrowAnsatz{A,N,H} <: AbstractAnsatz{A,Float64,N}
     hamiltonian::H
 end
 
 function RelativeJastrowAnsatz(hamiltonian)
     address = starting_address(hamiltonian)
     @assert address isa SingleComponentFockAddress
     N = cld(num_modes(address), 2)
-    return RelativeJastrowAnsatz{typeof(address),Float64,N,typeof(hamiltonian)}(hamiltonian)
+    return RelativeJastrowAnsatz{typeof(address),N,typeof(hamiltonian)}(hamiltonian)
 end
 
 Rimu.build_basis(rj::RelativeJastrowAnsatz) = build_basis(rj.hamiltonian)
@@ -88,10 +87,8 @@ function val_and_grad(rj::RelativeJastrowAnsatz, addr, params)
     M = num_parameters(rj)
 
     products = ntuple(i -> circshift_dot(o, o, i - 1), Val(M))
-    val = sum(1:M) do i
-        params[i] * circshift_dot(o, o, i - 1)
-    end
-    val = exp(-val)
+    exponent = dot(params, products)
+    val = exp(-exponent)
     grad = -SVector{M,Float64}(products .* val)
 
     return val, grad

test/Project.toml

@@ -0,0 +1,2 @@
+[deps]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

test/ansatz.jl

@@ -14,7 +14,7 @@ function check_ansatz(H, ansatz, params)
         end
 
         @testset "val_and_grad" begin
-            addrs = rand(build_basis(H), 5)
+            addrs = [starting_address(H); rand(build_basis(H), 10)]
             for addr in addrs
                 val1 = ansatz(addr, params)
                 val2, grad1 = val_and_grad(ansatz, addr, params)
@@ -47,6 +47,8 @@ end
         HubbardMom1D(BoseFS(10, 5 => 3); u=4),
         HubbardRealSpace(FermiFS2C((1,0,0,1), (1,1,0,0)); geometry=PeriodicBoundaries(2,2)),
     )
+        M = num_modes(starting_address(H))
+
         check_ansatz(H, GutzwillerAnsatz(H), [0.4])
         if H isa ExtendedHubbardReal1D
             check_ansatz(H, ExtendedGutzwillerAnsatz(H), [0.4, 0.5])
@@ -55,6 +57,9 @@ end
             check_ansatz(H, MultinomialAnsatz(H), [0.5])
             check_ansatz(H, GutzwillerAnsatz(H) + MultinomialAnsatz(H), [0.5, 0.4, 0.2])
         end
+        check_ansatz(H, JastrowAnsatz(H), rand(M * (M - 2)))
+        check_ansatz(H, RelativeJastrowAnsatz(H), rand(cld(M, 2)))
+        check_ansatz(H, DensityProfileAnsatz(H), rand(M))
     end
 end