From f35ce0593c2411649008407addf993e2d355a1b1 Mon Sep 17 00:00:00 2001
From: ACEsuit <christohortner@gmail.com>
Date: Tue, 14 Nov 2023 18:23:05 -0800
Subject: [PATCH] removed all allocations

---
 src/splines.jl |  8 +++++++-
 src/uface.jl   | 39 ++++++++++++++++++++++++++++++++++-----
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/src/splines.jl b/src/splines.jl
index 2383fd4..d22f672 100644
--- a/src/splines.jl
+++ b/src/splines.jl
@@ -5,11 +5,17 @@ struct SparseStaticArray{N, T}
 end
 
 
-struct SplineRadialsZ{SPL, N}
+struct SplineRadialsZ{SPL, N, LEN}
    _i2z::NTuple{N, Int}
    spl::NTuple{N, SPL}
 end
 
+SplineRadialsZ(_i2z::NTuple{N, Int}, spl::NTuple{N, SPL}, LEN
+              ) where {N, SPL} = 
+         SplineRadialsZ{SPL, N, LEN}(_i2z, spl)
+
+Base.length(basis::SplineRadialsZ{SPL, N, LEN}) where {SPL, N, LEN} = LEN
+
 struct SplineRadials{SPL, N}
    _i2z::NTuple{N, Int}
    spl::NTuple{N, SPL}
diff --git a/src/uface.jl b/src/uface.jl
index b1be9d9..fb5bca5 100644
--- a/src/uface.jl
+++ b/src/uface.jl
@@ -2,6 +2,7 @@ import Polynomials4ML
 import ACEpotentials
 using Interpolations, ObjectPools
 import SpheriCart
+import SpheriCart: SphericalHarmonics, compute! 
 import ACEpotentials.ACE1
 import ACEpotentials.ACE1: AtomicNumber
 using LinearAlgebra: norm 
@@ -40,15 +41,35 @@ function ACEbase.evaluate(ace::UFACE, Rs, Zs, zi)
 end
 
 
+_get_L(ybasis::SphericalHarmonics{L}) where {L} = L 
+_len_ylm(ybasis) = (_get_L(ybasis) + 1)^2
+
 function ACEbase.evaluate(ace::UFACE_inner, Rs, Zs)
-   
+   TF = eltype(eltype(Rs))
    rbasis = ace.rbasis 
+   NZ = length(rbasis._i2z)
+
    _spl(rbasis, z) = rbasis.spl[UltraFastACE._z2i(rbasis, z)]
    
    # embeddings 
-   Ez = reduce(vcat, [ SVector((z .== rbasis._i2z)...)' for z in Zs ])
-   Rn = reduce(vcat, [ _spl(rbasis, Zs[j])(norm(Rs[j]))' for j = 1:length(Rs) ])
-   Zlm = ace.ybasis(Rs)
+   # Ez = reduce(vcat, [ SVector((z .== rbasis._i2z)...)' for z in Zs ])
+   Ez = acquire!(ace.pool, :Ez, (length(Zs), NZ), UInt8)
+   fill!(Ez, 0)
+   for (j, z) in enumerate(Zs)
+      iz = _z2i(rbasis, z)
+      Ez[j, iz] = 1
+   end
+
+   # Rn = reduce(vcat, [ _spl(rbasis, Zs[j])(norm(Rs[j]))' for j = 1:length(Rs) ])
+   Rn = acquire!(ace.pool, :Rn, (length(Rs), length(rbasis)), TF)
+   for (j, z) in enumerate(Zs) 
+      spl_j = _spl(rbasis, z)
+      Rn[j, :] .= spl_j(norm(Rs[j]))
+   end
+
+   # Zlm = ace.ybasis(Rs)
+   Zlm = acquire!(ace.pool, :Zlm, (length(Rs), _len_ylm(ace.ybasis)), TF)
+   compute!(Zlm, ace.ybasis, Rs)
    
    # pooling 
    A = ace.abasis((Ez, Rn, Zlm))
@@ -56,6 +77,12 @@ function ACEbase.evaluate(ace::UFACE_inner, Rs, Zs)
    # n correlations
    φ = ace.aadot(A)   
 
+   # release the borrowed arrays 
+   release!(Zlm)
+   release!(Rn)
+   release!(Ez)
+   release!(A)
+
    return φ
 end
 
@@ -100,9 +127,11 @@ function uface_from_ace1_inner(mbpot, iz; n_spl_points = 100)
 
    # radial embedding
    Rn_basis = mbpot.pibasis.basis1p.J
+   LEN_Rn = length(Rn_basis.J)
    spl = make_radial_splines(Rn_basis, zlist; npoints = n_spl_points)
    rbasis_new = SplineRadialsZ(Int.(t_zlist), 
-                     ntuple(iz1 -> spl[(zlist[iz1], z0)], length(zlist)))
+                     ntuple(iz1 -> spl[(zlist[iz1], z0)], length(zlist)), 
+                     LEN_Rn)
    # P4ML style spec of radial embedding 
    spec2i_Rn = 1:length(Rn_basis.J)