includes forcenet embeddings

ocpmodels/common/utils.py

@@ -211,7 +211,7 @@ def setup_imports():
     datasets_folder = os.path.join(root_folder, "datasets")
     datasets_pattern = os.path.join(datasets_folder, "**", "*.py")
     model_folder = os.path.join(root_folder, "models")
-    model_pattern = os.path.join(model_folder, "**", "*.py")
+    model_pattern = os.path.join(model_folder, "*.py")
 
     importlib.import_module("ocpmodels.common.meter")
 

ocpmodels/models/cgcnn.py

@@ -7,12 +7,13 @@
 
 import torch
 import torch.nn as nn
+from torch_geometric.nn import MessagePassing, global_mean_pool, radius_graph
+from torch_geometric.nn.models.schnet import GaussianSmearing
+
 from ocpmodels.common.registry import registry
 from ocpmodels.common.utils import get_pbc_distances, radius_graph_pbc
-from ocpmodels.datasets.elemental_embeddings import EMBEDDINGS
 from ocpmodels.models.base import BaseModel
-from torch_geometric.nn import MessagePassing, global_mean_pool, radius_graph
-from torch_geometric.nn.models.schnet import GaussianSmearing
+from ocpmodels.models.utils.embeddings import EMBEDDINGS
 
 
 @registry.register_model("cgcnn")

ocpmodels/models/utils/embeddings/__init__.py

@@ -0,0 +1,9 @@
+__all__ = [
+    "ATOMIC_RADII",
+    "FORCENET_EMBEDDINGS",
+    "EMBEDDINGS",
+]
+
+from .atomic_radii import ATOMIC_RADII
+from .elemental_embeddings import EMBEDDINGS
+from .forcenet_embeddings import FORCENET_EMBEDDINGS

ocpmodels/models/utils/embeddings/atomic_radii.py

@@ -0,0 +1,108 @@
+"""
+Atomic radii in picometers
+
+NaN stored for unavailable parameters.
+"""
+ATOMIC_RADII = {
+    0: float("NaN"),
+    1: 25.0,
+    2: 120.0,
+    3: 145.0,
+    4: 105.0,
+    5: 85.0,
+    6: 70.0,
+    7: 65.0,
+    8: 60.0,
+    9: 50.0,
+    10: 160.0,
+    11: 180.0,
+    12: 150.0,
+    13: 125.0,
+    14: 110.0,
+    15: 100.0,
+    16: 100.0,
+    17: 100.0,
+    18: 71.0,
+    19: 220.0,
+    20: 180.0,
+    21: 160.0,
+    22: 140.0,
+    23: 135.0,
+    24: 140.0,
+    25: 140.0,
+    26: 140.0,
+    27: 135.0,
+    28: 135.0,
+    29: 135.0,
+    30: 135.0,
+    31: 130.0,
+    32: 125.0,
+    33: 115.0,
+    34: 115.0,
+    35: 115.0,
+    36: float("NaN"),
+    37: 235.0,
+    38: 200.0,
+    39: 180.0,
+    40: 155.0,
+    41: 145.0,
+    42: 145.0,
+    43: 135.0,
+    44: 130.0,
+    45: 135.0,
+    46: 140.0,
+    47: 160.0,
+    48: 155.0,
+    49: 155.0,
+    50: 145.0,
+    51: 145.0,
+    52: 140.0,
+    53: 140.0,
+    54: float("NaN"),
+    55: 260.0,
+    56: 215.0,
+    57: 195.0,
+    58: 185.0,
+    59: 185.0,
+    60: 185.0,
+    61: 185.0,
+    62: 185.0,
+    63: 185.0,
+    64: 180.0,
+    65: 175.0,
+    66: 175.0,
+    67: 175.0,
+    68: 175.0,
+    69: 175.0,
+    70: 175.0,
+    71: 175.0,
+    72: 155.0,
+    73: 145.0,
+    74: 135.0,
+    75: 135.0,
+    76: 130.0,
+    77: 135.0,
+    78: 135.0,
+    79: 135.0,
+    80: 150.0,
+    81: 190.0,
+    82: 180.0,
+    83: 160.0,
+    84: 190.0,
+    85: float("NaN"),
+    86: float("NaN"),
+    87: float("NaN"),
+    88: 215.0,
+    89: 195.0,
+    90: 180.0,
+    91: 180.0,
+    92: 175.0,
+    93: 175.0,
+    94: 175.0,
+    95: 175.0,
+    96: float("NaN"),
+    97: float("NaN"),
+    98: float("NaN"),
+    99: float("NaN"),
+    100: float("NaN"),
+}

ocpmodels/datasets/elemental_embeddings.py → ocpmodels/models/utils/embeddings/elemental_embeddings.py

@@ -3,6 +3,9 @@
 
 This source code is licensed under the MIT license found in the
 LICENSE file in the root directory of this source tree.
+
+
+Original CGCNN k-hot elemental embeddings.
 """
 
 EMBEDDINGS = {