{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from huggingface_hub import hf_hub_download\n",
    "from ase.calculators.calculator import Calculator\n",
    "# from mlip_arena.models import MLIP, MLIPCalculator, ModuleMLIP\n",
    "\n",
    "from mlip_arena.models.externals import MACE_MP_Medium\n",
    "\n",
    "from mlip_arena.models.utils import MLIPMap, MLIPEnum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "issubclass(MLIPEnum[\"MACE-MP(M)\"].value, Calculator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isinstance(MLIPEnum[\"MACE-MP(M)\"], MLIPEnum)# in MLIPEnum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['MACE-MP(M)', 'CHGNet', 'EquiformerV2(OC22)', 'eSCN(OC20)']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "MLIPEnum._member_names_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'MACE-MP(M)': mlip_arena.models.externals.MACE_MP_Medium,\n",
       " 'CHGNet': mlip_arena.models.externals.CHGNet,\n",
       " 'EquiformerV2(OC22)': mlip_arena.models.externals.EquiformerV2,\n",
       " 'eSCN(OC20)': mlip_arena.models.externals.eSCN}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "MLIPMap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLIPEnum.MACE-MP(M)\n",
      "MLIPEnum.CHGNet\n",
      "MLIPEnum.EquiformerV2(OC22)\n",
      "MLIPEnum.eSCN(OC20)\n"
     ]
    }
   ],
   "source": [
    "for mlip in MLIPEnum:\n",
    "    print(mlip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "mlip_arena.models.externals.MACE_MP_Medium"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "MLIPMap['MACE-MP(M)']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using Materials Project MACE for MACECalculator with /global/homes/c/cyrusyc/.cache/mace/5yyxdm76\n",
      "Selected GPU cuda:0 with 40338.06 MB free memory from 1 GPUs\n",
      "Default dtype float32 does not match model dtype float64, converting models to float32.\n"
     ]
    }
   ],
   "source": [
    "mace_mp = MACE_MP_Medium()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Select GPU cuda:0 with 40316.98 MB free memory from 1 GPUs\n",
      "CHGNet v0.3.0 initialized with 412,525 parameters\n",
      "CHGNet will run on cuda:0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Detected old config, converting to new format. Consider updating to avoid potential incompatibilities.\n",
      "WARNING:root:Skipping scheduler setup. No training set found.\n"
     ]
    }
   ],
   "source": [
    "from mlip_arena.models.externals import EquiformerV2, CHGNet\n",
    "\n",
    "chgnet = CHGNet()\n",
    "\n",
    "equiformer_v2 = EquiformerV2()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "fpath = hf_hub_download(\n",
    "    repo_id=\"cyrusyc/mace-universal\",\n",
    "    subfolder=\"pretrained\",\n",
    "    filename=\"2023-12-12-mace-128-L1_epoch-199.model\",\n",
    "    revision=None,  # TODO: Add revision\n",
    ")\n",
    "\n",
    "model = torch.load(fpath, map_location=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "module = ModuleMLIP(model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "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": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "module.save_pretrained(\n",
    "    \"mace\",\n",
    "    repo_id=\"atomind/MACE_MP_Medium\".lower().replace(\"_\", \"-\"),\n",
    "    push_to_hub=True\n",
    ")"
   ]
  },
  {
   "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,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}