Chemellia / AtomicGraphNets.jl

Atomic graph models for molecules and crystals in Julia
MIT License
62 stars 10 forks source link

smart way to share layers with GeometricFlux? #66

Open rkurchin opened 3 years ago

rkurchin commented 3 years ago

The packages are super similar; I was originally using GeometricFlux.jl as a core dependency for this package but decided to build some of my own types in ChemistryFeaturization to specialize on atomic structures a bit more. I wonder if there's nonetheless a smart way to share layer definitions, when applicable, between the packages? For instance, some layers are described here and @yuehhua has done awesome work implementing them!

Since it seems all those layers can act directly on matrices as well as on FeaturedGraph objects, it seems likely we could just import them and dispatch onto AtomGraph as well...

yuehhua commented 3 years ago

An explicit way is to implement the shared API in GeometricFlux, then any package can leverage these API in their packages. If you have some idea about this, I can implement API for sharing.

rkurchin commented 3 years ago

I think all that should be needed would be an additional dispatch within AtomicGraphNets that tells it how to handle our graph types within your layers. This is definitely something I want to look into, but I have a couple more urgent things on my plate right now so I may not get to it for a bit...

yuehhua commented 3 years ago

Yeah, I just want to collect some requirements from user to enhance GeometricFlux, and simultaneously help out others. So, currently, you may want a way or tool to handle graph types in layers?

rkurchin commented 3 years ago

I mean the specific AtomGraph or FeaturizedAtoms types that ChemistryFeaturization defines. It should actually be really simple, just telling the layer which fields to get the actual graph and feature information from, and how to put them into your layers. Glancing over this code, I think one thing that would make that easier would be if all the layers accepted a common signature so that the dispatches could be done the same way. It seems a lot of them accept FeaturedGraph (though I can't seem to find where that's defined, so I'm having trouble figuring out if there would be an easy way to just interconvert between my types and that), some accept (L, X) (this would be easy as AtomGraph explicitly stores the laplacian to avoid having to compute it every time), etc. Even cooler might be if there were an abstract type that all the convolutional layers inherited from and they used a common signature so that I could do all of this in basically one line! :)

yuehhua commented 3 years ago

The easiest way to support AtomGraph and FeaturizedAtoms type is to convert them into a FeaturedGraph. The definition of FeaturedGraph is in GraphSignals.jl. I make all of the graph convolutional layers support FeaturedGraph and it has functionality of a graph, e.g. nv for getting number of vertex or ne for getting number of edges. Maybe I can check the type definition of AtomGraph and FeaturizedAtoms and give you some suggestions.

rkurchin commented 3 years ago

Ah okay I'll take a look at that. I might play around with both options since I imagine converting to the FeaturedGraph type could have a performance penalty relative to just defining the dispatch, so probably I'll make that conversion as the default thing and then do specific dispatches if/as needed. Thanks!

yuehhua commented 3 years ago

Actually, converting into FeaturedGraph doesn't make a performance drop. As I am designing FeaturedGraph, I try to keep converting between graph representations as less as possible. I end up just keep the original graph representation as it is. The FeaturedGraph is just a wrapping around graph representation and features. The graph is converted at the moment of graph convolutional layers need it. The graph convolutional layer calls API likes adjacency_list or adjacency_matrix to convert original graph into a new representation, which is needed in their own computation. In such idea, the graph is just converted once before layer computation.