updates before modification

.gitignore

@@ -2,6 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
+tests/
 
 # C extensions
 *.so

mlip_arena/models/__init__.py

@@ -12,11 +12,6 @@ from torch_geometric.data import Data
 with open(os.path.join(os.path.dirname(__file__), "registry.yaml")) as f:
     REGISTRY = yaml.load(f, Loader=yaml.FullLoader)
 
-# class MLIPEnum(enum.Enum):
-#     for model, metadata in REGISTRY.items():
-#         model_class = getattr(importlib.import_module(model["module"]), model)
-#         self.setattr(model, model_class)
-
 
 class MLIP(
     nn.Module,
@@ -30,7 +25,7 @@ class MLIP(
 class ModuleMLIP(MLIP):
     def __init__(self, model: nn.Module, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-        self.register_module("model", model)
+        self.add_module("model", model)
 
     def forward(self, x):
         print("Forwarding...")
@@ -41,15 +36,12 @@ class ModuleMLIP(MLIP):
 
 class MLIPCalculator(Calculator):
     name: str
-    device: torch.device
-    model: MLIP
+    # device: torch.device
+    # model: MLIP
     implemented_properties: list[str] = ["energy", "forces", "stress"]
 
     def __init__(
         self,
-        # PyTorch
-        model_path: str | Path,
-        device: torch.device | None = None,
         # ASE Calculator
         restart=None,
         atoms=None,

mlip_arena/models/mace.py

@@ -1,3 +1,6 @@
+from typing import Optional, Tuple
+
+import numpy as np
 import torch
 from ase import Atoms
 from ase.calculators.calculator import all_changes
@@ -16,48 +19,33 @@ class MACE_MP_Medium(MLIPCalculator):
         directory=".",
         **kwargs,
     ):
-        # Download the pytorch model from huggingface to local and load it
-        # NOTE: this is not the ideal way to load the model, but it is the simplest
-        # way to do it for now. Ideally, if the model is the subclass of PyTorchModelHubMixin,
-        # we should be able to load it directly from the hub or local using MLIP class.
+        super().__init__(restart=restart, atoms=atoms, directory=directory, **kwargs)
+
+        self.name: str = self.__class__.__name__
+
         fpath = hf_hub_download(
             repo_id="cyrusyc/mace-universal",
             subfolder="pretrained",
             filename="2023-12-12-mace-128-L1_epoch-199.model",
-            revision=None,  # TODO: Add revision
+            revision="main",
         )
-        # module = ModuleMLIP(torch.load(fpath, map_location="cpu"))
-        print(torch.load(fpath, map_location="cpu"))
-        repo_id = f"atomind/{self.__class__.__name__}".replace("_", "-")
-        # module.save_pretrained(
-        #     save_directory=self.__class__.__name__,
-        #     repo_id=repo_id,
-        #     push_to_hub=True,
-        # )
 
-        super().__init__(
-            model_path=repo_id,
-            device=device,
-            restart=restart,
-            atoms=atoms,
-            directory=directory,
-            **kwargs,
+        self.device = device or torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
         )
 
-        # self.name: str = self.__class__.__name__
-        # self.device = device or torch.device(
-        #     "cuda" if torch.cuda.is_available() else "cpu"
-        # )
-        # self.model: MLIP = ModuleMLIP.from_pretrained(repo_id, map_location=self.device)
-        # self.implemented_properties = ["energy", "forces", "stress"]
+        self.model = torch.load(fpath, map_location=self.device)
+
+        self.implemented_properties = ["energy", "forces", "stress"]
+
+        # repo_id = f"atomind/{self.__class__.__name__}".lower().replace("_", "-")
 
-        self.display = "MACE-MP-0 (medium)"
-        self.version = "1.0.0"
-        self.implemented_properties = [
-            "energy",
-            "forces",
-            "stress",
-        ]
+        # model = ModuleMLIP(model=model)
+        # model.save_pretrained(
+        #     self.__class__.__name__.lower().replace("_", "-"),
+        #     repo_id=repo_id,
+        #     push_to_hub=True,
+        # )
 
     def calculate(
         self, atoms: Atoms, properties: list[str], system_changes: list = all_changes

mlip_arena/models/utils.py

@@ -1,15 +1,91 @@
+"""Utility functions for MLIP models."""
+
 import importlib
-import os
 from enum import Enum
+from typing import Any
+
+import torch
+from ase.calculators.calculator import Calculator
+from ase.calculators.mixing import SumCalculator
+from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
 
 from mlip_arena.models import REGISTRY
 
-MLIPEnum = Enum(
-    "MLIPEnum",
-    {
-        model: getattr(
-            importlib.import_module(f"{__package__}.{metadata['module']}"), model
-        )
-        for model, metadata in REGISTRY.items()
-    },
-)
+MLIPMap = {
+    model: getattr(
+        importlib.import_module(f"{__package__}.{metadata['module']}"), model
+    )
+    for model, metadata in REGISTRY.items()
+}
+
+
+class EXTMLIPEnum(Enum):
+    """
+    Enumeration class for EXTMLIP models.
+
+    Attributes:
+        M3GNet (str): M3GNet model.
+        CHGNet (str): CHGNet model.
+        MACE (str): MACE model.
+    """
+
+    M3GNet = "M3GNet"
+    CHGNet = "CHGNet"
+    MACE = "MACE"
+
+
+def get_freer_device() -> torch.device:
+    """Get the GPU with the most free memory.
+
+    Returns:
+        torch.device: The selected GPU device.
+
+    Raises:
+        ValueError: If no GPU is available.
+    """
+    device_count = torch.cuda.device_count()
+    if device_count == 0:
+        print("No GPU available. Using CPU.")
+        return torch.device("cpu")
+
+    mem_free = [
+        torch.cuda.get_device_properties(i).total_memory
+        - torch.cuda.memory_allocated(i)
+        for i in range(device_count)
+    ]
+
+    free_gpu_index = mem_free.index(max(mem_free))
+
+    print(
+        f"Selected GPU {free_gpu_index} with {mem_free[free_gpu_index] / 1024**2:.2f} MB free memory from {device_count} GPUs"
+    )
+
+    return torch.device(f"cuda:{free_gpu_index}")
+
+
+def external_ase_calculator(name: EXTMLIPEnum, **kwargs: Any) -> Calculator:
+    """Construct an ASE calculator from an external third-party MLIP packages"""
+
+    calculator = None
+    device = get_freer_device()
+
+    if name == EXTMLIPEnum.MACE:
+        from mace.calculators import mace_mp
+
+        calculator = mace_mp(device=str(device), **kwargs)
+
+    elif name == EXTMLIPEnum.CHGNet:
+        from chgnet.model.dynamics import CHGNetCalculator
+
+        calculator = CHGNetCalculator(use_device=str(device), **kwargs)
+
+    elif name == EXTMLIPEnum.M3GNet:
+        import matgl
+        from matgl.ext.ase import PESCalculator
+
+        potential = matgl.load_model("M3GNet-MP-2021.2.8-PES")
+        calculator = PESCalculator(potential, **kwargs)
+
+    calculator.__setattr__("name", name.value)
+
+    return calculator

mlip_arena/tasks/diatomics.py

@@ -1,123 +0,0 @@
-import covalent as ct
-import numpy as np
-import pandas as pd
-import torch
-from ase import Atoms
-from ase.calculators.calculator import Calculator
-from ase.data import chemical_symbols
-from matplotlib import pyplot as plt
-
-from mlip_arena.models import MLIPCalculator
-
-device = torch.device("cuda")
-
-
-@ct.electron
-def calculate_single_diatomic(
-    calculator: MLIPCalculator | Calculator,
-    atom1: str,
-    atom2: str,
-    rmin: float = 0.1,
-    rmax: float = 6.5,
-    npts: int = int(1e3),
-):
-    a = 2 * rmax
-
-    rs = np.linspace(rmin, rmax, npts)
-    e = np.zeros_like(rs)
-    f = np.zeros_like(rs)
-
-    da = atom1 + atom2
-
-    for i, r in enumerate(rs):
-
-        positions = [
-            [0, 0, 0],
-            [r, 0, 0],
-        ]
-
-        # Create the unit cell with two atoms
-        atoms = Atoms(da, positions=positions, cell=[a, a, a])
-
-        atoms.calc = calculator
-
-        e[i] = atoms.get_potential_energy()
-        f[i] = np.inner(np.array([1, 0, 0]), atoms.get_forces()[1])
-
-    return rs, e, f, da
-
-
-@ct.lattice
-def calculate_homonuclear_diatomics(calculator: MLIPCalculator | Calculator):
-
-    chemical_symbols.remove("X")
-
-    results = {}
-
-    for atom in chemical_symbols:
-        rs, e, f, da = calculate_single_diatomic(calculator, atom, atom)
-        results[da] = {"r": rs, "E": e, "F": f}
-
-    return results
-
-
-# with plt.style.context("default"):
-
-#     SMALL_SIZE = 6
-#     MEDIUM_SIZE = 8
-#     LARGE_SIZE = 10
-
-#     LINE_WIDTH = 1
-
-#     plt.rcParams.update(
-#         {
-#             "pgf.texsystem": "pdflatex",
-#             "font.family": "sans-serif",
-#             "text.usetex": True,
-#             "pgf.rcfonts": True,
-#             "figure.constrained_layout.use": True,
-#             "axes.labelsize": MEDIUM_SIZE,
-#             "axes.titlesize": MEDIUM_SIZE,
-#             "legend.frameon": False,
-#             "legend.fontsize": MEDIUM_SIZE,
-#             "legend.loc": "best",
-#             "lines.linewidth": LINE_WIDTH,
-#             "xtick.labelsize": SMALL_SIZE,
-#             "ytick.labelsize": SMALL_SIZE,
-#         }
-#     )
-
-#     fig, ax = plt.subplots(layout="constrained", figsize=(3, 2), dpi=300)
-
-#     color = "tab:red"
-#     ax.plot(rs, e, color=color, zorder=1)
-
-#     ax.axhline(ls="--", color=color, alpha=0.5, lw=0.5 * LINE_WIDTH)
-
-#     ylo, yhi = ax.get_ylim()
-#     ax.set(xlabel=r"r [$\AA]$", ylim=(max(-7, ylo), min(5, yhi)))
-#     ax.set_ylabel(ylabel="E [eV]", color=color)
-#     ax.tick_params(axis="y", labelcolor=color)
-#     ax.text(0.8, 0.85, da, fontsize=LARGE_SIZE, transform=ax.transAxes)
-
-#     color = "tab:blue"
-
-#     at = ax.twinx()
-#     at.plot(rs, f, color=color, zorder=0, lw=0.5 * LINE_WIDTH)
-
-#     at.axhline(ls="--", color=color, alpha=0.5, lw=0.5 * LINE_WIDTH)
-
-#     ylo, yhi = at.get_ylim()
-#     at.set(
-#         xlabel=r"r [$\AA]$",
-#         ylim=(max(-20, ylo), min(20, yhi)),
-#     )
-#     at.set_ylabel(ylabel="F [eV/$\AA$]", color=color)
-#     at.tick_params(axis="y", labelcolor=color)
-
-#     plt.show()
-
-
-if __name__ == "__main__":
-
-    local = ct.executor.LocalExecutor()

tests/hf_hub.ipynb

@@ -2,9 +2,18 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 1,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
    "source": [
     "import torch\n",
     "from huggingface_hub import hf_hub_download\n",
@@ -13,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -30,7 +39,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -39,23 +48,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "model.safetensors: 100%|██████████| 44.2M/44.2M [00:02<00:00, 20.6MB/s]\n"
-     ]
-    },
     {
      "data": {
       "text/plain": [
-       "CommitInfo(commit_url='https://huggingface.co/atomind/mace-mp-medium/commit/ef94f6bd9c7167bb28d594d5a3e7a5fbfbda2acb', commit_message='Push model using huggingface_hub.', commit_description='', oid='ef94f6bd9c7167bb28d594d5a3e7a5fbfbda2acb', pr_url=None, pr_revision=None, pr_num=None)"
+       "CommitInfo(commit_url='https://huggingface.co/atomind/mace-mp-medium/commit/eb12c5387b9e655d83a4e2e10c0f0779c3745227', commit_message='Push model using huggingface_hub.', commit_description='', oid='eb12c5387b9e655d83a4e2e10c0f0779c3745227', pr_url=None, pr_revision=None, pr_num=None)"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -68,6 +70,214 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "from mlip_arena.models.mace import MACE_MP_Medium\n",
+    "import torch\n",
+    "\n",
+    "calc = MACE_MP_Medium(device=torch.device(\"cuda\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ScaleShiftMACE(\n",
+       "  (node_embedding): LinearNodeEmbeddingBlock(\n",
+       "    (linear): Linear(89x0e -> 128x0e | 11392 weights)\n",
+       "  )\n",
+       "  (radial_embedding): RadialEmbeddingBlock(\n",
+       "    (bessel_fn): BesselBasis(r_max=6.0, num_basis=10, trainable=False)\n",
+       "    (cutoff_fn): PolynomialCutoff(p=5.0, r_max=6.0)\n",
+       "  )\n",
+       "  (spherical_harmonics): SphericalHarmonics()\n",
+       "  (atomic_energies_fn): AtomicEnergiesBlock(energies=[-3.6672, -1.3321, -3.4821, -4.7367, -7.7249, -8.4056, -7.3601, -7.2846, -4.8965, 0.0000, -2.7594, -2.8140, -4.8469, -7.6948, -6.9633, -4.6726, -2.8117, -0.0626, -2.6176, -5.3905, -7.8858, -10.2684, -8.6651, -9.2331, -8.3050, -7.0490, -5.5774, -5.1727, -3.2521, -1.2902, -3.5271, -4.7085, -3.9765, -3.8862, -2.5185, 6.7669, -2.5635, -4.9380, -10.1498, -11.8469, -12.1389, -8.7917, -8.7869, -7.7809, -6.8500, -4.8910, -2.0634, -0.6396, -2.7887, -3.8186, -3.5871, -2.8804, -1.6356, 9.8467, -2.7653, -4.9910, -8.9337, -8.7356, -8.0190, -8.2515, -7.5917, -8.1697, -13.5927, -18.5175, -7.6474, -8.1230, -7.6078, -6.8503, -7.8269, -3.5848, -7.4554, -12.7963, -14.1081, -9.3549, -11.3875, -9.6219, -7.3244, -5.3047, -2.3801, 0.2495, -2.3240, -3.7300, -3.4388, -5.0629, -11.0246, -12.2656, -13.8556, -14.9331, -15.2828])\n",
+       "  (interactions): ModuleList(\n",
+       "    (0): RealAgnosticResidualInteractionBlock(\n",
+       "      (linear_up): Linear(128x0e -> 128x0e | 16384 weights)\n",
+       "      (conv_tp): TensorProduct(128x0e x 1x0e+1x1o+1x2e+1x3o -> 128x0e+128x1o+128x2e+128x3o | 512 paths | 512 weights)\n",
+       "      (conv_tp_weights): FullyConnectedNet[10, 64, 64, 64, 512]\n",
+       "      (linear): Linear(128x0e+128x1o+128x2e+128x3o -> 128x0e+128x1o+128x2e+128x3o | 65536 weights)\n",
+       "      (skip_tp): FullyConnectedTensorProduct(128x0e x 89x0e -> 128x0e+128x1o | 1458176 paths | 1458176 weights)\n",
+       "      (reshape): reshape_irreps()\n",
+       "    )\n",
+       "    (1): RealAgnosticResidualInteractionBlock(\n",
+       "      (linear_up): Linear(128x0e+128x1o -> 128x0e+128x1o | 32768 weights)\n",
+       "      (conv_tp): TensorProduct(128x0e+128x1o x 1x0e+1x1o+1x2e+1x3o -> 256x0e+384x1o+384x2e+256x3o | 1280 paths | 1280 weights)\n",
+       "      (conv_tp_weights): FullyConnectedNet[10, 64, 64, 64, 1280]\n",
+       "      (linear): Linear(256x0e+384x1o+384x2e+256x3o -> 128x0e+128x1o+128x2e+128x3o | 163840 weights)\n",
+       "      (skip_tp): FullyConnectedTensorProduct(128x0e+128x1o x 89x0e -> 128x0e | 1458176 paths | 1458176 weights)\n",
+       "      (reshape): reshape_irreps()\n",
+       "    )\n",
+       "  )\n",
+       "  (products): ModuleList(\n",
+       "    (0): EquivariantProductBasisBlock(\n",
+       "      (symmetric_contractions): SymmetricContraction(\n",
+       "        (contractions): ModuleList(\n",
+       "          (0): Contraction(\n",
+       "            (contractions_weighting): ModuleList(\n",
+       "              (0-1): 2 x GraphModule()\n",
+       "            )\n",
+       "            (contractions_features): ModuleList(\n",
+       "              (0-1): 2 x GraphModule()\n",
+       "            )\n",
+       "            (weights): ParameterList(\n",
+       "                (0): Parameter containing: [torch.float64 of size 89x4x128 (cuda:0)]\n",
+       "                (1): Parameter containing: [torch.float64 of size 89x1x128 (cuda:0)]\n",
+       "            )\n",
+       "            (graph_opt_main): GraphModule()\n",
+       "          )\n",
+       "          (1): Contraction(\n",
+       "            (contractions_weighting): ModuleList(\n",
+       "              (0-1): 2 x GraphModule()\n",
+       "            )\n",
+       "            (contractions_features): ModuleList(\n",
+       "              (0-1): 2 x GraphModule()\n",
+       "            )\n",
+       "            (weights): ParameterList(\n",
+       "                (0): Parameter containing: [torch.float64 of size 89x6x128 (cuda:0)]\n",
+       "                (1): Parameter containing: [torch.float64 of size 89x1x128 (cuda:0)]\n",
+       "            )\n",
+       "            (graph_opt_main): GraphModule()\n",
+       "          )\n",
+       "        )\n",
+       "      )\n",
+       "      (linear): Linear(128x0e+128x1o -> 128x0e+128x1o | 32768 weights)\n",
+       "    )\n",
+       "    (1): EquivariantProductBasisBlock(\n",
+       "      (symmetric_contractions): SymmetricContraction(\n",
+       "        (contractions): ModuleList(\n",
+       "          (0): Contraction(\n",
+       "            (contractions_weighting): ModuleList(\n",
+       "              (0-1): 2 x GraphModule()\n",
+       "            )\n",
+       "            (contractions_features): ModuleList(\n",
+       "              (0-1): 2 x GraphModule()\n",
+       "            )\n",
+       "            (weights): ParameterList(\n",
+       "                (0): Parameter containing: [torch.float64 of size 89x4x128 (cuda:0)]\n",
+       "                (1): Parameter containing: [torch.float64 of size 89x1x128 (cuda:0)]\n",
+       "            )\n",
+       "            (graph_opt_main): GraphModule()\n",
+       "          )\n",
+       "        )\n",
+       "      )\n",
+       "      (linear): Linear(128x0e -> 128x0e | 16384 weights)\n",
+       "    )\n",
+       "  )\n",
+       "  (readouts): ModuleList(\n",
+       "    (0): LinearReadoutBlock(\n",
+       "      (linear): Linear(128x0e+128x1o -> 1x0e | 128 weights)\n",
+       "    )\n",
+       "    (1): NonLinearReadoutBlock(\n",
+       "      (linear_1): Linear(128x0e -> 16x0e | 2048 weights)\n",
+       "      (non_linearity): Activation [x] (16x0e -> 16x0e)\n",
+       "      (linear_2): Linear(16x0e -> 1x0e | 16 weights)\n",
+       "    )\n",
+       "  )\n",
+       "  (scale_shift): ScaleShiftBlock(scale=0.804154, shift=0.164097)\n",
+       ")"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "calc.model\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mlip_arena.models import MLIP\n",
+    "\n",
+    "model = MLIP.from_pretrained(\"atomind/mace-mp-medium\", map_location=\"cuda\", revision=\"main\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<generator object Module.modules at 0x7ff33915f920>"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.modules()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "AttributeError",
+     "evalue": "MLIP has no attribute `model`",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[8], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_submodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmodel\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/nn/modules/module.py:681\u001b[0m, in \u001b[0;36mModule.get_submodule\u001b[0;34m(self, target)\u001b[0m\n\u001b[1;32m    678\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m item \u001b[38;5;129;01min\u001b[39;00m atoms:\n\u001b[1;32m    680\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(mod, item):\n\u001b[0;32m--> 681\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m(mod\u001b[38;5;241m.\u001b[39m_get_name() \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m has no \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    682\u001b[0m                              \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mattribute `\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m item \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    684\u001b[0m     mod \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(mod, item)\n\u001b[1;32m    686\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(mod, torch\u001b[38;5;241m.\u001b[39mnn\u001b[38;5;241m.\u001b[39mModule):\n",
+      "\u001b[0;31mAttributeError\u001b[0m: MLIP has no attribute `model`"
+     ]
+    }
+   ],
+   "source": [
+    "model.get_submodule(\"model\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for name, param in model.named_parameters():\n",
+    "    print(name, param.data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(module)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,