MACE

[RylieWeaver]

An implementation of the MACE Model in HYDRA.

There are 2 key aspects of the MACE Architecture:
(1) Message passing and interaction blocks are equivariant to the O(3) (rotation/inversion) group and invariant to the T(3) (translation) group.
(2) Predictions are made in an n-body expansion, where n is the number of layers. This is done by creating multi-body interactions, then decoding them. Layer 1 will decode 1-body interactions, layer 2 will decode w-body interactions, and so on. So, for a 3-layer model predicting energy, there are 3 outputs for energy, one at each layer, and they are summed at the end. This requires some adjustment to the behavior from Base.py. Specifically, there is a multihead decoder adapted from HYDRA after each convolutional layer.

[allaffa]

Some files may have been added or changes by mistake.
Please check and explain.

[RylieWeaver]

@allaffa As you can see, the MACE repo had a large collection of utils files. I'd like your opinion on how to handle these. On one hand, leaving too much can create clutter. On the other hand, taking things out could reduce customization options / induce unforeseen errors later on.

I think there is a middle-ground, where I could take out ~50% of the utils, focusing on dataset/distributed utils. In terms of impact, I estimate that it would only make a difference to a curious MACE user that uses an unusual argument, since it's hard to account for all these when cleaning out a big repo. In terms of time, it would take a day or two so still be ready by Monday to rebase.

[RylieWeaver]

@allaffa Additionally, despite having torch in their name, I have place torch-ema and torchmetrics in requirements.txt. I did this because the index-url for requirements-torch.txt was incompatible with those packages. Is there a more preferable way to do this, or this is fine?

[allaffa]

Second round of review completed.

[allaffa]

Just a coupled of questions about why DimeNet++ is missing from some tests

[RylieWeaver]

@allaffa

There is something important to note wrt blocks.py in MACE:

Many of these blocks need to be changed in order to adapt to HYDRA (usually shaping and variable passing). Importantly, many of these "blocks" are similar, with small tweaks. For example:

    RealAgnosticAttResidualInteractionBlock,
    RealAgnosticInteractionBlock,
    RealAgnosticResidualInteractionBlock,
    ResidualElementDependentInteractionBlock,

The differences between these blocks are a residual connection or not, and similarly small adjustments. Given the similarities and in order to make it more tractable, I chose to adjust and implement one of these blocks, which I discern to be generally applicable and effective. For example, I chose RealAgnosticAttResidualInteractionBlock out of those 4.

TLDR: This should be generally effective to choose one block and move forward, however is important to note down to memory in case an impassioned MACE user comes around with a very specific choice of architecture.

[RylieWeaver]

Ready for next round of review

[allaffa]

@allaffa

Getting closer. :)

Yep :). Here’s a summary update for posterity, which I’ll also cover in the meeting:

1. Replacing torch_scatter operations with an import worked fine. The original implementation avoided it for dependency reasons, but we’re good as long as we can install torch_scatter.
2. I moved irreps_tools.py to the utils/model directory for broader applicability and created operations.py for a general function to get edge_vectors_and_lengths. We might revisit the naming or location.
3. I couldn’t find anything similar in Nequip to generalize further. For example, grep searches for "wigner," "U_matrix," and "Contraction" returned nothing. Nequip's InteractionBlock is also different.
4. In compile.py, I left all compilation-related functions to ensure speed and accuracy, even though some can be commented out without issue.

@RylieWeaver
Since the functionalities incompile.py can be commented out without issue and are provided only for the sake of speed-up, I think it would be beneficial to add a short description about this at the beginning of the file.

[allaffa]

I left a few more comments.

[allaffa]

Another round completed.

[allaffa]

@RylieWeaver
You MUST absolutely get rid of this comments once for all.
"temporarily" commenting things out like this is extremely dangerous.

[allaffa]

@RylieWeaver
Uncomment what needs to be kept, and remove everything else.

[allaffa]

Another round of review

[RylieWeaver]

@allaffa

Ready for next round. All tests should be running now.

[allaffa]

@RylieWeaver
Conflicts left to resolve

[RylieWeaver]

@RylieWeaver Conflicts left to resolve

Oh, I think it involves the typo fix just now. The struggle of multiple PRs, lol

[allaffa]

I left just a question about the torch seed set to 97.
For the rest, it looks good to me.

[RylieWeaver]

I left just a question about the torch seed set to 97. For the rest, it looks good to me.

Shall we have Pei take a look for a final check with some fresh eyes?