How to handle higher dimensional coordinate vectors

[tueboesen]

I am trying to incorporate your neural network on a protein system, however unlike your typical datasets where you work on atoms and each atom is represented by a single 3D coordinate and velocity vector, I work on compound residues consisting of multiple atoms.
What this means is that each of my residues is represented by 4 3D coordinate vectors, so a 12D coordinate vector essentially.

Given this, I'm wondering what would be the best strategy when using SphericalHarmonicEdgeAttrs, and RadialBasisEdgeEncoding on such a beast. I could of course just treat it as 4 separate 3D coordinate vectors and then just call each of the functions 4 times with the standard 1x0e+1x1o for the edge_sh, and standard edge_vector_length for the RadialBasisEdgeEncoding. I can then just concatenate the results and get some output which will then be 4x0e+4x1o and this will likely capture what I am looking for.

I am wondering whether there is a better way to do this? I believe there might be, but I don't quite have the experience or grasp on how these two functions work (SphericalHarmonicEdgeAttrs, RadialBasisEdgeEncoding) to know what would be the right way to do this.

[Linux-cpp-lisp]

Hi @tueboesen ,

There is definitely a better way to do this, and it is pretty straightforward, but I want to make sure I'm understanding right what you are saying first: the four vectors are all global coordinates, or they are all relative to the position of the residue position?

If it is a single position that has some vector-valued input features associated to it, there is a way to handle it, but if it is really 4 independent points, then just putting 4 particles into the graph is the correct approach.

[tueboesen]

I think I understand what you are suggesting, but I also think there is a fundamental problem in categorizing an amino acid by just one coordinate.
The problem with this approach is that you end up building a graph based on "wrong" or inconsistent distances.
Most amino acids have a long sidechain attached to it (though not all), as seen here: https://sciencenotes.org/wp-content/uploads/2018/09/AminoAcidSideChains-791x1024.png

This means that the orientation of these amino acids are fundamentally important for determining whether two amino acids are close together. Two amino acids with a centroid placed 15 angstrom apart, might easily be able to affect each other if their orientation is such that the side chains point towards each other and if they are of the right type of amino acid.

I realize now, that this is a problem not just for how to do the spherical_sh, but also for building the graph in general, so I probably need to rework it there as well.

To put the problem into a framework that is more familiar to NequIP, you can think of each amino acid as a different block in tetris, and in order to know whether two tetris blocks are close to each other you not only need to know the centroid positions of the tetris blocks, but also the orientations and the types of tetris blocks.

[Linux-cpp-lisp]

Right; this makes sense.

I guess what I'm not fully understanding from your messages is your goal with this embedding: what do you hope to gain from avoiding the "all-atom" (4 "atoms"/residue) graph representation? This is of course the easiest to deal with, since it identically matches the normal all-atom case.

If the goal is speed, it's not obvious to me that exchanging 4 "atoms" and a smaller model for a heavier model + embedding steps + bigger cutoff will gain you much.

[tueboesen]

The goal is indeed to move from atoms to residues. There are primarily two reasons for that, one is performance (speed and memory reasons), while the second one is to enforce stability inside amino acids. The last reason might not be super important for my current use case, since I really just want to get the best possible prediction of the energies, but if I were to do actual MD simulations with amino acids, this could also be important since we have suddenly moved from a description of light atoms to amino acids, and hence you would expect to be able to use much larger time-steps, since the light hydrogen atoms are no longer directly modeled.

An amino acid can almost perfectly be described by 4 coordinates and its type, whereas if you were to do the same with atoms, then each amino acid would consist of somewhere between 10-27 atoms. So on average we are talking about a 5 time reduction when moving to residues rather than atoms.

But yes you are right if the cutoff becomes much bigger then it might not be worth it. But maybe I can avoid that if I define the distance more sensible (like the actual shortest distance between amino acids). It is a non-trivial problem for sure, that I will try and think some more of in the coming days.

[Linux-cpp-lisp]

I see— so this is a coarse graining, just from ~27 atoms to ~4 "interaction sites". That makes a great deal of sense. My recommendation would then be to just use NequIP as usual on the graph of "interaction sites" as your atoms; this is simpler and doesn't require model architecture modifications.

My understanding is that internal rigidity assumptions are usually enforced in these sorts of biochemical systems using constraints in MD, which are independent of the potential being used. E.g. (first MD code's docs I found): https://hoomd-blue.readthedocs.io/en/stable/module-md-constrain.html

[Linux-cpp-lisp]

Of course, in a well-learned force field, the PES itself should constrain the internal DOF of the residue— as I understand it, the main purpose of constraints is to get a bigger time step by restricting hydrogen vibrations in all-atom systems. If you are already coarse-graining away such fast DOFs from the geometric representation of your system, then constraints may not be necessary.

Of course, it's also important to mention that the "other" approach—a single center with vector input features— does not inherantly restrict the internal degrees of freedom in any way.