Merge pull request #20 from ACEsuit/co/dynaa

WIP: Enable dynamic AA to avoid stackoverflow

.github/workflows/CI.yml

@@ -19,8 +19,8 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.0'
           - '1.9'
+          - '1.10'
           - 'nightly'
         os:
           - ubuntu-latest
@@ -33,6 +33,10 @@ jobs:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v1
+      - run: |
+          using Pkg
+          Pkg.pkg"registry add https://github.com/ACEsuit/ACEregistry"
+        shell: bash -c "julia --color=yes {0}"
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
   docs:

Project.toml

@@ -7,6 +7,7 @@ version = "0.0.2"
 ACE1 = "e3f9bc04-086e-409a-ba78-e9769fe067bb"
 ACE1x = "5cc4c08c-8782-4a30-af6d-550b302e9707"
 ACEbase = "14bae519-eb20-449c-a949-9c58ed33163e"
+Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChunkSplitters = "ae650224-84b6-46f8-82ea-d812ca08434e"
 DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07"
@@ -20,19 +21,21 @@ ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
 Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpheriCart = "5caf2b29-02d9-47a3-9434-5931c85ba645"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StaticPolynomials = "62e018b1-6e46-5407-a5a7-97d4fbcae734"
 StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b"
+WithAlloc = "fb1aa66a-603c-4c1d-9bc4-66947c7b08dd"
 
 [compat]
 ChunkSplitters = "2"
 Folds = "0.2"
 LoopVectorization = "0.12"
 NamedTupleTools = "0.14"
 ObjectPools = "0.3"
-Polynomials4ML = "0.2.7"
+Polynomials4ML = "0.2.11"
 julia = "1"
 
 [extras]

profile/profile_evaluate.jl

@@ -14,7 +14,7 @@ pot = model.potential
 mbpot = pot.components[2]
 
 # convert to UFACE format 
-uf_ace = UltraFastACE.uface_from_ace1(mbpot; n_spl_points = 100)
+uf_ace = UltraFastACE.uface_from_ace1(pot; n_spl_points = 100)
 
 ## ------------------------------------
 

src/UltraFastACE.jl

@@ -4,6 +4,9 @@ import ACEbase
 import ACEbase: evaluate, evaluate!, 
                 evaluate_ed, evaluate_ed!
 
+using Bumper, WithAlloc, StrideArrays      
+import WithAlloc: whatalloc          
+
 _i2z(obj, i::Integer) = obj._i2z[i] 
 
 function _z2i(obj, Z)
@@ -28,6 +31,7 @@ include("auxiliary.jl")
 include("pair.jl")
 
 include("uface.jl")
+include("uface_eval.jl")
 
 include("julip_calculator.jl")
 

src/auxiliary.jl

@@ -1,3 +1,5 @@
+
+
 #
 # some auxiliary functions for UF_ACE evaluation 
 #
@@ -35,11 +37,21 @@ _len_ylm(ybasis) = (_get_L(ybasis) + 1)^2
 
 function evaluate_ylm(ace, Rs)
    TF = eltype(eltype(Rs))
-   Zlm = acquire!(ace.pool, :Zlm, (length(Rs), _len_ylm(ace.ybasis)), TF)
+   Zlm = zeros(TF, length(Rs), _len_ylm(ace.ybasis))
+   evaluate_ylm!(Zlm, ace, Rs)
+   return Zlm 
+end
+
+function evaluate_ylm!(Zlm, ace, Rs)
    compute!(Zlm, ace.ybasis, Rs)
    return Zlm 
 end
 
+function whatalloc(::typeof(evaluate_ylm!), ace, Rs)
+   TF = eltype(eltype(Rs))
+   return (TF, length(Rs), _len_ylm(ace.ybasis))
+end
+
 function evaluate_ylm_ed(ace, Rs)
    TF = eltype(eltype(Rs))
    Zlm = acquire!(ace.pool, :Zlm, (length(Rs), _len_ylm(ace.ybasis)), TF)
@@ -48,15 +60,31 @@ 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 
 
 
+# ------------------------------ 
+#  element embedding 
 
 function embed_z(ace, Rs, Zs)
    TF = eltype(eltype(Rs))
-   Ez = acquire!(ace.pool, :Ez, (length(Zs), length(ace.rbasis)), TF)
+   # Ez = acquire!(ace.pool, :Ez, (length(Zs), length(ace.rbasis)), TF)
+   Ez = zeros(TF, length(Zs), length(ace.rbasis))
+   return embed_z!(Ez, ace, Rs, Zs)
+end
+
+
+function embed_z!(Ez, ace, Rs, Zs)
    fill!(Ez, 0)
    for (j, z) in enumerate(Zs)
       iz = _z2i(ace.rbasis, z)
@@ -66,3 +94,35 @@ function embed_z(ace, Rs, Zs)
 end
 
 
+function whatalloc(::typeof(embed_z!), ace, Rs, Zs)
+   TF = eltype(eltype(Rs))
+   return (TF, length(Zs), length(ace.rbasis))
+end
+
+
+# ------------------------------ 
+#  aadot via P4ML 
+
+using LinearAlgebra: dot 
+
+struct AADot{T, TAA}
+   cc::Vector{T} 
+   aabasis::TAA
+end
+
+function (aadot::AADot)(A)
+   @no_escape begin 
+      AA = @alloc(eltype(A), length(aadot.aabasis))
+      P4ML.evaluate!(AA, aadot.aabasis, A)
+      out = dot(aadot.cc, AA)
+   end
+   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

@@ -25,11 +25,17 @@ end
 function evaluate(ace, basis::SplineRadialsZ, 
                   Rs::AbstractVector{<: SVector}, Zs::AbstractVector)
    TF = eltype(eltype(Rs))                  
-   Rn = acquire!(ace.pool, :Rn, (length(Rs), length(basis)), TF)
+   Rn = zeros(TF, length(Rs), length(basis))
    evaluate!(Rn, basis, Rs, Zs)
    return Rn 
 end
 
+function whatalloc(::typeof(ACEbase.evaluate!), 
+                   basis::SplineRadialsZ, Rs, Zs)
+   TF = eltype(eltype(Rs))                  
+   return (TF, length(Rs), length(basis))
+end
+
 
 function evaluate!(out, basis::SplineRadialsZ, Rs, Zs)
    nX = length(Rs)
@@ -49,6 +55,8 @@ function evaluate!(out, basis::SplineRadialsZ, Rs, Zs)
 end
 
 
+
+
 function evaluate_ed(ace, rbasis, Rs, Zs)
    TF = eltype(eltype(Rs))                  
    Rn = acquire!(ace.pool, :Rn, (length(Rs), length(rbasis)), TF)
@@ -57,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)
@@ -79,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