Merge pull request #111 from ACEsuit/fix_hyperdual

Add _valtype support for RadialDecay with HyperDual input

Author: cortner

Project.toml

@@ -12,6 +12,7 @@ Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+HyperDualNumbers = "50ceba7f-c3ee-5a84-a6e8-3ad40456ec97"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
@@ -25,6 +26,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 WithAlloc = "fb1aa66a-603c-4c1d-9bc4-66947c7b08dd"
 
+
 [compat]
 ACEbase = "0.4.5"
 BenchmarkTools = "1"
@@ -34,6 +36,7 @@ Combinatorics = "1"
 DiffResults = "1.1.0"
 ForwardDiff = "0.10, 1.0"
 GPUArraysCore = "0.2.0"
+HyperDualNumbers = "4.0.10"
 KernelAbstractions = "0.9.34"
 LinearAlgebra = "1.10"
 LuxCore = "1.2.0"
@@ -49,10 +52,11 @@ WithAlloc = "0.1.0"
 julia = "1.10"
 
 [extras]
+HyperDualNumbers = "50ceba7f-c3ee-5a84-a6e8-3ad40456ec97"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "Printf", "Optimisers", "Zygote"]
+test = ["Test", "Printf", "Optimisers", "Zygote", "HyperDualNumbers"]

src/atomicorbitals/radialdecay.jl

@@ -32,14 +32,12 @@ end
 
 Base.length(basis::RadialDecay) = size(basis.ζ, 1)
 
-_valtype(::RadialDecay, T::Type{<: Real}) = T
-
-_valtype(::RadialDecay, T::Type{<: Real}, 
-         ps::Union{Nothing, @NamedTuple{}}, st) = T
+_valtype(::RadialDecay, T::Type{<: Number}, args...) = T
 
 _valtype(::RadialDecay, T::Type{<: Real}, 
          ps, st) = promote_type(T, eltype(ps.ζ), eltype(ps.D))
 
+
 _static_params(basis::RadialDecay) = (ζ = basis.ζ, D = basis.D)
 
 _init_luxparams(basis::RadialDecay) = 

src/testing.jl

@@ -11,9 +11,11 @@ import Polynomials4ML: _generate_input, _generate_batch
 using ChainRulesCore: rrule 
 
 using Test, ForwardDiff, Bumper, WithAlloc
-
+using HyperDualNumbers
+using HyperDualNumbers: Hyper
 using StaticArrays 
-using LinearAlgebra: norm, dot 
+using Random, LuxCore
+using LinearAlgebra: norm, dot, normalize
 
 using ACEbase.Testing: print_tf, println_slim, fdtest 
 
