Merge pull request #316 from ACEsuit/co/etpair

ET Pair Potential Model

Author: cortner

src/et_models/convert.jl

@@ -6,7 +6,7 @@ import EquivariantTensors as ET
 import Polynomials4ML as P4ML
 
 import ACEpotentials.Models: LearnableRnlrzzBasis, PolyEnvelope2sX, 
-         _i2z, GeneralizedAgnesiTransform 
+         _i2z, GeneralizedAgnesiTransform, PolyEnvelope1sR
 
 using LinearAlgebra: norm, dot 
 
@@ -68,6 +68,7 @@ function convert2et(model)
 end
 
 
+
 # In ET we currently store an edge xij as a NamedTuple, e.g, 
 #    xij = (𝐫ij = ..., zi = ..., zj = ...)
 # The NTtransform is a wrapper for mapping xij -> y 
@@ -85,15 +86,6 @@ function _convert_Rnl_learnable(basis; zlist = ChemicalSpecies.(basis._i2z),
    # number of species 
    NZ = length(zlist)
 
-   # species z -> index i mapping 
-   __z2i = let _i2z = (_i2z = zlist,) 
-      z -> _z2i(_i2z, z)
-   end
-
-   # __zz2i maps a `(Zi, Zj)` pair to a single index `a` representing 
-   # (Zi, Zj) in a flattened array
-   __zz2ii = (zi, zj) -> (__z2i(zi) - 1) * NZ + __z2i(zj)
-
    selector = let zlist = tuple(zlist...)
       xij -> ET.catcat2idx(zlist, xij.z0, xij.z1)
    end
@@ -124,8 +116,8 @@ function _convert_Rnl_learnable(basis; zlist = ChemicalSpecies.(basis._i2z),
       @assert env.x1 == -1
       @assert env.x2 == 1
    end
-
    et_env = y -> (1 - y^2)^2
+   # et_env = _convert_envelope(basis.envelopes)
 
    # the polynomial basis just stays the same 
    # but needs to be wrapped due to the envelope being applied 
@@ -205,4 +197,108 @@ function _convert_agnesi(rbasis::LearnableRnlrzzBasis)
    end
 
    return ET.NTtransformST(f_agnesi, st)
-end 
+end 
+
+
+function _convert_envelope(envelopes) 
+   TENV = typeof(envelopes[1]) 
+   for env in envelopes
+      @assert typeof(env) == TENV
+   end 
+
+   @show TENV 
+   return _convert_env_TENV(TENV, envelopes)
+end 
+
+function _convert_env_TENV(::Type{<: PolyEnvelope2sX}, envelopes) 
+   for env in envelopes
+      @assert env isa PolyEnvelope2sX
+      @assert env.p1 == env.p2 == 2 
+      @assert env.x1 == -1
+      @assert env.x2 == 1
+   end
+   return y -> (1 - y^2)^2
+end 
+
+function _convert_env_TENV(::Type{<: PolyEnvelope1sR}, envelopes) 
+   env1 = envelopes[1]
+   for env in envelopes
+      @assert env == env1 
+   end
+   f_env = (r, st) -> _eval_env_1sr(r, st.rcut, st.p)
+   refst = ( rcut = env1.rcut, p = env1.p )
+   return ET.st_transform(f_env, refst)
+end 
+
+function _eval_env_1sr(r, rcut, p)
+   _1 = one(r)
+   s = r / rcut 
+   return (s^(-p) - _1) * (_1 - s) * (s < _1)
+end
+
+function _convert_pair_envelope(envelopes) 
+   TENV = typeof(envelopes[1]) 
+   for env in envelopes
+      @assert typeof(env) == TENV
+   end 
+   env1 = envelopes[1]
+   @assert env1 isa PolyEnvelope1sR 
+   for env in envelopes
+      @assert env == env1 
+   end
+   refst = ( rcut = env1.rcut, p = env1.p ) 
+   f_env = ET.dp_transform( (x, st) -> _eval_env_1sr( norm(x.𝐫), st.rcut, st.p ), 
+                            refst )
+   return f_env                             
+end
+
+
+
+function convertpair(model)
+
+   # extract radial basis information 
+   basis = model.pairbasis 
+   zlist = ChemicalSpecies.(basis._i2z)
+   NZ = length(zlist)
+
+   # this construction is a little different from the Rnl basis for the 
+   # many-body model because the envelope takes a different input 
+   # and this makes life a little more complicated. 
+
+   # 1: polynomials without the envelope
+   # 
+   dp_agnesi = _convert_agnesi(basis)
+   polys = basis.polys
+   selector2 = let zlist = zlist
+      xij -> ET.catcat2idx(zlist, xij.z0, xij.z1)
+   end
+   et_linl = ET.SelectLinL(length(polys),         # indim
+                           length(basis),       # outdim
+                           NZ^2,                     # num (Zi,Zj) pairs
+                           selector2)
+   rbasis_1 = ET.EmbedDP(dp_agnesi, polys, et_linl)
+
+   # 2: envelope 
+   dp_envelope = _convert_pair_envelope(basis.envelopes)
+   # _env_r = _convert_envelope(basis.envelopes)
+   # dp_envelope = ET.dp_transform( (x, st) -> _env_r.f( norm(x.𝐫), st ), 
+   #                                 _env_r.refstate )
+
+   # 3. combine into the radial basis
+   rembed = ET.EdgeEmbed( EnvRBranchL(dp_envelope, rbasis_1) )
+
+   # 4. rembed provides the radial basis for the pair model, now we just 
+   #    need the readout layer which is similar to before.
+   selector1 = let zlist = zlist
+      x -> ET.cat2idx(zlist, x.z)
+   end
+   readout = ET.SelectLinL(
+                  length(basis), 
+                  1,                  # output dim (only one site energy per atom)
+                  NZ,     # number of categories = num species 
+                  selector1)
+
+   et_pair = ETPairModel(rembed, readout)
+
+   return et_pair
+end 

src/et_models/et_envbranch.jl

@@ -0,0 +1,54 @@
+
+
+using ConcreteStructs
+import Polynomials4ML: evaluate, evaluate_ed 
+import LuxCore: AbstractLuxContainerLayer
+import ChainRulesCore: NoTangent, rrule, unthunk
+
+"""
+   struct EnvRBranchL
+
+An auxiliary layer that is basically a branch layer needed to build  
+radial bases, with additional evaluate_ed functionality, needed for 
+Jacobians. 
+"""
+@concrete struct EnvRBranchL <: AbstractLuxContainerLayer{(:envelope, :rbasis)}
+   envelope 
+   rbasis 
+end 
+
+(l::EnvRBranchL)(X, ps, st) = _apply_envrbranchl(l, X, ps, st), st
+
+evaluate(l::EnvRBranchL, X, ps, st) = l(X, ps, st)
+
+function _apply_envrbranchl(l::EnvRBranchL, X, ps, st)
+   ee, _ = l.envelope(X, ps.envelope, st.envelope)
+   P, _ = l.rbasis(X, ps.rbasis, st.rbasis)
+   return ee .* P
+end
+
+function evaluate_ed(l::EnvRBranchL, X, ps, st)
+   (ee, d_ee), _ = evaluate_ed(l.envelope, X, ps.envelope, st.envelope)
+   (P, d_P), _ = evaluate_ed(l.rbasis,X, ps.rbasis, st.rbasis)
+
+   # product rule 
+   pP = ee .* P
+   ∂_pP = d_ee .* P .+ ee .* d_P
+
+   return (pP, ∂_pP), st
+end
+
+function rrule(::typeof(_apply_envrbranchl), 
+               l::EnvRBranchL, X, ps, st)
+
+   (P, dP), st = evaluate_ed(l, X, ps, st)
+
+   function _pb_embeddp(_∂P)
+      ∂P = unthunk(_∂P)
+      ∂X = dropdims( sum(∂P .* dP, dims = 2), dims = 2) 
+      return NoTangent(), NoTangent(), ∂X, NoTangent(), NoTangent()
+   end
+
+   return P, _pb_embeddp
+end
+

src/et_models/et_models.jl

@@ -1,9 +1,17 @@
 
 module ETModels 
 
+# utility layers : these should likely be moved into ET or be removed 
+#     if more convenient implementations can be found. 
+#  
+include("et_envbranch.jl")
+
+# ET based ACE model components 
 include("et_ace.jl")
 include("onebody.jl")
+include("et_pair.jl")
 
+# converstion utilities: convert from 0.8 style ACE models to ET based models
 include("convert.jl")
 
 

src/et_models/et_pair.jl

@@ -0,0 +1,66 @@
+#
+# This is a temporary model implementation needed due to the fact that 
+# ETACEModel has Rnl, Ylm hard-coded. In the future it could be tested 
+# whether the pair model could simply be taken as another ACE model 
+# with a single embedding rather than several, This would need generalization 
+# of a fair few methods in both ACEpotentials and EquivariantTensors.
+#
+
+
+import EquivariantTensors as ET 
+import Zygote 
+import LuxCore: AbstractLuxContainerLayer
+using ConcreteStructs: @concrete
+
+
+@concrete struct ETPairModel  <: AbstractLuxContainerLayer{(:rembed, :readout)}
+   rembed     # radial embedding layer = basis 
+   readout    # normally a selectlinl readout layer
+end 
+
+
+(l::ETPairModel)(X::ET.ETGraph, ps, st) = _apply_etpairmodel(l, X, ps, st), st 
+      
+      
+function _apply_etpairmodel(l::ETPairModel, X::ET.ETGraph, ps, st)      
+   # evaluate the basis 
+   𝔹 = site_basis(l, X, ps, st)
+
+   # readout layer 
+   φ, _ = l.readout((𝔹, X.node_data), ps.readout, st.readout)
+
+   return φ
+end
+
+# ----------------------------------------------------------- 
+
+
+function site_grads(l::ETPairModel, X::ET.ETGraph, ps, st)
+   ∂X = Zygote.gradient( X -> sum(_apply_etpairmodel(l, X, ps, st)), X)[1]
+   return ∂X
+end
+
+
+# ----------------------------------------------------------- 
+#    basis and jacobian evaluation 
+
+
+function site_basis(l::ETPairModel, X::ET.ETGraph, ps, st)      
+   # embed edges 
+   Rnl, _ = l.rembed(X, ps.rembed, st.rembed)
+
+   # the basis is obtain by summing over the neighbours of each node, 
+   # which is just a sum over the first dimension of Rnl 
+   𝔹 = dropdims(sum(Rnl, dims=1), dims=1)
+
+   return 𝔹
+end
+
+
+function site_basis_jacobian(l::ETPairModel, X::ET.ETGraph, ps, st)    
+   (R, ∂R), _ = ET.evaluate_ed(l.rembed, X, ps.rembed, st.rembed)
+   𝔹 = dropdims(sum(R, dims=1), dims=1)
+   # ∂𝔹 == ∂R
+   return 𝔹, ∂R
+end
+

test/etmodels/test_etbackend.jl

@@ -1,6 +1,6 @@
-# using Pkg; Pkg.activate(joinpath(@__DIR__(), ".."))
+# using Pkg; Pkg.activate(joinpath(@__DIR__(), "..", ".."))
 # using TestEnv; TestEnv.activate();
-# Pkg.develop(url = joinpath(@__DIR__(), "..", "..", "EquivariantTensors.jl"))
+# Pkg.develop(url = joinpath(@__DIR__(), "..", "..", "..", "EquivariantTensors.jl"))
 # Pkg.develop(url = joinpath(@__DIR__(), "..", "..", "Polynomials4ML.jl"))
 # Pkg.develop(url = joinpath(@__DIR__(), "..", "..", "DecoratedParticles"))
 
@@ -190,7 +190,7 @@ WW = et_ps_2.readout.W
 println_slim(@test 𝔹1 ≈ 𝔹2)
 Ei_a = [ dot(𝔹2[i, :], WW[1, :, iZ[i]])    for (i, iz) in enumerate(iZ) ]
 Ei_b = et_model_2(G, et_ps_2, et_st_2)[1][:]
-println(@test Ei_a ≈ Ei_b)
+println_slim(@test Ei_a ≈ Ei_b)
 
 ##
 
@@ -240,8 +240,6 @@ println_slim( @test abs(E1 - E4) / (abs(E1) + abs(E4) + 1e-7) < 1e-5 )
 
 ## 
 # gradients on GPU 
-# currently failing because somehow the transform is still 
-# accessing some Float64 values somewhere .... 
 
 @info("Check Evaluation of gradient on GPU")
 g1 = ETM.site_grads(et_model_2, G_32, ps_32_2, st_32_2)
@@ -268,6 +266,8 @@ println_slim( @test 𝔹1 ≈ 𝔹2 )
 
 println_slim( @test 𝔹1 ≈ 𝔹2 )
 err_jac = norm.(∂𝔹1 - ∂𝔹2) ./ (norm.(∂𝔹1) + norm.(∂𝔹2) .+ 0.1) 
-println_slim( @test maximum(err_jac) < 1e-5 )
+println_slim( @test maximum(err_jac) < 1e-4 )
+@show maximum(err_jac)
+@info("The jacobian error feels a bit large. This may need further investigation.")
 
-=# 
+=#

test/etmodels/test_etonebody.jl

@@ -139,6 +139,7 @@ G_dev = dev(G_32)
 E1, st = et_V0(G_32, ps_32, st_32)
 E2_dev, st_dev = et_V0(G_dev, ps_dev, st_dev)
 E2 = Array(E2_dev)
+# TODO: add E1 ≈ E2 test??
 
 g1 = ETM.site_grads(et_V0, G_32, ps_32, st_32)
 g2_dev = ETM.site_grads(et_V0, G_dev, ps_dev, st_dev)