Product basis

[CheukHinHoJerry]

This branch implements the ProdBasis without pooling for building the atomic orbital bases.

[cortner]

Hi Jerry - Thank thanks for this. I am reading it and it looks great so far. One question that came to mind which we should discuss: this seems to me to be a SparseProduct not a ProductBasis. I think a ProductBasis would be an object that takes x as input then produces several bases B_n and then combines them into a product B_n1n2....

Specifically we could introduce another object

struct ProductBasis{TBB, TP}
   BB::TBB     # e.g. BB = (B1, B2, B3)::Tuple{Exponentials, Chebyshev, SolidHarmonics}
   prod::TP    # e.g.  prod = SparseProduct(spec)
end

which manually implements what we discussed we want to do via Lux - so maybe we shouldn't do this and immediately focus on producing Lux layers?

What do you think?

[cortner]

To use this for Dexuan to work on the sampler we do need the pooling. The gradients could come in the next PR.

[CheukHinHoJerry]

Re the naming problem - Yes, I agree. Sorry for misnaming that. I will fix that soon.

Re Lux layers: Yes, I think we can focus on Lux layer. I think it provides more flexibility if we want to use a learnable radial embedding following by the Tucker Tensor format.

[CheukHinHoJerry]

To use this for Dexuan to work on the sampler we do need the pooling. The gradients could come in the next PR.

Maybe I can do the pooling in ACEpsi? If I remeber correctly you have finished the most technical part and what I need to do is mostly just copying.

[cortner]

so this is ready to merge from your end?

[CheukHinHoJerry]

Sorry for late reply. I just did a small clean up but there is one small issue:

For N = 1 the _prod_grad function is not working correctly so I guess this is why this function exist too in ACEcore.jl

@inline function _prod_grad(b::SVector{1, T}) where {T} 
   return b[1], SVector(one(T))
end

But it has inconsistent return with the general case. Could you please suggest how to fix this?

[cortner]

Interesting- can you please post a MWE sonI can reproduce and test?

[CheukHinHoJerry]

This is also in /test/test_sparseproduct.jl. If you change the for N = 2:5 in the second test to for N = 1:5 it fails.

using StaticArrays, ForwardDiff

prodgrad = Polynomials4ML._prod_grad

# special case when N = 1
for ntest = 1:10
   local v1, g
   b = rand(SVector{1, Float64})
   g = prodgrad(b)
   g1 = ForwardDiff.gradient(prod, b)
   print_tf(@test g1 ≈ SVector(g...)[2])
end

for N = 2:5 
   for ntest = 1:10
      local v1, g 
      b = rand(SVector{N, Float64})
      g = prodgrad(b.data, Val(N))
      g1 = ForwardDiff.gradient(prod, b)
      print_tf(@test g1 ≈ SVector(g...))
   end
end

[cortner]

I pushed a fix.

[cortner]

looks like the tests are passing. Ready to merge? Then you can get going on the next step :)

[CheukHinHoJerry]

looks like the tests are passing. Ready to merge? Then you can get going on the next step :)

Yes please - I will start working on BF pooling later today. I guess I should do that in ACEpsi.jl?

[cortner]

Maybe I can do the pooling in ACEpsi? If I remeber correctly you have finished the most technical part and what I need to do is mostly just copying.

Yes, I think doing this well will be the next step.