test both static and dynamic

src/auxiliary.jl

@@ -60,6 +60,18 @@ function evaluate_ylm_ed(ace, Rs)
    return Zlm, dZlm
 end
 
+function evaluate_ylm_ed!(Zlm, dZlm, ace, Rs)
+   compute_with_gradients!(Zlm, dZlm, ace.ybasis, Rs)
+   return Zlm, dZlm
+end
+
+function whatalloc(::typeof(evaluate_ylm_ed!), ace, Rs)
+   TF = eltype(eltype(Rs))
+   return (TF, length(Rs), _len_ylm(ace.ybasis)), 
+          (SVector{3, TF}, length(Rs), _len_ylm(ace.ybasis))
+end
+
+
 
 # ------------------------------ 
 #  element embedding 
@@ -107,3 +119,10 @@ function (aadot::AADot)(A)
    return out 
 end
 
+function eval_and_grad!(∇φ_A, aadot::AADot, A)
+   φ = aadot(A)
+   ∇φ_A_1 = P4ML._pb_evaluate(aadot.aabasis, aadot.cc, A)
+   ∇φ_A .= unwrap(∇φ_A_1)
+   release!(∇φ_A_1)
+   return φ
+end

src/ncorr.jl

@@ -28,4 +28,6 @@ function generate_AA_dot(spec, c)
    return Polynomial(dynamic_poly)
 end
 
-
+function eval_and_grad!(∇_A, aadot, A)
+   return evaluate_and_gradient!(∇_A, aadot, A)
+end

src/splines.jl

@@ -65,6 +65,11 @@ function evaluate_ed(ace, rbasis, Rs, Zs)
    return Rn, dRn 
 end
 
+function whatalloc(::typeof(evaluate_ed!), basis::SplineRadialsZ, Rs, Zs)
+   TF = eltype(eltype(Rs))
+   return (TF, length(Rs), length(basis)), 
+          (SVector{3, TF}, length(Rs), length(basis))
+end
 
 function evaluate_ed!(Rn, dRn, basis::SplineRadialsZ, Rs, Zs)
    nX = length(Rs)
@@ -87,6 +92,6 @@ function evaluate_ed!(Rn, dRn, basis::SplineRadialsZ, Rs, Zs)
          dRn[ij, n] = g[n] * 𝐫̂ij
       end
    end
-   return nothing 
+   return Rn, dRn 
 end
 

src/uface_eval.jl

@@ -26,28 +26,27 @@ function ACEbase.evaluate(ace::UFACE_inner, Rs, Zs)
    Zlm = @withalloc evaluate_ylm!(ace, Rs)
 
    # pooling 
-   A = ace.abasis((Ez, Rn, Zlm))
+   TA = promote_type(eltype(Ez), eltype(Rn), eltype(Zlm))
+   A = @alloc(TA, length(ace.abasis))
+   P4ML.evaluate!(A, ace.abasis, (Ez, Rn, Zlm))
 
    # n correlations
    φ = ace.aadot(A)
 
-   # release the borrowed arrays 
-   release!(A)
-
    end
 
    return φ
 end
 
 function evaluate_ed(ace::UFACE, Rs, Zs, z0)
-   ∇φ = acquire!(ace.pool, :out_dEs, (length(Rs),), eltype(Rs))
+   ∇φ = zeros(eltype(Rs), length(Rs))
    return evaluate_ed!(∇φ, ace, Rs, Zs, z0)
 end
 
 function evaluate_ed!(∇φ, ace::UFACE, Rs, Zs, z0)
    i_z0 = _z2i(ace, z0)
    ace_inner = ace.ace_inner[i_z0]
-   φ, _ = evaluate_ed!(∇φ, ace_inner, Rs, Zs)
+   φ, _∇φ = evaluate_ed!(∇φ, ace_inner, Rs, Zs)
    add_evaluate_d!(∇φ, ace.pairpot, Rs, Zs, z0)
    φ += ace.E0s[z0] + evaluate(ace.pairpot, Rs, Zs, z0)
    return φ, ∇φ
@@ -59,34 +58,37 @@ function ACEbase.evaluate_ed!(∇φ, ace::UFACE_inner, Rs, Zs)
    rbasis = ace.rbasis 
    NZ = length(rbasis._i2z)
 
+   @no_escape begin
+
    # embeddings    
    # element embedding  (there is no gradient)
-   Ez = embed_z(ace, Rs, Zs)
+   Ez = @withalloc embed_z!(ace, Rs, Zs)
 
    # radial embedding 
-   Rn, dRn = evaluate_ed(ace, rbasis, Rs, Zs)
+   Rn, dRn = @withalloc evaluate_ed!(rbasis, Rs, Zs)
 
    # angular embedding 
-   Zlm, dZlm = evaluate_ylm_ed(ace, Rs)
+   Zlm, dZlm = @withalloc evaluate_ylm_ed!(ace, Rs)
 
    # pooling 
-   A = ace.abasis((Ez, Rn, Zlm))
+   TA = promote_type(eltype(Ez), eltype(Rn), eltype(Zlm))
+   A = @alloc(TA, length(ace.abasis))
+   P4ML.evaluate!(A, ace.abasis, (Ez, Rn, Zlm))
 
    # n correlations - compute with gradient, do it in-place 
