simplifying more tests

src/Polynomials4ML.jl

@@ -23,6 +23,8 @@ import LuxCore: AbstractExplicitLayer, AbstractExplicitContainerLayer,
 using Random: AbstractRNG     
 
 function _generate_input end 
+function _generate_batch end 
+
 function natural_indices end   # could rename this get_spec or similar ... 
 function index end
 function orthpolybasis end

src/ace/sparseprodpool.jl

@@ -57,6 +57,21 @@ _valtype(basis::PooledSparseProduct, BB::Tuple) =
 _gradtype(basis::PooledSparseProduct, BB::Tuple) = 
       mapreduce(eltype, promote_type, BB)
 
+function _generate_input_1(basis::PooledSparseProduct{NB}) where {NB} 
+   NN = [ maximum(b[i] for b in basis.spec) for i = 1:NB ]
+   BB = ntuple(i -> randn(NN[i]), NB)
+   return BB 
+end 
+
+function _generate_input(basis::PooledSparseProduct{NB}; nX = rand(5:15)) where {NB} 
+   NN = [ maximum(b[i] for b in basis.spec) for i = 1:NB ]
+   BB = ntuple(i -> randn(nX, NN[i]), NB)
+   return BB 
+end 
+
+function _generate_batch(basis::PooledSparseProduct, args...; kwargs...) 
+   error("PooledSparseProduct is not implemented for batch inputs")
+end
 
 
 # ----------------------- evaluation kernels 
@@ -90,11 +105,19 @@ function whatalloc(evaluate!, basis::PooledSparseProduct{NB}, BB::TupVecMat) whe
    return (TV, nA) 
 end
 
+
 function evaluate!(A, basis::PooledSparseProduct{NB}, BB::TupVec) where {NB}
-   BB_batch = ntuple(i -> reshape(BB[i], (1, length(BB[i]))), NB)
-   return evaluate!(A, basis, BB_batch)
+   spec = basis.spec
+   @assert length(A) >= length(spec)
+   fill!(A, 0)
+   @inbounds for (iA, ϕ) in enumerate(spec)
+      b = ntuple(t -> BB[t][ϕ[t]], NB)
+      A[iA] = @fastmath(prod(b))
+   end
+   return A
 end
 
+
 function evaluate!(A, basis::PooledSparseProduct{NB}, BB::TupMat, 
                    nX = size(BB[1], 1)) where {NB}
    @assert all(B->size(B, 1) >= nX, BB)
@@ -139,21 +162,38 @@ end
 using StaticArrays
 using Base.Cartesian: @nexprs 
 
-# NOT AT ALL OBVIOUS HOW TO DO THIS ?!?!?
-# function whatalloc(::typeof(pullback_evaluate!), 
-#                    basis::PooledSparseProduct, BB::TupMat)
-
-# end
 
-function pullback_evaluate(∂A, basis::PooledSparseProduct{NB}, BB::TupMat) where {NB}
-   nX = size(BB[1], 1)
+function whatalloc(::typeof(pullback_evaluate!), 
+                   ∂A, basis::PooledSparseProduct{NB}, BB::TupMat) where  {NB}
    TA = promote_type(eltype.(BB)..., eltype(∂A))
-   ∂BB = ntuple(i -> zeros(TA, size(BB[i])...), NB)
-   pullback_evaluate!(∂BB, ∂A, basis, BB)
-   return ∂BB
+   return ntuple(i -> (TA, size(BB[i])...), NB)                   
 end
 
 
