FabianFuchsML / se3-transformer-public

code for the SE3 Transformers paper: https://arxiv.org/abs/2006.10503
475 stars 69 forks source link

non scalar output #10

Open Chen-Cai-OSU opened 3 years ago

Chen-Cai-OSU commented 3 years ago

Hi Fabian,

I was wondering is it easy for the current se3-transformer to support the output whose irreps type is not scalar? The property I want to predict is a 3x3 symmetric PSD matrix. In addition, I have some scalar features on each edge. Is it easy to achieve this by the current se3-transformer?

I asked you before but at that time there are some nasty change-of-basis matrices that I am not able to figure out. Just want to check with you whether this is now supported as the code has changed a lot since the last time I asked.

FabianFuchsML commented 3 years ago

Hey! So I did add a re-implementation of the nbody experiment since the last time we talked about this. This has a type-1 feature output, and can easily be augmented with type 2 features etc. So in theory it should be possible to construct a matrix output with that, but I haven't tried so myself.