-   ∂φ_∂A = acquire!(ace.pool, :∂A, size(A), TF)
-   φ = evaluate_and_gradient!(∂φ_∂A, ace.aadot, A)
+   ∂φ_∂A = @alloc(TA, length(A))
+   φ = eval_and_grad!(∂φ_∂A, ace.aadot, A)
 
    # backprop through A  =>  this part could be done more nicely I think
    ∂φ_∂Ez = BlackHole(TF) 
    # ∂φ_∂Ez = zeros(TF, size(Ez))
-   ∂φ_∂Rn = acquire!(ace.pool, :∂Rn, size(Rn), TF)
-   ∂φ_∂Zlm = acquire!(ace.pool, :∂Zlm, size(Zlm), TF)
+   ∂φ_∂Rn = @alloc(TF, size(Rn)...)
+   ∂φ_∂Zlm = @alloc(TF, size(Zlm)...)
    fill!(∂φ_∂Rn, zero(TF))
    fill!(∂φ_∂Zlm, zero(TF))
-   P4ML._pullback_evaluate!((∂φ_∂Ez, unwrap(∂φ_∂Rn), unwrap(∂φ_∂Zlm)), 
-                             unwrap(∂φ_∂A), 
-                             ace.abasis, 
-                             (unwrap(Ez), unwrap(Rn), unwrap(Zlm)); 
+   P4ML._pullback_evaluate!((∂φ_∂Ez, ∂φ_∂Rn, ∂φ_∂Zlm), 
+                             ∂φ_∂A, 
+                             ace.abasis, (Ez, Rn, Zlm); 
                              sizecheck=false)
 
    # backprop through the embeddings 
@@ -112,12 +114,7 @@ function ACEbase.evaluate_ed!(∇φ, ace::UFACE_inner, Rs, Zs)
       end
    end
 
-   # release the borrowed arrays 
-   release!(Zlm); release!(dZlm)
-   release!(Rn); release!(dRn)
-   release!(Ez)
-   release!(A)
-   release!(∂φ_∂Rn); release!(∂φ_∂Zlm); release!(∂φ_∂A)
+   end # no_escape
 
    return φ, ∇φ 
 end

test/test_import_ace1.jl

@@ -18,19 +18,38 @@ elements = [:Si,:O]
 # use totaldegree = 10, 12 for static aa 
 #     totaldegree > 15 for dynamic aa
 
-model = acemodel(; elements = elements, 
+@info("Dynamic Model: totaldegree = 15")
+model_dyn = acemodel(; elements = elements, 
                    order = 3, totaldegree = 15, 
                    Eref = Dict(:Si => -1.234, :O => -0.432))
-pot = model.potential
-pairpot = pot.components[1]
-mbpot = pot.components[2]
-pot1 = pot.components[3]
-@show length(mbpot.pibasis.inner[1])
+pot_dyn = model_dyn.potential
+pairpot_dyn = pot_dyn.components[1]
+mbpot_dyn = pot_dyn.components[2]
+pot1_dyn = pot_dyn.components[3]
+@show length(mbpot_dyn.pibasis.inner[1])
+
+
+@info("Static Model: totaldegree = 10")
+model_st = acemodel(; elements = elements, 
+                   order = 3, totaldegree = 10, 
+                   Eref = Dict(:Si => -1.234, :O => -0.432))
+pot_st = model_st.potential
+pairpot_st = pot_st.components[1]
+mbpot_st = pot_st.components[2]
+pot1_st = pot_st.components[3]
+@show length(mbpot_st.pibasis.inner[1])
+
+## 
+
+MODELS = [ (model_dyn, pot_dyn, pairpot_dyn, mbpot_dyn, pot1_dyn), 
+           (model_st,  pot_st,  pairpot_st,  mbpot_st,  pot1_st) ]
 
 ##
 # normalize the potential a bit so that all contributions are O(1) 
 # pot1 will be O(1) by construction  
 
+for (model, pot, pairpot, mbpot, pot1) in MODELS
+
 Nsample = 10_000
 pairpot.coeffs[:] = randn(length(pairpot.coeffs)) 
 a = sum( evaluate(pairpot, rand_env()...) for _ = 1:Nsample ) / Nsample
@@ -45,13 +64,13 @@ sum( evaluate(mbpot_, rand_env()...) for _ = 1:Nsample ) / Nsample
 
 pot.components[2] = mbpot = mbpot_
 
-
 # convert to UFACE format 
 uf_ace = UltraFastACE.uface_from_ace1(pot; n_spl_points = 10_000)
 
-## ------------------------------------
+# ------------------------------------
 
 @info("Test Consistency of ACE1 with UFACE")
+
 for ntest = 1:50 
    Rs, Zs, z0 = rand_env()
 
@@ -73,8 +92,9 @@ end
 println()
 
 
-## check gradient
+# check gradient ------------------------------
 
+@info("test gradients")
 for ntest = 1:30 
    Rs, Zs, z0 = rand_env()
    v1, dv1 = evaluate_ed(uf_ace, Rs, Zs, z0)
@@ -84,4 +104,7 @@ for ntest = 1:30
    dF = t -> (dV = evaluate_ed(uf_ace, Rs + t * U, Zs, z0)[2]; 
             sum( dot(dv, u) for (dv, u) in zip(dV, U) ) )
    print_tf(@test ACEbase.Testing.fdtest(F, dF, 0.0; verbose=false))
+end
+println()
+
 end