{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using custom NNPs with Auto3D\n",
    "\n",
    "Auto3D (>= 2.3.0) is compatible with any jitable NNPs. This notebook demonstrates how to wrapper and jit a custom NNP to a specific format that Auto3D can use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import os, sys\nroot = os.path.dirname(os.path.dirname(os.path.abspath(\"__file__\")))\n\nimport torch\nimport torchani\nimport Auto3D\nfrom Auto3D import Auto3DOptions, main\n\nprint(Auto3D.__version__)"
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Below is the template of the wrapper that you need to implement with your custom NNP.\n",
    "class userNNP(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(userNNP, self).__init__()\n",
    "        \"\"\"This is an example NNP model that can be used with Auto3D.\n",
    "        You can initialize an NNP model however you want,\n",
    "        just make sure that:\n",
    "            - It contains the coord_pad and species_pad attributes \n",
    "              (These values will be used when processing the molecules in batch.)\n",
    "            - The signature of the forward method is the same as below.\n",
    "        \"\"\"\n",
    "        # Here I constructed an example NNP using ANI2x.\n",
    "        # In your case, you can replace this with your own NNP model.\n",
    "        self.model = torchani.models.ANI2x(periodic_table_index=True)\n",
    "\n",
    "        self.coord_pad = 0  # int, the padding value for coordinates\n",
    "        self.species_pad = -1  # int, the padding value for species.\n",
    "\n",
    "    def forward(self,\n",
    "                species: torch.Tensor,\n",
    "                coords: torch.Tensor,\n",
    "                charges: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Your NNP should take species, coords, and charges as input\n",
    "        and return the energies of the molecules.\n",
    "\n",
    "        species contains the atomic numbers of the atoms in the molecule: [B, N]\n",
    "        where B is the batch size, N is the number of atoms in the largest molecule.\n",
    "        \n",
    "        coords contains the coordinates of the atoms in the molecule: [B, N, 3]\n",
    "        where B is the batch size, N is the number of atoms in the largest molecule,\n",
    "        and 3 represents the x, y, z coordinates.\n",
    "        \n",
    "        charges contains the molecular charges: [B]\n",
    "        \n",
    "        The forward function returns the energies of the molecules: [B],\n",
    "        output energy unit: eV\"\"\"\n",
    "\n",
    "        # an example for computing molecular energy, replace with your NNP model\n",
    "        energies = self.model((species, coords)).energies * 27.211386245988\n",
    "        return energies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/jack/miniconda3/envs/py39/lib/python3.9/site-packages/torchani/resources/\n"
     ]
    }
   ],
   "source": [
    "# initialize and jit the wrapper with your NNP model\n",
    "myNNP = userNNP()\n",
    "myNNP_jit = torch.jit.script(myNNP)\n",
    "\n",
    "# save the model to a file for later use\n",
    "model_path = os.path.join(root, 'myNNP.pt')\n",
    "torch.jit.save(myNNP_jit, model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# Now you can run Auto3D with your custom NNP model.\n# Simply pass the model_path to the optimizing_engine argument\n\nsmi_path = os.path.join(root, \"example/files/smiles.smi\")  # You can specify the path to your file here\nconfig = Auto3DOptions(path=smi_path, k=1, optimizing_engine=model_path, use_gpu=True, gpu_idx=0)\nout = main(config)\nprint(out)"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py39",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}