From f4f52412eef775fe431b7c2e952b02602e5cb810 Mon Sep 17 00:00:00 2001 From: cortner Date: Thu, 20 Apr 2023 07:10:17 -0700 Subject: [PATCH 01/79] starting reorganisation towards 3d --- src/ACEpsi.jl | 2 + src/bflow.jl | 28 +-- src/bflow3d.jl | 552 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 568 insertions(+), 14 deletions(-) create mode 100644 src/bflow3d.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index 6d3c2e5..6d46eb3 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -3,4 +3,6 @@ module ACEpsi include("bflow.jl") include("envelope.jl") +include("embeddings/rnlexample.jl") + end diff --git a/src/bflow.jl b/src/bflow.jl index 2b9286a..7f4bcc8 100644 --- a/src/bflow.jl +++ b/src/bflow.jl @@ -6,7 +6,7 @@ using ACEcore.Utils: gensparse using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr import ForwardDiff -mutable struct BFwf{T, TT, TPOLY, TE} +mutable struct BFwf1{T, TT, TPOLY, TE} trans::TT polys::TPOLY pooling::PooledSparseProduct{2} @@ -30,9 +30,9 @@ mutable struct BFwf{T, TT, TPOLY, TE} ∇AA::Array{T, 3} end -(Φ::BFwf)(args...) = evaluate(Φ, args...) +(Φ::BFwf1)(args...) = evaluate(Φ, args...) -function BFwf(Nel::Integer, polys; totdeg = length(polys), +function BFwf1(Nel::Integer, polys; totdeg = length(polys), ν = 3, T = Float64, trans = identity, sd_admissible = bb -> (true), @@ -65,7 +65,7 @@ function BFwf(Nel::Integer, polys; totdeg = length(polys), Q, _ = qr(randn(T, length(corr), Nel)) W = Matrix(Q) - return BFwf(trans, polys, pooling, corr, W, envelope, spec, + return BFwf1(trans, polys, pooling, corr, W, envelope, spec, zeros(T, Nel, length(polys)), zeros(T, Nel, length(polys)), zeros(T, Nel, length(polys)), @@ -128,7 +128,7 @@ end """ This function returns a nice version of spec. """ -function displayspec(wf::BFwf) +function displayspec(wf::BFwf1) K = length(wf.polys) spec1p = [ (k, σ) for σ in [1, 2, 3] for k in 1:K] spec1p = sort(spec1p, by = b -> b[1]) @@ -141,7 +141,7 @@ function displayspec(wf::BFwf) end -function assemble_A(wf::BFwf, X::AbstractVector, Σ, Pnn=nothing) +function assemble_A(wf::BFwf1, X::AbstractVector, Σ, Pnn=nothing) nX = length(X) # position embedding @@ -161,7 +161,7 @@ function assemble_A(wf::BFwf, X::AbstractVector, Σ, Pnn=nothing) return A end -function evaluate(wf::BFwf, X::AbstractVector, Σ, Pnn=nothing) +function evaluate(wf::BFwf1, X::AbstractVector, Σ, Pnn=nothing) nX = length(X) A = assemble_A(wf, X, Σ) AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) @@ -184,7 +184,7 @@ function evaluate(wf::BFwf, X::AbstractVector, Σ, Pnn=nothing) end -function gradp_evaluate(wf::BFwf, X::AbstractVector, Σ) +function gradp_evaluate(wf::BFwf1, X::AbstractVector, Σ) nX = length(X) A = assemble_A(wf, X, Σ) @@ -227,7 +227,7 @@ Base.setindex!(A::ZeroNoEffect, args...) = nothing Base.getindex(A::ZeroNoEffect, args...) = Bool(0) -function gradient(wf::BFwf, X, Σ) +function gradient(wf::BFwf1, X, Σ) nX = length(X) # ------ forward pass ----- @@ -326,7 +326,7 @@ end # ------------------ Laplacian implementation -function laplacian(wf::BFwf, X, Σ) +function laplacian(wf::BFwf1, X, Σ) A, ∇A, ΔA = _assemble_A_∇A_ΔA(wf, X, Σ) AA, ∇AA, ΔAA = _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) @@ -453,7 +453,7 @@ end # ------------------ gradp of Laplacian -function gradp_laplacian(wf::BFwf, X, Σ) +function gradp_laplacian(wf::BFwf1, X, Σ) # ---- gradp of Laplacian of Ψ ---- @@ -527,13 +527,13 @@ end -# ----------------- BFwf parameter wraging +# ----------------- BFwf1 parameter wraging -function get_params(U::BFwf) +function get_params(U::BFwf1) return (U.W, U.envelope.ξ) end -function set_params!(U::BFwf, para) +function set_params!(U::BFwf1, para) U.W = para[1] set_params!(U.envelope, para[2]) return U diff --git a/src/bflow3d.jl b/src/bflow3d.jl new file mode 100644 index 0000000..4344c9a --- /dev/null +++ b/src/bflow3d.jl @@ -0,0 +1,552 @@ + +using ACEcore, Polynomials4ML +using Polynomials4ML: OrthPolyBasis1D3T +using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! +using ACEcore.Utils: gensparse +using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr +import ForwardDiff + +# ---------------------------------------- +mutable struct RnlYlmEmbedding + bRnl::TR + bYlm::TY + pooling::PooledSparseProduct{3} +end + + +""" +This function return correct Si for pooling operation. +""" +function onehot!(Si, i, Σ) + Si .= 0 + for k = 1:length(Σ) + Si[k, spin2num(Σ[k])] = 1 + end + # set current electron to ϕ, also remove their contribution in the sum of ↑ or ↓ basis + Si[i, 1] = 1 + Si[i, 2] = 0 + Si[i, 3] = 0 +end + +""" +This function convert spin to corresponding integer value used in spec +""" +function spin2num(σ) + if σ == '↑' + return 2 + elseif σ == '↓' + return 3 + elseif σ == '∅' + return 1 + end + error("illegal spin char for spin2num") +end + +""" +This function convert num to corresponding spin string. +""" +function num2spin(σ) + if σ == 2 + return '↑' + elseif σ == 3 + return '↓' + elseif σ == 1 + return '∅' + end + error("illegal integer value for num2spin") +end + +function assemble_A(bA::RnlYlmEmbedding, X::AbstractVector, Σ) + + nX = length(X) + # position embedding + P = wf.P + Xt = wf.trans.(X) + evaluate!(P, wf.polys, Xt) # nX x dim(polys) + + A = wf.A # zeros(nX, length(wf.pooling)) + Ai = wf.Ai # zeros(length(wf.pooling)) + Si = wf.Si # zeros(Bool, nX, 2) + + for i = 1:nX + onehot!(Si, i, Σ) + ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) + A[i, :] .= Ai + end + return A +end + + + +# ---------------------------------------- + + +mutable struct BFwf{T, TT, TPOLY, TE} + embedding::TA + corr::SparseSymmProdDAG{T} + W::Matrix{T} + envelope::TE + spec::Vector{Vector{Int64}} # corr.spec TODO: this needs to be remove + # ---------------- Temporaries + P::Matrix{T} + ∂P::Matrix{T} + dP::Matrix{T} + Φ::Matrix{T} + ∂Φ::Matrix{T} + A::Matrix{T} + ∂A::Matrix{T} + Ai::Vector{T} + ∂Ai::Vector{T} + Si::Matrix{Bool} + ∂AA::Matrix{T} + ∂Si::Matrix{T} + ∇AA::Array{T, 3} +end + +(Φ::BFwf)(args...) = evaluate(Φ, args...) + +function BFwf(Nel::Integer, polys; totdeg = length(polys), + ν = 3, T = Float64, + trans = identity, + sd_admissible = bb -> (true), + envelope = envelopefcn(x -> sqrt(1 + x^2), rand())) + # 1-particle spec + K = length(polys) + spec1p = [ (k, σ) for σ in [1, 2, 3] for k in 1:K] # (1, 2, 3) = (∅, ↑, ↓); + spec1p = sort(spec1p, by = b -> b[1]) # sorting to prevent gensparse being confused + + pooling = PooledSparseProduct(spec1p) + # generate the many-particle spec + tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] + default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) + + specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, + minvv = fill(0, ν), + maxvv = fill(length(spec1p), ν), + ordered = true) + + + spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] + + # further restrict + spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] + + corr1 = SparseSymmProd(spec; T = Float64) + corr = corr1.dag + + # initial guess for weights + Q, _ = qr(randn(T, length(corr), Nel)) + W = Matrix(Q) + + return BFwf(trans, polys, pooling, corr, W, envelope, spec, + zeros(T, Nel, length(polys)), + zeros(T, Nel, length(polys)), + zeros(T, Nel, length(polys)), + zeros(T, Nel, Nel), + zeros(T, Nel, Nel), + zeros(T, Nel, length(pooling)), + zeros(T, Nel, length(pooling)), + zeros(T, length(pooling)), + zeros(T, length(pooling)), + zeros(Bool, Nel, 3), + zeros(T, Nel, length(corr)), + zeros(T, Nel, 3), + zeros(T, Nel, Nel, length(corr)) ) + +end + + +""" +This function returns a nice version of spec. +""" +function displayspec(wf::BFwf) + K = length(wf.polys) + spec1p = [ (k, σ) for σ in [1, 2, 3] for k in 1:K] + spec1p = sort(spec1p, by = b -> b[1]) + _getnicespec = l -> (l[1], num2spin(l[2])) + nicespec = [] + for k = 1:length(wf.spec) + push!(nicespec, _getnicespec.([spec1p[wf.spec[k][j]] for j = 1:length(wf.spec[k])])) + end + return nicespec +end + + + +function evaluate(wf::BFwf, X::AbstractVector, Σ, Pnn=nothing) + nX = length(X) + A = assemble_A(wf, X, Σ) + AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) + + # the only basis to be purified are those with same spin + # scan through all corr basis, if they comes from same spin, remove self interation by using basis + # from same spin + # first we have to construct coefficent for basis coming from same spin, that is in other words the coefficent + # matrix of the original polynomial basis, this will be pass from the argument Pnn + # === purification goes here === # + + # === # + Φ = wf.Φ + mul!(Φ, parent(AA), wf.W) # nX x nX + Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin + release!(AA) + + env = wf.envelope(X) + return 2 * logabsdet(Φ)[1] + 2 * log(abs(env)) +end + + +function gradp_evaluate(wf::BFwf, X::AbstractVector, Σ) + nX = length(X) + + A = assemble_A(wf, X, Σ) + AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) + Φ = wf.Φ + mul!(Φ, parent(AA), wf.W) + Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin + + + # ψ = log | det( Φ ) | + # ∂Φ = ∂ψ/∂Φ = Φ⁻ᵀ + ∂Φ = transpose(pinv(Φ)) + + # ∂W = ∂ψ / ∂W = ∂Φ * ∂_W( AA * W ) = ∂Φ * AA + # ∂Wij = ∑_ab ∂Φab * ∂_Wij( ∑_k AA_ak W_kb ) + # = ∑_ab ∂Φab * ∑_k δ_ik δ_bj AA_ak + # = ∑_a ∂Φaj AA_ai = ∂Φaj' * AA_ai + ∇p = transpose(parent(AA)) * ∂Φ + + release!(AA) + ∇p = ∇p * 2 + + + # ------ gradient of env (w.r.t. ξ) ----- + # ∂ = ∂/∂ξ + # r = ||x|| + # ∂(2 * logabs(env)) = ∂(2 * log(exp(-ξf(r)))) = ∂(-2ξf(r)) = -f(r) + ∇logabsenv = - 2 * wf.envelope.f(norm(X)) + + return (∇p = ∇p, ∇logabsenv = [∇logabsenv]) # TODO: return a named tuple (W = gradp, D = gradient w.r.t parameter of env) +end + + + +# ----------------------- gradient + +struct ZeroNoEffect end +Base.size(::ZeroNoEffect, ::Integer) = Inf +Base.setindex!(A::ZeroNoEffect, args...) = nothing +Base.getindex(A::ZeroNoEffect, args...) = Bool(0) + + +function gradient(wf::BFwf, X, Σ) + nX = length(X) + + # ------ forward pass ----- + + # position embedding (forward-mode) + # here we evaluate and differentiate at the same time, which is cheap + P = wf.P + dP = wf.dP + Xt = wf.trans.(X) + Polynomials4ML.evaluate_ed!(P, dP, wf.polys, Xt) + ∂Xt = ForwardDiff.derivative.(Ref(x -> wf.trans(x)), X) + @inbounds for k = 1:size(dP, 2) + @simd ivdep for i = 1:nX + dP[i, k] *= ∂Xt[i] + end + end + + # no gradients here - need to somehow "constify" this vector + # could use other packages for inspiration ... + + # pooling : need an elegant way to shift this loop into a kernel! + # e.g. by passing output indices to the pooling function. + + A = wf.A # zeros(nX, length(wf.pooling)) + Ai = wf.Ai # zeros(length(wf.pooling)) + Si = wf.Si # zeros(Bool, nX, 3) + + for i = 1:nX + onehot!(Si, i, Σ) + ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) + A[i, :] .= Ai + end + + # n-correlations + AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) + + # generalized orbitals + Φ = wf.Φ + mul!(Φ, parent(AA), wf.W) + + # the resulting matrix should contains two block each comes from each spin + Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] + + # envelope + env = wf.envelope(X) + + # and finally the wave function + # ψ = logabsdet(Φ)[1] + log(abs(env)) + + # ------ backward pass ------ + #∂Φ = ∂ψ / ∂Φ = Φ⁻ᵀ + ∂Φ = transpose(pinv(Φ)) + + # ∂AA = ∂ψ/∂AA = ∂ψ/∂Φ * ∂Φ/∂AA = ∂Φ * wf.W' + ∂AA = wf.∂AA + mul!(∂AA, ∂Φ, transpose(wf.W)) + + # ∂A = ∂ψ/∂A = ∂ψ/∂AA * ∂AA/∂A -> use custom pullback + ∂A = wf.∂A # zeros(size(A)) + ACEcore.pullback_arg!(∂A, ∂AA, wf.corr, parent(AA)) + release!(AA) + + # ∂P = ∂ψ/∂P = ∂ψ/∂A * ∂A/∂P -> use custom pullback + # but need to do some work here since multiple + # pullbacks can be combined here into a single one maybe? + ∂P = wf.∂P # zeros(size(P)) + fill!(∂P, 0) + ∂Si = wf.∂Si # zeros(size(Si)) # should use ZeroNoEffect here ?!??! + Si_ = zeros(nX, 3) + for i = 1:nX + onehot!(Si_, i, Σ) + # note this line ADDS the pullback into ∂P, not overwrite the content!! + ∂Ai = @view ∂A[i, :] + ACEcore._pullback_evalpool!((∂P, ∂Si), ∂Ai, wf.pooling, (P, Si_)) + end + + # ∂X = ∂ψ/∂X = ∂ψ/∂P * ∂P/∂X + # here we can now finally employ the dP=∂P/∂X that we already know. + # ∂ψ/∂Xi = ∑_k ∂ψ/∂Pik * ∂Pik/∂Xi + # = ∑_k ∂P[i, k] * dP[i, k] + g = zeros(nX) + @inbounds for k = 1:length(wf.polys) + @simd ivdep for i = 1:nX + g[i] += ∂P[i, k] * dP[i, k] + end + end + # g = sum(∂P .* dP, dims = 2)[:] + + # envelope + ∇env = ForwardDiff.gradient(wf.envelope, X) + g += ∇env / env + g = g * 2 + return g +end + + +# ------------------ Laplacian implementation + +function laplacian(wf::BFwf, X, Σ) + + A, ∇A, ΔA = _assemble_A_∇A_ΔA(wf, X, Σ) + AA, ∇AA, ΔAA = _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) + + Δψ = _laplacian_inner(AA, ∇AA, ΔAA, wf, Σ) + + # envelope + env = wf.envelope(X) + ∇env = ForwardDiff.gradient(wf.envelope, X) + Δenv = tr(ForwardDiff.hessian(wf.envelope, X)) + + # Δ(ln(env)) = Δenv / env - ∇env ⋅ ∇env / env ^ 2 + Δψ += Δenv / env - dot(∇env, ∇env) / env^2 + Δψ = Δψ * 2 + return Δψ +end + +function _assemble_A_∇A_ΔA(wf, X, Σ) + TX = eltype(X) + lenA = length(wf.pooling) + nX = length(X) + A = zeros(TX, nX, lenA) + ∇A = zeros(TX, nX, nX, lenA) + ΔA = zeros(TX, nX, lenA) + spec_A = wf.pooling.spec + + Xt = wf.trans.(X) + P, dP, ddP = Polynomials4ML.evaluate_ed2(wf.polys, Xt) + dtrans = x -> ForwardDiff.derivative(wf.trans, x) + ddtrans = x -> ForwardDiff.derivative(dtrans, x) + ∂Xt = dtrans.(X) + ∂∂Xt = ddtrans.(X) + @inbounds for k = 1:size(dP, 2) + @simd ivdep for i = 1:nX + dP[i, k], ddP[i, k] = ∂Xt[i] * dP[i, k], ∂∂Xt[i] * dP[i, k] + ∂Xt[i]^2 * ddP[i, k] + end + end + + Si_ = zeros(nX, 3) + Ai = zeros(length(wf.pooling)) + @inbounds for i = 1:nX # loop over orbital bases (which i becomes ∅) + fill!(Si_, 0) + onehot!(Si_, i, Σ) + ACEcore.evalpool!(Ai, wf.pooling, (P, Si_)) + @. A[i, :] .= Ai + for (iA, (k, σ)) in enumerate(spec_A) + for a = 1:nX + ∇A[a, i, iA] = dP[a, k] * Si_[a, σ] + ΔA[i, iA] += ddP[a, k] * Si_[a, σ] + end + end + end + return A, ∇A, ΔA +end + +function _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) + nX = size(A, 1) + AA = zeros(nX, length(wf.corr)) + ∇AA = wf.∇AA + ΔAA = zeros(nX, length(wf.corr)) + + @inbounds for iAA = 1:wf.corr.num1 + @. AA[:, iAA] .= A[:, iAA] + @. ∇AA[:, :, iAA] .= ∇A[:, :, iAA] + @. ΔAA[:, iAA] .= ΔA[:, iAA] + end + + lenAA = length(wf.corr) + @inbounds for iAA = wf.corr.num1+1:lenAA + k1, k2 = wf.corr.nodes[iAA] + for i = 1:nX + AA_k1 = AA[i, k1]; AA_k2 = AA[i, k2] + AA[i, iAA] = AA_k1 * AA_k2 + L = ΔAA[i, k1] * AA_k2 + L = muladd(ΔAA[i, k2], AA_k1, L) + @simd ivdep for a = 1:nX + ∇AA_k1 = ∇AA[a, i, k1]; ∇AA_k2 = ∇AA[a, i, k2] + L = muladd(2 * ∇AA_k1, ∇AA_k2, L) + g = ∇AA_k1 * AA_k2 + ∇AA[a, i, iAA] = muladd(∇AA_k2, AA_k1, g) + end + ΔAA[i, iAA] = L + end + end + return AA, ∇AA, ΔAA +end + + +function _laplacian_inner(AA, ∇AA, ΔAA, wf, Σ) + + # Δψ = Φ⁻ᵀ : ΔΦ - ∑ᵢ (Φ⁻ᵀ * Φᵢ)ᵀ : (Φ⁻ᵀ * Φᵢ) + # where Φᵢ = ∂_{xi} Φ + + nX = size(AA, 1) + + # the wf, and the first layer of derivatives + Φ = wf.Φ + mul!(Φ, parent(AA), wf.W) + Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin + Φ⁻ᵀ = transpose(pinv(Φ)) + + # first contribution to the laplacian + ΔΦ = ΔAA * wf.W + ΔΦ = ΔΦ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin + + Δψ = dot(Φ⁻ᵀ, ΔΦ) + + # the gradient contribution + # TODO: we can rework this into a single BLAS3 call + # which will also give us a single back-propagation + # ∇Φi = zeros(nX, nX) + ∇Φ_all = reshape(reshape(∇AA, nX*nX, :) * wf.W, nX, nX, nX) + Φ⁻¹∇Φi = zeros(nX, nX) + for i = 1:nX + ∇Φi = ∇Φ_all[i, :, :] .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] + mul!(Φ⁻¹∇Φi, transpose(Φ⁻ᵀ), ∇Φi) + Δψ -= dot(transpose(Φ⁻¹∇Φi), Φ⁻¹∇Φi) + end + + return Δψ +end + + +# ------------------ gradp of Laplacian + + +function gradp_laplacian(wf::BFwf, X, Σ) + + + # ---- gradp of Laplacian of Ψ ---- + + nX = length(X) + + A, ∇A, ΔA = _assemble_A_∇A_ΔA(wf, X, Σ) + AA, ∇AA, ΔAA = _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) + + + # the wf, and the first layer of derivatives + Φ = wf.Φ + mul!(Φ, parent(AA), wf.W) + Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin + + Φ⁻¹ = pinv(Φ) + Φ⁻ᵀ = transpose(Φ⁻¹) + + # first contribution to the laplacian + ΔΦ = ΔAA * wf.W + ΔΦ = ΔΦ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin + + # Δψ += dot(Φ⁻ᵀ, ΔΦ) ... this leads to the next two terms + ∂ΔΦ = Φ⁻ᵀ + ∇Δψ = transpose(ΔAA) * ∂ΔΦ + + ∂Φ = - Φ⁻ᵀ * transpose(ΔΦ) * Φ⁻ᵀ + # ∇Δψ += transpose(AA) * ∂Φ + + + # the gradient contribution + # TODO: we can rework this into a single BLAS3 call + # which will also give us a single back-propagation + # ∇Φi = zeros(nX, nX) + # Φ⁻¹∇Φi = zeros(nX, nX) + ∇Φ_all = reshape(reshape(∇AA, nX*nX, :) * wf.W, nX, nX, nX) + ∂∇Φ_all = zeros(nX, nX, nX) + + for i = 1:nX + ∇Φi = ∇Φ_all[i, :, :] .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] + ∇Φiᵀ = transpose(∇Φi) + # Δψ += - dot( [Φ⁻¹∇Φi]ᵀ, Φ⁻¹∇Φi ) + + # ∂∇Φi = - 2 * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ + ∂∇Φ_all[i, :, :] = - 2 * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ + # ∇Δψ += transpose(∇AA[i, :, :]) * ∂∇Φi + + ∂Φ += 2 * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ + # ∇Δψ += transpose(AA) * ∂Φ + end + + ∇Δψ += transpose(AA) * ∂Φ + + ∇Δψ += reshape( transpose(reshape(∇AA, nX*nX, :)) * reshape(∂∇Φ_all, nX*nX, nX), + size(∇Δψ) ) + + + # ---- gradp of Laplacian of env ---- + # ∂ = ∂/∂ξ + # r = ||x|| + # ∂(2 * Δ(logabs(env))) = ∂(2 * Δ(-ξf(r)) = -2 * Δf(r) + # Δf(r) = (n-1)/r f'(r) + f''(r) + r = norm(X) + f(r) = wf.envelope.f(r) + df(r) = ForwardDiff.derivative(f, r) + ddf(r) = ForwardDiff.derivative(r -> df(r), r) + Δf = ddf(r) + (length(X)-1) * df(r)/r + + return (∇Δψ = 2 * ∇Δψ, ∇Δlogabsenv = [2 * -Δf]) +end + + + +# ----------------- BFwf parameter wraging + +function get_params(U::BFwf) + return (U.W, U.envelope.ξ) +end + +function set_params!(U::BFwf, para) + U.W = para[1] + set_params!(U.envelope, para[2]) + return U +end From 82bc707fd7f28701e6f5b66bfd4531e7f955e228 Mon Sep 17 00:00:00 2001 From: cortner Date: Thu, 20 Apr 2023 07:11:02 -0700 Subject: [PATCH 02/79] initial example or atomic orbitals --- src/atomicorbitals/atomicorbitals.jl | 6 ++ src/atomicorbitals/rnlexample.jl | 116 +++++++++++++++++++++++++++ 2 files changed, 122 insertions(+) create mode 100644 src/atomicorbitals/atomicorbitals.jl create mode 100644 src/atomicorbitals/rnlexample.jl diff --git a/src/atomicorbitals/atomicorbitals.jl b/src/atomicorbitals/atomicorbitals.jl new file mode 100644 index 0000000..fa4a811 --- /dev/null +++ b/src/atomicorbitals/atomicorbitals.jl @@ -0,0 +1,6 @@ +module AtomicOrbitals + +include("rnlexample.jl") + + +end \ No newline at end of file diff --git a/src/atomicorbitals/rnlexample.jl b/src/atomicorbitals/rnlexample.jl new file mode 100644 index 0000000..a938c7d --- /dev/null +++ b/src/atomicorbitals/rnlexample.jl @@ -0,0 +1,116 @@ +using Polynomials4ML, ForwardDiff + +const NLM{T} = NamedTuple{(:n, :l, :m), Tuple{T, T, T}} +const NL{T} = NamedTuple{(:n, :l), Tuple{T, T}} + +struct RnlExample{TP, TI} + Pn::TP + spec::Vector{NL{TI}} +end + +Base.length(basis::RnlExample) = length(basis.spec) + + +# -------- Evaluation Code + +_alloc(basis::RnlExample, r::T) where {T <: Real} = + zeros(T, length(Rnl)) + +_alloc(basis::RnlExample, rr::Vector{T}) where {T <: Real} = + zeros(T, length(rr), length(basis)) + + +evaluate(basis::RnlExample, r::Number) = evaluate(basis, [r,])[:] + +function evaluate(basis::RnlExample, R::AbstractVector) + nR = length(R) + Pn = Polynomials4ML.evaluate(basis.Pn, R) + Rnl = _alloc(basis, R) + + maxL = maximum(b.l for b in basis.spec) + rL = ones(eltype(R), length(R), maxL+1) + + # @inbounds begin + + for l = 1:maxL + # @simd ivdep + for j = 1:nR + rL[j, l+1] = R[j] * rL[j, l] + end + end + + for (i, b) in enumerate(basis.spec) + # @simd ivdep + for j = 1:nR + Rnl[j, i] = Pn[j, b.n] * rL[j, b.l+1] + end + end + + # end + + return Rnl +end + + +function evaluate_ed(basis::RnlExample, R) + nR = length(R) + Pn, dPn = Polynomials4ML.evaluate_ed(basis.Pn, R) + Rnl = _alloc(basis, R) + dRnl = _alloc(basis, R) + + maxL = maximum(b.l for b in basis.spec) + rL = ones(eltype(R), length(R), maxL+1) + drL = zeros(eltype(R), length(R), maxL+1) + for l = 1:maxL + # @simd ivdep + for j = 1:nR + rL[j, l+1] = R[j] * rL[j, l] + drL[j, l+1] = l * rL[j, l] + end + end + + for (i, b) in enumerate(basis.spec) + # @simd ivdep + for j = 1:nR + Rnl[j, i] = Pn[j, b.n] * rL[j, b.l+1] + dRnl[j, i] = dPn[j, b.n] * rL[j, b.l+1] + Pn[j, b.n] * drL[j, b.l+1] + end + end + + return Rnl, dRnl +end + + +function evaluate_ed2(basis::RnlExample, R) + nR = length(R) + Pn, dPn, ddPn = Polynomials4ML.evaluate_ed2(basis.Pn, R) + Rnl = _alloc(basis, R) + dRnl = _alloc(basis, R) + ddRnl = _alloc(basis, R) + + maxL = maximum(b.l for b in basis.spec) + rL = ones(eltype(R), length(R), maxL+1) + drL = zeros(eltype(R), length(R), maxL+1) + ddrL = zeros(eltype(R), length(R), maxL+1) + for l = 1:maxL + # @simd ivdep + for j = 1:nR + rL[j, l+1] = R[j] * rL[j, l] # r^l + drL[j, l+1] = l * rL[j, l] # l * r^(l-1) + ddrL[j, l+1] = l * drL[j, l] # (l-1) * drL[j, l] # l * (l-1) * r^(l-2) + end + end + + for (i, b) in enumerate(basis.spec) + # @simd ivdep + for j = 1:nR + Rnl[j, i] = Pn[j, b.n] * rL[j, b.l+1] + dRnl[j, i] = dPn[j, b.n] * rL[j, b.l+1] + Pn[j, b.n] * drL[j, b.l+1] + ddRnl[j, i] = ( ddPn[j, b.n] * rL[j, b.l+1] + + 2 * dPn[j, b.n] * drL[j, b.l+1] + + Pn[j, b.n] * ddrL[j, b.l+1] ) + end + end + + return Rnl, dRnl, ddRnl +end From 30e16c86997114a18b6ffe5ba117c6fd74d2d0fa Mon Sep 17 00:00:00 2001 From: cortner Date: Thu, 20 Apr 2023 09:18:52 -0700 Subject: [PATCH 03/79] fix unit tests --- src/ACEpsi.jl | 2 +- src/atomicorbitals/rnlexample.jl | 8 ++++++++ test/runtests.jl | 4 +--- test/test_atomicorbitals.jl | 26 ++++++++++++++++++++++++++ test/test_bflow.jl | 10 +++++----- 5 files changed, 41 insertions(+), 9 deletions(-) create mode 100644 test/test_atomicorbitals.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index 6d46eb3..daaaee2 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -3,6 +3,6 @@ module ACEpsi include("bflow.jl") include("envelope.jl") -include("embeddings/rnlexample.jl") +include("atomicorbitals/atomicorbitals.jl") end diff --git a/src/atomicorbitals/rnlexample.jl b/src/atomicorbitals/rnlexample.jl index a938c7d..ed5c3d7 100644 --- a/src/atomicorbitals/rnlexample.jl +++ b/src/atomicorbitals/rnlexample.jl @@ -1,4 +1,5 @@ using Polynomials4ML, ForwardDiff +import Polynomials4ML: evaluate, evaluate_ed, evaluate_ed2 const NLM{T} = NamedTuple{(:n, :l, :m), Tuple{T, T, T}} const NL{T} = NamedTuple{(:n, :l), Tuple{T, T}} @@ -8,6 +9,13 @@ struct RnlExample{TP, TI} spec::Vector{NL{TI}} end +function RnlExample(totaldegree::Integer) + bPn = legendre_basis(totaldegree+1) + maxn = length(bPn) + spec = [ (n=n, l=l) for n = 1:maxn for l = 0:(totaldegree-n+1)] + return RnlExample(bPn, spec) +end + Base.length(basis::RnlExample) = length(basis.spec) diff --git a/test/runtests.jl b/test/runtests.jl index 62daed3..6925ae7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,8 +2,6 @@ using ACEpsi using Test @testset "ACEpsi.jl" begin - - @testset "BFwf" begin include("test_bflow.jl") end - + @testset "AtomicOrbitals" begin include("test_atomicorbitals.jl") end end diff --git a/test/test_atomicorbitals.jl b/test/test_atomicorbitals.jl new file mode 100644 index 0000000..0afff95 --- /dev/null +++ b/test/test_atomicorbitals.jl @@ -0,0 +1,26 @@ + + +using Polynomials4ML, ForwardDiff, Test, ACEpsi +using Polynomials4ML.Testing: println_slim +using Polynomials4ML: evaluate, evaluate_ed, evaluate_ed2 + +# -------------- RnlExample ----------------- +@info("Testing RnlExample basis") +bRnl = ACEpsi.AtomicOrbitals.RnlExample(5) + +rr = 2 * rand(10) .- 1 +Rnl = evaluate(bRnl, rr) +Rnl1, dRnl1 = evaluate_ed(bRnl, rr) +Rnl2, dRnl2, ddRnl2 = evaluate_ed2(bRnl, rr) + +fdRnl = vcat([ ForwardDiff.derivative(r -> evaluate(bRnl, [r,]), r) + for r in rr ]...) +fddRnl = vcat([ ForwardDiff.derivative(r -> evaluate_ed(bRnl, [r,])[2], r) + for r in rr ]...) + +println_slim(@test Rnl ≈ Rnl1 ≈ Rnl2) +println_slim(@test dRnl1 ≈ dRnl2 ≈ fdRnl) +println_slim(@test ddRnl2 ≈ fddRnl) + + +# -------------- **** ----------------- diff --git a/test/test_bflow.jl b/test/test_bflow.jl index 29511b3..f4d0e39 100644 --- a/test/test_bflow.jl +++ b/test/test_bflow.jl @@ -1,6 +1,6 @@ using Polynomials4ML, ACEcore, ACEbase, Printf, ACEpsi -using ACEpsi: BFwf, gradient, evaluate, laplacian +using ACEpsi: BFwf1, gradient, laplacian using LinearAlgebra #using Random #Random.seed!(123) @@ -99,7 +99,7 @@ function fdtest(F, Σ, dF, x::AbstractVector; h0 = 1.0, verbose=true) const ↑, ↓, ∅ = '↑','↓','∅' Nel = 5 polys = legendre_basis(8) -wf = BFwf(Nel, polys; ν=3) +wf = BFwf1(Nel, polys; ν=3) X = 2 * rand(Nel) .- 1 Σ = rand([↑, ↓], Nel) @@ -211,8 +211,8 @@ grad_test3(Fp, dFp, ξ0) @info("Test getting/setting parameters") -wf1 = BFwf(Nel, polys; ν=3) -wf2 = BFwf(Nel, polys; ν=3) +wf1 = BFwf1(Nel, polys; ν=3) +wf2 = BFwf1(Nel, polys; ν=3) @printf(" wf1 - wf2: %f \n", abs(wf1(X, Σ) - wf2(X, Σ))) param1 = ACEpsi.get_params(wf1) wf2 = ACEpsi.set_params!(wf2, param1) @@ -220,6 +220,6 @@ wf2 = ACEpsi.set_params!(wf2, param1) ## -@warn("removed compac test since json file is missing") +@warn("removed compat test since json file is missing") # @info("Test compatibility with ACESchrodinger") # Jerry: Not sure if this should be kept in the same file # include("compare_bflow.jl") From a2bf38de94a83615a9ed03b8a2dfb7c368d4c4ee Mon Sep 17 00:00:00 2001 From: cortner Date: Sun, 30 Apr 2023 22:56:18 -0700 Subject: [PATCH 04/79] several steps towards 3d A basis --- Project.toml | 8 +- src/ACEpsi.jl | 6 + src/atomicorbitals/atomicorbitals.jl | 3 + src/atomicorbitals/atorbbasis.jl | 127 +++++++++++++++ src/atomicorbitals/rnlexample.jl | 11 +- src/bflow3d.jl | 11 +- test/runtests.jl | 2 +- test/test_atomicorbitals.jl | 1 + test/test_atorbbasis.jl | 17 ++ test/test_bflow1.jl | 225 +++++++++++++++++++++++++++ 10 files changed, 397 insertions(+), 14 deletions(-) create mode 100644 src/atomicorbitals/atorbbasis.jl create mode 100644 test/test_atorbbasis.jl create mode 100644 test/test_bflow1.jl diff --git a/Project.toml b/Project.toml index a8e4139..cd6b365 100644 --- a/Project.toml +++ b/Project.toml @@ -13,17 +13,17 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d" Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [compat] +ACEbase = "0.2" ACEcore = "0.0.1" BenchmarkTools = "1" ForwardDiff = "0.10" +JSON = "0.21" ObjectPools = "0.0.2" -Polynomials4ML = "0.0.2" +Polynomials4ML = "0.0.6" julia = "1" -JSON = "0.21" -ACEbase = "0.2" - [extras] Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index daaaee2..5040aaa 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -1,5 +1,11 @@ module ACEpsi +const Spin = Char +const ↑ = '↑' +const ↓ = '↓' +spins = [↑, ↓] +export ↑, ↓ + include("bflow.jl") include("envelope.jl") diff --git a/src/atomicorbitals/atomicorbitals.jl b/src/atomicorbitals/atomicorbitals.jl index fa4a811..b089ed7 100644 --- a/src/atomicorbitals/atomicorbitals.jl +++ b/src/atomicorbitals/atomicorbitals.jl @@ -1,5 +1,8 @@ module AtomicOrbitals +import Polynomials4ML: degree + +include("atorbbasis.jl") include("rnlexample.jl") diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl new file mode 100644 index 0000000..9a0d7a7 --- /dev/null +++ b/src/atomicorbitals/atorbbasis.jl @@ -0,0 +1,127 @@ + +using ACEpsi: ↑, ↓, spins, Spin +using StaticArrays + +struct Nuc{T} + rr::SVector{3, T} + charge::T +end + + +# +# Ordering of the embedding +# nuc | 1 2 3 1 2 3 1 2 3 +# k | 1 1 1 2 2 2 2 2 2 +# + +const NTRNL1 = NamedTuple{(:n, :l, :m, :s), Tuple{Int, Int, Int, Spin}} +const NTRNLI = NamedTuple{(:I, :n, :l, :m, :s), Tuple{Int, Int, Int, Int, Spin}} + +mutable struct AtomicOrbitalsBasis{TR, TY, T} + bRnl::TR + bYlm::TY + spec1::Vector{NTRNL1} + spec::Vector{NTRNLI} + nuclei::Vector{Nuc{T}} + # pooling::PooledSparseProduct{3} +end + +function AtomicOrbitalsBasis(bRnl, bYlm, spec1::Vector{NTRNL1}, + nuclei::Vector{<: Nuc}) + spec = NTRNLI[] + basis = AtomicOrbitalsBasis(bRnl, bYlm, spec1, spec, eltype(nuclei)[]) + set_nuclei!(basis, nuclei) + return basis +end + + +function AtomicOrbitalsBasis(bRnl, bYlm; + totaldegree=3, + nuclei = Nuc{Float64}[], + ) + spec1 = make_nlms_spec(bRnl, bYlm; + totaldegree = totaldegree) + return AtomicOrbitalsBasis(bRnl, bYlm, spec1, nuclei) +end + + +""" +This constructs the specification of all the atomic orbitals for one +nucleus. + +* bRnl : radial basis +* Ylm : angular basis, assumed to be spherical harmonics +* admissible : a filter, default is a total degree +""" +function make_nlms_spec(bRnl, bYlm; + totaldegree::Integer = -1, + admissible = ( (br, by) -> degree(bRnl, br) + + degree(bYlm, by) <= totaldegree), + nnuc = 0) + + spec_Rnl = natural_indices(bRnl) + spec_Ylm = natural_indices(bYlm) + + spec1 = NTRNL1[] + for (iR, br) in enumerate(spec_Rnl), (iY, by) in enumerate(spec_Ylm) + if br.l != by.l + continue + end + if admissible(br, by) + for s in spins + push!(spec1, (n=br.n, l = br.l, m = by.m, s = s)) + end + end + end + return spec1 +end + + +function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) + basis.nuclei = copy(collect(nuclei)) + Nnuc = length(basis.nuclei) + + spec = NTRNLI[] + for b in basis.spec1, I = 1:Nnuc + push!(spec, (I = I, b...)) + end + + basis.spec = spec + + # bA = _build_pooling(spec1, Nnuc) + # basis.bA = bA + return nothing +end + + +# ------------ Evaluation kernels + +""" +This function return correct Si for pooling operation. +""" +function onehot_spin!(Si, i, Σ) + Si .= 0 + for k = 1:length(Σ) + Si[k, spin2num(Σ[k])] = 1 + end + # set current electron to ϕ, also remove their contribution in the sum of ↑ or ↓ basis + Si[i, 1] = 1 + Si[i, 2] = 0 + Si[i, 3] = 0 +end + + +# function evaluate(basis::AtomicOrbitalsBasis, +# X::AbstractVector{<: AbstractVector}, +# Σ) +# nuc = basis.nuclei +# Nnuc = length(nuc) +# Nel = length(X) +# VT = promote_type(eltype(nuc), eltype(X)) + +# # create all the shifted configurations +# XX = zeros(VT, Nel) +# Rnl = + +# end + diff --git a/src/atomicorbitals/rnlexample.jl b/src/atomicorbitals/rnlexample.jl index ed5c3d7..02979ef 100644 --- a/src/atomicorbitals/rnlexample.jl +++ b/src/atomicorbitals/rnlexample.jl @@ -1,5 +1,6 @@ using Polynomials4ML, ForwardDiff -import Polynomials4ML: evaluate, evaluate_ed, evaluate_ed2 +import Polynomials4ML: evaluate, evaluate_ed, evaluate_ed2, + natural_indices const NLM{T} = NamedTuple{(:n, :l, :m), Tuple{T, T, T}} const NL{T} = NamedTuple{(:n, :l), Tuple{T, T}} @@ -18,6 +19,10 @@ end Base.length(basis::RnlExample) = length(basis.spec) +natural_indices(basis::RnlExample) = copy(basis.spec) + +degree(basis::RnlExample, i::Integer) = degree(basis, spec[i]) +degree(basis::RnlExample, b::NamedTuple) = b.n + b.l # -------- Evaluation Code @@ -43,14 +48,14 @@ function evaluate(basis::RnlExample, R::AbstractVector) for l = 1:maxL # @simd ivdep for j = 1:nR - rL[j, l+1] = R[j] * rL[j, l] + rL[j, l+1] = R[j] * rL[j, l] # r^l end end for (i, b) in enumerate(basis.spec) # @simd ivdep for j = 1:nR - Rnl[j, i] = Pn[j, b.n] * rL[j, b.l+1] + Rnl[j, i] = Pn[j, b.n] * rL[j, b.l+1] # r^l * P_n -> degree l+n end end diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 4344c9a..c911046 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -7,15 +7,13 @@ using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr import ForwardDiff # ---------------------------------------- -mutable struct RnlYlmEmbedding - bRnl::TR - bYlm::TY - pooling::PooledSparseProduct{3} -end + + + """ -This function return correct Si for pooling operation. +This function returns correct Si for pooling operation. """ function onehot!(Si, i, Σ) Si .= 0 @@ -28,6 +26,7 @@ function onehot!(Si, i, Σ) Si[i, 3] = 0 end + """ This function convert spin to corresponding integer value used in spec """ diff --git a/test/runtests.jl b/test/runtests.jl index 6925ae7..89d47d8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,6 +2,6 @@ using ACEpsi using Test @testset "ACEpsi.jl" begin - @testset "BFwf" begin include("test_bflow.jl") end + @testset "BFwf" begin include("test_bflow1.jl") end @testset "AtomicOrbitals" begin include("test_atomicorbitals.jl") end end diff --git a/test/test_atomicorbitals.jl b/test/test_atomicorbitals.jl index 0afff95..34a6363 100644 --- a/test/test_atomicorbitals.jl +++ b/test/test_atomicorbitals.jl @@ -24,3 +24,4 @@ println_slim(@test ddRnl2 ≈ fddRnl) # -------------- **** ----------------- + diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl new file mode 100644 index 0000000..ba4b0b8 --- /dev/null +++ b/test/test_atorbbasis.jl @@ -0,0 +1,17 @@ + +using ACEpsi, Polynomials4ML, StaticArrays +using Polynomials4ML: natural_indices, degree +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc + + +## + +totdeg = 3 +bRnl = ACEpsi.AtomicOrbitals.RnlExample(totdeg) +bYlm = RYlmBasis(totdeg) +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] + +bAnlm = AtomicOrbitalsBasis(bRnl, bYlm; + totaldegree = totdeg, + nuclei = nuclei ) + diff --git a/test/test_bflow1.jl b/test/test_bflow1.jl new file mode 100644 index 0000000..f4d0e39 --- /dev/null +++ b/test/test_bflow1.jl @@ -0,0 +1,225 @@ + +using Polynomials4ML, ACEcore, ACEbase, Printf, ACEpsi +using ACEpsi: BFwf1, gradient, laplacian +using LinearAlgebra +#using Random +#Random.seed!(123) +## +function lap_test(f, Δf, X) + F = f(X) + ΔF = Δf(X) + nX = length(X) + n1, nF = size(F) + U = randn(n1) + V = randn(nF) ./ (1:nF).^2 + f0 = U' * F * V + Δf0 = U' * ΔF * V + EE = Matrix(I, (nX, nX)) + for h in sqrt(0.1).^((5:12)) + Δfh = 0.0 + for i = 1:nX + Δfh += (U'*f(X + h * EE[:, i])*V - f0) / h^2 + Δfh += (U'*f(X - h * EE[:, i])*V - f0) / h^2 + end + @printf(" %.1e | %.2e \n", h, abs(Δfh - Δf0)) + end +end + +function grad_test2(f, df, X::AbstractVector) + F = f(X) + ∇F = df(X) + nX = length(X) + EE = Matrix(I, (nX, nX)) + + for h in 0.1.^(3:12) + gh = [ (f(X + h * EE[:, i]) - F) / h for i = 1:nX ] + @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) + end +end + +function grad_test3(f, df, X::Float64) + F = f(X) + ∇F = df(X) # return as vector of size 1 + for h in 0.1.^(3:10) + gh = [ (f(X + h) - F) / h] + @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) + end +end + + +function grad_test(f, df, X) + F = f(X) + ∇F = df(X) + nX, nF = size(F) + U = randn(nX) + V = randn(nF) ./ (1:nF).^2 + f0 = U' * F * V + ∇f0 = [ U' * ∇F[i, :, :] * V for i = 1:nX ] + EE = Matrix(I, (Nel, Nel)) + for h in 0.1.^(2:10) + gh = [ (U'*f(X + h * EE[:, i])*V - f0) / h for i = 1:Nel ] + @printf(" %.1e | %.2e \n", h, norm(gh - ∇f0, Inf)) + end +end + +"This function should be removed later to test in a nicer way..." +function fdtest(F, Σ, dF, x::AbstractVector; h0 = 1.0, verbose=true) + errors = Float64[] + E = F(x, Σ) + dE = dF + # loop through finite-difference step-lengths + verbose && @printf("---------|----------- \n") + verbose && @printf(" h | error \n") + verbose && @printf("---------|----------- \n") + for p = 2:11 + h = 0.1^p + dEh = copy(dE) + for n = 1:length(dE) + x[n] += h + dEh[n] = (F(x, Σ) - E) / h + x[n] -= h + end + push!(errors, norm(dE - dEh, Inf)) + verbose && @printf(" %1.1e | %4.2e \n", h, errors[end]) + end + verbose && @printf("---------|----------- \n") + if minimum(errors) <= 1e-3 * maximum(errors) + verbose && println("passed") + return true + else + @warn("""It seems the finite-difference test has failed, which indicates + that there is an inconsistency between the function and gradient + evaluation. Please double-check this manually / visually. (It is + also possible that the function being tested is poorly scaled.)""") + return false + end + end +## + +const ↑, ↓, ∅ = '↑','↓','∅' +Nel = 5 +polys = legendre_basis(8) +wf = BFwf1(Nel, polys; ν=3) + +X = 2 * rand(Nel) .- 1 +Σ = rand([↑, ↓], Nel) + +wf(X, Σ) +g = gradient(wf, X, Σ) + +## + +using LinearAlgebra +using Printf +#using ACEbase.Testing: fdtest + +@info("Fd test of gradient w.r.t. X") +fdtest(wf, Σ, g, X) + + +# ## + +# todo: move this to a performance benchmark script +# using BenchmarkTools +# @btime $wf($X) +# @btime gradient($wf, $X) + +## + +A, ∇A, ΔA = ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ) + +@info("test ∇A") +grad_test(X -> ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ)[1], + X -> ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ)[2], + X) + +@info("test ΔA") +lap_test(X -> ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ)[1], + X -> ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ)[3], + X) + + +## + +function f_AA(X, Σ) + A, ∇A, ΔA = ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ) + AA, ∇AA, ΔAA = ACEpsi._assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) + return AA, ∇AA, ΔAA +end + +@info("test ∇A") +grad_test(X -> f_AA(X, Σ)[1], X -> f_AA(X, Σ)[2], X) + +@info("test ΔA") +lap_test(X -> f_AA(X, Σ)[1], X -> f_AA(X, Σ)[3], X) + + +## + +@info("test Δψ") +lap_test(X -> [wf(X, Σ);;], X -> [ACEpsi.laplacian(wf, X, Σ);;], X) + + +## + +@info("Test ∇ψ w.r.t. parameters") + +ACEpsi.gradp_evaluate(wf, X, Σ) + + +W0 = copy(wf.W) +w0 = W0[:] +Fp = w -> ( wf.W[:] .= w[:]; wf(X, Σ)) +dFp = w -> ( wf.W[:] .= w[:]; ACEpsi.gradp_evaluate(wf, X, Σ)[1][:] ) + +grad_test2(Fp, dFp, w0) + +## + +@info("test ∇env w.r.t. parameter") +Ξ0 = copy(wf.envelope.ξ) +ξ0 = Ξ0 + +Envp = w -> (wf.envelope.ξ = w; wf(X, Σ)) +dEnvp = w -> (wf.envelope.ξ = w; ACEpsi.gradp_evaluate(wf, X, Σ)[2]) + +grad_test3(Envp, dEnvp, ξ0) + +## + +@info("Test ∇Δψ w.r.t. parameters") + +Fp = w -> ( wf.W[:] .= w[:]; ACEpsi.laplacian(wf, X, Σ)) +dFp = w -> ( wf.W[:] .= w[:]; ACEpsi.gradp_laplacian(wf, X, Σ)[1][:] ) + +grad_test2(Fp, dFp, w0) + +## + +@info("Test ∇Δenv w.r.t. parameters") + +Ξ0 = copy(wf.envelope.ξ) +ξ0 = Ξ0 + + +Fp = w -> ( wf.envelope.ξ = w; ACEpsi.laplacian(wf, X, Σ)) +dFp = w -> (wf.envelope.ξ = w; ACEpsi.gradp_laplacian(wf, X, Σ)[2][:] ) + +grad_test3(Fp, dFp, ξ0) + +## + +@info("Test getting/setting parameters") + +wf1 = BFwf1(Nel, polys; ν=3) +wf2 = BFwf1(Nel, polys; ν=3) +@printf(" wf1 - wf2: %f \n", abs(wf1(X, Σ) - wf2(X, Σ))) +param1 = ACEpsi.get_params(wf1) +wf2 = ACEpsi.set_params!(wf2, param1) +@printf(" wf1 - wf2 with parameter wf1: %f \n", abs(wf1(X, Σ) - wf2(X, Σ))) + +## + +@warn("removed compat test since json file is missing") +# @info("Test compatibility with ACESchrodinger") # Jerry: Not sure if this should be kept in the same file +# include("compare_bflow.jl") From 86cc764a5d23c4552ba53ec60e4a47d7cd904cb3 Mon Sep 17 00:00:00 2001 From: cortner Date: Mon, 1 May 2023 14:10:35 -0700 Subject: [PATCH 05/79] lots of cleanup and drafting --- src/ACEpsi.jl | 10 +-- src/atomicorbitals/atorbbasis.jl | 140 ++++++++++++++++++++++++------- src/bflow.jl | 27 ------ src/spins.jl | 50 +++++++++++ test/test_atorbbasis.jl | 9 +- 5 files changed, 175 insertions(+), 61 deletions(-) create mode 100644 src/spins.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index 5040aaa..52b496f 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -1,14 +1,14 @@ module ACEpsi -const Spin = Char -const ↑ = '↑' -const ↓ = '↓' -spins = [↑, ↓] -export ↑, ↓ +# define spin symbols and some basic functionality +include("spins.jl") +# the old 1d backflow code, keep around for now... include("bflow.jl") include("envelope.jl") +# the new 3d backflow code include("atomicorbitals/atomicorbitals.jl") +# include("bflow3d.jl") end diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 9a0d7a7..dcaf077 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -1,6 +1,8 @@ -using ACEpsi: ↑, ↓, spins, Spin +using ACEpsi: ↑, ↓, spins, extspins, Spin, spin2idx, idx2spin +import Polynomials4ML using StaticArrays +using LinearAlgebra: norm struct Nuc{T} rr::SVector{3, T} @@ -14,23 +16,45 @@ end # k | 1 1 1 2 2 2 2 2 2 # -const NTRNL1 = NamedTuple{(:n, :l, :m, :s), Tuple{Int, Int, Int, Spin}} -const NTRNLI = NamedTuple{(:I, :n, :l, :m, :s), Tuple{Int, Int, Int, Int, Spin}} +const NTRNL1 = NamedTuple{(:n, :l, :m), Tuple{Int, Int, Int}} +const NTRNLIS = NamedTuple{(:I, :s, :n, :l, :m), Tuple{Int, Spin, Int, Int, Int}} mutable struct AtomicOrbitalsBasis{TR, TY, T} bRnl::TR bYlm::TY - spec1::Vector{NTRNL1} - spec::Vector{NTRNLI} - nuclei::Vector{Nuc{T}} - # pooling::PooledSparseProduct{3} + # ------- specification of the atomic orbitals Rnl * Ylm + spec1::Vector{NTRNL1} # human readable spec + spec1idx::Vector{Tuple{Int, Int}} # indices into Rnl, Ylm + nuclei::Vector{Nuc{T}} # nuclei (defines the shifted orbitals) + # ------- specification of the pooling operations A_{sInlm} + # spec::Vector{NTRNLIS} # human readable spec, definition is + # implicit in iteration protocol. +end + + +function _invmap(a::AbstractVector) + inva = Dict{eltype(a), Int}() + for i = 1:length(a) + inva[a[i]] = i + end + return inva end function AtomicOrbitalsBasis(bRnl, bYlm, spec1::Vector{NTRNL1}, nuclei::Vector{<: Nuc}) - spec = NTRNLI[] - basis = AtomicOrbitalsBasis(bRnl, bYlm, spec1, spec, eltype(nuclei)[]) - set_nuclei!(basis, nuclei) + + spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) + spec_Rnl = bRnl.spec; inv_Rnl = _invmap(spec_Rnl) + spec_Ylm = Polynomials4ML.natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) + + spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) + for (i, b) in enumerate(spec1) + spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) + end + + # spec = NTRNLIS[] + basis = AtomicOrbitalsBasis(bRnl, bYlm, spec1, spec1idx, eltype(nuclei)[]) + # set_nuclei!(basis, nuclei) return basis end @@ -68,9 +92,7 @@ function make_nlms_spec(bRnl, bYlm; continue end if admissible(br, by) - for s in spins - push!(spec1, (n=br.n, l = br.l, m = by.m, s = s)) - end + push!(spec1, (n=br.n, l = br.l, m = by.m)) end end return spec1 @@ -81,9 +103,9 @@ function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) basis.nuclei = copy(collect(nuclei)) Nnuc = length(basis.nuclei) - spec = NTRNLI[] - for b in basis.spec1, I = 1:Nnuc - push!(spec, (I = I, b...)) + spec = NTRNLIS[] + for b in basis.spec1, I = 1:Nnuc, s in extspins() + push!(spec, (I = I, s=s, b...)) end basis.spec = spec @@ -111,17 +133,79 @@ function onehot_spin!(Si, i, Σ) end -# function evaluate(basis::AtomicOrbitalsBasis, -# X::AbstractVector{<: AbstractVector}, -# Σ) -# nuc = basis.nuclei -# Nnuc = length(nuc) -# Nel = length(X) -# VT = promote_type(eltype(nuc), eltype(X)) - -# # create all the shifted configurations -# XX = zeros(VT, Nel) -# Rnl = +function proto_evaluate(basis::AtomicOrbitalsBasis, + X::AbstractVector{<: AbstractVector}, + Σ) + nuc = basis.nuclei + Nnuc = length(nuc) + Nel = length(X) + Nnlm = length(basis.spec1) + VT = promote_type(eltype(nuc[1].rr), eltype(X[1])) + @show VT -# end + # create all the shifted configurations + # this MUST be done in the format (I, i) + XX = zeros(VT, (Nnuc, Nel)) + xx = zeros(eltype(VT), (Nnuc, Nel)) + for I = 1:Nnuc, i = 1:Nel + XX[I, i] = X[i] - nuc[I].rr + xx[I, i] = norm(XX[I, i]) + end + + # evaluate the radial and angular components on all the shifted particles + Rnl = reshape(evaluate(basis.bRnl, xx[:]), (Nnuc, Nel)) + Ylm = reshape(evaluate(basis.bYlm, XX[:]), (Nnuc, Nel)) + + # evaluate all the atomic orbitals as ϕ_nlm = Rnl * Ylm + TA = promote_type(eltype(Rnl), eltype(Ylm)) + @show TA + ϕnlm = zeros(TA, (Nnuc, Nel, Nnlm)) + for (k, (iR, iY)) in enumerate(basis.spec1idx) + for i = 1:Nel, I = 1:Nnuc + ϕnlm[I, i, k] = Rnl[I, i, iR] * Ylm[I, i, iY] + end + end + + # evaluate the pooling operation + # spin I k = (nlm) + Aall = zeros(TA, (2, Nnuc, Nnlm)) + for k = 1:Nnlm + for i = 1:Nel + iσ = spin2num(Σ[i]) + for I = 1:Nnuc + Aall[iσ, I, k] += ϕnlm[I, i, k] + end + end + end + + # now correct the pooling Aall and write into A^(i) + # with do it with i leading so that the N-correlations can + # be parallelized over i + # + # A[i, :] = A-basis for electron i, with channels, s, I, k=nlm + # A[i, ∅, I, k] = ϕnlm[I, i, k] + # for σ = ↑ or ↓ we have + # A[i, σ, I, k] = ∑_{j ≂̸ i : Σ[j] == σ} ϕnlm[I, j, k] + # = ∑_{j : Σ[j] == σ} ϕnlm[I, j, k] - (Σ[i] == σ) * ϕnlm[I, i, k] + # + @assert spin2idx(↑) == 1 + @assert spin2idx(↓) == 2 + @assert spin2idx(∅) == 3 + A = zeros(TA, ((Nel, 3, Nnuc, Nnlm))) + for k = 1:Nnlm + for I = 1:Nnuc + for i = 1:Nel + A[i, 3, I, k] = ϕnlm[I, i, k] + end + for iσ = 1:2 + σ = idx2spin(iσ) + for i = 1:Nel + A[i, iσ, I, k] = Aall[iσ, I, k] - (Σ[i] == σ) * ϕnlm[I, i, k] + end + end + end + end + + return A +end diff --git a/src/bflow.jl b/src/bflow.jl index 7f4bcc8..bfea545 100644 --- a/src/bflow.jl +++ b/src/bflow.jl @@ -96,33 +96,6 @@ function onehot!(Si, i, Σ) Si[i, 3] = 0 end -""" -This function convert spin to corresponding integer value used in spec -""" -function spin2num(σ) - if σ == '↑' - return 2 - elseif σ == '↓' - return 3 - elseif σ == '∅' - return 1 - end - error("illegal spin char for spin2num") -end - -""" -This function convert num to corresponding spin string. -""" -function num2spin(σ) - if σ == 2 - return '↑' - elseif σ == 3 - return '↓' - elseif σ == 1 - return '∅' - end - error("illegal integer value for num2spin") -end """ diff --git a/src/spins.jl b/src/spins.jl new file mode 100644 index 0000000..6e1fad9 --- /dev/null +++ b/src/spins.jl @@ -0,0 +1,50 @@ + +using StaticArrays: SA + +export ↑, ↓, spins + + +# Define the spin types and variables +const Spin = Char +const ↑ = '↑' +const ↓ = '↓' +const ∅ = '∅' # this is only for internal use +_spins = SA[↑, ↓] +_extspins = SA[↑, ↓, ∅] + +spins() = _spins +extspins() = _extspins + + +""" +This function convert spin to corresponding integer value used in spec +""" +function spin2idx(σ) + if σ == ↑ + return 1 + elseif σ == ↓ + return 2 + elseif σ == ∅ + return 3 + end + error("illegal spin char for spin2idx") +end + +""" +This function convert idx to corresponding spin string. +""" +function idx2spin(i) + if i == 1 + return ↑ + elseif i == 2 + return ↓ + elseif i == 3 + return ∅ + end + error("illegal integer value for idx2spin") +end + + +# TODO : deprecate these +const spin2num = spin2idx +const num2spin = idx2spin diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index ba4b0b8..9830970 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -1,7 +1,7 @@ using ACEpsi, Polynomials4ML, StaticArrays using Polynomials4ML: natural_indices, degree -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, proto_evaluate ## @@ -15,3 +15,10 @@ bAnlm = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei ) +## + +Nel = 5 +X = randn(SVector{3, Float64}, Nel) +Σ = rand(spins(), Nel) + +proto_evaluate(bAnlm, X, Σ) \ No newline at end of file From 7ea915f2c344f1c63c48531a16652313bbbe0607 Mon Sep 17 00:00:00 2001 From: cortner Date: Mon, 1 May 2023 14:19:38 -0700 Subject: [PATCH 06/79] working but untested evaluation of A --- src/atomicorbitals/atorbbasis.jl | 52 ++++++++++++++------------------ test/test_atorbbasis.jl | 2 ++ 2 files changed, 25 insertions(+), 29 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index dcaf077..cf77b41 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -1,12 +1,12 @@ -using ACEpsi: ↑, ↓, spins, extspins, Spin, spin2idx, idx2spin +using ACEpsi: ↑, ↓, ∅, spins, extspins, Spin, spin2idx, idx2spin import Polynomials4ML using StaticArrays using LinearAlgebra: norm struct Nuc{T} rr::SVector{3, T} - charge::T + charge::T # should this be an integer? end @@ -54,7 +54,7 @@ function AtomicOrbitalsBasis(bRnl, bYlm, spec1::Vector{NTRNL1}, # spec = NTRNLIS[] basis = AtomicOrbitalsBasis(bRnl, bYlm, spec1, spec1idx, eltype(nuclei)[]) - # set_nuclei!(basis, nuclei) + set_nuclei!(basis, nuclei) return basis end @@ -101,37 +101,30 @@ end function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) basis.nuclei = copy(collect(nuclei)) + return nothing +end + + +function get_spec(basis::AtomicOrbitalsBasis) + spec = NTRNLIS[] Nnuc = length(basis.nuclei) - spec = NTRNLIS[] - for b in basis.spec1, I = 1:Nnuc, s in extspins() - push!(spec, (I = I, s=s, b...)) - end + spec = Array{NTRNLIS, 3}(undef, (3, Nnuc, length(basis.spec1))) - basis.spec = spec + for (k, nlm) in enumerate(basis.spec1) + for I = 1:Nnuc + for (is, s) in enumerate(extspins()) + spec[is, I, k] = (I = I, s=s, nlm...) + end + end + end - # bA = _build_pooling(spec1, Nnuc) - # basis.bA = bA - return nothing + return spec end # ------------ Evaluation kernels -""" -This function return correct Si for pooling operation. -""" -function onehot_spin!(Si, i, Σ) - Si .= 0 - for k = 1:length(Σ) - Si[k, spin2num(Σ[k])] = 1 - end - # set current electron to ϕ, also remove their contribution in the sum of ↑ or ↓ basis - Si[i, 1] = 1 - Si[i, 2] = 0 - Si[i, 3] = 0 -end - function proto_evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, @@ -140,7 +133,8 @@ function proto_evaluate(basis::AtomicOrbitalsBasis, Nnuc = length(nuc) Nel = length(X) Nnlm = length(basis.spec1) - VT = promote_type(eltype(nuc[1].rr), eltype(X[1])) + T = promote_type(eltype(nuc[1].rr), eltype(X[1])) + VT = SVector{3, T} @show VT # create all the shifted configurations @@ -153,8 +147,8 @@ function proto_evaluate(basis::AtomicOrbitalsBasis, end # evaluate the radial and angular components on all the shifted particles - Rnl = reshape(evaluate(basis.bRnl, xx[:]), (Nnuc, Nel)) - Ylm = reshape(evaluate(basis.bYlm, XX[:]), (Nnuc, Nel)) + Rnl = reshape(evaluate(basis.bRnl, xx[:]), (Nnuc, Nel, length(basis.bRnl))) + Ylm = reshape(evaluate(basis.bYlm, XX[:]), (Nnuc, Nel, length(basis.bYlm))) # evaluate all the atomic orbitals as ϕ_nlm = Rnl * Ylm TA = promote_type(eltype(Rnl), eltype(Ylm)) @@ -171,7 +165,7 @@ function proto_evaluate(basis::AtomicOrbitalsBasis, Aall = zeros(TA, (2, Nnuc, Nnlm)) for k = 1:Nnlm for i = 1:Nel - iσ = spin2num(Σ[i]) + iσ = spin2idx(Σ[i]) for I = 1:Nnuc Aall[iσ, I, k] += ϕnlm[I, i, k] end diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 9830970..6516c0d 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -15,6 +15,8 @@ bAnlm = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei ) +spec = ACEpsi.AtomicOrbitals.get_spec(bAnlm) + ## Nel = 5 From ab85417f1ed5b686f5aa301ceabb87b175cc79d8 Mon Sep 17 00:00:00 2001 From: cortner Date: Mon, 1 May 2023 22:10:19 -0700 Subject: [PATCH 07/79] first unit test for proto_evaluate passing --- Project.toml | 1 + src/atomicorbitals/atorbbasis.jl | 4 +++ test/test_atorbbasis.jl | 45 +++++++++++++++++++++++++++++--- 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index cd6b365..cb754d0 100644 --- a/Project.toml +++ b/Project.toml @@ -14,6 +14,7 @@ ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d" Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b" [compat] ACEbase = "0.2" diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index cf77b41..464e541 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -182,6 +182,10 @@ function proto_evaluate(basis::AtomicOrbitalsBasis, # A[i, σ, I, k] = ∑_{j ≂̸ i : Σ[j] == σ} ϕnlm[I, j, k] # = ∑_{j : Σ[j] == σ} ϕnlm[I, j, k] - (Σ[i] == σ) * ϕnlm[I, i, k] # + # + # TODO: discuss - this could be stored much more efficiently as a + # lazy array. Could that have advantages? + # @assert spin2idx(↑) == 1 @assert spin2idx(↓) == 2 @assert spin2idx(∅) == 3 diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 6516c0d..d208801 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -1,8 +1,8 @@ -using ACEpsi, Polynomials4ML, StaticArrays +using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, proto_evaluate - +using Polynomials4ML.Testing: print_tf ## @@ -23,4 +23,43 @@ Nel = 5 X = randn(SVector{3, Float64}, Nel) Σ = rand(spins(), Nel) -proto_evaluate(bAnlm, X, Σ) \ No newline at end of file +A = proto_evaluate(bAnlm, X, Σ) + + +## test the Evaluation: +using LinearAlgebra: norm +bYlm_ = RYlmBasis(totdeg) +Nnlm = length(bAnlm.spec1) +Nnuc = length(bAnlm.nuclei) +inv_nl = ACEpsi.AtomicOrbitals._invmap(bRnl.spec) +inv_lm = ACEpsi.AtomicOrbitals._invmap(natural_indices(bYlm)) + +for I = 1:Nnuc + XI = X .- Ref(bAnlm.nuclei[I].rr) + xI = norm.(XI) + Rnl = evaluate(bRnl, xI) + Ylm = evaluate(bYlm_, XI) + for k = 1:Nnlm + nlm = bAnlm.spec1[k] + iR = inv_nl[(n = nlm.n, l = nlm.l)] + iY = inv_lm[(l = nlm.l, m = nlm.m)] + + for i = 1:Nel + for (is, s) in enumerate(ACEpsi.extspins()) + a1 = A[i, is, I, k] + + if s in [↑, ↓] + a2 = sum( Rnl[j, iR] * Ylm[j, iY] * (Σ[j] == s) * (1 - (j == i)) for j = 1:Nel ) + else # s = ∅ + a2 = Rnl[i, iR] * Ylm[i, iY] + end + # println("(i=$i, σ=$s, I=$I, n=$(nlm.n), l=$(nlm.l), m=$(nlm.m)) -> ", abs(a1 - a2)) + print_tf(@test abs(a1 - a2) < 1e-12) + end + end + end +end +println() + +## + From dddcb2c712384ec255a15beb78729e119e2c412e Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 8 May 2023 09:11:24 +0800 Subject: [PATCH 08/79] test atorbbasis --- src/atomicorbitals/atorbbasis.jl | 184 ++++++++++++++++--------------- test/test_atorbbasis.jl | 47 ++++++-- 2 files changed, 131 insertions(+), 100 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 464e541..8083c1b 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -9,7 +9,6 @@ struct Nuc{T} charge::T # should this be an integer? end - # # Ordering of the embedding # nuc | 1 2 3 1 2 3 1 2 3 @@ -19,56 +18,12 @@ end const NTRNL1 = NamedTuple{(:n, :l, :m), Tuple{Int, Int, Int}} const NTRNLIS = NamedTuple{(:I, :s, :n, :l, :m), Tuple{Int, Spin, Int, Int, Int}} -mutable struct AtomicOrbitalsBasis{TR, TY, T} +mutable struct ProductBasis{NB, TR, TY} + sparsebasis::SparseProduct{NB} bRnl::TR bYlm::TY - # ------- specification of the atomic orbitals Rnl * Ylm - spec1::Vector{NTRNL1} # human readable spec - spec1idx::Vector{Tuple{Int, Int}} # indices into Rnl, Ylm - nuclei::Vector{Nuc{T}} # nuclei (defines the shifted orbitals) - # ------- specification of the pooling operations A_{sInlm} - # spec::Vector{NTRNLIS} # human readable spec, definition is - # implicit in iteration protocol. end - -function _invmap(a::AbstractVector) - inva = Dict{eltype(a), Int}() - for i = 1:length(a) - inva[a[i]] = i - end - return inva -end - -function AtomicOrbitalsBasis(bRnl, bYlm, spec1::Vector{NTRNL1}, - nuclei::Vector{<: Nuc}) - - spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) - spec_Rnl = bRnl.spec; inv_Rnl = _invmap(spec_Rnl) - spec_Ylm = Polynomials4ML.natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) - - spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) - for (i, b) in enumerate(spec1) - spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) - end - - # spec = NTRNLIS[] - basis = AtomicOrbitalsBasis(bRnl, bYlm, spec1, spec1idx, eltype(nuclei)[]) - set_nuclei!(basis, nuclei) - return basis -end - - -function AtomicOrbitalsBasis(bRnl, bYlm; - totaldegree=3, - nuclei = Nuc{Float64}[], - ) - spec1 = make_nlms_spec(bRnl, bYlm; - totaldegree = totaldegree) - return AtomicOrbitalsBasis(bRnl, bYlm, spec1, nuclei) -end - - """ This constructs the specification of all the atomic orbitals for one nucleus. @@ -99,65 +54,50 @@ function make_nlms_spec(bRnl, bYlm; end -function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) - basis.nuclei = copy(collect(nuclei)) - return nothing -end - - -function get_spec(basis::AtomicOrbitalsBasis) - spec = NTRNLIS[] - Nnuc = length(basis.nuclei) - - spec = Array{NTRNLIS, 3}(undef, (3, Nnuc, length(basis.spec1))) - - for (k, nlm) in enumerate(basis.spec1) - for I = 1:Nnuc - for (is, s) in enumerate(extspins()) - spec[is, I, k] = (I = I, s=s, nlm...) - end - end +function evaluate(basis::ProductBasis, X::AbstractVector{<: AbstractVector}, Σ) + Nel = length(X) + T = promote_type(eltype(X[1])) + VT = SVector{3, T} + @show VT + + # create all the shifted configurations + xx = zeros(eltype(VT), Nel) + for i = 1:Nel + xx[i] = norm(X[i]) end - return spec -end + # evaluate the radial and angular components on all the shifted particles + Rnl = reshape(evaluate(basis.bRnl, xx[:]), (Nel, length(basis.bRnl))) + Ylm = reshape(evaluate(basis.bYlm, X[:]), (Nel, length(basis.bYlm))) + # evaluate all the atomic orbitals as ϕ_nlm = Rnl * Ylm + ϕnlm = evaluate(basis.sparsebasis, (Rnl, Ylm)) -# ------------ Evaluation kernels + return ϕnlm +end +mutable struct AtomicOrbitalsBasis{NB, T} + prodbasis::ProductBasis{NB} + nuclei::Vector{Nuc{T}} # nuclei (defines the shifted orbitals) +end -function proto_evaluate(basis::AtomicOrbitalsBasis, - X::AbstractVector{<: AbstractVector}, - Σ) +function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) nuc = basis.nuclei Nnuc = length(nuc) - Nel = length(X) - Nnlm = length(basis.spec1) - T = promote_type(eltype(nuc[1].rr), eltype(X[1])) - VT = SVector{3, T} - @show VT - # create all the shifted configurations - # this MUST be done in the format (I, i) + XX = zeros(VT, (Nnuc, Nel)) - xx = zeros(eltype(VT), (Nnuc, Nel)) + for I = 1:Nnuc, i = 1:Nel XX[I, i] = X[i] - nuc[I].rr - xx[I, i] = norm(XX[I, i]) end - # evaluate the radial and angular components on all the shifted particles - Rnl = reshape(evaluate(basis.bRnl, xx[:]), (Nnuc, Nel, length(basis.bRnl))) - Ylm = reshape(evaluate(basis.bYlm, XX[:]), (Nnuc, Nel, length(basis.bYlm))) + Nnlm = length(basis.prodbasis.sparsebasis.spec) - # evaluate all the atomic orbitals as ϕ_nlm = Rnl * Ylm - TA = promote_type(eltype(Rnl), eltype(Ylm)) - @show TA ϕnlm = zeros(TA, (Nnuc, Nel, Nnlm)) - for (k, (iR, iY)) in enumerate(basis.spec1idx) - for i = 1:Nel, I = 1:Nnuc - ϕnlm[I, i, k] = Rnl[I, i, iR] * Ylm[I, i, iY] - end + + for I = 1:Nnuc + ϕnlm[I,:,:] = evaluate(basis.prodbasis, XX[I,:], Σ) end # evaluate the pooling operation @@ -207,3 +147,67 @@ function proto_evaluate(basis::AtomicOrbitalsBasis, return A end + + + + + + +function _invmap(a::AbstractVector) + inva = Dict{eltype(a), Int}() + for i = 1:length(a) + inva[a[i]] = i + end + return inva +end + +function _invmap(a::AbstractVector) + inva = Dict{eltype(a), Int}() + for i = 1:length(a) + inva[a[i]] = i + end + return inva +end + +""" + +function AtomicOrbitalsBasis(bRnl, bYlm; + totaldegree=3, + nuclei = Nuc{Float64}[], + ) + spec1 = make_nlms_spec(bRnl, bYlm; + totaldegree = totaldegree) + return AtomicOrbitalsBasis(bRnl, bYlm, spec1, nuclei) +end + + + + +function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) + basis.nuclei = copy(collect(nuclei)) + return nothing +end + + +function get_spec(basis::AtomicOrbitalsBasis) + spec = NTRNLIS[] + Nnuc = length(basis.nuclei) + + spec = Array{NTRNLIS, 3}(undef, (3, Nnuc, length(basis.spec1))) + + for (k, nlm) in enumerate(basis.spec1) + for I = 1:Nnuc + for (is, s) in enumerate(extspins()) + spec[is, I, k] = (I = I, s=s, nlm...) + end + end + end + + return spec +end + + +# ------------ Evaluation kernels + + +""" \ No newline at end of file diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index d208801..7a8a89d 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -1,29 +1,55 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test -using Polynomials4ML: natural_indices, degree +using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, proto_evaluate using Polynomials4ML.Testing: print_tf -## - totdeg = 3 bRnl = ACEpsi.AtomicOrbitals.RnlExample(totdeg) bYlm = RYlmBasis(totdeg) nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] -bAnlm = AtomicOrbitalsBasis(bRnl, bYlm; - totaldegree = totdeg, - nuclei = nuclei ) - -spec = ACEpsi.AtomicOrbitals.get_spec(bAnlm) +totaldegree = 4 +spec1 = make_nlms_spec(bRnl, bYlm; totaldegree = totaldegree) +spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) +spec_Rnl = bRnl.spec; inv_Rnl = _invmap(spec_Rnl) +spec_Ylm = Polynomials4ML.natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) -## +spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) +for (i, b) in enumerate(spec1) + spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) +end +sparsebasis = SparseProduct(spec1idx) +prodbasis = ProductBasis(sparsebasis,bRnl,bYlm) Nel = 5 X = randn(SVector{3, Float64}, Nel) Σ = rand(spins(), Nel) -A = proto_evaluate(bAnlm, X, Σ) +ϕnlm = evaluate(prodbasis, X, Σ) + +bAnlm = AtomicOrbitalsBasis(prodbasis, nuclei) +A = evaluate(bAnlm, X, Σ) + + +""" + + + + + + + + + +## not test +spec = ACEpsi.AtomicOrbitals.get_spec(bAnlm) + +## + + + + ## test the Evaluation: @@ -63,3 +89,4 @@ println() ## +""" \ No newline at end of file From 750a05501537fd6702a4c5af970f031cd4958f9b Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Sun, 7 May 2023 19:31:48 -0700 Subject: [PATCH 09/79] a small clean up for the draft for constructing A basis, depends on 0.1.0-dev polynomials4ML --- src/atomicorbitals/atorbbasis.jl | 24 ++++------ test/test_atorbbasis.jl | 81 +++++++++++++++----------------- 2 files changed, 47 insertions(+), 58 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 8083c1b..c7627e5 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -1,6 +1,8 @@ +import Polynomials4ML: evaluate using ACEpsi: ↑, ↓, ∅, spins, extspins, Spin, spin2idx, idx2spin -import Polynomials4ML +using Polynomials4ML: SparseProduct + using StaticArrays using LinearAlgebra: norm @@ -58,7 +60,6 @@ function evaluate(basis::ProductBasis, X::AbstractVector{<: AbstractVector}, Σ) Nel = length(X) T = promote_type(eltype(X[1])) VT = SVector{3, T} - @show VT # create all the shifted configurations xx = zeros(eltype(VT), Nel) @@ -84,7 +85,9 @@ end function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) nuc = basis.nuclei Nnuc = length(nuc) - + Nel = size(X, 1) + T = promote_type(eltype(X[1])) + VT = SVector{3, T} XX = zeros(VT, (Nnuc, Nel)) @@ -93,8 +96,8 @@ function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVecto end Nnlm = length(basis.prodbasis.sparsebasis.spec) - - ϕnlm = zeros(TA, (Nnuc, Nel, Nnlm)) + + ϕnlm = zeros(T, (Nnuc, Nel, Nnlm)) for I = 1:Nnuc ϕnlm[I,:,:] = evaluate(basis.prodbasis, XX[I,:], Σ) @@ -102,7 +105,7 @@ function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVecto # evaluate the pooling operation # spin I k = (nlm) - Aall = zeros(TA, (2, Nnuc, Nnlm)) + Aall = zeros(T, (2, Nnuc, Nnlm)) for k = 1:Nnlm for i = 1:Nel iσ = spin2idx(Σ[i]) @@ -129,7 +132,7 @@ function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVecto @assert spin2idx(↑) == 1 @assert spin2idx(↓) == 2 @assert spin2idx(∅) == 3 - A = zeros(TA, ((Nel, 3, Nnuc, Nnlm))) + A = zeros(T, ((Nel, 3, Nnuc, Nnlm))) for k = 1:Nnlm for I = 1:Nnuc for i = 1:Nel @@ -161,13 +164,6 @@ function _invmap(a::AbstractVector) return inva end -function _invmap(a::AbstractVector) - inva = Dict{eltype(a), Int}() - for i = 1:length(a) - inva[a[i]] = i - end - return inva -end """ diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 7a8a89d..9630bad 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -1,7 +1,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, proto_evaluate +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, _invmap, ProductBasis, evaluate using Polynomials4ML.Testing: print_tf totdeg = 3 @@ -20,7 +20,7 @@ for (i, b) in enumerate(spec1) spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) end sparsebasis = SparseProduct(spec1idx) -prodbasis = ProductBasis(sparsebasis,bRnl,bYlm) +prodbasis = ProductBasis(sparsebasis, bRnl, bYlm) Nel = 5 X = randn(SVector{3, Float64}, Nel) @@ -29,64 +29,57 @@ X = randn(SVector{3, Float64}, Nel) ϕnlm = evaluate(prodbasis, X, Σ) bAnlm = AtomicOrbitalsBasis(prodbasis, nuclei) -A = evaluate(bAnlm, X, Σ) - - -""" - - - +A = evaluate(bAnlm, X, Σ) -## not test -spec = ACEpsi.AtomicOrbitals.get_spec(bAnlm) +# ## not test +# spec = ACEpsi.AtomicOrbitals.get_spec(bAnlm) -## +# ## -## test the Evaluation: -using LinearAlgebra: norm -bYlm_ = RYlmBasis(totdeg) -Nnlm = length(bAnlm.spec1) -Nnuc = length(bAnlm.nuclei) -inv_nl = ACEpsi.AtomicOrbitals._invmap(bRnl.spec) -inv_lm = ACEpsi.AtomicOrbitals._invmap(natural_indices(bYlm)) +# ## test the Evaluation: +# using LinearAlgebra: norm +# bYlm_ = RYlmBasis(totdeg) +# Nnlm = length(bAnlm.spec1) +# Nnuc = length(bAnlm.nuclei) +# inv_nl = ACEpsi.AtomicOrbitals._invmap(bRnl.spec) +# inv_lm = ACEpsi.AtomicOrbitals._invmap(natural_indices(bYlm)) -for I = 1:Nnuc - XI = X .- Ref(bAnlm.nuclei[I].rr) - xI = norm.(XI) - Rnl = evaluate(bRnl, xI) - Ylm = evaluate(bYlm_, XI) - for k = 1:Nnlm - nlm = bAnlm.spec1[k] - iR = inv_nl[(n = nlm.n, l = nlm.l)] - iY = inv_lm[(l = nlm.l, m = nlm.m)] +# for I = 1:Nnuc +# XI = X .- Ref(bAnlm.nuclei[I].rr) +# xI = norm.(XI) +# Rnl = evaluate(bRnl, xI) +# Ylm = evaluate(bYlm_, XI) +# for k = 1:Nnlm +# nlm = bAnlm.spec1[k] +# iR = inv_nl[(n = nlm.n, l = nlm.l)] +# iY = inv_lm[(l = nlm.l, m = nlm.m)] - for i = 1:Nel - for (is, s) in enumerate(ACEpsi.extspins()) - a1 = A[i, is, I, k] +# for i = 1:Nel +# for (is, s) in enumerate(ACEpsi.extspins()) +# a1 = A[i, is, I, k] - if s in [↑, ↓] - a2 = sum( Rnl[j, iR] * Ylm[j, iY] * (Σ[j] == s) * (1 - (j == i)) for j = 1:Nel ) - else # s = ∅ - a2 = Rnl[i, iR] * Ylm[i, iY] - end - # println("(i=$i, σ=$s, I=$I, n=$(nlm.n), l=$(nlm.l), m=$(nlm.m)) -> ", abs(a1 - a2)) - print_tf(@test abs(a1 - a2) < 1e-12) - end - end - end -end -println() +# if s in [↑, ↓] +# a2 = sum( Rnl[j, iR] * Ylm[j, iY] * (Σ[j] == s) * (1 - (j == i)) for j = 1:Nel ) +# else # s = ∅ +# a2 = Rnl[i, iR] * Ylm[i, iY] +# end +# # println("(i=$i, σ=$s, I=$I, n=$(nlm.n), l=$(nlm.l), m=$(nlm.m)) -> ", abs(a1 - a2)) +# print_tf(@test abs(a1 - a2) < 1e-12) +# end +# end +# end +# end +# println() ## -""" \ No newline at end of file From caee2beb4512a80702440f002c096ab4eb941d29 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Sun, 7 May 2023 21:01:58 -0700 Subject: [PATCH 10/79] add backflowpooing, define sensible constructor and test evaluate of A basis --- src/ACEpsi.jl | 2 + src/atomicorbitals/atorbbasis.jl | 92 +++++++++++-------- src/backflowpooling.jl | 69 ++++++++++++++ test/test_atorbbasis.jl | 150 ++++++++++++++++--------------- 4 files changed, 205 insertions(+), 108 deletions(-) create mode 100644 src/backflowpooling.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index 52b496f..ad223ba 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -11,4 +11,6 @@ include("envelope.jl") include("atomicorbitals/atomicorbitals.jl") # include("bflow3d.jl") +include("backflowpooling.jl") + end diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index c7627e5..b297aba 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -3,6 +3,7 @@ import Polynomials4ML: evaluate using ACEpsi: ↑, ↓, ∅, spins, extspins, Spin, spin2idx, idx2spin using Polynomials4ML: SparseProduct +using Polynomials4ML using StaticArrays using LinearAlgebra: norm @@ -22,10 +23,26 @@ const NTRNLIS = NamedTuple{(:I, :s, :n, :l, :m), Tuple{Int, Spin, Int, Int, Int} mutable struct ProductBasis{NB, TR, TY} sparsebasis::SparseProduct{NB} + spec1::Vector{NTRNL1} bRnl::TR bYlm::TY end + +function ProductBasis(spec1, bRnl, bYlm) + spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) + spec_Rnl = bRnl.spec; inv_Rnl = _invmap(spec_Rnl) + spec_Ylm = Polynomials4ML.natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) + + spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) + for (i, b) in enumerate(spec1) + spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) + end + sparsebasis = SparseProduct(spec1idx) + return ProductBasis(sparsebasis, spec1, bRnl, bYlm) +end + + """ This constructs the specification of all the atomic orbitals for one nucleus. @@ -103,17 +120,23 @@ function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVecto ϕnlm[I,:,:] = evaluate(basis.prodbasis, XX[I,:], Σ) end + + + + return ϕnlm # evaluate the pooling operation # spin I k = (nlm) - Aall = zeros(T, (2, Nnuc, Nnlm)) - for k = 1:Nnlm - for i = 1:Nel - iσ = spin2idx(Σ[i]) - for I = 1:Nnuc - Aall[iσ, I, k] += ϕnlm[I, i, k] - end - end - end + + + # Aall = zeros(T, (2, Nnuc, Nnlm)) + # for k = 1:Nnlm + # for i = 1:Nel + # iσ = spin2idx(Σ[i]) + # for I = 1:Nnuc + # Aall[iσ, I, k] += ϕnlm[I, i, k] + # end + # end + # end # now correct the pooling Aall and write into A^(i) # with do it with i leading so that the N-correlations can @@ -129,25 +152,26 @@ function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVecto # TODO: discuss - this could be stored much more efficiently as a # lazy array. Could that have advantages? # - @assert spin2idx(↑) == 1 - @assert spin2idx(↓) == 2 - @assert spin2idx(∅) == 3 - A = zeros(T, ((Nel, 3, Nnuc, Nnlm))) - for k = 1:Nnlm - for I = 1:Nnuc - for i = 1:Nel - A[i, 3, I, k] = ϕnlm[I, i, k] - end - for iσ = 1:2 - σ = idx2spin(iσ) - for i = 1:Nel - A[i, iσ, I, k] = Aall[iσ, I, k] - (Σ[i] == σ) * ϕnlm[I, i, k] - end - end - end - end - return A + # @assert spin2idx(↑) == 1 + # @assert spin2idx(↓) == 2 + # @assert spin2idx(∅) == 3 + # A = zeros(T, ((Nel, 3, Nnuc, Nnlm))) + # for k = 1:Nnlm + # for I = 1:Nnuc + # for i = 1:Nel + # A[i, 3, I, k] = ϕnlm[I, i, k] + # end + # for iσ = 1:2 + # σ = idx2spin(iσ) + # for i = 1:Nel + # A[i, iσ, I, k] = Aall[iσ, I, k] - (Σ[i] == σ) * ϕnlm[I, i, k] + # end + # end + # end + # end + + # return A end @@ -165,20 +189,17 @@ function _invmap(a::AbstractVector) end -""" - function AtomicOrbitalsBasis(bRnl, bYlm; totaldegree=3, nuclei = Nuc{Float64}[], ) spec1 = make_nlms_spec(bRnl, bYlm; - totaldegree = totaldegree) - return AtomicOrbitalsBasis(bRnl, bYlm, spec1, nuclei) + totaldegree = totaldegree) + prodbasis = ProductBasis(spec1, bRnl, bYlm) + return AtomicOrbitalsBasis(prodbasis, nuclei) end - - function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) basis.nuclei = copy(collect(nuclei)) return nothing @@ -189,9 +210,9 @@ function get_spec(basis::AtomicOrbitalsBasis) spec = NTRNLIS[] Nnuc = length(basis.nuclei) - spec = Array{NTRNLIS, 3}(undef, (3, Nnuc, length(basis.spec1))) + spec = Array{NTRNLIS, 3}(undef, (3, Nnuc, length(basis.prodbasis.spec1))) - for (k, nlm) in enumerate(basis.spec1) + for (k, nlm) in enumerate(basis.prodbasis.spec1) for I = 1:Nnuc for (is, s) in enumerate(extspins()) spec[is, I, k] = (I = I, s=s, nlm...) @@ -206,4 +227,3 @@ end # ------------ Evaluation kernels -""" \ No newline at end of file diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl new file mode 100644 index 0000000..13502e5 --- /dev/null +++ b/src/backflowpooling.jl @@ -0,0 +1,69 @@ +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis + +mutable struct BackflowPooling + basis::AtomicOrbitalsBasis +end + +function evaluate(pooling::BackflowPooling, ϕnlm, Σ) + basis = pooling.basis + nuc = basis.nuclei + Nnuc = length(nuc) + Nnlm = length(basis.prodbasis.sparsebasis.spec) + Nel = length(Σ) + T = promote_type(eltype(ϕnlm)) + + # evaluate the pooling operation + # spin I k = (nlm) + + Aall = zeros(T, (2, Nnuc, Nnlm)) + for k = 1:Nnlm + for i = 1:Nel + iσ = spin2idx(Σ[i]) + for I = 1:Nnuc + Aall[iσ, I, k] += ϕnlm[I, i, k] + end + end + end + + # now correct the pooling Aall and write into A^(i) + # with do it with i leading so that the N-correlations can + # be parallelized over i + # + # A[i, :] = A-basis for electron i, with channels, s, I, k=nlm + # A[i, ∅, I, k] = ϕnlm[I, i, k] + # for σ = ↑ or ↓ we have + # A[i, σ, I, k] = ∑_{j ≂̸ i : Σ[j] == σ} ϕnlm[I, j, k] + # = ∑_{j : Σ[j] == σ} ϕnlm[I, j, k] - (Σ[i] == σ) * ϕnlm[I, i, k] + # + # + # TODO: discuss - this could be stored much more efficiently as a + # lazy array. Could that have advantages? + # + + @assert spin2idx(↑) == 1 + @assert spin2idx(↓) == 2 + @assert spin2idx(∅) == 3 + A = zeros(T, ((Nel, 3, Nnuc, Nnlm))) + for k = 1:Nnlm + for I = 1:Nnuc + for i = 1:Nel + A[i, 3, I, k] = ϕnlm[I, i, k] + end + for iσ = 1:2 + σ = idx2spin(iσ) + for i = 1:Nel + A[i, iσ, I, k] = Aall[iσ, I, k] - (Σ[i] == σ) * ϕnlm[I, i, k] + end + end + end + end + + return A +end + + + + + + + diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 9630bad..e1ad2f2 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -1,85 +1,91 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, _invmap, ProductBasis, evaluate +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate +using ACEpsi: BackflowPooling using Polynomials4ML.Testing: print_tf -totdeg = 3 -bRnl = ACEpsi.AtomicOrbitals.RnlExample(totdeg) -bYlm = RYlmBasis(totdeg) -nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] - -totaldegree = 4 -spec1 = make_nlms_spec(bRnl, bYlm; totaldegree = totaldegree) -spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) -spec_Rnl = bRnl.spec; inv_Rnl = _invmap(spec_Rnl) -spec_Ylm = Polynomials4ML.natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) - -spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) -for (i, b) in enumerate(spec1) - spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) -end -sparsebasis = SparseProduct(spec1idx) -prodbasis = ProductBasis(sparsebasis, bRnl, bYlm) - +# test configs +Rnldegree = 4 +Ylmdegree = 4 +totdegree = 8 Nel = 5 X = randn(SVector{3, Float64}, Nel) Σ = rand(spins(), Nel) -ϕnlm = evaluate(prodbasis, X, Σ) - -bAnlm = AtomicOrbitalsBasis(prodbasis, nuclei) - - - -A = evaluate(bAnlm, X, Σ) - - - -# ## not test -# spec = ACEpsi.AtomicOrbitals.get_spec(bAnlm) - -# ## - - - - - - -# ## test the Evaluation: -# using LinearAlgebra: norm -# bYlm_ = RYlmBasis(totdeg) -# Nnlm = length(bAnlm.spec1) -# Nnuc = length(bAnlm.nuclei) -# inv_nl = ACEpsi.AtomicOrbitals._invmap(bRnl.spec) -# inv_lm = ACEpsi.AtomicOrbitals._invmap(natural_indices(bYlm)) - -# for I = 1:Nnuc -# XI = X .- Ref(bAnlm.nuclei[I].rr) -# xI = norm.(XI) -# Rnl = evaluate(bRnl, xI) -# Ylm = evaluate(bYlm_, XI) -# for k = 1:Nnlm -# nlm = bAnlm.spec1[k] -# iR = inv_nl[(n = nlm.n, l = nlm.l)] -# iY = inv_lm[(l = nlm.l, m = nlm.m)] +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] +## + +# Defining AtomicOrbitalsBasis +bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +bYlm = RYlmBasis(Ylmdegree) +spec1 = make_nlms_spec(bRnl, bYlm; totaldegree = totdegree) + +# define basis and pooling operations +prodbasis = ProductBasis(spec1, bRnl, bYlm) +aobasis = AtomicOrbitalsBasis(prodbasis, nuclei) +pooling = BackflowPooling(aobasis) + +# we can also construct in this way... To be discuss which one we want / keep both +aobasis2 = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdegree, nuclei = nuclei, ) + +@info("Checking two type of construction are same") +for ntest = 1:30 + local X = randn(SVector{3, Float64}, Nel) + local Σ = rand(spins(), Nel) + print_tf(@test evaluate(aobasis, X, Σ) ≈ evaluate(aobasis2, X, Σ)) +end +println() -# for i = 1:Nel -# for (is, s) in enumerate(ACEpsi.extspins()) -# a1 = A[i, is, I, k] +@info("Test evaluate ProductBasis") +ϕnlm = evaluate(prodbasis, X, Σ) -# if s in [↑, ↓] -# a2 = sum( Rnl[j, iR] * Ylm[j, iY] * (Σ[j] == s) * (1 - (j == i)) for j = 1:Nel ) -# else # s = ∅ -# a2 = Rnl[i, iR] * Ylm[i, iY] -# end -# # println("(i=$i, σ=$s, I=$I, n=$(nlm.n), l=$(nlm.l), m=$(nlm.m)) -> ", abs(a1 - a2)) -# print_tf(@test abs(a1 - a2) < 1e-12) -# end -# end -# end -# end -# println() +@info("Test evaluate AtomicOrbitalsBasis") +bϕnlm = evaluate(aobasis, X, Σ) + +@info("Test BackflowPooling") +A = ACEpsi.evaluate(pooling, bϕnlm, Σ) + +println() + + +## +@info("Check get_spec is working") +spec = ACEpsi.AtomicOrbitals.get_spec(aobasis) + + +@info("Test evaluation by manual construction") +using LinearAlgebra: norm +bYlm_ = RYlmBasis(totdegree) +Nnlm = length(aobasis.prodbasis.sparsebasis.spec) +Nnuc = length(aobasis.nuclei) + +for I = 1:Nnuc + XI = X .- Ref(aobasis.nuclei[I].rr) + xI = norm.(XI) + Rnl = evaluate(bRnl, xI) + Ylm = evaluate(bYlm_, XI) + for k = 1:Nnlm + nlm = aobasis.prodbasis.sparsebasis.spec[k] + iR = nlm[1] + iY = nlm[2] + + for i = 1:Nel + for (is, s) in enumerate(ACEpsi.extspins()) + a1 = A[i, is, I, k] + + if s in [↑, ↓] + a2 = sum( Rnl[j, iR] * Ylm[j, iY] * (Σ[j] == s) * (1 - (j == i)) for j = 1:Nel ) + else # s = ∅ + a2 = Rnl[i, iR] * Ylm[i, iY] + end + # println("(i=$i, σ=$s, I=$I, n=$(nlm.n), l=$(nlm.l), m=$(nlm.m)) -> ", abs(a1 - a2)) + print_tf(@test a1 ≈ a2) + end + end + end +end +println() -## +# From efea037fe59319a90279f78cc955802eb9d952b3 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Sun, 7 May 2023 21:15:13 -0700 Subject: [PATCH 11/79] small clean ups --- src/atomicorbitals/atorbbasis.jl | 149 ++++++++++--------------------- test/test_atorbbasis.jl | 3 +- 2 files changed, 49 insertions(+), 103 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index b297aba..82e2323 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -21,28 +21,6 @@ end const NTRNL1 = NamedTuple{(:n, :l, :m), Tuple{Int, Int, Int}} const NTRNLIS = NamedTuple{(:I, :s, :n, :l, :m), Tuple{Int, Spin, Int, Int, Int}} -mutable struct ProductBasis{NB, TR, TY} - sparsebasis::SparseProduct{NB} - spec1::Vector{NTRNL1} - bRnl::TR - bYlm::TY -end - - -function ProductBasis(spec1, bRnl, bYlm) - spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) - spec_Rnl = bRnl.spec; inv_Rnl = _invmap(spec_Rnl) - spec_Ylm = Polynomials4ML.natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) - - spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) - for (i, b) in enumerate(spec1) - spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) - end - sparsebasis = SparseProduct(spec1idx) - return ProductBasis(sparsebasis, spec1, bRnl, bYlm) -end - - """ This constructs the specification of all the atomic orbitals for one nucleus. @@ -73,6 +51,38 @@ function make_nlms_spec(bRnl, bYlm; end +# Jerry: This is just a specific case of a general ProductBasis, this should go to Polynomials4ML later with a general implementation +# I will do that after reconfiming this is what we want +mutable struct ProductBasis{NB, TR, TY} + spec1::Vector{NTRNL1} + bRnl::TR + bYlm::TY + # ---- evaluation kernel from Polynomials4ML ---- + sparsebasis::SparseProduct{NB} +end + +function _invmap(a::AbstractVector) + inva = Dict{eltype(a), Int}() + for i = 1:length(a) + inva[a[i]] = i + end + return inva +end + +function ProductBasis(spec1, bRnl, bYlm) + spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) + spec_Rnl = bRnl.spec; inv_Rnl = _invmap(spec_Rnl) + spec_Ylm = Polynomials4ML.natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) + + spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) + for (i, b) in enumerate(spec1) + spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) + end + sparsebasis = SparseProduct(spec1idx) + return ProductBasis(sparsebasis, spec1, bRnl, bYlm) +end + + function evaluate(basis::ProductBasis, X::AbstractVector{<: AbstractVector}, Σ) Nel = length(X) T = promote_type(eltype(X[1])) @@ -99,107 +109,42 @@ mutable struct AtomicOrbitalsBasis{NB, T} nuclei::Vector{Nuc{T}} # nuclei (defines the shifted orbitals) end + +function AtomicOrbitalsBasis(bRnl, bYlm; + totaldegree=3, + nuclei = Nuc{Float64}[], + ) + spec1 = make_nlms_spec(bRnl, bYlm; + totaldegree = totaldegree) + prodbasis = ProductBasis(spec1, bRnl, bYlm) + return AtomicOrbitalsBasis(prodbasis, nuclei) +end + + function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) nuc = basis.nuclei Nnuc = length(nuc) Nel = size(X, 1) T = promote_type(eltype(X[1])) VT = SVector{3, T} - XX = zeros(VT, (Nnuc, Nel)) + # trans for I = 1:Nnuc, i = 1:Nel XX[I, i] = X[i] - nuc[I].rr end Nnlm = length(basis.prodbasis.sparsebasis.spec) - ϕnlm = zeros(T, (Nnuc, Nel, Nnlm)) - for I = 1:Nnuc + for I = 1:Nnuc ϕnlm[I,:,:] = evaluate(basis.prodbasis, XX[I,:], Σ) end - - - return ϕnlm - # evaluate the pooling operation - # spin I k = (nlm) - - - # Aall = zeros(T, (2, Nnuc, Nnlm)) - # for k = 1:Nnlm - # for i = 1:Nel - # iσ = spin2idx(Σ[i]) - # for I = 1:Nnuc - # Aall[iσ, I, k] += ϕnlm[I, i, k] - # end - # end - # end - - # now correct the pooling Aall and write into A^(i) - # with do it with i leading so that the N-correlations can - # be parallelized over i - # - # A[i, :] = A-basis for electron i, with channels, s, I, k=nlm - # A[i, ∅, I, k] = ϕnlm[I, i, k] - # for σ = ↑ or ↓ we have - # A[i, σ, I, k] = ∑_{j ≂̸ i : Σ[j] == σ} ϕnlm[I, j, k] - # = ∑_{j : Σ[j] == σ} ϕnlm[I, j, k] - (Σ[i] == σ) * ϕnlm[I, i, k] - # - # - # TODO: discuss - this could be stored much more efficiently as a - # lazy array. Could that have advantages? - # - - # @assert spin2idx(↑) == 1 - # @assert spin2idx(↓) == 2 - # @assert spin2idx(∅) == 3 - # A = zeros(T, ((Nel, 3, Nnuc, Nnlm))) - # for k = 1:Nnlm - # for I = 1:Nnuc - # for i = 1:Nel - # A[i, 3, I, k] = ϕnlm[I, i, k] - # end - # for iσ = 1:2 - # σ = idx2spin(iσ) - # for i = 1:Nel - # A[i, iσ, I, k] = Aall[iσ, I, k] - (Σ[i] == σ) * ϕnlm[I, i, k] - # end - # end - # end - # end - - # return A -end - - - - - - - -function _invmap(a::AbstractVector) - inva = Dict{eltype(a), Int}() - for i = 1:length(a) - inva[a[i]] = i - end - return inva end - -function AtomicOrbitalsBasis(bRnl, bYlm; - totaldegree=3, - nuclei = Nuc{Float64}[], - ) - spec1 = make_nlms_spec(bRnl, bYlm; - totaldegree = totaldegree) - prodbasis = ProductBasis(spec1, bRnl, bYlm) - return AtomicOrbitalsBasis(prodbasis, nuclei) -end - - +# ------------ utils for AtomicOrbitalsBasis ------------ function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) basis.nuclei = copy(collect(nuclei)) return nothing diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index e1ad2f2..2301a28 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -26,7 +26,7 @@ prodbasis = ProductBasis(spec1, bRnl, bYlm) aobasis = AtomicOrbitalsBasis(prodbasis, nuclei) pooling = BackflowPooling(aobasis) -# we can also construct in this way... To be discuss which one we want / keep both +# we can also construct in this way which wraps the definition of product basis inside aobasis2 = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdegree, nuclei = nuclei, ) @info("Checking two type of construction are same") @@ -35,6 +35,7 @@ for ntest = 1:30 local Σ = rand(spins(), Nel) print_tf(@test evaluate(aobasis, X, Σ) ≈ evaluate(aobasis2, X, Σ)) end + println() @info("Test evaluate ProductBasis") From 541359aeb76db18097d98163b5d0ff86ba76650a Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Mon, 15 May 2023 02:47:52 -0700 Subject: [PATCH 12/79] add lux implementation draft --- Project.toml | 11 +- src/ACEpsi.jl | 2 +- src/atomicorbitals/atorbbasis.jl | 46 +- src/backflowpooling.jl | 25 + src/bflow3d.jl | 961 +++++++++++++++++-------------- test/test_atorbbasis.jl | 31 +- test/test_bflow_lux.jl | 37 ++ 7 files changed, 656 insertions(+), 457 deletions(-) create mode 100644 test/test_bflow_lux.jl diff --git a/Project.toml b/Project.toml index cb754d0..6afc154 100644 --- a/Project.toml +++ b/Project.toml @@ -10,20 +10,23 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d" Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b" [compat] -ACEbase = "0.2" -ACEcore = "0.0.1" +ACEbase = "0.4.2" +ACEcore = "0.0.5" BenchmarkTools = "1" ForwardDiff = "0.10" JSON = "0.21" -ObjectPools = "0.0.2" -Polynomials4ML = "0.0.6" +ObjectPools = "0.2.1" +Polynomials4ML = "0.1.3" julia = "1" [extras] diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index ad223ba..f37be0c 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -9,7 +9,7 @@ include("envelope.jl") # the new 3d backflow code include("atomicorbitals/atomicorbitals.jl") -# include("bflow3d.jl") +include("bflow3d.jl") include("backflowpooling.jl") diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 82e2323..48101e5 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -2,6 +2,8 @@ import Polynomials4ML: evaluate using ACEpsi: ↑, ↓, ∅, spins, extspins, Spin, spin2idx, idx2spin using Polynomials4ML: SparseProduct +using LuxCore: AbstractExplicitLayer +using Random: AbstractRNG using Polynomials4ML using StaticArrays @@ -61,6 +63,8 @@ mutable struct ProductBasis{NB, TR, TY} sparsebasis::SparseProduct{NB} end +(PB::ProductBasis)(args...) = evaluate(PB, args...) + function _invmap(a::AbstractVector) inva = Dict{eltype(a), Int}() for i = 1:length(a) @@ -79,17 +83,16 @@ function ProductBasis(spec1, bRnl, bYlm) spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) end sparsebasis = SparseProduct(spec1idx) - return ProductBasis(sparsebasis, spec1, bRnl, bYlm) + return ProductBasis(spec1, bRnl, bYlm, sparsebasis) end function evaluate(basis::ProductBasis, X::AbstractVector{<: AbstractVector}, Σ) Nel = length(X) T = promote_type(eltype(X[1])) - VT = SVector{3, T} - # create all the shifted configurations - xx = zeros(eltype(VT), Nel) + # create all the shifted configurations + xx = zeros(eltype(T), Nel) for i = 1:Nel xx[i] = norm(X[i]) end @@ -109,6 +112,8 @@ mutable struct AtomicOrbitalsBasis{NB, T} nuclei::Vector{Nuc{T}} # nuclei (defines the shifted orbitals) end +(aobasis::AtomicOrbitalsBasis)(args...) = evaluate(aobasis, args...) + function AtomicOrbitalsBasis(bRnl, bYlm; totaldegree=3, @@ -172,3 +177,36 @@ end # ------------ Evaluation kernels +# ------------ connect with Lux +struct AtomicOrbitalsBasisLayer{TB} <: AbstractExplicitLayer + basis::TB + meta::Dict{String, Any} +end + +Base.length(l::AtomicOrbitalsBasisLayer) = length(l.basis) + +lux(basis::AtomicOrbitalsBasis) = AtomicOrbitalsBasisLayer(basis, Dict{String, Any}()) + +initialparameters(rng::AbstractRNG, l::AtomicOrbitalsBasisLayer) = _init_luxparams(rng, l.basis) + +initialstates(rng::AbstractRNG, l::AtomicOrbitalsBasisLayer) = _init_luxstate(rng, l.basis) + +(l::AtomicOrbitalsBasisLayer)(X, ps, st) = + evaluate(l.basis, X, st.Σ), st + + + +# ----- ObejctPools +# (l::AtomicOrbitalsBasisLayer)(args...) = evaluate(l, args...) + +# function evaluate(l::AtomicOrbitalsBasisLayer, x::SINGLE, ps, st) +# B = acquire!(st.cache, :B, (length(l.basis), ), _valtype(l.basis, x)) +# evaluate!(parent(B), l.basis, x) +# return B +# end + +# function evaluate(l::AtomicOrbitalsBasisLayer, X::AbstractArray{<: SINGLE}, ps, st) +# B = acquire!(st.cache[:Bbatch], (length(l.basis), length(X)), _valtype(l.basis, X[1])) +# evaluate!(parent(B), l.basis, X) +# return B +# end diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index 13502e5..2607bf0 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -1,9 +1,13 @@ using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis +using LuxCore: AbstractExplicitLayer +using Random: AbstractRNG mutable struct BackflowPooling basis::AtomicOrbitalsBasis end +(pooling::BackflowPooling)(args...) = evaluate(pooling, args...) + function evaluate(pooling::BackflowPooling, ϕnlm, Σ) basis = pooling.basis nuc = basis.nuclei @@ -63,7 +67,28 @@ end +# --------------------- connect with Lux + +struct BackflowPoolingLayer <: AbstractExplicitLayer + basis::BackflowPooling +end + +lux(basis::BackflowPooling) = BackflowPoolingLayer(basis) + +initialparameters(rng::AbstractRNG, l::BackflowPoolingLayer) = _init_luxparams(rng, l.basis) + +initialstates(rng::AbstractRNG, l::BackflowPoolingLayer) = _init_luxstate(rng, l.basis) +# This should be removed later and replace by ObejctPools +(l::BackflowPoolingLayer)(ϕnlm, ps, st) = + evaluate(l.basis, ϕnlm, st.Σ), st +# ----- ObejctPools +# (l::BackflowPoolingLayer)(args...) = evaluate(l, args...) +# function evaluate(l::BackflowPoolingLayer, ϕnlm_Σ::SINGLE, ps, st) +# B = acquire!(st.cache, :B, (length(l.basis), ), _valtype(l.basis, x)) +# evaluate!(parent(B), l.basis, ϕnlm_Σ[1], ϕnlm_Σ[2]) +# return B +# end diff --git a/src/bflow3d.jl b/src/bflow3d.jl index c911046..87b4a82 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -5,117 +5,125 @@ using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! using ACEcore.Utils: gensparse using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr import ForwardDiff +using ACEpsi.AtomicOrbitals: make_nlms_spec +using LuxCore: AbstractExplicitLayer +using Lux: Dense, Chain # ---------------------------------------- -""" -This function returns correct Si for pooling operation. -""" -function onehot!(Si, i, Σ) - Si .= 0 - for k = 1:length(Σ) - Si[k, spin2num(Σ[k])] = 1 - end - # set current electron to ϕ, also remove their contribution in the sum of ↑ or ↓ basis - Si[i, 1] = 1 - Si[i, 2] = 0 - Si[i, 3] = 0 +# """ +# This function returns correct Si for pooling operation. +# """ +# function onehot!(Si, i, Σ) +# Si .= 0 +# for k = 1:length(Σ) +# Si[k, spin2num(Σ[k])] = 1 +# end +# # set current electron to ϕ, also remove their contribution in the sum of ↑ or ↓ basis +# Si[i, 1] = 1 +# Si[i, 2] = 0 +# Si[i, 3] = 0 +# end + + +# """ +# This function convert spin to corresponding integer value used in spec +# """ +# function spin2num(σ) +# if σ == '↑' +# return 2 +# elseif σ == '↓' +# return 3 +# elseif σ == '∅' +# return 1 +# end +# error("illegal spin char for spin2num") +# end + +# """ +# This function convert num to corresponding spin string. +# """ +# function num2spin(σ) +# if σ == 2 +# return '↑' +# elseif σ == 3 +# return '↓' +# elseif σ == 1 +# return '∅' +# end +# error("illegal integer value for num2spin") +# end + +# function assemble_A(bA::RnlYlmEmbedding, X::AbstractVector, Σ) + +# nX = length(X) +# # position embedding +# P = wf.P +# Xt = wf.trans.(X) +# evaluate!(P, wf.polys, Xt) # nX x dim(polys) + +# A = wf.A # zeros(nX, length(wf.pooling)) +# Ai = wf.Ai # zeros(length(wf.pooling)) +# Si = wf.Si # zeros(Bool, nX, 2) + +# for i = 1:nX +# onehot!(Si, i, Σ) +# ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) +# A[i, :] .= Ai +# end +# return A +# end + +struct MaskLayer <: AbstractExplicitLayer + nX::Integer end +(l::MaskLayer)(Φ, ps, st) = Φ .* [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX], st -""" -This function convert spin to corresponding integer value used in spec -""" -function spin2num(σ) - if σ == '↑' - return 2 - elseif σ == '↓' - return 3 - elseif σ == '∅' - return 1 - end - error("illegal spin char for spin2num") -end +# TODO: maybe in here we can do fast-reshape with stride array? +struct ReshapeLayer <: AbstractExplicitLayer end -""" -This function convert num to corresponding spin string. -""" -function num2spin(σ) - if σ == 2 - return '↑' - elseif σ == 3 - return '↓' - elseif σ == 1 - return '∅' - end - error("illegal integer value for num2spin") +@inline function (f::ReshapeLayer)(x::AbstractArray{T, N}, ps, st::NamedTuple) where {T, N} + # (nX, 3, length(nuclei), length(spec1p)) -> (nX, 3 * length(nuclei) * length(spec1p)) each row go over the spin channel first, then nuclei, then spec1p + return reshape(x, (size(x, 1), prod(size(x)[2:end]))), st end -function assemble_A(bA::RnlYlmEmbedding, X::AbstractVector, Σ) - - nX = length(X) - # position embedding - P = wf.P - Xt = wf.trans.(X) - evaluate!(P, wf.polys, Xt) # nX x dim(polys) - - A = wf.A # zeros(nX, length(wf.pooling)) - Ai = wf.Ai # zeros(length(wf.pooling)) - Si = wf.Si # zeros(Bool, nX, 2) - - for i = 1:nX - onehot!(Si, i, Σ) - ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) - A[i, :] .= Ai - end - return A +struct TransposeLayer <: AbstractExplicitLayer end + +@inline function (f::TransposeLayer)(x::AbstractArray{T, N}, ps, st::NamedTuple) where {T, N} + return transpose(x), st end +struct logabsdetLayer <: AbstractExplicitLayer end +@inline function (f::logabsdetLayer)(Φ::AbstractArray{T, N}, ps, st::NamedTuple) where {T, N} + return 2 * logabsdet(Φ)[1], st +end -# ---------------------------------------- +function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, + ν = 3, T = Float64, + trans = identity, + sd_admissible = bb -> (true), + envelope = x -> x) + spec1p = make_nlms_spec(bRnl, bYlm; + totaldegree = totdeg) + # size(X) = (nX, 3); length(Σ) = nX + # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1 from totaldegree)) + aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) -mutable struct BFwf{T, TT, TPOLY, TE} - embedding::TA - corr::SparseSymmProdDAG{T} - W::Matrix{T} - envelope::TE - spec::Vector{Vector{Int64}} # corr.spec TODO: this needs to be remove - # ---------------- Temporaries - P::Matrix{T} - ∂P::Matrix{T} - dP::Matrix{T} - Φ::Matrix{T} - ∂Φ::Matrix{T} - A::Matrix{T} - ∂A::Matrix{T} - Ai::Vector{T} - ∂Ai::Vector{T} - Si::Matrix{Bool} - ∂AA::Matrix{T} - ∂Si::Matrix{T} - ∇AA::Array{T, 3} -end + # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1 from totaldegree)) + pooling = BackflowPooling(aobasis) -(Φ::BFwf)(args...) = evaluate(Φ, args...) - -function BFwf(Nel::Integer, polys; totdeg = length(polys), - ν = 3, T = Float64, - trans = identity, - sd_admissible = bb -> (true), - envelope = envelopefcn(x -> sqrt(1 + x^2), rand())) - # 1-particle spec - K = length(polys) - spec1p = [ (k, σ) for σ in [1, 2, 3] for k in 1:K] # (1, 2, 3) = (∅, ↑, ↓); - spec1p = sort(spec1p, by = b -> b[1]) # sorting to prevent gensparse being confused - - pooling = PooledSparseProduct(spec1p) - # generate the many-particle spec + # wrap into Lux layers, should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm ? + aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) + pooling_layer = ACEpsi.lux(pooling) + + # define sparse for n-correlations tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) @@ -131,421 +139,486 @@ function BFwf(Nel::Integer, polys; totdeg = length(polys), spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] corr1 = SparseSymmProd(spec; T = Float64) - corr = corr1.dag - - # initial guess for weights - Q, _ = qr(randn(T, length(corr), Nel)) - W = Matrix(Q) - - return BFwf(trans, polys, pooling, corr, W, envelope, spec, - zeros(T, Nel, length(polys)), - zeros(T, Nel, length(polys)), - zeros(T, Nel, length(polys)), - zeros(T, Nel, Nel), - zeros(T, Nel, Nel), - zeros(T, Nel, length(pooling)), - zeros(T, Nel, length(pooling)), - zeros(T, length(pooling)), - zeros(T, length(pooling)), - zeros(Bool, Nel, 3), - zeros(T, Nel, length(corr)), - zeros(T, Nel, 3), - zeros(T, Nel, Nel, length(corr)) ) - + # corr = corr1.dag + # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) + corr_layer = ACEcore.lux(corr1) + + # TODO: confirm what is the envelope and add trainable basis + # A Note here is that the the initial state of the chain must contain st.Σ in some layers... + # to be discuss whether we should put it in X directly since we do not differentiate w.r.t Σ and it is not trainable too + return Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = ReshapeLayer(), bAA = corr_layer, transpose = TransposeLayer(), Linear1 = Dense(length(corr1), Nel), Mask = MaskLayer(Nel), logabsdet = logabsdetLayer()) end -""" -This function returns a nice version of spec. -""" -function displayspec(wf::BFwf) - K = length(wf.polys) - spec1p = [ (k, σ) for σ in [1, 2, 3] for k in 1:K] - spec1p = sort(spec1p, by = b -> b[1]) - _getnicespec = l -> (l[1], num2spin(l[2])) - nicespec = [] - for k = 1:length(wf.spec) - push!(nicespec, _getnicespec.([spec1p[wf.spec[k][j]] for j = 1:length(wf.spec[k])])) - end - return nicespec -end - - - -function evaluate(wf::BFwf, X::AbstractVector, Σ, Pnn=nothing) - nX = length(X) - A = assemble_A(wf, X, Σ) - AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) +# ---------- old implementation + + +# mutable struct BFwf{T, TT, TPOLY, TE} +# embedding::TA +# corr::SparseSymmProdDAG{T} +# W::Matrix{T} +# envelope::TE +# spec::Vector{Vector{Int64}} # corr.spec TODO: this needs to be remove +# # ---------------- Temporaries +# P::Matrix{T} +# ∂P::Matrix{T} +# dP::Matrix{T} +# Φ::Matrix{T} +# ∂Φ::Matrix{T} +# A::Matrix{T} +# ∂A::Matrix{T} +# Ai::Vector{T} +# ∂Ai::Vector{T} +# Si::Matrix{Bool} +# ∂AA::Matrix{T} +# ∂Si::Matrix{T} +# ∇AA::Array{T, 3} +# end + +# (Φ::BFwf)(args...) = evaluate(Φ, args...) + +# function BFwf(Nel::Integer, polys; totdeg = length(polys), +# ν = 3, T = Float64, +# trans = identity, +# sd_admissible = bb -> (true), +# envelope = envelopefcn(x -> sqrt(1 + x^2), rand())) +# # 1-particle spec +# K = length(polys) +# spec1p = [ (k, σ) for σ in [1, 2, 3] for k in 1:K] # (1, 2, 3) = (∅, ↑, ↓); +# spec1p = sort(spec1p, by = b -> b[1]) # sorting to prevent gensparse being confused - # the only basis to be purified are those with same spin - # scan through all corr basis, if they comes from same spin, remove self interation by using basis - # from same spin - # first we have to construct coefficent for basis coming from same spin, that is in other words the coefficent - # matrix of the original polynomial basis, this will be pass from the argument Pnn - # === purification goes here === # +# pooling = PooledSparseProduct(spec1p) +# # generate the many-particle spec +# tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] +# default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) - # === # - Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) # nX x nX - Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin - release!(AA) - - env = wf.envelope(X) - return 2 * logabsdet(Φ)[1] + 2 * log(abs(env)) -end +# specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, +# minvv = fill(0, ν), +# maxvv = fill(length(spec1p), ν), +# ordered = true) + + +# spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] + +# # further restrict +# spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] + +# corr1 = SparseSymmProd(spec; T = Float64) +# corr = corr1.dag + +# # initial guess for weights +# Q, _ = qr(randn(T, length(corr), Nel)) +# W = Matrix(Q) + +# return BFwf(trans, polys, pooling, corr, W, envelope, spec, +# zeros(T, Nel, length(polys)), +# zeros(T, Nel, length(polys)), +# zeros(T, Nel, length(polys)), +# zeros(T, Nel, Nel), +# zeros(T, Nel, Nel), +# zeros(T, Nel, length(pooling)), +# zeros(T, Nel, length(pooling)), +# zeros(T, length(pooling)), +# zeros(T, length(pooling)), +# zeros(Bool, Nel, 3), +# zeros(T, Nel, length(corr)), +# zeros(T, Nel, 3), +# zeros(T, Nel, Nel, length(corr)) ) + +# end + + +# """ +# This function returns a nice version of spec. +# """ +# function displayspec(wf::BFwf) +# K = length(wf.polys) +# spec1p = [ (k, σ) for σ in [1, 2, 3] for k in 1:K] +# spec1p = sort(spec1p, by = b -> b[1]) +# _getnicespec = l -> (l[1], num2spin(l[2])) +# nicespec = [] +# for k = 1:length(wf.spec) +# push!(nicespec, _getnicespec.([spec1p[wf.spec[k][j]] for j = 1:length(wf.spec[k])])) +# end +# return nicespec +# end + + + +# function evaluate(wf::BFwf, X::AbstractVector, Σ, Pnn=nothing) +# nX = length(X) +# A = assemble_A(wf, X, Σ) +# AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) + +# # the only basis to be purified are those with same spin +# # scan through all corr basis, if they comes from same spin, remove self interation by using basis +# # from same spin +# # first we have to construct coefficent for basis coming from same spin, that is in other words the coefficent +# # matrix of the original polynomial basis, this will be pass from the argument Pnn +# # === purification goes here === # + +# # === # +# Φ = wf.Φ +# mul!(Φ, parent(AA), wf.W) # nX x nX +# Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin +# release!(AA) +# env = wf.envelope(X) +# return 2 * logabsdet(Φ)[1] + 2 * log(abs(env)) +# end -function gradp_evaluate(wf::BFwf, X::AbstractVector, Σ) - nX = length(X) + +# function gradp_evaluate(wf::BFwf, X::AbstractVector, Σ) +# nX = length(X) - A = assemble_A(wf, X, Σ) - AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) - Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) - Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin +# A = assemble_A(wf, X, Σ) +# AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) +# Φ = wf.Φ +# mul!(Φ, parent(AA), wf.W) +# Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin - # ψ = log | det( Φ ) | - # ∂Φ = ∂ψ/∂Φ = Φ⁻ᵀ - ∂Φ = transpose(pinv(Φ)) +# # ψ = log | det( Φ ) | +# # ∂Φ = ∂ψ/∂Φ = Φ⁻ᵀ +# ∂Φ = transpose(pinv(Φ)) - # ∂W = ∂ψ / ∂W = ∂Φ * ∂_W( AA * W ) = ∂Φ * AA - # ∂Wij = ∑_ab ∂Φab * ∂_Wij( ∑_k AA_ak W_kb ) - # = ∑_ab ∂Φab * ∑_k δ_ik δ_bj AA_ak - # = ∑_a ∂Φaj AA_ai = ∂Φaj' * AA_ai - ∇p = transpose(parent(AA)) * ∂Φ +# # ∂W = ∂ψ / ∂W = ∂Φ * ∂_W( AA * W ) = ∂Φ * AA +# # ∂Wij = ∑_ab ∂Φab * ∂_Wij( ∑_k AA_ak W_kb ) +# # = ∑_ab ∂Φab * ∑_k δ_ik δ_bj AA_ak +# # = ∑_a ∂Φaj AA_ai = ∂Φaj' * AA_ai +# ∇p = transpose(parent(AA)) * ∂Φ - release!(AA) - ∇p = ∇p * 2 +# release!(AA) +# ∇p = ∇p * 2 - # ------ gradient of env (w.r.t. ξ) ----- - # ∂ = ∂/∂ξ - # r = ||x|| - # ∂(2 * logabs(env)) = ∂(2 * log(exp(-ξf(r)))) = ∂(-2ξf(r)) = -f(r) - ∇logabsenv = - 2 * wf.envelope.f(norm(X)) +# # ------ gradient of env (w.r.t. ξ) ----- +# # ∂ = ∂/∂ξ +# # r = ||x|| +# # ∂(2 * logabs(env)) = ∂(2 * log(exp(-ξf(r)))) = ∂(-2ξf(r)) = -f(r) +# ∇logabsenv = - 2 * wf.envelope.f(norm(X)) - return (∇p = ∇p, ∇logabsenv = [∇logabsenv]) # TODO: return a named tuple (W = gradp, D = gradient w.r.t parameter of env) -end +# return (∇p = ∇p, ∇logabsenv = [∇logabsenv]) # TODO: return a named tuple (W = gradp, D = gradient w.r.t parameter of env) +# end -# ----------------------- gradient +# # ----------------------- gradient -struct ZeroNoEffect end -Base.size(::ZeroNoEffect, ::Integer) = Inf -Base.setindex!(A::ZeroNoEffect, args...) = nothing -Base.getindex(A::ZeroNoEffect, args...) = Bool(0) +# struct ZeroNoEffect end +# Base.size(::ZeroNoEffect, ::Integer) = Inf +# Base.setindex!(A::ZeroNoEffect, args...) = nothing +# Base.getindex(A::ZeroNoEffect, args...) = Bool(0) -function gradient(wf::BFwf, X, Σ) - nX = length(X) +# function gradient(wf::BFwf, X, Σ) +# nX = length(X) - # ------ forward pass ----- +# # ------ forward pass ----- - # position embedding (forward-mode) - # here we evaluate and differentiate at the same time, which is cheap - P = wf.P - dP = wf.dP - Xt = wf.trans.(X) - Polynomials4ML.evaluate_ed!(P, dP, wf.polys, Xt) - ∂Xt = ForwardDiff.derivative.(Ref(x -> wf.trans(x)), X) - @inbounds for k = 1:size(dP, 2) - @simd ivdep for i = 1:nX - dP[i, k] *= ∂Xt[i] - end - end +# # position embedding (forward-mode) +# # here we evaluate and differentiate at the same time, which is cheap +# P = wf.P +# dP = wf.dP +# Xt = wf.trans.(X) +# Polynomials4ML.evaluate_ed!(P, dP, wf.polys, Xt) +# ∂Xt = ForwardDiff.derivative.(Ref(x -> wf.trans(x)), X) +# @inbounds for k = 1:size(dP, 2) +# @simd ivdep for i = 1:nX +# dP[i, k] *= ∂Xt[i] +# end +# end - # no gradients here - need to somehow "constify" this vector - # could use other packages for inspiration ... +# # no gradients here - need to somehow "constify" this vector +# # could use other packages for inspiration ... - # pooling : need an elegant way to shift this loop into a kernel! - # e.g. by passing output indices to the pooling function. +# # pooling : need an elegant way to shift this loop into a kernel! +# # e.g. by passing output indices to the pooling function. - A = wf.A # zeros(nX, length(wf.pooling)) - Ai = wf.Ai # zeros(length(wf.pooling)) - Si = wf.Si # zeros(Bool, nX, 3) +# A = wf.A # zeros(nX, length(wf.pooling)) +# Ai = wf.Ai # zeros(length(wf.pooling)) +# Si = wf.Si # zeros(Bool, nX, 3) - for i = 1:nX - onehot!(Si, i, Σ) - ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) - A[i, :] .= Ai - end +# for i = 1:nX +# onehot!(Si, i, Σ) +# ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) +# A[i, :] .= Ai +# end - # n-correlations - AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) +# # n-correlations +# AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) - # generalized orbitals - Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) +# # generalized orbitals +# Φ = wf.Φ +# mul!(Φ, parent(AA), wf.W) - # the resulting matrix should contains two block each comes from each spin - Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] +# # the resulting matrix should contains two block each comes from each spin +# Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] - # envelope - env = wf.envelope(X) - - # and finally the wave function - # ψ = logabsdet(Φ)[1] + log(abs(env)) - - # ------ backward pass ------ - #∂Φ = ∂ψ / ∂Φ = Φ⁻ᵀ - ∂Φ = transpose(pinv(Φ)) - - # ∂AA = ∂ψ/∂AA = ∂ψ/∂Φ * ∂Φ/∂AA = ∂Φ * wf.W' - ∂AA = wf.∂AA - mul!(∂AA, ∂Φ, transpose(wf.W)) - - # ∂A = ∂ψ/∂A = ∂ψ/∂AA * ∂AA/∂A -> use custom pullback - ∂A = wf.∂A # zeros(size(A)) - ACEcore.pullback_arg!(∂A, ∂AA, wf.corr, parent(AA)) - release!(AA) - - # ∂P = ∂ψ/∂P = ∂ψ/∂A * ∂A/∂P -> use custom pullback - # but need to do some work here since multiple - # pullbacks can be combined here into a single one maybe? - ∂P = wf.∂P # zeros(size(P)) - fill!(∂P, 0) - ∂Si = wf.∂Si # zeros(size(Si)) # should use ZeroNoEffect here ?!??! - Si_ = zeros(nX, 3) - for i = 1:nX - onehot!(Si_, i, Σ) - # note this line ADDS the pullback into ∂P, not overwrite the content!! - ∂Ai = @view ∂A[i, :] - ACEcore._pullback_evalpool!((∂P, ∂Si), ∂Ai, wf.pooling, (P, Si_)) - end - - # ∂X = ∂ψ/∂X = ∂ψ/∂P * ∂P/∂X - # here we can now finally employ the dP=∂P/∂X that we already know. - # ∂ψ/∂Xi = ∑_k ∂ψ/∂Pik * ∂Pik/∂Xi - # = ∑_k ∂P[i, k] * dP[i, k] - g = zeros(nX) - @inbounds for k = 1:length(wf.polys) - @simd ivdep for i = 1:nX - g[i] += ∂P[i, k] * dP[i, k] - end - end - # g = sum(∂P .* dP, dims = 2)[:] - - # envelope - ∇env = ForwardDiff.gradient(wf.envelope, X) - g += ∇env / env - g = g * 2 - return g -end - - -# ------------------ Laplacian implementation - -function laplacian(wf::BFwf, X, Σ) - - A, ∇A, ΔA = _assemble_A_∇A_ΔA(wf, X, Σ) - AA, ∇AA, ΔAA = _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) - - Δψ = _laplacian_inner(AA, ∇AA, ΔAA, wf, Σ) - - # envelope - env = wf.envelope(X) - ∇env = ForwardDiff.gradient(wf.envelope, X) - Δenv = tr(ForwardDiff.hessian(wf.envelope, X)) - - # Δ(ln(env)) = Δenv / env - ∇env ⋅ ∇env / env ^ 2 - Δψ += Δenv / env - dot(∇env, ∇env) / env^2 - Δψ = Δψ * 2 - return Δψ -end - -function _assemble_A_∇A_ΔA(wf, X, Σ) - TX = eltype(X) - lenA = length(wf.pooling) - nX = length(X) - A = zeros(TX, nX, lenA) - ∇A = zeros(TX, nX, nX, lenA) - ΔA = zeros(TX, nX, lenA) - spec_A = wf.pooling.spec - - Xt = wf.trans.(X) - P, dP, ddP = Polynomials4ML.evaluate_ed2(wf.polys, Xt) - dtrans = x -> ForwardDiff.derivative(wf.trans, x) - ddtrans = x -> ForwardDiff.derivative(dtrans, x) - ∂Xt = dtrans.(X) - ∂∂Xt = ddtrans.(X) - @inbounds for k = 1:size(dP, 2) - @simd ivdep for i = 1:nX - dP[i, k], ddP[i, k] = ∂Xt[i] * dP[i, k], ∂∂Xt[i] * dP[i, k] + ∂Xt[i]^2 * ddP[i, k] - end - end - - Si_ = zeros(nX, 3) - Ai = zeros(length(wf.pooling)) - @inbounds for i = 1:nX # loop over orbital bases (which i becomes ∅) - fill!(Si_, 0) - onehot!(Si_, i, Σ) - ACEcore.evalpool!(Ai, wf.pooling, (P, Si_)) - @. A[i, :] .= Ai - for (iA, (k, σ)) in enumerate(spec_A) - for a = 1:nX - ∇A[a, i, iA] = dP[a, k] * Si_[a, σ] - ΔA[i, iA] += ddP[a, k] * Si_[a, σ] - end - end - end - return A, ∇A, ΔA -end - -function _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) - nX = size(A, 1) - AA = zeros(nX, length(wf.corr)) - ∇AA = wf.∇AA - ΔAA = zeros(nX, length(wf.corr)) - - @inbounds for iAA = 1:wf.corr.num1 - @. AA[:, iAA] .= A[:, iAA] - @. ∇AA[:, :, iAA] .= ∇A[:, :, iAA] - @. ΔAA[:, iAA] .= ΔA[:, iAA] - end - - lenAA = length(wf.corr) - @inbounds for iAA = wf.corr.num1+1:lenAA - k1, k2 = wf.corr.nodes[iAA] - for i = 1:nX - AA_k1 = AA[i, k1]; AA_k2 = AA[i, k2] - AA[i, iAA] = AA_k1 * AA_k2 - L = ΔAA[i, k1] * AA_k2 - L = muladd(ΔAA[i, k2], AA_k1, L) - @simd ivdep for a = 1:nX - ∇AA_k1 = ∇AA[a, i, k1]; ∇AA_k2 = ∇AA[a, i, k2] - L = muladd(2 * ∇AA_k1, ∇AA_k2, L) - g = ∇AA_k1 * AA_k2 - ∇AA[a, i, iAA] = muladd(∇AA_k2, AA_k1, g) - end - ΔAA[i, iAA] = L - end - end - return AA, ∇AA, ΔAA -end - - -function _laplacian_inner(AA, ∇AA, ΔAA, wf, Σ) - - # Δψ = Φ⁻ᵀ : ΔΦ - ∑ᵢ (Φ⁻ᵀ * Φᵢ)ᵀ : (Φ⁻ᵀ * Φᵢ) - # where Φᵢ = ∂_{xi} Φ - - nX = size(AA, 1) +# # envelope +# env = wf.envelope(X) + +# # and finally the wave function +# # ψ = logabsdet(Φ)[1] + log(abs(env)) + +# # ------ backward pass ------ +# #∂Φ = ∂ψ / ∂Φ = Φ⁻ᵀ +# ∂Φ = transpose(pinv(Φ)) + +# # ∂AA = ∂ψ/∂AA = ∂ψ/∂Φ * ∂Φ/∂AA = ∂Φ * wf.W' +# ∂AA = wf.∂AA +# mul!(∂AA, ∂Φ, transpose(wf.W)) + +# # ∂A = ∂ψ/∂A = ∂ψ/∂AA * ∂AA/∂A -> use custom pullback +# ∂A = wf.∂A # zeros(size(A)) +# ACEcore.pullback_arg!(∂A, ∂AA, wf.corr, parent(AA)) +# release!(AA) + +# # ∂P = ∂ψ/∂P = ∂ψ/∂A * ∂A/∂P -> use custom pullback +# # but need to do some work here since multiple +# # pullbacks can be combined here into a single one maybe? +# ∂P = wf.∂P # zeros(size(P)) +# fill!(∂P, 0) +# ∂Si = wf.∂Si # zeros(size(Si)) # should use ZeroNoEffect here ?!??! +# Si_ = zeros(nX, 3) +# for i = 1:nX +# onehot!(Si_, i, Σ) +# # note this line ADDS the pullback into ∂P, not overwrite the content!! +# ∂Ai = @view ∂A[i, :] +# ACEcore._pullback_evalpool!((∂P, ∂Si), ∂Ai, wf.pooling, (P, Si_)) +# end + +# # ∂X = ∂ψ/∂X = ∂ψ/∂P * ∂P/∂X +# # here we can now finally employ the dP=∂P/∂X that we already know. +# # ∂ψ/∂Xi = ∑_k ∂ψ/∂Pik * ∂Pik/∂Xi +# # = ∑_k ∂P[i, k] * dP[i, k] +# g = zeros(nX) +# @inbounds for k = 1:length(wf.polys) +# @simd ivdep for i = 1:nX +# g[i] += ∂P[i, k] * dP[i, k] +# end +# end +# # g = sum(∂P .* dP, dims = 2)[:] + +# # envelope +# ∇env = ForwardDiff.gradient(wf.envelope, X) +# g += ∇env / env +# g = g * 2 +# return g +# end + + +# # ------------------ Laplacian implementation + +# function laplacian(wf::BFwf, X, Σ) + +# A, ∇A, ΔA = _assemble_A_∇A_ΔA(wf, X, Σ) +# AA, ∇AA, ΔAA = _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) + +# Δψ = _laplacian_inner(AA, ∇AA, ΔAA, wf, Σ) + +# # envelope +# env = wf.envelope(X) +# ∇env = ForwardDiff.gradient(wf.envelope, X) +# Δenv = tr(ForwardDiff.hessian(wf.envelope, X)) + +# # Δ(ln(env)) = Δenv / env - ∇env ⋅ ∇env / env ^ 2 +# Δψ += Δenv / env - dot(∇env, ∇env) / env^2 +# Δψ = Δψ * 2 +# return Δψ +# end + +# function _assemble_A_∇A_ΔA(wf, X, Σ) +# TX = eltype(X) +# lenA = length(wf.pooling) +# nX = length(X) +# A = zeros(TX, nX, lenA) +# ∇A = zeros(TX, nX, nX, lenA) +# ΔA = zeros(TX, nX, lenA) +# spec_A = wf.pooling.spec + +# Xt = wf.trans.(X) +# P, dP, ddP = Polynomials4ML.evaluate_ed2(wf.polys, Xt) +# dtrans = x -> ForwardDiff.derivative(wf.trans, x) +# ddtrans = x -> ForwardDiff.derivative(dtrans, x) +# ∂Xt = dtrans.(X) +# ∂∂Xt = ddtrans.(X) +# @inbounds for k = 1:size(dP, 2) +# @simd ivdep for i = 1:nX +# dP[i, k], ddP[i, k] = ∂Xt[i] * dP[i, k], ∂∂Xt[i] * dP[i, k] + ∂Xt[i]^2 * ddP[i, k] +# end +# end + +# Si_ = zeros(nX, 3) +# Ai = zeros(length(wf.pooling)) +# @inbounds for i = 1:nX # loop over orbital bases (which i becomes ∅) +# fill!(Si_, 0) +# onehot!(Si_, i, Σ) +# ACEcore.evalpool!(Ai, wf.pooling, (P, Si_)) +# @. A[i, :] .= Ai +# for (iA, (k, σ)) in enumerate(spec_A) +# for a = 1:nX +# ∇A[a, i, iA] = dP[a, k] * Si_[a, σ] +# ΔA[i, iA] += ddP[a, k] * Si_[a, σ] +# end +# end +# end +# return A, ∇A, ΔA +# end + +# function _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) +# nX = size(A, 1) +# AA = zeros(nX, length(wf.corr)) +# ∇AA = wf.∇AA +# ΔAA = zeros(nX, length(wf.corr)) + +# @inbounds for iAA = 1:wf.corr.num1 +# @. AA[:, iAA] .= A[:, iAA] +# @. ∇AA[:, :, iAA] .= ∇A[:, :, iAA] +# @. ΔAA[:, iAA] .= ΔA[:, iAA] +# end + +# lenAA = length(wf.corr) +# @inbounds for iAA = wf.corr.num1+1:lenAA +# k1, k2 = wf.corr.nodes[iAA] +# for i = 1:nX +# AA_k1 = AA[i, k1]; AA_k2 = AA[i, k2] +# AA[i, iAA] = AA_k1 * AA_k2 +# L = ΔAA[i, k1] * AA_k2 +# L = muladd(ΔAA[i, k2], AA_k1, L) +# @simd ivdep for a = 1:nX +# ∇AA_k1 = ∇AA[a, i, k1]; ∇AA_k2 = ∇AA[a, i, k2] +# L = muladd(2 * ∇AA_k1, ∇AA_k2, L) +# g = ∇AA_k1 * AA_k2 +# ∇AA[a, i, iAA] = muladd(∇AA_k2, AA_k1, g) +# end +# ΔAA[i, iAA] = L +# end +# end +# return AA, ∇AA, ΔAA +# end + + +# function _laplacian_inner(AA, ∇AA, ΔAA, wf, Σ) + +# # Δψ = Φ⁻ᵀ : ΔΦ - ∑ᵢ (Φ⁻ᵀ * Φᵢ)ᵀ : (Φ⁻ᵀ * Φᵢ) +# # where Φᵢ = ∂_{xi} Φ + +# nX = size(AA, 1) - # the wf, and the first layer of derivatives - Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) - Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin - Φ⁻ᵀ = transpose(pinv(Φ)) +# # the wf, and the first layer of derivatives +# Φ = wf.Φ +# mul!(Φ, parent(AA), wf.W) +# Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin +# Φ⁻ᵀ = transpose(pinv(Φ)) - # first contribution to the laplacian - ΔΦ = ΔAA * wf.W - ΔΦ = ΔΦ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin +# # first contribution to the laplacian +# ΔΦ = ΔAA * wf.W +# ΔΦ = ΔΦ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin - Δψ = dot(Φ⁻ᵀ, ΔΦ) +# Δψ = dot(Φ⁻ᵀ, ΔΦ) - # the gradient contribution - # TODO: we can rework this into a single BLAS3 call - # which will also give us a single back-propagation - # ∇Φi = zeros(nX, nX) - ∇Φ_all = reshape(reshape(∇AA, nX*nX, :) * wf.W, nX, nX, nX) - Φ⁻¹∇Φi = zeros(nX, nX) - for i = 1:nX - ∇Φi = ∇Φ_all[i, :, :] .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] - mul!(Φ⁻¹∇Φi, transpose(Φ⁻ᵀ), ∇Φi) - Δψ -= dot(transpose(Φ⁻¹∇Φi), Φ⁻¹∇Φi) - end +# # the gradient contribution +# # TODO: we can rework this into a single BLAS3 call +# # which will also give us a single back-propagation +# # ∇Φi = zeros(nX, nX) +# ∇Φ_all = reshape(reshape(∇AA, nX*nX, :) * wf.W, nX, nX, nX) +# Φ⁻¹∇Φi = zeros(nX, nX) +# for i = 1:nX +# ∇Φi = ∇Φ_all[i, :, :] .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] +# mul!(Φ⁻¹∇Φi, transpose(Φ⁻ᵀ), ∇Φi) +# Δψ -= dot(transpose(Φ⁻¹∇Φi), Φ⁻¹∇Φi) +# end - return Δψ -end +# return Δψ +# end -# ------------------ gradp of Laplacian +# # ------------------ gradp of Laplacian -function gradp_laplacian(wf::BFwf, X, Σ) +# function gradp_laplacian(wf::BFwf, X, Σ) - # ---- gradp of Laplacian of Ψ ---- +# # ---- gradp of Laplacian of Ψ ---- - nX = length(X) +# nX = length(X) - A, ∇A, ΔA = _assemble_A_∇A_ΔA(wf, X, Σ) - AA, ∇AA, ΔAA = _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) +# A, ∇A, ΔA = _assemble_A_∇A_ΔA(wf, X, Σ) +# AA, ∇AA, ΔAA = _assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) - # the wf, and the first layer of derivatives - Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) - Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin +# # the wf, and the first layer of derivatives +# Φ = wf.Φ +# mul!(Φ, parent(AA), wf.W) +# Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin - Φ⁻¹ = pinv(Φ) - Φ⁻ᵀ = transpose(Φ⁻¹) +# Φ⁻¹ = pinv(Φ) +# Φ⁻ᵀ = transpose(Φ⁻¹) - # first contribution to the laplacian - ΔΦ = ΔAA * wf.W - ΔΦ = ΔΦ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin +# # first contribution to the laplacian +# ΔΦ = ΔAA * wf.W +# ΔΦ = ΔΦ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin - # Δψ += dot(Φ⁻ᵀ, ΔΦ) ... this leads to the next two terms - ∂ΔΦ = Φ⁻ᵀ - ∇Δψ = transpose(ΔAA) * ∂ΔΦ +# # Δψ += dot(Φ⁻ᵀ, ΔΦ) ... this leads to the next two terms +# ∂ΔΦ = Φ⁻ᵀ +# ∇Δψ = transpose(ΔAA) * ∂ΔΦ - ∂Φ = - Φ⁻ᵀ * transpose(ΔΦ) * Φ⁻ᵀ - # ∇Δψ += transpose(AA) * ∂Φ +# ∂Φ = - Φ⁻ᵀ * transpose(ΔΦ) * Φ⁻ᵀ +# # ∇Δψ += transpose(AA) * ∂Φ - # the gradient contribution - # TODO: we can rework this into a single BLAS3 call - # which will also give us a single back-propagation - # ∇Φi = zeros(nX, nX) - # Φ⁻¹∇Φi = zeros(nX, nX) - ∇Φ_all = reshape(reshape(∇AA, nX*nX, :) * wf.W, nX, nX, nX) - ∂∇Φ_all = zeros(nX, nX, nX) +# # the gradient contribution +# # TODO: we can rework this into a single BLAS3 call +# # which will also give us a single back-propagation +# # ∇Φi = zeros(nX, nX) +# # Φ⁻¹∇Φi = zeros(nX, nX) +# ∇Φ_all = reshape(reshape(∇AA, nX*nX, :) * wf.W, nX, nX, nX) +# ∂∇Φ_all = zeros(nX, nX, nX) - for i = 1:nX - ∇Φi = ∇Φ_all[i, :, :] .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] - ∇Φiᵀ = transpose(∇Φi) - # Δψ += - dot( [Φ⁻¹∇Φi]ᵀ, Φ⁻¹∇Φi ) +# for i = 1:nX +# ∇Φi = ∇Φ_all[i, :, :] .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] +# ∇Φiᵀ = transpose(∇Φi) +# # Δψ += - dot( [Φ⁻¹∇Φi]ᵀ, Φ⁻¹∇Φi ) - # ∂∇Φi = - 2 * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ - ∂∇Φ_all[i, :, :] = - 2 * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ - # ∇Δψ += transpose(∇AA[i, :, :]) * ∂∇Φi +# # ∂∇Φi = - 2 * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ +# ∂∇Φ_all[i, :, :] = - 2 * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ +# # ∇Δψ += transpose(∇AA[i, :, :]) * ∂∇Φi - ∂Φ += 2 * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ - # ∇Δψ += transpose(AA) * ∂Φ - end +# ∂Φ += 2 * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ * ∇Φiᵀ * Φ⁻ᵀ +# # ∇Δψ += transpose(AA) * ∂Φ +# end - ∇Δψ += transpose(AA) * ∂Φ +# ∇Δψ += transpose(AA) * ∂Φ - ∇Δψ += reshape( transpose(reshape(∇AA, nX*nX, :)) * reshape(∂∇Φ_all, nX*nX, nX), - size(∇Δψ) ) +# ∇Δψ += reshape( transpose(reshape(∇AA, nX*nX, :)) * reshape(∂∇Φ_all, nX*nX, nX), +# size(∇Δψ) ) - # ---- gradp of Laplacian of env ---- - # ∂ = ∂/∂ξ - # r = ||x|| - # ∂(2 * Δ(logabs(env))) = ∂(2 * Δ(-ξf(r)) = -2 * Δf(r) - # Δf(r) = (n-1)/r f'(r) + f''(r) - r = norm(X) - f(r) = wf.envelope.f(r) - df(r) = ForwardDiff.derivative(f, r) - ddf(r) = ForwardDiff.derivative(r -> df(r), r) - Δf = ddf(r) + (length(X)-1) * df(r)/r +# # ---- gradp of Laplacian of env ---- +# # ∂ = ∂/∂ξ +# # r = ||x|| +# # ∂(2 * Δ(logabs(env))) = ∂(2 * Δ(-ξf(r)) = -2 * Δf(r) +# # Δf(r) = (n-1)/r f'(r) + f''(r) +# r = norm(X) +# f(r) = wf.envelope.f(r) +# df(r) = ForwardDiff.derivative(f, r) +# ddf(r) = ForwardDiff.derivative(r -> df(r), r) +# Δf = ddf(r) + (length(X)-1) * df(r)/r - return (∇Δψ = 2 * ∇Δψ, ∇Δlogabsenv = [2 * -Δf]) -end +# return (∇Δψ = 2 * ∇Δψ, ∇Δlogabsenv = [2 * -Δf]) +# end -# ----------------- BFwf parameter wraging +# # ----------------- BFwf parameter wraging -function get_params(U::BFwf) - return (U.W, U.envelope.ξ) -end +# function get_params(U::BFwf) +# return (U.W, U.envelope.ξ) +# end -function set_params!(U::BFwf, para) - U.W = para[1] - set_params!(U.envelope, para[2]) - return U -end +# function set_params!(U::BFwf, para) +# U.W = para[1] +# set_params!(U.envelope, para[2]) +# return U +# end diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 2301a28..468c042 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -1,9 +1,10 @@ - using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate using ACEpsi: BackflowPooling using Polynomials4ML.Testing: print_tf +using LuxCore +using Random # test configs Rnldegree = 4 @@ -39,13 +40,13 @@ end println() @info("Test evaluate ProductBasis") -ϕnlm = evaluate(prodbasis, X, Σ) +ϕnlm = prodbasis(X, Σ) @info("Test evaluate AtomicOrbitalsBasis") -bϕnlm = evaluate(aobasis, X, Σ) +bϕnlm = aobasis(X, Σ) @info("Test BackflowPooling") -A = ACEpsi.evaluate(pooling, bϕnlm, Σ) +A = pooling(bϕnlm, Σ) println() @@ -90,3 +91,25 @@ println() # +@info("---------- Lux tests ----------") +aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) +pooling_layer = ACEpsi.lux(pooling) + +# set up null states +ps, st = LuxCore.setup(MersenneTwister(1234), aobasis_layer) +st1 = (Σ = Σ, ) + +@info("Checking layers working fine seperately") +print_tf(@test aobasis_layer(X, ps, st1)[1] ≈ bϕnlm) +print_tf(@test pooling_layer(bϕnlm, ps, st1)[1] ≈ A) +println() + +@info("Checking Chain towards A basis") +using Lux: Chain +tryChain = Chain(; aobasis = aobasis_layer, pooling = pooling_layer) +chain_ps, chain_st = LuxCore.setup(MersenneTwister(1234), tryChain) +chain_st = (aobasis = (Σ = Σ, ), pooling = (Σ = Σ, )) +# try Chain is as expected +print_tf(@test tryChain(X, chain_ps, chain_st)[1] ≈ A) +println() + diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl new file mode 100644 index 0000000..1dba9cc --- /dev/null +++ b/test/test_bflow_lux.jl @@ -0,0 +1,37 @@ +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate +using ACEpsi: BackflowPooling +using ACEpsi: BFwf_lux +using Polynomials4ML.Testing: print_tf +using LuxCore +using Random + +function setupBFState(rng, BFwf_chain, Σ) + chain_ps, chain_st = LuxCore.setup(rng, BFwf_chain) + layers = keys(chain_st) + nΣ = Tuple([(Σ = Σ, ) for _ = 1:length(layers)]) + chain_st = (; zip(layers, nΣ)...) + return chain_ps, chain_st +end + +Rnldegree = 4 +Ylmdegree = 4 +totdegree = 8 +Nel = 5 +X = randn(SVector{3, Float64}, Nel) +Σ = rand(spins(), Nel) + +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] +## + +# Defining AtomicOrbitalsBasis +bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +bYlm = RYlmBasis(Ylmdegree) + + + +BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) + +BFwf_chain(X, ps, st) \ No newline at end of file From eb7ee057b8b51a7af7f7a5f74fed28cc2d936592 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Mon, 15 May 2023 06:30:21 -0700 Subject: [PATCH 13/79] use WrappedLayer --- src/atomicorbitals/atorbbasis.jl | 2 +- src/bflow3d.jl | 29 ++++++----------------------- test/test_atorbbasis.jl | 1 + 3 files changed, 8 insertions(+), 24 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 48101e5..9a95351 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -152,7 +152,7 @@ end # ------------ utils for AtomicOrbitalsBasis ------------ function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) basis.nuclei = copy(collect(nuclei)) - return nothing + return nothing end diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 87b4a82..b3e2332 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -8,7 +8,7 @@ import ForwardDiff using ACEpsi.AtomicOrbitals: make_nlms_spec using LuxCore: AbstractExplicitLayer -using Lux: Dense, Chain +using Lux: Dense, Chain, WrappedFunction # ---------------------------------------- @@ -84,25 +84,6 @@ end (l::MaskLayer)(Φ, ps, st) = Φ .* [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX], st -# TODO: maybe in here we can do fast-reshape with stride array? -struct ReshapeLayer <: AbstractExplicitLayer end - -@inline function (f::ReshapeLayer)(x::AbstractArray{T, N}, ps, st::NamedTuple) where {T, N} - # (nX, 3, length(nuclei), length(spec1p)) -> (nX, 3 * length(nuclei) * length(spec1p)) each row go over the spin channel first, then nuclei, then spec1p - return reshape(x, (size(x, 1), prod(size(x)[2:end]))), st -end - -struct TransposeLayer <: AbstractExplicitLayer end - -@inline function (f::TransposeLayer)(x::AbstractArray{T, N}, ps, st::NamedTuple) where {T, N} - return transpose(x), st -end - -struct logabsdetLayer <: AbstractExplicitLayer end - -@inline function (f::logabsdetLayer)(Φ::AbstractArray{T, N}, ps, st::NamedTuple) where {T, N} - return 2 * logabsdet(Φ)[1], st -end function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, ν = 3, T = Float64, @@ -143,10 +124,12 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) corr_layer = ACEcore.lux(corr1) - # TODO: confirm what is the envelope and add trainable basis + # TODO: confirm what is the envelope and add trainable basis after bA # A Note here is that the the initial state of the chain must contain st.Σ in some layers... - # to be discuss whether we should put it in X directly since we do not differentiate w.r.t Σ and it is not trainable too - return Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = ReshapeLayer(), bAA = corr_layer, transpose = TransposeLayer(), Linear1 = Dense(length(corr1), Nel), Mask = MaskLayer(Nel), logabsdet = logabsdetLayer()) + # to be discuss whether we should put it in X directly since it is not trainable and is the same throughout the training + reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) + logabsdet_func = x -> 2 * logabsdet(x)[1] + return Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), bAA = corr_layer, transpose = WrappedFunction(transpose), Linear1 = Dense(length(corr1), Nel), Mask = MaskLayer(Nel), logabsdet = WrappedFunction(logabsdet_func)) end diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 468c042..8083b5f 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -111,5 +111,6 @@ chain_ps, chain_st = LuxCore.setup(MersenneTwister(1234), tryChain) chain_st = (aobasis = (Σ = Σ, ), pooling = (Σ = Σ, )) # try Chain is as expected print_tf(@test tryChain(X, chain_ps, chain_st)[1] ≈ A) + println() From 620c39a2d0f90294a8f4bea2908cf8e3f4c27a69 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 16 May 2023 09:39:45 -0700 Subject: [PATCH 14/79] clean up and move productbasis to P4ML --- src/atomicorbitals/atorbbasis.jl | 43 +++--------- src/bflow3d.jl | 109 +++++++++++++++---------------- test/test_atorbbasis.jl | 2 +- test/test_bflow1.jl | 2 +- test/test_bflow_lux.jl | 3 +- 5 files changed, 66 insertions(+), 93 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 9a95351..a70460a 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -1,7 +1,7 @@ import Polynomials4ML: evaluate using ACEpsi: ↑, ↓, ∅, spins, extspins, Spin, spin2idx, idx2spin -using Polynomials4ML: SparseProduct +using Polynomials4ML: SparseProduct, ProductBasis using LuxCore: AbstractExplicitLayer using Random: AbstractRNG @@ -53,18 +53,6 @@ function make_nlms_spec(bRnl, bYlm; end -# Jerry: This is just a specific case of a general ProductBasis, this should go to Polynomials4ML later with a general implementation -# I will do that after reconfiming this is what we want -mutable struct ProductBasis{NB, TR, TY} - spec1::Vector{NTRNL1} - bRnl::TR - bYlm::TY - # ---- evaluation kernel from Polynomials4ML ---- - sparsebasis::SparseProduct{NB} -end - -(PB::ProductBasis)(args...) = evaluate(PB, args...) - function _invmap(a::AbstractVector) inva = Dict{eltype(a), Int}() for i = 1:length(a) @@ -73,40 +61,25 @@ function _invmap(a::AbstractVector) return inva end +function dropnames(namedtuple::NamedTuple, names::Tuple{Vararg{Symbol}}) + keepnames = Base.diff_names(Base._nt_names(namedtuple), names) + return NamedTuple{keepnames}(namedtuple) +end + function ProductBasis(spec1, bRnl, bYlm) spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) spec_Rnl = bRnl.spec; inv_Rnl = _invmap(spec_Rnl) - spec_Ylm = Polynomials4ML.natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) + spec_Ylm = natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) for (i, b) in enumerate(spec1) - spec1idx[i] = (inv_Rnl[(n=b.n, l=b.l)], inv_Ylm[(l=b.l, m=b.m)]) + spec1idx[i] = (inv_Rnl[dropnames(b,(:m,))], inv_Ylm[(l=b.l, m=b.m)]) end sparsebasis = SparseProduct(spec1idx) return ProductBasis(spec1, bRnl, bYlm, sparsebasis) end -function evaluate(basis::ProductBasis, X::AbstractVector{<: AbstractVector}, Σ) - Nel = length(X) - T = promote_type(eltype(X[1])) - - # create all the shifted configurations - xx = zeros(eltype(T), Nel) - for i = 1:Nel - xx[i] = norm(X[i]) - end - - # evaluate the radial and angular components on all the shifted particles - Rnl = reshape(evaluate(basis.bRnl, xx[:]), (Nel, length(basis.bRnl))) - Ylm = reshape(evaluate(basis.bYlm, X[:]), (Nel, length(basis.bYlm))) - - # evaluate all the atomic orbitals as ϕ_nlm = Rnl * Ylm - ϕnlm = evaluate(basis.sparsebasis, (Rnl, Ylm)) - - return ϕnlm -end - mutable struct AtomicOrbitalsBasis{NB, T} prodbasis::ProductBasis{NB} nuclei::Vector{Nuc{T}} # nuclei (defines the shifted orbitals) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index b3e2332..a1c7bca 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -11,7 +11,61 @@ using LuxCore: AbstractExplicitLayer using Lux: Dense, Chain, WrappedFunction # ---------------------------------------- +struct MaskLayer <: AbstractExplicitLayer + nX::Integer +end + +(l::MaskLayer)(Φ, ps, st) = Φ .* [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX], st + + +function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, + ν = 3, T = Float64, + sd_admissible = bb -> (true), + envelope = x -> x) # enveolpe to be replaced by SJ-factor + + spec1p = make_nlms_spec(bRnl, bYlm; + totaldegree = totdeg) + + # size(X) = (nX, 3); length(Σ) = nX + aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) + pooling = BackflowPooling(aobasis) + + # define sparse for n-correlations + tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] + default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) + specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, + minvv = fill(0, ν), + maxvv = fill(length(spec1p), ν), + ordered = true) + spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] + + # further restrict + spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] + + # define n-correlation + corr1 = SparseSymmProd(spec; T = Float64) + + # ----------- Lux connections --------- + # Should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm for trainable radial basis later? + # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1)) + aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) + # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) + pooling_layer = ACEpsi.lux(pooling) + # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) + corr_layer = ACEcore.lux(corr1) + + # TODO: Add J-factor and add trainable basis later + + reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) + # Questions to discuss: + # 1. it seems that we do not need trans since the bases have already taken care of it? + # 2. Why we need a tranpose here??? Seems that the output from corr_layer is (length(spec), nX)??? + # 3. How do we define n-correlations if we use trainable basis? + return Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), + bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), + Mask = MaskLayer(Nel), logabsdet = WrappedFunction(x -> 2 * logabsdet(x)[1])) +end @@ -78,61 +132,6 @@ using Lux: Dense, Chain, WrappedFunction # return A # end -struct MaskLayer <: AbstractExplicitLayer - nX::Integer -end - -(l::MaskLayer)(Φ, ps, st) = Φ .* [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX], st - - -function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, - ν = 3, T = Float64, - trans = identity, - sd_admissible = bb -> (true), - envelope = x -> x) - - spec1p = make_nlms_spec(bRnl, bYlm; - totaldegree = totdeg) - # size(X) = (nX, 3); length(Σ) = nX - # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1 from totaldegree)) - aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) - - # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1 from totaldegree)) - pooling = BackflowPooling(aobasis) - - # wrap into Lux layers, should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm ? - aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) - pooling_layer = ACEpsi.lux(pooling) - - # define sparse for n-correlations - tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] - default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) - - specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, - minvv = fill(0, ν), - maxvv = fill(length(spec1p), ν), - ordered = true) - - - spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] - - # further restrict - spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] - - corr1 = SparseSymmProd(spec; T = Float64) - # corr = corr1.dag - # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) - corr_layer = ACEcore.lux(corr1) - - # TODO: confirm what is the envelope and add trainable basis after bA - # A Note here is that the the initial state of the chain must contain st.Σ in some layers... - # to be discuss whether we should put it in X directly since it is not trainable and is the same throughout the training - reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) - logabsdet_func = x -> 2 * logabsdet(x)[1] - return Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), bAA = corr_layer, transpose = WrappedFunction(transpose), Linear1 = Dense(length(corr1), Nel), Mask = MaskLayer(Nel), logabsdet = WrappedFunction(logabsdet_func)) -end - - # ---------- old implementation diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 8083b5f..1cafbce 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -2,7 +2,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate using ACEpsi: BackflowPooling -using Polynomials4ML.Testing: print_tf +using ACEbase.Testing: print_tf using LuxCore using Random diff --git a/test/test_bflow1.jl b/test/test_bflow1.jl index f4d0e39..87d5429 100644 --- a/test/test_bflow1.jl +++ b/test/test_bflow1.jl @@ -99,7 +99,7 @@ function fdtest(F, Σ, dF, x::AbstractVector; h0 = 1.0, verbose=true) const ↑, ↓, ∅ = '↑','↓','∅' Nel = 5 polys = legendre_basis(8) -wf = BFwf1(Nel, polys; ν=3) +wf = BFwf1(Nel, polys, trans = x -> atan(x); ν=3) X = 2 * rand(Nel) .- 1 Σ = rand([↑, ↓], Nel) diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 1dba9cc..872c0c9 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -3,7 +3,7 @@ using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate using ACEpsi: BackflowPooling using ACEpsi: BFwf_lux -using Polynomials4ML.Testing: print_tf +using ACEbase.Testing: print_tf using LuxCore using Random @@ -34,4 +34,5 @@ bYlm = RYlmBasis(Ylmdegree) BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) +@info("Test evaluate") BFwf_chain(X, ps, st) \ No newline at end of file From 0323f175b696d4477a8ee5af1d783b1ff7dbc26d Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 16 May 2023 20:14:46 -0700 Subject: [PATCH 15/79] minor fix --- src/atomicorbitals/atorbbasis.jl | 2 +- test/test_atorbbasis.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index a70460a..b9e4c2f 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -116,7 +116,7 @@ function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVecto ϕnlm = zeros(T, (Nnuc, Nel, Nnlm)) for I = 1:Nnuc - ϕnlm[I,:,:] = evaluate(basis.prodbasis, XX[I,:], Σ) + ϕnlm[I,:,:] = evaluate(basis.prodbasis, XX[I,:]) end return ϕnlm diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 1cafbce..3a7dedc 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -40,7 +40,7 @@ end println() @info("Test evaluate ProductBasis") -ϕnlm = prodbasis(X, Σ) +ϕnlm = prodbasis(X) @info("Test evaluate AtomicOrbitalsBasis") bϕnlm = aobasis(X, Σ) From fdc6bda70fbaa77ef9872f02613cb3d845caee2a Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 17 May 2023 21:37:42 +0800 Subject: [PATCH 16/79] add jastrow layer --- src/ACEpsi.jl | 1 + src/atomicorbitals/atorbbasis.jl | 8 +-- src/bflow3d.jl | 14 +++-- src/jastrow.jl | 59 +++++++++++++++++ test/runtests.jl | 3 +- test/test_bflow_lux.jl | 105 ++++++++++++++++++++++++++++--- 6 files changed, 170 insertions(+), 20 deletions(-) create mode 100644 src/jastrow.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index f37be0c..785dfb8 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -9,6 +9,7 @@ include("envelope.jl") # the new 3d backflow code include("atomicorbitals/atomicorbitals.jl") +include("jastrow.jl") include("bflow3d.jl") include("backflowpooling.jl") diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index b9e4c2f..4b123e6 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -40,13 +40,13 @@ function make_nlms_spec(bRnl, bYlm; spec_Rnl = natural_indices(bRnl) spec_Ylm = natural_indices(bYlm) - spec1 = NTRNL1[] + spec1 = [] for (iR, br) in enumerate(spec_Rnl), (iY, by) in enumerate(spec_Ylm) if br.l != by.l continue end if admissible(br, by) - push!(spec1, (n=br.n, l = br.l, m = by.m)) + push!(spec1, (br..., m = by.m)) end end return spec1 @@ -130,10 +130,10 @@ end function get_spec(basis::AtomicOrbitalsBasis) - spec = NTRNLIS[] + spec = [] Nnuc = length(basis.nuclei) - spec = Array{NTRNLIS, 3}(undef, (3, Nnuc, length(basis.prodbasis.spec1))) + spec = Array{3}(undef, (3, Nnuc, length(basis.prodbasis.spec1))) for (k, nlm) in enumerate(basis.prodbasis.spec1) for I = 1:Nnuc diff --git a/src/bflow3d.jl b/src/bflow3d.jl index a1c7bca..979d2b1 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -3,12 +3,12 @@ using ACEcore, Polynomials4ML using Polynomials4ML: OrthPolyBasis1D3T using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! using ACEcore.Utils: gensparse -using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr +using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det import ForwardDiff using ACEpsi.AtomicOrbitals: make_nlms_spec using LuxCore: AbstractExplicitLayer -using Lux: Dense, Chain, WrappedFunction +using Lux: Dense, Chain, WrappedFunction, BranchLayer # ---------------------------------------- struct MaskLayer <: AbstractExplicitLayer @@ -56,19 +56,21 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, corr_layer = ACEcore.lux(corr1) # TODO: Add J-factor and add trainable basis later + js = Jastrow(nuclei) + jastrow_layer = ACEpsi.lux(js) reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) # Questions to discuss: # 1. it seems that we do not need trans since the bases have already taken care of it? # 2. Why we need a tranpose here??? Seems that the output from corr_layer is (length(spec), nX)??? # 3. How do we define n-correlations if we use trainable basis? - return Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), - bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), - Mask = MaskLayer(Nel), logabsdet = WrappedFunction(x -> 2 * logabsdet(x)[1])) + BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), + bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), + Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) + return Chain(; branch = BranchLayer((js = jastrow_layer, bf = BFwf_chain, )), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ) end - # """ # This function returns correct Si for pooling operation. # """ diff --git a/src/jastrow.jl b/src/jastrow.jl new file mode 100644 index 0000000..29b3165 --- /dev/null +++ b/src/jastrow.jl @@ -0,0 +1,59 @@ +# This should probably be re-implemented properly and be made performant. +# the current version is just for testing +using ACEpsi.AtomicOrbitals: Nuc +using LuxCore: AbstractExplicitLayer +using Random: AbstractRNG + +mutable struct Jastrow{T} + nuclei::Vector{Nuc{T}} # nuclei +end + +(f::Jastrow)(args...) = evaluate(f, args...) + +## F_2(x) = -1/2\sum_{l=1}^L \sum_{i=1}^N Z_l|yi,l|+1/4\sum_{1\leq i (Σ = Σ, ))...) + else + rp_st = (; rp_st..., keys(nt)[i] => replace_namedtuples(nt[i], (;))) + end + end + return rp_st + end +end + +function setupBFState(rng, bf, Σ) + ps, st = LuxCore.setup(rng, bf) + rp_st = replace_namedtuples(st, (;)) + return ps, rp_st end Rnldegree = 4 @@ -29,10 +43,83 @@ nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) bYlm = RYlmBasis(Ylmdegree) - - BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) @info("Test evaluate") -BFwf_chain(X, ps, st) \ No newline at end of file +A1 = BFwf_chain(X, ps, st) + +# original implementation +function BFwf_lux2(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, + ν = 3, T = Float64, + sd_admissible = bb -> (true), + envelope = x -> x) # enveolpe to be replaced by SJ-factor + + spec1p = make_nlms_spec(bRnl, bYlm; + totaldegree = totdeg) + + # size(X) = (nX, 3); length(Σ) = nX + aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) + pooling = BackflowPooling(aobasis) + + # define sparse for n-correlations + tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] + default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) + + specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, + minvv = fill(0, ν), + maxvv = fill(length(spec1p), ν), + ordered = true) + spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] + + # further restrict + spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] + + # define n-correlation + corr1 = SparseSymmProd(spec; T = Float64) + + # ----------- Lux connections --------- + # Should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm for trainable radial basis later? + # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1)) + aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) + # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) + pooling_layer = ACEpsi.lux(pooling) + # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) + corr_layer = ACEcore.lux(corr1) + + # TODO: Add J-factor and add trainable basis later + js = Jastrow(nuclei) + jastrow_layer = ACEpsi.lux(js) + + reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) + # Questions to discuss: + # 1. it seems that we do not need trans since the bases have already taken care of it? + # 2. Why we need a tranpose here??? Seems that the output from corr_layer is (length(spec), nX)??? + # 3. How do we define n-correlations if we use trainable basis? + BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), + bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), + Mask = ACEpsi.MaskLayer(Nel), logabsdet = WrappedFunction(x -> det(x))) + return Chain(; branch = BranchLayer((js = jastrow_layer, bf = BFwf_chain, ))) + end + +using ACEcore, Polynomials4ML +using Polynomials4ML: OrthPolyBasis1D3T +using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! +using ACEcore.Utils: gensparse +using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det +import ForwardDiff +using ACEpsi.AtomicOrbitals: make_nlms_spec + +using LuxCore: AbstractExplicitLayer +using Lux: Dense, Chain, WrappedFunction, BranchLayer +using ACEpsi: Jastrow + +js = Jastrow(nuclei) +jatrow_layer = ACEpsi.lux(js) +branchlayer = BFwf_lux2(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +ps2, st2 = setupBFState(MersenneTwister(1234), branchlayer, Σ) + +A2 = branchlayer(X, ps2, st2) +@test(2 * logabsdet(A2[1][1] *A2[1][2])[1] == A1[1]) + + From bd9474e86b16c63a05c66c8d05d7704fe4e95c29 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Wed, 17 May 2023 09:13:54 -0700 Subject: [PATCH 17/79] clean up --- src/ACEpsi.jl | 2 + src/atomicorbitals/atorbbasis.jl | 4 +- src/lux_utils.jl | 22 +++++ test/test_bflow_lux.jl | 142 +++++++++++++------------------ 4 files changed, 86 insertions(+), 84 deletions(-) create mode 100644 src/lux_utils.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index 785dfb8..85dbd4b 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -14,4 +14,6 @@ include("bflow3d.jl") include("backflowpooling.jl") +# lux utils for bflow +include("lux_utils.jl") end diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 4b123e6..483da91 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -1,7 +1,7 @@ import Polynomials4ML: evaluate using ACEpsi: ↑, ↓, ∅, spins, extspins, Spin, spin2idx, idx2spin -using Polynomials4ML: SparseProduct, ProductBasis +using Polynomials4ML: SparseProduct, ProductBasis, _make_reqfields using LuxCore: AbstractExplicitLayer using Random: AbstractRNG @@ -76,7 +76,7 @@ function ProductBasis(spec1, bRnl, bYlm) spec1idx[i] = (inv_Rnl[dropnames(b,(:m,))], inv_Ylm[(l=b.l, m=b.m)]) end sparsebasis = SparseProduct(spec1idx) - return ProductBasis(spec1, bRnl, bYlm, sparsebasis) + return ProductBasis(spec1, bRnl, bYlm, sparsebasis, _make_reqfields()...) end diff --git a/src/lux_utils.jl b/src/lux_utils.jl new file mode 100644 index 0000000..2f94612 --- /dev/null +++ b/src/lux_utils.jl @@ -0,0 +1,22 @@ +using LuxCore + +function replace_namedtuples(nt, rp_st, Σ) + if length(nt) == 0 + return rp_st + else + for i in 1:length(nt) + if length(nt[i]) == 0 + rp_st = (; rp_st..., (; keys(nt)[i] => (Σ = Σ, ))...) + else + rp_st = (; rp_st..., keys(nt)[i] => replace_namedtuples(nt[i], (;), Σ)) + end + end + return rp_st + end +end + +function setupBFState(rng, bf, Σ) + ps, st = LuxCore.setup(rng, bf) + rp_st = replace_namedtuples(st, (;), Σ) + return ps, rp_st +end \ No newline at end of file diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 6cb16f1..1f6a1c0 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -1,34 +1,12 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate -using ACEpsi: BackflowPooling +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState using ACEpsi: BFwf_lux using ACEbase.Testing: print_tf using LuxCore using Random - -function replace_namedtuples(nt, rp_st) - if length(nt) == 0 - return rp_st - else - for i in 1:length(nt) - if length(nt[i]) == 0 - rp_st = (; rp_st..., (; keys(nt)[i] => (Σ = Σ, ))...) - else - rp_st = (; rp_st..., keys(nt)[i] => replace_namedtuples(nt[i], (;))) - end - end - return rp_st - end -end - -function setupBFState(rng, bf, Σ) - ps, st = LuxCore.setup(rng, bf) - rp_st = replace_namedtuples(st, (;)) - return ps, rp_st -end - Rnldegree = 4 Ylmdegree = 4 totdegree = 8 @@ -49,77 +27,77 @@ ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) @info("Test evaluate") A1 = BFwf_chain(X, ps, st) -# original implementation -function BFwf_lux2(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, - ν = 3, T = Float64, - sd_admissible = bb -> (true), - envelope = x -> x) # enveolpe to be replaced by SJ-factor +# # original implementation +# function BFwf_lux2(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, +# ν = 3, T = Float64, +# sd_admissible = bb -> (true), +# envelope = x -> x) # enveolpe to be replaced by SJ-factor - spec1p = make_nlms_spec(bRnl, bYlm; - totaldegree = totdeg) +# spec1p = make_nlms_spec(bRnl, bYlm; +# totaldegree = totdeg) - # size(X) = (nX, 3); length(Σ) = nX - aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) - pooling = BackflowPooling(aobasis) +# # size(X) = (nX, 3); length(Σ) = nX +# aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) +# pooling = BackflowPooling(aobasis) - # define sparse for n-correlations - tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] - default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) +# # define sparse for n-correlations +# tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] +# default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) - specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, - minvv = fill(0, ν), - maxvv = fill(length(spec1p), ν), - ordered = true) - spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] +# specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, +# minvv = fill(0, ν), +# maxvv = fill(length(spec1p), ν), +# ordered = true) +# spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] - # further restrict - spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] +# # further restrict +# spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] - # define n-correlation - corr1 = SparseSymmProd(spec; T = Float64) +# # define n-correlation +# corr1 = SparseSymmProd(spec; T = Float64) - # ----------- Lux connections --------- - # Should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm for trainable radial basis later? - # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1)) - aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) - # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) - pooling_layer = ACEpsi.lux(pooling) - # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) - corr_layer = ACEcore.lux(corr1) +# # ----------- Lux connections --------- +# # Should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm for trainable radial basis later? +# # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1)) +# aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) +# # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) +# pooling_layer = ACEpsi.lux(pooling) +# # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) +# corr_layer = ACEcore.lux(corr1) - # TODO: Add J-factor and add trainable basis later - js = Jastrow(nuclei) - jastrow_layer = ACEpsi.lux(js) +# # TODO: Add J-factor and add trainable basis later +# js = Jastrow(nuclei) +# jastrow_layer = ACEpsi.lux(js) - reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) - # Questions to discuss: - # 1. it seems that we do not need trans since the bases have already taken care of it? - # 2. Why we need a tranpose here??? Seems that the output from corr_layer is (length(spec), nX)??? - # 3. How do we define n-correlations if we use trainable basis? - BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), - bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), - Mask = ACEpsi.MaskLayer(Nel), logabsdet = WrappedFunction(x -> det(x))) - return Chain(; branch = BranchLayer((js = jastrow_layer, bf = BFwf_chain, ))) - end +# reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) +# # Questions to discuss: +# # 1. it seems that we do not need trans since the bases have already taken care of it? +# # 2. Why we need a tranpose here??? Seems that the output from corr_layer is (length(spec), nX)??? +# # 3. How do we define n-correlations if we use trainable basis? +# BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), +# bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), +# Mask = ACEpsi.MaskLayer(Nel), logabsdet = WrappedFunction(x -> det(x))) +# return Chain(; branch = BranchLayer((js = jastrow_layer, bf = BFwf_chain, ))) +# end -using ACEcore, Polynomials4ML -using Polynomials4ML: OrthPolyBasis1D3T -using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! -using ACEcore.Utils: gensparse -using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det -import ForwardDiff -using ACEpsi.AtomicOrbitals: make_nlms_spec +# using ACEcore, Polynomials4ML +# using Polynomials4ML: OrthPolyBasis1D3T +# using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! +# using ACEcore.Utils: gensparse +# using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det +# import ForwardDiff +# using ACEpsi.AtomicOrbitals: make_nlms_spec -using LuxCore: AbstractExplicitLayer -using Lux: Dense, Chain, WrappedFunction, BranchLayer -using ACEpsi: Jastrow +# using LuxCore: AbstractExplicitLayer +# using Lux: Dense, Chain, WrappedFunction, BranchLayer +# using ACEpsi: Jastrow -js = Jastrow(nuclei) -jatrow_layer = ACEpsi.lux(js) -branchlayer = BFwf_lux2(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) -ps2, st2 = setupBFState(MersenneTwister(1234), branchlayer, Σ) +# js = Jastrow(nuclei) +# jatrow_layer = ACEpsi.lux(js) +# branchlayer = BFwf_lux2(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +# ps2, st2 = setupBFState(MersenneTwister(1234), branchlayer, Σ) -A2 = branchlayer(X, ps2, st2) -@test(2 * logabsdet(A2[1][1] *A2[1][2])[1] == A1[1]) +# A2 = branchlayer(X, ps2, st2) +# @test(2 * logabsdet(A2[1][1] *A2[1][2])[1] == A1[1]) From cfcab47a7371b834a2f1bb10823b47ab56786888 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Thu, 18 May 2023 01:31:00 +0800 Subject: [PATCH 18/79] add test_rrule for js --- src/jastrow.jl | 13 +++++++++++++ test/test_bflow_lux.jl | 18 ++++++++++++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/src/jastrow.jl b/src/jastrow.jl index 29b3165..649b491 100644 --- a/src/jastrow.jl +++ b/src/jastrow.jl @@ -4,6 +4,8 @@ using ACEpsi.AtomicOrbitals: Nuc using LuxCore: AbstractExplicitLayer using Random: AbstractRNG +import ChainRules +import ChainRules: rrule, NoTangent mutable struct Jastrow{T} nuclei::Vector{Nuc{T}} # nuclei end @@ -57,3 +59,14 @@ initialstates(rng::AbstractRNG, l::JastrowLayer) = _init_luxstate(rng, l.basis) # This should be removed later and replace by ObejctPools (l::JastrowLayer)(X, ps, st) = evaluate(l.basis, X, st.Σ), st + +function rrule(::typeof(evaluate), js::Jastrow, X::AbstractVector, Σ::AbstractVector) + J = evaluate(js, X) + + function pb(dJ) + @assert dJ isa Real + return NoTangent(), NoTangent(), X, NoTangent() + end + + return J, pb + end \ No newline at end of file diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 6cb16f1..1c62bce 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -8,7 +8,7 @@ using LuxCore using Random -function replace_namedtuples(nt, rp_st) +function replace_namedtuples(nt, rp_st, Σ) if length(nt) == 0 return rp_st else @@ -25,7 +25,7 @@ end function setupBFState(rng, bf, Σ) ps, st = LuxCore.setup(rng, bf) - rp_st = replace_namedtuples(st, (;)) + rp_st = replace_namedtuples(st, (;), Σ) return ps, rp_st end @@ -49,6 +49,8 @@ ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) @info("Test evaluate") A1 = BFwf_chain(X, ps, st) + + # original implementation function BFwf_lux2(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, ν = 3, T = Float64, @@ -123,3 +125,15 @@ A2 = branchlayer(X, ps2, st2) @test(2 * logabsdet(A2[1][1] *A2[1][2])[1] == A1[1]) +using Lux +using Zygote +using ACEpsi:evaluate + +js = Jastrow(nuclei) +jatrow_layer = lux(js) +js_chain = Chain(; jatrow_layer) +ps, st = setupBFState(MersenneTwister(1234), js_chain, Σ) + +gs = Zygote.gradient(X -> js_chain(X, ps, st)[1], X) +Zygote.gradient(X -> evaluate(js, X, Σ),X) + From fdb23d561d46e435e8570194fd0ff3701f21b2e6 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Thu, 18 May 2023 01:38:20 +0800 Subject: [PATCH 19/79] minor fix --- Project.toml | 1 + src/jastrow.jl | 2 +- test/test_bflow_lux.jl | 81 ++---------------------------------------- 3 files changed, 5 insertions(+), 79 deletions(-) diff --git a/Project.toml b/Project.toml index 6afc154..7ab238a 100644 --- a/Project.toml +++ b/Project.toml @@ -7,6 +7,7 @@ version = "0.0.1" ACEbase = "14bae519-eb20-449c-a949-9c58ed33163e" ACEcore = "44c1e890-45d1-48ea-94d6-c2ea5b573f71" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/src/jastrow.jl b/src/jastrow.jl index 649b491..4939786 100644 --- a/src/jastrow.jl +++ b/src/jastrow.jl @@ -61,7 +61,7 @@ initialstates(rng::AbstractRNG, l::JastrowLayer) = _init_luxstate(rng, l.basis) evaluate(l.basis, X, st.Σ), st function rrule(::typeof(evaluate), js::Jastrow, X::AbstractVector, Σ::AbstractVector) - J = evaluate(js, X) + J = evaluate(js, X, Σ) function pb(dJ) @assert dJ isa Real diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 1c62bce..4e878bb 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -16,7 +16,7 @@ function replace_namedtuples(nt, rp_st, Σ) if length(nt[i]) == 0 rp_st = (; rp_st..., (; keys(nt)[i] => (Σ = Σ, ))...) else - rp_st = (; rp_st..., keys(nt)[i] => replace_namedtuples(nt[i], (;))) + rp_st = (; rp_st..., keys(nt)[i] => replace_namedtuples(nt[i], (;), Σ)) end end return rp_st @@ -49,91 +49,16 @@ ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) @info("Test evaluate") A1 = BFwf_chain(X, ps, st) - - -# original implementation -function BFwf_lux2(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, - ν = 3, T = Float64, - sd_admissible = bb -> (true), - envelope = x -> x) # enveolpe to be replaced by SJ-factor - - spec1p = make_nlms_spec(bRnl, bYlm; - totaldegree = totdeg) - - # size(X) = (nX, 3); length(Σ) = nX - aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) - pooling = BackflowPooling(aobasis) - - # define sparse for n-correlations - tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] - default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) - - specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, - minvv = fill(0, ν), - maxvv = fill(length(spec1p), ν), - ordered = true) - spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] - - # further restrict - spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] - - # define n-correlation - corr1 = SparseSymmProd(spec; T = Float64) - - # ----------- Lux connections --------- - # Should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm for trainable radial basis later? - # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1)) - aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) - # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) - pooling_layer = ACEpsi.lux(pooling) - # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) - corr_layer = ACEcore.lux(corr1) - - # TODO: Add J-factor and add trainable basis later - js = Jastrow(nuclei) - jastrow_layer = ACEpsi.lux(js) - - reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) - # Questions to discuss: - # 1. it seems that we do not need trans since the bases have already taken care of it? - # 2. Why we need a tranpose here??? Seems that the output from corr_layer is (length(spec), nX)??? - # 3. How do we define n-correlations if we use trainable basis? - BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), - bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), - Mask = ACEpsi.MaskLayer(Nel), logabsdet = WrappedFunction(x -> det(x))) - return Chain(; branch = BranchLayer((js = jastrow_layer, bf = BFwf_chain, ))) - end - -using ACEcore, Polynomials4ML -using Polynomials4ML: OrthPolyBasis1D3T -using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! -using ACEcore.Utils: gensparse -using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det -import ForwardDiff -using ACEpsi.AtomicOrbitals: make_nlms_spec - -using LuxCore: AbstractExplicitLayer -using Lux: Dense, Chain, WrappedFunction, BranchLayer using ACEpsi: Jastrow - -js = Jastrow(nuclei) -jatrow_layer = ACEpsi.lux(js) -branchlayer = BFwf_lux2(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) -ps2, st2 = setupBFState(MersenneTwister(1234), branchlayer, Σ) - -A2 = branchlayer(X, ps2, st2) -@test(2 * logabsdet(A2[1][1] *A2[1][2])[1] == A1[1]) - - using Lux using Zygote using ACEpsi:evaluate js = Jastrow(nuclei) -jatrow_layer = lux(js) +jatrow_layer = ACEpsi.lux(js) js_chain = Chain(; jatrow_layer) ps, st = setupBFState(MersenneTwister(1234), js_chain, Σ) gs = Zygote.gradient(X -> js_chain(X, ps, st)[1], X) -Zygote.gradient(X -> evaluate(js, X, Σ),X) +Zygote.gradient(X -> ACEpsi.evaluate(js, X, Σ),X) From 41a2447dba93d13347974d4af77c5fdf6ab4017a Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Thu, 18 May 2023 09:44:46 -0700 Subject: [PATCH 20/79] add backflow pooling rrule --- src/atomicorbitals/atorbbasis.jl | 8 +++--- src/backflowpooling.jl | 42 +++++++++++++++++++++++++++++++- test/test_atorbbasis.jl | 22 ++++++++++++++++- test/test_bflow_lux.jl | 4 ++- 4 files changed, 69 insertions(+), 7 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 4b123e6..b9ce674 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -1,7 +1,7 @@ import Polynomials4ML: evaluate using ACEpsi: ↑, ↓, ∅, spins, extspins, Spin, spin2idx, idx2spin -using Polynomials4ML: SparseProduct, ProductBasis +using Polynomials4ML: SparseProduct, ProductBasis, _make_reqfields using LuxCore: AbstractExplicitLayer using Random: AbstractRNG @@ -76,7 +76,7 @@ function ProductBasis(spec1, bRnl, bYlm) spec1idx[i] = (inv_Rnl[dropnames(b,(:m,))], inv_Ylm[(l=b.l, m=b.m)]) end sparsebasis = SparseProduct(spec1idx) - return ProductBasis(spec1, bRnl, bYlm, sparsebasis) + return ProductBasis(spec1, bRnl, bYlm, sparsebasis, _make_reqfields()...) end @@ -86,7 +86,7 @@ mutable struct AtomicOrbitalsBasis{NB, T} end (aobasis::AtomicOrbitalsBasis)(args...) = evaluate(aobasis, args...) - +Base.length(aobasis::AtomicOrbitalsBasis) = length(aobasis.prodbasis.spec1) * length(aobasis.nuclei) function AtomicOrbitalsBasis(bRnl, bYlm; totaldegree=3, @@ -133,7 +133,7 @@ function get_spec(basis::AtomicOrbitalsBasis) spec = [] Nnuc = length(basis.nuclei) - spec = Array{3}(undef, (3, Nnuc, length(basis.prodbasis.spec1))) + spec = Array{Any}(undef, (3, Nnuc, length(basis.prodbasis.spec1))) for (k, nlm) in enumerate(basis.prodbasis.spec1) for I = 1:Nnuc diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index 2607bf0..6ccfa6b 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -8,7 +8,7 @@ end (pooling::BackflowPooling)(args...) = evaluate(pooling, args...) -function evaluate(pooling::BackflowPooling, ϕnlm, Σ) +function evaluate(pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) basis = pooling.basis nuc = basis.nuclei Nnuc = length(nuc) @@ -65,8 +65,47 @@ function evaluate(pooling::BackflowPooling, ϕnlm, Σ) return A end +# --------------------- connect with ChainRule +function rrule(::typeof(evaluate), pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) + return _rrule_evaluate(pooling, ϕnlm, Σ) +end + +function _rrule_evaluate(pooling::BackflowPooling, ϕnlm, Σ) + A = pooling(ϕnlm, Σ) + return A, ∂A -> _pullback_evaluate(∂A, pooling, ϕnlm, Σ) +end + +function _pullback_evaluate(∂A, pooling::BackflowPooling, ϕnlm, Σ) + TA = promote_type(eltype.(ϕnlm)...) + ∂ϕnlm = zeros(TA, size(ϕnlm)) + _pullback_evaluate!(∂ϕnlm, ∂A, pooling, ϕnlm, Σ) + return ∂ϕnlm +end +function _pullback_evaluate!(∂ϕnlm, ∂A, pooling::BackflowPooling, ϕnlm, Σ) + Nnuc, Nel, Nnlm = size(ϕnlm) + basis = pooling.basis + + @assert Nnlm == length(basis.prodbasis.spec1) + @assert Nel == length(Σ) + @assert size(∂ϕnlm) == (Nnuc, Nel, Nnlm) + @assert size(∂A) == (Nel, 3, Nnuc, Nnlm) + + for I = 1:Nnuc + for i = 1:Nel + for k = 1:Nnlm + ∂ϕnlm[I, i, k] += ∂A[i, 3, I, k] + for ii = 1:Nel + ∂ϕnlm[I, i, k] += ∂A[ii, spin2idx(Σ[i]), I, k] .* (i != ii) + end + end + end + end + + return nothing +end + # --------------------- connect with Lux struct BackflowPoolingLayer <: AbstractExplicitLayer @@ -92,3 +131,4 @@ initialstates(rng::AbstractRNG, l::BackflowPoolingLayer) = _init_luxstate(rng, l # evaluate!(parent(B), l.basis, ϕnlm_Σ[1], ϕnlm_Σ[2]) # return B # end + diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 3a7dedc..866a081 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -2,7 +2,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate using ACEpsi: BackflowPooling -using ACEbase.Testing: print_tf +using ACEbase.Testing: print_tf, fdtest using LuxCore using Random @@ -90,6 +90,26 @@ end println() # +@info("---------- rrule tests ----------") +using LinearAlgebra: dot + +@info("BackFlowPooling rrule") +for ntest = 1:30 + local testϕnlm + testϕnlm = randn(size(bϕnlm)) + bdd = randn(size(bϕnlm)) + _BB(t) = testϕnlm + t * bdd + bA2 = pooling(testϕnlm, Σ) + u = randn(size(bA2)) + F(t) = dot(u, pooling(_BB(t), Σ)) + dF(t) = begin + val, pb = ACEpsi._rrule_evaluate(pooling, _BB(t), Σ) + ∂BB = pb(u) + return dot(∂BB, bdd) + end + print_tf(@test fdtest(F, dF, 0.0; verbose=false)) +end +println() @info("---------- Lux tests ----------") aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 4e878bb..5f3d1b2 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -49,6 +49,8 @@ ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) @info("Test evaluate") A1 = BFwf_chain(X, ps, st) + +# try with rrule using ACEpsi: Jastrow using Lux using Zygote @@ -60,5 +62,5 @@ js_chain = Chain(; jatrow_layer) ps, st = setupBFState(MersenneTwister(1234), js_chain, Σ) gs = Zygote.gradient(X -> js_chain(X, ps, st)[1], X) -Zygote.gradient(X -> ACEpsi.evaluate(js, X, Σ),X) +Zygote.gradient(X -> ACEpsi.evaluate(js, X, Σ), X) From 9b57699730eaa8cf806b51d11a44c0225e721f2f Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Fri, 19 May 2023 00:51:25 +0800 Subject: [PATCH 21/79] fix Zygote of jastrow --- src/jastrow.jl | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/src/jastrow.jl b/src/jastrow.jl index 4939786..2b08819 100644 --- a/src/jastrow.jl +++ b/src/jastrow.jl @@ -6,6 +6,8 @@ using Random: AbstractRNG import ChainRules import ChainRules: rrule, NoTangent +using Zygote: Buffer + mutable struct Jastrow{T} nuclei::Vector{Nuc{T}} # nuclei end @@ -20,7 +22,8 @@ function evaluate(f::Jastrow, X::AbstractVector, Σ) Nnuc = length(nuc) Nel = size(X, 1) T = promote_type(eltype(X[1])) - XN = zeros(T, (Nnuc, Nel)) + XX = zeros(T, (Nnuc, Nel)) + XN = Buffer(XX) C0 = (2-pi)/(12*pi) F2 = zero(T) @@ -59,14 +62,3 @@ initialstates(rng::AbstractRNG, l::JastrowLayer) = _init_luxstate(rng, l.basis) # This should be removed later and replace by ObejctPools (l::JastrowLayer)(X, ps, st) = evaluate(l.basis, X, st.Σ), st - -function rrule(::typeof(evaluate), js::Jastrow, X::AbstractVector, Σ::AbstractVector) - J = evaluate(js, X, Σ) - - function pb(dJ) - @assert dJ isa Real - return NoTangent(), NoTangent(), X, NoTangent() - end - - return J, pb - end \ No newline at end of file From ab8da0f47a7ff0b3bdfdc0aac2c595606918110f Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 23 May 2023 08:14:49 -0700 Subject: [PATCH 22/79] draft towards grad --- Project.toml | 3 ++ src/atomicorbitals/atorbbasis.jl | 16 +++++++ src/backflowpooling.jl | 10 ++++- test/test_atorbbasis.jl | 11 ++++- test/test_bflow_lux.jl | 34 ++++++++++----- test/try.jl | 74 ++++++++++++++++++++++++++++++++ 6 files changed, 133 insertions(+), 15 deletions(-) create mode 100644 test/try.jl diff --git a/Project.toml b/Project.toml index 7ab238a..af6f4c1 100644 --- a/Project.toml +++ b/Project.toml @@ -8,6 +8,8 @@ ACEbase = "14bae519-eb20-449c-a949-9c58ed33163e" ACEcore = "44c1e890-45d1-48ea-94d6-c2ea5b573f71" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -19,6 +21,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] ACEbase = "0.4.2" diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index b9ce674..077daae 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -9,6 +9,9 @@ using Polynomials4ML using StaticArrays using LinearAlgebra: norm +using ChainRulesCore +using ChainRulesCore: NoTangent + struct Nuc{T} rr::SVector{3, T} charge::T # should this be an integer? @@ -150,6 +153,18 @@ end # ------------ Evaluation kernels +# ------------ connect with ChainRulesCore + +# Placeholder for now, fix this later after making sure Zygote is done correct with Lux +function ChainRulesCore.rrule(::typeof(evaluate), basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) + val = evaluate(basis, X, Σ) + dB = similar(X) + function pb(dA) + return NoTangent(), NoTangent(), dB, NoTangent() + end + return val, pb +end + # ------------ connect with Lux struct AtomicOrbitalsBasisLayer{TB} <: AbstractExplicitLayer basis::TB @@ -169,6 +184,7 @@ initialstates(rng::AbstractRNG, l::AtomicOrbitalsBasisLayer) = _init_luxstate(rn + # ----- ObejctPools # (l::AtomicOrbitalsBasisLayer)(args...) = evaluate(l, args...) diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index 6ccfa6b..07bc326 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -1,6 +1,8 @@ using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis using LuxCore: AbstractExplicitLayer using Random: AbstractRNG +using ChainRulesCore +using ChainRulesCore: NoTangent mutable struct BackflowPooling basis::AtomicOrbitalsBasis @@ -66,8 +68,12 @@ function evaluate(pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) end # --------------------- connect with ChainRule -function rrule(::typeof(evaluate), pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) - return _rrule_evaluate(pooling, ϕnlm, Σ) +function ChainRulesCore.rrule(::typeof(evaluate), pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) + A = pooling(ϕnlm, Σ) + function pb(∂A) + return NoTangent(), NoTangent(), _pullback_evaluate(∂A, pooling, ϕnlm, Σ), NoTangent() + end + return A, pb end function _rrule_evaluate(pooling::BackflowPooling, ϕnlm, Σ) diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 866a081..d6c59c3 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -5,6 +5,7 @@ using ACEpsi: BackflowPooling using ACEbase.Testing: print_tf, fdtest using LuxCore using Random +using Zygote # test configs Rnldegree = 4 @@ -111,6 +112,13 @@ for ntest = 1:30 end println() +@info("Checking Zygote running correctly") +val, pb = Zygote.pullback(pooling, bϕnlm, Σ) +val1, pb1 = ACEpsi._rrule_evaluate(pooling, bϕnlm, Σ) +@assert val1 ≈ val1 +@assert pb1(val) ≈ pb(val)[1] # pb(val)[2] is for Σ with no pb + + @info("---------- Lux tests ----------") aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) pooling_layer = ACEpsi.lux(pooling) @@ -132,5 +140,4 @@ chain_st = (aobasis = (Σ = Σ, ), pooling = (Σ = Σ, )) # try Chain is as expected print_tf(@test tryChain(X, chain_ps, chain_st)[1] ≈ A) -println() - +println() \ No newline at end of file diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index b951609..685a25a 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -28,18 +28,30 @@ ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) @info("Test evaluate") A1 = BFwf_chain(X, ps, st) +@info("Test Zygote API") +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) + +y, st = Lux.apply(BFwf_chain, X, ps, st) + +## Pullback API to capture change in state +(l, st_), pb = pullback(p -> Lux.apply(BFwf_chain, X, p, st), ps) +gs = pb((one.(l), nothing))[1] + +# Jastrow: try with gradient +# using ACEpsi: Jastrow +# using Lux +# using Zygote +# using ACEpsi:evaluate + +# js = Jastrow(nuclei) +# jatrow_layer = ACEpsi.lux(js) +# js_chain = Chain(; jatrow_layer) +# ps, st = setupBFState(MersenneTwister(1234), js_chain, Σ) + +# gs = Zygote.gradient(X -> js_chain(X, ps, st)[1], X) +# Zygote.gradient(X -> ACEpsi.evaluate(js, X, Σ), X) -# try with rrule -using ACEpsi: Jastrow -using Lux -using Zygote -using ACEpsi:evaluate -js = Jastrow(nuclei) -jatrow_layer = ACEpsi.lux(js) -js_chain = Chain(; jatrow_layer) -ps, st = setupBFState(MersenneTwister(1234), js_chain, Σ) +# BackFlowPooling: Try with rrule -gs = Zygote.gradient(X -> js_chain(X, ps, st)[1], X) -Zygote.gradient(X -> ACEpsi.evaluate(js, X, Σ), X) diff --git a/test/try.jl b/test/try.jl new file mode 100644 index 0000000..4e94c5e --- /dev/null +++ b/test/try.jl @@ -0,0 +1,74 @@ +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow +using ACEbase.Testing: print_tf +using LuxCore +using Lux +using Zygote +using Random +using ACEcore.Utils: gensparse +using ACEcore: SparseSymmProd +using ACEcore + +Rnldegree = 4 +Ylmdegree = 4 +ν = 2 +totdeg = 10 +Nel = 5 +X = randn(SVector{3, Float64}, Nel) +Σ = rand(spins(), Nel) +sd_admissible = bb -> (true) + + +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] +## + +# Defining AtomicOrbitalsBasis +bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +bYlm = RYlmBasis(Ylmdegree) + +spec1p = make_nlms_spec(bRnl, bYlm; + totaldegree = totdeg) + +# size(X) = (nX, 3); length(Σ) = nX +aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) +pooling = BackflowPooling(aobasis) + +# define sparse for n-correlations +tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] +default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) + +specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, + minvv = fill(0, ν), + maxvv = fill(length(spec1p), ν), + ordered = true) +spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] + +# further restrict +spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] + +# define n-correlation +corr1 = SparseSymmProd(spec; T = Float64) + +reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) + + +pooling_layer = ACEpsi.lux(pooling) +# (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) +corr_layer = ACEcore.lux(corr1) + +pool2AAChain = Chain(; pooling = pooling_layer, reshape = WrappedFunction(reshape_func), +bAA = corr_layer) + + +# dummy input +ϕnlm = aobasis(X, Σ) + +ps, st = setupBFState(MersenneTwister(1234), pool2AAChain, Σ) + +y, st = Lux.apply(pool2AAChain, ϕnlm, ps, st) + +## Pullback API to capture change in state +(l, st_), pb = pullback(p -> Lux.apply(pool2AAChain, ϕnlm, p, st), ps) +gs = pb((one.(l), nothing))[1] From 0db7beabed8d4ad2758780e54f9ee206aaff0785 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 24 May 2023 01:19:55 +0800 Subject: [PATCH 23/79] test layers' pullback --- test/test2.jl | 113 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 113 insertions(+) create mode 100644 test/test2.jl diff --git a/test/test2.jl b/test/test2.jl new file mode 100644 index 0000000..2df9266 --- /dev/null +++ b/test/test2.jl @@ -0,0 +1,113 @@ +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow +using ACEbase.Testing: print_tf +using LuxCore +using Lux +using Zygote +using Random +using ACEcore.Utils: gensparse +using ACEcore: SparseSymmProd +using ACEcore + +Rnldegree = 4 +Ylmdegree = 4 +totdegree = 8 + +Nel = 5 +X = randn(SVector{3, Float64}, Nel) +Σ = rand(spins(), Nel) +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] +## + +# Defining AtomicOrbitalsBasis +bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +bYlm = RYlmBasis(Ylmdegree) +spec1 = make_nlms_spec(bRnl, bYlm; totaldegree = totdegree) + +# define basis and pooling operations +prodbasis = ProductBasis(spec1, bRnl, bYlm) +aobasis = AtomicOrbitalsBasis(prodbasis, nuclei) +pooling = BackflowPooling(aobasis) +ϕnlm = prodbasis(X) +bϕnlm = aobasis(X, Σ) +A = pooling(bϕnlm, Σ) +ν = 3 +tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] +default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) + +specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, + minvv = fill(0, ν), + maxvv = fill(length(spec1p), ν), + ordered = true) +spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] + +# further restrict +spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] + +# define n-correlation +corr1 = SparseSymmProd(spec; T = Float64) + +# pooling layer +pooling_layer = ACEpsi.lux(pooling) +pChain = Chain(; reshape_layer) +ps, st = setupBFState(MersenneTwister(1234), pChain, Σ) +y, st = Lux.apply(pChain, A, ps, st) +(l, st_), pb = pullback(ϕnlm -> Lux.apply(pChain, ϕnlm, ps, st), ϕnlm) +gs = pb((l, nothing))[1] + +# reshape layer +reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) +reshape_layer = WrappedFunction(reshape_func) +rChain = Chain(; reshape_layer) +ps, st = setupBFState(MersenneTwister(1234), rChain, Σ) +y, st = Lux.apply(rChain, A, ps, st) +(l, st_), pb = pullback(A -> Lux.apply(rChain, A, ps, st), A) +gs = pb((l, nothing))[1] + +A = l +# corr_layer +corr_layer = ACEcore.lux(corr1) +cChain = Chain(; corr_layer) +ps, st = setupBFState(MersenneTwister(1234), cChain, Σ) +y, st = Lux.apply(cChain, A, ps, st) +(l, st_), pb = pullback(A -> Lux.apply(cChain, A, ps, st), A) +gs = pb((l, nothing))[1] + +# transpose layer +ϕ = l +transpose_layer = WrappedFunction(transpose) +tChain = Chain(; transpose_layer) +ps, st = setupBFState(MersenneTwister(1234), tChain, Σ) +y, st = Lux.apply(tChain, ϕ, ps, st) +(l, st_), pb = pullback(ϕ -> Lux.apply(tChain, ϕ, ps, st), ϕ) +gs = pb((l, nothing))[1] + +# hidden layer +ϕ = l +hidden_layer = Dense(length(corr1), Nel) +hChain = Chain(; hidden_layer) +ps, st = setupBFState(MersenneTwister(1234), hChain, Σ) +y, st = Lux.apply(hChain, ϕ, ps, st) +(l, st_), pb = pullback(ϕ -> Lux.apply(hChain, ϕ, ps, st), ϕ) +gs = pb((l, nothing))[1] + +# Mask layer +ϕ = l +Mask = ACEpsi.MaskLayer(Nel) +mChain = Chain(; Mask) +ps, st = setupBFState(MersenneTwister(1234), mChain, Σ) +y, st = Lux.apply(mChain, ϕ, ps, st) +(l, st_), pb = pullback(ϕ -> Lux.apply(mChain, ϕ, ps, st), ϕ) +gs = pb((l, nothing))[1] + +# det layer +using LinearAlgebra:det +ϕ = l +det_layer = WrappedFunction(x -> det(x)) +dChain = Chain(; det_layer) +ps, st = setupBFState(MersenneTwister(1234), dChain, Σ) +y, st = Lux.apply(dChain, ϕ, ps, st) +(l, st_), pb = pullback(ϕ -> Lux.apply(dChain, ϕ, ps, st), ϕ) +gs = pb((l, nothing))[1] \ No newline at end of file From f5eae2c92ee16193170d20a141abf1bc2ccc83d1 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 24 May 2023 01:47:39 +0800 Subject: [PATCH 24/79] add aobasis layer' pull back --- test/test2.jl | 46 +++++++++++++++++++++++++++++++++++++--------- 1 file changed, 37 insertions(+), 9 deletions(-) diff --git a/test/test2.jl b/test/test2.jl index 2df9266..abdc5cf 100644 --- a/test/test2.jl +++ b/test/test2.jl @@ -34,27 +34,36 @@ pooling = BackflowPooling(aobasis) bϕnlm = aobasis(X, Σ) A = pooling(bϕnlm, Σ) ν = 3 -tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] +totdeg = 6 +tup2b = vv -> [ spec1[v] for v in vv[vv .> 0] ] default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, minvv = fill(0, ν), - maxvv = fill(length(spec1p), ν), + maxvv = fill(length(spec1), ν), ordered = true) spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] -# further restrict -spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] - # define n-correlation corr1 = SparseSymmProd(spec; T = Float64) + + +# aobasis layer +aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) +aChain = Chain(; aobasis_layer) +ps, st = setupBFState(MersenneTwister(1234), aChain, Σ) +y, st = Lux.apply(aChain, X, ps, st) +(l, st_), pb = pullback(X -> Lux.apply(aChain, X, ps, st), X) +gs = pb((l, nothing))[1] + # pooling layer +bϕnlm = l pooling_layer = ACEpsi.lux(pooling) -pChain = Chain(; reshape_layer) +pChain = Chain(; pooling_layer) ps, st = setupBFState(MersenneTwister(1234), pChain, Σ) -y, st = Lux.apply(pChain, A, ps, st) -(l, st_), pb = pullback(ϕnlm -> Lux.apply(pChain, ϕnlm, ps, st), ϕnlm) +y, st = Lux.apply(pChain, l, ps, st) +(l, st_), pb = pullback(ϕnlm -> Lux.apply(pChain, ϕnlm, ps, st), bϕnlm) gs = pb((l, nothing))[1] # reshape layer @@ -66,8 +75,15 @@ y, st = Lux.apply(rChain, A, ps, st) (l, st_), pb = pullback(A -> Lux.apply(rChain, A, ps, st), A) gs = pb((l, nothing))[1] -A = l +# first 3 layer +fChain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func)) +ps, st = setupBFState(MersenneTwister(1234), fChain, Σ) +y, st = Lux.apply(fChain, X, ps, st) +(l, st_), pb = pullback(X -> Lux.apply(fChain, X, ps, st), X) +gs = pb((l, nothing))[1] + # corr_layer +A = l corr_layer = ACEcore.lux(corr1) cChain = Chain(; corr_layer) ps, st = setupBFState(MersenneTwister(1234), cChain, Σ) @@ -75,6 +91,7 @@ y, st = Lux.apply(cChain, A, ps, st) (l, st_), pb = pullback(A -> Lux.apply(cChain, A, ps, st), A) gs = pb((l, nothing))[1] +corr_l = l # transpose layer ϕ = l transpose_layer = WrappedFunction(transpose) @@ -110,4 +127,15 @@ dChain = Chain(; det_layer) ps, st = setupBFState(MersenneTwister(1234), dChain, Σ) y, st = Lux.apply(dChain, ϕ, ps, st) (l, st_), pb = pullback(ϕ -> Lux.apply(dChain, ϕ, ps, st), ϕ) +gs = pb((l, nothing))[1] + +# last 4 layer +using LinearAlgebra:transpose +A = corr_l + +lChain = Chain(; transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), + Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) +ps, st = setupBFState(MersenneTwister(1234), lChain, Σ) +y, st = Lux.apply(lChain, A, ps, st) +(l, st_), pb = pullback(A -> Lux.apply(lChain, A, ps, st), A) gs = pb((l, nothing))[1] \ No newline at end of file From 240f7cd74331c2ad140b726aa9ade7f4d672a03e Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 24 May 2023 01:51:54 +0800 Subject: [PATCH 25/79] minor fix --- test/test2.jl | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/test/test2.jl b/test/test2.jl index abdc5cf..0a5a27e 100644 --- a/test/test2.jl +++ b/test/test2.jl @@ -138,4 +138,20 @@ lChain = Chain(; transpose_layer = WrappedFunction(transpose), hidden1 = Dense(l ps, st = setupBFState(MersenneTwister(1234), lChain, Σ) y, st = Lux.apply(lChain, A, ps, st) (l, st_), pb = pullback(A -> Lux.apply(lChain, A, ps, st), A) +gs = pb((l, nothing))[1] + + + + + +using ACEpsi: Jastrow + +js = Jastrow(nuclei) +jastrow_layer = ACEpsi.lux(js) +jastrow_layer1 = ACEpsi.lux(js) +bChain = Chain(; branch = BranchLayer((js = jastrow_layer, bf = jastrow_layer1, )), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ) + +ps, st = setupBFState(MersenneTwister(1234), bChain, Σ) +y, st = Lux.apply(bChain, X, ps, st) +(l, st_), pb = pullback(X -> Lux.apply(bChain, X, ps, st), X) gs = pb((l, nothing))[1] \ No newline at end of file From d5735db572ed1e0547fa090b2700ad91f1b03a4f Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 23 May 2023 11:15:43 -0700 Subject: [PATCH 26/79] Calling rrule correctly, fix correctness next --- src/backflowpooling.jl | 2 ++ src/bflow3d.jl | 4 ++-- test/test2.jl | 4 ++-- test/test_bflow_lux.jl | 4 +--- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index 07bc326..7b6f7b1 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -10,6 +10,8 @@ end (pooling::BackflowPooling)(args...) = evaluate(pooling, args...) +Base.length(pooling::BackflowPooling) = 3 * length(pooling.basis) # length(spin()) * length(1pbasis) + function evaluate(pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) basis = pooling.basis nuc = basis.nuclei diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 979d2b1..265150d 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -44,7 +44,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] # define n-correlation - corr1 = SparseSymmProd(spec; T = Float64) + corr1 = Polynomials4ML.SparseSymmProd(spec; T = Float64) # ----------- Lux connections --------- # Should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm for trainable radial basis later? @@ -53,7 +53,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) pooling_layer = ACEpsi.lux(pooling) # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) - corr_layer = ACEcore.lux(corr1) + corr_layer = Polynomials4ML.lux(corr1) # TODO: Add J-factor and add trainable basis later js = Jastrow(nuclei) diff --git a/test/test2.jl b/test/test2.jl index 0a5a27e..c1f39af 100644 --- a/test/test2.jl +++ b/test/test2.jl @@ -45,7 +45,7 @@ specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] # define n-correlation -corr1 = SparseSymmProd(spec; T = Float64) +corr1 = Polynomials4ML.SparseSymmProd(spec; T = Float64) @@ -84,7 +84,7 @@ gs = pb((l, nothing))[1] # corr_layer A = l -corr_layer = ACEcore.lux(corr1) +corr_layer = Polynomials4ML.lux(corr1) cChain = Chain(; corr_layer) ps, st = setupBFState(MersenneTwister(1234), cChain, Σ) y, st = Lux.apply(cChain, A, ps, st) diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 685a25a..e59e231 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -34,7 +34,7 @@ ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) y, st = Lux.apply(BFwf_chain, X, ps, st) ## Pullback API to capture change in state -(l, st_), pb = pullback(p -> Lux.apply(BFwf_chain, X, p, st), ps) +(l, st_), pb = pullback(x -> Lux.apply(BFwf_chain, x, ps, st), X) gs = pb((one.(l), nothing))[1] # Jastrow: try with gradient @@ -52,6 +52,4 @@ gs = pb((one.(l), nothing))[1] # Zygote.gradient(X -> ACEpsi.evaluate(js, X, Σ), X) -# BackFlowPooling: Try with rrule - From c35759f7fc229e2b7404d06fbe1dc166a7ea9778 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 24 May 2023 10:10:28 +0800 Subject: [PATCH 27/79] add rrule of rnlexample --- src/atomicorbitals/rnlexample.jl | 19 ++++++++++++++++++- test/test_atomicorbitals.jl | 29 +++++++++++++++++++++++++++-- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/src/atomicorbitals/rnlexample.jl b/src/atomicorbitals/rnlexample.jl index 02979ef..0b701c7 100644 --- a/src/atomicorbitals/rnlexample.jl +++ b/src/atomicorbitals/rnlexample.jl @@ -1,7 +1,8 @@ using Polynomials4ML, ForwardDiff import Polynomials4ML: evaluate, evaluate_ed, evaluate_ed2, natural_indices - +using ChainRulesCore +using ChainRulesCore: NoTangent const NLM{T} = NamedTuple{(:n, :l, :m), Tuple{T, T, T}} const NL{T} = NamedTuple{(:n, :l), Tuple{T, T}} @@ -127,3 +128,19 @@ function evaluate_ed2(basis::RnlExample, R) return Rnl, dRnl, ddRnl end + +using LinearAlgebra:dot + +function ChainRulesCore.rrule(::typeof(evaluate), basis::RnlExample, R) + A = evaluate(basis, R) + ∂R = similar(R) + dR = evaluate_ed(basis, R)[2] + function pb(∂A) + @assert size(∂A) == (length(R), length(basis)) + for i = 1:length(R) + ∂R[i] = dot(@view(∂A[i, :]), @view(dR[i, :])) + end + return NoTangent(), NoTangent(), ∂R + end + return A, pb +end \ No newline at end of file diff --git a/test/test_atomicorbitals.jl b/test/test_atomicorbitals.jl index 34a6363..64f0f46 100644 --- a/test/test_atomicorbitals.jl +++ b/test/test_atomicorbitals.jl @@ -1,7 +1,7 @@ using Polynomials4ML, ForwardDiff, Test, ACEpsi -using Polynomials4ML.Testing: println_slim +using Polynomials4ML.Testing: println_slim, print_tf using Polynomials4ML: evaluate, evaluate_ed, evaluate_ed2 # -------------- RnlExample ----------------- @@ -24,4 +24,29 @@ println_slim(@test ddRnl2 ≈ fddRnl) # -------------- **** ----------------- - +using ACEbase.Testing: fdtest +using Zygote + +@info("Test rrule") +using LinearAlgebra: dot + +for ntest = 1:30 + local rr + local uu + local Rnl + local u + + rr = 2 .* randn(10) .- 1 + uu = 2 .* randn(10) .- 1 + _rr(t) = rr + t * uu + Rnl = evaluate(bRnl, rr) + u = randn(size(Rnl)) + F(t) = dot(u, evaluate(bRnl, _rr(t))) + dF(t) = begin + val, pb = Zygote.pullback(evaluate, bRnl, _rr(t)) + ∂BB = pb(u)[2] # pb(u)[1] returns NoTangent() for basis argument + return sum( dot(∂BB[i], uu[i]) for i = 1:length(uu) ) + end + print_tf(@test fdtest(F, dF, 0.0; verbose = false)) +end +println() \ No newline at end of file From e278b2b656f77b8f8595a35a56bcafab7e49eb72 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Fri, 26 May 2023 23:48:59 -0700 Subject: [PATCH 28/79] add gradp test --- Project.toml | 2 +- src/atomicorbitals/atorbbasis.jl | 38 +++++++++++++++++--------------- src/bflow3d.jl | 2 +- src/lux_utils.jl | 5 +++-- test/test_atorbbasis.jl | 6 ++++- test/test_bflow_lux.jl | 23 +++++++++++++++++-- 6 files changed, 51 insertions(+), 25 deletions(-) diff --git a/Project.toml b/Project.toml index af6f4c1..4029611 100644 --- a/Project.toml +++ b/Project.toml @@ -16,6 +16,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d" +Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -30,7 +31,6 @@ BenchmarkTools = "1" ForwardDiff = "0.10" JSON = "0.21" ObjectPools = "0.2.1" -Polynomials4ML = "0.1.3" julia = "1" [extras] diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 077daae..265156e 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -11,6 +11,7 @@ using LinearAlgebra: norm using ChainRulesCore using ChainRulesCore: NoTangent +using Zygote struct Nuc{T} rr::SVector{3, T} @@ -103,26 +104,27 @@ end function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) - nuc = basis.nuclei + nuc = basis.nuclei Nnuc = length(nuc) Nel = size(X, 1) T = promote_type(eltype(X[1])) - VT = SVector{3, T} - XX = zeros(VT, (Nnuc, Nel)) - # trans - for I = 1:Nnuc, i = 1:Nel - XX[I, i] = X[i] - nuc[I].rr - end + # XX = zeros(VT, (Nnuc, Nel)) + + # # trans + # for I = 1:Nnuc, i = 1:Nel + # XX[I, i] = X[i] - nuc[I].rr + # end Nnlm = length(basis.prodbasis.sparsebasis.spec) - ϕnlm = zeros(T, (Nnuc, Nel, Nnlm)) + ϕnlm = Zygote.Buffer(zeros(T, (Nnuc, Nel, Nnlm))) + # Think how to prevent this extra FLOPS here while keeping it Zygote-friendly for I = 1:Nnuc - ϕnlm[I,:,:] = evaluate(basis.prodbasis, XX[I,:]) + ϕnlm[I,:,:] = evaluate(basis.prodbasis, map(x -> x - nuc[I].rr, X)) end - return ϕnlm + return copy(ϕnlm) end # ------------ utils for AtomicOrbitalsBasis ------------ @@ -156,14 +158,14 @@ end # ------------ connect with ChainRulesCore # Placeholder for now, fix this later after making sure Zygote is done correct with Lux -function ChainRulesCore.rrule(::typeof(evaluate), basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) - val = evaluate(basis, X, Σ) - dB = similar(X) - function pb(dA) - return NoTangent(), NoTangent(), dB, NoTangent() - end - return val, pb -end +# function ChainRulesCore.rrule(::typeof(evaluate), basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) +# val = evaluate(basis, X, Σ) +# dB = similar(X) +# function pb(dA) +# return NoTangent(), NoTangent(), dB, NoTangent() +# end +# return val, pb +# end # ------------ connect with Lux struct AtomicOrbitalsBasisLayer{TB} <: AbstractExplicitLayer diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 265150d..c3bec08 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -44,7 +44,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] # define n-correlation - corr1 = Polynomials4ML.SparseSymmProd(spec; T = Float64) + corr1 = Polynomials4ML.SparseSymmProd(spec) # ----------- Lux connections --------- # Should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm for trainable radial basis later? diff --git a/src/lux_utils.jl b/src/lux_utils.jl index 2f94612..1f2e83e 100644 --- a/src/lux_utils.jl +++ b/src/lux_utils.jl @@ -1,10 +1,11 @@ using LuxCore function replace_namedtuples(nt, rp_st, Σ) - if length(nt) == 0 - return rp_st + if :tmp in keys(nt) || length(nt) == 0 + return (; rp_st...) else for i in 1:length(nt) + # @show nt[i] if length(nt[i]) == 0 rp_st = (; rp_st..., (; keys(nt)[i] => (Σ = Σ, ))...) else diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index d6c59c3..7cdeecb 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -140,4 +140,8 @@ chain_st = (aobasis = (Σ = Σ, ), pooling = (Σ = Σ, )) # try Chain is as expected print_tf(@test tryChain(X, chain_ps, chain_st)[1] ≈ A) -println() \ No newline at end of file +println() + +val, pb = Zygote.pullback(evaluate, aobasis, X, Σ) + +pb(val) \ No newline at end of file diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index e59e231..96d40bc 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -6,8 +6,21 @@ using ACEbase.Testing: print_tf using LuxCore using Lux using Zygote +using Optimisers # mainly for the destrcuture(ps) function using Random +function grad_test2(f, df, X::AbstractVector) + F = f(X) + ∇F = df(X) + nX = length(X) + EE = Matrix(I, (nX, nX)) + + for h in 0.1.^(3:12) + gh = [ (f(X + h * EE[:, i]) - F) / h for i = 1:nX ] + @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) + end +end + Rnldegree = 4 Ylmdegree = 4 totdegree = 8 @@ -34,8 +47,14 @@ ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) y, st = Lux.apply(BFwf_chain, X, ps, st) ## Pullback API to capture change in state -(l, st_), pb = pullback(x -> Lux.apply(BFwf_chain, x, ps, st), X) -gs = pb((one.(l), nothing))[1] +gl = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1] + +@info("Test ∇ψ w.r.t. parameters") +W0, re = destructure(ps) +Fp = w -> BFwf_chain(X, re(w), st)[1] +dFp = w -> ( gl = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1]; destructure(gl)[1]) +grad_test2(Fp, dFp, W0) + # Jastrow: try with gradient # using ACEpsi: Jastrow From 061f2dc2cd3f3c24875e44592d5e00646993d4cc Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Mon, 29 May 2023 18:49:51 -0700 Subject: [PATCH 29/79] draft HyperDual --- src/atomicorbitals/rnlexample.jl | 4 ++-- src/lux_utils.jl | 1 + test/test_bflow_lux.jl | 31 +++++++++++++++++++++++++++---- 3 files changed, 30 insertions(+), 6 deletions(-) diff --git a/src/atomicorbitals/rnlexample.jl b/src/atomicorbitals/rnlexample.jl index 0b701c7..9ab46e2 100644 --- a/src/atomicorbitals/rnlexample.jl +++ b/src/atomicorbitals/rnlexample.jl @@ -27,10 +27,10 @@ degree(basis::RnlExample, b::NamedTuple) = b.n + b.l # -------- Evaluation Code -_alloc(basis::RnlExample, r::T) where {T <: Real} = +_alloc(basis::RnlExample, r::T) where T = zeros(T, length(Rnl)) -_alloc(basis::RnlExample, rr::Vector{T}) where {T <: Real} = +_alloc(basis::RnlExample, rr::Vector{T}) where T = zeros(T, length(rr), length(basis)) diff --git a/src/lux_utils.jl b/src/lux_utils.jl index 1f2e83e..78822a3 100644 --- a/src/lux_utils.jl +++ b/src/lux_utils.jl @@ -1,4 +1,5 @@ using LuxCore +using HyperDualNumbers function replace_namedtuples(nt, rp_st, Σ) if :tmp in keys(nt) || length(nt) == 0 diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 96d40bc..a80d8ca 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -6,8 +6,11 @@ using ACEbase.Testing: print_tf using LuxCore using Lux using Zygote -using Optimisers # mainly for the destrcuture(ps) function +using Optimisers # mainly for the destrcuture(ps) functionx using Random +using HyperDualNumbers: Hyper +using Printf +using LinearAlgebra function grad_test2(f, df, X::AbstractVector) F = f(X) @@ -26,6 +29,11 @@ Ylmdegree = 4 totdegree = 8 Nel = 5 X = randn(SVector{3, Float64}, Nel) + +# wrap it as HyperDualNumbers +x2dualwrtj(x, j) = SVector([Hyper(x[i], i == j, i == j, 0) for i = 1:3], 3) + +hX = [x2dualwrtj(x, 1) for x in X] Σ = rand(spins(), Nel) nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] @@ -40,21 +48,36 @@ ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) @info("Test evaluate") A1 = BFwf_chain(X, ps, st) +hA1 = BFwf_chain(hX, ps, st) + +@assert hA1[1].value ≈ A1[1] @info("Test Zygote API") ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) - y, st = Lux.apply(BFwf_chain, X, ps, st) -## Pullback API to capture change in state -gl = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1] +@info("Test ∇ψ w.r.t. X") + + + +## Pullback API to capture change in state +Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1] @info("Test ∇ψ w.r.t. parameters") W0, re = destructure(ps) Fp = w -> BFwf_chain(X, re(w), st)[1] dFp = w -> ( gl = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1]; destructure(gl)[1]) grad_test2(Fp, dFp, W0) +@info("Test Δψ w.r.t. X") + + + +@info("Test gradp Δψ") + + + + # Jastrow: try with gradient # using ACEpsi: Jastrow From 3bbe08fc4fdd9b76cf17d675960e5cfe949f8e3a Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 30 May 2023 01:12:31 -0700 Subject: [PATCH 30/79] minor fix --- src/atomicorbitals/atorbbasis.jl | 27 --------------------------- 1 file changed, 27 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 265156e..4449767 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -57,33 +57,6 @@ function make_nlms_spec(bRnl, bYlm; end -function _invmap(a::AbstractVector) - inva = Dict{eltype(a), Int}() - for i = 1:length(a) - inva[a[i]] = i - end - return inva -end - -function dropnames(namedtuple::NamedTuple, names::Tuple{Vararg{Symbol}}) - keepnames = Base.diff_names(Base._nt_names(namedtuple), names) - return NamedTuple{keepnames}(namedtuple) -end - -function ProductBasis(spec1, bRnl, bYlm) - spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) - spec_Rnl = bRnl.spec; inv_Rnl = _invmap(spec_Rnl) - spec_Ylm = natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) - - spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) - for (i, b) in enumerate(spec1) - spec1idx[i] = (inv_Rnl[dropnames(b,(:m,))], inv_Ylm[(l=b.l, m=b.m)]) - end - sparsebasis = SparseProduct(spec1idx) - return ProductBasis(spec1, bRnl, bYlm, sparsebasis, _make_reqfields()...) -end - - mutable struct AtomicOrbitalsBasis{NB, T} prodbasis::ProductBasis{NB} nuclei::Vector{Nuc{T}} # nuclei (defines the shifted orbitals) From dded5b120648a870dc2f586279f2b9452a25ac02 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 31 May 2023 21:41:49 +0800 Subject: [PATCH 31/79] minor change on test_bflow_lux --- test/test_bflow_lux.jl | 29 ++++++----------------------- 1 file changed, 6 insertions(+), 23 deletions(-) diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index a80d8ca..a202530 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -11,6 +11,7 @@ using Random using HyperDualNumbers: Hyper using Printf using LinearAlgebra +using ACEbase.Testing: fdtest function grad_test2(f, df, X::AbstractVector) F = f(X) @@ -31,7 +32,7 @@ Nel = 5 X = randn(SVector{3, Float64}, Nel) # wrap it as HyperDualNumbers -x2dualwrtj(x, j) = SVector([Hyper(x[i], i == j, i == j, 0) for i = 1:3], 3) +x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) hX = [x2dualwrtj(x, 1) for x in X] Σ = rand(spins(), Nel) @@ -56,9 +57,11 @@ hA1 = BFwf_chain(hX, ps, st) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) y, st = Lux.apply(BFwf_chain, X, ps, st) - +Zygote.gradient(x -> BFwf_chain(x, ps, st)[1], X)[1] @info("Test ∇ψ w.r.t. X") - +F(X) = BFwf_chain(X, ps, st)[1] +dF(X) = Zygote.gradient(x -> BFwf_chain(x, ps, st)[1], X)[1] +fdtest(F, dF, X, verbose = true) ## Pullback API to capture change in state @@ -71,27 +74,7 @@ grad_test2(Fp, dFp, W0) @info("Test Δψ w.r.t. X") - - @info("Test gradp Δψ") - - -# Jastrow: try with gradient -# using ACEpsi: Jastrow -# using Lux -# using Zygote -# using ACEpsi:evaluate - -# js = Jastrow(nuclei) -# jatrow_layer = ACEpsi.lux(js) -# js_chain = Chain(; jatrow_layer) -# ps, st = setupBFState(MersenneTwister(1234), js_chain, Σ) - -# gs = Zygote.gradient(X -> js_chain(X, ps, st)[1], X) -# Zygote.gradient(X -> ACEpsi.evaluate(js, X, Σ), X) - - - From 94f6a5e0276fa0aeb990f1d8d038b22acabcb3c1 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 31 May 2023 22:38:08 +0800 Subject: [PATCH 32/79] minor fix on hyper_real --- src/ACEpsi.jl | 3 ++ src/hyper.jl | 5 ++ test/test_bflow_lux.jl | 108 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 116 insertions(+) create mode 100644 src/hyper.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index 85dbd4b..bcd4eee 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -1,5 +1,8 @@ module ACEpsi +# define operation on HyperDualNumbers +include("hyper.jl") + # define spin symbols and some basic functionality include("spins.jl") diff --git a/src/hyper.jl b/src/hyper.jl new file mode 100644 index 0000000..bb4f9b8 --- /dev/null +++ b/src/hyper.jl @@ -0,0 +1,5 @@ +using HyperDualNumbers: Hyper +using Zygote: Buffer + +Base.real(x::Hyper{<:Number}) = Hyper(real(x.value), real(x.epsilon1), real(x.epsilon2), real(x.epsilon12)) + diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index a202530..a2d466a 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -78,3 +78,111 @@ grad_test2(Fp, dFp, W0) + + + + + + + + + +### test +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState,Jastrow +using ACEbase.Testing: print_tf +using LuxCore +using Lux +using Zygote +using Optimisers # mainly for the destrcuture(ps) functionx +using Random +using HyperDualNumbers: Hyper +using Printf +using LinearAlgebra +using ACEbase.Testing: fdtest + +function grad_test2(f, df, X::AbstractVector) + F = f(X) + ∇F = df(X) + nX = length(X) + EE = Matrix(I, (nX, nX)) + + for h in 0.1.^(3:12) + gh = [ (f(X + h * EE[:, i]) - F) / h for i = 1:nX ] + @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) + end +end + +Rnldegree = 4 +Ylmdegree = 4 +totdegree = 8 +Nel = 5 +X = randn(SVector{3, Float64}, Nel) + +# wrap it as HyperDualNumbers +x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) + +hX = [x2dualwrtj(x, 1) for x in X] +Σ = rand(spins(), Nel) + +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] +## + +# Defining AtomicOrbitalsBasis +bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +bYlm = RYlmBasis(Ylmdegree) + + +totdeg = totdegree +spec1p = make_nlms_spec(bRnl, bYlm; + totaldegree = totdeg) + +aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) +pooling = BackflowPooling(aobasis) + + +tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] +default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) +using ACEcore.Utils: gensparse +ν = 2 +specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, + minvv = fill(0, ν), + maxvv = fill(length(spec1p), ν), + ordered = true) +spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] +corr1 = Polynomials4ML.SparseSymmProd(spec) + +aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) +pooling_layer = ACEpsi.lux(pooling) +corr_layer = Polynomials4ML.lux(corr1) + +js = Jastrow(nuclei) +jastrow_layer = ACEpsi.lux(js) +reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) +bchain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), + bAA = corr_layer, transpose_layer = WrappedFunction(transpose)) + + #, hidden1 = Dense(length(corr1), Nel) + #, Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) + + +ps, st = setupBFState(MersenneTwister(1234), bchain, Σ) + +Zygote.jacobian(p -> bchain(hX, p, st)[1], ps)[1] + +Zygote.jacobian(x -> bchain(x, ps, st)[1], hX)[1] + +# the problem layer + +#bchain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), +# bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel)) + + +#ps, st = setupBFState(MersenneTwister(1234), bchain, Σ) + +#Zygote.jacobian(p -> bchain(hX, p, st)[1], ps)[1] + +#Zygote.jacobian(x -> bchain(x, ps, st)[1], hX)[1] + \ No newline at end of file From b96f307c0dca8b94c2134b4690e8166270ba203f Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 31 May 2023 22:51:59 +0800 Subject: [PATCH 33/79] minor fix --- src/hyper.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/hyper.jl b/src/hyper.jl index bb4f9b8..f74392f 100644 --- a/src/hyper.jl +++ b/src/hyper.jl @@ -1,5 +1,4 @@ using HyperDualNumbers: Hyper -using Zygote: Buffer Base.real(x::Hyper{<:Number}) = Hyper(real(x.value), real(x.epsilon1), real(x.epsilon2), real(x.epsilon12)) From fd4e9d2748534e4e233e91dfb34cd0c29e32c16a Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Wed, 31 May 2023 20:50:05 -0700 Subject: [PATCH 34/79] pass lap test and grad test with hypderdual --- Project.toml | 1 + test/test_bflow_lux.jl | 227 +++++++++++++++++++++++++---------------- 2 files changed, 138 insertions(+), 90 deletions(-) diff --git a/Project.toml b/Project.toml index 4029611..49d0cf4 100644 --- a/Project.toml +++ b/Project.toml @@ -11,6 +11,7 @@ ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +HyperDualNumbers = "50ceba7f-c3ee-5a84-a6e8-3ad40456ec97" JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index a2d466a..5afcc8a 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -6,7 +6,7 @@ using ACEbase.Testing: print_tf using LuxCore using Lux using Zygote -using Optimisers # mainly for the destrcuture(ps) functionx +using Optimisers # mainly for the destrcuture(ps) function using Random using HyperDualNumbers: Hyper using Printf @@ -34,7 +34,9 @@ X = randn(SVector{3, Float64}, Nel) # wrap it as HyperDualNumbers x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) -hX = [x2dualwrtj(x, 1) for x in X] +hX = [x2dualwrtj(x, 0) for x in X] +hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec + Σ = rand(spins(), Nel) nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] @@ -63,6 +65,23 @@ F(X) = BFwf_chain(X, ps, st)[1] dF(X) = Zygote.gradient(x -> BFwf_chain(x, ps, st)[1], X)[1] fdtest(F, dF, X, verbose = true) +@info("Test consistency with HyperDualNumbers") +for _ = 1:30 + local X = randn(SVector{3, Float64}, Nel) + local Σ = rand(spins(), Nel) + local hdF = [zeros(3) for _ = 1:Nel] + local hX = [x2dualwrtj(x, 0) for x in X] + for i = 1:3 + for j = 1:Nel + hX[j] = x2dualwrtj(X[j], i) # ∂Ψ/∂xj_{i} + hdF[j][i] = BFwf_chain(hX, ps, st)[1].epsilon1 + hX[j] = x2dualwrtj(X[j], 0) + end + end + print_tf(@test dF(X) ≈ hdF) +end +println() + ## Pullback API to capture change in state Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1] @@ -72,109 +91,137 @@ Fp = w -> BFwf_chain(X, re(w), st)[1] dFp = w -> ( gl = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1]; destructure(gl)[1]) grad_test2(Fp, dFp, W0) -@info("Test Δψ w.r.t. X") - -@info("Test gradp Δψ") - - - +@info("Test Δψ w.r.t. X using HyperDualNumbers") - - - - - - - -### test -using ACEpsi, Polynomials4ML, StaticArrays, Test -using Polynomials4ML: natural_indices, degree, SparseProduct -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate -using ACEpsi: BackflowPooling, BFwf_lux, setupBFState,Jastrow -using ACEbase.Testing: print_tf -using LuxCore -using Lux -using Zygote -using Optimisers # mainly for the destrcuture(ps) functionx -using Random -using HyperDualNumbers: Hyper -using Printf -using LinearAlgebra -using ACEbase.Testing: fdtest - -function grad_test2(f, df, X::AbstractVector) - F = f(X) - ∇F = df(X) - nX = length(X) - EE = Matrix(I, (nX, nX)) - - for h in 0.1.^(3:12) - gh = [ (f(X + h * EE[:, i]) - F) / h for i = 1:nX ] - @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) +X = [randn(3) for _ = 1:Nel] +hX = [x2dualwrtj(x, 0) for x in X] +Σ = rand(spins(), Nel) +F(x) = BFwf_chain(x, ps, st)[1] +function ΔF(x) + ΔΨ = 0.0 + hX = [x2dualwrtj(xx, 0) for xx in x] + for i = 1:3 + for j = 1:Nel + hX[j] = x2dualwrtj(x[j], i) # ∂Φ/∂xj_{i} + ΔΨ += BFwf_chain(hX, ps, st)[1].epsilon12 + hX[j] = x2dualwrtj(x[j], 0) + end end + return ΔΨ +end +Δ1 = ΔF(X) +f0 = F(X) + +for h in 0.1.^(1:8) + Δfh = 0.0 + for i = 1:Nel + for j = 1:3 + XΔX_add, XΔX_sub = deepcopy(X), deepcopy(X) + XΔX_add[i][j] += h + XΔX_sub[i][j] -= h + Δfh += (F(XΔX_add) - f0) / h^2 + Δfh += (F(XΔX_sub) - f0) / h^2 + end + end + @printf(" %.1e | %.2e \n", h, abs(Δfh - Δ1)) end -Rnldegree = 4 -Ylmdegree = 4 -totdegree = 8 -Nel = 5 -X = randn(SVector{3, Float64}, Nel) - -# wrap it as HyperDualNumbers -x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) - -hX = [x2dualwrtj(x, 1) for x in X] -Σ = rand(spins(), Nel) - -nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] -## - -# Defining AtomicOrbitalsBasis -bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) -bYlm = RYlmBasis(Ylmdegree) - - -totdeg = totdegree -spec1p = make_nlms_spec(bRnl, bYlm; - totaldegree = totdeg) - -aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) -pooling = BackflowPooling(aobasis) +# [ tr( hessian(xx -> gfun(xx).dot.W[i], xx) ) +# for i = 1:length(Δ1.dot.W) ] +@info("Test gradp Δψ") -tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] -default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) -using ACEcore.Utils: gensparse -ν = 2 -specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, - minvv = fill(0, ν), - maxvv = fill(length(spec1p), ν), - ordered = true) -spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] -corr1 = Polynomials4ML.SparseSymmProd(spec) -aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) -pooling_layer = ACEpsi.lux(pooling) -corr_layer = Polynomials4ML.lux(corr1) -js = Jastrow(nuclei) -jastrow_layer = ACEpsi.lux(js) -reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) -bchain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), - bAA = corr_layer, transpose_layer = WrappedFunction(transpose)) +### test +# using ACEpsi, Polynomials4ML, StaticArrays, Test +# using Polynomials4ML: natural_indices, degree, SparseProduct +# using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate +# using ACEpsi: BackflowPooling, BFwf_lux, setupBFState,Jastrow +# using ACEbase.Testing: print_tf +# using LuxCore +# using Lux +# using Zygote +# using Optimisers # mainly for the destrcuture(ps) functionx +# using Random +# using HyperDualNumbers: Hyper +# using Printf +# using LinearAlgebra +# using ACEbase.Testing: fdtest + +# function grad_test2(f, df, X::AbstractVector) +# F = f(X) +# ∇F = df(X) +# nX = length(X) +# EE = Matrix(I, (nX, nX)) + +# for h in 0.1.^(3:12) +# gh = [ (f(X + h * EE[:, i]) - F) / h for i = 1:nX ] +# @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) +# end +# end + +# Rnldegree = 4 +# Ylmdegree = 4 +# totdegree = 8 +# Nel = 5 +# X = randn(SVector{3, Float64}, Nel) + +# # wrap it as HyperDualNumbers +# x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) + +# hX = [x2dualwrtj(x, 1) for x in X] +# Σ = rand(spins(), Nel) + +# nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] +# ## + +# # Defining AtomicOrbitalsBasis +# bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +# bYlm = RYlmBasis(Ylmdegree) + + +# totdeg = totdegree +# spec1p = make_nlms_spec(bRnl, bYlm; +# totaldegree = totdeg) + +# aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) +# pooling = BackflowPooling(aobasis) + + +# tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] +# default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) +# using ACEcore.Utils: gensparse +# ν = 2 +# specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, +# minvv = fill(0, ν), +# maxvv = fill(length(spec1p), ν), +# ordered = true) +# spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] +# corr1 = Polynomials4ML.SparseSymmProd(spec) + +# aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) +# pooling_layer = ACEpsi.lux(pooling) +# corr_layer = Polynomials4ML.lux(corr1) + +# js = Jastrow(nuclei) +# jastrow_layer = ACEpsi.lux(js) +# reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) +# bchain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), +# bAA = corr_layer, transpose_layer = WrappedFunction(transpose)) - #, hidden1 = Dense(length(corr1), Nel) - #, Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) +# #, hidden1 = Dense(length(corr1), Nel) +# #, Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) -ps, st = setupBFState(MersenneTwister(1234), bchain, Σ) +# ps, st = setupBFState(MersenneTwister(1234), bchain, Σ) -Zygote.jacobian(p -> bchain(hX, p, st)[1], ps)[1] +# Zygote.jacobian(p -> bchain(hX, p, st)[1], ps)[1] -Zygote.jacobian(x -> bchain(x, ps, st)[1], hX)[1] +# Zygote.jacobian(x -> bchain(x, ps, st)[1], hX)[1] -# the problem layer +# # the problem layer #bchain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), # bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel)) From ed6a9710209de106dea3246c3ca16c1db842cb86 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Fri, 2 Jun 2023 23:54:53 +0800 Subject: [PATCH 35/79] rewrite dense layer --- src/bflow3d.jl | 30 ++++++++++++-- test/test_for_dense.jl | 90 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 4 deletions(-) create mode 100644 test/test_for_dense.jl diff --git a/src/bflow3d.jl b/src/bflow3d.jl index c3bec08..eb7141c 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -1,14 +1,17 @@ -using ACEcore, Polynomials4ML +using ACEcore, Polynomials4ML, Random using Polynomials4ML: OrthPolyBasis1D3T using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! using ACEcore.Utils: gensparse using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det import ForwardDiff using ACEpsi.AtomicOrbitals: make_nlms_spec - using LuxCore: AbstractExplicitLayer -using Lux: Dense, Chain, WrappedFunction, BranchLayer +using LuxCore +using Lux +using Lux: Chain, WrappedFunction, BranchLayer +using ChainRulesCore +using ChainRulesCore: NoTangent # ---------------------------------------- struct MaskLayer <: AbstractExplicitLayer @@ -17,6 +20,25 @@ end (l::MaskLayer)(Φ, ps, st) = Φ .* [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX], st +struct DenseLayer <: AbstractExplicitLayer + in_dim::Integer + out_dim::Integer +end + +function (l::DenseLayer)(x::AbstractMatrix, ps, st) + return ps.W * x, st +end + +LuxCore.initialparameters(rng::AbstractRNG, l::DenseLayer) = ( W = randn(rng, l.out_dim, l.in_dim), ) +LuxCore.initialstates(rng::AbstractRNG, l::DenseLayer) = NamedTuple() + +function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMatrix, ps, st) + val = l(x, ps, st) + function pb(A) + return NoTangent, NoTangent, NoTangent, A[1] * x', NoTangent + end + return val, pb +end function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, ν = 3, T = Float64, @@ -65,7 +87,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, # 2. Why we need a tranpose here??? Seems that the output from corr_layer is (length(spec), nX)??? # 3. How do we define n-correlations if we use trainable basis? BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), - bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), + bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = DenseLayer(length(corr1), Nel), Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) return Chain(; branch = BranchLayer((js = jastrow_layer, bf = BFwf_chain, )), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ) end diff --git a/test/test_for_dense.jl b/test/test_for_dense.jl new file mode 100644 index 0000000..83a58ac --- /dev/null +++ b/test/test_for_dense.jl @@ -0,0 +1,90 @@ + +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState,Jastrow +using ACEbase.Testing: print_tf +using LuxCore +using Lux +using Zygote +using Optimisers # mainly for the destrcuture(ps) functionx +using Random +using HyperDualNumbers: Hyper +using Printf +using LinearAlgebra +using ACEbase.Testing: fdtest + +Rnldegree = 2 +Ylmdegree = 2 +totdegree = 2 +Nel = 5 +X = randn(SVector{3, Float64}, Nel) + +# wrap it as HyperDualNumbers +x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) + +hX = [x2dualwrtj(x, 1) for x in X] +Σ = rand(spins(), Nel) + +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] +## + +# Defining AtomicOrbitalsBasis +bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +bYlm = RYlmBasis(Ylmdegree) + + +totdeg = totdegree +spec1p = make_nlms_spec(bRnl, bYlm; + totaldegree = totdeg) + +aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) +pooling = BackflowPooling(aobasis) + + +tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] +default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) +using ACEcore.Utils: gensparse +ν = 2 +specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, + minvv = fill(0, ν), + maxvv = fill(length(spec1p), ν), + ordered = true) +spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] +corr1 = Polynomials4ML.SparseSymmProd(spec) + +aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) +pooling_layer = ACEpsi.lux(pooling) +corr_layer = Polynomials4ML.lux(corr1) + +js = Jastrow(nuclei) +jastrow_layer = ACEpsi.lux(js) +reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) +bchain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), + bAA = corr_layer, transpose_layer = WrappedFunction(transpose)) + +ps, st = setupBFState(MersenneTwister(1234), bchain, Σ) + +y = bchain(hX, ps, st)[1] + +b = Chain(; b = ACEpsi.DenseLayer(length(corr1), Nel)) +ps, st = setupBFState(MersenneTwister(1234), b, Σ) + +Zygote.jacobian(p -> b(y, p, st)[1], ps)[1] + + +using ChainRulesCore +import ChainRules: rrule, NoTangent + + +function ChainRulesCore.rrule(::typeof(Lux.apply), l::ACEpsi.DenseLayer, x::AbstractMatrix, ps, st) + val = l(x, ps, st) + function pb(A) + return NoTangent, NoTangent, NoTangent, (W = A[1] * x',), NoTangent + end + return val, pb +end + + +Zygote.refresh() +Zygote.jacobian(p -> b(y, p, st)[1], ps)[1] From f446849deef1ee4fb76dc05b38ffd4eb82465ee5 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Sat, 3 Jun 2023 01:57:12 +0800 Subject: [PATCH 36/79] add gradp_lap --- src/bflow3d.jl | 4 +- src/hyper.jl | 95 +++++++++++++++++++++++++ test/test_bflow_lux.jl | 154 +++++++++++++++-------------------------- test/test_for_dense.jl | 90 ------------------------ 4 files changed, 154 insertions(+), 189 deletions(-) delete mode 100644 test/test_for_dense.jl diff --git a/src/bflow3d.jl b/src/bflow3d.jl index eb7141c..5c5add2 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -32,10 +32,10 @@ end LuxCore.initialparameters(rng::AbstractRNG, l::DenseLayer) = ( W = randn(rng, l.out_dim, l.in_dim), ) LuxCore.initialstates(rng::AbstractRNG, l::DenseLayer) = NamedTuple() -function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMatrix, ps, st) +function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMatrix{Hyper{T}}, ps, st) where T<:Number val = l(x, ps, st) function pb(A) - return NoTangent, NoTangent, NoTangent, A[1] * x', NoTangent + return NoTangent(), NoTangent(), NoTangent(), (W = A[1] * x',), NoTangent() end return val, pb end diff --git a/src/hyper.jl b/src/hyper.jl index f74392f..b0d4d50 100644 --- a/src/hyper.jl +++ b/src/hyper.jl @@ -2,3 +2,98 @@ using HyperDualNumbers: Hyper Base.real(x::Hyper{<:Number}) = Hyper(real(x.value), real(x.epsilon1), real(x.epsilon2), real(x.epsilon12)) +struct NTarr{NTT} + nt::NTT +end + +export array + +array(nt::NamedTuple) = NTarr(nt) + +# ------------------------------ +# 0 + +zero!(a::AbstractArray) = fill!(a, zero(eltype(a))) +zero!(a::Nothing) = nothing + +function zero!(nt::NamedTuple) + for k in keys(nt) + zero!(nt[k]) + end + return nt +end + +Base.zero(nt::NamedTuple) = zero!(deepcopy(nt)) + +Base.zero(nt::NTarr) = NTarr(zero(nt.nt)) + +# ------------------------------ +# + + + +function _add!(a1::AbstractArray, a2::AbstractArray) + a1[:] .= a1[:] .+ a2[:] + return nothing +end + +_add!(at::Nothing, args...) = nothing + +function _add!(nt1::NamedTuple, nt2) + for k in keys(nt1) + _add!(nt1[k], nt2[k]) + end + return nothing +end + +function _add(nt1::NamedTuple, nt2::NamedTuple) + nt = deepcopy(nt1) + _add!(nt, nt2) + return nt +end + +Base.:+(nt1::NTarr, nt2::NTarr) = NTarr(_add(nt1.nt, nt2.nt)) + +# ------------------------------ +# * + +_mul!(::Nothing, args... ) = nothing + +function _mul!(a::AbstractArray, λ::Number) + a[:] .= a[:] .* λ + return nothing +end + +function _mul!(nt::NamedTuple, λ::Number) + for k in keys(nt) + _mul!(nt[k], λ) + end + return nothing +end + +function _mul(nt::NamedTuple, λ::Number) + nt = deepcopy(nt) + _mul!(nt, λ) + return nt +end + +Base.:*(λ::Number, nt::NTarr) = NTarr(_mul(nt.nt, λ)) +Base.:*(nt::NTarr, λ::Number) = NTarr(_mul(nt.nt, λ)) + +# ------------------------------ +# map + +_map!(f, a::AbstractArray) = map!(f, a, a) + +_map!(f, ::Nothing) = nothing + +function _map!(f, nt::NamedTuple) + for k in keys(nt) + _map!(f, nt[k]) + end + return nothing +end + +function Base.map!(f, dest::NTarr, src::NTarr) + _map!(f, nt.nt) + return nt +end diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 5afcc8a..3baac47 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -84,13 +84,23 @@ println() ## Pullback API to capture change in state -Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1] +p = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1] +p, = destructure(p) @info("Test ∇ψ w.r.t. parameters") W0, re = destructure(ps) Fp = w -> BFwf_chain(X, re(w), st)[1] dFp = w -> ( gl = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1]; destructure(gl)[1]) grad_test2(Fp, dFp, W0) +@info("Test consistency with HyperDualNumbers") +hp = Zygote.gradient(p -> BFwf_chain(hX, p, st)[1], ps)[1] +hp, = destructure(hp) +P = similar(p) +for i = 1:length(P) + P[i] = hp[i].value +end + +@assert P ≈ p @info("Test Δψ w.r.t. X using HyperDualNumbers") @@ -131,105 +141,55 @@ end # for i = 1:length(Δ1.dot.W) ] @info("Test gradp Δψ") +g_bchain = xx -> Zygote.gradient(p -> BFwf_chain(xx, p, st)[1], ps)[1] +g_bchain(hX) +using ACEpsi: zero! +using HyperDualNumbers +function grad_lap(g_bchain, x) + function _mapadd!(f, dest::NamedTuple, src::NamedTuple) + for k in keys(dest) + _mapadd!(f, dest[k], src[k]) + end + return nothing + end + _mapadd!(f, dest::Nothing, src) = nothing + _mapadd!(f, dest::AbstractArray, src::AbstractArray) = + map!((s, d) -> d + f(s), dest, src, dest) -### test -# using ACEpsi, Polynomials4ML, StaticArrays, Test -# using Polynomials4ML: natural_indices, degree, SparseProduct -# using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate -# using ACEpsi: BackflowPooling, BFwf_lux, setupBFState,Jastrow -# using ACEbase.Testing: print_tf -# using LuxCore -# using Lux -# using Zygote -# using Optimisers # mainly for the destrcuture(ps) functionx -# using Random -# using HyperDualNumbers: Hyper -# using Printf -# using LinearAlgebra -# using ACEbase.Testing: fdtest - -# function grad_test2(f, df, X::AbstractVector) -# F = f(X) -# ∇F = df(X) -# nX = length(X) -# EE = Matrix(I, (nX, nX)) - -# for h in 0.1.^(3:12) -# gh = [ (f(X + h * EE[:, i]) - F) / h for i = 1:nX ] -# @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) -# end -# end - -# Rnldegree = 4 -# Ylmdegree = 4 -# totdegree = 8 -# Nel = 5 -# X = randn(SVector{3, Float64}, Nel) - -# # wrap it as HyperDualNumbers -# x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) - -# hX = [x2dualwrtj(x, 1) for x in X] -# Σ = rand(spins(), Nel) - -# nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] -# ## - -# # Defining AtomicOrbitalsBasis -# bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) -# bYlm = RYlmBasis(Ylmdegree) - - -# totdeg = totdegree -# spec1p = make_nlms_spec(bRnl, bYlm; -# totaldegree = totdeg) - -# aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) -# pooling = BackflowPooling(aobasis) - - -# tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] -# default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) -# using ACEcore.Utils: gensparse -# ν = 2 -# specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, -# minvv = fill(0, ν), -# maxvv = fill(length(spec1p), ν), -# ordered = true) -# spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] -# corr1 = Polynomials4ML.SparseSymmProd(spec) - -# aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) -# pooling_layer = ACEpsi.lux(pooling) -# corr_layer = Polynomials4ML.lux(corr1) - -# js = Jastrow(nuclei) -# jastrow_layer = ACEpsi.lux(js) -# reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) -# bchain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), -# bAA = corr_layer, transpose_layer = WrappedFunction(transpose)) - -# #, hidden1 = Dense(length(corr1), Nel) -# #, Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) - - -# ps, st = setupBFState(MersenneTwister(1234), bchain, Σ) - -# Zygote.jacobian(p -> bchain(hX, p, st)[1], ps)[1] - -# Zygote.jacobian(x -> bchain(x, ps, st)[1], hX)[1] - -# # the problem layer - -#bchain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), -# bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel)) - + Δ = zero!(g_bchain(x)) + hX = [x2dualwrtj(xx, 0) for xx in x] + for i = 1:3 + for j = 1:length(x) + hX[j] = x2dualwrtj(x[j], i) + _mapadd!(ε₁ε₂part, Δ, g_bchain(hX)) + hX[j] = x2dualwrtj(x[j], 0) + end + end + return Δ +end -#ps, st = setupBFState(MersenneTwister(1234), bchain, Σ) +function ΔF(x, ps) + ΔΨ = 0.0 + hX = [x2dualwrtj(xx, 0) for xx in x] + for i = 1:3 + for j = 1:Nel + hX[j] = x2dualwrtj(x[j], i) # ∂Φ/∂xj_{i} + ΔΨ += BFwf_chain(hX, ps, st)[1].epsilon12 + hX[j] = x2dualwrtj(x[j], 0) + end + end + return ΔΨ +end -#Zygote.jacobian(p -> bchain(hX, p, st)[1], ps)[1] +function ∇ΔF(x, ps) + g_bchain = xx -> Zygote.gradient(p -> BFwf_chain(xx, p, st)[1], ps)[1] + p, = destructure(grad_lap(g_bchain, x)) + return p +end -#Zygote.jacobian(x -> bchain(x, ps, st)[1], hX)[1] - \ No newline at end of file +W0, re = destructure(ps) +Fp = w -> ΔF(X, re(w)) +dFp = w -> ∇ΔF(X, re(w)) +fdtest(Fp, dFp, W0) \ No newline at end of file diff --git a/test/test_for_dense.jl b/test/test_for_dense.jl deleted file mode 100644 index 83a58ac..0000000 --- a/test/test_for_dense.jl +++ /dev/null @@ -1,90 +0,0 @@ - -using ACEpsi, Polynomials4ML, StaticArrays, Test -using Polynomials4ML: natural_indices, degree, SparseProduct -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate -using ACEpsi: BackflowPooling, BFwf_lux, setupBFState,Jastrow -using ACEbase.Testing: print_tf -using LuxCore -using Lux -using Zygote -using Optimisers # mainly for the destrcuture(ps) functionx -using Random -using HyperDualNumbers: Hyper -using Printf -using LinearAlgebra -using ACEbase.Testing: fdtest - -Rnldegree = 2 -Ylmdegree = 2 -totdegree = 2 -Nel = 5 -X = randn(SVector{3, Float64}, Nel) - -# wrap it as HyperDualNumbers -x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) - -hX = [x2dualwrtj(x, 1) for x in X] -Σ = rand(spins(), Nel) - -nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] -## - -# Defining AtomicOrbitalsBasis -bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) -bYlm = RYlmBasis(Ylmdegree) - - -totdeg = totdegree -spec1p = make_nlms_spec(bRnl, bYlm; - totaldegree = totdeg) - -aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) -pooling = BackflowPooling(aobasis) - - -tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] -default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) -using ACEcore.Utils: gensparse -ν = 2 -specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, - minvv = fill(0, ν), - maxvv = fill(length(spec1p), ν), - ordered = true) -spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] -corr1 = Polynomials4ML.SparseSymmProd(spec) - -aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) -pooling_layer = ACEpsi.lux(pooling) -corr_layer = Polynomials4ML.lux(corr1) - -js = Jastrow(nuclei) -jastrow_layer = ACEpsi.lux(js) -reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) -bchain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), - bAA = corr_layer, transpose_layer = WrappedFunction(transpose)) - -ps, st = setupBFState(MersenneTwister(1234), bchain, Σ) - -y = bchain(hX, ps, st)[1] - -b = Chain(; b = ACEpsi.DenseLayer(length(corr1), Nel)) -ps, st = setupBFState(MersenneTwister(1234), b, Σ) - -Zygote.jacobian(p -> b(y, p, st)[1], ps)[1] - - -using ChainRulesCore -import ChainRules: rrule, NoTangent - - -function ChainRulesCore.rrule(::typeof(Lux.apply), l::ACEpsi.DenseLayer, x::AbstractMatrix, ps, st) - val = l(x, ps, st) - function pb(A) - return NoTangent, NoTangent, NoTangent, (W = A[1] * x',), NoTangent - end - return val, pb -end - - -Zygote.refresh() -Zygote.jacobian(p -> b(y, p, st)[1], ps)[1] From b93614705427d68cd9b1b98b0763eb57330654da Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Sat, 3 Jun 2023 08:15:23 -0700 Subject: [PATCH 37/79] test clean ups --- test/test_bflow_lux.jl | 50 ++++++++++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 19 deletions(-) diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 3baac47..4812b53 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -1,17 +1,18 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate -using ACEpsi: BackflowPooling, BFwf_lux, setupBFState,Jastrow -using ACEbase.Testing: print_tf +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow +using ACEbase.Testing: print_tf, fd_test + using LuxCore using Lux using Zygote using Optimisers # mainly for the destrcuture(ps) function using Random -using HyperDualNumbers: Hyper using Printf using LinearAlgebra -using ACEbase.Testing: fdtest + +using HyperDualNumbers: Hyper function grad_test2(f, df, X::AbstractVector) F = f(X) @@ -30,41 +31,46 @@ Ylmdegree = 4 totdegree = 8 Nel = 5 X = randn(SVector{3, Float64}, Nel) +Σ = rand(spins(), Nel) +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] # wrap it as HyperDualNumbers x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) - hX = [x2dualwrtj(x, 0) for x in X] hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec -Σ = rand(spins(), Nel) - -nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] ## # Defining AtomicOrbitalsBasis bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) bYlm = RYlmBasis(Ylmdegree) +# setup state BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) +## + @info("Test evaluate") A1 = BFwf_chain(X, ps, st) hA1 = BFwf_chain(hX, ps, st) -@assert hA1[1].value ≈ A1[1] +print_tf(@test hA1[1].value ≈ A1[1]) + +println() + +## -@info("Test Zygote API") +@info("Test ∇ψ w.r.t. X") ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) y, st = Lux.apply(BFwf_chain, X, ps, st) -Zygote.gradient(x -> BFwf_chain(x, ps, st)[1], X)[1] -@info("Test ∇ψ w.r.t. X") F(X) = BFwf_chain(X, ps, st)[1] dF(X) = Zygote.gradient(x -> BFwf_chain(x, ps, st)[1], X)[1] fdtest(F, dF, X, verbose = true) +## + @info("Test consistency with HyperDualNumbers") for _ = 1:30 local X = randn(SVector{3, Float64}, Nel) @@ -82,17 +88,20 @@ for _ = 1:30 end println() +## -## Pullback API to capture change in state +@info("Test ∇ψ w.r.t. parameters") p = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1] p, = destructure(p) -@info("Test ∇ψ w.r.t. parameters") + W0, re = destructure(ps) Fp = w -> BFwf_chain(X, re(w), st)[1] dFp = w -> ( gl = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1]; destructure(gl)[1]) grad_test2(Fp, dFp, W0) -@info("Test consistency with HyperDualNumbers") +## + +@info("Test consistency when input isa HyperDualNumbers") hp = Zygote.gradient(p -> BFwf_chain(hX, p, st)[1], ps)[1] hp, = destructure(hp) P = similar(p) @@ -100,7 +109,11 @@ for i = 1:length(P) P[i] = hp[i].value end -@assert P ≈ p +print_tf(@test P ≈ p) + +println() + +## @info("Test Δψ w.r.t. X using HyperDualNumbers") @@ -137,10 +150,9 @@ for h in 0.1.^(1:8) @printf(" %.1e | %.2e \n", h, abs(Δfh - Δ1)) end -# [ tr( hessian(xx -> gfun(xx).dot.W[i], xx) ) -# for i = 1:length(Δ1.dot.W) ] +## -@info("Test gradp Δψ") +@info("Test gradp Δψ using HyperDualNumbers") g_bchain = xx -> Zygote.gradient(p -> BFwf_chain(xx, p, st)[1], ps)[1] g_bchain(hX) From 70d030cab777a4fdcfab3ee4e5451fdfc7c85734 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 5 Jun 2023 08:36:05 +0800 Subject: [PATCH 38/79] minor fix --- test/test_bflow_lux.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 4812b53..35fda7f 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -2,8 +2,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow -using ACEbase.Testing: print_tf, fd_test - +using ACEbase.Testing: print_tf, fdtest using LuxCore using Lux using Zygote From ede054a940113cefab038c0e3c5aec672d2dcb9d Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 5 Jun 2023 08:40:55 +0800 Subject: [PATCH 39/79] minor fix --- test/test2.jl | 157 -------------------------------------------------- 1 file changed, 157 deletions(-) delete mode 100644 test/test2.jl diff --git a/test/test2.jl b/test/test2.jl deleted file mode 100644 index c1f39af..0000000 --- a/test/test2.jl +++ /dev/null @@ -1,157 +0,0 @@ -using ACEpsi, Polynomials4ML, StaticArrays, Test -using Polynomials4ML: natural_indices, degree, SparseProduct -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate -using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow -using ACEbase.Testing: print_tf -using LuxCore -using Lux -using Zygote -using Random -using ACEcore.Utils: gensparse -using ACEcore: SparseSymmProd -using ACEcore - -Rnldegree = 4 -Ylmdegree = 4 -totdegree = 8 - -Nel = 5 -X = randn(SVector{3, Float64}, Nel) -Σ = rand(spins(), Nel) -nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] -## - -# Defining AtomicOrbitalsBasis -bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) -bYlm = RYlmBasis(Ylmdegree) -spec1 = make_nlms_spec(bRnl, bYlm; totaldegree = totdegree) - -# define basis and pooling operations -prodbasis = ProductBasis(spec1, bRnl, bYlm) -aobasis = AtomicOrbitalsBasis(prodbasis, nuclei) -pooling = BackflowPooling(aobasis) -ϕnlm = prodbasis(X) -bϕnlm = aobasis(X, Σ) -A = pooling(bϕnlm, Σ) -ν = 3 -totdeg = 6 -tup2b = vv -> [ spec1[v] for v in vv[vv .> 0] ] -default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) - -specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, - minvv = fill(0, ν), - maxvv = fill(length(spec1), ν), - ordered = true) -spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] - -# define n-correlation -corr1 = Polynomials4ML.SparseSymmProd(spec; T = Float64) - - - -# aobasis layer -aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) -aChain = Chain(; aobasis_layer) -ps, st = setupBFState(MersenneTwister(1234), aChain, Σ) -y, st = Lux.apply(aChain, X, ps, st) -(l, st_), pb = pullback(X -> Lux.apply(aChain, X, ps, st), X) -gs = pb((l, nothing))[1] - -# pooling layer -bϕnlm = l -pooling_layer = ACEpsi.lux(pooling) -pChain = Chain(; pooling_layer) -ps, st = setupBFState(MersenneTwister(1234), pChain, Σ) -y, st = Lux.apply(pChain, l, ps, st) -(l, st_), pb = pullback(ϕnlm -> Lux.apply(pChain, ϕnlm, ps, st), bϕnlm) -gs = pb((l, nothing))[1] - -# reshape layer -reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) -reshape_layer = WrappedFunction(reshape_func) -rChain = Chain(; reshape_layer) -ps, st = setupBFState(MersenneTwister(1234), rChain, Σ) -y, st = Lux.apply(rChain, A, ps, st) -(l, st_), pb = pullback(A -> Lux.apply(rChain, A, ps, st), A) -gs = pb((l, nothing))[1] - -# first 3 layer -fChain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func)) -ps, st = setupBFState(MersenneTwister(1234), fChain, Σ) -y, st = Lux.apply(fChain, X, ps, st) -(l, st_), pb = pullback(X -> Lux.apply(fChain, X, ps, st), X) -gs = pb((l, nothing))[1] - -# corr_layer -A = l -corr_layer = Polynomials4ML.lux(corr1) -cChain = Chain(; corr_layer) -ps, st = setupBFState(MersenneTwister(1234), cChain, Σ) -y, st = Lux.apply(cChain, A, ps, st) -(l, st_), pb = pullback(A -> Lux.apply(cChain, A, ps, st), A) -gs = pb((l, nothing))[1] - -corr_l = l -# transpose layer -ϕ = l -transpose_layer = WrappedFunction(transpose) -tChain = Chain(; transpose_layer) -ps, st = setupBFState(MersenneTwister(1234), tChain, Σ) -y, st = Lux.apply(tChain, ϕ, ps, st) -(l, st_), pb = pullback(ϕ -> Lux.apply(tChain, ϕ, ps, st), ϕ) -gs = pb((l, nothing))[1] - -# hidden layer -ϕ = l -hidden_layer = Dense(length(corr1), Nel) -hChain = Chain(; hidden_layer) -ps, st = setupBFState(MersenneTwister(1234), hChain, Σ) -y, st = Lux.apply(hChain, ϕ, ps, st) -(l, st_), pb = pullback(ϕ -> Lux.apply(hChain, ϕ, ps, st), ϕ) -gs = pb((l, nothing))[1] - -# Mask layer -ϕ = l -Mask = ACEpsi.MaskLayer(Nel) -mChain = Chain(; Mask) -ps, st = setupBFState(MersenneTwister(1234), mChain, Σ) -y, st = Lux.apply(mChain, ϕ, ps, st) -(l, st_), pb = pullback(ϕ -> Lux.apply(mChain, ϕ, ps, st), ϕ) -gs = pb((l, nothing))[1] - -# det layer -using LinearAlgebra:det -ϕ = l -det_layer = WrappedFunction(x -> det(x)) -dChain = Chain(; det_layer) -ps, st = setupBFState(MersenneTwister(1234), dChain, Σ) -y, st = Lux.apply(dChain, ϕ, ps, st) -(l, st_), pb = pullback(ϕ -> Lux.apply(dChain, ϕ, ps, st), ϕ) -gs = pb((l, nothing))[1] - -# last 4 layer -using LinearAlgebra:transpose -A = corr_l - -lChain = Chain(; transpose_layer = WrappedFunction(transpose), hidden1 = Dense(length(corr1), Nel), - Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) -ps, st = setupBFState(MersenneTwister(1234), lChain, Σ) -y, st = Lux.apply(lChain, A, ps, st) -(l, st_), pb = pullback(A -> Lux.apply(lChain, A, ps, st), A) -gs = pb((l, nothing))[1] - - - - - -using ACEpsi: Jastrow - -js = Jastrow(nuclei) -jastrow_layer = ACEpsi.lux(js) -jastrow_layer1 = ACEpsi.lux(js) -bChain = Chain(; branch = BranchLayer((js = jastrow_layer, bf = jastrow_layer1, )), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ) - -ps, st = setupBFState(MersenneTwister(1234), bChain, Σ) -y, st = Lux.apply(bChain, X, ps, st) -(l, st_), pb = pullback(X -> Lux.apply(bChain, X, ps, st), X) -gs = pb((l, nothing))[1] \ No newline at end of file From d653b3714c11fe1dfbc3cea21f8761e62c86f4be Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Mon, 5 Jun 2023 01:16:15 -0700 Subject: [PATCH 40/79] replacing ProductBasis and aobasis with lux layer, bflow lux tests passing --- src/atomicorbitals/atomicorbitals.jl | 2 +- src/atomicorbitals/atorbbasis.jl | 213 ++++++++++++++++----------- src/atomicorbitals/productbasis.jl | 40 +++++ src/atomicorbitals/rnlexample.jl | 2 +- src/backflowpooling.jl | 8 +- src/bflow3d.jl | 24 +-- test/test_bflow_lux.jl | 2 +- 7 files changed, 190 insertions(+), 101 deletions(-) create mode 100644 src/atomicorbitals/productbasis.jl diff --git a/src/atomicorbitals/atomicorbitals.jl b/src/atomicorbitals/atomicorbitals.jl index b089ed7..64c5b13 100644 --- a/src/atomicorbitals/atomicorbitals.jl +++ b/src/atomicorbitals/atomicorbitals.jl @@ -4,6 +4,6 @@ import Polynomials4ML: degree include("atorbbasis.jl") include("rnlexample.jl") - +include("productbasis.jl") end \ No newline at end of file diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 4449767..d6b2bc5 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -1,8 +1,8 @@ import Polynomials4ML: evaluate using ACEpsi: ↑, ↓, ∅, spins, extspins, Spin, spin2idx, idx2spin -using Polynomials4ML: SparseProduct, ProductBasis, _make_reqfields -using LuxCore: AbstractExplicitLayer +using Polynomials4ML: SparseProduct, _make_reqfields +using LuxCore: AbstractExplicitLayer, AbstractExplicitContainerLayer using Random: AbstractRNG using Polynomials4ML @@ -13,6 +13,8 @@ using ChainRulesCore using ChainRulesCore: NoTangent using Zygote +using Lux: Chain + struct Nuc{T} rr::SVector{3, T} charge::T # should this be an integer? @@ -57,120 +59,161 @@ function make_nlms_spec(bRnl, bYlm; end -mutable struct AtomicOrbitalsBasis{NB, T} - prodbasis::ProductBasis{NB} - nuclei::Vector{Nuc{T}} # nuclei (defines the shifted orbitals) -end +# mutable struct AtomicOrbitalsBasis{NB, T} +# prodbasis::ProductBasis{NB} +# nuclei::Vector{Nuc{T}} # nuclei (defines the shifted orbitals) +# end -(aobasis::AtomicOrbitalsBasis)(args...) = evaluate(aobasis, args...) -Base.length(aobasis::AtomicOrbitalsBasis) = length(aobasis.prodbasis.spec1) * length(aobasis.nuclei) - -function AtomicOrbitalsBasis(bRnl, bYlm; - totaldegree=3, - nuclei = Nuc{Float64}[], - ) - spec1 = make_nlms_spec(bRnl, bYlm; - totaldegree = totaldegree) - prodbasis = ProductBasis(spec1, bRnl, bYlm) - return AtomicOrbitalsBasis(prodbasis, nuclei) -end +# (aobasis::AtomicOrbitalsBasis)(args...) = evaluate(aobasis, args...) +# Base.length(aobasis::AtomicOrbitalsBasis) = length(aobasis.prodbasis.spec1) * length(aobasis.nuclei) + +# function AtomicOrbitalsBasis(bRnl, bYlm; +# totaldegree=3, +# nuclei = Nuc{Float64}[], +# ) +# spec1 = make_nlms_spec(bRnl, bYlm; +# totaldegree = totaldegree) +# prodbasis = ProductBasis(spec1, bRnl, bYlm) +# return AtomicOrbitalsBasis(prodbasis, nuclei) +# end -function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) - nuc = basis.nuclei - Nnuc = length(nuc) - Nel = size(X, 1) - T = promote_type(eltype(X[1])) +# function evaluate(basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) +# nuc = basis.nuclei +# Nnuc = length(nuc) +# Nel = size(X, 1) +# T = promote_type(eltype(X[1])) - # XX = zeros(VT, (Nnuc, Nel)) +# # XX = zeros(VT, (Nnuc, Nel)) - # # trans - # for I = 1:Nnuc, i = 1:Nel - # XX[I, i] = X[i] - nuc[I].rr - # end +# # # trans +# # for I = 1:Nnuc, i = 1:Nel +# # XX[I, i] = X[i] - nuc[I].rr +# # end - Nnlm = length(basis.prodbasis.sparsebasis.spec) - ϕnlm = Zygote.Buffer(zeros(T, (Nnuc, Nel, Nnlm))) +# Nnlm = length(basis.prodbasis.sparsebasis.spec) +# ϕnlm = Zygote.Buffer(zeros(T, (Nnuc, Nel, Nnlm))) - # Think how to prevent this extra FLOPS here while keeping it Zygote-friendly - for I = 1:Nnuc - ϕnlm[I,:,:] = evaluate(basis.prodbasis, map(x -> x - nuc[I].rr, X)) - end +# # Think how to prevent this extra FLOPS here while keeping it Zygote-friendly +# for I = 1:Nnuc +# ϕnlm[I,:,:] = evaluate(basis.prodbasis, map(x -> x - nuc[I].rr, X)) +# end - return copy(ϕnlm) -end +# return copy(ϕnlm) +# end -# ------------ utils for AtomicOrbitalsBasis ------------ -function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) - basis.nuclei = copy(collect(nuclei)) - return nothing -end +# # ------------ utils for AtomicOrbitalsBasis ------------ +# function set_nuclei!(basis::AtomicOrbitalsBasis, nuclei::AbstractVector{<: Nuc}) +# basis.nuclei = copy(collect(nuclei)) +# return nothing +# end -function get_spec(basis::AtomicOrbitalsBasis) - spec = [] - Nnuc = length(basis.nuclei) +# function get_spec(basis::AtomicOrbitalsBasis) +# spec = [] +# Nnuc = length(basis.nuclei) - spec = Array{Any}(undef, (3, Nnuc, length(basis.prodbasis.spec1))) +# spec = Array{Any}(undef, (3, Nnuc, length(basis.prodbasis.spec1))) - for (k, nlm) in enumerate(basis.prodbasis.spec1) - for I = 1:Nnuc - for (is, s) in enumerate(extspins()) - spec[is, I, k] = (I = I, s=s, nlm...) - end - end - end +# for (k, nlm) in enumerate(basis.prodbasis.spec1) +# for I = 1:Nnuc +# for (is, s) in enumerate(extspins()) +# spec[is, I, k] = (I = I, s=s, nlm...) +# end +# end +# end - return spec -end +# return spec +# end -# ------------ Evaluation kernels +# # ------------ Evaluation kernels -# ------------ connect with ChainRulesCore +# # ------------ connect with ChainRulesCore -# Placeholder for now, fix this later after making sure Zygote is done correct with Lux -# function ChainRulesCore.rrule(::typeof(evaluate), basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) -# val = evaluate(basis, X, Σ) -# dB = similar(X) -# function pb(dA) -# return NoTangent(), NoTangent(), dB, NoTangent() -# end -# return val, pb +# # Placeholder for now, fix this later after making sure Zygote is done correct with Lux +# # function ChainRulesCore.rrule(::typeof(evaluate), basis::AtomicOrbitalsBasis, X::AbstractVector{<: AbstractVector}, Σ) +# # val = evaluate(basis, X, Σ) +# # dB = similar(X) +# # function pb(dA) +# # return NoTangent(), NoTangent(), dB, NoTangent() +# # end +# # return val, pb +# # end + +# # ------------ connect with Lux +# struct AtomicOrbitalsBasisLayer{TB} <: AbstractExplicitLayer +# basis::TB +# meta::Dict{String, Any} # end -# ------------ connect with Lux -struct AtomicOrbitalsBasisLayer{TB} <: AbstractExplicitLayer - basis::TB +# Base.length(l::AtomicOrbitalsBasisLayer) = length(l.basis) + +# lux(basis::AtomicOrbitalsBasis) = AtomicOrbitalsBasisLayer(basis, Dict{String, Any}()) + +# initialparameters(rng::AbstractRNG, l::AtomicOrbitalsBasisLayer) = _init_luxparams(rng, l.basis) + +# initialstates(rng::AbstractRNG, l::AtomicOrbitalsBasisLayer) = _init_luxstate(rng, l.basis) + +# (l::AtomicOrbitalsBasisLayer)(X, ps, st) = +# evaluate(l.basis, X, st.Σ), st + + +# This can be done using ObjectPools, but for simplicity I didn't do that for now since I +# don't want lux layers storing ObjectPools stuffs + +struct AtomicOrbitalsBasisLayer{L, T} <: AbstractExplicitContainerLayer{(:prodbasis, )} + prodbasis::L + nuclei::Vector{Nuc{T}} meta::Dict{String, Any} end -Base.length(l::AtomicOrbitalsBasisLayer) = length(l.basis) +Base.length(l::AtomicOrbitalsBasisLayer) = length(l.prodbasis.layer.ϕnlms.basis.spec) * length(l.nuclei) + +function AtomicOrbitalsBasisLayer(prodbasis, nuclei) + return AtomicOrbitalsBasisLayer(prodbasis, nuclei, Dict{String, Any}()) +end -lux(basis::AtomicOrbitalsBasis) = AtomicOrbitalsBasisLayer(basis, Dict{String, Any}()) +function evaluate(l::AtomicOrbitalsBasisLayer, X, ps, st) + nuc = l.nuclei + Nnuc = length(nuc) + Nel = size(X, 1) + T = promote_type(eltype(X[1])) + Nnlm = length(l.prodbasis.layers.ϕnlms.basis.spec) -initialparameters(rng::AbstractRNG, l::AtomicOrbitalsBasisLayer) = _init_luxparams(rng, l.basis) + ϕnlm = Zygote.Buffer(zeros(T, (Nnuc, Nel, Nnlm))) + for I = 1:Nnuc + ϕnlm[I,:,:], _ = l.prodbasis(map(x -> x - nuc[I].rr, X), ps, st) + end + + return copy(ϕnlm), st +end -initialstates(rng::AbstractRNG, l::AtomicOrbitalsBasisLayer) = _init_luxstate(rng, l.basis) (l::AtomicOrbitalsBasisLayer)(X, ps, st) = - evaluate(l.basis, X, st.Σ), st + evaluate(l, X, ps, st) +# ------------ utils for AtomicOrbitalsBasis ------------ +function set_nuclei!(basis::AtomicOrbitalsBasisLayer, nuclei::AbstractVector{<: Nuc}) + basis.nuclei = copy(collect(nuclei)) + return nothing +end - -# ----- ObejctPools -# (l::AtomicOrbitalsBasisLayer)(args...) = evaluate(l, args...) +function get_spec(l::AtomicOrbitalsBasisLayer) + spec = [] + Nnuc = length(l.nuclei) -# function evaluate(l::AtomicOrbitalsBasisLayer, x::SINGLE, ps, st) -# B = acquire!(st.cache, :B, (length(l.basis), ), _valtype(l.basis, x)) -# evaluate!(parent(B), l.basis, x) -# return B -# end + spec = Array{Any}(undef, (3, Nnuc, length(l.prodbasis.layers.ϕnlms.basis.spec))) -# function evaluate(l::AtomicOrbitalsBasisLayer, X::AbstractArray{<: SINGLE}, ps, st) -# B = acquire!(st.cache[:Bbatch], (length(l.basis), length(X)), _valtype(l.basis, X[1])) -# evaluate!(parent(B), l.basis, X) -# return B -# end + for (k, nlm) in enumerate(l.prodbasis.layers.ϕnlms.basis.spec) + for I = 1:Nnuc + for (is, s) in enumerate(extspins()) + spec[is, I, k] = (I = I, s=s, nlm...) + end + end + end + + return spec +end \ No newline at end of file diff --git a/src/atomicorbitals/productbasis.jl b/src/atomicorbitals/productbasis.jl new file mode 100644 index 0000000..0f98337 --- /dev/null +++ b/src/atomicorbitals/productbasis.jl @@ -0,0 +1,40 @@ +using Lux: WrappedFunction +using Lux +using Polynomials4ML: SparseProduct + +function _invmap(a::AbstractVector) + inva = Dict{eltype(a), Int}() + for i = 1:length(a) + inva[a[i]] = i + end + return inva +end + +function dropnames(namedtuple::NamedTuple, names::Tuple{Vararg{Symbol}}) + keepnames = Base.diff_names(Base._nt_names(namedtuple), names) + return NamedTuple{keepnames}(namedtuple) +end + +function ProductBasisLayer(spec1, bRnl, bYlm) + spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) + spec_Rnl = natural_indices(bRnl); inv_Rnl = _invmap(spec_Rnl) + spec_Ylm = natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) + + spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) + for (i, b) in enumerate(spec1) + spec1idx[i] = (inv_Rnl[dropnames(b,(:m,))], inv_Ylm[(l=b.l, m=b.m)]) + end + sparsebasis = SparseProduct(spec1idx) + + # wrap into lux layers + l_Rn = Polynomials4ML.lux(bRnl) + l_Ylm = Polynomials4ML.lux(bYlm) + l_ϕnlm = Polynomials4ML.lux(sparsebasis) + + # formming model with Lux Chain + _norm(x) = norm.(x) + + l_xnx = Lux.Parallel(nothing; normx = WrappedFunction(_norm), x = WrappedFunction(identity)) + l_embed = Lux.Parallel(nothing; Rn = l_Rn, Ylm = l_Ylm) + return Chain(; xnx = l_xnx, embed = l_embed, ϕnlms = l_ϕnlm) + end \ No newline at end of file diff --git a/src/atomicorbitals/rnlexample.jl b/src/atomicorbitals/rnlexample.jl index 9ab46e2..e9a3d68 100644 --- a/src/atomicorbitals/rnlexample.jl +++ b/src/atomicorbitals/rnlexample.jl @@ -6,7 +6,7 @@ using ChainRulesCore: NoTangent const NLM{T} = NamedTuple{(:n, :l, :m), Tuple{T, T, T}} const NL{T} = NamedTuple{(:n, :l), Tuple{T, T}} -struct RnlExample{TP, TI} +struct RnlExample{TP, TI} <: Polynomials4ML.AbstractPoly4MLBasis Pn::TP spec::Vector{NL{TI}} end diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index 7b6f7b1..b79949f 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -1,11 +1,11 @@ -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis +using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasisLayer using LuxCore: AbstractExplicitLayer using Random: AbstractRNG using ChainRulesCore using ChainRulesCore: NoTangent mutable struct BackflowPooling - basis::AtomicOrbitalsBasis + basis::AtomicOrbitalsBasisLayer end (pooling::BackflowPooling)(args...) = evaluate(pooling, args...) @@ -16,7 +16,7 @@ function evaluate(pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) basis = pooling.basis nuc = basis.nuclei Nnuc = length(nuc) - Nnlm = length(basis.prodbasis.sparsebasis.spec) + Nnlm = length(basis.prodbasis.layers.ϕnlms.basis.spec) Nel = length(Σ) T = promote_type(eltype(ϕnlm)) @@ -95,7 +95,7 @@ function _pullback_evaluate!(∂ϕnlm, ∂A, pooling::BackflowPooling, ϕnlm, Σ Nnuc, Nel, Nnlm = size(ϕnlm) basis = pooling.basis - @assert Nnlm == length(basis.prodbasis.spec1) + @assert Nnlm == length(basis.prodbasis.layers.ϕnlms.basis.spec) @assert Nel == length(Σ) @assert size(∂ϕnlm) == (Nnuc, Nel, Nnlm) @assert size(∂A) == (Nel, 3, Nnuc, Nnlm) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 5c5add2..4ded405 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -14,12 +14,15 @@ using ChainRulesCore using ChainRulesCore: NoTangent # ---------------------------------------- +# ----------------- custom layers ------------------ struct MaskLayer <: AbstractExplicitLayer nX::Integer end (l::MaskLayer)(Φ, ps, st) = Φ .* [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX], st +## + struct DenseLayer <: AbstractExplicitLayer in_dim::Integer out_dim::Integer @@ -29,6 +32,7 @@ function (l::DenseLayer)(x::AbstractMatrix, ps, st) return ps.W * x, st end +# Jerry: Maybe we should use Glorot Uniform if we have no idea about what we should use? LuxCore.initialparameters(rng::AbstractRNG, l::DenseLayer) = ( W = randn(rng, l.out_dim, l.in_dim), ) LuxCore.initialstates(rng::AbstractRNG, l::DenseLayer) = NamedTuple() @@ -40,6 +44,10 @@ function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMat return val, pb end +## + +# ----------------- custom layers ------------------ + function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, ν = 3, T = Float64, sd_admissible = bb -> (true), @@ -49,8 +57,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, totaldegree = totdeg) # size(X) = (nX, 3); length(Σ) = nX - aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) - pooling = BackflowPooling(aobasis) + # aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) # define sparse for n-correlations tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] @@ -69,23 +76,22 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, corr1 = Polynomials4ML.SparseSymmProd(spec) # ----------- Lux connections --------- - # Should we break down the aobasis again into x -> (norm(x), x) -> (Rln, Ylm) -> ϕnlm for trainable radial basis later? # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1)) - aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) + prodbasis_layer = ACEpsi.AtomicOrbitals.ProductBasisLayer(spec1p, bRnl, bYlm) + aobasis_layer = ACEpsi.AtomicOrbitals.AtomicOrbitalsBasisLayer(prodbasis_layer, nuclei) + # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) + pooling = BackflowPooling(aobasis_layer) pooling_layer = ACEpsi.lux(pooling) + # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) corr_layer = Polynomials4ML.lux(corr1) - # TODO: Add J-factor and add trainable basis later js = Jastrow(nuclei) jastrow_layer = ACEpsi.lux(js) reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) - # Questions to discuss: - # 1. it seems that we do not need trans since the bases have already taken care of it? - # 2. Why we need a tranpose here??? Seems that the output from corr_layer is (length(spec), nX)??? - # 3. How do we define n-correlations if we use trainable basis? + BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = DenseLayer(length(corr1), Nel), Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 35fda7f..49fea3e 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -1,6 +1,6 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate +using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow using ACEbase.Testing: print_tf, fdtest using LuxCore From d1fd930d748c7d33bafd599d771b60325297847d Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Mon, 5 Jun 2023 01:42:00 -0700 Subject: [PATCH 41/79] fix atorbbasis test --- src/atomicorbitals/atorbbasis.jl | 6 +-- test/test_atorbbasis.jl | 68 ++++++++------------------------ 2 files changed, 19 insertions(+), 55 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index d6b2bc5..1083e9c 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -201,13 +201,13 @@ function set_nuclei!(basis::AtomicOrbitalsBasisLayer, nuclei::AbstractVector{<: end -function get_spec(l::AtomicOrbitalsBasisLayer) +function get_spec(l::AtomicOrbitalsBasisLayer, spec1p) spec = [] Nnuc = length(l.nuclei) - spec = Array{Any}(undef, (3, Nnuc, length(l.prodbasis.layers.ϕnlms.basis.spec))) + spec = Array{Any}(undef, (3, Nnuc, length(spec1p))) - for (k, nlm) in enumerate(l.prodbasis.layers.ϕnlms.basis.spec) + for (k, nlm) in enumerate(spec1p) for I = 1:Nnuc for (is, s) in enumerate(extspins()) spec[is, I, k] = (I = I, s=s, nlm...) diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 7cdeecb..b74912d 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -1,12 +1,14 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate +using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate, AtomicOrbitalsBasisLayer +using ACEpsi: extspins using ACEpsi: BackflowPooling using ACEbase.Testing: print_tf, fdtest using LuxCore using Random using Zygote + # test configs Rnldegree = 4 Ylmdegree = 4 @@ -21,30 +23,20 @@ nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] # Defining AtomicOrbitalsBasis bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) bYlm = RYlmBasis(Ylmdegree) -spec1 = make_nlms_spec(bRnl, bYlm; totaldegree = totdegree) +spec1p = make_nlms_spec(bRnl, bYlm; totaldegree = totdegree) # define basis and pooling operations -prodbasis = ProductBasis(spec1, bRnl, bYlm) -aobasis = AtomicOrbitalsBasis(prodbasis, nuclei) -pooling = BackflowPooling(aobasis) - -# we can also construct in this way which wraps the definition of product basis inside -aobasis2 = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdegree, nuclei = nuclei, ) +prodbasis_layer = ACEpsi.AtomicOrbitals.ProductBasisLayer(spec1p, bRnl, bYlm) +aobasis_layer = ACEpsi.AtomicOrbitals.AtomicOrbitalsBasisLayer(prodbasis_layer, nuclei) -@info("Checking two type of construction are same") -for ntest = 1:30 - local X = randn(SVector{3, Float64}, Nel) - local Σ = rand(spins(), Nel) - print_tf(@test evaluate(aobasis, X, Σ) ≈ evaluate(aobasis2, X, Σ)) -end +pooling = BackflowPooling(aobasis_layer) +pooling_layer = ACEpsi.lux(pooling) println() -@info("Test evaluate ProductBasis") -ϕnlm = prodbasis(X) - @info("Test evaluate AtomicOrbitalsBasis") -bϕnlm = aobasis(X, Σ) +ps, st = LuxCore.setup(MersenneTwister(1234), aobasis_layer) +bϕnlm, st = aobasis_layer(X, ps, st) @info("Test BackflowPooling") A = pooling(bϕnlm, Σ) @@ -54,22 +46,22 @@ println() ## @info("Check get_spec is working") -spec = ACEpsi.AtomicOrbitals.get_spec(aobasis) +spec = ACEpsi.AtomicOrbitals.get_spec(aobasis_layer, spec1p) @info("Test evaluation by manual construction") using LinearAlgebra: norm bYlm_ = RYlmBasis(totdegree) -Nnlm = length(aobasis.prodbasis.sparsebasis.spec) -Nnuc = length(aobasis.nuclei) +Nnlm = length(aobasis_layer.prodbasis.layers.ϕnlms.basis.spec) +Nnuc = length(aobasis_layer.nuclei) for I = 1:Nnuc - XI = X .- Ref(aobasis.nuclei[I].rr) + XI = X .- Ref(aobasis_layer.nuclei[I].rr) xI = norm.(XI) Rnl = evaluate(bRnl, xI) Ylm = evaluate(bYlm_, XI) for k = 1:Nnlm - nlm = aobasis.prodbasis.sparsebasis.spec[k] + nlm = aobasis_layer.prodbasis.layers.ϕnlms.basis.spec[k] iR = nlm[1] iY = nlm[2] @@ -116,32 +108,4 @@ println() val, pb = Zygote.pullback(pooling, bϕnlm, Σ) val1, pb1 = ACEpsi._rrule_evaluate(pooling, bϕnlm, Σ) @assert val1 ≈ val1 -@assert pb1(val) ≈ pb(val)[1] # pb(val)[2] is for Σ with no pb - - -@info("---------- Lux tests ----------") -aobasis_layer = ACEpsi.AtomicOrbitals.lux(aobasis) -pooling_layer = ACEpsi.lux(pooling) - -# set up null states -ps, st = LuxCore.setup(MersenneTwister(1234), aobasis_layer) -st1 = (Σ = Σ, ) - -@info("Checking layers working fine seperately") -print_tf(@test aobasis_layer(X, ps, st1)[1] ≈ bϕnlm) -print_tf(@test pooling_layer(bϕnlm, ps, st1)[1] ≈ A) -println() - -@info("Checking Chain towards A basis") -using Lux: Chain -tryChain = Chain(; aobasis = aobasis_layer, pooling = pooling_layer) -chain_ps, chain_st = LuxCore.setup(MersenneTwister(1234), tryChain) -chain_st = (aobasis = (Σ = Σ, ), pooling = (Σ = Σ, )) -# try Chain is as expected -print_tf(@test tryChain(X, chain_ps, chain_st)[1] ≈ A) - -println() - -val, pb = Zygote.pullback(evaluate, aobasis, X, Σ) - -pb(val) \ No newline at end of file +@assert pb1(val) ≈ pb(val)[1] # pb(val)[2] is for Σ with no pb \ No newline at end of file From 52a12014356fa1c1350f1267b21d1df7e614904d Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Wed, 7 Jun 2023 07:55:29 -0700 Subject: [PATCH 42/79] fix type instability in bfpooling, more performance fix soon... --- src/atomicorbitals/productbasis.jl | 4 +- src/backflowpooling.jl | 70 ++++++++++++----------- test/test_atorbbasis.jl | 5 +- test/test_bflow_lux.jl | 17 ++++++ test/test_matmul.jl | 92 ++++++++++++++++++++++++++++++ 5 files changed, 153 insertions(+), 35 deletions(-) create mode 100644 test/test_matmul.jl diff --git a/src/atomicorbitals/productbasis.jl b/src/atomicorbitals/productbasis.jl index 0f98337..4073913 100644 --- a/src/atomicorbitals/productbasis.jl +++ b/src/atomicorbitals/productbasis.jl @@ -1,6 +1,6 @@ using Lux: WrappedFunction using Lux -using Polynomials4ML: SparseProduct +using Polynomials4ML: SparseProduct, AbstractPoly4MLBasis function _invmap(a::AbstractVector) inva = Dict{eltype(a), Int}() @@ -15,7 +15,7 @@ function dropnames(namedtuple::NamedTuple, names::Tuple{Vararg{Symbol}}) return NamedTuple{keepnames}(namedtuple) end -function ProductBasisLayer(spec1, bRnl, bYlm) +function ProductBasisLayer(spec1::Vector, bRnl::AbstractPoly4MLBasis, bYlm::AbstractPoly4MLBasis) spec1idx = Vector{Tuple{Int, Int}}(undef, length(spec1)) spec_Rnl = natural_indices(bRnl); inv_Rnl = _invmap(spec_Rnl) spec_Ylm = natural_indices(bYlm); inv_Ylm = _invmap(spec_Ylm) diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index b79949f..f613954 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -3,35 +3,43 @@ using LuxCore: AbstractExplicitLayer using Random: AbstractRNG using ChainRulesCore using ChainRulesCore: NoTangent +using Polynomials4ML: _make_reqfields, @reqfields, POOL, TMP, META +using ObjectPools: acquire! mutable struct BackflowPooling basis::AtomicOrbitalsBasisLayer + @reqfields +end + +function BackflowPooling(basis::AtomicOrbitalsBasisLayer) + return BackflowPooling(basis, _make_reqfields()...) end (pooling::BackflowPooling)(args...) = evaluate(pooling, args...) Base.length(pooling::BackflowPooling) = 3 * length(pooling.basis) # length(spin()) * length(1pbasis) -function evaluate(pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) - basis = pooling.basis - nuc = basis.nuclei - Nnuc = length(nuc) - Nnlm = length(basis.prodbasis.layers.ϕnlms.basis.spec) +function evaluate(pooling::BackflowPooling, ϕnlm::AbstractArray, Σ::AbstractVector) + Nnuc, _, Nnlm = size(ϕnlm) Nel = length(Σ) T = promote_type(eltype(ϕnlm)) # evaluate the pooling operation # spin I k = (nlm) - Aall = zeros(T, (2, Nnuc, Nnlm)) - for k = 1:Nnlm - for i = 1:Nel - iσ = spin2idx(Σ[i]) + Aall = acquire!(pooling.pool, :Aall, (2, Nnuc, Nnlm), T) + fill!(Aall, 0) + + @inbounds begin + for k = 1:Nnlm for I = 1:Nnuc - Aall[iσ, I, k] += ϕnlm[I, i, k] + @simd ivdep for i = 1:Nel + iσ = spin2idx(Σ[i]) + Aall[iσ, I, k] += ϕnlm[I, i, k] + end end end - end + end # inbounds # now correct the pooling Aall and write into A^(i) # with do it with i leading so that the N-correlations can @@ -51,22 +59,29 @@ function evaluate(pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) @assert spin2idx(↑) == 1 @assert spin2idx(↓) == 2 @assert spin2idx(∅) == 3 - A = zeros(T, ((Nel, 3, Nnuc, Nnlm))) - for k = 1:Nnlm - for I = 1:Nnuc - for i = 1:Nel - A[i, 3, I, k] = ϕnlm[I, i, k] - end - for iσ = 1:2 - σ = idx2spin(iσ) - for i = 1:Nel - A[i, iσ, I, k] = Aall[iσ, I, k] - (Σ[i] == σ) * ϕnlm[I, i, k] + + A = acquire!(pooling.tmp, :Aall, (Nel, 3, Nnuc, Nnlm), T) + fill!(A, 0) + + @inbounds begin + for k = 1:Nnlm + for I = 1:Nnuc + @simd ivdep for i = 1:Nel + A[i, 3, I, k] = ϕnlm[I, i, k] + end + @simd ivdep for iσ = 1:2 + σ = idx2spin(iσ) + for i = 1:Nel + A[i, iσ, I, k] = Aall[iσ, I, k] - (Σ[i] == σ) * ϕnlm[I, i, k] + end end end end - end + end # inbounds - return A + release!(Aall) + + return A end # --------------------- connect with ChainRule @@ -131,12 +146,3 @@ initialstates(rng::AbstractRNG, l::BackflowPoolingLayer) = _init_luxstate(rng, l (l::BackflowPoolingLayer)(ϕnlm, ps, st) = evaluate(l.basis, ϕnlm, st.Σ), st -# ----- ObejctPools -# (l::BackflowPoolingLayer)(args...) = evaluate(l, args...) - -# function evaluate(l::BackflowPoolingLayer, ϕnlm_Σ::SINGLE, ps, st) -# B = acquire!(st.cache, :B, (length(l.basis), ), _valtype(l.basis, x)) -# evaluate!(parent(B), l.basis, ϕnlm_Σ[1], ϕnlm_Σ[2]) -# return B -# end - diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index b74912d..58ee23a 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -7,7 +7,7 @@ using ACEbase.Testing: print_tf, fdtest using LuxCore using Random using Zygote - +using Cthulhu # test configs Rnldegree = 4 @@ -33,6 +33,9 @@ pooling = BackflowPooling(aobasis_layer) pooling_layer = ACEpsi.lux(pooling) println() +@info("Test evaluate ProductBasisLayer") +ps1, st1 = LuxCore.setup(MersenneTwister(1234), prodbasis_layer) +bϕnlm, st1 = prodbasis_layer(X, ps1, st1) @info("Test evaluate AtomicOrbitalsBasis") ps, st = LuxCore.setup(MersenneTwister(1234), aobasis_layer) diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 49fea3e..dee7d79 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -10,6 +10,8 @@ using Optimisers # mainly for the destrcuture(ps) function using Random using Printf using LinearAlgebra +using BenchmarkTools +using Cthulhu using HyperDualNumbers: Hyper @@ -58,6 +60,18 @@ print_tf(@test hA1[1].value ≈ A1[1]) println() +## +F(X) = BFwf_chain(X, ps, st)[1] + +# @profview let F = F, X = X +# for i = 1:10_000 +# F(X) +# end +# end + +@btime F(X) + + ## @info("Test ∇ψ w.r.t. X") @@ -120,6 +134,8 @@ X = [randn(3) for _ = 1:Nel] hX = [x2dualwrtj(x, 0) for x in X] Σ = rand(spins(), Nel) F(x) = BFwf_chain(x, ps, st)[1] + + function ΔF(x) ΔΨ = 0.0 hX = [x2dualwrtj(xx, 0) for xx in x] @@ -132,6 +148,7 @@ function ΔF(x) end return ΔΨ end + Δ1 = ΔF(X) f0 = F(X) diff --git a/test/test_matmul.jl b/test/test_matmul.jl new file mode 100644 index 0000000..ffa3473 --- /dev/null +++ b/test/test_matmul.jl @@ -0,0 +1,92 @@ +using LinearAlgebra +using BenchmarkTools +using ACEpsi, Polynomials4ML +using ACEpsi: setupBFState, Nuc, BFwf_lux +using ACEpsi.AtomicOrbitals: make_nlms_spec +using Polynomials4ML.Utils: gensparse +using Lux +using Random +using StaticArrays + +function BFwf_lux_AA(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, + ν = 3, T = Float64, + sd_admissible = bb -> (true), + envelope = x -> x) # enveolpe to be replaced by SJ-factor + + spec1p = make_nlms_spec(bRnl, bYlm; + totaldegree = totdeg) + + # size(X) = (nX, 3); length(Σ) = nX + # aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) + + # define sparse for n-correlations + tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] + default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) + + specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, + minvv = fill(0, ν), + maxvv = fill(length(spec1p), ν), + ordered = true) + spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] + + # further restrict + spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] + + # define n-correlation + corr1 = Polynomials4ML.SparseSymmProd(spec) + + # ----------- Lux connections --------- + # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1)) + prodbasis_layer = ACEpsi.AtomicOrbitals.ProductBasisLayer(spec1p, bRnl, bYlm) + aobasis_layer = ACEpsi.AtomicOrbitals.AtomicOrbitalsBasisLayer(prodbasis_layer, nuclei) + + # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) + pooling = ACEpsi.BackflowPooling(aobasis_layer) + pooling_layer = ACEpsi.lux(pooling) + + reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) + + # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) + corr_layer = Polynomials4ML.lux(corr1) + return Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), + bAA = corr_layer, transpose_layer = WrappedFunction(transpose)) +end + +Rnldegree = 4 +Ylmdegree = 4 +totdegree = 8 +Nel = 5 +X = randn(SVector{3, Float64}, Nel) +Σ = rand(spins(), Nel) +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] + +## + +# Defining AtomicOrbitalsBasis +bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +bYlm = RYlmBasis(Ylmdegree) + +# setup state +BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) +F(X) = BFwf_chain(X, ps, st)[1] + +# Define a shorter Chain up to AA +BFwf_AA_chain = BFwf_lux_AA(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +ps2, st2 = setupBFState(MersenneTwister(1234), BFwf_AA_chain, Σ) +F2(X) = BFwf_AA_chain(X, ps2, st2)[1] + +AA = F2(X) + +# getting weight from dense_layer +W = ps.branch.bf.hidden1.W +typeof(AA) # Transpose{Float64, ObjectPools.FlexCachedArray{Float64,....}} +typeof(W) +AA_mat = Matrix(AA) # Matrix{Float64} +@assert W * AA ≈ W * AA_mat + +@btime F(X) # 46.671 μs +@btime F2(X) # 22.180 μs +@btime W * AA # 37.171 μs +@btime W * AA_mat # 9.400 μs +@btime W * Matrix(AA) # 18.500 μs \ No newline at end of file From 6171a0f30329030c7683672f8ad68d240a205c33 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 7 Jun 2023 22:56:48 +0800 Subject: [PATCH 43/79] draft of vmc --- src/backflowpooling.jl | 8 +++---- vmc/Eloc.jl | 32 +++++++++++++++++++++++++++ vmc/gradient.jl | 50 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 86 insertions(+), 4 deletions(-) create mode 100644 vmc/Eloc.jl create mode 100644 vmc/gradient.jl diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index b79949f..edd8804 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -14,16 +14,16 @@ Base.length(pooling::BackflowPooling) = 3 * length(pooling.basis) # length(spin( function evaluate(pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) basis = pooling.basis - nuc = basis.nuclei + nuc = basis.nuclei Nnuc = length(nuc) - Nnlm = length(basis.prodbasis.layers.ϕnlms.basis.spec) + Nnlm = size(ϕnlm, 3) Nel = length(Σ) T = promote_type(eltype(ϕnlm)) # evaluate the pooling operation - # spin I k = (nlm) + # spin I k = (nlm) - Aall = zeros(T, (2, Nnuc, Nnlm)) + Aall = zeros(T, ((2, Nnuc, Nnlm))) for k = 1:Nnlm for i = 1:Nel iσ = spin2idx(Σ[i]) diff --git a/vmc/Eloc.jl b/vmc/Eloc.jl new file mode 100644 index 0000000..68ca16c --- /dev/null +++ b/vmc/Eloc.jl @@ -0,0 +1,32 @@ +# INTERFACE FOR HAMILTIANS H ψ -> H(psi, X) +struct SumH{TK, TT, TE} + K::TK + Vext::TT + Vee::TE +end + +K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) + +Vext(wf, X::AbstractVector, ps, st) = sum(1/norm(nuclei[i].rr - X[j]) + for i = 1:length(nuclei) for j in 1 : length(X)) + +Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1 : length(X)-1 for j = i+1 : length(X)) + +H = SumH(K, Vext, Vee) + +(H::SumH)(wf, X::AbstractVector, ps, st) = + H.K(wf, X, ps, st) + (H.Vext(wf, X, ps, st) + H.Vee(wf, X, ps, st)) * evaluate(wf, X, ps, st) + + +# evaluate local energy with SumH + +""" +E_loc = E_pot - 1/4 ∇²ᵣ ϕ(r) - 1/8 (∇ᵣ ϕ)²(r) +https://arxiv.org/abs/2105.08351 +""" + +function Elocal(H::SumH, wf, X::AbstractVector, ps, st) + gra = gradient(wf, X, ps, st) + val = H.Vext(wf, X, ps, st) + H.Vee(wf, X, ps, st) - 1/4 * laplacian(wf, X, ps, st) - 1/8 * gra' * gra + return val +end \ No newline at end of file diff --git a/vmc/gradient.jl b/vmc/gradient.jl new file mode 100644 index 0000000..f574b36 --- /dev/null +++ b/vmc/gradient.jl @@ -0,0 +1,50 @@ +x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) + +function _gradlap(g_bchain, x) + function _mapadd!(f, dest::NamedTuple, src::NamedTuple) + for k in keys(dest) + _mapadd!(f, dest[k], src[k]) + end + return nothing + end + _mapadd!(f, dest::Nothing, src) = nothing + _mapadd!(f, dest::AbstractArray, src::AbstractArray) = + map!((s, d) -> d + f(s), dest, src, dest) + + Δ = zero!(g_bchain(x)) + hX = [x2dualwrtj(xx, 0) for xx in x] + for i = 1:3 + for j = 1:length(x) + hX[j] = x2dualwrtj(x[j], i) + _mapadd!(ε₁ε₂part, Δ, g_bchain(hX)) + hX[j] = x2dualwrtj(x[j], 0) + end + end + return Δ +end + +evaluate(wf, X::AbstractVector, ps, st) = wf(X, ps, st) + +gradient(wf, x, ps, st) = Zygote.gradient(x -> wf(x, ps, st)[1], x)[1] + +grad_params(wf, x, ps, st) = Zygote.gradient(p -> wf(x, p, st)[1], ps)[1] + +function laplacian(wf, x, ps, st) + ΔΨ = 0.0 + hX = [x2dualwrtj(xx, 0) for xx in x] + Nel = length(x) + for i = 1:3 + for j = 1:Nel + hX[j] = x2dualwrtj(x[j], i) # ∂Φ/∂xj_{i} + ΔΨ += wf(hX, ps, st)[1].epsilon12 + hX[j] = x2dualwrtj(x[j], 0) + end + end + return ΔΨ +end + +function grad_lap(wf, x, ps, st) + g_bchain = xx -> Zygote.gradient(p -> wf(xx, p, st)[1], ps)[1] + p, = destructure(_gradlap(g_bchain, x)) + return p +end \ No newline at end of file From b69f3bb9fbd76b292ffa24965e537bf8a0af1672 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Sat, 10 Jun 2023 12:37:09 +0800 Subject: [PATCH 44/79] add metropolis --- src/ACEpsi.jl | 3 + {vmc => src/vmc}/Eloc.jl | 11 +--- {vmc => src/vmc}/gradient.jl | 4 +- src/vmc/metropolis.jl | 115 +++++++++++++++++++++++++++++++++++ src/vmc/vmc.jl | 7 +++ test/test_bflow_lux.jl | 1 + 6 files changed, 131 insertions(+), 10 deletions(-) rename {vmc => src/vmc}/Eloc.jl (64%) rename {vmc => src/vmc}/gradient.jl (94%) create mode 100644 src/vmc/metropolis.jl create mode 100644 src/vmc/vmc.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index bcd4eee..36fa5a5 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -19,4 +19,7 @@ include("backflowpooling.jl") # lux utils for bflow include("lux_utils.jl") + +# vmc +include("vmc/vmc.jl") end diff --git a/vmc/Eloc.jl b/src/vmc/Eloc.jl similarity index 64% rename from vmc/Eloc.jl rename to src/vmc/Eloc.jl index 68ca16c..8ad8f17 100644 --- a/vmc/Eloc.jl +++ b/src/vmc/Eloc.jl @@ -1,3 +1,5 @@ +export SumH + # INTERFACE FOR HAMILTIANS H ψ -> H(psi, X) struct SumH{TK, TT, TE} K::TK @@ -5,15 +7,6 @@ struct SumH{TK, TT, TE} Vee::TE end -K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) - -Vext(wf, X::AbstractVector, ps, st) = sum(1/norm(nuclei[i].rr - X[j]) - for i = 1:length(nuclei) for j in 1 : length(X)) - -Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1 : length(X)-1 for j = i+1 : length(X)) - -H = SumH(K, Vext, Vee) - (H::SumH)(wf, X::AbstractVector, ps, st) = H.K(wf, X, ps, st) + (H.Vext(wf, X, ps, st) + H.Vee(wf, X, ps, st)) * evaluate(wf, X, ps, st) diff --git a/vmc/gradient.jl b/src/vmc/gradient.jl similarity index 94% rename from vmc/gradient.jl rename to src/vmc/gradient.jl index f574b36..c13a2ce 100644 --- a/vmc/gradient.jl +++ b/src/vmc/gradient.jl @@ -1,3 +1,5 @@ +using Zygote + x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) function _gradlap(g_bchain, x) @@ -23,7 +25,7 @@ function _gradlap(g_bchain, x) return Δ end -evaluate(wf, X::AbstractVector, ps, st) = wf(X, ps, st) +evaluate(wf, X::AbstractVector, ps, st) = wf(X, ps, st)[1] gradient(wf, x, ps, st) = Zygote.gradient(x -> wf(x, ps, st)[1], x)[1] diff --git a/src/vmc/metropolis.jl b/src/vmc/metropolis.jl new file mode 100644 index 0000000..73c2ff4 --- /dev/null +++ b/src/vmc/metropolis.jl @@ -0,0 +1,115 @@ +using StatsBase + +""" +`MHSampler` +Metropolis-Hastings sampling algorithm. +""" + +mutable struct MHSampler + Nel::Int + Δt::Float64 # step size (of Gaussian proposal) + burnin::Int # burn-in iterations + lag::Int # iterations between successive samples + N_batch::Int # batch size + nchains::Int # Number of chains + Ψ # many-body wavefunction for sampling + x0::Any # initial sampling + walkerType::String # walker type: "unbiased", "Langevin" + bc::String # boundary condition + type::Int64 # move how many electron one time +end + +MHSampler(Ψ, Nel; Δt = 0.1, + burnin = 100, + lag = 10, + N_batch = 1, + nchains = 1000, + x0 = [], + wT = "unbiased", + bc = "periodic", + type = 3) = + MHSampler(Nel, Δt, burnin, lag, N_batch, nchains, Ψ, x0, wT, bc, type) + + +""" +unbiased random walk: R_n+1 = R_n + Δ⋅Wn +biased random walk: R_n+1 = R_n + Δ⋅Wn + Δ⋅∇(log Ψ)(R_n) +""" + + + +function MHstep(r0, + Ψx0, + Nels::Int, + sam::MHSampler) + rand_sample(X::AbstractVector, u::Int, Δt::AbstractFloat) = begin + Y = copy(X) + ind = sample(1:length(X), u, replace=false) + Y[ind] = X[ind] + Δt * randn(SVector{3, Float64}, u) + return Y + end + rp = rand_sample.(r0, Ref(sam.type), Ref(sam.Δt)) + Ψxp = ACEpsi.vmc.evaluate.(Ref(sam.Ψ), rp, Ref(ps), Ref(st)) + accprob = accfcn(r0, rp, Ψx0, Ψxp, sam) + u = rand(sam.nchains) + acc = u .<= accprob[:] + r = acc .* rp + (1 .- acc) .* r0 + Ψ = acc .* Ψxp + (1 .- acc) .* Ψx0 + return r, Ψ, acc +end + +""" +acceptance rate for log|Ψ| +ψₜ₊₁²/ψₜ² = exp((log|Ψₜ₊₁|^2-log |ψₜ|^2)) +""" + +function accfcn(r0, rp, Ψx0, Ψxp, sam::MHSampler) + acc = exp.(Ψxp .- Ψx0) + return acc +end + +"""============== Metropolis sampling algorithm ============ +type = "restart" +""" +function sampler_restart(sam::MHSampler, ps, st) + r0 = [randn(SVector{3, Float64}, Nel) for _ = 1:sam.nchains] + Ψx0 = ACEpsi.vmc.evaluate.(Ref(sam.Ψ), r0, Ref(ps), Ref(st)) + acc = [] + for _ = 1 : sam.burnin + r0, Ψx0, a = MHstep(r0, Ψx0, sam.Nel, sam); + push!(acc,a) + end + return r0, Ψx0, mean(mean(acc)) +end + +""" +type = "continue" +start from the previous sampling x0 +""" +function sampler(sam::MHSampler, ps, st) + if isempty(sam.x0) + r0, Ψx0, = sampler_restart(sam, ps, st); + else + r0 = sam.x0 + Ψx0 = ACEpsi.vmc.evaluate.(Ref(sam.Ψ), r0, Ref(ps), Ref(st)) + end + acc = [] + for i = 1:sam.lag + r0, Ψx0, a = MHstep(r0, Ψx0, sam.Nel, sam); + push!(acc, a) + end + return r0, Ψx0, mean(mean(acc)) +end + + + +""" +Rayleigh quotient by VMC using Metropolis sampling +""" +function rq_MC(Ψ, sam::MHSampler, ham::SumH, ps, st) + r, ~, acc = sampler(sam, ps, st); + Eloc = Elocal.(Ref(ham), Ref(Ψ), r, Ref(sam.Σ)) + val = sum(Eloc) / length(Eloc) + var = sqrt(sum((Eloc .-val).^2)/(length(Eloc)*(length(Eloc)-1))) + return val, var, acc +end \ No newline at end of file diff --git a/src/vmc/vmc.jl b/src/vmc/vmc.jl new file mode 100644 index 0000000..9dafc21 --- /dev/null +++ b/src/vmc/vmc.jl @@ -0,0 +1,7 @@ +module vmc + +include("Eloc.jl") +include("gradient.jl") +include("metropolis.jl") + +end \ No newline at end of file diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index dee7d79..c6c99a9 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -2,6 +2,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow +using ACEpsi.vmc: evaluate, gradient, laplacian, grad_params using ACEbase.Testing: print_tf, fdtest using LuxCore using Lux From 0eb5311d9df3429bd51ca1f4e37b46d35be21dac Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 12 Jun 2023 10:29:40 +0800 Subject: [PATCH 45/79] remove trans layer --- src/bflow3d.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 4ded405..63ab4f6 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -29,7 +29,7 @@ struct DenseLayer <: AbstractExplicitLayer end function (l::DenseLayer)(x::AbstractMatrix, ps, st) - return ps.W * x, st + return x * ps.W, st end # Jerry: Maybe we should use Glorot Uniform if we have no idea about what we should use? @@ -39,7 +39,7 @@ LuxCore.initialstates(rng::AbstractRNG, l::DenseLayer) = NamedTuple() function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMatrix{Hyper{T}}, ps, st) where T<:Number val = l(x, ps, st) function pb(A) - return NoTangent(), NoTangent(), NoTangent(), (W = A[1] * x',), NoTangent() + return NoTangent(), NoTangent(), NoTangent(), (W = x * A[1]',), NoTangent() end return val, pb end @@ -93,7 +93,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), - bAA = corr_layer, transpose_layer = WrappedFunction(transpose), hidden1 = DenseLayer(length(corr1), Nel), + bAA = corr_layer, hidden1 = DenseLayer(Nel,length(corr1)), Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) return Chain(; branch = BranchLayer((js = jastrow_layer, bf = BFwf_chain, )), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ) end From beaf70d576bedc731bb46ffc95c690a4b2d81968 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 12 Jun 2023 12:42:17 +0800 Subject: [PATCH 46/79] example of gaussian lux --- Project.toml | 1 + src/vmc/gradient.jl | 32 +----- src/vmc/metropolis.jl | 60 ++++++++--- src/vmc/optimiser.jl | 57 ++++++++++ test/test_bflow1.jl | 225 --------------------------------------- test/test_bflow_lux1.jl | 228 ++++++++++++++++++++++++++++++++++++++++ test/test_vmc.jl | 55 ++++++++++ 7 files changed, 389 insertions(+), 269 deletions(-) create mode 100644 src/vmc/optimiser.jl delete mode 100644 test/test_bflow1.jl create mode 100644 test/test_bflow_lux1.jl create mode 100644 test/test_vmc.jl diff --git a/Project.toml b/Project.toml index 49d0cf4..af2c47c 100644 --- a/Project.toml +++ b/Project.toml @@ -22,6 +22,7 @@ Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/src/vmc/gradient.jl b/src/vmc/gradient.jl index c13a2ce..d5cb390 100644 --- a/src/vmc/gradient.jl +++ b/src/vmc/gradient.jl @@ -1,32 +1,8 @@ using Zygote +using HyperDualNumbers: Hyper x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) -function _gradlap(g_bchain, x) - function _mapadd!(f, dest::NamedTuple, src::NamedTuple) - for k in keys(dest) - _mapadd!(f, dest[k], src[k]) - end - return nothing - end - _mapadd!(f, dest::Nothing, src) = nothing - _mapadd!(f, dest::AbstractArray, src::AbstractArray) = - map!((s, d) -> d + f(s), dest, src, dest) - - Δ = zero!(g_bchain(x)) - hX = [x2dualwrtj(xx, 0) for xx in x] - for i = 1:3 - for j = 1:length(x) - hX[j] = x2dualwrtj(x[j], i) - _mapadd!(ε₁ε₂part, Δ, g_bchain(hX)) - hX[j] = x2dualwrtj(x[j], 0) - end - end - return Δ -end - -evaluate(wf, X::AbstractVector, ps, st) = wf(X, ps, st)[1] - gradient(wf, x, ps, st) = Zygote.gradient(x -> wf(x, ps, st)[1], x)[1] grad_params(wf, x, ps, st) = Zygote.gradient(p -> wf(x, p, st)[1], ps)[1] @@ -44,9 +20,3 @@ function laplacian(wf, x, ps, st) end return ΔΨ end - -function grad_lap(wf, x, ps, st) - g_bchain = xx -> Zygote.gradient(p -> wf(xx, p, st)[1], ps)[1] - p, = destructure(_gradlap(g_bchain, x)) - return p -end \ No newline at end of file diff --git a/src/vmc/metropolis.jl b/src/vmc/metropolis.jl index 73c2ff4..17bdb1b 100644 --- a/src/vmc/metropolis.jl +++ b/src/vmc/metropolis.jl @@ -1,10 +1,11 @@ using StatsBase - +using StaticArrays +using Optimisers +export MHSampler """ `MHSampler` Metropolis-Hastings sampling algorithm. """ - mutable struct MHSampler Nel::Int Δt::Float64 # step size (of Gaussian proposal) @@ -36,12 +37,12 @@ unbiased random walk: R_n+1 = R_n + Δ⋅Wn biased random walk: R_n+1 = R_n + Δ⋅Wn + Δ⋅∇(log Ψ)(R_n) """ - +eval(wf, X::AbstractVector, ps, st) = wf(X, ps, st)[1] function MHstep(r0, Ψx0, Nels::Int, - sam::MHSampler) + sam::MHSampler, ps, st) rand_sample(X::AbstractVector, u::Int, Δt::AbstractFloat) = begin Y = copy(X) ind = sample(1:length(X), u, replace=false) @@ -49,8 +50,8 @@ function MHstep(r0, return Y end rp = rand_sample.(r0, Ref(sam.type), Ref(sam.Δt)) - Ψxp = ACEpsi.vmc.evaluate.(Ref(sam.Ψ), rp, Ref(ps), Ref(st)) - accprob = accfcn(r0, rp, Ψx0, Ψxp, sam) + Ψxp = eval.(Ref(sam.Ψ), rp, Ref(ps), Ref(st)) + accprob = accfcn(Ψx0, Ψxp) u = rand(sam.nchains) acc = u .<= accprob[:] r = acc .* rp + (1 .- acc) .* r0 @@ -63,7 +64,7 @@ acceptance rate for log|Ψ| ψₜ₊₁²/ψₜ² = exp((log|Ψₜ₊₁|^2-log |ψₜ|^2)) """ -function accfcn(r0, rp, Ψx0, Ψxp, sam::MHSampler) +function accfcn(Ψx0, Ψxp) acc = exp.(Ψxp .- Ψx0) return acc end @@ -71,12 +72,13 @@ end """============== Metropolis sampling algorithm ============ type = "restart" """ + function sampler_restart(sam::MHSampler, ps, st) - r0 = [randn(SVector{3, Float64}, Nel) for _ = 1:sam.nchains] - Ψx0 = ACEpsi.vmc.evaluate.(Ref(sam.Ψ), r0, Ref(ps), Ref(st)) + r0 = [randn(SVector{3, Float64}, sam.Nel) for _ = 1:sam.nchains] + Ψx0 = eval.(Ref(sam.Ψ), r0, Ref(ps), Ref(st)) acc = [] for _ = 1 : sam.burnin - r0, Ψx0, a = MHstep(r0, Ψx0, sam.Nel, sam); + r0, Ψx0, a = MHstep(r0, Ψx0, sam.Nel, sam, ps, st); push!(acc,a) end return r0, Ψx0, mean(mean(acc)) @@ -91,11 +93,11 @@ function sampler(sam::MHSampler, ps, st) r0, Ψx0, = sampler_restart(sam, ps, st); else r0 = sam.x0 - Ψx0 = ACEpsi.vmc.evaluate.(Ref(sam.Ψ), r0, Ref(ps), Ref(st)) + Ψx0 = eval.(Ref(sam.Ψ), r0, Ref(ps), Ref(st)) end acc = [] for i = 1:sam.lag - r0, Ψx0, a = MHstep(r0, Ψx0, sam.Nel, sam); + r0, Ψx0, a = MHstep(r0, Ψx0, sam.Nel, sam, ps, st); push!(acc, a) end return r0, Ψx0, mean(mean(acc)) @@ -112,4 +114,36 @@ function rq_MC(Ψ, sam::MHSampler, ham::SumH, ps, st) val = sum(Eloc) / length(Eloc) var = sqrt(sum((Eloc .-val).^2)/(length(Eloc)*(length(Eloc)-1))) return val, var, acc -end \ No newline at end of file +end + +function Eloc_Exp_TV_clip(wf, ps, st, + sam::MHSampler, + ham::SumH; + clip = 20.) + x, x0, acc = sampler(sam, ps, st) + Eloc = Elocal.(Ref(ham), Ref(wf), x, Ref(ps), Ref(st)) + val = sum(Eloc) / length(Eloc) + var = sqrt(sum((Eloc .-val).^2)/(length(Eloc)*(length(Eloc) -1))) + ΔE = Eloc .- median( Eloc ) + a = clip * mean( abs.(ΔE) ) + ind = findall(x -> abs(x) > a, ΔE) + ΔE[ind] = (a * sign.(ΔE) .* (1 .+ log.((1 .+(abs.(ΔE)/a).^2)/2)))[ind] + E_clip = median(Eloc) .+ ΔE + return val, var, E_clip, x, x0, acc +end + +function params(a::NamedTuple) + p,= destructure(a) + return p +end + +function grad(wf, x, ps, st, E); + dy = grad_params.(Ref(wf), x, Ref(ps), Ref(st)); + N = length(x) + p = params.(dy) + _,t = destructure(dy[1]) + g = 1/N * sum( p .* E) - 1/(N^2) * sum(E) * sum(p) + g = t(g) + return g; +end + diff --git a/src/vmc/optimiser.jl b/src/vmc/optimiser.jl new file mode 100644 index 0000000..7f4fc12 --- /dev/null +++ b/src/vmc/optimiser.jl @@ -0,0 +1,57 @@ +export VMC + +""" +Minimizing the Rayleigh quotient by VMC +tol: tolrance +MaxIter: maximum iteration number +lr: learning rates +lr_dc_pow: rate to decrease 'lr' -- lr / (1 + k / lr_dc) +""" +mutable struct VMC <: AbstractOptimizer + tol::Float64 + MaxIter::Int + lr::Float64 + lr_dc::Float64 +end + +vmc(MaxIter::Int, lr::Float64; tol = 1.0e-3, lr_dc = 50.0) = + VMC(tol, MaxIter, lr, lr_dc); + +""" +decreasing the learning rate +""" +function InverseLR(ν, lr, lr_dc) + return lr / (1 + ν / lr_dc), ν+1 +end + +function gd_GradientByVMC( + sam::AbstractSampler, + ham::SumH, + wf) + res = 1.0; + λ₀ = 0.; + g0 = 0; + α = opt.lr; + err_opt = zeros(opt.MaxIter) + x0, ~, acc = sampler_restart(sam); + verbose && @printf(" k | 𝔼[E_L] | V[E_L] | res | LR |accRate| Δt | regular\n") + νk = ν; + for k = 1 : opt.MaxIter + sam.x0 = x0; + sam.Ψ = wf; + + # optimize + g0, mₜ, vₜ, c, λ₀, σ, x0, acc = adam(mₜ, vₜ, c, ν, U, sam, ham, clip, ηₜ = α; Γ = Γ); + res = norm(g0); + set_params!(U, c) + verbose && @printf(" %3.d | %.5f | %.5f | %.5f | %.5f | %.3f | %.3f | %.4f \n", + k, λ₀, σ, res, α, acc, sam.Δt, norm(Γ)); + err_opt[k] = λ₀ + if res < opt.tol + break; + end + end + set_params!(U, c) + return U, λ₀, σ, err_opt, err_avg, mₜ, vₜ +end + diff --git a/test/test_bflow1.jl b/test/test_bflow1.jl deleted file mode 100644 index 87d5429..0000000 --- a/test/test_bflow1.jl +++ /dev/null @@ -1,225 +0,0 @@ - -using Polynomials4ML, ACEcore, ACEbase, Printf, ACEpsi -using ACEpsi: BFwf1, gradient, laplacian -using LinearAlgebra -#using Random -#Random.seed!(123) -## -function lap_test(f, Δf, X) - F = f(X) - ΔF = Δf(X) - nX = length(X) - n1, nF = size(F) - U = randn(n1) - V = randn(nF) ./ (1:nF).^2 - f0 = U' * F * V - Δf0 = U' * ΔF * V - EE = Matrix(I, (nX, nX)) - for h in sqrt(0.1).^((5:12)) - Δfh = 0.0 - for i = 1:nX - Δfh += (U'*f(X + h * EE[:, i])*V - f0) / h^2 - Δfh += (U'*f(X - h * EE[:, i])*V - f0) / h^2 - end - @printf(" %.1e | %.2e \n", h, abs(Δfh - Δf0)) - end -end - -function grad_test2(f, df, X::AbstractVector) - F = f(X) - ∇F = df(X) - nX = length(X) - EE = Matrix(I, (nX, nX)) - - for h in 0.1.^(3:12) - gh = [ (f(X + h * EE[:, i]) - F) / h for i = 1:nX ] - @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) - end -end - -function grad_test3(f, df, X::Float64) - F = f(X) - ∇F = df(X) # return as vector of size 1 - for h in 0.1.^(3:10) - gh = [ (f(X + h) - F) / h] - @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) - end -end - - -function grad_test(f, df, X) - F = f(X) - ∇F = df(X) - nX, nF = size(F) - U = randn(nX) - V = randn(nF) ./ (1:nF).^2 - f0 = U' * F * V - ∇f0 = [ U' * ∇F[i, :, :] * V for i = 1:nX ] - EE = Matrix(I, (Nel, Nel)) - for h in 0.1.^(2:10) - gh = [ (U'*f(X + h * EE[:, i])*V - f0) / h for i = 1:Nel ] - @printf(" %.1e | %.2e \n", h, norm(gh - ∇f0, Inf)) - end -end - -"This function should be removed later to test in a nicer way..." -function fdtest(F, Σ, dF, x::AbstractVector; h0 = 1.0, verbose=true) - errors = Float64[] - E = F(x, Σ) - dE = dF - # loop through finite-difference step-lengths - verbose && @printf("---------|----------- \n") - verbose && @printf(" h | error \n") - verbose && @printf("---------|----------- \n") - for p = 2:11 - h = 0.1^p - dEh = copy(dE) - for n = 1:length(dE) - x[n] += h - dEh[n] = (F(x, Σ) - E) / h - x[n] -= h - end - push!(errors, norm(dE - dEh, Inf)) - verbose && @printf(" %1.1e | %4.2e \n", h, errors[end]) - end - verbose && @printf("---------|----------- \n") - if minimum(errors) <= 1e-3 * maximum(errors) - verbose && println("passed") - return true - else - @warn("""It seems the finite-difference test has failed, which indicates - that there is an inconsistency between the function and gradient - evaluation. Please double-check this manually / visually. (It is - also possible that the function being tested is poorly scaled.)""") - return false - end - end -## - -const ↑, ↓, ∅ = '↑','↓','∅' -Nel = 5 -polys = legendre_basis(8) -wf = BFwf1(Nel, polys, trans = x -> atan(x); ν=3) - -X = 2 * rand(Nel) .- 1 -Σ = rand([↑, ↓], Nel) - -wf(X, Σ) -g = gradient(wf, X, Σ) - -## - -using LinearAlgebra -using Printf -#using ACEbase.Testing: fdtest - -@info("Fd test of gradient w.r.t. X") -fdtest(wf, Σ, g, X) - - -# ## - -# todo: move this to a performance benchmark script -# using BenchmarkTools -# @btime $wf($X) -# @btime gradient($wf, $X) - -## - -A, ∇A, ΔA = ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ) - -@info("test ∇A") -grad_test(X -> ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ)[1], - X -> ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ)[2], - X) - -@info("test ΔA") -lap_test(X -> ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ)[1], - X -> ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ)[3], - X) - - -## - -function f_AA(X, Σ) - A, ∇A, ΔA = ACEpsi._assemble_A_∇A_ΔA(wf, X, Σ) - AA, ∇AA, ΔAA = ACEpsi._assemble_AA_∇AA_ΔAA(A, ∇A, ΔA, wf) - return AA, ∇AA, ΔAA -end - -@info("test ∇A") -grad_test(X -> f_AA(X, Σ)[1], X -> f_AA(X, Σ)[2], X) - -@info("test ΔA") -lap_test(X -> f_AA(X, Σ)[1], X -> f_AA(X, Σ)[3], X) - - -## - -@info("test Δψ") -lap_test(X -> [wf(X, Σ);;], X -> [ACEpsi.laplacian(wf, X, Σ);;], X) - - -## - -@info("Test ∇ψ w.r.t. parameters") - -ACEpsi.gradp_evaluate(wf, X, Σ) - - -W0 = copy(wf.W) -w0 = W0[:] -Fp = w -> ( wf.W[:] .= w[:]; wf(X, Σ)) -dFp = w -> ( wf.W[:] .= w[:]; ACEpsi.gradp_evaluate(wf, X, Σ)[1][:] ) - -grad_test2(Fp, dFp, w0) - -## - -@info("test ∇env w.r.t. parameter") -Ξ0 = copy(wf.envelope.ξ) -ξ0 = Ξ0 - -Envp = w -> (wf.envelope.ξ = w; wf(X, Σ)) -dEnvp = w -> (wf.envelope.ξ = w; ACEpsi.gradp_evaluate(wf, X, Σ)[2]) - -grad_test3(Envp, dEnvp, ξ0) - -## - -@info("Test ∇Δψ w.r.t. parameters") - -Fp = w -> ( wf.W[:] .= w[:]; ACEpsi.laplacian(wf, X, Σ)) -dFp = w -> ( wf.W[:] .= w[:]; ACEpsi.gradp_laplacian(wf, X, Σ)[1][:] ) - -grad_test2(Fp, dFp, w0) - -## - -@info("Test ∇Δenv w.r.t. parameters") - -Ξ0 = copy(wf.envelope.ξ) -ξ0 = Ξ0 - - -Fp = w -> ( wf.envelope.ξ = w; ACEpsi.laplacian(wf, X, Σ)) -dFp = w -> (wf.envelope.ξ = w; ACEpsi.gradp_laplacian(wf, X, Σ)[2][:] ) - -grad_test3(Fp, dFp, ξ0) - -## - -@info("Test getting/setting parameters") - -wf1 = BFwf1(Nel, polys; ν=3) -wf2 = BFwf1(Nel, polys; ν=3) -@printf(" wf1 - wf2: %f \n", abs(wf1(X, Σ) - wf2(X, Σ))) -param1 = ACEpsi.get_params(wf1) -wf2 = ACEpsi.set_params!(wf2, param1) -@printf(" wf1 - wf2 with parameter wf1: %f \n", abs(wf1(X, Σ) - wf2(X, Σ))) - -## - -@warn("removed compat test since json file is missing") -# @info("Test compatibility with ACESchrodinger") # Jerry: Not sure if this should be kept in the same file -# include("compare_bflow.jl") diff --git a/test/test_bflow_lux1.jl b/test/test_bflow_lux1.jl new file mode 100644 index 0000000..a61b2d5 --- /dev/null +++ b/test/test_bflow_lux1.jl @@ -0,0 +1,228 @@ +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow +using ACEpsi.vmc: evaluate, gradient, laplacian, grad_params +using ACEbase.Testing: print_tf, fdtest +using LuxCore +using Lux +using Zygote +using Optimisers # mainly for the destrcuture(ps) function +using Random +using Printf +using LinearAlgebra +using BenchmarkTools + +using HyperDualNumbers: Hyper + +function grad_test2(f, df, X::AbstractVector) + F = f(X) + ∇F = df(X) + nX = length(X) + EE = Matrix(I, (nX, nX)) + + for h in 0.1.^(3:12) + gh = [ (f(X + h * EE[:, i]) - F) / h for i = 1:nX ] + @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) + end +end + +Rnldegree = n1 = 4 +Ylmdegree = 4 +totdegree = 8 +Nel = 5 +X = randn(SVector{3, Float64}, Nel) +Σ = rand(spins(), Nel) +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] + +# wrap it as HyperDualNumbers +x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) +hX = [x2dualwrtj(x, 0) for x in X] +hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec + +## + +# Defining AtomicOrbitalsBasis +n2 = 3 +Pn = Polynomials4ML.legendre_basis(n1+1) +spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] +ζ = rand(length(spec)) +Dn = GaussianBasis(ζ) +bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) +bYlm = RYlmBasis(Ylmdegree) + +# setup state +BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) + +## + +@info("Test evaluate") +A1 = BFwf_chain(X, ps, st) +hA1 = BFwf_chain(hX, ps, st) + +print_tf(@test hA1[1].value ≈ A1[1]) + +println() + +## +F(X) = BFwf_chain(X, ps, st)[1] + +# @profview let F = F, X = X +# for i = 1:10_000 +# F(X) +# end +# end + +@btime F(X) + + +## + +@info("Test ∇ψ w.r.t. X") +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) +y, st = Lux.apply(BFwf_chain, X, ps, st) + +F(X) = BFwf_chain(X, ps, st)[1] +dF(X) = Zygote.gradient(x -> BFwf_chain(x, ps, st)[1], X)[1] +fdtest(F, dF, X, verbose = true) + +## + +@info("Test consistency with HyperDualNumbers") +for _ = 1:30 + local X = randn(SVector{3, Float64}, Nel) + local Σ = rand(spins(), Nel) + local hdF = [zeros(3) for _ = 1:Nel] + local hX = [x2dualwrtj(x, 0) for x in X] + for i = 1:3 + for j = 1:Nel + hX[j] = x2dualwrtj(X[j], i) # ∂Ψ/∂xj_{i} + hdF[j][i] = BFwf_chain(hX, ps, st)[1].epsilon1 + hX[j] = x2dualwrtj(X[j], 0) + end + end + print_tf(@test dF(X) ≈ hdF) +end +println() + +## + +@info("Test ∇ψ w.r.t. parameters") +p = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1] +p, = destructure(p) + +W0, re = destructure(ps) +Fp = w -> BFwf_chain(X, re(w), st)[1] +dFp = w -> ( gl = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1]; destructure(gl)[1]) +grad_test2(Fp, dFp, W0) + +## + +@info("Test consistency when input isa HyperDualNumbers") +hp = Zygote.gradient(p -> BFwf_chain(hX, p, st)[1], ps)[1] +hp, = destructure(hp) +P = similar(p) +for i = 1:length(P) + P[i] = hp[i].value +end + +print_tf(@test P ≈ p) + +println() + +## + +@info("Test Δψ w.r.t. X using HyperDualNumbers") + +X = [randn(3) for _ = 1:Nel] +hX = [x2dualwrtj(x, 0) for x in X] +Σ = rand(spins(), Nel) +F(x) = BFwf_chain(x, ps, st)[1] + + +function ΔF(x) + ΔΨ = 0.0 + hX = [x2dualwrtj(xx, 0) for xx in x] + for i = 1:3 + for j = 1:Nel + hX[j] = x2dualwrtj(x[j], i) # ∂Φ/∂xj_{i} + ΔΨ += BFwf_chain(hX, ps, st)[1].epsilon12 + hX[j] = x2dualwrtj(x[j], 0) + end + end + return ΔΨ +end + +Δ1 = ΔF(X) +f0 = F(X) + +for h in 0.1.^(1:8) + Δfh = 0.0 + for i = 1:Nel + for j = 1:3 + XΔX_add, XΔX_sub = deepcopy(X), deepcopy(X) + XΔX_add[i][j] += h + XΔX_sub[i][j] -= h + Δfh += (F(XΔX_add) - f0) / h^2 + Δfh += (F(XΔX_sub) - f0) / h^2 + end + end + @printf(" %.1e | %.2e \n", h, abs(Δfh - Δ1)) +end + +## + +@info("Test gradp Δψ using HyperDualNumbers") +g_bchain = xx -> Zygote.gradient(p -> BFwf_chain(xx, p, st)[1], ps)[1] +g_bchain(hX) + +using ACEpsi: zero! +using HyperDualNumbers + +function grad_lap(g_bchain, x) + function _mapadd!(f, dest::NamedTuple, src::NamedTuple) + for k in keys(dest) + _mapadd!(f, dest[k], src[k]) + end + return nothing + end + _mapadd!(f, dest::Nothing, src) = nothing + _mapadd!(f, dest::AbstractArray, src::AbstractArray) = + map!((s, d) -> d + f(s), dest, src, dest) + + Δ = zero!(g_bchain(x)) + hX = [x2dualwrtj(xx, 0) for xx in x] + for i = 1:3 + for j = 1:length(x) + hX[j] = x2dualwrtj(x[j], i) + _mapadd!(ε₁ε₂part, Δ, g_bchain(hX)) + hX[j] = x2dualwrtj(x[j], 0) + end + end + return Δ +end + +function ΔF(x, ps) + ΔΨ = 0.0 + hX = [x2dualwrtj(xx, 0) for xx in x] + for i = 1:3 + for j = 1:Nel + hX[j] = x2dualwrtj(x[j], i) # ∂Φ/∂xj_{i} + ΔΨ += BFwf_chain(hX, ps, st)[1].epsilon12 + hX[j] = x2dualwrtj(x[j], 0) + end + end + return ΔΨ +end + +function ∇ΔF(x, ps) + g_bchain = xx -> Zygote.gradient(p -> BFwf_chain(xx, p, st)[1], ps)[1] + p, = destructure(grad_lap(g_bchain, x)) + return p +end + +W0, re = destructure(ps) +Fp = w -> ΔF(X, re(w)) +dFp = w -> ∇ΔF(X, re(w)) +fdtest(Fp, dFp, W0) \ No newline at end of file diff --git a/test/test_vmc.jl b/test/test_vmc.jl new file mode 100644 index 0000000..afae0ea --- /dev/null +++ b/test/test_vmc.jl @@ -0,0 +1,55 @@ +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow +using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, Eloc_Exp_TV_clip, grad +using ACEbase.Testing: print_tf, fdtest +using LuxCore +using Lux +using Zygote +using Optimisers # mainly for the destrcuture(ps) function +using Random +using Printf +using LinearAlgebra +using BenchmarkTools + +using HyperDualNumbers: Hyper + + +Rnldegree = 4 +Ylmdegree = 4 +totdegree = 8 +Nel = 5 +X = randn(SVector{3, Float64}, Nel) +Σ = rand(spins(), Nel) +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] + +# wrap it as HyperDualNumbers +x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) +hX = [x2dualwrtj(x, 0) for x in X] +hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec + +## + +# Defining AtomicOrbitalsBasis +bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +bYlm = RYlmBasis(Ylmdegree) + +# setup state +wf = BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) + +K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) +Vext(wf, X::AbstractVector, ps, st) = sum(1/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) +Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) + + +ham = SumH(K, Vext, Vee) +sam = MHSampler(wf, Nel) + +λ₀, σ, E_clip, x_clip, x0, acc = Eloc_Exp_TV_clip(wf, ps, st, sam, ham) +g = grad(wf, x_clip, ps, st, E_clip) + +# Optimization +st_opt = Optimisers.setup(Optimisers.Adam(0.0001), ps) +st_opt, ps = Optimisers.update(st_opt, ps, g) \ No newline at end of file From eb58a49a29242786c9eabd61e01f99ca918bc59a Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 12 Jun 2023 14:03:14 +0800 Subject: [PATCH 47/79] add vmc --- src/ACEpsi.jl | 2 +- src/vmc/metropolis.jl | 6 ++--- src/vmc/optimiser.jl | 61 +++++-------------------------------------- src/vmc/vmc.jl | 57 ++++++++++++++++++++++++++++++++++++---- test/test_vmc.jl | 11 +++----- 5 files changed, 65 insertions(+), 72 deletions(-) diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index 36fa5a5..cb2934b 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -21,5 +21,5 @@ include("backflowpooling.jl") include("lux_utils.jl") # vmc -include("vmc/vmc.jl") +include("vmc/optimiser.jl") end diff --git a/src/vmc/metropolis.jl b/src/vmc/metropolis.jl index 17bdb1b..20d1fa6 100644 --- a/src/vmc/metropolis.jl +++ b/src/vmc/metropolis.jl @@ -120,7 +120,7 @@ function Eloc_Exp_TV_clip(wf, ps, st, sam::MHSampler, ham::SumH; clip = 20.) - x, x0, acc = sampler(sam, ps, st) + x, ~, acc = sampler(sam, ps, st) Eloc = Elocal.(Ref(ham), Ref(wf), x, Ref(ps), Ref(st)) val = sum(Eloc) / length(Eloc) var = sqrt(sum((Eloc .-val).^2)/(length(Eloc)*(length(Eloc) -1))) @@ -129,7 +129,7 @@ function Eloc_Exp_TV_clip(wf, ps, st, ind = findall(x -> abs(x) > a, ΔE) ΔE[ind] = (a * sign.(ΔE) .* (1 .+ log.((1 .+(abs.(ΔE)/a).^2)/2)))[ind] E_clip = median(Eloc) .+ ΔE - return val, var, E_clip, x, x0, acc + return val, var, E_clip, x, acc end function params(a::NamedTuple) @@ -142,7 +142,7 @@ function grad(wf, x, ps, st, E); N = length(x) p = params.(dy) _,t = destructure(dy[1]) - g = 1/N * sum( p .* E) - 1/(N^2) * sum(E) * sum(p) + g = -1/N * sum( p .* E) + 1/(N^2) * sum(E) * sum(p) g = t(g) return g; end diff --git a/src/vmc/optimiser.jl b/src/vmc/optimiser.jl index 7f4fc12..114f99f 100644 --- a/src/vmc/optimiser.jl +++ b/src/vmc/optimiser.jl @@ -1,57 +1,8 @@ -export VMC +module vmc -""" -Minimizing the Rayleigh quotient by VMC -tol: tolrance -MaxIter: maximum iteration number -lr: learning rates -lr_dc_pow: rate to decrease 'lr' -- lr / (1 + k / lr_dc) -""" -mutable struct VMC <: AbstractOptimizer - tol::Float64 - MaxIter::Int - lr::Float64 - lr_dc::Float64 -end - -vmc(MaxIter::Int, lr::Float64; tol = 1.0e-3, lr_dc = 50.0) = - VMC(tol, MaxIter, lr, lr_dc); - -""" -decreasing the learning rate -""" -function InverseLR(ν, lr, lr_dc) - return lr / (1 + ν / lr_dc), ν+1 -end - -function gd_GradientByVMC( - sam::AbstractSampler, - ham::SumH, - wf) - res = 1.0; - λ₀ = 0.; - g0 = 0; - α = opt.lr; - err_opt = zeros(opt.MaxIter) - x0, ~, acc = sampler_restart(sam); - verbose && @printf(" k | 𝔼[E_L] | V[E_L] | res | LR |accRate| Δt | regular\n") - νk = ν; - for k = 1 : opt.MaxIter - sam.x0 = x0; - sam.Ψ = wf; - - # optimize - g0, mₜ, vₜ, c, λ₀, σ, x0, acc = adam(mₜ, vₜ, c, ν, U, sam, ham, clip, ηₜ = α; Γ = Γ); - res = norm(g0); - set_params!(U, c) - verbose && @printf(" %3.d | %.5f | %.5f | %.5f | %.5f | %.3f | %.3f | %.4f \n", - k, λ₀, σ, res, α, acc, sam.Δt, norm(Γ)); - err_opt[k] = λ₀ - if res < opt.tol - break; - end - end - set_params!(U, c) - return U, λ₀, σ, err_opt, err_avg, mₜ, vₜ -end +include("Eloc.jl") +include("gradient.jl") +include("metropolis.jl") +include("vmc.jl") +end \ No newline at end of file diff --git a/src/vmc/vmc.jl b/src/vmc/vmc.jl index 9dafc21..a7cadab 100644 --- a/src/vmc/vmc.jl +++ b/src/vmc/vmc.jl @@ -1,7 +1,54 @@ -module vmc +export VMC +using Printf +using LinearAlgebra +using Optimisers +""" +Minimizing the Rayleigh quotient by VMC +tol: tolrance +MaxIter: maximum iteration number +lr: learning rates +lr_dc_pow: rate to decrease 'lr' -- lr / (1 + k / lr_dc) +""" +mutable struct VMC + tol::Float64 + MaxIter::Int + lr::Float64 + lr_dc::Float64 +end -include("Eloc.jl") -include("gradient.jl") -include("metropolis.jl") +VMC(MaxIter::Int, lr::Float64; tol = 1.0e-3, lr_dc = 50.0) = VMC(tol, MaxIter, lr, lr_dc); + +function gd_GradientByVMC(opt::VMC, + sam::MHSampler, + ham::SumH, + wf, ps, st) + res = 1.0; + λ₀ = 0.; + α = opt.lr; + err_opt = zeros(opt.MaxIter) + x0, ~, acc = sampler_restart(sam, ps, st); + verbose = true + verbose && @printf(" k | 𝔼[E_L] | V[E_L] | res | LR |accRate| Δt \n") + for k = 1 : opt.MaxIter + sam.x0 = x0; + sam.Ψ = wf; + + # optimize + λ₀, σ, E, x0, acc = Eloc_Exp_TV_clip(wf, ps, st, sam, ham) + g = grad(wf, x0, ps, st, E) + + # Optimization + st_opt = Optimisers.setup(Optimisers.Adam(α), ps) + st_opt, ps = Optimisers.update(st_opt, ps, g) + + res = norm(destructure(g)[1]); + verbose && @printf(" %3.d | %.5f | %.5f | %.5f | %.5f | %.3f | %.3f \n", + k, λ₀, σ, res, α, acc, sam.Δt); + err_opt[k] = λ₀ + if res < opt.tol + break; + end + end + return wf, λ₀, σ, err_opt, ps +end -end \ No newline at end of file diff --git a/test/test_vmc.jl b/test/test_vmc.jl index afae0ea..a5256d3 100644 --- a/test/test_vmc.jl +++ b/test/test_vmc.jl @@ -2,7 +2,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow -using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, Eloc_Exp_TV_clip, grad +using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC using ACEbase.Testing: print_tf, fdtest using LuxCore using Lux @@ -46,10 +46,5 @@ Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 ham = SumH(K, Vext, Vee) sam = MHSampler(wf, Nel) - -λ₀, σ, E_clip, x_clip, x0, acc = Eloc_Exp_TV_clip(wf, ps, st, sam, ham) -g = grad(wf, x_clip, ps, st, E_clip) - -# Optimization -st_opt = Optimisers.setup(Optimisers.Adam(0.0001), ps) -st_opt, ps = Optimisers.update(st_opt, ps, g) \ No newline at end of file +opt = VMC(100, 0.01) +gd_GradientByVMC(opt,sam,ham,wf,ps,st) \ No newline at end of file From 2b1784c834a0a0d6486021a4c8eb1a83950ee74c Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 12 Jun 2023 14:36:10 +0800 Subject: [PATCH 48/79] add example of He --- src/vmc/metropolis.jl | 4 ++-- src/vmc/vmc.jl | 21 ++++++++++++++++++--- test/test_vmc.jl | 7 +++---- 3 files changed, 23 insertions(+), 9 deletions(-) diff --git a/src/vmc/metropolis.jl b/src/vmc/metropolis.jl index 20d1fa6..2dc0b67 100644 --- a/src/vmc/metropolis.jl +++ b/src/vmc/metropolis.jl @@ -28,7 +28,7 @@ MHSampler(Ψ, Nel; Δt = 0.1, x0 = [], wT = "unbiased", bc = "periodic", - type = 3) = + type = 1) = MHSampler(Nel, Δt, burnin, lag, N_batch, nchains, Ψ, x0, wT, bc, type) @@ -142,7 +142,7 @@ function grad(wf, x, ps, st, E); N = length(x) p = params.(dy) _,t = destructure(dy[1]) - g = -1/N * sum( p .* E) + 1/(N^2) * sum(E) * sum(p) + g = 1/N * sum( p .* E) - 1/(N^2) * sum(E) * sum(p) g = t(g) return g; end diff --git a/src/vmc/vmc.jl b/src/vmc/vmc.jl index a7cadab..cd532f2 100644 --- a/src/vmc/vmc.jl +++ b/src/vmc/vmc.jl @@ -16,23 +16,38 @@ mutable struct VMC lr_dc::Float64 end +function InverseLR(ν, lr, lr_dc) + return lr / (1 + ν / lr_dc), ν+1 +end + VMC(MaxIter::Int, lr::Float64; tol = 1.0e-3, lr_dc = 50.0) = VMC(tol, MaxIter, lr, lr_dc); function gd_GradientByVMC(opt::VMC, sam::MHSampler, ham::SumH, - wf, ps, st) + wf, ps, st, ν = 1, verbose = true, accMCMC = [10, [0.45, 0.55]]) res = 1.0; λ₀ = 0.; α = opt.lr; err_opt = zeros(opt.MaxIter) x0, ~, acc = sampler_restart(sam, ps, st); - verbose = true + verbose && @printf("Initialize MCMC: Δt = %.2f, accRate = %.4f \n", sam.Δt, acc) + acc_step = accMCMC[1]; + acc_opt = zeros(acc_step) + acc_range = accMCMC[2]; verbose && @printf(" k | 𝔼[E_L] | V[E_L] | res | LR |accRate| Δt \n") for k = 1 : opt.MaxIter sam.x0 = x0; sam.Ψ = wf; - + acc_opt[mod(k,acc_step)+1] = acc + if mod(k,acc_step) == 0 + if mean(acc_opt) < acc_range[1] + sam.Δt = sam.Δt * exp(1/10 * (mean(acc_opt) - acc_range[1])/acc_range[1]) + elseif mean(acc_opt) > acc_range[2] + sam.Δt = sam.Δt * exp(1/10 * (mean(acc_opt) - acc_range[2])/acc_range[2]) + end + end + α, ν = InverseLR(ν, opt.lr, opt.lr_dc) # optimize λ₀, σ, E, x0, acc = Eloc_Exp_TV_clip(wf, ps, st, sam, ham) g = grad(wf, x0, ps, st, E) diff --git a/test/test_vmc.jl b/test/test_vmc.jl index a5256d3..ab4aefe 100644 --- a/test/test_vmc.jl +++ b/test/test_vmc.jl @@ -19,10 +19,10 @@ using HyperDualNumbers: Hyper Rnldegree = 4 Ylmdegree = 4 totdegree = 8 -Nel = 5 +Nel = 2 X = randn(SVector{3, Float64}, Nel) -Σ = rand(spins(), Nel) -nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] +Σ = [↑, ↓] +nuclei = [ Nuc(3 * zeros(SVector{3, Float64}), 2.0)] # wrap it as HyperDualNumbers x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) @@ -43,7 +43,6 @@ K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) Vext(wf, X::AbstractVector, ps, st) = sum(1/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) - ham = SumH(K, Vext, Vee) sam = MHSampler(wf, Nel) opt = VMC(100, 0.01) From a42f67fe1354e3172c296a842d1314525b47c6d2 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Tue, 13 Jun 2023 21:14:21 +0800 Subject: [PATCH 49/79] add example of harmonic_oscillator --- src/jastrow.jl | 2 +- src/vmc/vmc.jl | 2 +- test/{test_vmc.jl => harmonic_oscillator.jl} | 31 ++++++++++++-------- test/test_bflow_lux1.jl | 25 ++++++++-------- 4 files changed, 34 insertions(+), 26 deletions(-) rename test/{test_vmc.jl => harmonic_oscillator.jl} (60%) diff --git a/src/jastrow.jl b/src/jastrow.jl index 2b08819..0a0e5a2 100644 --- a/src/jastrow.jl +++ b/src/jastrow.jl @@ -42,7 +42,7 @@ function evaluate(f::Jastrow, X::AbstractVector, Σ) for i = 1:Nnuc, j = 1:Nel-1, k = j+1: Nel F3 += C0 * nuc[i].charge * XN[i, j] * XN[i, k] * log(XN[i, j]^2 + XN[i, k]^2) end - return exp(F2 + F3) + return 1 # exp(F2 + F3) end diff --git a/src/vmc/vmc.jl b/src/vmc/vmc.jl index cd532f2..f60b4ab 100644 --- a/src/vmc/vmc.jl +++ b/src/vmc/vmc.jl @@ -53,7 +53,7 @@ function gd_GradientByVMC(opt::VMC, g = grad(wf, x0, ps, st, E) # Optimization - st_opt = Optimisers.setup(Optimisers.Adam(α), ps) + st_opt = Optimisers.setup(Optimisers.AdamW(α), ps) st_opt, ps = Optimisers.update(st_opt, ps, g) res = norm(destructure(g)[1]); diff --git a/test/test_vmc.jl b/test/harmonic_oscillator.jl similarity index 60% rename from test/test_vmc.jl rename to test/harmonic_oscillator.jl index ab4aefe..383ff6e 100644 --- a/test/test_vmc.jl +++ b/test/harmonic_oscillator.jl @@ -16,13 +16,13 @@ using BenchmarkTools using HyperDualNumbers: Hyper -Rnldegree = 4 +n1 = Rnldegree = 4 Ylmdegree = 4 -totdegree = 8 -Nel = 2 +totdegree = 4 +Nel = 1 X = randn(SVector{3, Float64}, Nel) -Σ = [↑, ↓] -nuclei = [ Nuc(3 * zeros(SVector{3, Float64}), 2.0)] +Σ = [↑] +nuclei = [ Nuc(zeros(SVector{3, Float64}), 1.0)] # wrap it as HyperDualNumbers x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) @@ -32,18 +32,25 @@ hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec ## # Defining AtomicOrbitalsBasis -bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +n2 = 2 +Pn = Polynomials4ML.legendre_basis(n1+1) +spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] +ζ = rand(length(spec)) +Dn = GaussianBasis(ζ) +bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) bYlm = RYlmBasis(Ylmdegree) # setup state -wf = BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +wf = BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 1) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) -Vext(wf, X::AbstractVector, ps, st) = sum(1/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) -Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) + +Vext(wf, X::AbstractVector, ps, st) = 1/2 * sum([X[i][j]^2 for i = 1:length(X) for j = 1:3]) #sum(1/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) + +Vee(wf, X::AbstractVector, ps, st) = 0 # sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) ham = SumH(K, Vext, Vee) -sam = MHSampler(wf, Nel) -opt = VMC(100, 0.01) -gd_GradientByVMC(opt,sam,ham,wf,ps,st) \ No newline at end of file +sam = MHSampler(wf, Nel, Δt = 1.0, burnin = 1000) +opt = VMC(1000, 0.05, lr_dc = 10) +gd_GradientByVMC(opt, sam, ham, wf, ps, st) \ No newline at end of file diff --git a/test/test_bflow_lux1.jl b/test/test_bflow_lux1.jl index a61b2d5..7f9b0df 100644 --- a/test/test_bflow_lux1.jl +++ b/test/test_bflow_lux1.jl @@ -120,16 +120,17 @@ grad_test2(Fp, dFp, W0) ## @info("Test consistency when input isa HyperDualNumbers") -hp = Zygote.gradient(p -> BFwf_chain(hX, p, st)[1], ps)[1] -hp, = destructure(hp) -P = similar(p) -for i = 1:length(P) - P[i] = hp[i].value -end +#hp = Zygote.gradient(p -> BFwf_chain(hX, p, st)[1], ps)[1] -print_tf(@test P ≈ p) +#hp, = destructure(hp) +#P = similar(p) +#for i = 1:length(P) +# P[i] = hp[i].value +#end -println() +#print_tf(@test P ≈ p) + +#println() ## @@ -222,7 +223,7 @@ function ∇ΔF(x, ps) return p end -W0, re = destructure(ps) -Fp = w -> ΔF(X, re(w)) -dFp = w -> ∇ΔF(X, re(w)) -fdtest(Fp, dFp, W0) \ No newline at end of file +#W0, re = destructure(ps) +#Fp = w -> ΔF(X, re(w)) +#dFp = w -> ∇ΔF(X, re(w)) +#fdtest(Fp, dFp, W0) \ No newline at end of file From 8d0dd49256a8451a3aec9a09feb647ff43ebcbb8 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 14 Jun 2023 04:00:01 +0800 Subject: [PATCH 50/79] minor fix --- src/vmc/vmc.jl | 2 +- test/harmonic_oscillator.jl | 4 ++-- test/test_bflow_lux1.jl | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/vmc/vmc.jl b/src/vmc/vmc.jl index f60b4ab..82eaa66 100644 --- a/src/vmc/vmc.jl +++ b/src/vmc/vmc.jl @@ -64,6 +64,6 @@ function gd_GradientByVMC(opt::VMC, break; end end - return wf, λ₀, σ, err_opt, ps + return wf, err_opt, ps end diff --git a/test/harmonic_oscillator.jl b/test/harmonic_oscillator.jl index 383ff6e..2b88719 100644 --- a/test/harmonic_oscillator.jl +++ b/test/harmonic_oscillator.jl @@ -52,5 +52,5 @@ Vee(wf, X::AbstractVector, ps, st) = 0 # sum(1/norm(X[i]-X[j]) for i = 1:length( ham = SumH(K, Vext, Vee) sam = MHSampler(wf, Nel, Δt = 1.0, burnin = 1000) -opt = VMC(1000, 0.05, lr_dc = 10) -gd_GradientByVMC(opt, sam, ham, wf, ps, st) \ No newline at end of file +opt = VMC(100, 0.05, lr_dc = 10) +wf, err_opt, ps = gd_GradientByVMC(opt, sam, ham, wf, ps, st) \ No newline at end of file diff --git a/test/test_bflow_lux1.jl b/test/test_bflow_lux1.jl index 7f9b0df..3fe29c5 100644 --- a/test/test_bflow_lux1.jl +++ b/test/test_bflow_lux1.jl @@ -2,7 +2,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow -using ACEpsi.vmc: evaluate, gradient, laplacian, grad_params +using ACEpsi.vmc: gradient, laplacian, grad_params using ACEbase.Testing: print_tf, fdtest using LuxCore using Lux From 5aec160c74e913145dd8057508a83d7cfbf05e06 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 14 Jun 2023 10:37:07 +0800 Subject: [PATCH 51/79] minor fix --- src/vmc/metropolis.jl | 9 +++------ test/harmonic_oscillator.jl | 2 +- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/src/vmc/metropolis.jl b/src/vmc/metropolis.jl index 2dc0b67..756b51e 100644 --- a/src/vmc/metropolis.jl +++ b/src/vmc/metropolis.jl @@ -43,13 +43,10 @@ function MHstep(r0, Ψx0, Nels::Int, sam::MHSampler, ps, st) - rand_sample(X::AbstractVector, u::Int, Δt::AbstractFloat) = begin - Y = copy(X) - ind = sample(1:length(X), u, replace=false) - Y[ind] = X[ind] + Δt * randn(SVector{3, Float64}, u) - return Y + rand_sample(X::AbstractVector, Nels::Int, Δt::AbstractFloat) = begin + return X + Δt * randn(SVector{3, Float64}, Nels) end - rp = rand_sample.(r0, Ref(sam.type), Ref(sam.Δt)) + rp = rand_sample.(r0, Ref(Nels), Ref(sam.Δt)) Ψxp = eval.(Ref(sam.Ψ), rp, Ref(ps), Ref(st)) accprob = accfcn(Ψx0, Ψxp) u = rand(sam.nchains) diff --git a/test/harmonic_oscillator.jl b/test/harmonic_oscillator.jl index 2b88719..13aacdc 100644 --- a/test/harmonic_oscillator.jl +++ b/test/harmonic_oscillator.jl @@ -52,5 +52,5 @@ Vee(wf, X::AbstractVector, ps, st) = 0 # sum(1/norm(X[i]-X[j]) for i = 1:length( ham = SumH(K, Vext, Vee) sam = MHSampler(wf, Nel, Δt = 1.0, burnin = 1000) -opt = VMC(100, 0.05, lr_dc = 10) +opt = VMC(200, 0.05, lr_dc = 10) wf, err_opt, ps = gd_GradientByVMC(opt, sam, ham, wf, ps, st) \ No newline at end of file From 10c5c6876d8c9d2809661aad40e3a7e98e78aaa9 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Thu, 15 Jun 2023 05:24:40 -0700 Subject: [PATCH 52/79] fix BFwf and performacen fix --- src/bflow3d.jl | 10 +++++----- test/test_bflow_lux.jl | 6 +++--- test/test_bflow_lux1.jl | 2 +- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 63ab4f6..76ad137 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -29,17 +29,17 @@ struct DenseLayer <: AbstractExplicitLayer end function (l::DenseLayer)(x::AbstractMatrix, ps, st) - return x * ps.W, st + return Matrix(x) * ps.W, st end # Jerry: Maybe we should use Glorot Uniform if we have no idea about what we should use? LuxCore.initialparameters(rng::AbstractRNG, l::DenseLayer) = ( W = randn(rng, l.out_dim, l.in_dim), ) LuxCore.initialstates(rng::AbstractRNG, l::DenseLayer) = NamedTuple() -function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMatrix{Hyper{T}}, ps, st) where T<:Number +function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMatrix, ps, st) val = l(x, ps, st) function pb(A) - return NoTangent(), NoTangent(), NoTangent(), (W = x * A[1]',), NoTangent() + return NoTangent(), NoTangent(), A[1] * ps.W', (W = x' * A[1],), NoTangent() end return val, pb end @@ -93,9 +93,9 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), - bAA = corr_layer, hidden1 = DenseLayer(Nel,length(corr1)), + bAA = corr_layer, hidden1 = DenseLayer(Nel, length(corr1)), Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) - return Chain(; branch = BranchLayer((js = jastrow_layer, bf = BFwf_chain, )), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ) + return Chain(; branch = BranchLayer(; js = jastrow_layer, bf = BFwf_chain, ), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ) end diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index c6c99a9..49e56f0 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -2,7 +2,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow -using ACEpsi.vmc: evaluate, gradient, laplacian, grad_params +using ACEpsi.vmc: gradient, laplacian, grad_params using ACEbase.Testing: print_tf, fdtest using LuxCore using Lux @@ -12,7 +12,7 @@ using Random using Printf using LinearAlgebra using BenchmarkTools -using Cthulhu +# using Cthulhu using HyperDualNumbers: Hyper @@ -70,7 +70,7 @@ F(X) = BFwf_chain(X, ps, st)[1] # end # end -@btime F(X) +# @btime F(X) ## diff --git a/test/test_bflow_lux1.jl b/test/test_bflow_lux1.jl index 3fe29c5..0ff566b 100644 --- a/test/test_bflow_lux1.jl +++ b/test/test_bflow_lux1.jl @@ -74,7 +74,7 @@ F(X) = BFwf_chain(X, ps, st)[1] # end # end -@btime F(X) +# @btime F(X) ## From 83e65d9448b070321edd596a38f4466142952018 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Thu, 15 Jun 2023 07:17:29 -0700 Subject: [PATCH 53/79] fix dense layer with HyperDualsNumberExt with Octavian --- src/ACEpsi.jl | 4 + src/bflow3d.jl | 3 +- src/ext/HyperDualNumbersExt.jl | 235 +++++++++++++++++++++++++++++++++ 3 files changed, 241 insertions(+), 1 deletion(-) create mode 100644 src/ext/HyperDualNumbersExt.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index cb2934b..bd5d79a 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -22,4 +22,8 @@ include("lux_utils.jl") # vmc include("vmc/optimiser.jl") + +# HyperDualNumbersExt +include("ext/HyperDualNumbersExt.jl") + end diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 76ad137..7d08fc8 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -6,6 +6,7 @@ using ACEcore.Utils: gensparse using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det import ForwardDiff using ACEpsi.AtomicOrbitals: make_nlms_spec +using ACEpsi using LuxCore: AbstractExplicitLayer using LuxCore using Lux @@ -29,7 +30,7 @@ struct DenseLayer <: AbstractExplicitLayer end function (l::DenseLayer)(x::AbstractMatrix, ps, st) - return Matrix(x) * ps.W, st + return parent(x) * ps.W, st end # Jerry: Maybe we should use Glorot Uniform if we have no idea about what we should use? diff --git a/src/ext/HyperDualNumbersExt.jl b/src/ext/HyperDualNumbersExt.jl new file mode 100644 index 0000000..000f54c --- /dev/null +++ b/src/ext/HyperDualNumbersExt.jl @@ -0,0 +1,235 @@ +module HyperDualNumbersExt + +using HyperDualNumbers: Hyper + +using Octavian: ArrayInterface, + @turbo, @tturbo, + One, Zero, + indices, static +import Octavian: real_rep, _matmul!, _matmul_serial! + +real_rep(a::AbstractArray{DualT}) where {T,DualT<:Hyper{T}} = + reinterpret(reshape, T, a) +_view1(B::AbstractMatrix) = @view(B[1, :]) +_view1(B::AbstractArray{<:Any,3}) = @view(B[1, :, :]) + +for AbstractVectorOrMatrix in (:AbstractVector, :AbstractMatrix) + # multiplication of dual vector/matrix by standard matrix from the left + @eval function _matmul!( + _C::$(AbstractVectorOrMatrix){DualT}, + A::AbstractMatrix, + _B::$(AbstractVectorOrMatrix){DualT}, + α, + β = Zero(), + nthread::Nothing = nothing, + MKN = nothing, + contig_axis = nothing + ) where {DualT<:Hyper} + B = real_rep(_B) + C = real_rep(_C) + + @tturbo for n ∈ indices((C, B), 3), + m ∈ indices((C, A), (2, 1)), + l in indices((C, B), 1) + + Cₗₘₙ = zero(eltype(C)) + for k ∈ indices((A, B), 2) + Cₗₘₙ += A[m, k] * B[l, k, n] + end + C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] + end + + _C + end + + # multiplication of dual matrix by standard vector/matrix from the right + @eval @inline function _matmul!( + _C::$(AbstractVectorOrMatrix){DualT}, + _A::AbstractMatrix{DualT}, + B::$(AbstractVectorOrMatrix), + α = One(), + β = Zero(), + nthread::Nothing = nothing, + MKN = nothing + ) where {T,DualT<:Hyper{T}} + if Bool(ArrayInterface.is_dense(_C)) && + Bool(ArrayInterface.is_column_major(_C)) && + Bool(ArrayInterface.is_dense(_A)) && + Bool(ArrayInterface.is_column_major(_A)) + # we can avoid the reshape and call the standard method + A = reinterpret(T, _A) + C = reinterpret(T, _C) + _matmul!(C, A, B, α, β, nthread, nothing) + else + # we cannot use the standard method directly + A = real_rep(_A) + C = real_rep(_C) + + @tturbo for n ∈ indices((C, B), (3, 2)), + m ∈ indices((C, A), 2), + l in indices((C, A), 1) + + Cₗₘₙ = zero(eltype(C)) + for k ∈ indices((A, B), (3, 1)) + Cₗₘₙ += A[l, m, k] * B[k, n] + end + C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] + end + end + + _C + end + + @eval @inline function _matmul!( + _C::$(AbstractVectorOrMatrix){DualT}, + _A::AbstractMatrix{DualT}, + _B::$(AbstractVectorOrMatrix){DualT}, + α = One(), + β = Zero(), + nthread::Nothing = nothing, + MKN = nothing, + contig = nothing + ) where {T,DualT<:Hyper{T}} + A = real_rep(_A) + C = real_rep(_C) + B = real_rep(_B) + if Bool(ArrayInterface.is_dense(_C)) && + Bool(ArrayInterface.is_column_major(_C)) && + Bool(ArrayInterface.is_dense(_A)) && + Bool(ArrayInterface.is_column_major(_A)) + # we can avoid the reshape and call the standard method + Ar = reinterpret(T, _A) + Cr = reinterpret(T, _C) + _matmul!(Cr, Ar, _view1(B), α, β, nthread, nothing) + else + # we cannot use the standard method directly + @tturbo for n ∈ indices((C, B), 3), + m ∈ indices((C, A), 2), + l in indices((C, A), 1) + + Cₗₘₙ = zero(eltype(C)) + for k ∈ indices((A, B), (3, 2)) + Cₗₘₙ += A[l, m, k] * B[1, k, n] + end + C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] + end + end + Pstatic = static(P) + @tturbo for n ∈ indices((B, C), 3), m ∈ indices((A, C), 2), p ∈ 1:Pstatic + Cₚₘₙ = zero(eltype(C)) + for k ∈ indices((A, B), (3, 2)) + Cₚₘₙ += A[1, m, k] * B[p+1, k, n] + end + C[p+1, m, n] = C[p+1, m, n] + α * Cₚₘₙ + end + _C + end + + # multiplication of dual vector/matrix by standard matrix from the left + @eval function _matmul_serial!( + _C::$(AbstractVectorOrMatrix){DualT}, + A::AbstractMatrix, + _B::$(AbstractVectorOrMatrix){DualT}, + α, + β, + MKN + ) where {DualT<:Hyper} + B = real_rep(_B) + C = real_rep(_C) + + @turbo for n ∈ indices((C, B), 3), + m ∈ indices((C, A), (2, 1)), + l in indices((C, B), 1) + + Cₗₘₙ = zero(eltype(C)) + for k ∈ indices((A, B), 2) + Cₗₘₙ += A[m, k] * B[l, k, n] + end + C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] + end + + _C + end + +# multiplication of dual matrix by standard vector/matrix from the right + @eval @inline function _matmul_serial!( + _C::$(AbstractVectorOrMatrix){DualT}, + _A::AbstractMatrix{DualT}, + B::$(AbstractVectorOrMatrix), + α, + β, + MKN + ) where {T,DualT<:Hyper{T}} + if Bool(ArrayInterface.is_dense(_C)) && + Bool(ArrayInterface.is_column_major(_C)) && + Bool(ArrayInterface.is_dense(_A)) && + Bool(ArrayInterface.is_column_major(_A)) + # we can avoid the reshape and call the standard method + A = reinterpret(T, _A) + C = reinterpret(T, _C) + _matmul_serial!(C, A, B, α, β, nothing) + else + # we cannot use the standard method directly + A = real_rep(_A) + C = real_rep(_C) + + @turbo for n ∈ indices((C, B), (3, 2)), + m ∈ indices((C, A), 2), + l in indices((C, A), 1) + + Cₗₘₙ = zero(eltype(C)) + for k ∈ indices((A, B), (3, 1)) + Cₗₘₙ += A[l, m, k] * B[k, n] + end + C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] + end + end + + _C + end + + @eval @inline function _matmul_serial!( + _C::$(AbstractVectorOrMatrix){DualT}, + _A::AbstractMatrix{DualT}, + _B::$(AbstractVectorOrMatrix){DualT}, + α, + β, + MKN + ) where {DualT<:Hyper} + A = real_rep(_A) + C = real_rep(_C) + B = real_rep(_B) + if Bool(ArrayInterface.is_dense(_C)) && + Bool(ArrayInterface.is_column_major(_C)) && + Bool(ArrayInterface.is_dense(_A)) && + Bool(ArrayInterface.is_column_major(_A)) + # we can avoid the reshape and call the standard method + Ar = reinterpret(T, _A) + Cr = reinterpret(T, _C) + _matmul_serial!(Cr, Ar, _view1(B), α, β, nothing) + else + # we cannot use the standard method directly + @turbo for n ∈ indices((C, B), 3), + m ∈ indices((C, A), 2), + l in indices((C, A), 1) + + Cₗₘₙ = zero(eltype(C)) + for k ∈ indices((A, B), (3, 2)) + Cₗₘₙ += A[l, m, k] * B[1, k, n] + end + C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] + end + end + Pstatic = static(P) + @turbo for n ∈ indices((B, C), 3), m ∈ indices((A, C), 2), p ∈ 1:Pstatic + Cₚₘₙ = zero(eltype(C)) + for k ∈ indices((A, B), (3, 2)) + Cₚₘₙ += A[1, m, k] * B[p+1, k, n] + end + C[p+1, m, n] = C[p+1, m, n] + α * Cₚₘₙ + end + _C + end +end # for + +end # module \ No newline at end of file From fd76d44d8c28607751d35299f446232f40ee9664 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Tue, 20 Jun 2023 13:13:06 +0800 Subject: [PATCH 54/79] minor fix --- Project.toml | 1 + src/bflow3d.jl | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index af2c47c..4f5e5db 100644 --- a/Project.toml +++ b/Project.toml @@ -17,6 +17,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d" +Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 7d08fc8..1504498 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -30,7 +30,7 @@ struct DenseLayer <: AbstractExplicitLayer end function (l::DenseLayer)(x::AbstractMatrix, ps, st) - return parent(x) * ps.W, st + return Matrix(parent(x) * ps.W), st end # Jerry: Maybe we should use Glorot Uniform if we have no idea about what we should use? @@ -40,7 +40,7 @@ LuxCore.initialstates(rng::AbstractRNG, l::DenseLayer) = NamedTuple() function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMatrix, ps, st) val = l(x, ps, st) function pb(A) - return NoTangent(), NoTangent(), A[1] * ps.W', (W = x' * A[1],), NoTangent() + return NoTangent(), NoTangent(), Matrix(A[1] * ps.W'), (W = x' * A[1],), NoTangent() end return val, pb end From 6c27a5aa7f4fb0a22346236c192d2af7a7f9bbc4 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 20 Jun 2023 06:41:22 -0700 Subject: [PATCH 55/79] remove Octavian extension --- src/ACEpsi.jl | 3 - src/ext/HyperDualNumbersExt.jl | 235 --------------------------------- 2 files changed, 238 deletions(-) delete mode 100644 src/ext/HyperDualNumbersExt.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index bd5d79a..51ee78c 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -23,7 +23,4 @@ include("lux_utils.jl") # vmc include("vmc/optimiser.jl") -# HyperDualNumbersExt -include("ext/HyperDualNumbersExt.jl") - end diff --git a/src/ext/HyperDualNumbersExt.jl b/src/ext/HyperDualNumbersExt.jl deleted file mode 100644 index 000f54c..0000000 --- a/src/ext/HyperDualNumbersExt.jl +++ /dev/null @@ -1,235 +0,0 @@ -module HyperDualNumbersExt - -using HyperDualNumbers: Hyper - -using Octavian: ArrayInterface, - @turbo, @tturbo, - One, Zero, - indices, static -import Octavian: real_rep, _matmul!, _matmul_serial! - -real_rep(a::AbstractArray{DualT}) where {T,DualT<:Hyper{T}} = - reinterpret(reshape, T, a) -_view1(B::AbstractMatrix) = @view(B[1, :]) -_view1(B::AbstractArray{<:Any,3}) = @view(B[1, :, :]) - -for AbstractVectorOrMatrix in (:AbstractVector, :AbstractMatrix) - # multiplication of dual vector/matrix by standard matrix from the left - @eval function _matmul!( - _C::$(AbstractVectorOrMatrix){DualT}, - A::AbstractMatrix, - _B::$(AbstractVectorOrMatrix){DualT}, - α, - β = Zero(), - nthread::Nothing = nothing, - MKN = nothing, - contig_axis = nothing - ) where {DualT<:Hyper} - B = real_rep(_B) - C = real_rep(_C) - - @tturbo for n ∈ indices((C, B), 3), - m ∈ indices((C, A), (2, 1)), - l in indices((C, B), 1) - - Cₗₘₙ = zero(eltype(C)) - for k ∈ indices((A, B), 2) - Cₗₘₙ += A[m, k] * B[l, k, n] - end - C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] - end - - _C - end - - # multiplication of dual matrix by standard vector/matrix from the right - @eval @inline function _matmul!( - _C::$(AbstractVectorOrMatrix){DualT}, - _A::AbstractMatrix{DualT}, - B::$(AbstractVectorOrMatrix), - α = One(), - β = Zero(), - nthread::Nothing = nothing, - MKN = nothing - ) where {T,DualT<:Hyper{T}} - if Bool(ArrayInterface.is_dense(_C)) && - Bool(ArrayInterface.is_column_major(_C)) && - Bool(ArrayInterface.is_dense(_A)) && - Bool(ArrayInterface.is_column_major(_A)) - # we can avoid the reshape and call the standard method - A = reinterpret(T, _A) - C = reinterpret(T, _C) - _matmul!(C, A, B, α, β, nthread, nothing) - else - # we cannot use the standard method directly - A = real_rep(_A) - C = real_rep(_C) - - @tturbo for n ∈ indices((C, B), (3, 2)), - m ∈ indices((C, A), 2), - l in indices((C, A), 1) - - Cₗₘₙ = zero(eltype(C)) - for k ∈ indices((A, B), (3, 1)) - Cₗₘₙ += A[l, m, k] * B[k, n] - end - C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] - end - end - - _C - end - - @eval @inline function _matmul!( - _C::$(AbstractVectorOrMatrix){DualT}, - _A::AbstractMatrix{DualT}, - _B::$(AbstractVectorOrMatrix){DualT}, - α = One(), - β = Zero(), - nthread::Nothing = nothing, - MKN = nothing, - contig = nothing - ) where {T,DualT<:Hyper{T}} - A = real_rep(_A) - C = real_rep(_C) - B = real_rep(_B) - if Bool(ArrayInterface.is_dense(_C)) && - Bool(ArrayInterface.is_column_major(_C)) && - Bool(ArrayInterface.is_dense(_A)) && - Bool(ArrayInterface.is_column_major(_A)) - # we can avoid the reshape and call the standard method - Ar = reinterpret(T, _A) - Cr = reinterpret(T, _C) - _matmul!(Cr, Ar, _view1(B), α, β, nthread, nothing) - else - # we cannot use the standard method directly - @tturbo for n ∈ indices((C, B), 3), - m ∈ indices((C, A), 2), - l in indices((C, A), 1) - - Cₗₘₙ = zero(eltype(C)) - for k ∈ indices((A, B), (3, 2)) - Cₗₘₙ += A[l, m, k] * B[1, k, n] - end - C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] - end - end - Pstatic = static(P) - @tturbo for n ∈ indices((B, C), 3), m ∈ indices((A, C), 2), p ∈ 1:Pstatic - Cₚₘₙ = zero(eltype(C)) - for k ∈ indices((A, B), (3, 2)) - Cₚₘₙ += A[1, m, k] * B[p+1, k, n] - end - C[p+1, m, n] = C[p+1, m, n] + α * Cₚₘₙ - end - _C - end - - # multiplication of dual vector/matrix by standard matrix from the left - @eval function _matmul_serial!( - _C::$(AbstractVectorOrMatrix){DualT}, - A::AbstractMatrix, - _B::$(AbstractVectorOrMatrix){DualT}, - α, - β, - MKN - ) where {DualT<:Hyper} - B = real_rep(_B) - C = real_rep(_C) - - @turbo for n ∈ indices((C, B), 3), - m ∈ indices((C, A), (2, 1)), - l in indices((C, B), 1) - - Cₗₘₙ = zero(eltype(C)) - for k ∈ indices((A, B), 2) - Cₗₘₙ += A[m, k] * B[l, k, n] - end - C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] - end - - _C - end - -# multiplication of dual matrix by standard vector/matrix from the right - @eval @inline function _matmul_serial!( - _C::$(AbstractVectorOrMatrix){DualT}, - _A::AbstractMatrix{DualT}, - B::$(AbstractVectorOrMatrix), - α, - β, - MKN - ) where {T,DualT<:Hyper{T}} - if Bool(ArrayInterface.is_dense(_C)) && - Bool(ArrayInterface.is_column_major(_C)) && - Bool(ArrayInterface.is_dense(_A)) && - Bool(ArrayInterface.is_column_major(_A)) - # we can avoid the reshape and call the standard method - A = reinterpret(T, _A) - C = reinterpret(T, _C) - _matmul_serial!(C, A, B, α, β, nothing) - else - # we cannot use the standard method directly - A = real_rep(_A) - C = real_rep(_C) - - @turbo for n ∈ indices((C, B), (3, 2)), - m ∈ indices((C, A), 2), - l in indices((C, A), 1) - - Cₗₘₙ = zero(eltype(C)) - for k ∈ indices((A, B), (3, 1)) - Cₗₘₙ += A[l, m, k] * B[k, n] - end - C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] - end - end - - _C - end - - @eval @inline function _matmul_serial!( - _C::$(AbstractVectorOrMatrix){DualT}, - _A::AbstractMatrix{DualT}, - _B::$(AbstractVectorOrMatrix){DualT}, - α, - β, - MKN - ) where {DualT<:Hyper} - A = real_rep(_A) - C = real_rep(_C) - B = real_rep(_B) - if Bool(ArrayInterface.is_dense(_C)) && - Bool(ArrayInterface.is_column_major(_C)) && - Bool(ArrayInterface.is_dense(_A)) && - Bool(ArrayInterface.is_column_major(_A)) - # we can avoid the reshape and call the standard method - Ar = reinterpret(T, _A) - Cr = reinterpret(T, _C) - _matmul_serial!(Cr, Ar, _view1(B), α, β, nothing) - else - # we cannot use the standard method directly - @turbo for n ∈ indices((C, B), 3), - m ∈ indices((C, A), 2), - l in indices((C, A), 1) - - Cₗₘₙ = zero(eltype(C)) - for k ∈ indices((A, B), (3, 2)) - Cₗₘₙ += A[l, m, k] * B[1, k, n] - end - C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] - end - end - Pstatic = static(P) - @turbo for n ∈ indices((B, C), 3), m ∈ indices((A, C), 2), p ∈ 1:Pstatic - Cₚₘₙ = zero(eltype(C)) - for k ∈ indices((A, B), (3, 2)) - Cₚₘₙ += A[1, m, k] * B[p+1, k, n] - end - C[p+1, m, n] = C[p+1, m, n] + α * Cₚₘₙ - end - _C - end -end # for - -end # module \ No newline at end of file From afa90ab9857f19e40327d0dc9002c49fd79bc678 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 26 Jun 2023 10:08:42 +0800 Subject: [PATCH 56/79] minor change of js factor --- src/bflow3d.jl | 2 +- src/jastrow.jl | 62 ++++++++++++++++++++++++------------- src/vmc/metropolis.jl | 8 ++--- src/vmc/vmc.jl | 8 ++++- test/He.jl | 48 ++++++++++++++++++++++++++++ test/harmonic_oscillator.jl | 29 ++++++++--------- 6 files changed, 114 insertions(+), 43 deletions(-) create mode 100644 test/He.jl diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 1504498..202e51b 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -88,7 +88,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) corr_layer = Polynomials4ML.lux(corr1) - js = Jastrow(nuclei) + js = JPauliNet(nuclei) jastrow_layer = ACEpsi.lux(js) reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) diff --git a/src/jastrow.jl b/src/jastrow.jl index 0a0e5a2..124ad01 100644 --- a/src/jastrow.jl +++ b/src/jastrow.jl @@ -1,11 +1,9 @@ # This should probably be re-implemented properly and be made performant. # the current version is just for testing using ACEpsi.AtomicOrbitals: Nuc +using LuxCore using LuxCore: AbstractExplicitLayer using Random: AbstractRNG - -import ChainRules -import ChainRules: rrule, NoTangent using Zygote: Buffer mutable struct Jastrow{T} @@ -17,32 +15,22 @@ end ## F_2(x) = -1/2\sum_{l=1}^L \sum_{i=1}^N Z_l|yi,l|+1/4\sum_{1\leq i Date: Mon, 26 Jun 2023 10:09:55 +0800 Subject: [PATCH 57/79] add example of jastrow factor --- test/try.jl | 85 ++++++++++++++++++----------------------------------- 1 file changed, 28 insertions(+), 57 deletions(-) diff --git a/test/try.jl b/test/try.jl index 4e94c5e..5c53a49 100644 --- a/test/try.jl +++ b/test/try.jl @@ -1,74 +1,45 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct -using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasis, Nuc, make_nlms_spec, ProductBasis, evaluate -using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow -using ACEbase.Testing: print_tf +using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, JPauliNet +using ACEbase.Testing: print_tf, fdtest +using ACEpsi.vmc: gradient, laplacian, grad_params +using ACEbase.Testing: print_tf, fdtest using LuxCore using Lux using Zygote +using Optimisers # mainly for the destrcuture(ps) function using Random -using ACEcore.Utils: gensparse -using ACEcore: SparseSymmProd -using ACEcore +using Printf +using LinearAlgebra +using BenchmarkTools + +using HyperDualNumbers: Hyper + -Rnldegree = 4 -Ylmdegree = 4 -ν = 2 -totdeg = 10 Nel = 5 X = randn(SVector{3, Float64}, Nel) Σ = rand(spins(), Nel) -sd_admissible = bb -> (true) - - nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] -## - -# Defining AtomicOrbitalsBasis -bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) -bYlm = RYlmBasis(Ylmdegree) - -spec1p = make_nlms_spec(bRnl, bYlm; - totaldegree = totdeg) - -# size(X) = (nX, 3); length(Σ) = nX -aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) -pooling = BackflowPooling(aobasis) - -# define sparse for n-correlations -tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] -default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) - -specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, - minvv = fill(0, ν), - maxvv = fill(length(spec1p), ν), - ordered = true) -spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] - -# further restrict -spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] - -# define n-correlation -corr1 = SparseSymmProd(spec; T = Float64) - -reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) - - -pooling_layer = ACEpsi.lux(pooling) -# (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) -corr_layer = ACEcore.lux(corr1) -pool2AAChain = Chain(; pooling = pooling_layer, reshape = WrappedFunction(reshape_func), -bAA = corr_layer) +# wrap it as HyperDualNumbers +x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) +hX = [x2dualwrtj(x, 0) for x in X] +hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec +js = JPauliNet(nuclei) +jastrow_layer = ACEpsi.lux(js) +ps, st = LuxCore.setup(MersenneTwister(1234), jastrow_layer) +st = (Σ = Σ,) -# dummy input -ϕnlm = aobasis(X, Σ) +A1 = jastrow_layer(X, ps, st) +hA1 = jastrow_layer(hX, ps, st) +print_tf(@test hA1[1].value ≈ A1[1]) -ps, st = setupBFState(MersenneTwister(1234), pool2AAChain, Σ) +@info("Test ∇ψ w.r.t. X") +y, st = Lux.apply(jastrow_layer, X, ps, st) -y, st = Lux.apply(pool2AAChain, ϕnlm, ps, st) +F(X) = jastrow_layer(X, ps, st)[1] +dF(X) = Zygote.gradient(x -> jastrow_layer(x, ps, st)[1], X)[1] +fdtest(F, dF, X, verbose = true) -## Pullback API to capture change in state -(l, st_), pb = pullback(p -> Lux.apply(pool2AAChain, ϕnlm, p, st), ps) -gs = pb((one.(l), nothing))[1] From 406bee8a81796a50b36e370c2d5676f19795d6f0 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 26 Jun 2023 21:44:15 +0800 Subject: [PATCH 58/79] minor fix --- src/backflowpooling.jl | 6 ++--- src/bflow3d.jl | 2 +- src/jastrow.jl | 20 ++++++++++------ test/Be.jl | 53 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 70 insertions(+), 11 deletions(-) create mode 100644 test/Be.jl diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index 0882d84..9a4fdff 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -99,7 +99,7 @@ function _rrule_evaluate(pooling::BackflowPooling, ϕnlm, Σ) end function _pullback_evaluate(∂A, pooling::BackflowPooling, ϕnlm, Σ) - TA = promote_type(eltype.(ϕnlm)...) + TA = eltype(ϕnlm) ∂ϕnlm = zeros(TA, size(ϕnlm)) _pullback_evaluate!(∂ϕnlm, ∂A, pooling, ϕnlm, Σ) return ∂ϕnlm @@ -108,9 +108,9 @@ end function _pullback_evaluate!(∂ϕnlm, ∂A, pooling::BackflowPooling, ϕnlm, Σ) Nnuc, Nel, Nnlm = size(ϕnlm) - basis = pooling.basis + #basis = pooling.basis - @assert Nnlm == length(basis.prodbasis.layers.ϕnlms.basis.spec) + #@assert Nnlm == length(basis.prodbasis.layers.ϕnlms.basis.spec) @assert Nel == length(Σ) @assert size(∂ϕnlm) == (Nnuc, Nel, Nnlm) @assert size(∂A) == (Nel, 3, Nnuc, Nnlm) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 202e51b..1504498 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -88,7 +88,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) corr_layer = Polynomials4ML.lux(corr1) - js = JPauliNet(nuclei) + js = Jastrow(nuclei) jastrow_layer = ACEpsi.lux(js) reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) diff --git a/src/jastrow.jl b/src/jastrow.jl index 124ad01..d509201 100644 --- a/src/jastrow.jl +++ b/src/jastrow.jl @@ -15,22 +15,28 @@ end ## F_2(x) = -1/2\sum_{l=1}^L \sum_{i=1}^N Z_l|yi,l|+1/4\sum_{1\leq i Date: Mon, 26 Jun 2023 23:02:33 +0800 Subject: [PATCH 59/79] minor fix --- src/backflowpooling.jl | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index 9a4fdff..ae3caa3 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -27,7 +27,7 @@ function evaluate(pooling::BackflowPooling, ϕnlm::AbstractArray, Σ::AbstractVe # evaluate the pooling operation # spin I k = (nlm) - Aall = acquire!(pooling.pool, :Aall, (2, Nnuc, Nnlm), T) + Aall = acquire!(pooling.tmp, :Aall, (2, Nnuc, Nnlm), T) fill!(Aall, 0) @inbounds begin @@ -60,7 +60,7 @@ function evaluate(pooling::BackflowPooling, ϕnlm::AbstractArray, Σ::AbstractVe @assert spin2idx(↓) == 2 @assert spin2idx(∅) == 3 - A = acquire!(pooling.tmp, :Aall, (Nel, 3, Nnuc, Nnlm), T) + A = acquire!(pooling.pool, :Aall, (Nel, 3, Nnuc, Nnlm), T) fill!(A, 0) @inbounds begin @@ -100,7 +100,8 @@ end function _pullback_evaluate(∂A, pooling::BackflowPooling, ϕnlm, Σ) TA = eltype(ϕnlm) - ∂ϕnlm = zeros(TA, size(ϕnlm)) + ∂ϕnlm = acquire!(pooling.pool, :∂ϕnlm, size(ϕnlm), TA) + fill!(∂ϕnlm, zero(TA)) _pullback_evaluate!(∂ϕnlm, ∂A, pooling, ϕnlm, Σ) return ∂ϕnlm end From 04cd372abf1b69b8ed5fb137df9843710265c128 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 28 Jun 2023 20:48:03 +0800 Subject: [PATCH 60/79] minor fix --- src/vmc/metropolis.jl | 2 +- test/Be.jl | 12 ++++-------- test/He.jl | 7 ++++--- 3 files changed, 9 insertions(+), 12 deletions(-) diff --git a/src/vmc/metropolis.jl b/src/vmc/metropolis.jl index 7b1b929..f130af6 100644 --- a/src/vmc/metropolis.jl +++ b/src/vmc/metropolis.jl @@ -44,7 +44,7 @@ function MHstep(r0, Nels::Int, sam::MHSampler, ps, st) rand_sample(X::AbstractVector, Nels::Int, Δt::AbstractFloat) = begin - return X + Δt * randn(SVector{3, Float64}, Nels) + return X + Δt * randn(SVector{3, eltype(X[1])}, Nels) end rp = rand_sample.(r0, Ref(Nels), Ref(sam.Δt)) Ψxp = eval.(Ref(sam.Ψ), rp, Ref(ps), Ref(st)) diff --git a/test/Be.jl b/test/Be.jl index 721da22..0632291 100644 --- a/test/Be.jl +++ b/test/Be.jl @@ -43,11 +43,7 @@ Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) ham = SumH(K, Vext, Vee) -sam = MHSampler(wf, Nel, Δt = 0.5, burnin = 1000) -opt = VMC(10, 0.015, lr_dc = 100) - -using BenchmarkTools -@btime wf(X, ps, st) -@btime gradient(wf, X, ps, st) - -@profview wf, err_opt, ps = gd_GradientByVMC(opt, sam, ham, wf, ps, st) \ No newline at end of file +sam = MHSampler(wf, Nel, Δt = 0.4, burnin = 1000, nchains = 2000) +opt = VMC(500, 0.1, lr_dc = 10) +wf, err_opt, ps = gd_GradientByVMC(opt, sam, ham, wf, ps, st) +#Eref = -14.667 \ No newline at end of file diff --git a/test/He.jl b/test/He.jl index 9cdd635..13433a9 100644 --- a/test/He.jl +++ b/test/He.jl @@ -43,6 +43,7 @@ Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) ham = SumH(K, Vext, Vee) -sam = MHSampler(wf, Nel, Δt = 0.6, burnin = 1000) -opt = VMC(1000, 0.1, lr_dc = 50) -wf, err_opt, ps = gd_GradientByVMC(opt, sam, ham, wf, ps, st) \ No newline at end of file +sam = MHSampler(wf, Nel, Δt = 0.3, burnin = 1000, nchains = 2000) +opt = VMC(500, 0.1, lr_dc = 50) +wf, err_opt, ps = gd_GradientByVMC(opt, sam, ham, wf, ps, st) +#Eref = -2.9037 From 00126e4e3aaac4307f822a12931f0d5cb2fc07b9 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Wed, 28 Jun 2023 23:17:49 +0800 Subject: [PATCH 61/79] small change of spec --- src/bflow3d.jl | 56 +++++++++++++++++++++++++++++++++++--------------- src/jastrow.jl | 14 ++++--------- test/He.jl | 11 +++++++--- 3 files changed, 51 insertions(+), 30 deletions(-) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 1504498..a9463ba 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -6,6 +6,7 @@ using ACEcore.Utils: gensparse using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det import ForwardDiff using ACEpsi.AtomicOrbitals: make_nlms_spec +using ACEpsi: ↑, ↓, ∅, spins, extspins, Spin, spin2idx, idx2spin using ACEpsi using LuxCore: AbstractExplicitLayer using LuxCore @@ -47,44 +48,58 @@ end ## +function get_spec(nuclei, spec1p) + spec = [] + Nnuc = length(nuclei) + + spec = Array{Any}(undef, (3, Nnuc, length(spec1p))) + + for (k, nlm) in enumerate(spec1p) + for I = 1:Nnuc + for (is, s) in enumerate(extspins()) + spec[is, I, k] = (s=s, I = I, nlm...) + end + end + end + + return spec[:] +end + # ----------------- custom layers ------------------ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, ν = 3, T = Float64, - sd_admissible = bb -> (true), - envelope = x -> x) # enveolpe to be replaced by SJ-factor + sd_admissible = bb -> prod(b.s != '∅' for b in bb) == 0) spec1p = make_nlms_spec(bRnl, bYlm; totaldegree = totdeg) - # size(X) = (nX, 3); length(Σ) = nX - # aobasis = AtomicOrbitalsBasis(bRnl, bYlm; totaldegree = totdeg, nuclei = nuclei, ) + # ----------- Lux connections --------- + # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1)) + prodbasis_layer = ACEpsi.AtomicOrbitals.ProductBasisLayer(spec1p, bRnl, bYlm) + aobasis_layer = ACEpsi.AtomicOrbitals.AtomicOrbitalsBasisLayer(prodbasis_layer, nuclei) + # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) + pooling = BackflowPooling(aobasis_layer) + pooling_layer = ACEpsi.lux(pooling) + + spec1p = get_spec(nuclei, spec1p) # define sparse for n-correlations tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ] - default_admissible = bb -> (length(bb) == 0) || (sum(b[1] - 1 for b in bb ) <= totdeg) + default_admissible = bb -> (length(bb) == 0) || (sum(b.n1 - 1 for b in bb ) <= totdeg) specAA = gensparse(; NU = ν, tup2b = tup2b, admissible = default_admissible, minvv = fill(0, ν), maxvv = fill(length(spec1p), ν), ordered = true) spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))] - + # further restrict spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] - + # define n-correlation corr1 = Polynomials4ML.SparseSymmProd(spec) - # ----------- Lux connections --------- - # AtomicOrbitalsBasis: (X, Σ) -> (length(nuclei), nX, length(spec1)) - prodbasis_layer = ACEpsi.AtomicOrbitals.ProductBasisLayer(spec1p, bRnl, bYlm) - aobasis_layer = ACEpsi.AtomicOrbitals.AtomicOrbitalsBasisLayer(prodbasis_layer, nuclei) - - # BackFlowPooling: (length(nuclei), nX, length(spec1 from totaldegree)) -> (nX, 3, length(nuclei), length(spec1)) - pooling = BackflowPooling(aobasis_layer) - pooling_layer = ACEpsi.lux(pooling) - # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) corr_layer = Polynomials4ML.lux(corr1) @@ -96,9 +111,16 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), bAA = corr_layer, hidden1 = DenseLayer(Nel, length(corr1)), Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) - return Chain(; branch = BranchLayer(; js = jastrow_layer, bf = BFwf_chain, ), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ) + return Chain(; branch = BranchLayer(; js = jastrow_layer, bf = BFwf_chain, ), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ), spec, spec1p end +function displayspec(spec, spec1p) + nicespec = [] + for k = 1:length(spec) + push!(nicespec, ([spec1p[spec[k][j]] for j = 1:length(spec[k])])) + end + return nicespec +end # """ # This function returns correct Si for pooling operation. diff --git a/src/jastrow.jl b/src/jastrow.jl index d509201..ee5d9c2 100644 --- a/src/jastrow.jl +++ b/src/jastrow.jl @@ -15,12 +15,9 @@ end ## F_2(x) = -1/2\sum_{l=1}^L \sum_{i=1}^N Z_l|yi,l|+1/4\sum_{1\leq i Date: Thu, 29 Jun 2023 22:38:54 -0700 Subject: [PATCH 62/79] minor fix --- src/bflow3d.jl | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index a9463ba..537fd53 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -31,7 +31,7 @@ struct DenseLayer <: AbstractExplicitLayer end function (l::DenseLayer)(x::AbstractMatrix, ps, st) - return Matrix(parent(x) * ps.W), st + return parent(x) * ps.W, st end # Jerry: Maybe we should use Glorot Uniform if we have no idea about what we should use? @@ -41,7 +41,7 @@ LuxCore.initialstates(rng::AbstractRNG, l::DenseLayer) = NamedTuple() function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMatrix, ps, st) val = l(x, ps, st) function pb(A) - return NoTangent(), NoTangent(), Matrix(A[1] * ps.W'), (W = x' * A[1],), NoTangent() + return NoTangent(), NoTangent(), A[1] * ps.W', (W = x' * A[1],), NoTangent() end return val, pb end @@ -108,9 +108,10 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) + _det = x -> size(x) == (1, 1) ? x[1,1] : det(x) BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), bAA = corr_layer, hidden1 = DenseLayer(Nel, length(corr1)), - Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> det(x))) + Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> _det(x))) return Chain(; branch = BranchLayer(; js = jastrow_layer, bf = BFwf_chain, ), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ), spec, spec1p end From a71dc3187ca0afb4675928993b23138fc2a29312 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Fri, 30 Jun 2023 15:10:43 -0700 Subject: [PATCH 63/79] first step towards non-allocating lux interface... --- src/bflow3d.jl | 2 +- src/lux_utils.jl | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 537fd53..173f907 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -101,7 +101,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, corr1 = Polynomials4ML.SparseSymmProd(spec) # (nX, 3, length(nuclei), length(spec1 from totaldegree)) -> (nX, length(spec)) - corr_layer = Polynomials4ML.lux(corr1) + corr_layer = Polynomials4ML.lux(corr1; use_cache = false) js = Jastrow(nuclei) jastrow_layer = ACEpsi.lux(js) diff --git a/src/lux_utils.jl b/src/lux_utils.jl index 78822a3..bc02c93 100644 --- a/src/lux_utils.jl +++ b/src/lux_utils.jl @@ -2,8 +2,8 @@ using LuxCore using HyperDualNumbers function replace_namedtuples(nt, rp_st, Σ) - if :tmp in keys(nt) || length(nt) == 0 - return (; rp_st...) + if :tmp in keys(nt) || :pool in keys(nt) || length(nt) == 0 + return (; nt...) else for i in 1:length(nt) # @show nt[i] From 10385762363b670eab77957cccbb40b608f5518b Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Sun, 9 Jul 2023 20:54:12 +0800 Subject: [PATCH 64/79] add sr --- src/ACEpsi.jl | 2 +- src/bflow3d.jl | 2 +- src/vmc/opt.jl | 14 +++++++ src/vmc/optimiser.jl | 8 ---- src/vmc/optimisers/adamw.jl | 19 +++++++++ src/vmc/optimisers/sr.jl | 45 ++++++++++++++++++++ src/vmc/vmc.jl | 82 ++++++++++++++----------------------- src/vmc/vmc_utils.jl | 23 +++++++++++ test/Be.jl | 36 +++++++++++----- test/H2.jl | 64 +++++++++++++++++++++++++++++ test/He.jl | 33 ++++++++++----- 11 files changed, 245 insertions(+), 83 deletions(-) create mode 100644 src/vmc/opt.jl delete mode 100644 src/vmc/optimiser.jl create mode 100644 src/vmc/optimisers/adamw.jl create mode 100644 src/vmc/optimisers/sr.jl create mode 100644 src/vmc/vmc_utils.jl create mode 100644 test/H2.jl diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl index 51ee78c..73b6440 100644 --- a/src/ACEpsi.jl +++ b/src/ACEpsi.jl @@ -21,6 +21,6 @@ include("backflowpooling.jl") include("lux_utils.jl") # vmc -include("vmc/optimiser.jl") +include("vmc/opt.jl") end diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 173f907..7eb25b4 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -108,7 +108,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, reshape_func = x -> reshape(x, (size(x, 1), prod(size(x)[2:end]))) - _det = x -> size(x) == (1, 1) ? x[1,1] : det(x) + _det = x -> size(x) == (1, 1) ? x[1,1] : det(Matrix(x)) BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), bAA = corr_layer, hidden1 = DenseLayer(Nel, length(corr1)), Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> _det(x))) diff --git a/src/vmc/opt.jl b/src/vmc/opt.jl new file mode 100644 index 0000000..858409a --- /dev/null +++ b/src/vmc/opt.jl @@ -0,0 +1,14 @@ +module vmc + +abstract type opt end + +include("Eloc.jl") +include("gradient.jl") +include("metropolis.jl") + +include("vmc_utils.jl") +include("vmc.jl") +include("optimisers/adamw.jl") +include("optimisers/sr.jl") + +end \ No newline at end of file diff --git a/src/vmc/optimiser.jl b/src/vmc/optimiser.jl deleted file mode 100644 index 114f99f..0000000 --- a/src/vmc/optimiser.jl +++ /dev/null @@ -1,8 +0,0 @@ -module vmc - -include("Eloc.jl") -include("gradient.jl") -include("metropolis.jl") -include("vmc.jl") - -end \ No newline at end of file diff --git a/src/vmc/optimisers/adamw.jl b/src/vmc/optimisers/adamw.jl new file mode 100644 index 0000000..b38355e --- /dev/null +++ b/src/vmc/optimisers/adamw.jl @@ -0,0 +1,19 @@ +using Optimisers + + +mutable struct AdamW <: opt + β::Tuple + γ::Number + ϵ::Number +end + +AdamW() = AdamW((9f-1, 9.99f-1), 0.0, eps()) + +function Optimization(type::AdamW, wf, ps, st, sam::MHSampler, ham::SumH, α) + λ₀, σ, E, x0, acc = Eloc_Exp_TV_clip(wf, ps, st, sam, ham) + g = grad(wf, x0, ps, st, E) + st_opt = Optimisers.setup(Optimisers.AdamW(α, type.β, type.γ, type.ϵ), ps) + st_opt, ps = Optimisers.update(st_opt, ps, g) + res = norm(destructure(g)[1]) + return ps, acc, λ₀, res, σ +end \ No newline at end of file diff --git a/src/vmc/optimisers/sr.jl b/src/vmc/optimisers/sr.jl new file mode 100644 index 0000000..1c24a90 --- /dev/null +++ b/src/vmc/optimisers/sr.jl @@ -0,0 +1,45 @@ +using Optimisers + + +mutable struct SR <: opt + ϵ1::Number + ϵ2::Number +end + +SR() = SR(0., 0.01) + + +_destructure(ps) = destructure(ps)[1] + +function Optimization(type::SR, wf, ps, st, sam::MHSampler, ham::SumH, α) + ϵ1 = type.ϵ1 + ϵ2 = type.ϵ2 + + λ₀, σ, E, x0, acc = Eloc_Exp_TV_clip(wf, ps, st, sam, ham) + dy_ps = grad_params.(Ref(wf), x0, Ref(ps), Ref(st)) + dy = _destructure.(dy_ps) + O = 1/2 * dy + Ō = mean(O) + ΔO = [O[i] - Ō for i = 1:length(O)] + + g0 = 2.0 * sum(ΔO .*E)/ length(x0) + + S = sum([ΔO[i] * ΔO[i]' for i = 1:length(O)])/length(x0) + S = S + ϵ1 * Diagonal(diag(S)) + ϵ2 * Diagonal(diag(one(S))) + + g = S \ g0 + res = norm(g) + + p, s = destructure(ps) + p = p - α * g + ps = s(p) + + return ps, acc, λ₀, res, σ +end + + + + + + + \ No newline at end of file diff --git a/src/vmc/vmc.jl b/src/vmc/vmc.jl index 3cebde1..0c42b2f 100644 --- a/src/vmc/vmc.jl +++ b/src/vmc/vmc.jl @@ -2,74 +2,52 @@ export VMC using Printf using LinearAlgebra using Optimisers -#using Plots -""" -Minimizing the Rayleigh quotient by VMC -tol: tolrance -MaxIter: maximum iteration number -lr: learning rates -lr_dc_pow: rate to decrease 'lr' -- lr / (1 + k / lr_dc) -""" + mutable struct VMC tol::Float64 MaxIter::Int lr::Float64 lr_dc::Float64 + type::opt end -function InverseLR(ν, lr, lr_dc) - return lr / (1 + ν / lr_dc), ν+1 -end - -VMC(MaxIter::Int, lr::Float64; tol = 1.0e-3, lr_dc = 50.0) = VMC(tol, MaxIter, lr, lr_dc); +VMC(MaxIter::Int, lr::Float64, type; tol = 1.0e-3, lr_dc = 50.0) = VMC(tol, MaxIter, lr, lr_dc, type); -function gd_GradientByVMC(opt::VMC, - sam::MHSampler, - ham::SumH, - wf, ps, st, ν = 1, verbose = true, accMCMC = [10, [0.45, 0.55]]) - res = 1.0; - λ₀ = 0.; - α = opt.lr; - err_opt = zeros(opt.MaxIter) - x0, ~, acc = sampler_restart(sam, ps, st); - #x = reduce(vcat,reduce(vcat,x0)) - #display(histogram(x,xlim = (-40,40))) - verbose && @printf("Initialize MCMC: Δt = %.2f, accRate = %.4f \n", sam.Δt, acc) - acc_step = accMCMC[1]; +function gd_GradientByVMC(opt_vmc::VMC, sam::MHSampler, ham::SumH, + wf, ps, st, + ν = 1, verbose = true, accMCMC = [10, [0.45, 0.55]]) + + res, λ₀, α = 1.0, 0., opt_vmc.lr + err_opt = zeros(opt_vmc.MaxIter) + + x0, ~, acc = sampler_restart(sam, ps, st) + acc_step, acc_range = accMCMC acc_opt = zeros(acc_step) - acc_range = accMCMC[2]; + + verbose && @printf("Initialize MCMC: Δt = %.2f, accRate = %.4f \n", sam.Δt, acc) verbose && @printf(" k | 𝔼[E_L] | V[E_L] | res | LR |accRate| Δt \n") - for k = 1 : opt.MaxIter - sam.x0 = x0; - sam.Ψ = wf; + for k = 1 : opt_vmc.MaxIter + sam.x0 = x0 + + # adjust Δt acc_opt[mod(k,acc_step)+1] = acc - if mod(k,acc_step) == 0 - if mean(acc_opt) < acc_range[1] - sam.Δt = sam.Δt * exp(1/10 * (mean(acc_opt) - acc_range[1])/acc_range[1]) - elseif mean(acc_opt) > acc_range[2] - sam.Δt = sam.Δt * exp(1/10 * (mean(acc_opt) - acc_range[2])/acc_range[2]) - end - end - α, ν = InverseLR(ν, opt.lr, opt.lr_dc) - # optimize - λ₀, σ, E, x0, acc = Eloc_Exp_TV_clip(wf, ps, st, sam, ham) - g = grad(wf, x0, ps, st, E) - #if k % 10 == 0 - # x = reduce(vcat,reduce(vcat,x0)) - # display(histogram(x,xlim = (-40,40))) - #end - # Optimization - st_opt = Optimisers.setup(Optimisers.AdamW(α), ps) - st_opt, ps = Optimisers.update(st_opt, ps, g) + sam.Δt = acc_adjust(k, sam.Δt, acc_opt, acc_range, acc_step) + + # adjust learning rate + α, ν = InverseLR(ν, opt_vmc.lr, opt_vmc.lr_dc) - res = norm(destructure(g)[1]); - verbose && @printf(" %3.d | %.5f | %.5f | %.5f | %.5f | %.3f | %.3f \n", - k, λ₀, σ, res, α, acc, sam.Δt); + # optimization + ps, acc, λ₀, res, σ = Optimization(opt_vmc.type, wf, ps, st, sam::MHSampler, ham::SumH, α) + + # err + verbose && @printf(" %3.d | %.5f | %.5f | %.5f | %.5f | %.3f | %.3f \n", k, λ₀, σ, res, α, acc, sam.Δt) err_opt[k] = λ₀ - if res < opt.tol + + if res < opt_vmc.tol break; end end return wf, err_opt, ps end + diff --git a/src/vmc/vmc_utils.jl b/src/vmc/vmc_utils.jl new file mode 100644 index 0000000..df57e2c --- /dev/null +++ b/src/vmc/vmc_utils.jl @@ -0,0 +1,23 @@ +using LinearAlgebra + +function InverseLR(ν, lr, lr_dc) + return lr / (1 + ν / lr_dc), ν+1 +end + +function acc_adjust(k::Int, Δt::Number, acc_opt::AbstractVector, acc_range::AbstractVector, acc_step::Int) + if mod(k, acc_step) == 0 + if mean(acc_opt) < acc_range[1] + Δt = Δt * exp(1/10 * (mean(acc_opt) - acc_range[1])/acc_range[1]) + elseif mean(acc_opt) > acc_range[2] + Δt = Δt * exp(1/10 * (mean(acc_opt) - acc_range[2])/acc_range[2]) + end + end + return Δt +end + + + + + + + \ No newline at end of file diff --git a/test/Be.jl b/test/Be.jl index 0632291..efdaa98 100644 --- a/test/Be.jl +++ b/test/Be.jl @@ -14,9 +14,9 @@ using LinearAlgebra using BenchmarkTools using HyperDualNumbers: Hyper -n1 = Rnldegree = 6 +n1 = Rnldegree = 5 Ylmdegree = 2 -totdegree = 6 +totdegree = 5 Nel = 4 X = randn(SVector{3, Float64}, Nel) Σ = [↑,↑,↓,↓] @@ -27,23 +27,39 @@ nuclei = [ Nuc(zeros(SVector{3, Float64}), Nel * 1.0)] n2 = 1 Pn = Polynomials4ML.legendre_basis(n1+1) spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] -ζ = rand(length(spec)) +ζ = 10 * rand(length(spec)) Dn = SlaterBasis(ζ) bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) bYlm = RYlmBasis(Ylmdegree) # setup state -wf = BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +wf, spec, spec1p = BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +displayspec(spec, spec1p) + ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) +p, = destructure(ps) +length(p) K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) - Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) - Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) ham = SumH(K, Vext, Vee) -sam = MHSampler(wf, Nel, Δt = 0.4, burnin = 1000, nchains = 2000) -opt = VMC(500, 0.1, lr_dc = 10) -wf, err_opt, ps = gd_GradientByVMC(opt, sam, ham, wf, ps, st) -#Eref = -14.667 \ No newline at end of file +sam = MHSampler(wf, Nel, Δt = 0.5, burnin = 1000, nchains = 2000) + +opt_vmc = VMC(3000, 0.1, AdamW(), lr_dc = 100) +wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st) + +err = err_opt +per = 0.2 +err1 = zero(err) +for i = 1:length(err) + err1[i] = mean(err[Int(ceil(i-per * i)):i]) +end +err1 + +Eref = -14.667 + + +# using Plots +# plot(abs.(err1 .- Eref), w = 3, yscale=:log10) diff --git a/test/H2.jl b/test/H2.jl new file mode 100644 index 0000000..660ad56 --- /dev/null +++ b/test/H2.jl @@ -0,0 +1,64 @@ +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec +using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC +using ACEbase.Testing: print_tf, fdtest +using LuxCore +using Lux +using Zygote +using Optimisers # mainly for the destrcuture(ps) function +using Random +using Printf +using LinearAlgebra +using BenchmarkTools +using HyperDualNumbers: Hyper + +n1 = Rnldegree = 4 +Ylmdegree = 2 +totdegree = 4 +Nel = 2 +X = randn(SVector{3, Float64}, Nel) +Σ = [↑,↓] +nuclei = [ Nuc(SVector(0.0,0.0,-0.7), 1.0), Nuc(SVector(0.0,0.0,0.7), 1.0)] +## + +# Defining AtomicOrbitalsBasis +n2 = 1 +Pn = Polynomials4ML.legendre_basis(n1+1) +spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] +ζ = 10 * rand(length(spec)) +Dn = SlaterBasis(ζ) +bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) +bYlm = RYlmBasis(Ylmdegree) + +# setup state +wf, spec, spec1p = BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +displayspec(spec, spec1p) + +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) +p, = destructure(ps) +length(p) + +K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) +Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) +Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) + +ham = SumH(K, Vext, Vee) +sam = MHSampler(wf, Nel, Δt = 0.5, burnin = 1000, nchains = 2000) + +opt_vmc = VMC(3000, 0.1, AdamW(), lr_dc = 100) +wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st) + +err = err_opt +per = 0.2 +err1 = zero(err) +for i = 1:length(err) + err1[i] = mean(err[Int(ceil(i-per * i)):i]) +end +err1 + +Eref = -1.8887557142857143 + +# using Plots +plot(abs.(err1 .- Eref), w = 3, yscale=:log10) \ No newline at end of file diff --git a/test/He.jl b/test/He.jl index af22d5b..d57e448 100644 --- a/test/He.jl +++ b/test/He.jl @@ -2,7 +2,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec -using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC +using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC, AdamW, SR using ACEbase.Testing: print_tf, fdtest using LuxCore using Lux @@ -14,9 +14,9 @@ using LinearAlgebra using BenchmarkTools using HyperDualNumbers: Hyper -n1 = Rnldegree = 6 +n1 = Rnldegree = 4 Ylmdegree = 2 -totdegree = 6 +totdegree = 4 Nel = 2 X = randn(SVector{3, Float64}, Nel) Σ = [↑,↓] @@ -27,7 +27,7 @@ nuclei = [ Nuc(zeros(SVector{3, Float64}), 2.0)] n2 = 1 Pn = Polynomials4ML.legendre_basis(n1+1) spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] -ζ = rand(length(spec)) +ζ = 10 * rand(length(spec)) Dn = SlaterBasis(ζ) bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) bYlm = RYlmBasis(Ylmdegree) @@ -37,18 +37,29 @@ wf, spec, spec1p = BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; displayspec(spec, spec1p) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) - p, = destructure(ps) length(p) K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) - Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) - Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) ham = SumH(K, Vext, Vee) -sam = MHSampler(wf, Nel, Δt = 2.0, burnin = 1000, nchains = 2000) -opt = VMC(500, 0.1, lr_dc = 50) -wf, err_opt, ps = gd_GradientByVMC(opt, sam, ham, wf, ps, st) -#Eref = -2.9037 +sam = MHSampler(wf, Nel, Δt = 0.5, burnin = 1000, nchains = 2000) + +opt_vmc = VMC(3000, 0.1, SR(), lr_dc = 100) +wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st) + + + +err = err_opt +per = 0.2 +err1 = zero(err) +for i = 1:length(err) + err1[i] = mean(err[Int(ceil(i-per * i)):i]) +end + +Eref = -2.9037 + +# using Plots +# plot(abs.(err1 .- Eref), w = 3, yscale=:log10) \ No newline at end of file From 4d2b0da0f7ac36281f66aff3875297505570b958 Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Mon, 10 Jul 2023 14:55:06 +0800 Subject: [PATCH 65/79] minor fix --- src/vmc/vmc.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/vmc/vmc.jl b/src/vmc/vmc.jl index 0c42b2f..de5483a 100644 --- a/src/vmc/vmc.jl +++ b/src/vmc/vmc.jl @@ -37,7 +37,7 @@ function gd_GradientByVMC(opt_vmc::VMC, sam::MHSampler, ham::SumH, α, ν = InverseLR(ν, opt_vmc.lr, opt_vmc.lr_dc) # optimization - ps, acc, λ₀, res, σ = Optimization(opt_vmc.type, wf, ps, st, sam::MHSampler, ham::SumH, α) + ps, acc, λ₀, res, σ = Optimization(opt_vmc.type, wf, ps, st, sam, ham, α) # err verbose && @printf(" %3.d | %.5f | %.5f | %.5f | %.5f | %.3f | %.3f \n", k, λ₀, σ, res, α, acc, sam.Δt) From bbfd1f737938d5d745f26a0f503dbb982486133d Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Fri, 28 Jul 2023 15:06:28 +0800 Subject: [PATCH 66/79] minor change of sr --- Project.toml | 1 + src/vmc/optimisers/sr.jl | 21 +++++++++++++-------- test/He.jl | 2 +- test/test_bflow_lux1.jl | 2 +- 4 files changed, 16 insertions(+), 10 deletions(-) diff --git a/Project.toml b/Project.toml index 4f5e5db..d218521 100644 --- a/Project.toml +++ b/Project.toml @@ -23,6 +23,7 @@ Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/src/vmc/optimisers/sr.jl b/src/vmc/optimisers/sr.jl index 1c24a90..149501a 100644 --- a/src/vmc/optimisers/sr.jl +++ b/src/vmc/optimisers/sr.jl @@ -1,5 +1,6 @@ using Optimisers - +using Statistics +# stochastic reconfiguration mutable struct SR <: opt ϵ1::Number @@ -15,16 +16,20 @@ function Optimization(type::SR, wf, ps, st, sam::MHSampler, ham::SumH, α) ϵ1 = type.ϵ1 ϵ2 = type.ϵ2 + # E: N × 1 λ₀, σ, E, x0, acc = Eloc_Exp_TV_clip(wf, ps, st, sam, ham) dy_ps = grad_params.(Ref(wf), x0, Ref(ps), Ref(st)) - dy = _destructure.(dy_ps) - O = 1/2 * dy - Ō = mean(O) - ΔO = [O[i] - Ō for i = 1:length(O)] + # O: M × N + O = 1/2 * reshape(_destructure(dy_ps), (length(_destructure(ps)),sam.nchains)) + # ⟨O⟩: M x 1 + Ō = mean(O,dims =2) + ΔO = O .- Ō - g0 = 2.0 * sum(ΔO .*E)/ length(x0) + # g0: M x 1 + g0 = 2.0 * ΔO * E/sam.nchains - S = sum([ΔO[i] * ΔO[i]' for i = 1:length(O)])/length(x0) + # S: M x M + S = cov(O',dims=1) S = S + ϵ1 * Diagonal(diag(S)) + ϵ2 * Diagonal(diag(one(S))) g = S \ g0 @@ -36,7 +41,7 @@ function Optimization(type::SR, wf, ps, st, sam::MHSampler, ham::SumH, α) return ps, acc, λ₀, res, σ end - + diff --git a/test/He.jl b/test/He.jl index d57e448..860ac98 100644 --- a/test/He.jl +++ b/test/He.jl @@ -47,7 +47,7 @@ Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 ham = SumH(K, Vext, Vee) sam = MHSampler(wf, Nel, Δt = 0.5, burnin = 1000, nchains = 2000) -opt_vmc = VMC(3000, 0.1, SR(), lr_dc = 100) +opt_vmc = VMC(3000, 0.0015, SR(1e-4, 0.015), lr_dc = 50) wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st) diff --git a/test/test_bflow_lux1.jl b/test/test_bflow_lux1.jl index 0ff566b..420b6cf 100644 --- a/test/test_bflow_lux1.jl +++ b/test/test_bflow_lux1.jl @@ -52,7 +52,7 @@ bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) bYlm = RYlmBasis(Ylmdegree) # setup state -BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) ## From d09e01716f59e9afdc513aa45241b481221c2f97 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 8 Aug 2023 02:18:08 -0700 Subject: [PATCH 67/79] remove ACEcore dependency, add compat bounds of P4ML and Obejctpools --- Project.toml | 5 +- src/bflow.jl | 47 +++++---- src/bflow3d.jl | 26 ++--- test/test_bflow.jl | 15 ++- test/test_bflow_lux.jl | 43 ++++---- test/test_bflow_lux1.jl | 229 ---------------------------------------- 6 files changed, 71 insertions(+), 294 deletions(-) delete mode 100644 test/test_bflow_lux1.jl diff --git a/Project.toml b/Project.toml index d218521..7d183ab 100644 --- a/Project.toml +++ b/Project.toml @@ -5,7 +5,6 @@ version = "0.0.1" [deps] ACEbase = "14bae519-eb20-449c-a949-9c58ed33163e" -ACEcore = "44c1e890-45d1-48ea-94d6-c2ea5b573f71" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" @@ -30,11 +29,11 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] ACEbase = "0.4.2" -ACEcore = "0.0.5" BenchmarkTools = "1" ForwardDiff = "0.10" JSON = "0.21" -ObjectPools = "0.2.1" +ObjectPools = "0.3.1" +Polynomials4ML = "0.2.2" julia = "1" [extras] diff --git a/src/bflow.jl b/src/bflow.jl index bfea545..cc240ec 100644 --- a/src/bflow.jl +++ b/src/bflow.jl @@ -1,19 +1,20 @@ -using ACEcore, Polynomials4ML -using Polynomials4ML: OrthPolyBasis1D3T -using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! -using ACEcore.Utils: gensparse +using Polynomials4ML +using Polynomials4ML: OrthPolyBasis1D3T, PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! +using Polynomials4ML.Utils: gensparse using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr +using ObjectPools: unwrap + import ForwardDiff mutable struct BFwf1{T, TT, TPOLY, TE} trans::TT polys::TPOLY pooling::PooledSparseProduct{2} - corr::SparseSymmProdDAG{T} + corr::SparseSymmProdDAG W::Matrix{T} envelope::TE - spec::Vector{Vector{Int64}} # corr.spec TODO: this needs to be remove + spec::AbstractArray # ---------------- Temporaries P::Matrix{T} ∂P::Matrix{T} @@ -58,8 +59,10 @@ function BFwf1(Nel::Integer, polys; totdeg = length(polys), # further restrict spec = [t for t in spec if sd_admissible([spec1p[t[j]] for j = 1:length(t)])] - corr1 = SparseSymmProd(spec; T = Float64) - corr = corr1.dag + corr1 = SparseSymmProdDAG(spec) + corr = corr1 + + spec = Tuple.(spec) # initial guess for weights Q, _ = qr(randn(T, length(corr), Nel)) @@ -128,7 +131,7 @@ function assemble_A(wf::BFwf1, X::AbstractVector, Σ, Pnn=nothing) for i = 1:nX onehot!(Si, i, Σ) - ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) + Polynomials4ML.evaluate!(Ai, wf.pooling, (unwrap(P), Si)) A[i, :] .= Ai end return A @@ -137,7 +140,7 @@ end function evaluate(wf::BFwf1, X::AbstractVector, Σ, Pnn=nothing) nX = length(X) A = assemble_A(wf, X, Σ) - AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) + AA = Polynomials4ML.evaluate(wf.corr, A) # nX x length(wf.corr) # the only basis to be purified are those with same spin # scan through all corr basis, if they comes from same spin, remove self interation by using basis @@ -148,7 +151,7 @@ function evaluate(wf::BFwf1, X::AbstractVector, Σ, Pnn=nothing) # === # Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) # nX x nX + mul!(Φ, unwrap(AA), wf.W) # nX x nX Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin release!(AA) @@ -161,9 +164,9 @@ function gradp_evaluate(wf::BFwf1, X::AbstractVector, Σ) nX = length(X) A = assemble_A(wf, X, Σ) - AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) + AA = Polynomials4ML.evaluate(wf.corr, A) # nX x length(wf.corr) Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) + mul!(Φ, unwrap(AA), wf.W) Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin @@ -175,7 +178,7 @@ function gradp_evaluate(wf::BFwf1, X::AbstractVector, Σ) # ∂Wij = ∑_ab ∂Φab * ∂_Wij( ∑_k AA_ak W_kb ) # = ∑_ab ∂Φab * ∑_k δ_ik δ_bj AA_ak # = ∑_a ∂Φaj AA_ai = ∂Φaj' * AA_ai - ∇p = transpose(parent(AA)) * ∂Φ + ∇p = transpose(unwrap(AA)) * ∂Φ release!(AA) ∇p = ∇p * 2 @@ -230,16 +233,16 @@ function gradient(wf::BFwf1, X, Σ) for i = 1:nX onehot!(Si, i, Σ) - ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) + Polynomials4ML.evaluate!(Ai, wf.pooling, (unwrap(P), Si)) A[i, :] .= Ai end # n-correlations - AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) + AA = Polynomials4ML.evaluate(wf.corr, A) # nX x length(wf.corr) # generalized orbitals Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) + mul!(Φ, unwrap(AA), wf.W) # the resulting matrix should contains two block each comes from each spin Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] @@ -260,7 +263,7 @@ function gradient(wf::BFwf1, X, Σ) # ∂A = ∂ψ/∂A = ∂ψ/∂AA * ∂AA/∂A -> use custom pullback ∂A = wf.∂A # zeros(size(A)) - ACEcore.pullback_arg!(∂A, ∂AA, wf.corr, parent(AA)) + Polynomials4ML.pullback_arg!(∂A, ∂AA, wf.corr, unwrap(AA)) release!(AA) # ∂P = ∂ψ/∂P = ∂ψ/∂A * ∂A/∂P -> use custom pullback @@ -274,7 +277,7 @@ function gradient(wf::BFwf1, X, Σ) onehot!(Si_, i, Σ) # note this line ADDS the pullback into ∂P, not overwrite the content!! ∂Ai = @view ∂A[i, :] - ACEcore._pullback_evalpool!((∂P, ∂Si), ∂Ai, wf.pooling, (P, Si_)) + Polynomials4ML._pullback_evaluate!((∂P, ∂Si), ∂Ai, wf.pooling, (P, Si_)) end # ∂X = ∂ψ/∂X = ∂ψ/∂P * ∂P/∂X @@ -343,7 +346,7 @@ function _assemble_A_∇A_ΔA(wf, X, Σ) @inbounds for i = 1:nX # loop over orbital bases (which i becomes ∅) fill!(Si_, 0) onehot!(Si_, i, Σ) - ACEcore.evalpool!(Ai, wf.pooling, (P, Si_)) + Polynomials4ML.evaluate!(Ai, wf.pooling, (P, Si_)) @. A[i, :] .= Ai for (iA, (k, σ)) in enumerate(spec_A) for a = 1:nX @@ -397,7 +400,7 @@ function _laplacian_inner(AA, ∇AA, ΔAA, wf, Σ) # the wf, and the first layer of derivatives Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) + mul!(Φ, unwrap(AA), wf.W) Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin Φ⁻ᵀ = transpose(pinv(Φ)) @@ -439,7 +442,7 @@ function gradp_laplacian(wf::BFwf1, X, Σ) # the wf, and the first layer of derivatives Φ = wf.Φ - mul!(Φ, parent(AA), wf.W) + mul!(Φ, unwrap(AA), wf.W) Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin Φ⁻¹ = pinv(Φ) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 7eb25b4..db0ea04 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -1,8 +1,8 @@ -using ACEcore, Polynomials4ML, Random +using Polynomials4ML, Random using Polynomials4ML: OrthPolyBasis1D3T -using ACEcore: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! -using ACEcore.Utils: gensparse +using Polynomials4ML: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! +using Polynomials4ML.Utils: gensparse using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det import ForwardDiff using ACEpsi.AtomicOrbitals: make_nlms_spec @@ -31,7 +31,7 @@ struct DenseLayer <: AbstractExplicitLayer end function (l::DenseLayer)(x::AbstractMatrix, ps, st) - return parent(x) * ps.W, st + return unwrap(x) * ps.W, st end # Jerry: Maybe we should use Glorot Uniform if we have no idea about what we should use? @@ -180,7 +180,7 @@ end # for i = 1:nX # onehot!(Si, i, Σ) -# ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) +# ACEcore.evalpool!(Ai, wf.pooling, (unwrap(P), Si)) # A[i, :] .= Ai # end # return A @@ -295,7 +295,7 @@ end # # === # # Φ = wf.Φ -# mul!(Φ, parent(AA), wf.W) # nX x nX +# mul!(Φ, unwrap(AA), wf.W) # nX x nX # Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin # release!(AA) @@ -310,7 +310,7 @@ end # A = assemble_A(wf, X, Σ) # AA = ACEcore.evaluate(wf.corr, A) # nX x length(wf.corr) # Φ = wf.Φ -# mul!(Φ, parent(AA), wf.W) +# mul!(Φ, unwrap(AA), wf.W) # Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin @@ -322,7 +322,7 @@ end # # ∂Wij = ∑_ab ∂Φab * ∂_Wij( ∑_k AA_ak W_kb ) # # = ∑_ab ∂Φab * ∑_k δ_ik δ_bj AA_ak # # = ∑_a ∂Φaj AA_ai = ∂Φaj' * AA_ai -# ∇p = transpose(parent(AA)) * ∂Φ +# ∇p = transpose(unwrap(AA)) * ∂Φ # release!(AA) # ∇p = ∇p * 2 @@ -377,7 +377,7 @@ end # for i = 1:nX # onehot!(Si, i, Σ) -# ACEcore.evalpool!(Ai, wf.pooling, (parent(P), Si)) +# ACEcore.evalpool!(Ai, wf.pooling, (unwrap(P), Si)) # A[i, :] .= Ai # end @@ -386,7 +386,7 @@ end # # generalized orbitals # Φ = wf.Φ -# mul!(Φ, parent(AA), wf.W) +# mul!(Φ, unwrap(AA), wf.W) # # the resulting matrix should contains two block each comes from each spin # Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] @@ -407,7 +407,7 @@ end # # ∂A = ∂ψ/∂A = ∂ψ/∂AA * ∂AA/∂A -> use custom pullback # ∂A = wf.∂A # zeros(size(A)) -# ACEcore.pullback_arg!(∂A, ∂AA, wf.corr, parent(AA)) +# ACEcore.pullback_arg!(∂A, ∂AA, wf.corr, unwrap(AA)) # release!(AA) # # ∂P = ∂ψ/∂P = ∂ψ/∂A * ∂A/∂P -> use custom pullback @@ -544,7 +544,7 @@ end # # the wf, and the first layer of derivatives # Φ = wf.Φ -# mul!(Φ, parent(AA), wf.W) +# mul!(Φ, unwrap(AA), wf.W) # Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin # Φ⁻ᵀ = transpose(pinv(Φ)) @@ -586,7 +586,7 @@ end # # the wf, and the first layer of derivatives # Φ = wf.Φ -# mul!(Φ, parent(AA), wf.W) +# mul!(Φ, unwrap(AA), wf.W) # Φ = Φ .* [Σ[i] == Σ[j] for j = 1:nX, i = 1:nX] # the resulting matrix should contains two block each comes from each spin # Φ⁻¹ = pinv(Φ) diff --git a/test/test_bflow.jl b/test/test_bflow.jl index f4d0e39..7ab5c43 100644 --- a/test/test_bflow.jl +++ b/test/test_bflow.jl @@ -1,10 +1,9 @@ -using Polynomials4ML, ACEcore, ACEbase, Printf, ACEpsi +using Polynomials4ML, ACEbase, Printf, ACEpsi using ACEpsi: BFwf1, gradient, laplacian using LinearAlgebra -#using Random -#Random.seed!(123) -## + + function lap_test(f, Δf, X) F = f(X) ΔF = Δf(X) @@ -63,7 +62,7 @@ function grad_test(f, df, X) end "This function should be removed later to test in a nicer way..." -function fdtest(F, Σ, dF, x::AbstractVector; h0 = 1.0, verbose=true) +function _fdtest(F, Σ, dF, x::AbstractVector; h0 = 1.0, verbose=true) errors = Float64[] E = F(x, Σ) dE = dF @@ -99,7 +98,7 @@ function fdtest(F, Σ, dF, x::AbstractVector; h0 = 1.0, verbose=true) const ↑, ↓, ∅ = '↑','↓','∅' Nel = 5 polys = legendre_basis(8) -wf = BFwf1(Nel, polys; ν=3) +wf = BFwf1(Nel, polys; ν=3, trans = x -> 2/pi * atan(x)) X = 2 * rand(Nel) .- 1 Σ = rand([↑, ↓], Nel) @@ -111,10 +110,10 @@ g = gradient(wf, X, Σ) using LinearAlgebra using Printf -#using ACEbase.Testing: fdtest +#using ACEbase.Testing: _fdtest @info("Fd test of gradient w.r.t. X") -fdtest(wf, Σ, g, X) +_fdtest(wf, Σ, g, X) # ## diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 49e56f0..2d4cd57 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -12,7 +12,6 @@ using Random using Printf using LinearAlgebra using BenchmarkTools -# using Cthulhu using HyperDualNumbers: Hyper @@ -28,10 +27,10 @@ function grad_test2(f, df, X::AbstractVector) end end -Rnldegree = 4 -Ylmdegree = 4 -totdegree = 8 -Nel = 5 +Rnldegree = n1 = 3 +Ylmdegree = 3 +totdegree = 4 +Nel = 3 X = randn(SVector{3, Float64}, Nel) Σ = rand(spins(), Nel) nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] @@ -44,11 +43,16 @@ hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec ## # Defining AtomicOrbitalsBasis -bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +n2 = 2 +Pn = Polynomials4ML.legendre_basis(n1+1) +spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] +ζ = rand(length(spec)) +Dn = GaussianBasis(ζ) +bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) bYlm = RYlmBasis(Ylmdegree) # setup state -BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) ## @@ -116,16 +120,17 @@ grad_test2(Fp, dFp, W0) ## @info("Test consistency when input isa HyperDualNumbers") -hp = Zygote.gradient(p -> BFwf_chain(hX, p, st)[1], ps)[1] -hp, = destructure(hp) -P = similar(p) -for i = 1:length(P) - P[i] = hp[i].value -end +#hp = Zygote.gradient(p -> BFwf_chain(hX, p, st)[1], ps)[1] -print_tf(@test P ≈ p) +#hp, = destructure(hp) +#P = similar(p) +#for i = 1:length(P) +# P[i] = hp[i].value +#end -println() +#print_tf(@test P ≈ p) + +#println() ## @@ -218,7 +223,7 @@ function ∇ΔF(x, ps) return p end -W0, re = destructure(ps) -Fp = w -> ΔF(X, re(w)) -dFp = w -> ∇ΔF(X, re(w)) -fdtest(Fp, dFp, W0) \ No newline at end of file +#W0, re = destructure(ps) +#Fp = w -> ΔF(X, re(w)) +#dFp = w -> ∇ΔF(X, re(w)) +#fdtest(Fp, dFp, W0) \ No newline at end of file diff --git a/test/test_bflow_lux1.jl b/test/test_bflow_lux1.jl deleted file mode 100644 index 420b6cf..0000000 --- a/test/test_bflow_lux1.jl +++ /dev/null @@ -1,229 +0,0 @@ -using ACEpsi, Polynomials4ML, StaticArrays, Test -using Polynomials4ML: natural_indices, degree, SparseProduct -using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate -using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow -using ACEpsi.vmc: gradient, laplacian, grad_params -using ACEbase.Testing: print_tf, fdtest -using LuxCore -using Lux -using Zygote -using Optimisers # mainly for the destrcuture(ps) function -using Random -using Printf -using LinearAlgebra -using BenchmarkTools - -using HyperDualNumbers: Hyper - -function grad_test2(f, df, X::AbstractVector) - F = f(X) - ∇F = df(X) - nX = length(X) - EE = Matrix(I, (nX, nX)) - - for h in 0.1.^(3:12) - gh = [ (f(X + h * EE[:, i]) - F) / h for i = 1:nX ] - @printf(" %.1e | %.2e \n", h, norm(gh - ∇F, Inf)) - end -end - -Rnldegree = n1 = 4 -Ylmdegree = 4 -totdegree = 8 -Nel = 5 -X = randn(SVector{3, Float64}, Nel) -Σ = rand(spins(), Nel) -nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] - -# wrap it as HyperDualNumbers -x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) -hX = [x2dualwrtj(x, 0) for x in X] -hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec - -## - -# Defining AtomicOrbitalsBasis -n2 = 3 -Pn = Polynomials4ML.legendre_basis(n1+1) -spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] -ζ = rand(length(spec)) -Dn = GaussianBasis(ζ) -bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) -bYlm = RYlmBasis(Ylmdegree) - -# setup state -BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) -ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) - -## - -@info("Test evaluate") -A1 = BFwf_chain(X, ps, st) -hA1 = BFwf_chain(hX, ps, st) - -print_tf(@test hA1[1].value ≈ A1[1]) - -println() - -## -F(X) = BFwf_chain(X, ps, st)[1] - -# @profview let F = F, X = X -# for i = 1:10_000 -# F(X) -# end -# end - -# @btime F(X) - - -## - -@info("Test ∇ψ w.r.t. X") -ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) -y, st = Lux.apply(BFwf_chain, X, ps, st) - -F(X) = BFwf_chain(X, ps, st)[1] -dF(X) = Zygote.gradient(x -> BFwf_chain(x, ps, st)[1], X)[1] -fdtest(F, dF, X, verbose = true) - -## - -@info("Test consistency with HyperDualNumbers") -for _ = 1:30 - local X = randn(SVector{3, Float64}, Nel) - local Σ = rand(spins(), Nel) - local hdF = [zeros(3) for _ = 1:Nel] - local hX = [x2dualwrtj(x, 0) for x in X] - for i = 1:3 - for j = 1:Nel - hX[j] = x2dualwrtj(X[j], i) # ∂Ψ/∂xj_{i} - hdF[j][i] = BFwf_chain(hX, ps, st)[1].epsilon1 - hX[j] = x2dualwrtj(X[j], 0) - end - end - print_tf(@test dF(X) ≈ hdF) -end -println() - -## - -@info("Test ∇ψ w.r.t. parameters") -p = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1] -p, = destructure(p) - -W0, re = destructure(ps) -Fp = w -> BFwf_chain(X, re(w), st)[1] -dFp = w -> ( gl = Zygote.gradient(p -> BFwf_chain(X, p, st)[1], ps)[1]; destructure(gl)[1]) -grad_test2(Fp, dFp, W0) - -## - -@info("Test consistency when input isa HyperDualNumbers") -#hp = Zygote.gradient(p -> BFwf_chain(hX, p, st)[1], ps)[1] - -#hp, = destructure(hp) -#P = similar(p) -#for i = 1:length(P) -# P[i] = hp[i].value -#end - -#print_tf(@test P ≈ p) - -#println() - -## - -@info("Test Δψ w.r.t. X using HyperDualNumbers") - -X = [randn(3) for _ = 1:Nel] -hX = [x2dualwrtj(x, 0) for x in X] -Σ = rand(spins(), Nel) -F(x) = BFwf_chain(x, ps, st)[1] - - -function ΔF(x) - ΔΨ = 0.0 - hX = [x2dualwrtj(xx, 0) for xx in x] - for i = 1:3 - for j = 1:Nel - hX[j] = x2dualwrtj(x[j], i) # ∂Φ/∂xj_{i} - ΔΨ += BFwf_chain(hX, ps, st)[1].epsilon12 - hX[j] = x2dualwrtj(x[j], 0) - end - end - return ΔΨ -end - -Δ1 = ΔF(X) -f0 = F(X) - -for h in 0.1.^(1:8) - Δfh = 0.0 - for i = 1:Nel - for j = 1:3 - XΔX_add, XΔX_sub = deepcopy(X), deepcopy(X) - XΔX_add[i][j] += h - XΔX_sub[i][j] -= h - Δfh += (F(XΔX_add) - f0) / h^2 - Δfh += (F(XΔX_sub) - f0) / h^2 - end - end - @printf(" %.1e | %.2e \n", h, abs(Δfh - Δ1)) -end - -## - -@info("Test gradp Δψ using HyperDualNumbers") -g_bchain = xx -> Zygote.gradient(p -> BFwf_chain(xx, p, st)[1], ps)[1] -g_bchain(hX) - -using ACEpsi: zero! -using HyperDualNumbers - -function grad_lap(g_bchain, x) - function _mapadd!(f, dest::NamedTuple, src::NamedTuple) - for k in keys(dest) - _mapadd!(f, dest[k], src[k]) - end - return nothing - end - _mapadd!(f, dest::Nothing, src) = nothing - _mapadd!(f, dest::AbstractArray, src::AbstractArray) = - map!((s, d) -> d + f(s), dest, src, dest) - - Δ = zero!(g_bchain(x)) - hX = [x2dualwrtj(xx, 0) for xx in x] - for i = 1:3 - for j = 1:length(x) - hX[j] = x2dualwrtj(x[j], i) - _mapadd!(ε₁ε₂part, Δ, g_bchain(hX)) - hX[j] = x2dualwrtj(x[j], 0) - end - end - return Δ -end - -function ΔF(x, ps) - ΔΨ = 0.0 - hX = [x2dualwrtj(xx, 0) for xx in x] - for i = 1:3 - for j = 1:Nel - hX[j] = x2dualwrtj(x[j], i) # ∂Φ/∂xj_{i} - ΔΨ += BFwf_chain(hX, ps, st)[1].epsilon12 - hX[j] = x2dualwrtj(x[j], 0) - end - end - return ΔΨ -end - -function ∇ΔF(x, ps) - g_bchain = xx -> Zygote.gradient(p -> BFwf_chain(xx, p, st)[1], ps)[1] - p, = destructure(grad_lap(g_bchain, x)) - return p -end - -#W0, re = destructure(ps) -#Fp = w -> ΔF(X, re(w)) -#dFp = w -> ∇ΔF(X, re(w)) -#fdtest(Fp, dFp, W0) \ No newline at end of file From 90c070a61ce95cab27031cefcf23097d498eaa7e Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 8 Aug 2023 02:19:13 -0700 Subject: [PATCH 68/79] add CI for julia 1.9 --- .github/workflows/CI.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 076759e..18efdb8 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -20,6 +20,7 @@ jobs: version: - '1.0' - '1.8' + - '1.9' - 'nightly' os: - ubuntu-latest From 3a8186bf74d359e882cfaa470c14e7102fb13617 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 8 Aug 2023 02:30:37 -0700 Subject: [PATCH 69/79] clean up --- {test => examples/3D}/Be.jl | 0 {test => examples/3D}/H2.jl | 0 {test => examples/3D}/He.jl | 0 {test => examples/3D}/harmonic_oscillator.jl | 0 profile/profile_bflow1.jl | 2 +- test/compare_bflow.jl | 34 -------------------- test/test_admissiible.jl | 20 ------------ test/test_bflow.jl | 4 --- test/test_bflow_lux.jl | 8 ++--- {test => tmp}/test_matmul.jl | 0 {test => tmp}/try.jl | 0 11 files changed, 5 insertions(+), 63 deletions(-) rename {test => examples/3D}/Be.jl (100%) rename {test => examples/3D}/H2.jl (100%) rename {test => examples/3D}/He.jl (100%) rename {test => examples/3D}/harmonic_oscillator.jl (100%) delete mode 100644 test/compare_bflow.jl delete mode 100644 test/test_admissiible.jl rename {test => tmp}/test_matmul.jl (100%) rename {test => tmp}/try.jl (100%) diff --git a/test/Be.jl b/examples/3D/Be.jl similarity index 100% rename from test/Be.jl rename to examples/3D/Be.jl diff --git a/test/H2.jl b/examples/3D/H2.jl similarity index 100% rename from test/H2.jl rename to examples/3D/H2.jl diff --git a/test/He.jl b/examples/3D/He.jl similarity index 100% rename from test/He.jl rename to examples/3D/He.jl diff --git a/test/harmonic_oscillator.jl b/examples/3D/harmonic_oscillator.jl similarity index 100% rename from test/harmonic_oscillator.jl rename to examples/3D/harmonic_oscillator.jl diff --git a/profile/profile_bflow1.jl b/profile/profile_bflow1.jl index f3c3f42..2673150 100644 --- a/profile/profile_bflow1.jl +++ b/profile/profile_bflow1.jl @@ -1,6 +1,6 @@ -using Polynomials4ML, ACEcore, ACEpsi, ACEbase, Printf +using Polynomials4ML, ACEpsi, ACEbase, Printf using ACEpsi: BFwf, gradient, evaluate, laplacian using LinearAlgebra using BenchmarkTools diff --git a/test/compare_bflow.jl b/test/compare_bflow.jl deleted file mode 100644 index 6fd5175..0000000 --- a/test/compare_bflow.jl +++ /dev/null @@ -1,34 +0,0 @@ -using Polynomials4ML, ACEcore, ACEbase -using ACEpsi: BFwf, gradient, evaluate, envelopefcn -using JSON -using Printf -using LinearAlgebra - -const ↑, ↓, ∅ = '↑','↓','∅' -using JSON - -# == test configs == -data_ortho = JSON.parse(open("test/ACESchrodingerRef/orthopolyweights.json")) # import weights from ACESchrodinger.jl -data = JSON.parse(open("test/ACESchrodingerRef/bftest.json")) # import input electron position and parameter of model -ww = data_ortho[1]["ww"] -xx = data_ortho[1]["tdf"] -X = data[1]["X"] -PP = data[1]["P"] -refval = data[1]["refval"] -ww = Float64.(ww) -xx = Float64.(xx) -Σ = [↑, ↑, ↓, ↓, ↓]; -Nel = 5 -WW = DiscreteWeights(xx, ww) -polys = orthpolybasis(10, WW) - -wf = BFwf(Nel, polys; ν=2, totdeg = 10, trans = atan, envelope = envelopefcn(x -> sqrt(x^2 + 1), 0.5)) -# == - -for i = 1:5 - wf.W[:, i] = PP[i][2:end] # the first entry of PP[i] is the extra constant -end -println("ACESchrodinger.BFwf - ACEpsi.BFwf: ", wf(X, Σ) - refval) -spec1p = [ (k, σ) for σ in [1, 2, 3] for k in 1:length(polys) ] # (1, 2, 3) = (∅, ↑, ↓); -spec1p = sort(spec1p, by = b->b[1]) -@show displayspec(wf.spec, spec1p) \ No newline at end of file diff --git a/test/test_admissiible.jl b/test/test_admissiible.jl deleted file mode 100644 index 104865d..0000000 --- a/test/test_admissiible.jl +++ /dev/null @@ -1,20 +0,0 @@ -using Polynomials4ML, ACEcore, ACEpsi, ACEbase, Printf -using ACEpsi: BFwf, gradient, evaluate, laplacian, envelopefcn, displayspec -using LinearAlgebra -using BenchmarkTools -using JSON -const ↑, ↓, ∅ = '↑','↓','∅' -Σ = [↑, ↑, ↑, ↓, ↓]; -Nel = 5 -polys = legendre_basis(16) -MaxDeg = [16, 5] - - -test_ad = bb -> (@show bb; @show all([bb[i][1] < MaxDeg[i] for i = 1:length(bb)]); (length(bb) == 0 || all([bb[i][1] <= MaxDeg[length(bb)] for i = 1:length(bb)]))) -wf = BFwf(Nel, polys; ν=2, sd_admissible = test_ad) -@show length(wf.polys) -LL = displayspec(wf) -@show LL -for i = 1:length(LL) - @show LL[i] -end diff --git a/test/test_bflow.jl b/test/test_bflow.jl index 7ab5c43..fc8fdc8 100644 --- a/test/test_bflow.jl +++ b/test/test_bflow.jl @@ -110,7 +110,6 @@ g = gradient(wf, X, Σ) using LinearAlgebra using Printf -#using ACEbase.Testing: _fdtest @info("Fd test of gradient w.r.t. X") _fdtest(wf, Σ, g, X) @@ -219,6 +218,3 @@ wf2 = ACEpsi.set_params!(wf2, param1) ## -@warn("removed compat test since json file is missing") -# @info("Test compatibility with ACESchrodinger") # Jerry: Not sure if this should be kept in the same file -# include("compare_bflow.jl") diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index 2d4cd57..a96b8c3 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -27,10 +27,10 @@ function grad_test2(f, df, X::AbstractVector) end end -Rnldegree = n1 = 3 +Rnldegree = n1 = 2 Ylmdegree = 3 -totdegree = 4 -Nel = 3 +totdegree = 3 +Nel = 2 X = randn(SVector{3, Float64}, Nel) Σ = rand(spins(), Nel) nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] @@ -176,7 +176,7 @@ end @info("Test gradp Δψ using HyperDualNumbers") g_bchain = xx -> Zygote.gradient(p -> BFwf_chain(xx, p, st)[1], ps)[1] -g_bchain(hX) +# g_bchain(hX) using ACEpsi: zero! using HyperDualNumbers diff --git a/test/test_matmul.jl b/tmp/test_matmul.jl similarity index 100% rename from test/test_matmul.jl rename to tmp/test_matmul.jl diff --git a/test/try.jl b/tmp/try.jl similarity index 100% rename from test/try.jl rename to tmp/try.jl From 8428569436e25ed7054d9b6c8adc2089f435f4ea Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 8 Aug 2023 03:27:22 -0700 Subject: [PATCH 70/79] remove CI of J1.0, added and fixed all tests --- .github/workflows/CI.yml | 1 - src/bflow.jl | 15 ++++++++++++++- test/runtests.jl | 3 ++- test/test_atorbbasis.jl | 9 +++++++-- test/test_bflow.jl | 7 +++---- 5 files changed, 26 insertions(+), 9 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 18efdb8..0cc9b80 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -18,7 +18,6 @@ jobs: fail-fast: false matrix: version: - - '1.0' - '1.8' - '1.9' - 'nightly' diff --git a/src/bflow.jl b/src/bflow.jl index cc240ec..07b5971 100644 --- a/src/bflow.jl +++ b/src/bflow.jl @@ -91,7 +91,7 @@ This function return correct Si for pooling operation. function onehot!(Si, i, Σ) Si .= 0 for k = 1:length(Σ) - Si[k, spin2num(Σ[k])] = 1 + Si[k, spin2num1d(Σ[k])] = 1 end # set current electron to ϕ, also remove their contribution in the sum of ↑ or ↓ basis Si[i, 1] = 1 @@ -99,6 +99,19 @@ function onehot!(Si, i, Σ) Si[i, 3] = 0 end +""" +This function convert spin to corresponding integer value used in spec +""" +function spin2num1d(σ) + if σ == '↑' + return 2 + elseif σ == '↓' + return 3 + elseif σ == '∅' + return 1 + end + error("illegal spin char for spin2num") +end """ diff --git a/test/runtests.jl b/test/runtests.jl index af09249..84817bc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,7 +2,8 @@ using ACEpsi using Test @testset "ACEpsi.jl" begin - #@testset "BFwf" begin include("test_bflow1.jl") end + @testset "BFwf" begin include("test_bflow.jl") end @testset "BFwf_lux" begin include("test_bflow_lux.jl") end @testset "AtomicOrbitals" begin include("test_atomicorbitals.jl") end + @testset "AtomicOrbitalsBasis" begin include("test_atorbbasis.jl") end end diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl index 58ee23a..bd552e7 100644 --- a/test/test_atorbbasis.jl +++ b/test/test_atorbbasis.jl @@ -7,7 +7,6 @@ using ACEbase.Testing: print_tf, fdtest using LuxCore using Random using Zygote -using Cthulhu # test configs Rnldegree = 4 @@ -21,7 +20,13 @@ nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] ## # Defining AtomicOrbitalsBasis -bRnl = ACEpsi.AtomicOrbitals.RnlExample(Rnldegree) +n1 = 5 +n2 = 1 +Pn = Polynomials4ML.legendre_basis(n1+1) +spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] +ζ = 10 * rand(length(spec)) +Dn = SlaterBasis(ζ) +bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) bYlm = RYlmBasis(Ylmdegree) spec1p = make_nlms_spec(bRnl, bYlm; totaldegree = totdegree) diff --git a/test/test_bflow.jl b/test/test_bflow.jl index fc8fdc8..f8d306f 100644 --- a/test/test_bflow.jl +++ b/test/test_bflow.jl @@ -98,9 +98,8 @@ function _fdtest(F, Σ, dF, x::AbstractVector; h0 = 1.0, verbose=true) const ↑, ↓, ∅ = '↑','↓','∅' Nel = 5 polys = legendre_basis(8) -wf = BFwf1(Nel, polys; ν=3, trans = x -> 2/pi * atan(x)) - -X = 2 * rand(Nel) .- 1 +wf = BFwf1(Nel, polys; ν = 3) +X = randn(Nel) Σ = rand([↑, ↓], Nel) wf(X, Σ) @@ -112,7 +111,7 @@ using LinearAlgebra using Printf @info("Fd test of gradient w.r.t. X") -_fdtest(wf, Σ, g, X) +@test _fdtest(wf, Σ, g, X) # ## From 2cde5c6498ccb520ec6a3de96768b599f575d81a Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 8 Aug 2023 10:05:03 -0700 Subject: [PATCH 71/79] swap out dense layer by LinearLayer in P4ML --- src/bflow3d.jl | 31 ++----------------------------- 1 file changed, 2 insertions(+), 29 deletions(-) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index db0ea04..4fc10f5 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -1,7 +1,6 @@ using Polynomials4ML, Random -using Polynomials4ML: OrthPolyBasis1D3T -using Polynomials4ML: PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! +using Polynomials4ML: OrthPolyBasis1D3T, LinearLayer, PooledSparseProduct, SparseSymmProdDAG, SparseSymmProd, release! using Polynomials4ML.Utils: gensparse using LinearAlgebra: qr, I, logabsdet, pinv, mul!, dot , tr, det import ForwardDiff @@ -22,32 +21,6 @@ struct MaskLayer <: AbstractExplicitLayer end (l::MaskLayer)(Φ, ps, st) = Φ .* [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX], st - -## - -struct DenseLayer <: AbstractExplicitLayer - in_dim::Integer - out_dim::Integer -end - -function (l::DenseLayer)(x::AbstractMatrix, ps, st) - return unwrap(x) * ps.W, st -end - -# Jerry: Maybe we should use Glorot Uniform if we have no idea about what we should use? -LuxCore.initialparameters(rng::AbstractRNG, l::DenseLayer) = ( W = randn(rng, l.out_dim, l.in_dim), ) -LuxCore.initialstates(rng::AbstractRNG, l::DenseLayer) = NamedTuple() - -function ChainRulesCore.rrule(::typeof(Lux.apply), l::DenseLayer, x::AbstractMatrix, ps, st) - val = l(x, ps, st) - function pb(A) - return NoTangent(), NoTangent(), A[1] * ps.W', (W = x' * A[1],), NoTangent() - end - return val, pb -end - -## - function get_spec(nuclei, spec1p) spec = [] Nnuc = length(nuclei) @@ -110,7 +83,7 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, _det = x -> size(x) == (1, 1) ? x[1,1] : det(Matrix(x)) BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), - bAA = corr_layer, hidden1 = DenseLayer(Nel, length(corr1)), + bAA = corr_layer, hidden1 = LinearLayer(length(corr1), Nel), Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> _det(x))) return Chain(; branch = BranchLayer(; js = jastrow_layer, bf = BFwf_chain, ), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ), spec, spec1p end From a2e16eddba2a4786fc244168f1688d93849076d1 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 8 Aug 2023 15:20:33 -0700 Subject: [PATCH 72/79] add benchmark --- benchmarks/benchmark_bflow_lux.jl | 58 +++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 benchmarks/benchmark_bflow_lux.jl diff --git a/benchmarks/benchmark_bflow_lux.jl b/benchmarks/benchmark_bflow_lux.jl new file mode 100644 index 0000000..a7c0002 --- /dev/null +++ b/benchmarks/benchmark_bflow_lux.jl @@ -0,0 +1,58 @@ +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow +using ACEpsi.vmc: gradient, laplacian, grad_params +using ACEbase.Testing: print_tf, fdtest +using LuxCore +using Lux +using Zygote +using Optimisers # mainly for the destrcuture(ps) function +using Random +using Printf +using LinearAlgebra +using BenchmarkTools +using HyperDualNumbers: Hyper + +Rnldegree = n1 = 2 +Ylmdegree = 3 +totdegree = 3 +Nel = 2 +X = randn(SVector{3, Float64}, Nel) +Σ = rand(spins(), Nel) +nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ] + +# wrap it as HyperDualNumbers +x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3]) +hX = [x2dualwrtj(x, 0) for x in X] +hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec + +## + +# Defining AtomicOrbitalsBasis +n2 = 2 +Pn = Polynomials4ML.legendre_basis(n1+1) +spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] +ζ = rand(length(spec)) +Dn = GaussianBasis(ζ) +bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) +bYlm = RYlmBasis(Ylmdegree) + +# setup state +BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) + +## + +@info("Test evaluate") + +@btime BFwf_chain($X, $ps, $st) +@btime gradient($BFwf_chain, $X, $ps, $st) +@btime laplacian($BFwf_chain, $X, $ps, $st) + + +@profview let BFwf_chain = BFwf_chain, X = X, ps = ps, st = st + for i = 1:10_000 + BFwf_chain(X, ps, st) + end +end \ No newline at end of file From fde91cbb92f7140999895ffcb4c24cb860f5e220 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 8 Aug 2023 22:23:59 -0700 Subject: [PATCH 73/79] remove consistency test with ACESchrodinger --- test/ACESchrodingerRef/bftest.json | 1 - test/ACESchrodingerRef/orthopolyweights.json | 1 - 2 files changed, 2 deletions(-) delete mode 100644 test/ACESchrodingerRef/bftest.json delete mode 100644 test/ACESchrodingerRef/orthopolyweights.json diff --git a/test/ACESchrodingerRef/bftest.json b/test/ACESchrodingerRef/bftest.json deleted file mode 100644 index 53b37fc..0000000 --- a/test/ACESchrodingerRef/bftest.json +++ /dev/null @@ -1 +0,0 @@ -[{"X":[0.865674388682704,0.5593067117763364,0.7174253143031734,0.7441985665675926,0.5142646432433161],"P":[[0.0,-0.007273160835590364,0.016835134978959394,-0.05518907618250436,-0.013991696863045475,0.07727358901069811,-0.012208833323334363,0.00734169703086515,-0.010748681761954645,-0.004935955469421341,-0.002051589510234265,-0.007324816327756153,-0.0012150238303724283,-0.0041011296845646,-0.005622401475235145,-0.00033823061468851224,0.0018562963486714315,-0.0021593636943063716,-0.0012960853248113044,-9.257597403594705e-5,0.00015680119226319837,0.0010062258941460165,0.00020016178871157811,-0.001984019458548015,-0.0007327566145175912,0.0002311968723983094,0.0012465424624981008,-0.0002594792759211216,0.00024282324758588263,-0.0009735266760982919,0.0012701285395614488,-0.03120219973811509,-0.08762242223320123,-0.037451458620417544,0.015111096088431522,-0.012946126612891346,0.000558694985035626,0.001817837590677066,-0.010421408920535407,0.0034639646609274455,0.00079623752475448,-0.0048933534161406216,-0.00390375609672264,-0.004111600580489995,-0.0023695204420709254,0.0025733813382050785,-0.0028635128228424334,-7.909063891721852e-5,0.0029030381101404932,-0.0011625503460157905,-0.0013615615980300479,0.00021332697092271057,-0.00109786953861817,-0.0013941167446012129,-0.001490221295716501,-0.0001204696890062787,0.00015039779402193146,0.0015989019630892932,-0.0005821101180100074,0.0002013464602539328,0.0006329917978056687,-0.006915002618838396,-0.007600718250933177,-0.03858143382430678,-0.0175540066031235,-0.03318559508361785,0.008162432125145763,0.0021664489646505657,0.006806617189295271,-0.01451556369214969,-0.0086280196386515,-0.004818462018391713,0.0033186913105087367,-0.0005654313147579877,0.0022253635879058787,0.0001690587595510904,-0.002728984530868199,-0.0001950587651646309,-0.0011279545714458744,-0.0022226762547764706,0.0004916217375389001,0.0017909659462801437,0.0006744511701872369,-0.001551745035556326,-0.002895207427547411,-0.0006219214244312627,0.0009350389808923484,0.00016914793131505497,-0.0003452791792570605,-0.0003786816803868763,-0.01724391291318526,0.0015827677190570928,0.019830661143682055,-0.0005066998330636165,0.00020379950526259723,0.0034856005396233227,-0.000837864448065468,-0.00020115503547858517,-0.002460130404689392,0.0034914041464041315,0.0012819208111355205,0.0005390847087198669,0.00033834247158049996,0.00480614497351621,-0.006664690149082625,-0.0007453626878329316,0.00481376099061003,-0.0015549819968863534,-0.0008326186090949145,0.0018187397433018623,-0.001052692255136756,-0.0018153414744282587,-0.0003622576389311419,-0.00046909665822451006,-0.0006824731487953,-0.0007674462395043458,0.001152487981805814,0.00043849707611555937,0.0003175049893582542,-0.019043660671018557,-7.236050985906757e-5,-0.006225567859219423,0.0026038963224999487,-0.002046061988076383,-0.008413650104816227,-0.006267158027232268,0.001275688517479131,0.0015130214146165486,0.002907855885308389,0.0030036266149898513,0.0035278519117723705,-0.0017120549385533433,-0.0030989746933174816,0.0007666331347008639,0.0015189226547655697,0.0025970957249300064,-0.0020262326327436463,-0.0008831971546530497,0.00097209645752064,0.001798515115113,0.0005735720344046434,-0.0007705128076893775,0.0011114996248305841,0.00043085091688689853,0.0015350106998117316,-0.01020334057149607,-0.006671474048155625,-0.006691990485942631,-0.006925657242282331,-0.010209286713439432,-0.004499314807209286,-0.00349994458874293,0.004758877985775156,0.002439109930616411,0.0020473949096695148,-0.005942465920297297,0.0016042183709269917,-0.0012249775262228415,0.00210158086447522,0.00023977686102514238,0.0008672776381895977,0.002307169063277749,0.001541764645656077,0.0016122659784287463,-0.0005156455435007184,0.00021686290322525572,-0.0004924140087772139,-0.00041270045705382236,6.449869413117129e-5,-0.0009506759863607153,0.00015188792094967778,-0.014407666723961108,0.0027924612931630303,0.0028203461532720804,0.0038875131881751265,0.0046656720472962095,0.003121756952884793,0.00430676164985073,-0.00029668458285071677,-9.718441204005839e-5,-0.0032593092302272315,-0.0013437136275793354,7.475862453128584e-5,0.003180114218341685,-0.001050077202390746,0.002276099273747626,0.0001039980097248826,-0.00016037775522721693,-0.002060579020860672,0.00029309920297411837,0.0006729751887474921,0.0013929414579501016,-0.0007953082056825871,-0.0002467560848502057,0.00030325451502967254,1.1477998920399384e-5,-0.0030218974762138856,-0.0006020517613651615,-0.002648457204957628,0.0009168403790160222,-6.209893505350881e-5,-0.0009408360972925062,-0.0006902561796251357,-0.0006267627356047118,-0.0018815950845586325,0.001547085182744839,0.000779925182126374,0.0006746602298233217,0.0006488089538048761,0.0015598699777386066,-0.0003086159427842343,0.0009293670883072223,-0.001067785662864171,-0.0019115196768754617,-0.0003368689546790447,0.0005845479023218864,0.00023023192073515575,0.0013122170365125702,0.00022606709814607,-0.00031715989281457554,0.0016211182624186814,-0.0010148110061284302,0.0026628512691860397,0.003848579387554078,-0.005605156257441095,-0.0033637582833332272,-0.0011445573172373024,-0.0025159277350477885,-0.001585619490199665,0.00029816093712908853,1.341637085903415e-5,0.0004384131197410367,0.0013383489126430812,-0.0024946429295095323,0.00096931791419475,-0.00025100078458357067,0.0018234251508331873,-0.004305014808757479,-0.006040190820391921,0.0017377368014513215,-0.0015739783203371424,-0.003239626275420088,0.00223614893273366,-0.0018881738863567076,-0.0001427796704715549,-0.0012777487600203947,-0.002531410787272651,-0.0008390546235009237,0.0016817300735055617,0.0021300202652746376,0.00035628850132810394,-0.0009123377635161678,0.000618372094015173,0.0003406999378993869,0.001146975716821958,-0.0013172196655281774,-0.00015263422279787665,0.004042234168403429,0.0037452560245207402,-0.0007213939922429654,-0.0017252343168027335,0.0038763017419173105,0.0006787826109075411,-0.002340911386781388,-0.003326384460095757,0.00020595580781043344,0.0008929421361456529,-0.0015282751325236027,-0.0002532217139647597,0.0011252257475698126,0.0006188154666556426,0.000953459215436979,0.004475627653389925,-0.0032812080234037727,-0.0006571692913155281,-0.0013525487398722332,0.00194060137273032,-0.001620802951929255,-0.003499570981797633,0.0013122309270225882,-0.001942392898749902,-0.0008416427522398282,-0.0003937568759779808,-0.0014885866863977805,0.0011429989853829776,-0.0027537698924907012,-0.000851885168886742,-0.0032782921770560546,0.0032033836275734916,-0.0008045983269435984,-0.0005743553522711756,0.0026320272590839248,-0.0007435890545893948,0.0021025899252853334,0.00012260003559063013,-0.0010216120209657707,0.0007236231626159587,0.0021342688345096826,-0.00043697061634049135,-0.0009202521266326297,0.002879209006791936,-0.002140183433497032,0.0012134517692996014,-0.0015333183624764112,-0.0014382263138208152,0.00029763766886444507,-0.0009531660521184255,-0.0030912111743179337,-0.0003490169388219783,0.00042498437790248664,0.002130058338787467,0.0008957200665292495,-0.00037629573277469765,-0.002196930423843293,0.001466177008307038,0.0011736716521662536,-0.0016482242593388794,-0.0006018734840737325,0.001365207627867828,0.0023229609284834674,0.0022045806817354847,0.0018295904215531792,0.00022096611277266307,0.000903689171183338,-0.0012449085315198315,0.0005612445670031732,-0.0005762671502525746,-0.0022401928230521757],[0.0,-0.1415014645357159,0.03389722562049187,0.023969457534646535,0.0025789586766420767,0.013696179679857092,0.0004874807053287539,0.0022719882019194095,0.0027875258779416377,-0.011640730117812643,0.000941077665747104,0.005427135057118886,-0.007537562580797538,-0.004060026476873175,-0.0016552026056099159,-0.00032880518695875555,0.003718385943319974,0.0004029200354172692,-0.0013051176707870667,0.002439405725676743,0.0024972680098622377,-0.0017040225074165846,-0.0007543674847935729,-0.0007526713922301543,0.0004676740036212627,0.000778633211695623,-0.00028153576981056877,0.0010093164611364228,-0.0013198341530317496,6.359782165865776e-5,0.00042617370429111334,-0.06416247614801514,-0.05189577002323411,-0.020171621608307554,-0.0015242825292753033,0.013559890962133081,-0.00838203941094169,0.004836822934091367,-0.0008460654654378996,-0.007200788787867077,0.0034928318232435186,-0.0038034822741939344,0.0015878310378883353,0.0049451699995633445,0.0017350919453703291,-0.0017048357659562123,-0.0006153755612234325,0.0006945811921420385,-0.0038500833247343985,-0.0021193343060258094,0.0012146748354543754,2.1452229103609626e-5,-0.0012620399744945842,-0.0012228058700368208,-0.0014592033806210979,-0.0006834663476424457,-0.0009550066268569219,-0.0010835891300964228,-0.00010903423324790761,0.0006824577594390323,0.0005083487887822505,-0.033625345647078315,-0.032606979368444235,-0.0011417473414791505,0.0014443128564567342,0.017516251977753056,0.0012301184713490356,-2.9154193785463e-6,0.0017577224352692171,0.003653596685516202,0.00023624488454955346,-0.005923242198205954,0.0016740754963943463,0.0009757378573581285,0.0035911253805730875,0.0025650766807768744,-0.002073173176902431,-0.00030656236121647055,0.0004551403777908367,0.00015230199459836165,-0.00032510222335411035,-0.00024121898907764573,-0.0025802716141693784,0.00028484129938606827,-0.0007395547767799947,-0.0017197557176612584,0.0011972500231066186,0.000641256648915454,0.0003212171122332008,6.222078125770343e-5,-0.01731598416682352,0.008431158003897874,-0.009183858926086213,0.010929844500984033,0.008966934939013101,-0.009537970619365253,0.005757335743731189,-0.005865994219580617,0.004662518646578089,0.0038620501661730693,-0.0006168612883295073,-0.0019255752063450653,-0.0009215986974067855,0.0004839036480230309,0.0017514838331171608,-0.0016028696652950355,-0.0023923985606763205,-0.00036322863478577483,0.0017569820158912077,-0.0003767069453539647,0.0010454414246169365,0.00035013873988570294,0.0007107553790174087,-0.0008070956008169248,0.0009916482969493832,-0.00022074220067070668,0.0006318929074324685,-0.0007172779253403688,0.002891895359074585,0.004077198149168721,0.003284046969469344,-0.0006428534733165632,-0.008736284584531075,-0.00171509465192322,-0.0014401423693541971,-0.003641494027137925,0.002115925923916403,0.0033065335912449466,0.001208199738344879,-0.004226872902893575,0.0002449078414662611,0.0011423199213950479,-0.0012395864430411324,-0.0008738265032233334,0.0007083171010477472,-0.002123536794353635,0.0003612651905357475,-0.0014442473426582903,0.0004654383171098462,-0.00013659926288265266,-1.8936471649763285e-5,0.0007984005565635783,-0.0013169108708772034,-0.0005885073948266743,-4.827692111346017e-5,0.011551030582564722,-0.007112840538444509,5.5344534124386515e-5,0.0029736351126296275,-0.0006095507011469122,-0.0022727056287102317,0.0009589307011599023,0.004103178694479104,0.0042263223063950745,0.004622334458406624,-0.0001882531493094113,0.0003893512457331008,-0.00012795604771846403,-0.00024398719775172622,-0.00039977983796185204,-0.0012038849640740223,0.0010660724838580714,-0.0006960713096831122,6.34316512798159e-5,-0.0006516188817970727,0.0003890438561351152,-0.00032124553636511953,-0.0010239406095511024,0.0008595547264888112,-0.0004144223915150102,-0.0007578193188752278,0.0004967101170510563,0.0020909959048182763,-0.009355120140933933,-0.0008980593622959161,-0.0004574787831058427,9.743256267623031e-5,-0.004466375242725531,-0.0019794750014059926,-0.0027257323849125725,-0.0006720835337924623,-0.0024556313173180427,-0.0015335596306337803,-0.00032581872238008453,-0.0011599460344633354,-0.001211789767266311,-0.0010063867816401564,0.001735351279051978,0.0008116851164913951,0.0006291254751834421,0.0005378787812638067,0.0002672681452831687,0.0011770798195594497,0.0008020829210705136,-0.00043979814011494653,0.00035531133077323235,-0.00010827529786401046,0.0019433175143146336,-0.003093826080194325,-0.001487235016792136,0.002719009869361629,0.0025709665029451404,-0.0003353466722323951,-0.00044087809769537524,-0.001245792587253495,-0.002094183795271219,0.0004256034652395511,0.001161230349992884,0.004431291641480512,-0.0009633169671075653,0.002496383196031251,-0.00043758554159473097,0.0009687964115699736,0.00022181374778521053,-0.0002987706966402035,0.0007845514605978501,0.00022116849945458478,-0.0033070389014662267,0.004464086415957261,-0.003493931516370567,0.007708158346440653,-0.000589266003260251,-0.0015641110202145882,-0.002631726665078621,-0.002965835266920086,0.0008853309505307582,0.0009687690136981908,-0.0007507214597020755,-0.0008568549146337563,0.0006632181541579786,-0.0004833078325781586,-0.0005910619677281819,-0.0004826278418138978,-0.00019975616258931573,0.0004476934784564265,-0.0002021350888062385,-0.00036466348087067795,0.004283655598752125,-0.0028421946766189784,0.003593317783542888,0.0020459939287765846,0.00046718310113147687,0.003046468262816395,-0.00032920942992507965,0.0005443835212580235,-0.0041374119225665385,-8.10555056969186e-6,-0.002208170701622244,0.0002552947833695971,-0.0008145383461805036,-3.290341649087815e-5,-0.00034519257318151844,-0.0004570810547912395,-0.0012706805719087822,0.0010866498813543461,0.0004163710588487817,0.0011249396511859365,0.0003881711945219168,0.0009846422412594388,-0.0012042746999701285,-0.0020179831211198174,0.0004147143364935039,-0.0002380366188330622,-0.001913689400311321,0.002074737638719054,-0.0007647655994505971,-4.995452494484828e-5,0.00017041213878006095,0.0007414918572514411,-0.0006333078471490984,-0.0006942484942779821,-0.0016052332884127322,-0.0028935671643886447,0.0010183277298312337,-0.002163335912843784,-0.0003381440849861135,-0.0005022036905737118,-0.001062907643126416,-0.0012750005046449785,-0.00040014264086854007,0.0002353322515053723,-0.001673088881577215,-0.0007378558464321916,0.0007474990947378235,-0.0010324884572754205,-0.0022876920390454856,-0.0024018236692271017,0.0002700832593224041,-0.0018461956442216913,-0.0007187736794087258,0.0008001844891502288,0.0002207302826278403,-5.501999688902741e-5,7.208716865763236e-5,0.00013385539543922452,0.00033188035875307195,0.00023481013905259752,-0.0003766028049852175,0.0009561643056896107,-0.0013338769215958646,0.0007116849028602084,0.002329535787956947,0.0011229278933551979,-0.0004733601238994762,7.431131037837821e-6,-0.0007216073749576395,0.000529714588039374,-0.00039780257411345903,-0.0006131875286034484,-0.00214512239736763,0.0016627360135492277,0.0011382452794842244,6.423837861671132e-5,4.630065694135529e-5,7.802738027131558e-5,0.0019354199815806057,-0.0010431401026664692,-0.0010036577568998315,-0.0007202916468236336,-0.0007820606723768914,-6.7903850190233715e-6,0.0014193085961399115,0.0008515175723127755,5.387625947839787e-5,-0.0002910719631016457,-0.0005685460068819297,0.0008558063573472564,0.0002852396426840457],[0.0,-0.04402805754322954,-0.011142958961571527,-0.08355185615758533,0.007693057836078031,0.008122049065450365,0.014620488905996001,0.013363446056571266,-0.003909192028721517,-0.016326309914148164,-0.0009084404511918344,-0.000909721499046237,-0.013045523953146959,0.007944791735203387,0.005237424863455664,0.009071443576143916,-0.0034181574157606757,-0.000520000238696533,0.0008654754867416369,0.0004901290187083446,0.00011367666729697018,-0.0019599735293295063,0.0027663935251146746,0.001853162249835104,-0.0003921160809921587,-0.0008333880797629452,-0.0011039630020117331,-0.0005380107882431915,0.0022818292851276363,-0.001408490290052521,0.002372085705611642,-0.06501171552417831,0.03931962838066117,0.06671429363118775,0.04110530626743748,-0.016468920504775393,0.05310087125963052,-0.02156144368731037,0.005416850888297784,-0.00666128483988743,4.403586415969488e-5,-0.005247110590249385,0.005303421958830963,0.0033519452413559594,0.0004935586089650928,0.0011158361580500085,0.0012066733990109494,8.396064438458496e-5,-0.002590415228648016,0.0025816381050334097,0.002647505292393702,-1.5064779562123335e-5,0.0006931461755098325,-0.0011143324107808922,-0.0015215268555085041,-0.0005886855001973576,-0.0024173864641307465,0.0013905041919884602,0.0004926344598753639,-0.00020832945619871546,0.0005620741216666262,0.009783053438224228,0.08697588820022967,-0.004879529526639961,0.007542049198434136,0.03395731855284114,0.0012835565510867614,-0.009375665969310403,0.010148602202229908,0.0006842783448509161,0.004823590277254532,0.008260571872242954,-0.0033154065756634166,-0.0024863028459776036,0.00327411103440807,-0.001703641358365943,-0.002452753235582638,-0.0003952651975097854,-0.0002143458539703357,4.713807205955768e-5,0.000706683133291111,-0.0013234438411582358,0.0006065750590348115,0.000417423763979268,0.00038964776003984586,-0.0015594685766375203,0.0006788145188039321,0.0012682394085866164,-0.0010839331160297466,-0.000942795045233883,-0.0033156086330046438,0.03204771084118612,-0.006811021836947997,0.012649033563862877,0.006662179493776104,-0.007401725213088014,0.004520970556574989,-0.00159477324750003,0.002382057750515347,-0.004979235186786568,-0.00010056772852928918,0.0014916866676000094,0.0035232939692925462,-0.002270295806538275,-6.34926173722664e-5,0.0005768314239320663,-0.0011290210984718353,-0.003583493988041029,0.004257678732062989,-0.002019148804893395,-0.00048054189884131986,-0.001915026372407728,0.002226001017782774,0.0014360832467830821,-0.001735464787315388,0.0013434057579571543,-0.0013057855682751937,-0.0018876222034419265,-0.010258239959763157,-0.003074382724101865,-0.007447637565539284,-0.0050137566004224915,0.012100434782207037,0.012104921299818178,0.0044915070193090115,1.2555084738962484e-5,-0.006631439879646405,0.0018492364775181095,-0.0006667463849852597,-0.0006769973392898922,-0.00407935434449415,-0.0001324921889132793,-0.00012770304314821735,-0.0020790062014052067,0.00014192073468858067,0.0009680434531452223,0.001206538736714719,-0.0021861293815579577,0.0005333290902102261,-0.00018666261399317058,-0.000361888195218511,-7.909513539306612e-5,0.0002822125825621509,0.0006889852074675617,0.0009057106931889674,-0.03932349434282124,0.0013937456408790441,0.00779167772664732,0.017228823527280477,0.008307462897016827,-0.004276165621028891,0.003931943480262218,-0.002021544886301114,-0.003051605184114296,0.001652990394435074,-0.0019125164883423559,0.00038060930282128987,0.003343192451703172,0.0038836801418774963,0.0017505816599185623,0.0002864246939532309,-0.0014133675128177306,-0.0014283417161188175,0.00016087011763094182,-0.00035983113464459806,0.0005161047686018014,-0.0011238035739506056,0.00047112103207757716,-0.00045293544575067613,-0.00021089288594467178,0.0002148471261830107,-0.0006493604475529478,0.004842703847682515,-0.007813490082150334,0.005381696453756595,-0.0037133052007379133,0.0008532236705529915,-0.00011489177779547837,-0.0015738134977480492,-0.005373261318093009,-0.001869732349100704,7.388011166879513e-5,-0.0009350600458858547,3.013646394007428e-5,0.0002447844079466207,0.00013505054923966734,0.0012969418066866964,-0.0010376854446819991,-0.0011024019622741027,0.0015575854137105018,-1.187748855846456e-5,-0.001334325434432206,-0.0021191347699637997,-0.0015671937894741327,-0.0004955098365748804,0.00018539770164681111,0.0008493915249244286,-0.008175331207764591,-0.010719814223587898,0.0003931661254402116,-0.0017763814290830164,0.003311708935961154,-0.0004330927517330433,0.00017423467351135674,0.002587179395049181,-0.0010821699950854303,0.0008329475910821957,-0.00129349657377583,-0.001429025370821752,0.0007293106349703788,-0.0005955565586078021,-0.002739854580265446,0.0002384200220269314,0.0012496581397800947,0.0016121812292173568,0.0005730163943619895,-0.0002897753484847473,-0.006771639747668315,-0.004766772912445124,0.0003774473805386835,-0.003401937408364154,0.00112752753245001,-0.0017589574503238396,0.000734786784821611,-0.00037199693866514776,2.0491239172717422e-5,-0.002984910932501567,-0.0008654803560837526,-2.6186409968423833e-5,-0.0002225502414050418,-0.0009788027821722909,-0.000621392688684772,-0.0008905316315846754,-0.00017526186728466584,-0.0003341539379904442,0.0009629319081178466,-0.0017473020626600415,0.001184374913867703,-0.0018548226658740682,-0.00021183693867368953,-0.0009751885968951757,-0.0016107053724142718,0.0006604633275967079,0.0014161881691609245,0.0002800127948441434,-0.0004888999839426727,0.006143275067142836,0.00248242927134876,-0.002714376808795551,0.0030490190186851117,-0.0010341942651478529,0.0009150098676030077,0.00014690156554332965,0.001189141581451135,0.0008642327512319273,-0.0006276276801628641,0.00416342244039392,0.00461253216666197,0.0027302865057217056,0.0023642623900998953,0.000797301537349976,-0.0009496142539453641,-0.0026206892723266138,0.003054583381817993,-0.00031724197060036723,-8.379334445510498e-6,0.0005259166337449302,0.001001251562591875,0.0006049032540461225,-0.0004296440216764725,0.0015870454564155814,-0.00210698971181611,0.0017150192162836407,-0.000238657034065277,0.003058733680938888,0.0002708723922554791,-0.0012818049518332948,0.0015919010307478773,0.00048291670407832994,-0.001471674265555419,0.0010580793900541606,-3.261602460114984e-6,-0.00011315579085891213,0.0002142754613707886,0.001882378165803877,0.0019234805354636725,3.3777079039617604e-6,-0.0013696407432787835,-0.0006706186549689767,0.0023200499084230814,-0.0005432800541000589,-2.13892999411673e-5,0.002058854114475561,0.0006472978712984316,0.0007772420388506206,0.0007646638851072101,-0.0009493846413989477,-0.0009312409924910824,0.0006246346717898026,0.002058491285870386,-0.0035527847626472455,0.00033566961151770445,-0.0019841343292748403,0.0031847107298774084,-0.0019893279813137164,0.00038620783357522573,0.0012180373007744153,-0.0007132656701880948,-0.0009309917255598434,-0.0010838891750669986,-0.0005407783805422664,-1.0571773331642237e-5,0.002801791411167163,-0.0006579644426444493,0.0007738916239531171,-0.001071636884062243,-0.002202576366653914,-0.0022274201750426144,0.001438460330324162,0.0006302693683711359,-0.0008433664136932255,0.000678236164329376,-0.001075976524111195,0.0010435724722433009,-0.0005049818177228774,-0.0010366845334077895,-0.0018339613790994609,0.0008185203684755713],[0.0,0.08386970740551974,-0.0009704306216677576,-0.04592588008463851,0.0028685666346850015,0.011697220999670856,-0.011191614938387965,0.012759375798334523,0.0076885813293625805,-0.010441836843530634,-0.005853702178546021,-0.0038763203121232473,0.009362360450358098,-0.009967831820895345,-0.005260418881412472,-0.0029681372515559606,8.475607294543344e-5,-0.001989187371111607,-0.0009281661927614702,0.0001423988562177603,-0.003236642353563714,-0.002708530300814982,0.001669006292384854,0.0009584925945143888,-0.0011277269311636157,-0.0008964077731046439,0.0006807807387834202,-0.0008432977158589945,8.144121919870023e-5,0.0012541733191314473,0.0017023748281976617,0.08394985986960388,0.08089151516794894,0.048173035412213114,-0.031157233611531644,-0.007280896370582254,0.01535271010143602,0.01172848156913997,-0.008483364758833118,0.0040047205334004425,-0.009007200173387518,0.005639961608809443,0.00684173345333467,0.004341328180363217,0.0032977148482211086,-0.001626259062349758,-0.00011662957063659573,0.003799955976918868,0.0015279553421839642,-0.00039753522937658177,0.0017416131611761685,0.0002762383757838782,-0.002642735215462036,9.957383032887576e-5,-0.0011690958117569698,-0.0008643196858167544,0.0033429550029164534,0.0007656252491542861,0.0018589178554346917,-0.00023470163787052775,-4.5564349221181865e-5,-0.04708021499221474,-0.02461094734598668,-0.008965477341672684,-0.007238796642957464,-0.012613851043720021,-0.0036892716098293112,0.010458235730608256,-0.005975186899136275,0.014523669855506233,-0.00344622262700508,-0.003364473371098094,0.0009159395955967794,-0.006676497212212839,0.0019342230146693405,0.0004574917667275668,-0.001126378219379552,-0.00010582558083096318,0.0019879815957745336,0.0012152978009537761,-0.001854326006266154,0.0008535902428968038,-0.0027617665414098468,0.001740128007730472,-0.0002144713668677936,-0.003187755930819152,0.0007492188428399304,-0.00021293309234032254,-0.00019522761989200044,-0.0007258970198273656,-0.014956631260894006,-0.014480930921992862,0.005641218962234258,-0.016808840708572068,-0.0027254874054495402,-0.00928719566835116,0.0074661220139690294,0.001354033247278211,0.004933222472434316,0.000987488797272827,-0.004374381275238683,0.0002622807611444714,0.004583898115007326,0.00015870122766491036,0.0010636754461488618,0.0037628891724539935,0.0001523011513075129,-0.002532603804604314,0.00018930804895944342,0.0012530453535384242,-0.002319025756432803,0.00330161666828842,-0.00035504507401571154,-0.0006148453459930125,-0.001898693630071422,-0.0001813342417193562,0.0015752583259168306,-0.002187683131135863,-0.006025965351132038,-0.0001260552501196851,-0.010788184092271438,-0.01774803408872944,-0.010878135830473519,0.0023293875211294907,0.003211729248191183,0.0008676369685897532,-0.0025474869227369735,0.0042323901403919255,-0.005934769198670274,-0.0005773801029342905,0.004644567522015575,-0.00046124697349072223,-0.0017707804336112441,0.00206902921673181,0.0015139748320817678,-0.0029059826751051897,-0.0020962922733804962,-0.0008349597503149107,0.0035615462406532183,0.0004176653669464563,-0.0003107523755837688,0.0006899792019573006,0.00010419563912716009,0.0012295531762552387,-0.0009410602877643348,0.019513975111578274,-0.01398339413236611,-0.0050144803279762425,0.0032568604999901295,-0.0056279839545375155,-0.00022393440317839425,0.0022137768427113232,0.0013132735840286312,-0.0007332576873981875,0.003497472897330408,-0.0027800740407611184,-0.0012935280003663551,0.0010061008694466967,-0.0004508868870308679,-0.0033446699911708587,-0.0006889293869881519,-0.001852824275169578,0.001007650874079352,-0.0013693880271834007,0.001536880444284109,-0.002694917146410994,0.0011672250527397746,-0.0013247867941921139,0.0008738391224414858,-0.00015240992519973702,0.000930517877203611,-0.012112942480814192,-0.00491030075213,0.000911178940268876,0.004560419912531687,0.005811907896141531,-0.0013727069424115826,0.00045811054767689886,0.001673653858994744,0.0035808693978879717,0.0025178801041032994,0.002302246490716388,0.0050685925324770315,-0.0004221223125181713,-0.000648624705839004,-0.0012792170216180795,-0.001196919171934974,0.00044336963581870204,0.00065154515712236,0.0018339618059007633,0.0005134357838113822,-0.0004902313504978116,-5.72285690893406e-5,-0.0005687510353394116,-0.0025195276719083056,-0.0014450799819706599,-0.005393020128644126,-0.0006524545761657728,-0.0030716438559565675,-0.005321526498393636,-0.002714876438099475,0.007895331999310798,-0.003997047883807325,-0.0007582136620135889,-0.005355539027361355,0.0005753451144473372,-0.0006081422687335,0.001063966442888716,-0.003270757296706621,0.0006574548363340465,0.0015234035605945413,0.0008334992041817737,-0.0012248997152901717,0.0006968252411365405,0.0010735212052443376,-0.0022243402796182266,-0.0013276234743278847,-0.007905413262602525,-0.004283701426925043,-0.002943687529171942,-0.0010329011274543152,-0.0025809997575823975,-0.002561857953901345,-0.0024551518204565664,0.00048602991876391555,-0.00015879011746549327,-0.0019373293385056033,0.0015290705749877022,-0.0013502782511744598,-0.004327981073774667,0.004505639459318477,0.0008116708736342237,0.0007135010956583237,-0.0015816834097709634,-0.001198832377448572,-0.0005161310398885224,-0.0009804922047427346,0.0006744337214528709,-0.0005888197078642406,0.008197202656616784,0.0035149639205623567,-0.004284564573678459,0.0028658598650236344,-0.0012150575437991323,0.0015004292298309004,-0.002884749841116353,-0.003972969937024104,-0.0013145004493110043,-0.0010340417811295574,-0.0034854643798388233,-0.0015646718558845584,-0.00022542987125176702,0.0001580010509333612,-0.0008046643104319553,0.0008595865618353582,-0.002076309364368041,0.004503731623806539,9.000338536135329e-5,0.0014719216541901005,-0.00023809721975739665,-0.001067825665623421,0.004834179169484606,-8.45209963501152e-5,-0.0020270461624701175,0.0010993332726971427,0.0019282846905620042,0.0008428080251332878,0.000548688975593386,-0.0017879236797723576,-0.0005074793605964643,-0.000820108446125938,0.0005038678104534605,0.0026446923957827865,0.0025042487004194387,0.00363096703343181,0.00025430171332200733,0.001067397218429686,0.0011887892775203177,-0.0014597409694551888,-0.0003873545732979744,0.00022294337065004153,0.001782709120448477,-0.0009657157168968443,-0.0023430486582372797,0.00283446920402062,0.0016249986097077351,-0.0004595559487800244,-0.004396999360243408,0.0022798959646779357,0.0007396997008099389,-0.0018600194770316447,0.0004488038002656552,-0.0004904332718951571,-0.0017400405192807304,0.002406710143100841,-0.0008618779699613526,-0.001193172129634847,0.0006317976673705298,0.00017966878857362205,4.8320864076113224e-5,-0.0008682564303648513,0.0005057191142357151,-0.0013995181609154559,-0.0009428957286025687,0.0009318933104878499,0.0018077042096996628,-0.0015292130806268092,0.001024262942456958,0.005212085913413507,-0.001933593782977388,0.0008320796960362342,0.0010255566339031693,-0.0024562652697863703,0.00027756092078103274,-0.0012496514255268569,0.005542501308233257,0.0011621541630538568,0.0013907263942731648,0.0015505810059378332,-0.0002642823414091117,-0.002013810485467708,-0.002121666676426142,-0.0013219769509399452,0.0015649422329395966,-0.002046314408365764,-9.74143649550335e-5,-0.0032759648441764507,0.0010095299905392905],[0.0,-0.029408724152558186,-0.07961485608575579,0.02429575161434835,0.012695165514308115,0.034068488481416366,0.004738260547612065,-0.014810662920320124,0.016466818076198042,-0.01599656483421443,0.0011903852857444446,0.006565802667208097,0.0042835390856049245,0.0019173407089454434,0.0006695561076745098,0.0009603816933978828,0.0015027008696173873,-0.00036629233732104577,-0.00013771907219151924,-0.0016196597311155367,-0.0008159605539326211,0.0021723680038954114,-0.002202971410170692,0.0002067039000915997,-0.0013505835324627415,0.00016975100363364763,0.0009294012772919252,-0.0017563765489161222,7.043996297012429e-5,-0.0005201823221314795,-0.0016720491394588781,-0.05019820961521885,0.030601506208131325,-0.026487825613379132,-0.00694758312979253,0.018854068705853363,-0.018853861732462732,0.007725252848314819,0.0015886333072159325,0.00025206450947173404,0.005136216682454122,-0.002760569053680654,0.0026998815750210203,0.00043926049717113143,-0.01138589353137207,-0.00045346931438390846,-0.003970880561849017,0.0018195884621474837,-0.0008243499865267709,-0.00118079344414198,-0.003542703097294454,0.0012291106496810644,0.0029148895627155,0.0025033476044724503,0.0009441961879264032,0.0012477960242394262,-0.0008551776661664053,0.001203813520334674,0.0008883090443908005,0.00017518775429020524,-0.0015112861710620369,0.011133823347798675,0.023148877744152908,0.0001452976963559497,-0.030071051914032556,-0.01039970998333628,0.007089975319574339,-0.0009275967441523235,0.010190114734163959,0.0038978273242153898,-0.0011939769276209148,0.001479477436782824,0.010303052040057464,-0.002378504695930744,-0.0024474054107217133,-0.0015574102653112926,0.0003037315375750598,-0.005150203865968845,0.0018444588288032819,-0.0010183470886383607,-0.0010469442861759865,-0.002177309984464359,-0.0006847821561775338,0.0006051240195543805,0.0002787617778268732,0.0015175161580350505,-0.0019739281151331875,0.0004627477855440448,0.00040091359001521727,0.0006306187330422004,0.07040621307290723,0.007867536843038318,0.01535579495109694,0.00805449132171127,0.015169656785107797,0.00770683970861356,-0.0014789689777540828,-0.002155866634194098,-0.003998601038954876,0.003701042632020009,-0.00011269424268033158,-0.009525295150184385,0.002304384646366992,0.0029793110186764424,-0.0015338397656715645,-0.002507088053703426,-0.0013140568792070653,0.0018741206698632176,-0.00213422617179473,0.0006258227207557781,-0.0016322421574461112,-0.0030104996643087823,0.0006628968099637554,0.0008337753278824599,0.0015701119337084581,0.00017972646297288597,9.115018399893803e-6,-0.0010693639994815225,-0.0036565907275016266,0.01048954721176878,0.009979825068175179,-0.0021501523876782747,-0.0002571947193638018,-0.0005231113734642835,-0.008381791053287072,0.0036110062181667984,0.01236811658984308,-0.002363760834755232,-0.002049340736992869,0.0034190399802136937,3.095013064701948e-5,0.0006881654687300348,-0.0006459685874635523,5.981300927238139e-5,0.0009929929669522516,-0.0007034139220815446,0.0025947507006204916,-0.0012465678930106438,0.0007298349856393977,-0.001311239016098477,-0.0002709496829527328,0.0003573570583260926,0.00017650325155420783,0.0004312778200459983,0.0002596828077437762,0.007668832630592485,0.0005668885227145245,-0.007445005445540795,-0.002180049158424807,0.004189576444394635,0.007336798795207091,-0.010760264137570938,0.004354057937301244,-0.0021417276148978774,0.00020343679966305575,-0.0011667986148276442,-0.001521548187795829,-0.004072542747398399,-0.0033692565513772188,0.0015418630815053921,0.00012751746427295814,-0.0015908183644575793,0.0014867070682234214,-0.00023476601173898662,2.1392727250567297e-5,0.0015439300440425665,0.0006302750553393902,-0.0011745832126825054,-0.0011917268128131756,0.00016698934751635427,-0.001246909903159716,0.010608565691786264,-0.004123101721741894,-0.003980937183441734,0.00052464458878169,0.004077728857357225,0.0050046230587418895,-0.0007116415348767724,-0.003524126603693882,0.004896585005985333,0.002278768126393503,0.0005344721996156696,0.003961012418138575,0.00025395796499933545,0.0011120121797384098,0.0005336754795107502,-0.0008495722132742464,-0.002129987392663864,-2.3393446854187832e-5,0.00037732769694367486,0.00010266718689826554,0.0001049838555660403,-0.00047712254635120404,-0.0006418516622543081,0.0009096422703404578,0.002196005396419578,0.007097008536745168,2.35839755921466e-6,0.0003232356109497796,0.004112241635083961,-0.002445571490645541,-0.005084662960157807,0.0012746625108936813,-0.004128282947508196,0.00019133088889600327,0.0018843123826957659,-0.004069364823367721,7.853215200029877e-5,-0.0012148357243647121,0.0004326981650370287,0.0019524404007283242,0.00046665980679603397,0.0023328425160715513,0.001506668866074644,-0.0006260524576469025,0.000751630045850698,0.0009400209322299511,0.0023567058073696093,-0.0021051219255385964,-0.0024695747483556496,-0.0007619377632184466,0.005174178860544337,-0.004239008687786166,0.003322399382719688,-0.004274322055981807,-0.0012111416082931455,0.0026293109799787186,0.0007014015943274976,0.0014882587750042902,0.0022473986538776927,-0.0007260684797390137,-0.0018601255192169858,-0.0025226276740071005,-2.823778125084083e-5,-0.0007475326349349818,-0.0001664865480409446,0.0013555876672334247,0.004517889021508121,-0.003080990384298615,-0.002944226629132014,0.00039422473970278555,-0.0018302926283059334,-0.0009153153730956427,-0.0009196996627754447,0.001182214608830942,0.0010125654826699834,-0.002922331161292741,0.0009748063584593677,0.0013127571762226987,0.0006242878896491025,0.0019185374649353874,0.0003799275583028376,0.0015817012209455483,-7.06858787727385e-5,-0.0003773612494805833,0.0005187857570435263,-0.005941911086557807,0.005347178996690836,-0.0043680735836008665,0.0022371845811890775,0.0003764863510579694,0.0003505903455593586,-0.0022425222468880953,-0.001444535117892683,0.0013382419869145035,0.000965932589276706,-0.00023727747402093013,-3.0971534290143194e-5,-0.00010206701238134864,-0.0007092478170961671,-0.0008259484893786002,0.00015204223937540336,0.0050102217489241254,2.097848647156118e-6,0.0063745142053332916,0.004323809882398249,0.0010068062753886603,-0.0005253075282202657,-0.002365265947834636,-0.0008032393134943724,-7.157078372089952e-5,0.0014606474433231562,-0.0003292806653449654,-9.743060576889866e-5,-0.0009870641358958533,-0.0033757541402940606,0.003445804863635187,-0.004400916716619389,0.00034731670621810816,-0.0019689166585952333,-0.0021550350170804517,0.0024851583069887754,-0.0013706845149323353,-0.0005462649105470952,0.0005302940174428687,-0.0013759041802922235,-0.0006509915292246132,-0.001729754701923505,0.000409807659427466,0.00014575905523010274,-0.0014846375025426633,0.0018183882442078555,-0.002363980622635043,-0.0010981201010358212,-0.0019384169361203842,3.7819805769491953e-5,-0.00039790593744764316,-0.0015708995366575157,-0.0007170786048079367,-0.000461402403979508,0.002020057264371874,-0.000500210646455958,0.00011327896585982094,-0.0006232406487397653,0.00016235786985906031,-0.00045079067914616624,-0.0020334490888162665,-0.00023768097308445344,0.0007160088437790237,-0.0012011295069093027,-0.0015064500353723694,-0.0016361780988608278,-0.002101848634816187,0.0007902590684578514,-0.001362987501192283,0.0016024013681653398,-0.00039855329053647943,0.0006006279265734773]],"refval":-39.00416165865923}] \ No newline at end of file diff --git a/test/ACESchrodingerRef/orthopolyweights.json b/test/ACESchrodingerRef/orthopolyweights.json deleted file mode 100644 index 7d619e8..0000000 --- a/test/ACESchrodingerRef/orthopolyweights.json +++ /dev/null @@ -1 +0,0 @@ -[{"ww":[0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335,0.0033333333333333335],"tdf":[-1.5655603390389137,-1.5550883635269477,-1.5446163880149817,-1.5341444125030157,-1.5236724369910497,-1.5132004614790837,-1.5027284859671177,-1.4922565104551517,-1.481784534943186,-1.47131255943122,-1.460840583919254,-1.450368608407288,-1.439896632895322,-1.429424657383356,-1.41895268187139,-1.408480706359424,-1.398008730847458,-1.387536755335492,-1.377064779823526,-1.36659280431156,-1.356120828799594,-1.345648853287628,-1.335176877775662,-1.324704902263696,-1.3142329267517303,-1.3037609512397643,-1.2932889757277983,-1.2828170002158323,-1.2723450247038663,-1.2618730491919004,-1.2514010736799344,-1.2409290981679684,-1.2304571226560024,-1.2199851471440364,-1.2095131716320704,-1.1990411961201044,-1.1885692206081384,-1.1780972450961724,-1.1676252695842064,-1.1571532940722404,-1.1466813185602747,-1.1362093430483087,-1.1257373675363427,-1.1152653920243767,-1.1047934165124107,-1.0943214410004447,-1.0838494654884787,-1.0733774899765127,-1.0629055144645467,-1.0524335389525807,-1.0419615634406147,-1.0314895879286488,-1.0210176124166828,-1.0105456369047168,-1.0000736613927508,-0.9896016858807849,-0.9791297103688189,-0.9686577348568529,-0.9581857593448869,-0.9477137838329209,-0.937241808320955,-0.9267698328089891,-0.9162978572970231,-0.9058258817850571,-0.8953539062730911,-0.8848819307611251,-0.8744099552491591,-0.8639379797371932,-0.8534660042252272,-0.8429940287132612,-0.8325220532012952,-0.8220500776893293,-0.8115781021773633,-0.8011061266653973,-0.7906341511534313,-0.7801621756414654,-0.7696902001294994,-0.7592182246175334,-0.7487462491055674,-0.7382742735936014,-0.7278022980816354,-0.7173303225696694,-0.7068583470577035,-0.6963863715457376,-0.6859143960337716,-0.6754424205218056,-0.6649704450098396,-0.6544984694978736,-0.6440264939859076,-0.6335545184739416,-0.6230825429619756,-0.6126105674500097,-0.6021385919380438,-0.5916666164260778,-0.5811946409141118,-0.5707226654021458,-0.5602506898901798,-0.5497787143782138,-0.5393067388662478,-0.5288347633542819,-0.5183627878423159,-0.5078908123303499,-0.49741883681838395,-0.48694686130641796,-0.47647488579445196,-0.466002910282486,-0.45553093477052004,-0.44505895925855404,-0.43458698374658805,-0.4241150082346221,-0.4136430327226561,-0.40317105721069013,-0.39269908169872414,-0.3822271061867582,-0.3717551306747922,-0.3612831551628262,-0.3508111796508602,-0.3403392041388943,-0.3298672286269283,-0.3193952531149623,-0.30892327760299637,-0.2984513020910304,-0.2879793265790644,-0.2775073510670984,-0.26703537555513246,-0.25656340004316647,-0.24609142453120048,-0.2356194490192345,-0.22514747350726852,-0.21467549799530256,-0.20420352248333656,-0.1937315469713706,-0.1832595714594046,-0.17278759594743864,-0.16231562043547265,-0.1518436449235067,-0.1413716694115407,-0.13089969389957473,-0.12042771838760874,-0.10995574287564276,-0.09948376736367678,-0.0890117918517108,-0.07853981633974483,-0.06806784082777885,-0.05759586531581288,-0.0471238898038469,-0.03665191429188092,-0.026179938779914945,-0.015707963267948967,-0.005235987755982989,0.005235987755982989,0.015707963267948967,0.026179938779914945,0.03665191429188092,0.0471238898038469,0.05759586531581288,0.06806784082777885,0.07853981633974483,0.0890117918517108,0.09948376736367678,0.10995574287564276,0.12042771838760874,0.13089969389957473,0.1413716694115407,0.1518436449235067,0.16231562043547265,0.17278759594743864,0.1832595714594046,0.1937315469713706,0.20420352248333656,0.21467549799530256,0.22514747350726852,0.2356194490192345,0.24609142453120048,0.25656340004316647,0.26703537555513246,0.2775073510670984,0.2879793265790644,0.2984513020910304,0.30892327760299637,0.3193952531149623,0.3298672286269283,0.3403392041388943,0.3508111796508602,0.3612831551628262,0.3717551306747922,0.3822271061867582,0.39269908169872414,0.40317105721069013,0.4136430327226561,0.4241150082346221,0.43458698374658805,0.44505895925855404,0.45553093477052004,0.466002910282486,0.47647488579445196,0.48694686130641796,0.49741883681838395,0.5078908123303499,0.5183627878423159,0.5288347633542819,0.5393067388662478,0.5497787143782138,0.5602506898901798,0.5707226654021458,0.5811946409141118,0.5916666164260778,0.6021385919380438,0.6126105674500097,0.6230825429619756,0.6335545184739416,0.6440264939859076,0.6544984694978736,0.6649704450098396,0.6754424205218056,0.6859143960337716,0.6963863715457376,0.7068583470577035,0.7173303225696694,0.7278022980816354,0.7382742735936014,0.7487462491055674,0.7592182246175334,0.7696902001294994,0.7801621756414654,0.7906341511534313,0.8011061266653973,0.8115781021773633,0.8220500776893293,0.8325220532012952,0.8429940287132612,0.8534660042252272,0.8639379797371932,0.8744099552491591,0.8848819307611251,0.8953539062730911,0.9058258817850571,0.9162978572970231,0.9267698328089891,0.937241808320955,0.9477137838329209,0.9581857593448869,0.9686577348568529,0.9791297103688189,0.9896016858807849,1.0000736613927508,1.0105456369047168,1.0210176124166828,1.0314895879286488,1.0419615634406147,1.0524335389525807,1.0629055144645467,1.0733774899765127,1.0838494654884787,1.0943214410004447,1.1047934165124107,1.1152653920243767,1.1257373675363427,1.1362093430483087,1.1466813185602747,1.1571532940722404,1.1676252695842064,1.1780972450961724,1.1885692206081384,1.1990411961201044,1.2095131716320704,1.2199851471440364,1.2304571226560024,1.2409290981679684,1.2514010736799344,1.2618730491919004,1.2723450247038663,1.2828170002158323,1.2932889757277983,1.3037609512397643,1.3142329267517303,1.324704902263696,1.335176877775662,1.345648853287628,1.356120828799594,1.36659280431156,1.377064779823526,1.387536755335492,1.398008730847458,1.408480706359424,1.41895268187139,1.429424657383356,1.439896632895322,1.450368608407288,1.460840583919254,1.47131255943122,1.481784534943186,1.4922565104551517,1.5027284859671177,1.5132004614790837,1.5236724369910497,1.5341444125030157,1.5446163880149817,1.5550883635269477,1.5655603390389137]}] \ No newline at end of file From d6278771e755f7261d8cf44c10689deff9d4124e Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 8 Aug 2023 23:07:33 -0700 Subject: [PATCH 74/79] quick fix on pullback performance --- src/bflow3d.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 4fc10f5..7c60e64 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -85,9 +85,10 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), bAA = corr_layer, hidden1 = LinearLayer(length(corr1), Nel), Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> _det(x))) - return Chain(; branch = BranchLayer(; js = jastrow_layer, bf = BFwf_chain, ), prod = WrappedFunction(x -> prod(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ), spec, spec1p + return Chain(; branch = BranchLayer(; js = jastrow_layer, bf = BFwf_chain, ), prod = WrappedFunction(x -> x[1] * x[2]), logabs = WrappedFunction(x -> 2 * log(abs(x))) ), spec, spec1p end + function displayspec(spec, spec1p) nicespec = [] for k = 1:length(spec) From 30a1a3d7f4ca140034a837783c2d13ed6f17eb6c Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Sun, 3 Sep 2023 20:10:32 +0800 Subject: [PATCH 75/79] add sto-ng example --- Project.toml | 1 - examples/3D/Be.jl | 2 +- examples/3D/harmonic_oscillator.jl | 12 +++--- examples/3D/{H2.jl => sto-3g_H.jl} | 51 +++++++++++------------ examples/3D/sto-6g_H.jl | 63 ++++++++++++++++++++++++++++ examples/3D/sto-6g_H10.jl | 67 ++++++++++++++++++++++++++++++ src/jastrow.jl | 4 +- src/lux_utils.jl | 8 ++-- src/vmc/metropolis.jl | 17 +++++--- src/vmc/optimisers/adamw.jl | 7 ++-- 10 files changed, 182 insertions(+), 50 deletions(-) rename examples/3D/{H2.jl => sto-3g_H.jl} (62%) create mode 100644 examples/3D/sto-6g_H.jl create mode 100644 examples/3D/sto-6g_H10.jl diff --git a/Project.toml b/Project.toml index 7d183ab..bec3ab1 100644 --- a/Project.toml +++ b/Project.toml @@ -33,7 +33,6 @@ BenchmarkTools = "1" ForwardDiff = "0.10" JSON = "0.21" ObjectPools = "0.3.1" -Polynomials4ML = "0.2.2" julia = "1" [extras] diff --git a/examples/3D/Be.jl b/examples/3D/Be.jl index efdaa98..c441ae8 100644 --- a/examples/3D/Be.jl +++ b/examples/3D/Be.jl @@ -47,7 +47,7 @@ Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 ham = SumH(K, Vext, Vee) sam = MHSampler(wf, Nel, Δt = 0.5, burnin = 1000, nchains = 2000) -opt_vmc = VMC(3000, 0.1, AdamW(), lr_dc = 100) +opt_vmc = VMC(3000, 0.1, ACEpsi.vmc.adamW(), lr_dc = 100) wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st) err = err_opt diff --git a/examples/3D/harmonic_oscillator.jl b/examples/3D/harmonic_oscillator.jl index 8c9e382..fbdfba9 100644 --- a/examples/3D/harmonic_oscillator.jl +++ b/examples/3D/harmonic_oscillator.jl @@ -1,7 +1,7 @@ using ACEpsi, Polynomials4ML, StaticArrays, Test using Polynomials4ML: natural_indices, degree, SparseProduct using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate -using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC using ACEbase.Testing: print_tf, fdtest using LuxCore @@ -35,12 +35,11 @@ bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) bYlm = RYlmBasis(Ylmdegree) # setup state -wf = BFwf_chain = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) +wf, spec, spec1p = BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) wf(X, ps, st) - K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) Vext(wf, X::AbstractVector, ps, st) = 1/2 * sum([X[i][j]^2 for i = 1:length(X) for j = 1:3]) #sum(1/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) @@ -48,6 +47,7 @@ Vext(wf, X::AbstractVector, ps, st) = 1/2 * sum([X[i][j]^2 for i = 1:length(X) f Vee(wf, X::AbstractVector, ps, st) = 0 # sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) ham = SumH(K, Vext, Vee) -sam = MHSampler(wf, Nel, Δt = 0.3, burnin = 10) -opt = VMC(10, 0.05, lr_dc = 30) -wf, err_opt, ps = gd_GradientByVMC(opt, sam, ham, wf, ps, st) \ No newline at end of file +sam = MHSampler(wf, Nel, nuclei, Δt = 1.0, burnin = 1000, nchains = 2000) + +opt_vmc = VMC(100, 0.05, ACEpsi.vmc.adamW(); lr_dc = 30.0) +wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st) diff --git a/examples/3D/H2.jl b/examples/3D/sto-3g_H.jl similarity index 62% rename from examples/3D/H2.jl rename to examples/3D/sto-3g_H.jl index 660ad56..d37652c 100644 --- a/examples/3D/H2.jl +++ b/examples/3D/sto-3g_H.jl @@ -14,27 +14,32 @@ using LinearAlgebra using BenchmarkTools using HyperDualNumbers: Hyper -n1 = Rnldegree = 4 -Ylmdegree = 2 +n1 = Rnldegree = 1 +Ylmdegree = 0 totdegree = 4 -Nel = 2 +Nel = 1 X = randn(SVector{3, Float64}, Nel) -Σ = [↑,↓] -nuclei = [ Nuc(SVector(0.0,0.0,-0.7), 1.0), Nuc(SVector(0.0,0.0,0.7), 1.0)] +Σ = [↑] +nuclei = [Nuc(SVector(0.0,0.0,0.0), 1.0)] ## # Defining AtomicOrbitalsBasis n2 = 1 Pn = Polynomials4ML.legendre_basis(n1+1) -spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] -ζ = 10 * rand(length(spec)) -Dn = SlaterBasis(ζ) +spec = [(n1 = 1, n2 = 1, l = 0) ] + +ζ = [0.3425250914E+01, 0.6239137298E+00, 0.1688554040E+00] +ζ = reshape(ζ, 1, length(ζ)) +D = [0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00] +D = reshape(D, 1, length(D)) + +Dn = STO_NG((ζ, D)) bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) bYlm = RYlmBasis(Ylmdegree) # setup state -wf, spec, spec1p = BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2) -displayspec(spec, spec1p) +ord = 1 +wf, spec, spec1p = BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = ord) ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) p, = destructure(ps) @@ -42,23 +47,17 @@ length(p) K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) -Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) +function Vee(wf, X::AbstractVector, ps, st) + nX = length(X) + if nX <=1 + return 0 + else + return sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) + end +end ham = SumH(K, Vext, Vee) -sam = MHSampler(wf, Nel, Δt = 0.5, burnin = 1000, nchains = 2000) +sam = MHSampler(wf, Nel, nuclei, Δt = 1.0, burnin = 1000, nchains = 2000) -opt_vmc = VMC(3000, 0.1, AdamW(), lr_dc = 100) +opt_vmc = VMC(3000, 0.1, ACEpsi.vmc.adamW(); lr_dc = 100.0) wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st) - -err = err_opt -per = 0.2 -err1 = zero(err) -for i = 1:length(err) - err1[i] = mean(err[Int(ceil(i-per * i)):i]) -end -err1 - -Eref = -1.8887557142857143 - -# using Plots -plot(abs.(err1 .- Eref), w = 3, yscale=:log10) \ No newline at end of file diff --git a/examples/3D/sto-6g_H.jl b/examples/3D/sto-6g_H.jl new file mode 100644 index 0000000..96dc451 --- /dev/null +++ b/examples/3D/sto-6g_H.jl @@ -0,0 +1,63 @@ +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec +using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC +using ACEbase.Testing: print_tf, fdtest +using LuxCore +using Lux +using Zygote +using Optimisers # mainly for the destrcuture(ps) function +using Random +using Printf +using LinearAlgebra +using BenchmarkTools +using HyperDualNumbers: Hyper + +n1 = Rnldegree = 1 +Ylmdegree = 0 +totdegree = 4 +Nel = 1 +X = randn(SVector{3, Float64}, Nel) +Σ = [↑] +nuclei = [Nuc(SVector(0.0,0.0,0.0), 1.0)] +## + +# Defining AtomicOrbitalsBasis +n2 = 1 +Pn = Polynomials4ML.legendre_basis(n1+1) +spec = [(n1 = 1, n2 = 1, l = 0) ] + +ζ = [0.3552322122E+02, 0.6513143725E+01, 0.1822142904E+01,0.6259552659E+00, 0.2430767471E+00, 0.1001124280E+00] +ζ = reshape(ζ, 1, length(ζ)) +D = [0.9163596281E-02, 0.4936149294E-01,0.1685383049E+00,0.3705627997E+00, 0.4164915298E+00, 0.1303340841E+00] +D = reshape(D, 1, length(D)) + +Dn = STO_NG((ζ, D)) +bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) +bYlm = RYlmBasis(Ylmdegree) + +# setup state +ord = 1 +wf, spec, spec1p = BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = ord) + +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) +p, = destructure(ps) +length(p) + +K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) +Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) +function Vee(wf, X::AbstractVector, ps, st) + nX = length(X) + if nX <=1 + return 0 + else + return sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) + end +end + +ham = SumH(K, Vext, Vee) +sam = MHSampler(wf, Nel, nuclei, Δt = 1.0, burnin = 1000, nchains = 2000) + +opt_vmc = VMC(3000, 0.1, ACEpsi.vmc.adamW(); lr_dc = 100.0) +wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st) diff --git a/examples/3D/sto-6g_H10.jl b/examples/3D/sto-6g_H10.jl new file mode 100644 index 0000000..980e8b5 --- /dev/null +++ b/examples/3D/sto-6g_H10.jl @@ -0,0 +1,67 @@ +using ACEpsi, Polynomials4ML, StaticArrays, Test +using Polynomials4ML: natural_indices, degree, SparseProduct +using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate +using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec +using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC +using ACEbase.Testing: print_tf, fdtest +using LuxCore +using Lux +using Zygote +using Optimisers # mainly for the destrcuture(ps) function +using Random +using Printf +using LinearAlgebra +using BenchmarkTools +using HyperDualNumbers: Hyper + +n1 = Rnldegree = 1 +Ylmdegree = 0 +totdegree = 20 +Nel = 10 +X = randn(SVector{3, Float64}, Nel) +Σ = [↑,↑,↑,↑,↑,↓,↓,↓,↓,↓] +nuclei = [Nuc(SVector(0.0,0.0,-4.5), 1.0),Nuc(SVector(0.0,0.0,-3.5), 1.0), +Nuc(SVector(0.0,0.0,-2.5), 1.0),Nuc(SVector(0.0,0.0,-1.5), 1.0), +Nuc(SVector(0.0,0.0,-0.5), 1.0),Nuc(SVector(0.0,0.0,0.5), 1.0), +Nuc(SVector(0.0,0.0,1.5), 1.0),Nuc(SVector(0.0,0.0,2.5), 1.0), +Nuc(SVector(0.0,0.0,3.5), 1.0),Nuc(SVector(0.0,0.0,4.5), 1.0)] +## + +# Defining AtomicOrbitalsBasis +n2 = 1 +Pn = Polynomials4ML.legendre_basis(n1+1) +spec = [(n1 = 1, n2 = 1, l = 0) ] + +ζ = [0.3552322122E+02, 0.6513143725E+01, 0.1822142904E+01,0.6259552659E+00, 0.2430767471E+00, 0.1001124280E+00] +ζ = reshape(ζ, 1, length(ζ)) +D = [0.9163596281E-02, 0.4936149294E-01,0.1685383049E+00,0.3705627997E+00, 0.4164915298E+00, 0.1303340841E+00] +D = reshape(D, 1, length(D)) + +Dn = STO_NG((ζ, D)) +bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) +bYlm = RYlmBasis(Ylmdegree) + +# setup state +ord = 2 +wf, spec, spec1p = BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = ord) + +ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) +p, = destructure(ps) +length(p) + +K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st) +Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X)) +function Vee(wf, X::AbstractVector, ps, st) + nX = length(X) + if nX <=1 + return 0 + else + return sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X)) + end +end + +ham = SumH(K, Vext, Vee) +sam = MHSampler(wf, Nel, nuclei, Δt = 0.5, burnin = 1, nchains = 2000) + +opt_vmc = VMC(10, 0.1, ACEpsi.vmc.adamW(); lr_dc = 100.0) +@profview wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st) diff --git a/src/jastrow.jl b/src/jastrow.jl index ee5d9c2..e110103 100644 --- a/src/jastrow.jl +++ b/src/jastrow.jl @@ -30,7 +30,7 @@ function evaluate(f::Jastrow, X::AbstractVector, Σ) # trans - return exp(γ) + return 1.0 end @@ -71,7 +71,7 @@ function evaluate(f::JPauliNet, X::AbstractVector, Σ) γ += -(1/4) / (1+norm(X[i] - X[j])) end end - return exp(γ) + return 1.0 end struct JPauliNetLayer <: AbstractExplicitLayer diff --git a/src/lux_utils.jl b/src/lux_utils.jl index bc02c93..0282c8a 100644 --- a/src/lux_utils.jl +++ b/src/lux_utils.jl @@ -2,12 +2,12 @@ using LuxCore using HyperDualNumbers function replace_namedtuples(nt, rp_st, Σ) - if :tmp in keys(nt) || :pool in keys(nt) || length(nt) == 0 + if :tmp in keys(nt) || :pool in keys(nt) || length(nt) == 0 || :ζ in keys(nt) return (; nt...) else - for i in 1:length(nt) + for i in 1:length(keys(nt)) # @show nt[i] - if length(nt[i]) == 0 + if length(nt[i]) == 0 rp_st = (; rp_st..., (; keys(nt)[i] => (Σ = Σ, ))...) else rp_st = (; rp_st..., keys(nt)[i] => replace_namedtuples(nt[i], (;), Σ)) @@ -21,4 +21,4 @@ function setupBFState(rng, bf, Σ) ps, st = LuxCore.setup(rng, bf) rp_st = replace_namedtuples(st, (;), Σ) return ps, rp_st -end \ No newline at end of file +end diff --git a/src/vmc/metropolis.jl b/src/vmc/metropolis.jl index f130af6..1ec0959 100644 --- a/src/vmc/metropolis.jl +++ b/src/vmc/metropolis.jl @@ -2,13 +2,16 @@ using StatsBase using StaticArrays using Optimisers export MHSampler +using ACEpsi.AtomicOrbitals: Nuc + """ `MHSampler` Metropolis-Hastings sampling algorithm. """ -mutable struct MHSampler +mutable struct MHSampler{T} Nel::Int - Δt::Float64 # step size (of Gaussian proposal) + nuclei::Vector{Nuc{T}} + Δt::Real # step size (of Gaussian proposal) burnin::Int # burn-in iterations lag::Int # iterations between successive samples N_batch::Int # batch size @@ -17,10 +20,10 @@ mutable struct MHSampler x0::Any # initial sampling walkerType::String # walker type: "unbiased", "Langevin" bc::String # boundary condition - type::Int64 # move how many electron one time + type::Int # move how many electron one time end -MHSampler(Ψ, Nel; Δt = 0.1, +MHSampler(Ψ, Nel, nuclei; Δt = 0.1, burnin = 100, lag = 10, N_batch = 1, @@ -29,7 +32,7 @@ MHSampler(Ψ, Nel; Δt = 0.1, wT = "unbiased", bc = "periodic", type = 1) = - MHSampler(Nel, Δt, burnin, lag, N_batch, nchains, Ψ, x0, wT, bc, type) + MHSampler(Nel, nuclei, Δt, burnin, lag, N_batch, nchains, Ψ, x0, wT, bc, type) """ @@ -71,7 +74,9 @@ type = "restart" """ function sampler_restart(sam::MHSampler, ps, st) - r0 = [sam.Δt * randn(SVector{3, Float64}, sam.Nel) for _ = 1:sam.nchains] + r = [[sam.nuclei[i].rr for j = 1:Int(ceil(sam.nuclei[i].charge))] for i = 1:length(sam.nuclei)] + r = reduce(vcat,r) + r0 = [sam.Δt * randn(SVector{3, Float64}, sam.Nel) + r for _ = 1:sam.nchains] Ψx0 = eval.(Ref(sam.Ψ), r0, Ref(ps), Ref(st)) acc = [] for _ = 1 : sam.burnin diff --git a/src/vmc/optimisers/adamw.jl b/src/vmc/optimisers/adamw.jl index b38355e..5eb968e 100644 --- a/src/vmc/optimisers/adamw.jl +++ b/src/vmc/optimisers/adamw.jl @@ -1,15 +1,14 @@ using Optimisers - -mutable struct AdamW <: opt +mutable struct adamW <: opt β::Tuple γ::Number ϵ::Number end -AdamW() = AdamW((9f-1, 9.99f-1), 0.0, eps()) +adamW() = adamW((9f-1, 9.99f-1), 0.0, eps()) -function Optimization(type::AdamW, wf, ps, st, sam::MHSampler, ham::SumH, α) +function Optimization(type::adamW, wf, ps, st, sam::MHSampler, ham::SumH, α) λ₀, σ, E, x0, acc = Eloc_Exp_TV_clip(wf, ps, st, sam, ham) g = grad(wf, x0, ps, st, E) st_opt = Optimisers.setup(Optimisers.AdamW(α, type.β, type.γ, type.ϵ), ps) From cf557a2bd879a60ae4d55f612bb61066ee39e705 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Tue, 5 Sep 2023 01:14:30 -0700 Subject: [PATCH 76/79] quick fix --- src/atomicorbitals/atorbbasis.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 1083e9c..d22df6b 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -184,7 +184,8 @@ function evaluate(l::AtomicOrbitalsBasisLayer, X, ps, st) ϕnlm = Zygote.Buffer(zeros(T, (Nnuc, Nel, Nnlm))) for I = 1:Nnuc - ϕnlm[I,:,:], _ = l.prodbasis(map(x -> x - nuc[I].rr, X), ps, st) + # ϕnlm[I,:,:], _ = l.prodbasis(map(x -> x - nuc[I].rr, X), ps, st) + ϕnlm[I,:,:], _ = l.prodbasis(.-(X, Ref(nuc[I].rr)), ps, st) end return copy(ϕnlm), st From c779077ab904c81949a99ad214dc0b35631dd895 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Sat, 9 Sep 2023 14:49:00 -0700 Subject: [PATCH 77/79] hand written pb for aobassis and more performance fix --- src/atomicorbitals/atorbbasis.jl | 48 +++++++++++++++++++++++++------- src/backflowpooling.jl | 2 +- src/bflow3d.jl | 20 +++++++++++-- 3 files changed, 57 insertions(+), 13 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index d22df6b..da4d71b 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -5,7 +5,8 @@ using Polynomials4ML: SparseProduct, _make_reqfields using LuxCore: AbstractExplicitLayer, AbstractExplicitContainerLayer using Random: AbstractRNG -using Polynomials4ML +using Polynomials4ML: _make_reqfields, @reqfields, POOL, TMP, META + using StaticArrays using LinearAlgebra: norm @@ -13,7 +14,8 @@ using ChainRulesCore using ChainRulesCore: NoTangent using Zygote -using Lux: Chain +using Lux: Chain, apply +using ObjectPools: acquire! struct Nuc{T} rr::SVector{3, T} @@ -166,13 +168,13 @@ end struct AtomicOrbitalsBasisLayer{L, T} <: AbstractExplicitContainerLayer{(:prodbasis, )} prodbasis::L nuclei::Vector{Nuc{T}} - meta::Dict{String, Any} + @reqfields() end Base.length(l::AtomicOrbitalsBasisLayer) = length(l.prodbasis.layer.ϕnlms.basis.spec) * length(l.nuclei) function AtomicOrbitalsBasisLayer(prodbasis, nuclei) - return AtomicOrbitalsBasisLayer(prodbasis, nuclei, Dict{String, Any}()) + return AtomicOrbitalsBasisLayer(prodbasis, nuclei, _make_reqfields()...) end function evaluate(l::AtomicOrbitalsBasisLayer, X, ps, st) @@ -181,16 +183,42 @@ function evaluate(l::AtomicOrbitalsBasisLayer, X, ps, st) Nel = size(X, 1) T = promote_type(eltype(X[1])) Nnlm = length(l.prodbasis.layers.ϕnlms.basis.spec) - - ϕnlm = Zygote.Buffer(zeros(T, (Nnuc, Nel, Nnlm))) + # acquire FlexArray/FlexArrayCached from state + ϕnlm = acquire!(l.pool, :ϕnlm, (Nnuc, Nel, Nnlm), T) + # inplace evaluation X for I = 1:Nnuc - # ϕnlm[I,:,:], _ = l.prodbasis(map(x -> x - nuc[I].rr, X), ps, st) - ϕnlm[I,:,:], _ = l.prodbasis(.-(X, Ref(nuc[I].rr)), ps, st) + X .-= Ref(nuc[I].rr) + ϕnlm[I,:,:], _ = l.prodbasis(X, ps, st) + X .+= Ref(nuc[I].rr) end - - return copy(ϕnlm), st + # ϕnlm should be released in the next layer + return ϕnlm, st end +function ChainRulesCore.rrule(::typeof(apply), l::AtomicOrbitalsBasisLayer, X::Vector{SVector{3, T}}, ps, st) where T + val = evaluate(l, X, ps, st) + nuc = l.nuclei + Nnuc = length(nuc) + Nel = size(X, 1) + function pb(dϕnlm) # dA is of a tuple (dAmat, st), dAmat is of size (Nnuc, Nel, Nnlm) + # first we pullback up to each Xts, which should be of size (Nnuc, Nel, 3) + dXts = [] + for I = 1:Nnuc + # inplace trans X + X .-= Ref(nuc[I].rr) + # pullback of productbasis[I], now I used productbasis but generalized to specified atom-dependent basis later + # pbI : (Nel, Nnlm) -> vector of length Nel of SVector{3, T} + _out, pbI = Zygote.pullback(l.prodbasis, X, ps, st) + # write to dXts + push!(dXts, pbI((dϕnlm[1][I,:,:], _out[2]))[1]) # out[2] is the state + # get back to original X + X .+= Ref(nuc[I].rr) + end + # finally sum all contributions from different I channel, reduces to vector of length Nel of SVector{3, T} again + return NoTangent(), NoTangent(), sum(dXts), NoTangent(), NoTangent() + end + return val, pb +end (l::AtomicOrbitalsBasisLayer)(X, ps, st) = evaluate(l, X, ps, st) diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl index ae3caa3..7e4b60b 100644 --- a/src/backflowpooling.jl +++ b/src/backflowpooling.jl @@ -80,6 +80,7 @@ function evaluate(pooling::BackflowPooling, ϕnlm::AbstractArray, Σ::AbstractVe end # inbounds release!(Aall) + release!(ϕnlm) return A end @@ -142,7 +143,6 @@ initialparameters(rng::AbstractRNG, l::BackflowPoolingLayer) = _init_luxparams(r initialstates(rng::AbstractRNG, l::BackflowPoolingLayer) = _init_luxstate(rng, l.basis) - # This should be removed later and replace by ObejctPools (l::BackflowPoolingLayer)(ϕnlm, ps, st) = evaluate(l.basis, ϕnlm, st.Σ), st diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 7c60e64..66b0265 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -17,10 +17,26 @@ using ChainRulesCore: NoTangent # ----------------- custom layers ------------------ struct MaskLayer <: AbstractExplicitLayer - nX::Integer + nX::Int64 +end + +(l::MaskLayer)(Φ, ps, st) = begin + T = eltype(Φ) + A::Matrix{Bool} = [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX] + val::Matrix{T} = Φ .* A + return val, st +end + +function ChainRulesCore.rrule(::typeof(LuxCore.apply), l::MaskLayer, Φ, ps, st) + T = eltype(Φ) + A::Matrix{Bool} = [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX] + val::Matrix{T} = Φ .* A + function pb(dΦ) + return NoTangent(), NoTangent(), dΦ[1] .* A, NoTangent(), NoTangent() + end + return (val, st), pb end -(l::MaskLayer)(Φ, ps, st) = Φ .* [st.Σ[i] == st.Σ[j] for j = 1:l.nX, i = 1:l.nX], st function get_spec(nuclei, spec1p) spec = [] Nnuc = length(nuclei) From 97f6d47a6020827864331ed6aa6629095fd4c230 Mon Sep 17 00:00:00 2001 From: cheukhinhojerry Date: Sat, 9 Sep 2023 23:14:55 -0700 Subject: [PATCH 78/79] hand written pb w.r.t parameter --- src/atomicorbitals/atorbbasis.jl | 18 +++++++++++++----- src/bflow3d.jl | 5 +++-- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index da4d71b..013cadc 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -195,27 +195,35 @@ function evaluate(l::AtomicOrbitalsBasisLayer, X, ps, st) return ϕnlm, st end -function ChainRulesCore.rrule(::typeof(apply), l::AtomicOrbitalsBasisLayer, X::Vector{SVector{3, T}}, ps, st) where T +function ChainRulesCore.rrule(::typeof(apply), l::AtomicOrbitalsBasisLayer{L, T}, X::Vector{SVector{3, TX}}, ps, st) where {L, T, TX} val = evaluate(l, X, ps, st) nuc = l.nuclei Nnuc = length(nuc) Nel = size(X, 1) function pb(dϕnlm) # dA is of a tuple (dAmat, st), dAmat is of size (Nnuc, Nel, Nnlm) # first we pullback up to each Xts, which should be of size (Nnuc, Nel, 3) - dXts = [] + dXts = Vector{SVector{3, TX}}[] + dps = deepcopy(ps) + for t in 1:length(dps.embed.Rn.ζ) + dps.embed.Rn.ζ[t] = zero(TX) + end for I = 1:Nnuc # inplace trans X X .-= Ref(nuc[I].rr) # pullback of productbasis[I], now I used productbasis but generalized to specified atom-dependent basis later # pbI : (Nel, Nnlm) -> vector of length Nel of SVector{3, T} - _out, pbI = Zygote.pullback(l.prodbasis, X, ps, st) + _out, pbI = Zygote.pullback(l.prodbasis::L, X, ps, st) # write to dXts - push!(dXts, pbI((dϕnlm[1][I,:,:], _out[2]))[1]) # out[2] is the state + Xts, _dp = pbI((dϕnlm[1][I,:,:], _out[2])) + push!(dXts, Xts) # out[2] is the state + for t in 1:length(dps.embed.Rn.ζ) + dps.embed.Rn.ζ[t] += _dp.embed.Rn.ζ[t] + end # get back to original X X .+= Ref(nuc[I].rr) end # finally sum all contributions from different I channel, reduces to vector of length Nel of SVector{3, T} again - return NoTangent(), NoTangent(), sum(dXts), NoTangent(), NoTangent() + return NoTangent(), NoTangent(), sum(dXts), dps, NoTangent() end return val, pb end diff --git a/src/bflow3d.jl b/src/bflow3d.jl index 66b0265..64754ac 100644 --- a/src/bflow3d.jl +++ b/src/bflow3d.jl @@ -100,8 +100,9 @@ function BFwf_lux(Nel::Integer, bRnl, bYlm, nuclei; totdeg = 15, _det = x -> size(x) == (1, 1) ? x[1,1] : det(Matrix(x)) BFwf_chain = Chain(; ϕnlm = aobasis_layer, bA = pooling_layer, reshape = WrappedFunction(reshape_func), bAA = corr_layer, hidden1 = LinearLayer(length(corr1), Nel), - Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> _det(x))) - return Chain(; branch = BranchLayer(; js = jastrow_layer, bf = BFwf_chain, ), prod = WrappedFunction(x -> x[1] * x[2]), logabs = WrappedFunction(x -> 2 * log(abs(x))) ), spec, spec1p + Mask = ACEpsi.MaskLayer(Nel), det = WrappedFunction(x -> _det(x)), logabs = WrappedFunction(x -> 2 * log(abs(x))) ) + # return Chain(; branch = BranchLayer(; js = jastrow_layer, bf = BFwf_chain, ), prod = WrappedFunction(x -> x[1] * x[2]), logabs = WrappedFunction(x -> 2 * log(abs(x))) ), spec, spec1p + return BFwf_chain, spec, spec1p end From 99f5653b3322a82932b4efafdded1c025527123e Mon Sep 17 00:00:00 2001 From: DexuanZhou Date: Sun, 10 Sep 2023 14:52:33 +0800 Subject: [PATCH 79/79] add pb for gradp of sto-ng --- examples/3D/sto-6g_H10.jl | 4 +++ src/atomicorbitals/atorbbasis.jl | 52 +++++++++++++++++++++----------- test/test_bflow_lux.jl | 8 +++-- 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/examples/3D/sto-6g_H10.jl b/examples/3D/sto-6g_H10.jl index 980e8b5..03b9e6c 100644 --- a/examples/3D/sto-6g_H10.jl +++ b/examples/3D/sto-6g_H10.jl @@ -60,6 +60,10 @@ function Vee(wf, X::AbstractVector, ps, st) end end +using BenchmarkTools +@btime wf(X, ps, st) +@btime gradient(X, ps, st) + ham = SumH(K, Vext, Vee) sam = MHSampler(wf, Nel, nuclei, Δt = 0.5, burnin = 1, nchains = 2000) diff --git a/src/atomicorbitals/atorbbasis.jl b/src/atomicorbitals/atorbbasis.jl index 013cadc..52e1a4e 100644 --- a/src/atomicorbitals/atorbbasis.jl +++ b/src/atomicorbitals/atorbbasis.jl @@ -199,31 +199,47 @@ function ChainRulesCore.rrule(::typeof(apply), l::AtomicOrbitalsBasisLayer{L, T} val = evaluate(l, X, ps, st) nuc = l.nuclei Nnuc = length(nuc) - Nel = size(X, 1) function pb(dϕnlm) # dA is of a tuple (dAmat, st), dAmat is of size (Nnuc, Nel, Nnlm) # first we pullback up to each Xts, which should be of size (Nnuc, Nel, 3) dXts = Vector{SVector{3, TX}}[] dps = deepcopy(ps) - for t in 1:length(dps.embed.Rn.ζ) - dps.embed.Rn.ζ[t] = zero(TX) - end - for I = 1:Nnuc - # inplace trans X - X .-= Ref(nuc[I].rr) - # pullback of productbasis[I], now I used productbasis but generalized to specified atom-dependent basis later - # pbI : (Nel, Nnlm) -> vector of length Nel of SVector{3, T} - _out, pbI = Zygote.pullback(l.prodbasis::L, X, ps, st) - # write to dXts - Xts, _dp = pbI((dϕnlm[1][I,:,:], _out[2])) - push!(dXts, Xts) # out[2] is the state + if :ζ in keys(dps.embed.Rn) for t in 1:length(dps.embed.Rn.ζ) - dps.embed.Rn.ζ[t] += _dp.embed.Rn.ζ[t] - end + dps.embed.Rn.ζ[t] = zero(TX) + end + for I = 1:Nnuc + # inplace trans X + X .-= Ref(nuc[I].rr) + # pullback of productbasis[I], now I used productbasis but generalized to specified atom-dependent basis later + # pbI : (Nel, Nnlm) -> vector of length Nel of SVector{3, T} + _out, pbI = Zygote.pullback(l.prodbasis::L, X, ps, st) + # write to dXts + Xts, _dp = pbI((dϕnlm[1][I,:,:], _out[2])) + push!(dXts, Xts) # out[2] is the state + for t in 1:length(dps.embed.Rn.ζ) + dps.embed.Rn.ζ[t] += _dp.embed.Rn.ζ[t] + end # get back to original X - X .+= Ref(nuc[I].rr) + X .+= Ref(nuc[I].rr) + end + # finally sum all contributions from different I channel, reduces to vector of length Nel of SVector{3, T} again + return NoTangent(), NoTangent(), sum(dXts), dps, NoTangent() + else + for I = 1:Nnuc + # inplace trans X + X .-= Ref(nuc[I].rr) + # pullback of productbasis[I], now I used productbasis but generalized to specified atom-dependent basis later + # pbI : (Nel, Nnlm) -> vector of length Nel of SVector{3, T} + _out, pbI = Zygote.pullback(l.prodbasis::L, X, ps, st) + # write to dXts + Xts, = pbI((dϕnlm[1][I,:,:], _out[2])) + push!(dXts, Xts) # out[2] is the state + # get back to original X + X .+= Ref(nuc[I].rr) + end + # finally sum all contributions from different I channel, reduces to vector of length Nel of SVector{3, T} again + return NoTangent(), NoTangent(), sum(dXts), NoTangent(), NoTangent() end - # finally sum all contributions from different I channel, reduces to vector of length Nel of SVector{3, T} again - return NoTangent(), NoTangent(), sum(dXts), dps, NoTangent() end return val, pb end diff --git a/test/test_bflow_lux.jl b/test/test_bflow_lux.jl index a96b8c3..e0d7d05 100644 --- a/test/test_bflow_lux.jl +++ b/test/test_bflow_lux.jl @@ -135,8 +135,8 @@ grad_test2(Fp, dFp, W0) ## @info("Test Δψ w.r.t. X using HyperDualNumbers") - -X = [randn(3) for _ = 1:Nel] +X = randn(SVector{3, Float64}, Nel) +XX = [Vector(x) for x in X] hX = [x2dualwrtj(x, 0) for x in X] Σ = rand(spins(), Nel) F(x) = BFwf_chain(x, ps, st)[1] @@ -162,9 +162,11 @@ for h in 0.1.^(1:8) Δfh = 0.0 for i = 1:Nel for j = 1:3 - XΔX_add, XΔX_sub = deepcopy(X), deepcopy(X) + XΔX_add, XΔX_sub = deepcopy(XX), deepcopy(XX) XΔX_add[i][j] += h XΔX_sub[i][j] -= h + XΔX_add = [SVector{3, Float64}(x) for x in XΔX_add] + XΔX_sub = [SVector{3, Float64}(x) for x in XΔX_sub] Δfh += (F(XΔX_add) - f0) / h^2 Δfh += (F(XΔX_sub) - f0) / h^2 end