AtomicOrbitalBasis and BackflowPooling

[CheukHinHoJerry]

This branch implements the construction of the A basis for bflow wf.
@cortner It would be great if you can confirm this is what we need?

This code is not yet fully optimized and there is a lot to be done to make it more performant.

Also, currently the ProductBasis is a specific case and we should just implements a general version in Polynomials4ML. I will do that after confirming this is what we want.

CC @DexuanZhou

[cortner]

Thank you! This does look correct but there are so many changes it is hard to read. Can you maybe write a brief summary of how you re-organized the code and I'll judge from that?

[CheukHinHoJerry]

@cortner Sure. The original AtomicOrbitaBasis output the A basis directly. @DexuanZhou and I have made the following changes :

Implements the ProductBasis with SparseProduct:
The product basis is a mapping which maps x to the product of Rnl and Ylm with a given spec1.

Modify the AtomicOrbitalBasis as a wrapper of the ProductBasis:
It takes in x and apply different translation depends on nuclei defined during construction. And then it uses ProductBasis to evaluate the product on different translated x. The output of AtomicOrbitalBasis is a tensor of $\phi_{nlm}$ of different translations.

Move the BackFlowPooling out of
AtomicOrbitalRadius:
The pooling operation is moved out as a separate structure from the original AtomicOrbitalBasis. Which takes the tensor of $\phi_{nlm}$ from AtomicOrbitalBasis to the pooled A basis.

After implementing the above, I assmeble all the things by calling ProductBasis inside AtomicOrbitalBasis, and then following by a separate Pooling layer to construct the A basis.

[cortner]

I'll try to review this asap

[CheukHinHoJerry]

Maybe I can clean this up once first and give a summary before your review? I will do this within next 12 hours.

[cortner]

good, because I can't do it today or tomorrow :)

[CheukHinHoJerry]

This is moved to

#20 (comment)

I think we can delete this branch and also the PR.

CC @cortner @DexuanZhou

[CheukHinHoJerry]

Close and reference to #22