@@ -213,4 +215,49 @@ function test_withalloc(basis::AbstractP4MLBasis;
 end
 
 
+# Construct a static vector of HyperDual numbers from
+#   x : point coordinates
+#   v : first direction
+#   w : second direction
+make_hyper(x::SVector{N,T}, v::SVector{N,T}, w::SVector{N,T}) where {N,T} =
+   SVector{N}(Hyper.(x, v, w, zero(T)))
+
+# Extract components from an array of HyperDual numbers
+_val(A) = HyperDualNumbers.value.(A)  # function value
+_e1(A)  = HyperDualNumbers.eps1.(A)   # first directional derivative
+_e2(A)  = HyperDualNumbers.eps2.(A)   # second directional derivative
+
+# Project each gradient vector g in G onto direction v
+# G is a vector of static vectors (gradients for each basis function)
+_dir(G::AbstractVector{<:StaticVector{D,T}}, v::StaticVector{D,T}) where {D,T} =
+   map(g -> dot(g, v), G)
+
+function test_hyperdual_consistency(basis; rng=Random.default_rng(), rtol=1e-10, atol=1e-12)
+   @info("HyperDual test")
+   ps, st = LuxCore.setup(rng, basis)
+
+   x  = _generate_input(basis)
+   v = normalize(rand(rng, SVector{length(x)}))
+   w = normalize(rand(rng, SVector{length(x)}))
+   xh = make_hyper(x, v, w)
+
+   P  = evaluate(basis, x)
+   P_ed, G = evaluate_ed(basis, x)
+
+   Ph = evaluate(basis, xh)
+   @test isapprox(_val(Ph), P; rtol=rtol, atol=atol)
+   @test isapprox(_e1(Ph), _dir(G, v); rtol=rtol, atol=atol)
+   @test isapprox(_e2(Ph), _dir(G, w); rtol=rtol, atol=atol)
+
+   Pp  = evaluate(basis, x, ps, st)
+   Pp_ed, Gp = evaluate_ed(basis, x, ps, st)
+
+   Php = evaluate(basis, xh, ps, st)
+   @test isapprox(_val(Php), Pp; rtol=rtol, atol=atol)
+   @test isapprox(_e1(Php), _dir(Gp, v); rtol=rtol, atol=atol)
+   @test isapprox(_e2(Php), _dir(Gp, w); rtol=rtol, atol=atol)
+
+   @test isapprox(P_ed, P; rtol=rtol, atol=atol)
+   @test isapprox(Pp_ed, Pp; rtol=rtol, atol=atol)
+end
 end

src/utils/utils.jl

@@ -1,7 +1,8 @@
 module Utils 
 
 include("sparse.jl")
-# include("hyper.jl")
+
+include("hyper.jl")
 
 include("linl.jl")
 

test/test_atorbrad.jl

@@ -1,72 +1,34 @@
-using LinearAlgebra, StaticArrays, Test, Printf, Polynomials4ML
-using Polynomials4ML: evaluate, evaluate_ed
-using Polynomials4ML.Testing: print_tf, println_slim 
-using ForwardDiff
-using ACEbase.Testing: fdtest
-
+using Test
+using StaticArrays
+using LinearAlgebra
+using LuxCore
 import Polynomials4ML as P4ML
+using Polynomials4ML: evaluate, evaluate_ed
 
-##
-
-@info("Testing GaussianBasis")
-basis = P4ML._rand_gaussian_basis()
-
-@info("      correctness of evaluation")
-x = P4ML._generate_input(basis)
-P = evaluate(basis, x)
-P1, dP1 = evaluate_ed(basis, x)
-
-P4ML.Testing.test_evaluate_xx(basis)
-P4ML.Testing.test_chainrules(basis)
-
-# Test is broken - reshape is causing this, hence single-input test is turned off
-P4ML.Testing.test_withalloc(basis; allowed_allocs = 0, single=false)
-
-# Still fails with single=true, but btime doesn't show the allocation
-# even stranger, re-evaluating `_reshape` into P4ML also makes the allocation 
-# disappear. Probably best to drop this for now, and revisit in a few months. 
-# P4ML.Testing.test_withalloc(basis; allowed_allocs = 0, single=true)
-
-##
-# these are scripts to replicate and check this allocation problem. 
-# strangely it doesn't occur for the other bases. Only for AtomicOrbtials. 
-#
-# using BenchmarkTools
-
-# P, dP = evaluate_ed(basis, x)
-# @btime P4ML.evaluate_ed!($P, $dP, $basis, $x)
-
-# @profview let basis=basis, X=x, P=P, dP=dP 
-#     for _ = 1:1_000_000 
-#         P4ML.evaluate_ed!(P, dP, basis, X)
-#     end
-# end
-
-##
-
-@info("Testing SlaterBasis")
-basis = P4ML._rand_slater_basis()
-
-@info("      correctness of evaluation")
-x = P4ML._generate_input(basis)
-P = evaluate(basis, x)
-P1, dP1 = evaluate_ed(basis, x)
 
-P4ML.Testing.test_evaluate_xx(basis)
-P4ML.Testing.test_chainrules(basis)
-P4ML.Testing.test_withalloc(basis; allowed_allocs = 0, single=false)
+@testset "GaussianBasis + HyperDual matches evaluate/evaluate_ed" begin
+    basis = P4ML._rand_gaussian_basis()
+    P4ML.Testing.test_hyperdual_consistency(basis)
+    P4ML.Testing.test_evaluate_xx(basis)
+    P4ML.Testing.test_chainrules(basis)
+    P4ML.Testing.test_withalloc(basis; allowed_allocs = 0, single=false)
+end
 
-##
+@testset "SlaterBasis + HyperDual matches evaluate/evaluate_ed" begin
+    basis = P4ML._rand_slater_basis()
+    P4ML.Testing.test_hyperdual_consistency(basis)
+    P4ML.Testing.test_evaluate_xx(basis)
+    P4ML.Testing.test_chainrules(basis)
+    P4ML.Testing.test_withalloc(basis; allowed_allocs = 0, single=false)
+end
 
-@info("Testing STOBasis")
-basis = P4ML._rand_sto_basis()
+@testset "STOBasis + HyperDual matches evaluate/evaluate_ed" begin
+    basis = P4ML._rand_sto_basis()
+    P4ML.Testing.test_hyperdual_consistency(basis)
+    P4ML.Testing.test_evaluate_xx(basis)
+    P4ML.Testing.test_chainrules(basis)
+    P4ML.Testing.test_withalloc(basis; allowed_allocs = 0, single=false)
+end
 
-@info("      correctness of evaluation")
-x = P4ML._generate_input(basis)
-P = evaluate(basis, x)
-P1, dP1 = evaluate_ed(basis, x)
 
 
-P4ML.Testing.test_evaluate_xx(basis)
-P4ML.Testing.test_chainrules(basis)
-P4ML.Testing.test_withalloc(basis; allowed_allocs = 0, single=false)