Implement ViSNet

.github/workflows/CI.yaml

@@ -22,13 +22,14 @@ defaults:
 
 jobs:
   test:
-    name: Test on ${{ matrix.os }}, Python ${{ matrix.python-version }}
+    name: Test on ${{ matrix.os }}, Python ${{ matrix.python-version }}, Env ${{ matrix.deps }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         os: [macOS-latest, ubuntu-latest]
         python-version: ["3.10", "3.11"]
+        deps: ["devtools/conda-envs/test_env.yaml", "devtools/conda-envs/test_env-nightly.yaml"]
 
     steps:
     - name: Checkout Repository
@@ -45,7 +46,7 @@ jobs:
     - name: Setup Micromamba
       uses: mamba-org/setup-micromamba@v1
       with:
-        environment-file: devtools/conda-envs/test_env.yaml
+        environment-file: ${{ matrix.deps }}
         environment-name: test
         create-args: >-
           python==${{ matrix.python-version }}
@@ -68,5 +69,5 @@ jobs:
       with:
         file: ./coverage.xml
         flags: unittests
-        name: codecov-${{ matrix.os }}-py${{ matrix.python-version }}
+        name: codecov-${{ matrix.os }}-py${{ matrix.python-version }}-env${{ matrix.deps }}
         fail_ci_if_error: false

devtools/conda-envs/test_env-nightly.yaml

@@ -0,0 +1,24 @@
+name: test
+channels:
+  - conda-forge
+dependencies:
+  - pip
+  - pytorch
+  - pytorch_cluster
+  - pytorch_scatter
+  - pytorch_sparse
+  - numpy
+  - h5py
+  - e3nn
+  - dgllife
+  - dgl
+  - rdkit
+  - ase
+  # testing dependencies
+  - pytest
+  - pytest-cov
+  - codecov
+  - pydantic >=1.10.8,<2.0.0a0
+
+  - pip:
+    - pyg-nightly

devtools/conda-envs/test_env.yaml

@@ -18,4 +18,4 @@ dependencies:
   - pytest
   - pytest-cov
   - codecov
-  - pydantic >=1.10.8,<2.0.0a0
+  - pydantic >=1.10.8,<2.0.0a0

mtenn/config.py

@@ -5,7 +5,6 @@
 from pydantic import BaseModel, Field, root_validator
 import random
 from typing import Callable, ClassVar
-
 import mtenn
 import numpy as np
 import torch
@@ -39,7 +38,7 @@ class ModelType(StringEnum):
     schnet = "schnet"
     e3nn = "e3nn"
     INVALID = "INVALID"
-
+    visnet = "visnet"
 
 class StrategyConfig(StringEnum):
     """
@@ -749,3 +748,126 @@ def _build(self, mtenn_params={}):
             pred_readout=pred_readout,
             comb_readout=comb_readout,
         )
+
+
+class ViSNetModelConfig(ModelConfigBase):
+    """
+    Class for constructing a VisNet ML model. Default values here are the default values
+    given in PyG.
+    """
+
+    model_type: ModelType = Field(ModelType.visnet, const=True)
+    lmax: int = Field(1, description="The maximum degree of the spherical harmonics.")
+    vecnorm_type: str | None = Field(
+        None, description="The type of normalization to apply to the vectors."
+    )
+    trainable_vecnorm: bool = Field(
+        False, description="Whether the normalization weights are trainable."
+    )
+    num_heads: int = Field(8, description="The number of attention heads.")
+    num_layers: int = Field(6, description="The number of layers in the network.")
+    hidden_channels: int = Field(
+        128, description="The number of hidden channels in the node embeddings."
+    )
+    num_rbf: int = Field(32, description="The number of radial basis functions.")
+    trainable_rbf: bool = Field(
+        False, description="Whether the radial basis function parameters are trainable."
+    )
+    max_z: int = Field(100, description="The maximum atomic numbers.")
+    cutoff: float = Field(5.0, description="The cutoff distance.")
+    max_num_neighbors: int = Field(
+        32, 
+        description="The maximum number of neighbors considered for each atom."
+    )
+    vertex: bool = Field(
+        False, 
+        description="Whether to use vertex geometric features."
+    )
+    atomref: list[float] | None = Field(
+        None,
+        description=(
+            "Reference values for single-atom properties. Should have length max_z"
+            )    
+    )
+    reduce_op: str = Field(
+        "sum", 
+        description="The type of reduction operation to apply. ['sum', 'mean']"
+    )
+    mean: float = Field(0.0, description="The mean of the output distribution.")
+    std: float = Field(1.0, description="The standard deviation of the output distribution.")
+    derivative: bool = Field(
+        False, 
+        description="Whether to compute the derivative of the output with respect to the positions."
+    )
+
+    @root_validator(pre=False)
+    def validate(cls, values):
+        # Make sure the grouped stuff is properly assigned
+        ModelConfigBase._check_grouped(values)
+
+        # Make sure atomref length is correct (this is required by PyG)
+        atomref = values["atomref"]
+        if (atomref is not None) and (len(atomref) != values["max_z"]):
+            raise ValueError(f"atomref length must match max_z. (Expected {values['max_z']}, got {len(atomref)})")
+
+        return values
+
+
+
+    def _build(self, mtenn_params={}):
+        """
+        Build an MTENN ViSNet Model from this config.
+
+        Parameters
+        ----------
+        mtenn_params: dict
+            Dict giving the MTENN Readout. This will be passed by the `build` method in
+            the abstract base class
+
+        Returns
+        -------
+        mtenn.model.Model
+            MTENN ViSNet Model/GroupedModel
+        """
+        # Create an MTENN ViSNet model from PyG ViSNet model
+
+        from mtenn.conversion_utils.visnet import HAS_VISNET
+        if HAS_VISNET:
+            from mtenn.conversion_utils import ViSNet
+
+            model = ViSNet(
+                lmax=self.lmax,
+                vecnorm_type=self.vecnorm_type,
+                trainable_vecnorm=self.trainable_vecnorm,
+                num_heads=self.num_heads,
+                num_layers=self.num_layers,
+                hidden_channels=self.hidden_channels,
+                num_rbf=self.num_rbf,
+                trainable_rbf=self.trainable_rbf,
+                max_z=self.max_z,
+                cutoff=self.cutoff,
+                max_num_neighbors=self.max_num_neighbors,
+                vertex=self.vertex,
+                reduce_op=self.reduce_op,
+                mean=self.mean,
+                std=self.std,
+                derivative=self.derivative,
+                atomref=self.atomref,
+            )
+            combination = mtenn_params.get("combination", None)
+            pred_readout = mtenn_params.get("pred_readout", None)
+            comb_readout = mtenn_params.get("comb_readout", None)
+
+            return ViSNet.get_model(
+                model=model,
+                grouped=self.grouped,
+                fix_device=True,
+                strategy=self.strategy,
+                combination=combination,
+                pred_readout=pred_readout,
+                comb_readout=comb_readout,
+            )
+
+        else:
+            raise ImportError("ViSNet not found. Is your PyG >=2.5.0? Refer to issue #42.")
+

mtenn/conversion_utils/__init__.py

@@ -2,4 +2,9 @@
 from .gat import GAT
 from .schnet import SchNet
 
-__all__ = ["E3NN", "GAT", "SchNet"]
+# refer to issue #42
+from .visnet import HAS_VISNET
+if HAS_VISNET:
+    from .visnet import ViSNet
+
+__all__ = ["E3NN", "GAT", "SchNet", "ViSNet"]