torchmd-net

torchmd-net copied to clipboard

Allowing vector output from TensorNet. This will allow the use of EquivariantScalar or EquivariantVectorOutput from TensorNet.

Resolves #297.

Thanks a lot! Can you add the equivariance test, please? I understood what you meant in the issue (it just confused me to see them denotes as 'forces' and not specifying further :) )

I should be able to add it today. Sorry for the delay.

Don’t rush! Any day will be fine.

Thank you!

On Sun, 3 Mar 2024 at 18:51, Nikhil Shenoy @.***> wrote:

I should be able to add it today. Sorry for the delay.

— Reply to this email directly, view it on GitHub https://github.com/torchmd/torchmd-net/pull/301#issuecomment-1975243767, or unsubscribe https://github.com/notifications/unsubscribe-auth/ANJMOAYKVMROBKLKKWDTYNTYWNPLTAVCNFSM6AAAAABECTM5SWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNZVGI2DGNZWG4 . You are receiving this because you commented.Message ID: @.***>

Can you add the equivariance test, please?

Added an equivariance test.

Hey! Sorry for the delay. Revisiting this:

Your approach seems correct @shenoynikhil, however I think I would not proceed by adding a new model that is called equivariant-tensornet. This makes me think, @RaulPPelaez, that we should get rid of the is_equivariant flag, and directly modify the places where the ET is used (yaml example files, for example), and directly include there EquivariantScalar, instead of Scalar, as output model, where it is used. What do you think? I don't know how many things this change would impact, but in any case I think it is worth doing it, the flag does not make much sense at this point...

For me, I would say it is fine if TensorNet defaults to returning x, v instead of x, None. Even though TensorNet uses Scalar (not EquivariantScalar), for energy predictions, I have the impression that nothing would happen. As I said above at some point, v comes from A, which has already been used by learnable parameters to generate x. Using Scalar, no new learnable things would act independently on v, and therefore we should not get an error pointing to not using parameters that are trainable in the forward.

What do you think guys?

I would not proceed by adding a new model that is called equivariant-tensornet

I agree because it kind of implies that there is a non-equivariant tensornet. I would be more comfortable with tensornet setting is_equivariant (and thus selecting the Equivariant* output models) if vector_output is set to true. This option could come from train.py as "tensornet-vector-output" maybe. It is not standard for us to refer to different flavors of a model with different names, but I am not totally against it personally. In this particular case I do not think it is a good idea, imagine I have "vector-tensornet" on one hand and "tensornet-whatever" in another and I would like "vector-tensornet-whatever", there would be no way to get that.

For me, I would say it is fine if TensorNet defaults to returning x, v instead of x, None. Even though TensorNet uses Scalar (not EquivariantScalar), for energy predictions, I have the impression that nothing would happen

I do not like this because even if autograd eats it it is extra unnecessary work. I believe gating the functionality behind an user-provided function or a model name is best in this case.

we should get rid of the is_equivariant flag, and directly modify the places where the ET is used (yaml example files, for example), and directly include there EquivariantScalar, instead of Scalar, as output model, where it is used.

I agree with this, but I believe it is out of the scope of this PR. It would be great if you open an issue so we have it in our list... For retrocompatibility create_model can simply compare "model_name" with "equivariant-transformer" to decide whether to transform Scalar into EquivariantScalar.