adjust requirements.txt installation and use hidden_dim for sizing mo…

…re effectively in InteractionBlock

.github/workflows/CI.yml

@@ -40,8 +40,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install --upgrade -r requirements.txt -r requirements-dev.txt
-          python -m pip install --upgrade -r requirements-torch.txt --index-url https://download.pytorch.org/whl/cpu
-          python -m pip install --upgrade -r requirements-torch2.txt
+          python -m pip install --upgrade -r requirements-torch.txt --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
           python -m pip install --upgrade -r requirements-pyg.txt --find-links https://data.pyg.org/whl/torch-2.0.1+cpu.html
           python -m pip install --upgrade -r requirements-deepspeed.txt
       - name: Format black

hydragnn/utils/model/mace_utils/modules/blocks.py

@@ -272,12 +272,19 @@ def __repr__(self):
         )
 
 
+###########################################################################################
+# NOTE: Below is one of the many possible Interaction Blocks in the MACE architecture.
+#       Since there are adaptations to the original code in order to be integrated with
+#       the HydraGNN framework, and the changes between blocks are relatively minor, we've
+#       elected to adapt one general-purpose block here. Users can access the other blocks
+#       and adapt similarly from the original MACE code (linked with date at the top).
+###########################################################################################
 @compile_mode("script")
 class RealAgnosticAttResidualInteractionBlock(InteractionBlock):
     def _setup(self) -> None:
         self.node_feats_down_irreps = o3.Irreps(
-            "64x0e"
-        )  # Interesting, this seems to be a required shaping
+            [(o3.Irreps(self.hidden_irreps).count(o3.Irrep(0, 1)), (0, 1))]
+        )
         # First linear
         self.linear_up = o3.Linear(
             self.node_feats_irreps,
@@ -313,9 +320,12 @@ def _setup(self) -> None:
         )
         # The following specifies the network architecture for embedding l=0 (scalar) irreps
         ## It is worth double-checking, but I believe this means that type 0 (scalar) irreps
-        ## are being embedded by 3 layers of size 256 and the output dim, then activated.
+        # are being embedded by 3 layers of size self.hidden_dim (scalar irreps) and the
+        # output dim, then activated.
         self.conv_tp_weights = nn.FullyConnectedNet(
-            [input_dim] + 3 * [256] + [self.conv_tp.weight_numel],
+            [input_dim]
+            + 3 * [o3.Irreps(self.hidden_irreps).count(o3.Irrep(0, 1))]
+            + [self.conv_tp.weight_numel],
             torch.nn.functional.silu,
         )
 

requirements-torch.txt

@@ -1,4 +1,6 @@
 torch==2.0.1 
 torchvision 
 torchaudio
-
+e3nn==0.5.1
+torch-ema==0.3
+torchmetrics==1.4.0