+# function pullback_evaluate(∂A, basis::PooledSparseProduct{NB}, BB::TupMat) where {NB}
+#    nX = size(BB[1], 1)
+#    TA = promote_type(eltype.(BB)..., eltype(∂A))
+#    ∂BB = ntuple(i -> zeros(TA, size(BB[i])...), NB)
+#    pullback_evaluate!(∂BB, ∂A, basis, BB)
+#    return ∂BB
+# end
+
+# the next three method definitions ensure that we can use the 
+# WithAlloc stuff with the pullback_evaluate! function.
+# TODO: this should probably be replaced with a loop that generates  
+# the code up to a large-ish NB. 
+
+pullback_evaluate!(∂B1, ∂A, basis::PooledSparseProduct{1}, BB::TupMat) = 
+         pullback_evaluate!((∂B1,), ∂A, basis, BB)
+
+pullback_evaluate!(∂B1, ∂B2, ∂A, basis::PooledSparseProduct{2}, BB::TupMat) = 
+         pullback_evaluate!((∂B1, ∂B2,), ∂A, basis, BB)
+
+pullback_evaluate!(∂B1, ∂B2, ∂B3, ∂A, basis::PooledSparseProduct{3}, BB::TupMat) = 
+         pullback_evaluate!((∂B1, ∂B2, ∂B3,), ∂A, basis, BB)
+
+
+
 function pullback_evaluate!(∂BB, # output 
                             ∂A, basis::PooledSparseProduct{NB}, BB::TupMat # inputs 
                             ) where {NB}

src/sparseproduct.jl

@@ -40,7 +40,7 @@ _out_size(basis::SparseProduct, BB::Tuple{AbstractMatrix, AbstractMatrix}) = (si
 
 # ----------------------- evaluation kernels 
 
-function whatalloc(::typeof(evaluate!), basis::SparseProduct, BB) where {T}
+function whatalloc(::typeof(evaluate!), basis::SparseProduct, BB)
    VT = _valtype(basis, BB)
    return (VT, _out_size(basis, BB)...)
 end

src/testing.jl

@@ -2,7 +2,10 @@ module Testing
 
 using Polynomials4ML: evaluate!, evaluate_ed!, evaluate_ed2!, 
                evaluate, evaluate_d, evaluate_ed, evaluate_dd, evaluate_ed2, 
-               _generate_input, AbstractP4MLBasis
+               pullback_evaluate!,
+               AbstractP4MLBasis, AbstractP4MLTensor, AbstractP4MLLayer
+
+import Polynomials4ML: _generate_input, _generate_batch
 
 using Test, ForwardDiff, Bumper, WithAlloc
 
@@ -116,7 +119,8 @@ end
 
 # ------------------------ Test allocations using the WithAlloc interface
 
-function _allocations_inner(basis, x; ed = true, ed2 = true)
+function _allocations_inner(basis::AbstractP4MLBasis, x; 
+                            ed = true, ed2 = true)
    @no_escape begin 
       P = @withalloc evaluate!(basis, x)
       s = sum(P)
@@ -133,40 +137,113 @@ function _allocations_inner(basis, x; ed = true, ed2 = true)
    return s 
 end
 
+_generate_batch(basis::AbstractP4MLBasis; nbatch = rand(7:16)) = 
+         [ _generate_input(basis) for _ = 1:nbatch ]
+
 function test_withalloc(basis::AbstractP4MLBasis; 
             generate_x = () -> _generate_input(basis),
+            generate_batch = () -> _generate_batch(basis),
             allowed_allocs = 0, 
-            nbatch = 16, 
+            ed = true, ed2 = true, 
             kwargs...) 
-   X = [ generate_x() for _ = 1:nbatch ]
+   X = generate_batch()
    x = generate_x()
-   match_all = true 
    for Y in (x, X, )
-      nalloc1 = _allocations_inner(basis, Y; kwargs...)
-      nalloc = @allocated ( _allocations_inner(basis, Y; kwargs...) )
+      nalloc1 = _allocations_inner(basis, Y; ed=ed, ed2=ed2)
+      nalloc = @allocated ( _allocations_inner(basis, Y; ed=ed, ed2=ed2) )
+      println("nalloc = $nalloc (allowed = $allowed_allocs)")
       print_tf(@test nalloc <= allowed_allocs)
-      P1, dP1, ddP1 = evaluate_ed2(basis, Y)
+      if ed2 
+         P1, dP1, ddP1 = evaluate_ed2(basis, Y)
+      end
       @no_escape begin 
          P2 = @withalloc evaluate!(basis, Y)
          P3, dP3 = @withalloc evaluate_ed!(basis, Y)
-         P4, dP4, ddP4 = @withalloc evaluate_ed2!(basis, Y)
-         match_P1P2 = P1 ≈ P2 ≈ P3 ≈ P4
-         match_dP1dP2 = dP1 ≈ dP3 ≈ dP4
-         match_ddP1ddP2 = ddP1 ≈ ddP4
+         if ed2 
+            P4, dP4, ddP4 = @withalloc evaluate_ed2!(basis, Y)
+         end
+         match_P1P2 = P1 ≈ P2 ≈ P3
+         match_dP1dP2 = dP1 ≈ dP3
+         match_all = match_P1P2 & match_dP1dP2
          print_tf(@test match_P1P2)
          print_tf(@test match_dP1dP2)
-         print_tf(@test match_ddP1ddP2)
-      end
-      if nalloc > allowed_allocs 
-         println("nalloc = $nalloc > $allowed_allocs (allowed)")
+         if ed2
+            match_P1P2 = match_P1P2 & (P1 ≈ P4)
+            match_dP1dP2 = match_dP1dP2 & (dP3 ≈ dP4)
+            match_ddP1ddP2 = ddP1 ≈ ddP4
+            math_all = match_P1P2 & match_dP1dP2 & match_ddP1ddP2
+            print_tf(@test match_ddP1ddP2)
+         end
       end
-      if !match_P1P2 || !match_dP1dP2 || !match_ddP1ddP2
+      if !math_all 
          println("standard withalloc evaluations don't match")
       end
    end
    return nothing 
 end
 
+function _allocations_inner(basis::AbstractP4MLTensor, x; 
+                              pb = true)
+   @no_escape begin 
+      P = @withalloc evaluate!(basis, x)
+      s = sum(P)
+      if pb 
+         T = eltype(P) 
+         sz = size(P)
+         ∂P = @alloc(T, sz...)
+         fill!(∂P, zero(T))
+         ∂x = @withalloc pullback_evaluate!(∂P, basis, x)
+      end 
+      nothing 
+   end
+   return s 
+end
+
+
+function test_withalloc(basis::AbstractP4MLTensor; 
+            generate_single = () -> _generate_input(basis),
+            generate_batch = () -> _generate_batch(basis),
+            allowed_allocs = 0, 
+            pb = true, 
+            batch = true, 
+            single = true, 
+            kwargs...) 
+
+   if single 
+      X = generate_single()
+      nalloc_pre = _allocations_inner(basis, X; pb=pb)
+      nalloc_pre = _allocations_inner(basis, X; pb=pb)
+      nalloc = @allocated ( _allocations_inner(basis, X; pb=pb) )
+      println("single: nalloc = $nalloc (allowed = $allowed_allocs)")
+      println_slim(@test nalloc <= allowed_allocs)
+      A1 = evaluate(basis, X)
+      @no_escape begin 
+         A2 = @withalloc evaluate!(basis, X)
+         match_A1A2 = A1 ≈ A2
+      end
+      if !match_A1A2
+         println("single: standard withalloc evaluations don't match")
+      end      
+   end 
+
+   if batch 
+      X = generate_batch()
+      nalloc_pre = _allocations_inner(basis, X; pb=pb)
+      nalloc = @allocated ( _allocations_inner(basis, X; pb=pb) )
+      println("batch: nalloc = $nalloc (allowed = $allowed_allocs)")
+      println_slim(@test nalloc <= allowed_allocs)
+      A1 = evaluate(basis, X)
+      @no_escape begin 
+         A2 = @withalloc evaluate!(basis, X)
+         match_A1A2 = A1 ≈ A2
+      end
+      if !match_A1A2
+         println("batch: standard withalloc evaluations don't match")
+      end
+   end 
+
+   return nothing 
+end
 
 
 

test/ace/test_sparseprodpool.jl

@@ -1,6 +1,7 @@
 
 using BenchmarkTools, Test, Polynomials4ML, ChainRulesCore
-using Polynomials4ML: PooledSparseProduct, evaluate, evaluate!
+using Polynomials4ML: PooledSparseProduct, evaluate, evaluate!, 
+         _generate_input, _generate_input_1
 using Polynomials4ML.Testing: test_withalloc
 using ACEbase.Testing: fdtest, println_slim, print_tf 
 
@@ -23,16 +24,6 @@ function _generate_basis(; order=3, len = 50)
 end
 
 
-function _rand_input1(basis::PooledSparseProduct{ORDER}) where {ORDER} 
-   NN = [ maximum(b[i] for b in basis.spec) for i = 1:ORDER ]
-   BB = ntuple(i -> randn(NN[i]), ORDER)
-end
-
-function _rand_input(basis::PooledSparseProduct{ORDER}; nX = rand(5:15)) where {ORDER} 
-   NN = [ maximum(b[i] for b in basis.spec) for i = 1:ORDER ]
-   BB = ntuple(i -> randn(nX, NN[i]), ORDER)
-end
-
 ##
 
 @info("Test evaluation with a single input (no pooling)")
@@ -42,7 +33,7 @@ for ntest = 1:30
 
    order = mod1(ntest, 4)
    basis = _generate_basis(; order=order)
-   BB = _rand_input1(basis)
+   BB = _generate_input_1(basis)
    A1 = test_evaluate(basis, BB)
    A2 = evaluate(basis, BB)
    print_tf(@test A1 ≈ A2)
@@ -59,12 +50,11 @@ for ntest = 1:30
 
    order = mod1(ntest, 4)
    basis = _generate_basis(; order=order)
-   bBB = _rand_input(basis)
+   bBB = _generate_input(basis)
    bA1 = test_evaluate(basis, bBB)
    bA2 = evaluate(basis, bBB)
    bA3 = copy(bA2)
    evaluate!(bA3, basis, bBB)
-
    print_tf(@test bA1 ≈ bA2 ≈ bA3)
 end
 
@@ -73,13 +63,12 @@ println()
 
 ##
 
-using StrideArrays
-
 @info("    testing withalloc")
 basis = _generate_basis(; order=2)
-BB = PtrArray.(_rand_input1(basis))
-bBB = _rand_input(basis)
-# test_withalloc(basis; ed = false, ed2 = false)
+BB = _generate_input_1(basis)
+bBB = _generate_input(basis)
+test_withalloc(basis; batch=false)
+
 
 ##
 
@@ -91,9 +80,9 @@ for ntest = 1:30
    local bBB, bA2, u, basis, nX 
    order = mod1(ntest, 4)
    basis = _generate_basis(; order=order)
-   bBB = _rand_input(basis)
+   bBB = _generate_input(basis)
    nX = size(bBB[1], 1)
-   bUU = _rand_input(basis; nX = nX) # same shape and type as bBB 
+   bUU = _generate_input(basis; nX = nX) # same shape and type as bBB 
    _BB(t) = ntuple(i -> bBB[i] + t * bUU[i], order)
    bA2 = evaluate(basis, bBB)
    u = randn(size(bA2))
@@ -210,3 +199,4 @@ end
 # A, ∂A = P4ML.pfwd_evaluate(basis, BB, ΔBB)
 # @btime Polynomials4ML.pfwd_evaluate!($(unwrap(A)), $(unwrap(∂A)), $basis, $BB, $ΔBB)
 
+

test/test_sphericart.jl

@@ -130,13 +130,15 @@ println()
 
 ##
 
+@info("Test Racah normalization")
 racah = real_solidharmonics(4; normalisation = :racah)
 for ntest = 1:20 
    𝐫 = @SVector randn(3) 
    Z1 = racah(𝐫)
    Z2 = explicit_rsh(𝐫)
    print_tf(@test Z1 ≈ Z2) 
 end
+println() 
 
 ##